Coverage for HARK / interpolation.py: 96%
1541 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-07 05:16 +0000
1"""
2Custom interpolation methods for representing approximations to functions.
3It also includes wrapper classes to enforce standard methods across classes.
4Each interpolation class must have a distance() method that compares itself to
5another instance; this is used in HARK.core's solve() method to check for solution
6convergence. The interpolator classes currently in this module inherit their
7distance method from MetricObject.
8"""
10import warnings
11from copy import deepcopy
13import numpy as np
14from scipy.interpolate import CubicHermiteSpline
15from HARK.metric import MetricObject
16from HARK.rewards import CRRAutility, CRRAutilityP, CRRAutilityPP
17from numba import njit
20def _isscalar(x):
21 """
22 Check whether x is if a scalar type, or 0-dim.
24 Parameters
25 ----------
26 x : anything
27 An input to be checked for scalar-ness.
29 Returns
30 -------
31 is_scalar : boolean
32 True if the input is a scalar, False otherwise.
33 """
34 return np.isscalar(x) or hasattr(x, "shape") and x.shape == ()
37def _check_grid_dimensions(dimension, *args):
38 if dimension == 1:
39 if len(args[0]) != len(args[1]):
40 raise ValueError("Grid dimensions of x and f(x) do not match")
41 elif dimension == 2:
42 if args[0].shape != (args[1].size, args[2].size):
43 raise ValueError("Grid dimensions of x, y and f(x, y) do not match")
44 elif dimension == 3:
45 if args[0].shape != (args[1].size, args[2].size, args[3].size):
46 raise ValueError("Grid dimensions of x, y, z and f(x, y, z) do not match")
47 elif dimension == 4:
48 if args[0].shape != (args[1].size, args[2].size, args[3].size, args[4].size):
49 raise ValueError("Grid dimensions of x, y, z and f(x, y, z) do not match")
50 else:
51 raise ValueError("Dimension should be between 1 and 4 inclusive.")
54def _check_flatten(dimension, *args):
55 if dimension == 1:
56 if isinstance(args[0], np.ndarray) and args[0].shape != args[0].flatten().shape:
57 warnings.warn("input not of the size (n, ), attempting to flatten")
58 return False
59 else:
60 return True
63class HARKinterpolator1D(MetricObject):
64 """
65 A wrapper class for 1D interpolation methods in HARK.
66 """
68 distance_criteria = []
70 def __call__(self, x):
71 """
72 Evaluates the interpolated function at the given input.
74 Parameters
75 ----------
76 x : np.array or float
77 Real values to be evaluated in the interpolated function.
79 Returns
80 -------
81 y : np.array or float
82 The interpolated function evaluated at x: y = f(x), with the same
83 shape as x.
84 """
85 z = np.asarray(x)
86 return (self._evaluate(z.flatten())).reshape(z.shape)
88 def derivative(self, x):
89 """
90 Evaluates the derivative of the interpolated function at the given input.
92 Parameters
93 ----------
94 x : np.array or float
95 Real values to be evaluated in the interpolated function.
97 Returns
98 -------
99 dydx : np.array or float
100 The interpolated function's first derivative evaluated at x:
101 dydx = f'(x), with the same shape as x.
102 """
103 z = np.asarray(x)
104 return (self._der(z.flatten())).reshape(z.shape)
106 derivativeX = derivative # alias
108 def eval_with_derivative(self, x):
109 """
110 Evaluates the interpolated function and its derivative at the given input.
112 Parameters
113 ----------
114 x : np.array or float
115 Real values to be evaluated in the interpolated function.
117 Returns
118 -------
119 y : np.array or float
120 The interpolated function evaluated at x: y = f(x), with the same
121 shape as x.
122 dydx : np.array or float
123 The interpolated function's first derivative evaluated at x:
124 dydx = f'(x), with the same shape as x.
125 """
126 z = np.asarray(x)
127 y, dydx = self._evalAndDer(z.flatten())
128 return y.reshape(z.shape), dydx.reshape(z.shape)
130 def _evaluate(self, x):
131 """
132 Interpolated function evaluator, to be defined in subclasses.
133 """
134 raise NotImplementedError()
136 def _der(self, x):
137 """
138 Interpolated function derivative evaluator, to be defined in subclasses.
139 """
140 raise NotImplementedError()
142 def _evalAndDer(self, x):
143 """
144 Interpolated function and derivative evaluator, to be defined in subclasses.
145 """
146 raise NotImplementedError()
149class HARKinterpolator2D(MetricObject):
150 """
151 A wrapper class for 2D interpolation methods in HARK.
152 """
154 distance_criteria = []
156 def __call__(self, x, y):
157 """
158 Evaluates the interpolated function at the given input.
160 Parameters
161 ----------
162 x : np.array or float
163 Real values to be evaluated in the interpolated function.
164 y : np.array or float
165 Real values to be evaluated in the interpolated function; must be
166 the same size as x.
168 Returns
169 -------
170 fxy : np.array or float
171 The interpolated function evaluated at x,y: fxy = f(x,y), with the
172 same shape as x and y.
173 """
174 xa = np.asarray(x)
175 ya = np.asarray(y)
176 return (self._evaluate(xa.flatten(), ya.flatten())).reshape(xa.shape)
178 def derivativeX(self, x, y):
179 """
180 Evaluates the partial derivative of interpolated function with respect
181 to x (the first argument) at the given input.
183 Parameters
184 ----------
185 x : np.array or float
186 Real values to be evaluated in the interpolated function.
187 y : np.array or float
188 Real values to be evaluated in the interpolated function; must be
189 the same size as x.
191 Returns
192 -------
193 dfdx : np.array or float
194 The derivative of the interpolated function with respect to x, eval-
195 uated at x,y: dfdx = f_x(x,y), with the same shape as x and y.
196 """
197 xa = np.asarray(x)
198 ya = np.asarray(y)
199 return (self._derX(xa.flatten(), ya.flatten())).reshape(xa.shape)
201 def derivativeY(self, x, y):
202 """
203 Evaluates the partial derivative of interpolated function with respect
204 to y (the second argument) at the given input.
206 Parameters
207 ----------
208 x : np.array or float
209 Real values to be evaluated in the interpolated function.
210 y : np.array or float
211 Real values to be evaluated in the interpolated function; must be
212 the same size as x.
214 Returns
215 -------
216 dfdy : np.array or float
217 The derivative of the interpolated function with respect to y, eval-
218 uated at x,y: dfdx = f_y(x,y), with the same shape as x and y.
219 """
220 xa = np.asarray(x)
221 ya = np.asarray(y)
222 return (self._derY(xa.flatten(), ya.flatten())).reshape(xa.shape)
224 def _evaluate(self, x, y):
225 """
226 Interpolated function evaluator, to be defined in subclasses.
227 """
228 raise NotImplementedError()
230 def _derX(self, x, y):
231 """
232 Interpolated function x-derivative evaluator, to be defined in subclasses.
233 """
234 raise NotImplementedError()
236 def _derY(self, x, y):
237 """
238 Interpolated function y-derivative evaluator, to be defined in subclasses.
239 """
240 raise NotImplementedError()
243class HARKinterpolator3D(MetricObject):
244 """
245 A wrapper class for 3D interpolation methods in HARK.
246 """
248 distance_criteria = []
250 def __call__(self, x, y, z):
251 """
252 Evaluates the interpolated function at the given input.
254 Parameters
255 ----------
256 x : np.array or float
257 Real values to be evaluated in the interpolated function.
258 y : np.array or float
259 Real values to be evaluated in the interpolated function; must be
260 the same size as x.
261 z : np.array or float
262 Real values to be evaluated in the interpolated function; must be
263 the same size as x.
265 Returns
266 -------
267 fxyz : np.array or float
268 The interpolated function evaluated at x,y,z: fxyz = f(x,y,z), with
269 the same shape as x, y, and z.
270 """
271 xa = np.asarray(x)
272 ya = np.asarray(y)
273 za = np.asarray(z)
274 return (self._evaluate(xa.flatten(), ya.flatten(), za.flatten())).reshape(
275 xa.shape
276 )
278 def derivativeX(self, x, y, z):
279 """
280 Evaluates the partial derivative of the interpolated function with respect
281 to x (the first argument) at the given input.
283 Parameters
284 ----------
285 x : np.array or float
286 Real values to be evaluated in the interpolated function.
287 y : np.array or float
288 Real values to be evaluated in the interpolated function; must be
289 the same size as x.
290 z : np.array or float
291 Real values to be evaluated in the interpolated function; must be
292 the same size as x.
294 Returns
295 -------
296 dfdx : np.array or float
297 The derivative with respect to x of the interpolated function evaluated
298 at x,y,z: dfdx = f_x(x,y,z), with the same shape as x, y, and z.
299 """
300 xa = np.asarray(x)
301 ya = np.asarray(y)
302 za = np.asarray(z)
303 return (self._derX(xa.flatten(), ya.flatten(), za.flatten())).reshape(xa.shape)
305 def derivativeY(self, x, y, z):
306 """
307 Evaluates the partial derivative of the interpolated function with respect
308 to y (the second argument) at the given input.
310 Parameters
311 ----------
312 x : np.array or float
313 Real values to be evaluated in the interpolated function.
314 y : np.array or float
315 Real values to be evaluated in the interpolated function; must be
316 the same size as x.
317 z : np.array or float
318 Real values to be evaluated in the interpolated function; must be
319 the same size as x.
321 Returns
322 -------
323 dfdy : np.array or float
324 The derivative with respect to y of the interpolated function evaluated
325 at x,y,z: dfdy = f_y(x,y,z), with the same shape as x, y, and z.
326 """
327 xa = np.asarray(x)
328 ya = np.asarray(y)
329 za = np.asarray(z)
330 return (self._derY(xa.flatten(), ya.flatten(), za.flatten())).reshape(xa.shape)
332 def derivativeZ(self, x, y, z):
333 """
334 Evaluates the partial derivative of the interpolated function with respect
335 to z (the third argument) at the given input.
337 Parameters
338 ----------
339 x : np.array or float
340 Real values to be evaluated in the interpolated function.
341 y : np.array or float
342 Real values to be evaluated in the interpolated function; must be
343 the same size as x.
344 z : np.array or float
345 Real values to be evaluated in the interpolated function; must be
346 the same size as x.
348 Returns
349 -------
350 dfdz : np.array or float
351 The derivative with respect to z of the interpolated function evaluated
352 at x,y,z: dfdz = f_z(x,y,z), with the same shape as x, y, and z.
353 """
354 xa = np.asarray(x)
355 ya = np.asarray(y)
356 za = np.asarray(z)
357 return (self._derZ(xa.flatten(), ya.flatten(), za.flatten())).reshape(xa.shape)
359 def _evaluate(self, x, y, z):
360 """
361 Interpolated function evaluator, to be defined in subclasses.
362 """
363 raise NotImplementedError()
365 def _derX(self, x, y, z):
366 """
367 Interpolated function x-derivative evaluator, to be defined in subclasses.
368 """
369 raise NotImplementedError()
371 def _derY(self, x, y, z):
372 """
373 Interpolated function y-derivative evaluator, to be defined in subclasses.
374 """
375 raise NotImplementedError()
377 def _derZ(self, x, y, z):
378 """
379 Interpolated function y-derivative evaluator, to be defined in subclasses.
380 """
381 raise NotImplementedError()
384class HARKinterpolator4D(MetricObject):
385 """
386 A wrapper class for 4D interpolation methods in HARK.
387 """
389 distance_criteria = []
391 def __call__(self, w, x, y, z):
392 """
393 Evaluates the interpolated function at the given input.
395 Parameters
396 ----------
397 w : np.array or float
398 Real values to be evaluated in the interpolated function.
399 x : np.array or float
400 Real values to be evaluated in the interpolated function; must be
401 the same size as w.
402 y : np.array or float
403 Real values to be evaluated in the interpolated function; must be
404 the same size as w.
405 z : np.array or float
406 Real values to be evaluated in the interpolated function; must be
407 the same size as w.
409 Returns
410 -------
411 fwxyz : np.array or float
412 The interpolated function evaluated at w,x,y,z: fwxyz = f(w,x,y,z),
413 with the same shape as w, x, y, and z.
414 """
415 wa = np.asarray(w)
416 xa = np.asarray(x)
417 ya = np.asarray(y)
418 za = np.asarray(z)
419 return (
420 self._evaluate(wa.flatten(), xa.flatten(), ya.flatten(), za.flatten())
421 ).reshape(wa.shape)
423 def derivativeW(self, w, x, y, z):
424 """
425 Evaluates the partial derivative with respect to w (the first argument)
426 of the interpolated function at the given input.
428 Parameters
429 ----------
430 w : np.array or float
431 Real values to be evaluated in the interpolated function.
432 x : np.array or float
433 Real values to be evaluated in the interpolated function; must be
434 the same size as w.
435 y : np.array or float
436 Real values to be evaluated in the interpolated function; must be
437 the same size as w.
438 z : np.array or float
439 Real values to be evaluated in the interpolated function; must be
440 the same size as w.
442 Returns
443 -------
444 dfdw : np.array or float
445 The derivative with respect to w of the interpolated function eval-
446 uated at w,x,y,z: dfdw = f_w(w,x,y,z), with the same shape as inputs.
447 """
448 wa = np.asarray(w)
449 xa = np.asarray(x)
450 ya = np.asarray(y)
451 za = np.asarray(z)
452 return (
453 self._derW(wa.flatten(), xa.flatten(), ya.flatten(), za.flatten())
454 ).reshape(wa.shape)
456 def derivativeX(self, w, x, y, z):
457 """
458 Evaluates the partial derivative with respect to x (the second argument)
459 of the interpolated function at the given input.
461 Parameters
462 ----------
463 w : np.array or float
464 Real values to be evaluated in the interpolated function.
465 x : np.array or float
466 Real values to be evaluated in the interpolated function; must be
467 the same size as w.
468 y : np.array or float
469 Real values to be evaluated in the interpolated function; must be
470 the same size as w.
471 z : np.array or float
472 Real values to be evaluated in the interpolated function; must be
473 the same size as w.
475 Returns
476 -------
477 dfdx : np.array or float
478 The derivative with respect to x of the interpolated function eval-
479 uated at w,x,y,z: dfdx = f_x(w,x,y,z), with the same shape as inputs.
480 """
481 wa = np.asarray(w)
482 xa = np.asarray(x)
483 ya = np.asarray(y)
484 za = np.asarray(z)
485 return (
486 self._derX(wa.flatten(), xa.flatten(), ya.flatten(), za.flatten())
487 ).reshape(wa.shape)
489 def derivativeY(self, w, x, y, z):
490 """
491 Evaluates the partial derivative with respect to y (the third argument)
492 of the interpolated function at the given input.
494 Parameters
495 ----------
496 w : np.array or float
497 Real values to be evaluated in the interpolated function.
498 x : np.array or float
499 Real values to be evaluated in the interpolated function; must be
500 the same size as w.
501 y : np.array or float
502 Real values to be evaluated in the interpolated function; must be
503 the same size as w.
504 z : np.array or float
505 Real values to be evaluated in the interpolated function; must be
506 the same size as w.
508 Returns
509 -------
510 dfdy : np.array or float
511 The derivative with respect to y of the interpolated function eval-
512 uated at w,x,y,z: dfdy = f_y(w,x,y,z), with the same shape as inputs.
513 """
514 wa = np.asarray(w)
515 xa = np.asarray(x)
516 ya = np.asarray(y)
517 za = np.asarray(z)
518 return (
519 self._derY(wa.flatten(), xa.flatten(), ya.flatten(), za.flatten())
520 ).reshape(wa.shape)
522 def derivativeZ(self, w, x, y, z):
523 """
524 Evaluates the partial derivative with respect to z (the fourth argument)
525 of the interpolated function at the given input.
527 Parameters
528 ----------
529 w : np.array or float
530 Real values to be evaluated in the interpolated function.
531 x : np.array or float
532 Real values to be evaluated in the interpolated function; must be
533 the same size as w.
534 y : np.array or float
535 Real values to be evaluated in the interpolated function; must be
536 the same size as w.
537 z : np.array or float
538 Real values to be evaluated in the interpolated function; must be
539 the same size as w.
541 Returns
542 -------
543 dfdz : np.array or float
544 The derivative with respect to z of the interpolated function eval-
545 uated at w,x,y,z: dfdz = f_z(w,x,y,z), with the same shape as inputs.
546 """
547 wa = np.asarray(w)
548 xa = np.asarray(x)
549 ya = np.asarray(y)
550 za = np.asarray(z)
551 return (
552 self._derZ(wa.flatten(), xa.flatten(), ya.flatten(), za.flatten())
553 ).reshape(wa.shape)
555 def _evaluate(self, w, x, y, z):
556 """
557 Interpolated function evaluator, to be defined in subclasses.
558 """
559 raise NotImplementedError()
561 def _derW(self, w, x, y, z):
562 """
563 Interpolated function w-derivative evaluator, to be defined in subclasses.
564 """
565 raise NotImplementedError()
567 def _derX(self, w, x, y, z):
568 """
569 Interpolated function w-derivative evaluator, to be defined in subclasses.
570 """
571 raise NotImplementedError()
573 def _derY(self, w, x, y, z):
574 """
575 Interpolated function w-derivative evaluator, to be defined in subclasses.
576 """
577 raise NotImplementedError()
579 def _derZ(self, w, x, y, z):
580 """
581 Interpolated function w-derivative evaluator, to be defined in subclasses.
582 """
583 raise NotImplementedError()
586class IdentityFunction(MetricObject):
587 """
588 A fairly trivial interpolator that simply returns one of its arguments. Useful for avoiding
589 numeric error in extreme cases.
591 Parameters
592 ----------
593 i_dim : int
594 Index of the dimension on which the identity is defined. ``f(*x) = x[i]``
595 n_dims : int
596 Total number of input dimensions for this function.
597 """
599 distance_criteria = ["i_dim"]
601 def __init__(self, i_dim=0, n_dims=1):
602 self.i_dim = i_dim
603 self.n_dims = n_dims
605 def __call__(self, *args):
606 """
607 Evaluate the identity function.
608 """
609 return args[self.i_dim]
611 def derivative(self, *args):
612 """
613 Returns the derivative of the function with respect to the first dimension.
614 """
615 if self.i_dim == 0:
616 return np.ones_like(args[0])
617 else:
618 return np.zeros_like(args[0])
620 def derivativeX(self, *args):
621 """
622 Returns the derivative of the function with respect to the X dimension.
623 This is the first input whenever n_dims < 4 and the second input otherwise.
624 """
625 if self.n_dims >= 4:
626 j = 1
627 else:
628 j = 0
629 if self.i_dim == j:
630 return np.ones_like(args[0])
631 else:
632 return np.zeros_like(args[0])
634 def derivativeY(self, *args):
635 """
636 Returns the derivative of the function with respect to the Y dimension.
637 This is the second input whenever n_dims < 4 and the third input otherwise.
638 """
639 if self.n_dims >= 4:
640 j = 2
641 else:
642 j = 1
643 if self.i_dim == j:
644 return np.ones_like(args[0])
645 else:
646 return np.zeros_like(args[0])
648 def derivativeZ(self, *args):
649 """
650 Returns the derivative of the function with respect to the Z dimension.
651 This is the third input whenever n_dims < 4 and the fourth input otherwise.
652 """
653 if self.n_dims >= 4:
654 j = 3
655 else:
656 j = 2
657 if self.i_dim == j:
658 return np.ones_like(args[0])
659 else:
660 return np.zeros_like(args[0])
662 def derivativeW(self, *args):
663 """
664 Returns the derivative of the function with respect to the W dimension.
665 This should only exist when n_dims >= 4.
666 """
667 if self.n_dims >= 4:
668 j = 0
669 else:
670 assert False, (
671 "Derivative with respect to W can't be called when n_dims < 4!"
672 )
673 if self.i_dim == j:
674 return np.ones_like(args[0])
675 else:
676 return np.zeros_like(args[0])
679class ConstantFunction(MetricObject):
680 """
681 A class for representing trivial functions that return the same real output for any input. This
682 is convenient for models where an object might be a (non-trivial) function, but in some variations
683 that object is just a constant number. Rather than needing to make a (Bi/Tri/Quad)-
684 LinearInterpolation with trivial state grids and the same f_value in every entry, ConstantFunction
685 allows the user to quickly make a constant/trivial function. This comes up, e.g., in models
686 with endogenous pricing of insurance contracts; a contract's premium might depend on some state
687 variables of the individual, but in some variations the premium of a contract is just a number.
689 Parameters
690 ----------
691 value : float
692 The constant value that the function returns.
693 """
695 convergence_criteria = ["value"]
697 def __init__(self, value):
698 self.value = float(value)
700 def __call__(self, *args):
701 """
702 Evaluate the constant function. The first input must exist and should be an array.
703 Returns an array of identical shape to args[0] (if it exists).
704 """
705 if (
706 len(args) > 0
707 ): # If there is at least one argument, return appropriately sized array
708 if _isscalar(args[0]):
709 return self.value
710 else:
711 shape = args[0].shape
712 return self.value * np.ones(shape)
713 else: # Otherwise, return a single instance of the constant value
714 return self.value
716 def _der(self, *args):
717 """
718 Evaluate the derivative of the function. The first input must exist and should be an array.
719 Returns an array of identical shape to args[0] (if it exists). This is an array of zeros.
720 """
721 if len(args) > 0:
722 if _isscalar(args[0]):
723 return 0.0
724 else:
725 shape = args[0].shape
726 return np.zeros(shape)
727 else:
728 return 0.0
730 def eval_with_derivative(self, x):
731 val = self(x)
732 der = self._der(x)
733 return val, der
735 # All other derivatives are also zero everywhere, so these methods just point to derivative
736 derivative = _der
737 derivativeX = derivative
738 derivativeY = derivative
739 derivativeZ = derivative
740 derivativeW = derivative
741 derivativeXX = derivative
744class LinearInterp(HARKinterpolator1D):
745 """
746 A "from scratch" 1D linear interpolation class. Allows for linear or decay
747 extrapolation (approaching a limiting linear function from below).
749 NOTE: When no input is given for the limiting linear function, linear
750 extrapolation is used above the highest gridpoint.
752 Parameters
753 ----------
754 x_list : np.array
755 List of x values composing the grid.
756 y_list : np.array
757 List of y values, representing f(x) at the points in x_list.
758 intercept_limit : float
759 Intercept of limiting linear function.
760 slope_limit : float
761 Slope of limiting linear function.
762 lower_extrap : bool
763 Indicator for whether lower extrapolation is allowed. False means
764 f(x) = NaN for x < min(x_list); True means linear extrapolation.
765 pre_compute : bool
766 Indicator for whether interpolation coefficients should be pre-computed
767 and stored as attributes of self (default False). More memory will be used,
768 and instantiation will take slightly longer, but later evaluation will
769 be faster due to less arithmetic.
770 indexer : function or None (default)
771 If provided, a custom function that identifies the index of the interpolant
772 segment for each query point. Should return results identically to the
773 default behavior of np.maximum(np.searchsorted(self.x_list[:-1], x), 1).
774 WARNING: User is responsible for verifying that their custom indexer is
775 actually correct versus default behavior.
776 """
778 distance_criteria = ["x_list", "y_list"]
780 def __init__(
781 self,
782 x_list,
783 y_list,
784 intercept_limit=None,
785 slope_limit=None,
786 lower_extrap=False,
787 pre_compute=False,
788 indexer=None,
789 ):
790 # Make the basic linear spline interpolation
791 self.x_list = (
792 np.array(x_list)
793 if _check_flatten(1, x_list)
794 else np.array(x_list).flatten()
795 )
796 self.y_list = (
797 np.array(y_list)
798 if _check_flatten(1, y_list)
799 else np.array(y_list).flatten()
800 )
801 _check_grid_dimensions(1, self.y_list, self.x_list)
802 self.lower_extrap = lower_extrap
803 self.x_n = self.x_list.size
804 self.indexer = indexer
806 # Make a decay extrapolation
807 if intercept_limit is not None and slope_limit is not None:
808 slope_at_top = (y_list[-1] - y_list[-2]) / (x_list[-1] - x_list[-2])
809 level_diff = intercept_limit + slope_limit * x_list[-1] - y_list[-1]
810 slope_diff = slope_limit - slope_at_top
811 # If the model that can handle uncertainty has been calibrated with
812 # with uncertainty set to zero, the 'extrapolation' will blow up
813 # Guard against that and nearby problems by testing slope equality
814 if not np.isclose(slope_limit, slope_at_top, atol=1e-15):
815 self.decay_extrap_A = level_diff
816 self.decay_extrap_B = -slope_diff / level_diff
817 self.intercept_limit = intercept_limit
818 self.slope_limit = slope_limit
819 self.decay_extrap = True
820 else:
821 self.decay_extrap = False
822 else:
823 self.decay_extrap = False
825 # Calculate interpolation coefficients now rather than at evaluation time
826 if pre_compute:
827 self.slopes = (self.y_list[1:] - self.y_list[:-1]) / (
828 self.x_list[1:] - self.x_list[:-1]
829 )
830 self.intercepts = self.y_list[:-1] - self.slopes * self.x_list[:-1]
832 def _evalOrDer(self, x, _eval, _Der):
833 """
834 Returns the level and/or first derivative of the function at each value in
835 x. Only called internally by HARKinterpolator1D.eval_and_der (etc).
837 Parameters
838 ----------
839 x_list : scalar or np.array
840 Set of points where we want to evlauate the interpolated function and/or its derivative..
841 _eval : boolean
842 Indicator for whether to evalute the level of the interpolated function.
843 _Der : boolean
844 Indicator for whether to evaluate the derivative of the interpolated function.
846 Returns
847 -------
848 A list including the level and/or derivative of the interpolated function where requested.
849 """
850 if self.indexer is None:
851 i = np.maximum(np.searchsorted(self.x_list[:-1], x), 1)
852 else:
853 i = self.indexer(x)
855 if hasattr(self, "slopes"):
856 # Coefficients were pre-computed, use those
857 j = i - 1
858 dydx = self.slopes[j]
859 if _eval:
860 y = self.intercepts[j] + dydx * x
862 else:
863 # Find relative weights between endpoints and evaluate interpolation
864 alpha = (x - self.x_list[i - 1]) / (self.x_list[i] - self.x_list[i - 1])
866 if _eval:
867 y = (1.0 - alpha) * self.y_list[i - 1] + alpha * self.y_list[i]
868 if _Der:
869 dydx = (self.y_list[i] - self.y_list[i - 1]) / (
870 self.x_list[i] - self.x_list[i - 1]
871 )
873 if not self.lower_extrap:
874 below_lower_bound = x < self.x_list[0]
876 if _eval:
877 y[below_lower_bound] = np.nan
878 if _Der:
879 dydx[below_lower_bound] = np.nan
881 if self.decay_extrap:
882 above_upper_bound = x > self.x_list[-1]
883 x_temp = x[above_upper_bound] - self.x_list[-1]
885 if _eval:
886 y[above_upper_bound] = (
887 self.intercept_limit
888 + self.slope_limit * x[above_upper_bound]
889 - self.decay_extrap_A * np.exp(-self.decay_extrap_B * x_temp)
890 )
892 if _Der:
893 dydx[above_upper_bound] = (
894 self.slope_limit
895 + self.decay_extrap_B
896 * self.decay_extrap_A
897 * np.exp(-self.decay_extrap_B * x_temp)
898 )
900 output = []
901 if _eval:
902 output += [
903 y,
904 ]
905 if _Der:
906 output += [
907 dydx,
908 ]
910 return output
912 def _evaluate(self, x, return_indices=False):
913 """
914 Returns the level of the interpolated function at each value in x. Only
915 called internally by HARKinterpolator1D.__call__ (etc).
916 """
917 return self._evalOrDer(x, True, False)[0]
919 def _der(self, x):
920 """
921 Returns the first derivative of the interpolated function at each value
922 in x. Only called internally by HARKinterpolator1D.derivative (etc).
923 """
924 return self._evalOrDer(x, False, True)[0]
926 def _evalAndDer(self, x):
927 """
928 Returns the level and first derivative of the function at each value in
929 x. Only called internally by HARKinterpolator1D.eval_and_der (etc).
930 """
931 y, dydx = self._evalOrDer(x, True, True)
933 return y, dydx
936class CubicInterp(HARKinterpolator1D):
937 """
938 An interpolating function using piecewise cubic splines. Matches level and
939 slope of 1D function at gridpoints, smoothly interpolating in between.
940 Extrapolation above highest gridpoint approaches a limiting linear function
941 if desired (linear extrapolation also enabled.)
943 NOTE: When no input is given for the limiting linear function, linear
944 extrapolation is used above the highest gridpoint.
946 Parameters
947 ----------
948 x_list : np.array
949 List of x values composing the grid.
950 y_list : np.array
951 List of y values, representing f(x) at the points in x_list.
952 dydx_list : np.array
953 List of dydx values, representing f'(x) at the points in x_list
954 intercept_limit : float
955 Intercept of limiting linear function.
956 slope_limit : float
957 Slope of limiting linear function.
958 lower_extrap : boolean
959 Indicator for whether lower extrapolation is allowed. False means
960 f(x) = NaN for x < min(x_list); True means linear extrapolation.
961 """
963 distance_criteria = ["x_list", "y_list", "dydx_list"]
965 def __init__(
966 self,
967 x_list,
968 y_list,
969 dydx_list,
970 intercept_limit=None,
971 slope_limit=None,
972 lower_extrap=False,
973 ):
974 self.x_list = (
975 np.asarray(x_list)
976 if _check_flatten(1, x_list)
977 else np.array(x_list).flatten()
978 )
979 self.y_list = (
980 np.asarray(y_list)
981 if _check_flatten(1, y_list)
982 else np.array(y_list).flatten()
983 )
984 self.dydx_list = (
985 np.asarray(dydx_list)
986 if _check_flatten(1, dydx_list)
987 else np.array(dydx_list).flatten()
988 )
989 _check_grid_dimensions(1, self.y_list, self.x_list)
990 _check_grid_dimensions(1, self.dydx_list, self.x_list)
992 self.n = len(x_list)
994 # Define lower extrapolation as linear function (or just NaN)
995 if lower_extrap:
996 self.coeffs = [[y_list[0], dydx_list[0], 0, 0]]
997 else:
998 self.coeffs = [[np.nan, np.nan, np.nan, np.nan]]
1000 # Calculate interpolation coefficients on segments mapped to [0,1]
1001 for i in range(self.n - 1):
1002 x0 = x_list[i]
1003 y0 = y_list[i]
1004 x1 = x_list[i + 1]
1005 y1 = y_list[i + 1]
1006 Span = x1 - x0
1007 dydx0 = dydx_list[i] * Span
1008 dydx1 = dydx_list[i + 1] * Span
1010 temp = [
1011 y0,
1012 dydx0,
1013 3 * (y1 - y0) - 2 * dydx0 - dydx1,
1014 2 * (y0 - y1) + dydx0 + dydx1,
1015 ]
1016 self.coeffs.append(temp)
1018 # Calculate extrapolation coefficients as a decay toward limiting function y = mx+b
1019 if slope_limit is None and intercept_limit is None:
1020 slope_limit = dydx_list[-1]
1021 intercept_limit = y_list[-1] - slope_limit * x_list[-1]
1022 gap = slope_limit * x1 + intercept_limit - y1
1023 slope = slope_limit - dydx_list[self.n - 1]
1024 if (gap != 0) and (slope <= 0):
1025 temp = [intercept_limit, slope_limit, gap, slope / gap]
1026 elif slope > 0:
1027 temp = [
1028 intercept_limit,
1029 slope_limit,
1030 0,
1031 0,
1032 ] # fixing a problem when slope is positive
1033 else:
1034 temp = [intercept_limit, slope_limit, gap, 0]
1035 self.coeffs.append(temp)
1036 self.coeffs = np.array(self.coeffs)
1038 def _evaluate(self, x):
1039 """
1040 Returns the level of the interpolated function at each value in x. Only
1041 called internally by HARKinterpolator1D.__call__ (etc).
1042 """
1044 m = len(x)
1045 pos = np.searchsorted(self.x_list, x, side="right")
1046 y = np.zeros(m)
1047 if y.size > 0:
1048 out_bot = pos == 0
1049 out_top = pos == self.n
1050 in_bnds = np.logical_not(np.logical_or(out_bot, out_top))
1052 # Do the "in bounds" evaluation points
1053 i = pos[in_bnds]
1054 coeffs_in = self.coeffs[i, :]
1055 alpha = (x[in_bnds] - self.x_list[i - 1]) / (
1056 self.x_list[i] - self.x_list[i - 1]
1057 )
1058 y[in_bnds] = coeffs_in[:, 0] + alpha * (
1059 coeffs_in[:, 1] + alpha * (coeffs_in[:, 2] + alpha * coeffs_in[:, 3])
1060 )
1062 # Do the "out of bounds" evaluation points
1063 y[out_bot] = self.coeffs[0, 0] + self.coeffs[0, 1] * (
1064 x[out_bot] - self.x_list[0]
1065 )
1066 alpha = x[out_top] - self.x_list[self.n - 1]
1067 y[out_top] = (
1068 self.coeffs[self.n, 0]
1069 + x[out_top] * self.coeffs[self.n, 1]
1070 - self.coeffs[self.n, 2] * np.exp(alpha * self.coeffs[self.n, 3])
1071 )
1073 return y
1075 def _der(self, x):
1076 """
1077 Returns the first derivative of the interpolated function at each value
1078 in x. Only called internally by HARKinterpolator1D.derivative (etc).
1079 """
1081 m = len(x)
1082 pos = np.searchsorted(self.x_list, x, side="right")
1083 dydx = np.zeros(m)
1084 if dydx.size > 0:
1085 out_bot = pos == 0
1086 out_top = pos == self.n
1087 in_bnds = np.logical_not(np.logical_or(out_bot, out_top))
1089 # Do the "in bounds" evaluation points
1090 i = pos[in_bnds]
1091 coeffs_in = self.coeffs[i, :]
1092 alpha = (x[in_bnds] - self.x_list[i - 1]) / (
1093 self.x_list[i] - self.x_list[i - 1]
1094 )
1095 dydx[in_bnds] = (
1096 coeffs_in[:, 1]
1097 + alpha * (2 * coeffs_in[:, 2] + alpha * 3 * coeffs_in[:, 3])
1098 ) / (self.x_list[i] - self.x_list[i - 1])
1100 # Do the "out of bounds" evaluation points
1101 dydx[out_bot] = self.coeffs[0, 1]
1102 alpha = x[out_top] - self.x_list[self.n - 1]
1103 dydx[out_top] = self.coeffs[self.n, 1] - self.coeffs[
1104 self.n, 2
1105 ] * self.coeffs[self.n, 3] * np.exp(alpha * self.coeffs[self.n, 3])
1106 return dydx
1108 def _evalAndDer(self, x):
1109 """
1110 Returns the level and first derivative of the function at each value in
1111 x. Only called internally by HARKinterpolator1D.eval_and_der (etc).
1112 """
1113 m = len(x)
1114 pos = np.searchsorted(self.x_list, x, side="right")
1115 y = np.zeros(m)
1116 dydx = np.zeros(m)
1117 if y.size > 0:
1118 out_bot = pos == 0
1119 out_top = pos == self.n
1120 in_bnds = np.logical_not(np.logical_or(out_bot, out_top))
1122 # Do the "in bounds" evaluation points
1123 i = pos[in_bnds]
1124 coeffs_in = self.coeffs[i, :]
1125 alpha = (x[in_bnds] - self.x_list[i - 1]) / (
1126 self.x_list[i] - self.x_list[i - 1]
1127 )
1128 y[in_bnds] = coeffs_in[:, 0] + alpha * (
1129 coeffs_in[:, 1] + alpha * (coeffs_in[:, 2] + alpha * coeffs_in[:, 3])
1130 )
1131 dydx[in_bnds] = (
1132 coeffs_in[:, 1]
1133 + alpha * (2 * coeffs_in[:, 2] + alpha * 3 * coeffs_in[:, 3])
1134 ) / (self.x_list[i] - self.x_list[i - 1])
1136 # Do the "out of bounds" evaluation points
1137 y[out_bot] = self.coeffs[0, 0] + self.coeffs[0, 1] * (
1138 x[out_bot] - self.x_list[0]
1139 )
1140 dydx[out_bot] = self.coeffs[0, 1]
1141 alpha = x[out_top] - self.x_list[self.n - 1]
1142 y[out_top] = (
1143 self.coeffs[self.n, 0]
1144 + x[out_top] * self.coeffs[self.n, 1]
1145 - self.coeffs[self.n, 2] * np.exp(alpha * self.coeffs[self.n, 3])
1146 )
1147 dydx[out_top] = self.coeffs[self.n, 1] - self.coeffs[
1148 self.n, 2
1149 ] * self.coeffs[self.n, 3] * np.exp(alpha * self.coeffs[self.n, 3])
1150 return y, dydx
1153class CubicHermiteInterp(HARKinterpolator1D):
1154 """
1155 An interpolating function using piecewise cubic splines. Matches level and
1156 slope of 1D function at gridpoints, smoothly interpolating in between.
1157 Extrapolation above highest gridpoint approaches a limiting linear function
1158 if desired (linear extrapolation also enabled.)
1160 NOTE: When no input is given for the limiting linear function, linear
1161 extrapolation is used above the highest gridpoint.
1163 Parameters
1164 ----------
1165 x_list : np.array
1166 List of x values composing the grid.
1167 y_list : np.array
1168 List of y values, representing f(x) at the points in x_list.
1169 dydx_list : np.array
1170 List of dydx values, representing f'(x) at the points in x_list
1171 intercept_limit : float
1172 Intercept of limiting linear function.
1173 slope_limit : float
1174 Slope of limiting linear function.
1175 lower_extrap : boolean
1176 Indicator for whether lower extrapolation is allowed. False means
1177 f(x) = NaN for x < min(x_list); True means linear extrapolation.
1178 """
1180 distance_criteria = ["x_list", "y_list", "dydx_list"]
1182 def __init__(
1183 self,
1184 x_list,
1185 y_list,
1186 dydx_list,
1187 intercept_limit=None,
1188 slope_limit=None,
1189 lower_extrap=False,
1190 ):
1191 self.x_list = (
1192 np.asarray(x_list)
1193 if _check_flatten(1, x_list)
1194 else np.array(x_list).flatten()
1195 )
1196 self.y_list = (
1197 np.asarray(y_list)
1198 if _check_flatten(1, y_list)
1199 else np.array(y_list).flatten()
1200 )
1201 self.dydx_list = (
1202 np.asarray(dydx_list)
1203 if _check_flatten(1, dydx_list)
1204 else np.array(dydx_list).flatten()
1205 )
1206 _check_grid_dimensions(1, self.y_list, self.x_list)
1207 _check_grid_dimensions(1, self.dydx_list, self.x_list)
1209 self.n = len(x_list)
1211 self._chs = CubicHermiteSpline(
1212 self.x_list, self.y_list, self.dydx_list, extrapolate=None
1213 )
1214 self.coeffs = np.flip(self._chs.c.T, 1)
1216 # Define lower extrapolation as linear function (or just NaN)
1217 if lower_extrap:
1218 temp = np.array([self.y_list[0], self.dydx_list[0], 0, 0])
1219 else:
1220 temp = np.array([np.nan, np.nan, np.nan, np.nan])
1222 self.coeffs = np.vstack((temp, self.coeffs))
1224 x1 = self.x_list[self.n - 1]
1225 y1 = self.y_list[self.n - 1]
1227 # Calculate extrapolation coefficients as a decay toward limiting function y = mx+b
1228 if slope_limit is None and intercept_limit is None:
1229 slope_limit = self.dydx_list[-1]
1230 intercept_limit = self.y_list[-1] - slope_limit * self.x_list[-1]
1231 gap = slope_limit * x1 + intercept_limit - y1
1232 slope = slope_limit - self.dydx_list[self.n - 1]
1233 if (gap != 0) and (slope <= 0):
1234 temp = np.array([intercept_limit, slope_limit, gap, slope / gap])
1235 elif slope > 0:
1236 # fixing a problem when slope is positive
1237 temp = np.array([intercept_limit, slope_limit, 0, 0])
1238 else:
1239 temp = np.array([intercept_limit, slope_limit, gap, 0])
1240 self.coeffs = np.vstack((self.coeffs, temp))
1242 def out_of_bounds(self, x):
1243 out_bot = x < self.x_list[0]
1244 out_top = x > self.x_list[-1]
1245 return out_bot, out_top
1247 def _evaluate(self, x):
1248 """
1249 Returns the level of the interpolated function at each value in x. Only
1250 called internally by HARKinterpolator1D.__call__ (etc).
1251 """
1252 out_bot, out_top = self.out_of_bounds(x)
1254 return self._eval_helper(x, out_bot, out_top)
1256 def _eval_helper(self, x, out_bot, out_top):
1257 y = self._chs(x)
1259 # Do the "out of bounds" evaluation points
1260 if any(out_bot):
1261 y[out_bot] = self.coeffs[0, 0] + self.coeffs[0, 1] * (
1262 x[out_bot] - self.x_list[0]
1263 )
1265 if any(out_top):
1266 alpha = x[out_top] - self.x_list[self.n - 1]
1267 y[out_top] = (
1268 self.coeffs[self.n, 0]
1269 + x[out_top] * self.coeffs[self.n, 1]
1270 - self.coeffs[self.n, 2] * np.exp(alpha * self.coeffs[self.n, 3])
1271 )
1273 return y
1275 def _der(self, x):
1276 """
1277 Returns the first derivative of the interpolated function at each value
1278 in x. Only called internally by HARKinterpolator1D.derivative (etc).
1279 """
1280 out_bot, out_top = self.out_of_bounds(x)
1282 return self._der_helper(x, out_bot, out_top)
1284 def _der_helper(self, x, out_bot, out_top):
1285 dydx = self._chs(x, nu=1)
1287 # Do the "out of bounds" evaluation points
1288 if any(out_bot):
1289 dydx[out_bot] = self.coeffs[0, 1]
1290 if any(out_top):
1291 alpha = x[out_top] - self.x_list[self.n - 1]
1292 dydx[out_top] = self.coeffs[self.n, 1] - self.coeffs[
1293 self.n, 2
1294 ] * self.coeffs[self.n, 3] * np.exp(alpha * self.coeffs[self.n, 3])
1296 return dydx
1298 def _evalAndDer(self, x):
1299 """
1300 Returns the level and first derivative of the function at each value in
1301 x. Only called internally by HARKinterpolator1D.eval_and_der (etc).
1302 """
1303 out_bot, out_top = self.out_of_bounds(x)
1304 y = self._eval_helper(x, out_bot, out_top)
1305 dydx = self._der_helper(x, out_bot, out_top)
1306 return y, dydx
1308 def der_interp(self, nu=1):
1309 """
1310 Construct a new piecewise polynomial representing the derivative.
1311 See `scipy` for additional documentation:
1312 https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.CubicHermiteSpline.html
1313 """
1314 return self._chs.derivative(nu)
1316 def antider_interp(self, nu=1):
1317 """
1318 Construct a new piecewise polynomial representing the antiderivative.
1320 Antiderivative is also the indefinite integral of the function,
1321 and derivative is its inverse operation.
1323 See `scipy` for additional documentation:
1324 https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.CubicHermiteSpline.html
1325 """
1326 return self._chs.antiderivative(nu)
1328 def integrate(self, a, b, extrapolate=False):
1329 """
1330 Compute a definite integral over a piecewise polynomial.
1332 See `scipy` for additional documentation:
1333 https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.CubicHermiteSpline.html
1334 """
1335 return self._chs.integrate(a, b, extrapolate)
1337 def roots(self, discontinuity=True, extrapolate=False):
1338 """
1339 Find real roots of the the piecewise polynomial.
1341 See `scipy` for additional documentation:
1342 https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.CubicHermiteSpline.html
1343 """
1344 return self._chs.roots(discontinuity, extrapolate)
1346 def solve(self, y=0.0, discontinuity=True, extrapolate=False):
1347 """
1348 Find real solutions of the the equation ``pp(x) == y``.
1350 See `scipy` for additional documentation:
1351 https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.CubicHermiteSpline.html
1352 """
1353 return self._chs.solve(y, discontinuity, extrapolate)
1356class BilinearInterp(HARKinterpolator2D):
1357 """
1358 Bilinear full (or tensor) grid interpolation of a function f(x,y).
1360 Parameters
1361 ----------
1362 f_values : numpy.array
1363 An array of size (x_n,y_n) such that f_values[i,j] = f(x_list[i],y_list[j])
1364 x_list : numpy.array
1365 An array of x values, with length designated x_n.
1366 y_list : numpy.array
1367 An array of y values, with length designated y_n.
1368 xSearchFunc : function
1369 An optional function that returns the reference location for x values:
1370 indices = xSearchFunc(x_list,x). Default is np.searchsorted
1371 ySearchFunc : function
1372 An optional function that returns the reference location for y values:
1373 indices = ySearchFunc(y_list,y). Default is np.searchsorted
1374 """
1376 distance_criteria = ["x_list", "y_list", "f_values"]
1378 def __init__(self, f_values, x_list, y_list, xSearchFunc=None, ySearchFunc=None):
1379 self.f_values = f_values
1380 self.x_list = (
1381 np.array(x_list)
1382 if _check_flatten(1, x_list)
1383 else np.array(x_list).flatten()
1384 )
1385 self.y_list = (
1386 np.array(y_list)
1387 if _check_flatten(1, y_list)
1388 else np.array(y_list).flatten()
1389 )
1390 _check_grid_dimensions(2, self.f_values, self.x_list, self.y_list)
1391 self.x_n = x_list.size
1392 self.y_n = y_list.size
1393 if xSearchFunc is None:
1394 xSearchFunc = np.searchsorted
1395 if ySearchFunc is None:
1396 ySearchFunc = np.searchsorted
1397 self.xSearchFunc = xSearchFunc
1398 self.ySearchFunc = ySearchFunc
1400 def _evaluate(self, x, y):
1401 """
1402 Returns the level of the interpolated function at each value in x,y.
1403 Only called internally by HARKinterpolator2D.__call__ (etc).
1404 """
1405 x_pos = self.xSearchFunc(self.x_list, x)
1406 x_pos[x_pos < 1] = 1
1407 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
1408 y_pos = self.ySearchFunc(self.y_list, y)
1409 y_pos[y_pos < 1] = 1
1410 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
1411 alpha = (x - self.x_list[x_pos - 1]) / (
1412 self.x_list[x_pos] - self.x_list[x_pos - 1]
1413 )
1414 beta = (y - self.y_list[y_pos - 1]) / (
1415 self.y_list[y_pos] - self.y_list[y_pos - 1]
1416 )
1417 f = (
1418 (1 - alpha) * (1 - beta) * self.f_values[x_pos - 1, y_pos - 1]
1419 + (1 - alpha) * beta * self.f_values[x_pos - 1, y_pos]
1420 + alpha * (1 - beta) * self.f_values[x_pos, y_pos - 1]
1421 + alpha * beta * self.f_values[x_pos, y_pos]
1422 )
1423 return f
1425 def _derX(self, x, y):
1426 """
1427 Returns the derivative with respect to x of the interpolated function
1428 at each value in x,y. Only called internally by HARKinterpolator2D.derivativeX.
1429 """
1430 x_pos = self.xSearchFunc(self.x_list, x)
1431 x_pos[x_pos < 1] = 1
1432 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
1433 y_pos = self.ySearchFunc(self.y_list, y)
1434 y_pos[y_pos < 1] = 1
1435 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
1436 beta = (y - self.y_list[y_pos - 1]) / (
1437 self.y_list[y_pos] - self.y_list[y_pos - 1]
1438 )
1439 dfdx = (
1440 (
1441 (1 - beta) * self.f_values[x_pos, y_pos - 1]
1442 + beta * self.f_values[x_pos, y_pos]
1443 )
1444 - (
1445 (1 - beta) * self.f_values[x_pos - 1, y_pos - 1]
1446 + beta * self.f_values[x_pos - 1, y_pos]
1447 )
1448 ) / (self.x_list[x_pos] - self.x_list[x_pos - 1])
1449 return dfdx
1451 def _derY(self, x, y):
1452 """
1453 Returns the derivative with respect to y of the interpolated function
1454 at each value in x,y. Only called internally by HARKinterpolator2D.derivativeY.
1455 """
1456 x_pos = self.xSearchFunc(self.x_list, x)
1457 x_pos[x_pos < 1] = 1
1458 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
1459 y_pos = self.ySearchFunc(self.y_list, y)
1460 y_pos[y_pos < 1] = 1
1461 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
1462 alpha = (x - self.x_list[x_pos - 1]) / (
1463 self.x_list[x_pos] - self.x_list[x_pos - 1]
1464 )
1465 dfdy = (
1466 (
1467 (1 - alpha) * self.f_values[x_pos - 1, y_pos]
1468 + alpha * self.f_values[x_pos, y_pos]
1469 )
1470 - (
1471 (1 - alpha) * self.f_values[x_pos - 1, y_pos - 1]
1472 + alpha * self.f_values[x_pos, y_pos - 1]
1473 )
1474 ) / (self.y_list[y_pos] - self.y_list[y_pos - 1])
1475 return dfdy
1478class TrilinearInterp(HARKinterpolator3D):
1479 """
1480 Trilinear full (or tensor) grid interpolation of a function f(x,y,z).
1482 Parameters
1483 ----------
1484 f_values : numpy.array
1485 An array of size (x_n,y_n,z_n) such that f_values[i,j,k] =
1486 f(x_list[i],y_list[j],z_list[k])
1487 x_list : numpy.array
1488 An array of x values, with length designated x_n.
1489 y_list : numpy.array
1490 An array of y values, with length designated y_n.
1491 z_list : numpy.array
1492 An array of z values, with length designated z_n.
1493 xSearchFunc : function
1494 An optional function that returns the reference location for x values:
1495 indices = xSearchFunc(x_list,x). Default is np.searchsorted
1496 ySearchFunc : function
1497 An optional function that returns the reference location for y values:
1498 indices = ySearchFunc(y_list,y). Default is np.searchsorted
1499 zSearchFunc : function
1500 An optional function that returns the reference location for z values:
1501 indices = zSearchFunc(z_list,z). Default is np.searchsorted
1502 """
1504 distance_criteria = ["f_values", "x_list", "y_list", "z_list"]
1506 def __init__(
1507 self,
1508 f_values,
1509 x_list,
1510 y_list,
1511 z_list,
1512 xSearchFunc=None,
1513 ySearchFunc=None,
1514 zSearchFunc=None,
1515 ):
1516 self.f_values = f_values
1517 self.x_list = (
1518 np.array(x_list)
1519 if _check_flatten(1, x_list)
1520 else np.array(x_list).flatten()
1521 )
1522 self.y_list = (
1523 np.array(y_list)
1524 if _check_flatten(1, y_list)
1525 else np.array(y_list).flatten()
1526 )
1527 self.z_list = (
1528 np.array(z_list)
1529 if _check_flatten(1, z_list)
1530 else np.array(z_list).flatten()
1531 )
1532 _check_grid_dimensions(3, self.f_values, self.x_list, self.y_list, self.z_list)
1533 self.x_n = x_list.size
1534 self.y_n = y_list.size
1535 self.z_n = z_list.size
1536 if xSearchFunc is None:
1537 xSearchFunc = np.searchsorted
1538 if ySearchFunc is None:
1539 ySearchFunc = np.searchsorted
1540 if zSearchFunc is None:
1541 zSearchFunc = np.searchsorted
1542 self.xSearchFunc = xSearchFunc
1543 self.ySearchFunc = ySearchFunc
1544 self.zSearchFunc = zSearchFunc
1546 def _evaluate(self, x, y, z):
1547 """
1548 Returns the level of the interpolated function at each value in x,y,z.
1549 Only called internally by HARKinterpolator3D.__call__ (etc).
1550 """
1551 x_pos = self.xSearchFunc(self.x_list, x)
1552 x_pos[x_pos < 1] = 1
1553 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
1554 y_pos = self.ySearchFunc(self.y_list, y)
1555 y_pos[y_pos < 1] = 1
1556 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
1557 z_pos = self.zSearchFunc(self.z_list, z)
1558 z_pos[z_pos < 1] = 1
1559 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
1560 alpha = (x - self.x_list[x_pos - 1]) / (
1561 self.x_list[x_pos] - self.x_list[x_pos - 1]
1562 )
1563 beta = (y - self.y_list[y_pos - 1]) / (
1564 self.y_list[y_pos] - self.y_list[y_pos - 1]
1565 )
1566 gamma = (z - self.z_list[z_pos - 1]) / (
1567 self.z_list[z_pos] - self.z_list[z_pos - 1]
1568 )
1569 f = (
1570 (1 - alpha)
1571 * (1 - beta)
1572 * (1 - gamma)
1573 * self.f_values[x_pos - 1, y_pos - 1, z_pos - 1]
1574 + (1 - alpha)
1575 * (1 - beta)
1576 * gamma
1577 * self.f_values[x_pos - 1, y_pos - 1, z_pos]
1578 + (1 - alpha)
1579 * beta
1580 * (1 - gamma)
1581 * self.f_values[x_pos - 1, y_pos, z_pos - 1]
1582 + (1 - alpha) * beta * gamma * self.f_values[x_pos - 1, y_pos, z_pos]
1583 + alpha
1584 * (1 - beta)
1585 * (1 - gamma)
1586 * self.f_values[x_pos, y_pos - 1, z_pos - 1]
1587 + alpha * (1 - beta) * gamma * self.f_values[x_pos, y_pos - 1, z_pos]
1588 + alpha * beta * (1 - gamma) * self.f_values[x_pos, y_pos, z_pos - 1]
1589 + alpha * beta * gamma * self.f_values[x_pos, y_pos, z_pos]
1590 )
1591 return f
1593 def _derX(self, x, y, z):
1594 """
1595 Returns the derivative with respect to x of the interpolated function
1596 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeX.
1597 """
1598 x_pos = self.xSearchFunc(self.x_list, x)
1599 x_pos[x_pos < 1] = 1
1600 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
1601 y_pos = self.ySearchFunc(self.y_list, y)
1602 y_pos[y_pos < 1] = 1
1603 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
1604 z_pos = self.zSearchFunc(self.z_list, z)
1605 z_pos[z_pos < 1] = 1
1606 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
1607 beta = (y - self.y_list[y_pos - 1]) / (
1608 self.y_list[y_pos] - self.y_list[y_pos - 1]
1609 )
1610 gamma = (z - self.z_list[z_pos - 1]) / (
1611 self.z_list[z_pos] - self.z_list[z_pos - 1]
1612 )
1613 dfdx = (
1614 (
1615 (1 - beta) * (1 - gamma) * self.f_values[x_pos, y_pos - 1, z_pos - 1]
1616 + (1 - beta) * gamma * self.f_values[x_pos, y_pos - 1, z_pos]
1617 + beta * (1 - gamma) * self.f_values[x_pos, y_pos, z_pos - 1]
1618 + beta * gamma * self.f_values[x_pos, y_pos, z_pos]
1619 )
1620 - (
1621 (1 - beta)
1622 * (1 - gamma)
1623 * self.f_values[x_pos - 1, y_pos - 1, z_pos - 1]
1624 + (1 - beta) * gamma * self.f_values[x_pos - 1, y_pos - 1, z_pos]
1625 + beta * (1 - gamma) * self.f_values[x_pos - 1, y_pos, z_pos - 1]
1626 + beta * gamma * self.f_values[x_pos - 1, y_pos, z_pos]
1627 )
1628 ) / (self.x_list[x_pos] - self.x_list[x_pos - 1])
1629 return dfdx
1631 def _derY(self, x, y, z):
1632 """
1633 Returns the derivative with respect to y of the interpolated function
1634 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeY.
1635 """
1636 x_pos = self.xSearchFunc(self.x_list, x)
1637 x_pos[x_pos < 1] = 1
1638 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
1639 y_pos = self.ySearchFunc(self.y_list, y)
1640 y_pos[y_pos < 1] = 1
1641 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
1642 z_pos = self.zSearchFunc(self.z_list, z)
1643 z_pos[z_pos < 1] = 1
1644 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
1645 alpha = (x - self.x_list[x_pos - 1]) / (
1646 self.x_list[x_pos] - self.x_list[x_pos - 1]
1647 )
1648 gamma = (z - self.z_list[z_pos - 1]) / (
1649 self.z_list[z_pos] - self.z_list[z_pos - 1]
1650 )
1651 dfdy = (
1652 (
1653 (1 - alpha) * (1 - gamma) * self.f_values[x_pos - 1, y_pos, z_pos - 1]
1654 + (1 - alpha) * gamma * self.f_values[x_pos - 1, y_pos, z_pos]
1655 + alpha * (1 - gamma) * self.f_values[x_pos, y_pos, z_pos - 1]
1656 + alpha * gamma * self.f_values[x_pos, y_pos, z_pos]
1657 )
1658 - (
1659 (1 - alpha)
1660 * (1 - gamma)
1661 * self.f_values[x_pos - 1, y_pos - 1, z_pos - 1]
1662 + (1 - alpha) * gamma * self.f_values[x_pos - 1, y_pos - 1, z_pos]
1663 + alpha * (1 - gamma) * self.f_values[x_pos, y_pos - 1, z_pos - 1]
1664 + alpha * gamma * self.f_values[x_pos, y_pos - 1, z_pos]
1665 )
1666 ) / (self.y_list[y_pos] - self.y_list[y_pos - 1])
1667 return dfdy
1669 def _derZ(self, x, y, z):
1670 """
1671 Returns the derivative with respect to z of the interpolated function
1672 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeZ.
1673 """
1674 x_pos = self.xSearchFunc(self.x_list, x)
1675 x_pos[x_pos < 1] = 1
1676 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
1677 y_pos = self.ySearchFunc(self.y_list, y)
1678 y_pos[y_pos < 1] = 1
1679 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
1680 z_pos = self.zSearchFunc(self.z_list, z)
1681 z_pos[z_pos < 1] = 1
1682 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
1683 alpha = (x - self.x_list[x_pos - 1]) / (
1684 self.x_list[x_pos] - self.x_list[x_pos - 1]
1685 )
1686 beta = (y - self.y_list[y_pos - 1]) / (
1687 self.y_list[y_pos] - self.y_list[y_pos - 1]
1688 )
1689 dfdz = (
1690 (
1691 (1 - alpha) * (1 - beta) * self.f_values[x_pos - 1, y_pos - 1, z_pos]
1692 + (1 - alpha) * beta * self.f_values[x_pos - 1, y_pos, z_pos]
1693 + alpha * (1 - beta) * self.f_values[x_pos, y_pos - 1, z_pos]
1694 + alpha * beta * self.f_values[x_pos, y_pos, z_pos]
1695 )
1696 - (
1697 (1 - alpha)
1698 * (1 - beta)
1699 * self.f_values[x_pos - 1, y_pos - 1, z_pos - 1]
1700 + (1 - alpha) * beta * self.f_values[x_pos - 1, y_pos, z_pos - 1]
1701 + alpha * (1 - beta) * self.f_values[x_pos, y_pos - 1, z_pos - 1]
1702 + alpha * beta * self.f_values[x_pos, y_pos, z_pos - 1]
1703 )
1704 ) / (self.z_list[z_pos] - self.z_list[z_pos - 1])
1705 return dfdz
1708class QuadlinearInterp(HARKinterpolator4D):
1709 """
1710 Quadlinear full (or tensor) grid interpolation of a function f(w,x,y,z).
1712 Parameters
1713 ----------
1714 f_values : numpy.array
1715 An array of size (w_n,x_n,y_n,z_n) such that f_values[i,j,k,l] =
1716 f(w_list[i],x_list[j],y_list[k],z_list[l])
1717 w_list : numpy.array
1718 An array of x values, with length designated w_n.
1719 x_list : numpy.array
1720 An array of x values, with length designated x_n.
1721 y_list : numpy.array
1722 An array of y values, with length designated y_n.
1723 z_list : numpy.array
1724 An array of z values, with length designated z_n.
1725 wSearchFunc : function
1726 An optional function that returns the reference location for w values:
1727 indices = wSearchFunc(w_list,w). Default is np.searchsorted
1728 xSearchFunc : function
1729 An optional function that returns the reference location for x values:
1730 indices = xSearchFunc(x_list,x). Default is np.searchsorted
1731 ySearchFunc : function
1732 An optional function that returns the reference location for y values:
1733 indices = ySearchFunc(y_list,y). Default is np.searchsorted
1734 zSearchFunc : function
1735 An optional function that returns the reference location for z values:
1736 indices = zSearchFunc(z_list,z). Default is np.searchsorted
1737 """
1739 distance_criteria = ["f_values", "w_list", "x_list", "y_list", "z_list"]
1741 def __init__(
1742 self,
1743 f_values,
1744 w_list,
1745 x_list,
1746 y_list,
1747 z_list,
1748 wSearchFunc=None,
1749 xSearchFunc=None,
1750 ySearchFunc=None,
1751 zSearchFunc=None,
1752 ):
1753 self.f_values = f_values
1754 self.w_list = (
1755 np.array(w_list)
1756 if _check_flatten(1, w_list)
1757 else np.array(w_list).flatten()
1758 )
1759 self.x_list = (
1760 np.array(x_list)
1761 if _check_flatten(1, x_list)
1762 else np.array(x_list).flatten()
1763 )
1764 self.y_list = (
1765 np.array(y_list)
1766 if _check_flatten(1, y_list)
1767 else np.array(y_list).flatten()
1768 )
1769 self.z_list = (
1770 np.array(z_list)
1771 if _check_flatten(1, z_list)
1772 else np.array(z_list).flatten()
1773 )
1774 _check_grid_dimensions(
1775 4, self.f_values, self.w_list, self.x_list, self.y_list, self.z_list
1776 )
1777 self.w_n = w_list.size
1778 self.x_n = x_list.size
1779 self.y_n = y_list.size
1780 self.z_n = z_list.size
1781 if wSearchFunc is None:
1782 wSearchFunc = np.searchsorted
1783 if xSearchFunc is None:
1784 xSearchFunc = np.searchsorted
1785 if ySearchFunc is None:
1786 ySearchFunc = np.searchsorted
1787 if zSearchFunc is None:
1788 zSearchFunc = np.searchsorted
1789 self.wSearchFunc = wSearchFunc
1790 self.xSearchFunc = xSearchFunc
1791 self.ySearchFunc = ySearchFunc
1792 self.zSearchFunc = zSearchFunc
1794 def _evaluate(self, w, x, y, z):
1795 """
1796 Returns the level of the interpolated function at each value in x,y,z.
1797 Only called internally by HARKinterpolator4D.__call__ (etc).
1798 """
1799 w_pos = self.wSearchFunc(self.w_list, w)
1800 w_pos[w_pos < 1] = 1
1801 w_pos[w_pos > self.w_n - 1] = self.w_n - 1
1802 x_pos = self.xSearchFunc(self.x_list, x)
1803 x_pos[x_pos < 1] = 1
1804 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
1805 y_pos = self.ySearchFunc(self.y_list, y)
1806 y_pos[y_pos < 1] = 1
1807 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
1808 z_pos = self.zSearchFunc(self.z_list, z)
1809 z_pos[z_pos < 1] = 1
1810 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
1811 i = w_pos # for convenience
1812 j = x_pos
1813 k = y_pos
1814 l = z_pos
1815 alpha = (w - self.w_list[i - 1]) / (self.w_list[i] - self.w_list[i - 1])
1816 beta = (x - self.x_list[j - 1]) / (self.x_list[j] - self.x_list[j - 1])
1817 gamma = (y - self.y_list[k - 1]) / (self.y_list[k] - self.y_list[k - 1])
1818 delta = (z - self.z_list[l - 1]) / (self.z_list[l] - self.z_list[l - 1])
1819 f = (1 - alpha) * (
1820 (1 - beta)
1821 * (
1822 (1 - gamma) * (1 - delta) * self.f_values[i - 1, j - 1, k - 1, l - 1]
1823 + (1 - gamma) * delta * self.f_values[i - 1, j - 1, k - 1, l]
1824 + gamma * (1 - delta) * self.f_values[i - 1, j - 1, k, l - 1]
1825 + gamma * delta * self.f_values[i - 1, j - 1, k, l]
1826 )
1827 + beta
1828 * (
1829 (1 - gamma) * (1 - delta) * self.f_values[i - 1, j, k - 1, l - 1]
1830 + (1 - gamma) * delta * self.f_values[i - 1, j, k - 1, l]
1831 + gamma * (1 - delta) * self.f_values[i - 1, j, k, l - 1]
1832 + gamma * delta * self.f_values[i - 1, j, k, l]
1833 )
1834 ) + alpha * (
1835 (1 - beta)
1836 * (
1837 (1 - gamma) * (1 - delta) * self.f_values[i, j - 1, k - 1, l - 1]
1838 + (1 - gamma) * delta * self.f_values[i, j - 1, k - 1, l]
1839 + gamma * (1 - delta) * self.f_values[i, j - 1, k, l - 1]
1840 + gamma * delta * self.f_values[i, j - 1, k, l]
1841 )
1842 + beta
1843 * (
1844 (1 - gamma) * (1 - delta) * self.f_values[i, j, k - 1, l - 1]
1845 + (1 - gamma) * delta * self.f_values[i, j, k - 1, l]
1846 + gamma * (1 - delta) * self.f_values[i, j, k, l - 1]
1847 + gamma * delta * self.f_values[i, j, k, l]
1848 )
1849 )
1850 return f
1852 def _derW(self, w, x, y, z):
1853 """
1854 Returns the derivative with respect to w of the interpolated function
1855 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeW.
1856 """
1857 w_pos = self.wSearchFunc(self.w_list, w)
1858 w_pos[w_pos < 1] = 1
1859 w_pos[w_pos > self.w_n - 1] = self.w_n - 1
1860 x_pos = self.xSearchFunc(self.x_list, x)
1861 x_pos[x_pos < 1] = 1
1862 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
1863 y_pos = self.ySearchFunc(self.y_list, y)
1864 y_pos[y_pos < 1] = 1
1865 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
1866 z_pos = self.zSearchFunc(self.z_list, z)
1867 z_pos[z_pos < 1] = 1
1868 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
1869 i = w_pos # for convenience
1870 j = x_pos
1871 k = y_pos
1872 l = z_pos
1873 beta = (x - self.x_list[j - 1]) / (self.x_list[j] - self.x_list[j - 1])
1874 gamma = (y - self.y_list[k - 1]) / (self.y_list[k] - self.y_list[k - 1])
1875 delta = (z - self.z_list[l - 1]) / (self.z_list[l] - self.z_list[l - 1])
1876 dfdw = (
1877 (
1878 (1 - beta)
1879 * (1 - gamma)
1880 * (1 - delta)
1881 * self.f_values[i, j - 1, k - 1, l - 1]
1882 + (1 - beta) * (1 - gamma) * delta * self.f_values[i, j - 1, k - 1, l]
1883 + (1 - beta) * gamma * (1 - delta) * self.f_values[i, j - 1, k, l - 1]
1884 + (1 - beta) * gamma * delta * self.f_values[i, j - 1, k, l]
1885 + beta * (1 - gamma) * (1 - delta) * self.f_values[i, j, k - 1, l - 1]
1886 + beta * (1 - gamma) * delta * self.f_values[i, j, k - 1, l]
1887 + beta * gamma * (1 - delta) * self.f_values[i, j, k, l - 1]
1888 + beta * gamma * delta * self.f_values[i, j, k, l]
1889 )
1890 - (
1891 (1 - beta)
1892 * (1 - gamma)
1893 * (1 - delta)
1894 * self.f_values[i - 1, j - 1, k - 1, l - 1]
1895 + (1 - beta)
1896 * (1 - gamma)
1897 * delta
1898 * self.f_values[i - 1, j - 1, k - 1, l]
1899 + (1 - beta)
1900 * gamma
1901 * (1 - delta)
1902 * self.f_values[i - 1, j - 1, k, l - 1]
1903 + (1 - beta) * gamma * delta * self.f_values[i - 1, j - 1, k, l]
1904 + beta
1905 * (1 - gamma)
1906 * (1 - delta)
1907 * self.f_values[i - 1, j, k - 1, l - 1]
1908 + beta * (1 - gamma) * delta * self.f_values[i - 1, j, k - 1, l]
1909 + beta * gamma * (1 - delta) * self.f_values[i - 1, j, k, l - 1]
1910 + beta * gamma * delta * self.f_values[i - 1, j, k, l]
1911 )
1912 ) / (self.w_list[i] - self.w_list[i - 1])
1913 return dfdw
1915 def _derX(self, w, x, y, z):
1916 """
1917 Returns the derivative with respect to x of the interpolated function
1918 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeX.
1919 """
1920 w_pos = self.wSearchFunc(self.w_list, w)
1921 w_pos[w_pos < 1] = 1
1922 w_pos[w_pos > self.w_n - 1] = self.w_n - 1
1923 x_pos = self.xSearchFunc(self.x_list, x)
1924 x_pos[x_pos < 1] = 1
1925 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
1926 y_pos = self.ySearchFunc(self.y_list, y)
1927 y_pos[y_pos < 1] = 1
1928 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
1929 z_pos = self.zSearchFunc(self.z_list, z)
1930 z_pos[z_pos < 1] = 1
1931 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
1932 i = w_pos # for convenience
1933 j = x_pos
1934 k = y_pos
1935 l = z_pos
1936 alpha = (w - self.w_list[i - 1]) / (self.w_list[i] - self.w_list[i - 1])
1937 gamma = (y - self.y_list[k - 1]) / (self.y_list[k] - self.y_list[k - 1])
1938 delta = (z - self.z_list[l - 1]) / (self.z_list[l] - self.z_list[l - 1])
1939 dfdx = (
1940 (
1941 (1 - alpha)
1942 * (1 - gamma)
1943 * (1 - delta)
1944 * self.f_values[i - 1, j, k - 1, l - 1]
1945 + (1 - alpha) * (1 - gamma) * delta * self.f_values[i - 1, j, k - 1, l]
1946 + (1 - alpha) * gamma * (1 - delta) * self.f_values[i - 1, j, k, l - 1]
1947 + (1 - alpha) * gamma * delta * self.f_values[i - 1, j, k, l]
1948 + alpha * (1 - gamma) * (1 - delta) * self.f_values[i, j, k - 1, l - 1]
1949 + alpha * (1 - gamma) * delta * self.f_values[i, j, k - 1, l]
1950 + alpha * gamma * (1 - delta) * self.f_values[i, j, k, l - 1]
1951 + alpha * gamma * delta * self.f_values[i, j, k, l]
1952 )
1953 - (
1954 (1 - alpha)
1955 * (1 - gamma)
1956 * (1 - delta)
1957 * self.f_values[i - 1, j - 1, k - 1, l - 1]
1958 + (1 - alpha)
1959 * (1 - gamma)
1960 * delta
1961 * self.f_values[i - 1, j - 1, k - 1, l]
1962 + (1 - alpha)
1963 * gamma
1964 * (1 - delta)
1965 * self.f_values[i - 1, j - 1, k, l - 1]
1966 + (1 - alpha) * gamma * delta * self.f_values[i - 1, j - 1, k, l]
1967 + alpha
1968 * (1 - gamma)
1969 * (1 - delta)
1970 * self.f_values[i, j - 1, k - 1, l - 1]
1971 + alpha * (1 - gamma) * delta * self.f_values[i, j - 1, k - 1, l]
1972 + alpha * gamma * (1 - delta) * self.f_values[i, j - 1, k, l - 1]
1973 + alpha * gamma * delta * self.f_values[i, j - 1, k, l]
1974 )
1975 ) / (self.x_list[j] - self.x_list[j - 1])
1976 return dfdx
1978 def _derY(self, w, x, y, z):
1979 """
1980 Returns the derivative with respect to y of the interpolated function
1981 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeY.
1982 """
1983 w_pos = self.wSearchFunc(self.w_list, w)
1984 w_pos[w_pos < 1] = 1
1985 w_pos[w_pos > self.w_n - 1] = self.w_n - 1
1986 x_pos = self.xSearchFunc(self.x_list, x)
1987 x_pos[x_pos < 1] = 1
1988 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
1989 y_pos = self.ySearchFunc(self.y_list, y)
1990 y_pos[y_pos < 1] = 1
1991 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
1992 z_pos = self.zSearchFunc(self.z_list, z)
1993 z_pos[z_pos < 1] = 1
1994 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
1995 i = w_pos # for convenience
1996 j = x_pos
1997 k = y_pos
1998 l = z_pos
1999 alpha = (w - self.w_list[i - 1]) / (self.w_list[i] - self.w_list[i - 1])
2000 beta = (x - self.x_list[j - 1]) / (self.x_list[j] - self.x_list[j - 1])
2001 delta = (z - self.z_list[l - 1]) / (self.z_list[l] - self.z_list[l - 1])
2002 dfdy = (
2003 (
2004 (1 - alpha)
2005 * (1 - beta)
2006 * (1 - delta)
2007 * self.f_values[i - 1, j - 1, k, l - 1]
2008 + (1 - alpha) * (1 - beta) * delta * self.f_values[i - 1, j - 1, k, l]
2009 + (1 - alpha) * beta * (1 - delta) * self.f_values[i - 1, j, k, l - 1]
2010 + (1 - alpha) * beta * delta * self.f_values[i - 1, j, k, l]
2011 + alpha * (1 - beta) * (1 - delta) * self.f_values[i, j - 1, k, l - 1]
2012 + alpha * (1 - beta) * delta * self.f_values[i, j - 1, k, l]
2013 + alpha * beta * (1 - delta) * self.f_values[i, j, k, l - 1]
2014 + alpha * beta * delta * self.f_values[i, j, k, l]
2015 )
2016 - (
2017 (1 - alpha)
2018 * (1 - beta)
2019 * (1 - delta)
2020 * self.f_values[i - 1, j - 1, k - 1, l - 1]
2021 + (1 - alpha)
2022 * (1 - beta)
2023 * delta
2024 * self.f_values[i - 1, j - 1, k - 1, l]
2025 + (1 - alpha)
2026 * beta
2027 * (1 - delta)
2028 * self.f_values[i - 1, j, k - 1, l - 1]
2029 + (1 - alpha) * beta * delta * self.f_values[i - 1, j, k - 1, l]
2030 + alpha
2031 * (1 - beta)
2032 * (1 - delta)
2033 * self.f_values[i, j - 1, k - 1, l - 1]
2034 + alpha * (1 - beta) * delta * self.f_values[i, j - 1, k - 1, l]
2035 + alpha * beta * (1 - delta) * self.f_values[i, j, k - 1, l - 1]
2036 + alpha * beta * delta * self.f_values[i, j, k - 1, l]
2037 )
2038 ) / (self.y_list[k] - self.y_list[k - 1])
2039 return dfdy
2041 def _derZ(self, w, x, y, z):
2042 """
2043 Returns the derivative with respect to z of the interpolated function
2044 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeZ.
2045 """
2046 w_pos = self.wSearchFunc(self.w_list, w)
2047 w_pos[w_pos < 1] = 1
2048 w_pos[w_pos > self.w_n - 1] = self.w_n - 1
2049 x_pos = self.xSearchFunc(self.x_list, x)
2050 x_pos[x_pos < 1] = 1
2051 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
2052 y_pos = self.ySearchFunc(self.y_list, y)
2053 y_pos[y_pos < 1] = 1
2054 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
2055 z_pos = self.zSearchFunc(self.z_list, z)
2056 z_pos[z_pos < 1] = 1
2057 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
2058 i = w_pos # for convenience
2059 j = x_pos
2060 k = y_pos
2061 l = z_pos
2062 alpha = (w - self.w_list[i - 1]) / (self.w_list[i] - self.w_list[i - 1])
2063 beta = (x - self.x_list[j - 1]) / (self.x_list[j] - self.x_list[j - 1])
2064 gamma = (y - self.y_list[k - 1]) / (self.y_list[k] - self.y_list[k - 1])
2065 dfdz = (
2066 (
2067 (1 - alpha)
2068 * (1 - beta)
2069 * (1 - gamma)
2070 * self.f_values[i - 1, j - 1, k - 1, l]
2071 + (1 - alpha) * (1 - beta) * gamma * self.f_values[i - 1, j - 1, k, l]
2072 + (1 - alpha) * beta * (1 - gamma) * self.f_values[i - 1, j, k - 1, l]
2073 + (1 - alpha) * beta * gamma * self.f_values[i - 1, j, k, l]
2074 + alpha * (1 - beta) * (1 - gamma) * self.f_values[i, j - 1, k - 1, l]
2075 + alpha * (1 - beta) * gamma * self.f_values[i, j - 1, k, l]
2076 + alpha * beta * (1 - gamma) * self.f_values[i, j, k - 1, l]
2077 + alpha * beta * gamma * self.f_values[i, j, k, l]
2078 )
2079 - (
2080 (1 - alpha)
2081 * (1 - beta)
2082 * (1 - gamma)
2083 * self.f_values[i - 1, j - 1, k - 1, l - 1]
2084 + (1 - alpha)
2085 * (1 - beta)
2086 * gamma
2087 * self.f_values[i - 1, j - 1, k, l - 1]
2088 + (1 - alpha)
2089 * beta
2090 * (1 - gamma)
2091 * self.f_values[i - 1, j, k - 1, l - 1]
2092 + (1 - alpha) * beta * gamma * self.f_values[i - 1, j, k, l - 1]
2093 + alpha
2094 * (1 - beta)
2095 * (1 - gamma)
2096 * self.f_values[i, j - 1, k - 1, l - 1]
2097 + alpha * (1 - beta) * gamma * self.f_values[i, j - 1, k, l - 1]
2098 + alpha * beta * (1 - gamma) * self.f_values[i, j, k - 1, l - 1]
2099 + alpha * beta * gamma * self.f_values[i, j, k, l - 1]
2100 )
2101 ) / (self.z_list[l] - self.z_list[l - 1])
2102 return dfdz
2105class LowerEnvelope(HARKinterpolator1D):
2106 """
2107 The lower envelope of a finite set of 1D functions, each of which can be of
2108 any class that has the methods __call__, derivative, and eval_with_derivative.
2109 Generally: it combines HARKinterpolator1Ds.
2111 Parameters
2112 ----------
2113 *functions : function
2114 Any number of real functions; often instances of HARKinterpolator1D
2115 nan_bool : boolean
2116 An indicator for whether the solver should exclude NA's when
2117 forming the lower envelope
2118 """
2120 distance_criteria = ["functions"]
2122 def __init__(self, *functions, nan_bool=True):
2123 if nan_bool:
2124 self.compare = np.nanmin
2125 self.argcompare = np.nanargmin
2126 else:
2127 self.compare = np.min
2128 self.argcompare = np.argmin
2130 self.functions = []
2131 for function in functions:
2132 self.functions.append(function)
2133 self.funcCount = len(self.functions)
2135 def _evaluate(self, x):
2136 """
2137 Returns the level of the function at each value in x as the minimum among
2138 all of the functions. Only called internally by HARKinterpolator1D.__call__.
2139 """
2140 m = len(x)
2141 fx = np.zeros((m, self.funcCount))
2142 for j in range(self.funcCount):
2143 fx[:, j] = self.functions[j](x)
2144 y = self.compare(fx, axis=1)
2145 return y
2147 def _der(self, x):
2148 """
2149 Returns the first derivative of the function at each value in x. Only
2150 called internally by HARKinterpolator1D.derivative.
2151 """
2152 y, dydx = self._evalAndDer(x)
2153 return dydx # Sadly, this is the fastest / most convenient way...
2155 def _evalAndDer(self, x):
2156 """
2157 Returns the level and first derivative of the function at each value in
2158 x. Only called internally by HARKinterpolator1D.eval_and_der.
2159 """
2160 m = len(x)
2161 fx = np.zeros((m, self.funcCount))
2162 for j in range(self.funcCount):
2163 fx[:, j] = self.functions[j](x)
2164 i = self.argcompare(fx, axis=1)
2165 y = fx[np.arange(m), i]
2166 dydx = np.zeros_like(y)
2167 for j in np.unique(i):
2168 c = i == j
2169 dydx[c] = self.functions[j].derivative(x[c])
2170 return y, dydx
2173class UpperEnvelope(HARKinterpolator1D):
2174 """
2175 The upper envelope of a finite set of 1D functions, each of which can be of
2176 any class that has the methods __call__, derivative, and eval_with_derivative.
2177 Generally: it combines HARKinterpolator1Ds.
2179 Parameters
2180 ----------
2181 *functions : function
2182 Any number of real functions; often instances of HARKinterpolator1D
2183 nan_bool : boolean
2184 An indicator for whether the solver should exclude NA's when forming
2185 the lower envelope.
2186 """
2188 distance_criteria = ["functions"]
2190 def __init__(self, *functions, nan_bool=True):
2191 if nan_bool:
2192 self.compare = np.nanmax
2193 self.argcompare = np.nanargmax
2194 else:
2195 self.compare = np.max
2196 self.argcompare = np.argmax
2197 self.functions = []
2198 for function in functions:
2199 self.functions.append(function)
2200 self.funcCount = len(self.functions)
2202 def _evaluate(self, x):
2203 """
2204 Returns the level of the function at each value in x as the maximum among
2205 all of the functions. Only called internally by HARKinterpolator1D.__call__.
2206 """
2207 m = len(x)
2208 fx = np.zeros((m, self.funcCount))
2209 for j in range(self.funcCount):
2210 fx[:, j] = self.functions[j](x)
2211 y = self.compare(fx, axis=1)
2212 return y
2214 def _der(self, x):
2215 """
2216 Returns the first derivative of the function at each value in x. Only
2217 called internally by HARKinterpolator1D.derivative.
2218 """
2219 y, dydx = self._evalAndDer(x)
2220 return dydx # Sadly, this is the fastest / most convenient way...
2222 def _evalAndDer(self, x):
2223 """
2224 Returns the level and first derivative of the function at each value in
2225 x. Only called internally by HARKinterpolator1D.eval_and_der.
2226 """
2227 m = len(x)
2228 fx = np.zeros((m, self.funcCount))
2229 for j in range(self.funcCount):
2230 fx[:, j] = self.functions[j](x)
2231 i = self.argcompare(fx, axis=1)
2232 y = fx[np.arange(m), i]
2233 dydx = np.zeros_like(y)
2234 for j in np.unique(i):
2235 c = i == j
2236 dydx[c] = self.functions[j].derivative(x[c])
2237 return y, dydx
2240class LowerEnvelope2D(HARKinterpolator2D):
2241 """
2242 The lower envelope of a finite set of 2D functions, each of which can be of
2243 any class that has the methods __call__, derivativeX, and derivativeY.
2244 Generally: it combines HARKinterpolator2Ds.
2246 Parameters
2247 ----------
2248 *functions : function
2249 Any number of real functions; often instances of HARKinterpolator2D
2250 nan_bool : boolean
2251 An indicator for whether the solver should exclude NA's when forming
2252 the lower envelope.
2253 """
2255 distance_criteria = ["functions"]
2257 def __init__(self, *functions, nan_bool=True):
2258 if nan_bool:
2259 self.compare = np.nanmin
2260 self.argcompare = np.nanargmin
2261 else:
2262 self.compare = np.min
2263 self.argcompare = np.argmin
2264 self.functions = []
2265 for function in functions:
2266 self.functions.append(function)
2267 self.funcCount = len(self.functions)
2269 def _evaluate(self, x, y):
2270 """
2271 Returns the level of the function at each value in (x,y) as the minimum
2272 among all of the functions. Only called internally by
2273 HARKinterpolator2D.__call__.
2274 """
2275 m = len(x)
2276 temp = np.zeros((m, self.funcCount))
2277 for j in range(self.funcCount):
2278 temp[:, j] = self.functions[j](x, y)
2279 f = self.compare(temp, axis=1)
2280 return f
2282 def _derX(self, x, y):
2283 """
2284 Returns the first derivative of the function with respect to X at each
2285 value in (x,y). Only called internally by HARKinterpolator2D._derX.
2286 """
2287 m = len(x)
2288 temp = np.zeros((m, self.funcCount))
2289 for j in range(self.funcCount):
2290 temp[:, j] = self.functions[j](x, y)
2291 i = self.argcompare(temp, axis=1)
2292 dfdx = np.zeros_like(x)
2293 for j in np.unique(i):
2294 c = i == j
2295 dfdx[c] = self.functions[j].derivativeX(x[c], y[c])
2296 return dfdx
2298 def _derY(self, x, y):
2299 """
2300 Returns the first derivative of the function with respect to Y at each
2301 value in (x,y). Only called internally by HARKinterpolator2D._derY.
2302 """
2303 m = len(x)
2304 temp = np.zeros((m, self.funcCount))
2305 for j in range(self.funcCount):
2306 temp[:, j] = self.functions[j](x, y)
2307 i = self.argcompare(temp, axis=1)
2308 y = temp[np.arange(m), i]
2309 dfdy = np.zeros_like(x)
2310 for j in np.unique(i):
2311 c = i == j
2312 dfdy[c] = self.functions[j].derivativeY(x[c], y[c])
2313 return dfdy
2316class LowerEnvelope3D(HARKinterpolator3D):
2317 """
2318 The lower envelope of a finite set of 3D functions, each of which can be of
2319 any class that has the methods __call__, derivativeX, derivativeY, and
2320 derivativeZ. Generally: it combines HARKinterpolator2Ds.
2322 Parameters
2323 ----------
2324 *functions : function
2325 Any number of real functions; often instances of HARKinterpolator3D
2326 nan_bool : boolean
2327 An indicator for whether the solver should exclude NA's when forming
2328 the lower envelope.
2329 """
2331 distance_criteria = ["functions"]
2333 def __init__(self, *functions, nan_bool=True):
2334 if nan_bool:
2335 self.compare = np.nanmin
2336 self.argcompare = np.nanargmin
2337 else:
2338 self.compare = np.min
2339 self.argcompare = np.argmin
2340 self.functions = []
2341 for function in functions:
2342 self.functions.append(function)
2343 self.funcCount = len(self.functions)
2345 def _evaluate(self, x, y, z):
2346 """
2347 Returns the level of the function at each value in (x,y,z) as the minimum
2348 among all of the functions. Only called internally by
2349 HARKinterpolator3D.__call__.
2350 """
2351 m = len(x)
2352 temp = np.zeros((m, self.funcCount))
2353 for j in range(self.funcCount):
2354 temp[:, j] = self.functions[j](x, y, z)
2355 f = self.compare(temp, axis=1)
2356 return f
2358 def _derX(self, x, y, z):
2359 """
2360 Returns the first derivative of the function with respect to X at each
2361 value in (x,y,z). Only called internally by HARKinterpolator3D._derX.
2362 """
2363 m = len(x)
2364 temp = np.zeros((m, self.funcCount))
2365 for j in range(self.funcCount):
2366 temp[:, j] = self.functions[j](x, y, z)
2367 i = self.argcompare(temp, axis=1)
2368 dfdx = np.zeros_like(x)
2369 for j in np.unique(i):
2370 c = i == j
2371 dfdx[c] = self.functions[j].derivativeX(x[c], y[c], z[c])
2372 return dfdx
2374 def _derY(self, x, y, z):
2375 """
2376 Returns the first derivative of the function with respect to Y at each
2377 value in (x,y,z). Only called internally by HARKinterpolator3D._derY.
2378 """
2379 m = len(x)
2380 temp = np.zeros((m, self.funcCount))
2381 for j in range(self.funcCount):
2382 temp[:, j] = self.functions[j](x, y, z)
2383 i = self.argcompare(temp, axis=1)
2384 y = temp[np.arange(m), i]
2385 dfdy = np.zeros_like(x)
2386 for j in np.unique(i):
2387 c = i == j
2388 dfdy[c] = self.functions[j].derivativeY(x[c], y[c], z[c])
2389 return dfdy
2391 def _derZ(self, x, y, z):
2392 """
2393 Returns the first derivative of the function with respect to Z at each
2394 value in (x,y,z). Only called internally by HARKinterpolator3D._derZ.
2395 """
2396 m = len(x)
2397 temp = np.zeros((m, self.funcCount))
2398 for j in range(self.funcCount):
2399 temp[:, j] = self.functions[j](x, y, z)
2400 i = self.argcompare(temp, axis=1)
2401 y = temp[np.arange(m), i]
2402 dfdz = np.zeros_like(x)
2403 for j in np.unique(i):
2404 c = i == j
2405 dfdz[c] = self.functions[j].derivativeZ(x[c], y[c], z[c])
2406 return dfdz
2409class VariableLowerBoundFunc2D(MetricObject):
2410 """
2411 A class for representing a function with two real inputs whose lower bound
2412 in the first input depends on the second input. Useful for managing curved
2413 natural borrowing constraints, as occurs in the persistent shocks model.
2415 Parameters
2416 ----------
2417 func : function
2418 A function f: (R_+ x R) --> R representing the function of interest
2419 shifted by its lower bound in the first input.
2420 lowerBound : function
2421 The lower bound in the first input of the function of interest, as
2422 a function of the second input.
2423 """
2425 distance_criteria = ["func", "lowerBound"]
2427 def __init__(self, func, lowerBound):
2428 self.func = func
2429 self.lowerBound = lowerBound
2431 def __call__(self, x, y):
2432 """
2433 Evaluate the function at given state space points.
2435 Parameters
2436 ----------
2437 x : np.array
2438 First input values.
2439 y : np.array
2440 Second input values; should be of same shape as x.
2442 Returns
2443 -------
2444 f_out : np.array
2445 Function evaluated at (x,y), of same shape as inputs.
2446 """
2447 xShift = self.lowerBound(y)
2448 f_out = self.func(x - xShift, y)
2449 return f_out
2451 def derivativeX(self, x, y):
2452 """
2453 Evaluate the first derivative with respect to x of the function at given
2454 state space points.
2456 Parameters
2457 ----------
2458 x : np.array
2459 First input values.
2460 y : np.array
2461 Second input values; should be of same shape as x.
2463 Returns
2464 -------
2465 dfdx_out : np.array
2466 First derivative of function with respect to the first input,
2467 evaluated at (x,y), of same shape as inputs.
2468 """
2469 xShift = self.lowerBound(y)
2470 dfdx_out = self.func.derivativeX(x - xShift, y)
2471 return dfdx_out
2473 def derivativeY(self, x, y):
2474 """
2475 Evaluate the first derivative with respect to y of the function at given
2476 state space points.
2478 Parameters
2479 ----------
2480 x : np.array
2481 First input values.
2482 y : np.array
2483 Second input values; should be of same shape as x.
2485 Returns
2486 -------
2487 dfdy_out : np.array
2488 First derivative of function with respect to the second input,
2489 evaluated at (x,y), of same shape as inputs.
2490 """
2491 xShift, xShiftDer = self.lowerBound.eval_with_derivative(y)
2492 dfdy_out = self.func.derivativeY(
2493 x - xShift, y
2494 ) - xShiftDer * self.func.derivativeX(x - xShift, y)
2495 return dfdy_out
2498class VariableLowerBoundFunc3D(MetricObject):
2499 """
2500 A class for representing a function with three real inputs whose lower bound
2501 in the first input depends on the second input. Useful for managing curved
2502 natural borrowing constraints.
2504 Parameters
2505 ----------
2506 func : function
2507 A function f: (R_+ x R^2) --> R representing the function of interest
2508 shifted by its lower bound in the first input.
2509 lowerBound : function
2510 The lower bound in the first input of the function of interest, as
2511 a function of the second input.
2512 """
2514 distance_criteria = ["func", "lowerBound"]
2516 def __init__(self, func, lowerBound):
2517 self.func = func
2518 self.lowerBound = lowerBound
2520 def __call__(self, x, y, z):
2521 """
2522 Evaluate the function at given state space points.
2524 Parameters
2525 ----------
2526 x : np.array
2527 First input values.
2528 y : np.array
2529 Second input values; should be of same shape as x.
2530 z : np.array
2531 Third input values; should be of same shape as x.
2533 Returns
2534 -------
2535 f_out : np.array
2536 Function evaluated at (x,y,z), of same shape as inputs.
2537 """
2538 xShift = self.lowerBound(y)
2539 f_out = self.func(x - xShift, y, z)
2540 return f_out
2542 def derivativeX(self, x, y, z):
2543 """
2544 Evaluate the first derivative with respect to x of the function at given
2545 state space points.
2547 Parameters
2548 ----------
2549 x : np.array
2550 First input values.
2551 y : np.array
2552 Second input values; should be of same shape as x.
2553 z : np.array
2554 Third input values; should be of same shape as x.
2556 Returns
2557 -------
2558 dfdx_out : np.array
2559 First derivative of function with respect to the first input,
2560 evaluated at (x,y,z), of same shape as inputs.
2561 """
2562 xShift = self.lowerBound(y)
2563 dfdx_out = self.func.derivativeX(x - xShift, y, z)
2564 return dfdx_out
2566 def derivativeY(self, x, y, z):
2567 """
2568 Evaluate the first derivative with respect to y of the function at given
2569 state space points.
2571 Parameters
2572 ----------
2573 x : np.array
2574 First input values.
2575 y : np.array
2576 Second input values; should be of same shape as x.
2577 z : np.array
2578 Third input values; should be of same shape as x.
2580 Returns
2581 -------
2582 dfdy_out : np.array
2583 First derivative of function with respect to the second input,
2584 evaluated at (x,y,z), of same shape as inputs.
2585 """
2586 xShift, xShiftDer = self.lowerBound.eval_with_derivative(y)
2587 dfdy_out = self.func.derivativeY(
2588 x - xShift, y, z
2589 ) - xShiftDer * self.func.derivativeX(x - xShift, y, z)
2590 return dfdy_out
2592 def derivativeZ(self, x, y, z):
2593 """
2594 Evaluate the first derivative with respect to z of the function at given
2595 state space points.
2597 Parameters
2598 ----------
2599 x : np.array
2600 First input values.
2601 y : np.array
2602 Second input values; should be of same shape as x.
2603 z : np.array
2604 Third input values; should be of same shape as x.
2606 Returns
2607 -------
2608 dfdz_out : np.array
2609 First derivative of function with respect to the third input,
2610 evaluated at (x,y,z), of same shape as inputs.
2611 """
2612 xShift = self.lowerBound(y)
2613 dfdz_out = self.func.derivativeZ(x - xShift, y, z)
2614 return dfdz_out
2617class LinearInterpOnInterp1D(HARKinterpolator2D):
2618 """
2619 A 2D interpolator that linearly interpolates among a list of 1D interpolators.
2621 Parameters
2622 ----------
2623 xInterpolators : [HARKinterpolator1D]
2624 A list of 1D interpolations over the x variable. The nth element of
2625 xInterpolators represents f(x,y_values[n]).
2626 y_values: numpy.array
2627 An array of y values equal in length to xInterpolators.
2628 """
2630 distance_criteria = ["xInterpolators", "y_list"]
2632 def __init__(self, xInterpolators, y_values):
2633 self.xInterpolators = xInterpolators
2634 self.y_list = y_values
2635 self.y_n = y_values.size
2637 def _evaluate(self, x, y):
2638 """
2639 Returns the level of the interpolated function at each value in x,y.
2640 Only called internally by HARKinterpolator2D.__call__ (etc).
2641 """
2642 m = len(x)
2643 y_pos = np.searchsorted(self.y_list, y)
2644 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
2645 y_pos[y_pos < 1] = 1
2646 f = np.zeros(m) + np.nan
2647 if y.size > 0:
2648 for i in range(1, self.y_n):
2649 c = y_pos == i
2650 if np.any(c):
2651 alpha = (y[c] - self.y_list[i - 1]) / (
2652 self.y_list[i] - self.y_list[i - 1]
2653 )
2654 f[c] = (1 - alpha) * self.xInterpolators[i - 1](
2655 x[c]
2656 ) + alpha * self.xInterpolators[i](x[c])
2657 return f
2659 def _derX(self, x, y):
2660 """
2661 Returns the derivative with respect to x of the interpolated function
2662 at each value in x,y. Only called internally by HARKinterpolator2D.derivativeX.
2663 """
2664 m = len(x)
2665 y_pos = np.searchsorted(self.y_list, y)
2666 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
2667 y_pos[y_pos < 1] = 1
2668 dfdx = np.zeros(m) + np.nan
2669 if y.size > 0:
2670 for i in range(1, self.y_n):
2671 c = y_pos == i
2672 if np.any(c):
2673 alpha = (y[c] - self.y_list[i - 1]) / (
2674 self.y_list[i] - self.y_list[i - 1]
2675 )
2676 dfdx[c] = (1 - alpha) * self.xInterpolators[i - 1]._der(
2677 x[c]
2678 ) + alpha * self.xInterpolators[i]._der(x[c])
2679 return dfdx
2681 def _derY(self, x, y):
2682 """
2683 Returns the derivative with respect to y of the interpolated function
2684 at each value in x,y. Only called internally by HARKinterpolator2D.derivativeY.
2685 """
2686 m = len(x)
2687 y_pos = np.searchsorted(self.y_list, y)
2688 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
2689 y_pos[y_pos < 1] = 1
2690 dfdy = np.zeros(m) + np.nan
2691 if y.size > 0:
2692 for i in range(1, self.y_n):
2693 c = y_pos == i
2694 if np.any(c):
2695 dfdy[c] = (
2696 self.xInterpolators[i](x[c]) - self.xInterpolators[i - 1](x[c])
2697 ) / (self.y_list[i] - self.y_list[i - 1])
2698 return dfdy
2701class BilinearInterpOnInterp1D(HARKinterpolator3D):
2702 """
2703 A 3D interpolator that bilinearly interpolates among a list of lists of 1D
2704 interpolators.
2706 Constructor for the class, generating an approximation to a function of
2707 the form f(x,y,z) using interpolations over f(x,y_0,z_0) for a fixed grid
2708 of y_0 and z_0 values.
2710 Parameters
2711 ----------
2712 xInterpolators : [[HARKinterpolator1D]]
2713 A list of lists of 1D interpolations over the x variable. The i,j-th
2714 element of xInterpolators represents f(x,y_values[i],z_values[j]).
2715 y_values: numpy.array
2716 An array of y values equal in length to xInterpolators.
2717 z_values: numpy.array
2718 An array of z values equal in length to xInterpolators[0].
2719 """
2721 distance_criteria = ["xInterpolators", "y_list", "z_list"]
2723 def __init__(self, xInterpolators, y_values, z_values):
2724 self.xInterpolators = xInterpolators
2725 self.y_list = y_values
2726 self.y_n = y_values.size
2727 self.z_list = z_values
2728 self.z_n = z_values.size
2730 def _evaluate(self, x, y, z):
2731 """
2732 Returns the level of the interpolated function at each value in x,y,z.
2733 Only called internally by HARKinterpolator3D.__call__ (etc).
2735 Optimized to avoid nested loops by processing all unique (i,j) combinations
2736 with vectorized operations.
2737 """
2738 m = len(x)
2739 y_pos = np.searchsorted(self.y_list, y)
2740 y_pos = np.clip(y_pos, 1, self.y_n - 1)
2741 z_pos = np.searchsorted(self.z_list, z)
2742 z_pos = np.clip(z_pos, 1, self.z_n - 1)
2744 f = np.full(m, np.nan)
2746 # Find unique combinations of (y_pos, z_pos) to avoid redundant computations
2747 unique_pairs = np.unique(np.column_stack((y_pos, z_pos)), axis=0)
2749 for i, j in unique_pairs:
2750 c = (i == y_pos) & (j == z_pos)
2751 alpha = (y[c] - self.y_list[i - 1]) / (self.y_list[i] - self.y_list[i - 1])
2752 beta = (z[c] - self.z_list[j - 1]) / (self.z_list[j] - self.z_list[j - 1])
2753 f[c] = (
2754 (1 - alpha) * (1 - beta) * self.xInterpolators[i - 1][j - 1](x[c])
2755 + (1 - alpha) * beta * self.xInterpolators[i - 1][j](x[c])
2756 + alpha * (1 - beta) * self.xInterpolators[i][j - 1](x[c])
2757 + alpha * beta * self.xInterpolators[i][j](x[c])
2758 )
2759 return f
2761 def _derX(self, x, y, z):
2762 """
2763 Returns the derivative with respect to x of the interpolated function
2764 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeX.
2766 Optimized to avoid nested loops by processing unique (i,j) combinations.
2767 """
2768 m = len(x)
2769 y_pos = np.searchsorted(self.y_list, y)
2770 y_pos = np.clip(y_pos, 1, self.y_n - 1)
2771 z_pos = np.searchsorted(self.z_list, z)
2772 z_pos = np.clip(z_pos, 1, self.z_n - 1)
2774 dfdx = np.full(m, np.nan)
2776 # Find unique combinations to avoid redundant computations
2777 unique_pairs = np.unique(np.column_stack((y_pos, z_pos)), axis=0)
2779 for i, j in unique_pairs:
2780 c = (i == y_pos) & (j == z_pos)
2781 alpha = (y[c] - self.y_list[i - 1]) / (self.y_list[i] - self.y_list[i - 1])
2782 beta = (z[c] - self.z_list[j - 1]) / (self.z_list[j] - self.z_list[j - 1])
2783 dfdx[c] = (
2784 (1 - alpha) * (1 - beta) * self.xInterpolators[i - 1][j - 1]._der(x[c])
2785 + (1 - alpha) * beta * self.xInterpolators[i - 1][j]._der(x[c])
2786 + alpha * (1 - beta) * self.xInterpolators[i][j - 1]._der(x[c])
2787 + alpha * beta * self.xInterpolators[i][j]._der(x[c])
2788 )
2789 return dfdx
2791 def _derY(self, x, y, z):
2792 """
2793 Returns the derivative with respect to y of the interpolated function
2794 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeY.
2796 Optimized to avoid nested loops by processing unique (i,j) combinations.
2797 """
2798 m = len(x)
2799 y_pos = np.searchsorted(self.y_list, y)
2800 y_pos = np.clip(y_pos, 1, self.y_n - 1)
2801 z_pos = np.searchsorted(self.z_list, z)
2802 z_pos = np.clip(z_pos, 1, self.z_n - 1)
2804 dfdy = np.full(m, np.nan)
2806 # Find unique combinations to avoid redundant computations
2807 unique_pairs = np.unique(np.column_stack((y_pos, z_pos)), axis=0)
2809 for i, j in unique_pairs:
2810 c = (i == y_pos) & (j == z_pos)
2811 beta = (z[c] - self.z_list[j - 1]) / (self.z_list[j] - self.z_list[j - 1])
2812 dfdy[c] = (
2813 (
2814 (1 - beta) * self.xInterpolators[i][j - 1](x[c])
2815 + beta * self.xInterpolators[i][j](x[c])
2816 )
2817 - (
2818 (1 - beta) * self.xInterpolators[i - 1][j - 1](x[c])
2819 + beta * self.xInterpolators[i - 1][j](x[c])
2820 )
2821 ) / (self.y_list[i] - self.y_list[i - 1])
2822 return dfdy
2824 def _derZ(self, x, y, z):
2825 """
2826 Returns the derivative with respect to z of the interpolated function
2827 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeZ.
2829 Optimized to avoid nested loops by processing unique (i,j) combinations.
2830 """
2831 m = len(x)
2832 y_pos = np.searchsorted(self.y_list, y)
2833 y_pos = np.clip(y_pos, 1, self.y_n - 1)
2834 z_pos = np.searchsorted(self.z_list, z)
2835 z_pos = np.clip(z_pos, 1, self.z_n - 1)
2837 dfdz = np.full(m, np.nan)
2839 # Find unique combinations to avoid redundant computations
2840 unique_pairs = np.unique(np.column_stack((y_pos, z_pos)), axis=0)
2842 for i, j in unique_pairs:
2843 c = (i == y_pos) & (j == z_pos)
2844 alpha = (y[c] - self.y_list[i - 1]) / (self.y_list[i] - self.y_list[i - 1])
2845 dfdz[c] = (
2846 (
2847 (1 - alpha) * self.xInterpolators[i - 1][j](x[c])
2848 + alpha * self.xInterpolators[i][j](x[c])
2849 )
2850 - (
2851 (1 - alpha) * self.xInterpolators[i - 1][j - 1](x[c])
2852 + alpha * self.xInterpolators[i][j - 1](x[c])
2853 )
2854 ) / (self.z_list[j] - self.z_list[j - 1])
2855 return dfdz
2858class TrilinearInterpOnInterp1D(HARKinterpolator4D):
2859 """
2860 A 4D interpolator that trilinearly interpolates among a list of lists of 1D interpolators.
2862 Constructor for the class, generating an approximation to a function of
2863 the form f(w,x,y,z) using interpolations over f(w,x_0,y_0,z_0) for a fixed
2864 grid of y_0 and z_0 values.
2866 Parameters
2867 ----------
2868 wInterpolators : [[[HARKinterpolator1D]]]
2869 A list of lists of lists of 1D interpolations over the x variable.
2870 The i,j,k-th element of wInterpolators represents f(w,x_values[i],y_values[j],z_values[k]).
2871 x_values: numpy.array
2872 An array of x values equal in length to wInterpolators.
2873 y_values: numpy.array
2874 An array of y values equal in length to wInterpolators[0].
2875 z_values: numpy.array
2876 An array of z values equal in length to wInterpolators[0][0]
2877 """
2879 distance_criteria = ["wInterpolators", "x_list", "y_list", "z_list"]
2881 def __init__(self, wInterpolators, x_values, y_values, z_values):
2882 self.wInterpolators = wInterpolators
2883 self.x_list = x_values
2884 self.x_n = x_values.size
2885 self.y_list = y_values
2886 self.y_n = y_values.size
2887 self.z_list = z_values
2888 self.z_n = z_values.size
2890 def _evaluate(self, w, x, y, z):
2891 """
2892 Returns the level of the interpolated function at each value in w,x,y,z.
2893 Only called internally by HARKinterpolator4D.__call__ (etc).
2894 """
2895 m = len(x)
2896 x_pos = np.searchsorted(self.x_list, x)
2897 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
2898 y_pos = np.searchsorted(self.y_list, y)
2899 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
2900 y_pos[y_pos < 1] = 1
2901 z_pos = np.searchsorted(self.z_list, z)
2902 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
2903 z_pos[z_pos < 1] = 1
2904 f = np.zeros(m) + np.nan
2905 for i in range(1, self.x_n):
2906 for j in range(1, self.y_n):
2907 for k in range(1, self.z_n):
2908 c = np.logical_and(
2909 np.logical_and(i == x_pos, j == y_pos), k == z_pos
2910 )
2911 if np.any(c):
2912 alpha = (x[c] - self.x_list[i - 1]) / (
2913 self.x_list[i] - self.x_list[i - 1]
2914 )
2915 beta = (y[c] - self.y_list[j - 1]) / (
2916 self.y_list[j] - self.y_list[j - 1]
2917 )
2918 gamma = (z[c] - self.z_list[k - 1]) / (
2919 self.z_list[k] - self.z_list[k - 1]
2920 )
2921 f[c] = (
2922 (1 - alpha)
2923 * (1 - beta)
2924 * (1 - gamma)
2925 * self.wInterpolators[i - 1][j - 1][k - 1](w[c])
2926 + (1 - alpha)
2927 * (1 - beta)
2928 * gamma
2929 * self.wInterpolators[i - 1][j - 1][k](w[c])
2930 + (1 - alpha)
2931 * beta
2932 * (1 - gamma)
2933 * self.wInterpolators[i - 1][j][k - 1](w[c])
2934 + (1 - alpha)
2935 * beta
2936 * gamma
2937 * self.wInterpolators[i - 1][j][k](w[c])
2938 + alpha
2939 * (1 - beta)
2940 * (1 - gamma)
2941 * self.wInterpolators[i][j - 1][k - 1](w[c])
2942 + alpha
2943 * (1 - beta)
2944 * gamma
2945 * self.wInterpolators[i][j - 1][k](w[c])
2946 + alpha
2947 * beta
2948 * (1 - gamma)
2949 * self.wInterpolators[i][j][k - 1](w[c])
2950 + alpha * beta * gamma * self.wInterpolators[i][j][k](w[c])
2951 )
2952 return f
2954 def _derW(self, w, x, y, z):
2955 """
2956 Returns the derivative with respect to w of the interpolated function
2957 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeW.
2958 """
2959 m = len(x)
2960 x_pos = np.searchsorted(self.x_list, x)
2961 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
2962 y_pos = np.searchsorted(self.y_list, y)
2963 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
2964 y_pos[y_pos < 1] = 1
2965 z_pos = np.searchsorted(self.z_list, z)
2966 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
2967 z_pos[z_pos < 1] = 1
2968 dfdw = np.zeros(m) + np.nan
2969 for i in range(1, self.x_n):
2970 for j in range(1, self.y_n):
2971 for k in range(1, self.z_n):
2972 c = np.logical_and(
2973 np.logical_and(i == x_pos, j == y_pos), k == z_pos
2974 )
2975 if np.any(c):
2976 alpha = (x[c] - self.x_list[i - 1]) / (
2977 self.x_list[i] - self.x_list[i - 1]
2978 )
2979 beta = (y[c] - self.y_list[j - 1]) / (
2980 self.y_list[j] - self.y_list[j - 1]
2981 )
2982 gamma = (z[c] - self.z_list[k - 1]) / (
2983 self.z_list[k] - self.z_list[k - 1]
2984 )
2985 dfdw[c] = (
2986 (1 - alpha)
2987 * (1 - beta)
2988 * (1 - gamma)
2989 * self.wInterpolators[i - 1][j - 1][k - 1]._der(w[c])
2990 + (1 - alpha)
2991 * (1 - beta)
2992 * gamma
2993 * self.wInterpolators[i - 1][j - 1][k]._der(w[c])
2994 + (1 - alpha)
2995 * beta
2996 * (1 - gamma)
2997 * self.wInterpolators[i - 1][j][k - 1]._der(w[c])
2998 + (1 - alpha)
2999 * beta
3000 * gamma
3001 * self.wInterpolators[i - 1][j][k]._der(w[c])
3002 + alpha
3003 * (1 - beta)
3004 * (1 - gamma)
3005 * self.wInterpolators[i][j - 1][k - 1]._der(w[c])
3006 + alpha
3007 * (1 - beta)
3008 * gamma
3009 * self.wInterpolators[i][j - 1][k]._der(w[c])
3010 + alpha
3011 * beta
3012 * (1 - gamma)
3013 * self.wInterpolators[i][j][k - 1]._der(w[c])
3014 + alpha
3015 * beta
3016 * gamma
3017 * self.wInterpolators[i][j][k]._der(w[c])
3018 )
3019 return dfdw
3021 def _derX(self, w, x, y, z):
3022 """
3023 Returns the derivative with respect to x of the interpolated function
3024 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeX.
3025 """
3026 m = len(x)
3027 x_pos = np.searchsorted(self.x_list, x)
3028 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
3029 y_pos = np.searchsorted(self.y_list, y)
3030 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
3031 y_pos[y_pos < 1] = 1
3032 z_pos = np.searchsorted(self.z_list, z)
3033 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
3034 z_pos[z_pos < 1] = 1
3035 dfdx = np.zeros(m) + np.nan
3036 for i in range(1, self.x_n):
3037 for j in range(1, self.y_n):
3038 for k in range(1, self.z_n):
3039 c = np.logical_and(
3040 np.logical_and(i == x_pos, j == y_pos), k == z_pos
3041 )
3042 if np.any(c):
3043 beta = (y[c] - self.y_list[j - 1]) / (
3044 self.y_list[j] - self.y_list[j - 1]
3045 )
3046 gamma = (z[c] - self.z_list[k - 1]) / (
3047 self.z_list[k] - self.z_list[k - 1]
3048 )
3049 dfdx[c] = (
3050 (
3051 (1 - beta)
3052 * (1 - gamma)
3053 * self.wInterpolators[i][j - 1][k - 1](w[c])
3054 + (1 - beta)
3055 * gamma
3056 * self.wInterpolators[i][j - 1][k](w[c])
3057 + beta
3058 * (1 - gamma)
3059 * self.wInterpolators[i][j][k - 1](w[c])
3060 + beta * gamma * self.wInterpolators[i][j][k](w[c])
3061 )
3062 - (
3063 (1 - beta)
3064 * (1 - gamma)
3065 * self.wInterpolators[i - 1][j - 1][k - 1](w[c])
3066 + (1 - beta)
3067 * gamma
3068 * self.wInterpolators[i - 1][j - 1][k](w[c])
3069 + beta
3070 * (1 - gamma)
3071 * self.wInterpolators[i - 1][j][k - 1](w[c])
3072 + beta * gamma * self.wInterpolators[i - 1][j][k](w[c])
3073 )
3074 ) / (self.x_list[i] - self.x_list[i - 1])
3075 return dfdx
3077 def _derY(self, w, x, y, z):
3078 """
3079 Returns the derivative with respect to y of the interpolated function
3080 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeY.
3081 """
3082 m = len(x)
3083 x_pos = np.searchsorted(self.x_list, x)
3084 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
3085 y_pos = np.searchsorted(self.y_list, y)
3086 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
3087 y_pos[y_pos < 1] = 1
3088 z_pos = np.searchsorted(self.z_list, z)
3089 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
3090 z_pos[z_pos < 1] = 1
3091 dfdy = np.zeros(m) + np.nan
3092 for i in range(1, self.x_n):
3093 for j in range(1, self.y_n):
3094 for k in range(1, self.z_n):
3095 c = np.logical_and(
3096 np.logical_and(i == x_pos, j == y_pos), k == z_pos
3097 )
3098 if np.any(c):
3099 alpha = (x[c] - self.x_list[i - 1]) / (
3100 self.x_list[i] - self.x_list[i - 1]
3101 )
3102 gamma = (z[c] - self.z_list[k - 1]) / (
3103 self.z_list[k] - self.z_list[k - 1]
3104 )
3105 dfdy[c] = (
3106 (
3107 (1 - alpha)
3108 * (1 - gamma)
3109 * self.wInterpolators[i - 1][j][k - 1](w[c])
3110 + (1 - alpha)
3111 * gamma
3112 * self.wInterpolators[i - 1][j][k](w[c])
3113 + alpha
3114 * (1 - gamma)
3115 * self.wInterpolators[i][j][k - 1](w[c])
3116 + alpha * gamma * self.wInterpolators[i][j][k](w[c])
3117 )
3118 - (
3119 (1 - alpha)
3120 * (1 - gamma)
3121 * self.wInterpolators[i - 1][j - 1][k - 1](w[c])
3122 + (1 - alpha)
3123 * gamma
3124 * self.wInterpolators[i - 1][j - 1][k](w[c])
3125 + alpha
3126 * (1 - gamma)
3127 * self.wInterpolators[i][j - 1][k - 1](w[c])
3128 + alpha * gamma * self.wInterpolators[i][j - 1][k](w[c])
3129 )
3130 ) / (self.y_list[j] - self.y_list[j - 1])
3131 return dfdy
3133 def _derZ(self, w, x, y, z):
3134 """
3135 Returns the derivative with respect to z of the interpolated function
3136 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeZ.
3137 """
3138 m = len(x)
3139 x_pos = np.searchsorted(self.x_list, x)
3140 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
3141 y_pos = np.searchsorted(self.y_list, y)
3142 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
3143 y_pos[y_pos < 1] = 1
3144 z_pos = np.searchsorted(self.z_list, z)
3145 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
3146 z_pos[z_pos < 1] = 1
3147 dfdz = np.zeros(m) + np.nan
3148 for i in range(1, self.x_n):
3149 for j in range(1, self.y_n):
3150 for k in range(1, self.z_n):
3151 c = np.logical_and(
3152 np.logical_and(i == x_pos, j == y_pos), k == z_pos
3153 )
3154 if np.any(c):
3155 alpha = (x[c] - self.x_list[i - 1]) / (
3156 self.x_list[i] - self.x_list[i - 1]
3157 )
3158 beta = (y[c] - self.y_list[j - 1]) / (
3159 self.y_list[j] - self.y_list[j - 1]
3160 )
3161 dfdz[c] = (
3162 (
3163 (1 - alpha)
3164 * (1 - beta)
3165 * self.wInterpolators[i - 1][j - 1][k](w[c])
3166 + (1 - alpha)
3167 * beta
3168 * self.wInterpolators[i - 1][j][k](w[c])
3169 + alpha
3170 * (1 - beta)
3171 * self.wInterpolators[i][j - 1][k](w[c])
3172 + alpha * beta * self.wInterpolators[i][j][k](w[c])
3173 )
3174 - (
3175 (1 - alpha)
3176 * (1 - beta)
3177 * self.wInterpolators[i - 1][j - 1][k - 1](w[c])
3178 + (1 - alpha)
3179 * beta
3180 * self.wInterpolators[i - 1][j][k - 1](w[c])
3181 + alpha
3182 * (1 - beta)
3183 * self.wInterpolators[i][j - 1][k - 1](w[c])
3184 + alpha * beta * self.wInterpolators[i][j][k - 1](w[c])
3185 )
3186 ) / (self.z_list[k] - self.z_list[k - 1])
3187 return dfdz
3190class LinearInterpOnInterp2D(HARKinterpolator3D):
3191 """
3192 A 3D interpolation method that linearly interpolates between "layers" of
3193 arbitrary 2D interpolations. Useful for models with two endogenous state
3194 variables and one exogenous state variable when solving with the endogenous
3195 grid method. NOTE: should not be used if an exogenous 3D grid is used, will
3196 be significantly slower than TrilinearInterp.
3198 Constructor for the class, generating an approximation to a function of
3199 the form f(x,y,z) using interpolations over f(x,y,z_0) for a fixed grid
3200 of z_0 values.
3202 Parameters
3203 ----------
3204 xyInterpolators : [HARKinterpolator2D]
3205 A list of 2D interpolations over the x and y variables. The nth
3206 element of xyInterpolators represents f(x,y,z_values[n]).
3207 z_values: numpy.array
3208 An array of z values equal in length to xyInterpolators.
3209 """
3211 distance_criteria = ["xyInterpolators", "z_list"]
3213 def __init__(self, xyInterpolators, z_values):
3214 self.xyInterpolators = xyInterpolators
3215 self.z_list = z_values
3216 self.z_n = z_values.size
3218 def _evaluate(self, x, y, z):
3219 """
3220 Returns the level of the interpolated function at each value in x,y,z.
3221 Only called internally by HARKinterpolator3D.__call__ (etc).
3222 """
3223 m = len(x)
3224 z_pos = np.searchsorted(self.z_list, z)
3225 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
3226 z_pos[z_pos < 1] = 1
3227 f = np.zeros(m) + np.nan
3228 if x.size > 0:
3229 for i in range(1, self.z_n):
3230 c = z_pos == i
3231 if np.any(c):
3232 alpha = (z[c] - self.z_list[i - 1]) / (
3233 self.z_list[i] - self.z_list[i - 1]
3234 )
3235 f[c] = (1 - alpha) * self.xyInterpolators[i - 1](
3236 x[c], y[c]
3237 ) + alpha * self.xyInterpolators[i](x[c], y[c])
3238 return f
3240 def _derX(self, x, y, z):
3241 """
3242 Returns the derivative with respect to x of the interpolated function
3243 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeX.
3244 """
3245 m = len(x)
3246 z_pos = np.searchsorted(self.z_list, z)
3247 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
3248 z_pos[z_pos < 1] = 1
3249 dfdx = np.zeros(m) + np.nan
3250 if x.size > 0:
3251 for i in range(1, self.z_n):
3252 c = z_pos == i
3253 if np.any(c):
3254 alpha = (z[c] - self.z_list[i - 1]) / (
3255 self.z_list[i] - self.z_list[i - 1]
3256 )
3257 dfdx[c] = (1 - alpha) * self.xyInterpolators[i - 1].derivativeX(
3258 x[c], y[c]
3259 ) + alpha * self.xyInterpolators[i].derivativeX(x[c], y[c])
3260 return dfdx
3262 def _derY(self, x, y, z):
3263 """
3264 Returns the derivative with respect to y of the interpolated function
3265 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeY.
3266 """
3267 m = len(x)
3268 z_pos = np.searchsorted(self.z_list, z)
3269 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
3270 z_pos[z_pos < 1] = 1
3271 dfdy = np.zeros(m) + np.nan
3272 if x.size > 0:
3273 for i in range(1, self.z_n):
3274 c = z_pos == i
3275 if np.any(c):
3276 alpha = (z[c] - self.z_list[i - 1]) / (
3277 self.z_list[i] - self.z_list[i - 1]
3278 )
3279 dfdy[c] = (1 - alpha) * self.xyInterpolators[i - 1].derivativeY(
3280 x[c], y[c]
3281 ) + alpha * self.xyInterpolators[i].derivativeY(x[c], y[c])
3282 return dfdy
3284 def _derZ(self, x, y, z):
3285 """
3286 Returns the derivative with respect to z of the interpolated function
3287 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeZ.
3288 """
3289 m = len(x)
3290 z_pos = np.searchsorted(self.z_list, z)
3291 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
3292 z_pos[z_pos < 1] = 1
3293 dfdz = np.zeros(m) + np.nan
3294 if x.size > 0:
3295 for i in range(1, self.z_n):
3296 c = z_pos == i
3297 if np.any(c):
3298 dfdz[c] = (
3299 self.xyInterpolators[i](x[c], y[c])
3300 - self.xyInterpolators[i - 1](x[c], y[c])
3301 ) / (self.z_list[i] - self.z_list[i - 1])
3302 return dfdz
3305class BilinearInterpOnInterp2D(HARKinterpolator4D):
3306 """
3307 A 4D interpolation method that bilinearly interpolates among "layers" of
3308 arbitrary 2D interpolations. Useful for models with two endogenous state
3309 variables and two exogenous state variables when solving with the endogenous
3310 grid method. NOTE: should not be used if an exogenous 4D grid is used, will
3311 be significantly slower than QuadlinearInterp.
3313 Constructor for the class, generating an approximation to a function of
3314 the form f(w,x,y,z) using interpolations over f(w,x,y_0,z_0) for a fixed
3315 grid of y_0 and z_0 values.
3317 Parameters
3318 ----------
3319 wxInterpolators : [[HARKinterpolator2D]]
3320 A list of lists of 2D interpolations over the w and x variables.
3321 The i,j-th element of wxInterpolators represents
3322 f(w,x,y_values[i],z_values[j]).
3323 y_values: numpy.array
3324 An array of y values equal in length to wxInterpolators.
3325 z_values: numpy.array
3326 An array of z values equal in length to wxInterpolators[0].
3327 """
3329 distance_criteria = ["wxInterpolators", "y_list", "z_list"]
3331 def __init__(self, wxInterpolators, y_values, z_values):
3332 self.wxInterpolators = wxInterpolators
3333 self.y_list = y_values
3334 self.y_n = y_values.size
3335 self.z_list = z_values
3336 self.z_n = z_values.size
3338 def _evaluate(self, w, x, y, z):
3339 """
3340 Returns the level of the interpolated function at each value in x,y,z.
3341 Only called internally by HARKinterpolator4D.__call__ (etc).
3342 """
3343 m = len(x)
3344 y_pos = np.searchsorted(self.y_list, y)
3345 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
3346 y_pos[y_pos < 1] = 1
3347 z_pos = np.searchsorted(self.z_list, z)
3348 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
3349 z_pos[z_pos < 1] = 1
3350 f = np.zeros(m) + np.nan
3351 for i in range(1, self.y_n):
3352 for j in range(1, self.z_n):
3353 c = np.logical_and(i == y_pos, j == z_pos)
3354 if np.any(c):
3355 alpha = (y[c] - self.y_list[i - 1]) / (
3356 self.y_list[i] - self.y_list[i - 1]
3357 )
3358 beta = (z[c] - self.z_list[j - 1]) / (
3359 self.z_list[j] - self.z_list[j - 1]
3360 )
3361 f[c] = (
3362 (1 - alpha)
3363 * (1 - beta)
3364 * self.wxInterpolators[i - 1][j - 1](w[c], x[c])
3365 + (1 - alpha)
3366 * beta
3367 * self.wxInterpolators[i - 1][j](w[c], x[c])
3368 + alpha
3369 * (1 - beta)
3370 * self.wxInterpolators[i][j - 1](w[c], x[c])
3371 + alpha * beta * self.wxInterpolators[i][j](w[c], x[c])
3372 )
3373 return f
3375 def _derW(self, w, x, y, z):
3376 """
3377 Returns the derivative with respect to w of the interpolated function
3378 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeW.
3379 """
3380 # This may look strange, as we call the derivativeX() method to get the
3381 # derivative with respect to w, but that's just a quirk of 4D interpolations
3382 # beginning with w rather than x. The derivative wrt the first dimension
3383 # of an element of wxInterpolators is the w-derivative of the main function.
3384 m = len(x)
3385 y_pos = np.searchsorted(self.y_list, y)
3386 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
3387 y_pos[y_pos < 1] = 1
3388 z_pos = np.searchsorted(self.z_list, z)
3389 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
3390 z_pos[z_pos < 1] = 1
3391 dfdw = np.zeros(m) + np.nan
3392 for i in range(1, self.y_n):
3393 for j in range(1, self.z_n):
3394 c = np.logical_and(i == y_pos, j == z_pos)
3395 if np.any(c):
3396 alpha = (y[c] - self.y_list[i - 1]) / (
3397 self.y_list[i] - self.y_list[i - 1]
3398 )
3399 beta = (z[c] - self.z_list[j - 1]) / (
3400 self.z_list[j] - self.z_list[j - 1]
3401 )
3402 dfdw[c] = (
3403 (1 - alpha)
3404 * (1 - beta)
3405 * self.wxInterpolators[i - 1][j - 1].derivativeX(w[c], x[c])
3406 + (1 - alpha)
3407 * beta
3408 * self.wxInterpolators[i - 1][j].derivativeX(w[c], x[c])
3409 + alpha
3410 * (1 - beta)
3411 * self.wxInterpolators[i][j - 1].derivativeX(w[c], x[c])
3412 + alpha
3413 * beta
3414 * self.wxInterpolators[i][j].derivativeX(w[c], x[c])
3415 )
3416 return dfdw
3418 def _derX(self, w, x, y, z):
3419 """
3420 Returns the derivative with respect to x of the interpolated function
3421 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeX.
3422 """
3423 # This may look strange, as we call the derivativeY() method to get the
3424 # derivative with respect to x, but that's just a quirk of 4D interpolations
3425 # beginning with w rather than x. The derivative wrt the second dimension
3426 # of an element of wxInterpolators is the x-derivative of the main function.
3427 m = len(x)
3428 y_pos = np.searchsorted(self.y_list, y)
3429 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
3430 y_pos[y_pos < 1] = 1
3431 z_pos = np.searchsorted(self.z_list, z)
3432 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
3433 z_pos[z_pos < 1] = 1
3434 dfdx = np.zeros(m) + np.nan
3435 for i in range(1, self.y_n):
3436 for j in range(1, self.z_n):
3437 c = np.logical_and(i == y_pos, j == z_pos)
3438 if np.any(c):
3439 alpha = (y[c] - self.y_list[i - 1]) / (
3440 self.y_list[i] - self.y_list[i - 1]
3441 )
3442 beta = (z[c] - self.z_list[j - 1]) / (
3443 self.z_list[j] - self.z_list[j - 1]
3444 )
3445 dfdx[c] = (
3446 (1 - alpha)
3447 * (1 - beta)
3448 * self.wxInterpolators[i - 1][j - 1].derivativeY(w[c], x[c])
3449 + (1 - alpha)
3450 * beta
3451 * self.wxInterpolators[i - 1][j].derivativeY(w[c], x[c])
3452 + alpha
3453 * (1 - beta)
3454 * self.wxInterpolators[i][j - 1].derivativeY(w[c], x[c])
3455 + alpha
3456 * beta
3457 * self.wxInterpolators[i][j].derivativeY(w[c], x[c])
3458 )
3459 return dfdx
3461 def _derY(self, w, x, y, z):
3462 """
3463 Returns the derivative with respect to y of the interpolated function
3464 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeY.
3465 """
3466 m = len(x)
3467 y_pos = np.searchsorted(self.y_list, y)
3468 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
3469 y_pos[y_pos < 1] = 1
3470 z_pos = np.searchsorted(self.z_list, z)
3471 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
3472 z_pos[z_pos < 1] = 1
3473 dfdy = np.zeros(m) + np.nan
3474 for i in range(1, self.y_n):
3475 for j in range(1, self.z_n):
3476 c = np.logical_and(i == y_pos, j == z_pos)
3477 if np.any(c):
3478 beta = (z[c] - self.z_list[j - 1]) / (
3479 self.z_list[j] - self.z_list[j - 1]
3480 )
3481 dfdy[c] = (
3482 (
3483 (1 - beta) * self.wxInterpolators[i][j - 1](w[c], x[c])
3484 + beta * self.wxInterpolators[i][j](w[c], x[c])
3485 )
3486 - (
3487 (1 - beta) * self.wxInterpolators[i - 1][j - 1](w[c], x[c])
3488 + beta * self.wxInterpolators[i - 1][j](w[c], x[c])
3489 )
3490 ) / (self.y_list[i] - self.y_list[i - 1])
3491 return dfdy
3493 def _derZ(self, w, x, y, z):
3494 """
3495 Returns the derivative with respect to z of the interpolated function
3496 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeZ.
3497 """
3498 m = len(x)
3499 y_pos = np.searchsorted(self.y_list, y)
3500 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
3501 y_pos[y_pos < 1] = 1
3502 z_pos = np.searchsorted(self.z_list, z)
3503 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
3504 z_pos[z_pos < 1] = 1
3505 dfdz = np.zeros(m) + np.nan
3506 for i in range(1, self.y_n):
3507 for j in range(1, self.z_n):
3508 c = np.logical_and(i == y_pos, j == z_pos)
3509 if np.any(c):
3510 alpha = (y[c] - self.y_list[i - 1]) / (
3511 self.y_list[i] - self.y_list[i - 1]
3512 )
3513 dfdz[c] = (
3514 (
3515 (1 - alpha) * self.wxInterpolators[i - 1][j](w[c], x[c])
3516 + alpha * self.wxInterpolators[i][j](w[c], x[c])
3517 )
3518 - (
3519 (1 - alpha) * self.wxInterpolators[i - 1][j - 1](w[c], x[c])
3520 + alpha * self.wxInterpolators[i][j - 1](w[c], x[c])
3521 )
3522 ) / (self.z_list[j] - self.z_list[j - 1])
3523 return dfdz
3526class Curvilinear2DInterp(HARKinterpolator2D):
3527 """
3528 A 2D interpolation method for curvilinear or "warped grid" interpolation, as
3529 in White (2015). Used for models with two endogenous states that are solved
3530 with the endogenous grid method. Allows multiple function outputs, but all of
3531 the interpolated functions must share a common curvilinear grid.
3533 Parameters
3534 ----------
3535 f_values: numpy.array or [numpy.array]
3536 One or more 2D arrays of function values such that f_values[i,j] =
3537 f(x_values[i,j],y_values[i,j]).
3538 x_values: numpy.array
3539 A 2D array of x values of the same shape as f_values.
3540 y_values: numpy.array
3541 A 2D array of y values of the same shape as f_values.
3542 """
3544 distance_criteria = ["f_values", "x_values", "y_values"]
3546 def __init__(self, f_values, x_values, y_values):
3547 if isinstance(f_values, list):
3548 N_funcs = len(f_values)
3549 multi = True
3550 else:
3551 N_funcs = 1
3552 multi = False
3553 my_shape = x_values.shape
3554 if not (my_shape == y_values.shape):
3555 raise ValueError("y_values must have the same shape as x_values!")
3556 if multi:
3557 for n in range(N_funcs):
3558 if not (my_shape == f_values[n].shape):
3559 raise ValueError(
3560 "Each element of f_values must have the same shape as x_values!"
3561 )
3562 else:
3563 if not (my_shape == f_values.shape):
3564 raise ValueError("f_values must have the same shape as x_values!")
3566 if multi:
3567 self.f_values = f_values
3568 else:
3569 self.f_values = [f_values]
3570 self.x_values = x_values
3571 self.y_values = y_values
3572 self.x_n = my_shape[0]
3573 self.y_n = my_shape[1]
3574 self.N_funcs = N_funcs
3575 self.multi = multi
3576 self.update_polarity()
3578 def __call__(self, x, y):
3579 """
3580 Modification of HARKinterpolator2D.__call__ to account for multiple outputs.
3581 """
3582 xa = np.asarray(x)
3583 ya = np.asarray(y)
3584 S = xa.shape
3585 fa = self._evaluate(xa.flatten(), ya.flatten())
3586 output = [fa[n].reshape(S) for n in range(self.N_funcs)]
3587 if self.multi:
3588 return output
3589 else:
3590 return output[0]
3592 def derivativeX(self, x, y):
3593 """
3594 Modification of HARKinterpolator2D.derivativeX to account for multiple outputs.
3595 """
3596 xa = np.asarray(x)
3597 ya = np.asarray(y)
3598 S = xa.shape
3599 dfdxa = self._derX(xa.flatten(), ya.flatten())
3600 output = [dfdxa[n].reshape(S) for n in range(self.N_funcs)]
3601 if self.multi:
3602 return output
3603 else:
3604 return output[0]
3606 def derivativeY(self, x, y):
3607 """
3608 Modification of HARKinterpolator2D.derivativeY to account for multiple outputs.
3609 """
3610 xa = np.asarray(x)
3611 ya = np.asarray(y)
3612 S = xa.shape
3613 dfdya = self._derY(xa.flatten(), ya.flatten())
3614 output = [dfdya[n].reshape(S) for n in range(self.N_funcs)]
3615 if self.multi:
3616 return output
3617 else:
3618 return output[0]
3620 def update_polarity(self):
3621 """
3622 Fills in the polarity attribute of the interpolation, determining whether
3623 the "plus" (True) or "minus" (False) solution of the system of equations
3624 should be used for each sector. Needs to be called in __init__.
3626 Parameters
3627 ----------
3628 none
3630 Returns
3631 -------
3632 none
3633 """
3634 # Grab a point known to be inside each sector: the midway point between
3635 # the lower left and upper right vertex of each sector
3636 x_temp = 0.5 * (
3637 self.x_values[0 : (self.x_n - 1), 0 : (self.y_n - 1)]
3638 + self.x_values[1 : self.x_n, 1 : self.y_n]
3639 )
3640 y_temp = 0.5 * (
3641 self.y_values[0 : (self.x_n - 1), 0 : (self.y_n - 1)]
3642 + self.y_values[1 : self.x_n, 1 : self.y_n]
3643 )
3644 size = (self.x_n - 1) * (self.y_n - 1)
3645 x_temp = np.reshape(x_temp, size)
3646 y_temp = np.reshape(y_temp, size)
3647 y_pos = np.tile(np.arange(0, self.y_n - 1), self.x_n - 1)
3648 x_pos = np.reshape(
3649 np.tile(np.arange(0, self.x_n - 1), (self.y_n - 1, 1)).transpose(), size
3650 )
3652 # Set the polarity of all sectors to "plus", then test each sector
3653 self.polarity = np.ones((self.x_n - 1, self.y_n - 1), dtype=bool)
3654 alpha, beta = self.find_coords(x_temp, y_temp, x_pos, y_pos)
3655 polarity = np.logical_and(
3656 np.logical_and(alpha > 0, alpha < 1), np.logical_and(beta > 0, beta < 1)
3657 )
3659 # Update polarity: if (alpha,beta) not in the unit square, then that
3660 # sector must use the "minus" solution instead
3661 self.polarity = np.reshape(polarity, (self.x_n - 1, self.y_n - 1))
3663 def find_sector(self, x, y):
3664 """
3665 Finds the quadrilateral "sector" for each (x,y) point in the input.
3666 Only called as a subroutine of _evaluate(), etc. Uses a numba helper
3667 function below to accelerate computation.
3669 Parameters
3670 ----------
3671 x : np.array
3672 Values whose sector should be found.
3673 y : np.array
3674 Values whose sector should be found. Should be same size as x.
3676 Returns
3677 -------
3678 x_pos : np.array
3679 Sector x-coordinates for each point of the input, of the same size.
3680 y_pos : np.array
3681 Sector y-coordinates for each point of the input, of the same size.
3682 """
3683 x_pos, y_pos = find_sector_numba(x, y, self.x_values, self.y_values)
3684 return x_pos, y_pos
3686 def find_coords(self, x, y, x_pos, y_pos):
3687 """
3688 Calculates the relative coordinates (alpha,beta) for each point (x,y),
3689 given the sectors (x_pos,y_pos) in which they reside. Only called as
3690 a subroutine of _evaluate(), etc. Uses a numba helper function to acc-
3691 elerate computation, and has a "backup method" for when the math fails.
3693 Parameters
3694 ----------
3695 x : np.array
3696 Values whose sector should be found.
3697 y : np.array
3698 Values whose sector should be found. Should be same size as x.
3699 x_pos : np.array
3700 Sector x-coordinates for each point in (x,y), of the same size.
3701 y_pos : np.array
3702 Sector y-coordinates for each point in (x,y), of the same size.
3704 Returns
3705 -------
3706 alpha : np.array
3707 Relative "horizontal" position of the input in their respective sectors.
3708 beta : np.array
3709 Relative "vertical" position of the input in their respective sectors.
3710 """
3711 alpha, beta = find_coords_numba(
3712 x, y, x_pos, y_pos, self.x_values, self.y_values, self.polarity
3713 )
3715 # Alternate method if there are sectors that are "too regular"
3716 # These points weren't able to identify coordinates
3717 z = np.logical_or(np.isnan(alpha), np.isnan(beta))
3718 if np.any(z):
3719 ii = x_pos[z]
3720 jj = y_pos[z]
3721 xA = self.x_values[ii, jj]
3722 xB = self.x_values[ii + 1, jj]
3723 xC = self.x_values[ii, jj + 1]
3724 xD = self.x_values[ii + 1, jj + 1]
3725 yA = self.y_values[ii, jj]
3726 yB = self.y_values[ii + 1, jj]
3727 yC = self.y_values[ii, jj + 1]
3728 # yD = self.y_values[ii + 1, jj + 1]
3729 b = xB - xA
3730 f = yB - yA
3731 kappa = f / b
3732 int_bot = yA - kappa * xA
3733 int_top = yC - kappa * xC
3734 int_these = y[z] - kappa * x[z]
3735 beta_temp = (int_these - int_bot) / (int_top - int_bot)
3736 x_left = beta_temp * xC + (1.0 - beta_temp) * xA
3737 x_right = beta_temp * xD + (1.0 - beta_temp) * xB
3738 alpha_temp = (x[z] - x_left) / (x_right - x_left)
3739 beta[z] = beta_temp
3740 alpha[z] = alpha_temp
3742 return alpha, beta
3744 def _evaluate(self, x, y):
3745 """
3746 Returns the level of the interpolated function at each value in x,y.
3747 Only called internally by __call__ (etc).
3748 """
3749 x_pos, y_pos = self.find_sector(x, y)
3750 alpha, beta = self.find_coords(x, y, x_pos, y_pos)
3752 # Get weights on each vertex
3753 alpha_C = 1.0 - alpha
3754 beta_C = 1.0 - beta
3755 wA = alpha_C * beta_C
3756 wB = alpha * beta_C
3757 wC = alpha_C * beta
3758 wD = alpha * beta
3760 # Evaluate each function by bilinear interpolation
3761 f = []
3762 for n in range(self.N_funcs):
3763 f_n = (
3764 0.0
3765 + wA * self.f_values[n][x_pos, y_pos]
3766 + wB * self.f_values[n][x_pos + 1, y_pos]
3767 + wC * self.f_values[n][x_pos, y_pos + 1]
3768 + wD * self.f_values[n][x_pos + 1, y_pos + 1]
3769 )
3770 f.append(f_n)
3771 return f
3773 def _derX(self, x, y):
3774 """
3775 Returns the derivative with respect to x of the interpolated function
3776 at each value in x,y. Only called internally by derivativeX.
3777 """
3778 x_pos, y_pos = self.find_sector(x, y)
3779 alpha, beta = self.find_coords(x, y, x_pos, y_pos)
3781 # Get four corners data for each point
3782 xA = self.x_values[x_pos, y_pos]
3783 xB = self.x_values[x_pos + 1, y_pos]
3784 xC = self.x_values[x_pos, y_pos + 1]
3785 xD = self.x_values[x_pos + 1, y_pos + 1]
3786 yA = self.y_values[x_pos, y_pos]
3787 yB = self.y_values[x_pos + 1, y_pos]
3788 yC = self.y_values[x_pos, y_pos + 1]
3789 yD = self.y_values[x_pos + 1, y_pos + 1]
3791 # Calculate components of the alpha,beta --> x,y delta translation matrix
3792 alpha_C = 1 - alpha
3793 beta_C = 1 - beta
3794 alpha_x = beta_C * (xB - xA) + beta * (xD - xC)
3795 alpha_y = beta_C * (yB - yA) + beta * (yD - yC)
3796 beta_x = alpha_C * (xC - xA) + alpha * (xD - xB)
3797 beta_y = alpha_C * (yC - yA) + alpha * (yD - yB)
3799 # Invert the delta translation matrix into x,y --> alpha,beta
3800 det = alpha_x * beta_y - beta_x * alpha_y
3801 x_alpha = beta_y / det
3802 x_beta = -alpha_y / det
3804 # Calculate the derivative of f w.r.t. alpha and beta for each function
3805 dfdx = []
3806 for n in range(self.N_funcs):
3807 fA = self.f_values[n][x_pos, y_pos]
3808 fB = self.f_values[n][x_pos + 1, y_pos]
3809 fC = self.f_values[n][x_pos, y_pos + 1]
3810 fD = self.f_values[n][x_pos + 1, y_pos + 1]
3811 dfda = beta_C * (fB - fA) + beta * (fD - fC)
3812 dfdb = alpha_C * (fC - fA) + alpha * (fD - fB)
3814 # Calculate the derivative with respect to x
3815 dfdx_n = x_alpha * dfda + x_beta * dfdb
3816 dfdx.append(dfdx_n)
3817 return dfdx
3819 def _derY(self, x, y):
3820 """
3821 Returns the derivative with respect to y of the interpolated function
3822 at each value in x,y. Only called internally by derivativeY.
3823 """
3824 x_pos, y_pos = self.find_sector(x, y)
3825 alpha, beta = self.find_coords(x, y, x_pos, y_pos)
3827 # Get four corners data for each point
3828 xA = self.x_values[x_pos, y_pos]
3829 xB = self.x_values[x_pos + 1, y_pos]
3830 xC = self.x_values[x_pos, y_pos + 1]
3831 xD = self.x_values[x_pos + 1, y_pos + 1]
3832 yA = self.y_values[x_pos, y_pos]
3833 yB = self.y_values[x_pos + 1, y_pos]
3834 yC = self.y_values[x_pos, y_pos + 1]
3835 yD = self.y_values[x_pos + 1, y_pos + 1]
3837 # Calculate components of the alpha,beta --> x,y delta translation matrix
3838 alpha_C = 1 - alpha
3839 beta_C = 1 - beta
3840 alpha_x = beta_C * (xB - xA) + beta * (xD - xC)
3841 alpha_y = beta_C * (yB - yA) + beta * (yD - yC)
3842 beta_x = alpha_C * (xC - xA) + alpha * (xD - xB)
3843 beta_y = alpha_C * (yC - yA) + alpha * (yD - yB)
3845 # Invert the delta translation matrix into x,y --> alpha,beta
3846 det = alpha_x * beta_y - beta_x * alpha_y
3847 y_alpha = -beta_x / det
3848 y_beta = alpha_x / det
3850 # Calculate the derivative of f w.r.t. alpha and beta for each function
3851 dfdy = []
3852 for n in range(self.N_funcs):
3853 fA = self.f_values[n][x_pos, y_pos]
3854 fB = self.f_values[n][x_pos + 1, y_pos]
3855 fC = self.f_values[n][x_pos, y_pos + 1]
3856 fD = self.f_values[n][x_pos + 1, y_pos + 1]
3857 dfda = beta_C * (fB - fA) + beta * (fD - fC)
3858 dfdb = alpha_C * (fC - fA) + alpha * (fD - fB)
3860 # Calculate the derivative with respect to y
3861 dfdy_n = y_alpha * dfda + y_beta * dfdb
3862 dfdy.append(dfdy_n)
3863 return dfdy
3866# Define a function that checks whether a set of points violates a linear boundary
3867# defined by (x1,y1) and (x2,y2), where the latter is *COUNTER CLOCKWISE* from the
3868# former. Returns 1 if the point is outside the boundary and 0 otherwise.
3869@njit
3870def boundary_check(xq, yq, x1, y1, x2, y2): # pragma: no cover
3871 return int((y2 - y1) * xq - (x2 - x1) * yq > x1 * y2 - y1 * x2)
3874# Define a numba helper function for finding the sector in the irregular grid
3875@njit
3876def find_sector_numba(X_query, Y_query, X_values, Y_values): # pragma: no cover
3877 # Initialize the sector guess
3878 M = X_query.size
3879 x_n = X_values.shape[0]
3880 y_n = X_values.shape[1]
3881 ii = int(x_n / 2)
3882 jj = int(y_n / 2)
3883 top_ii = x_n - 2
3884 top_jj = y_n - 2
3886 # Initialize the output arrays
3887 X_pos = np.empty(M, dtype=np.int32)
3888 Y_pos = np.empty(M, dtype=np.int32)
3890 # Identify the correct sector for each point to be evaluated
3891 max_loops = x_n + y_n
3892 for m in range(M):
3893 found = False
3894 loops = 0
3895 while not found and loops < max_loops:
3896 # Get coordinates for the four vertices: (xA,yA),...,(xD,yD)
3897 x0 = X_query[m]
3898 y0 = Y_query[m]
3899 xA = X_values[ii, jj]
3900 xB = X_values[ii + 1, jj]
3901 xC = X_values[ii, jj + 1]
3902 xD = X_values[ii + 1, jj + 1]
3903 yA = Y_values[ii, jj]
3904 yB = Y_values[ii + 1, jj]
3905 yC = Y_values[ii, jj + 1]
3906 yD = Y_values[ii + 1, jj + 1]
3908 # Check the "bounding box" for the sector: is this guess plausible?
3909 D = int(y0 < np.minimum(yA, yB))
3910 R = int(x0 > np.maximum(xB, xD))
3911 U = int(y0 > np.maximum(yC, yD))
3912 L = int(x0 < np.minimum(xA, xC))
3914 # Check which boundaries are violated (and thus where to look next)
3915 in_box = np.all(np.logical_not(np.array([D, R, U, L])))
3916 if in_box:
3917 D = boundary_check(x0, y0, xA, yA, xB, yB)
3918 R = boundary_check(x0, y0, xB, yB, xD, yD)
3919 U = boundary_check(x0, y0, xD, yD, xC, yC)
3920 L = boundary_check(x0, y0, xC, yC, xA, yA)
3922 # Update the sector guess based on the violations
3923 ii_next = np.maximum(np.minimum(ii - L + R, top_ii), 0)
3924 jj_next = np.maximum(np.minimum(jj - D + U, top_jj), 0)
3926 # Check whether sector guess changed and go to next iteration
3927 found = (ii == ii_next) and (jj == jj_next)
3928 ii = ii_next
3929 jj = jj_next
3930 loops += 1
3932 # Put the final sector guess into the output array
3933 X_pos[m] = ii
3934 Y_pos[m] = jj
3936 # Return the output
3937 return X_pos, Y_pos
3940# Define a numba helper function for finding relative coordinates within sector
3941@njit
3942def find_coords_numba(
3943 X_query, Y_query, X_pos, Y_pos, X_values, Y_values, polarity
3944): # pragma: no cover
3945 M = X_query.size
3946 alpha = np.empty(M)
3947 beta = np.empty(M)
3949 # Calculate relative coordinates in the sector for each point
3950 for m in range(M):
3951 try:
3952 x0 = X_query[m]
3953 y0 = Y_query[m]
3954 ii = X_pos[m]
3955 jj = Y_pos[m]
3956 xA = X_values[ii, jj]
3957 xB = X_values[ii + 1, jj]
3958 xC = X_values[ii, jj + 1]
3959 xD = X_values[ii + 1, jj + 1]
3960 yA = Y_values[ii, jj]
3961 yB = Y_values[ii + 1, jj]
3962 yC = Y_values[ii, jj + 1]
3963 yD = Y_values[ii + 1, jj + 1]
3964 p = 2.0 * polarity[ii, jj] - 1.0
3965 a = xA
3966 b = xB - xA
3967 c = xC - xA
3968 d = xA - xB - xC + xD
3969 e = yA
3970 f = yB - yA
3971 g = yC - yA
3972 h = yA - yB - yC + yD
3973 denom = d * g - h * c
3974 mu = (h * b - d * f) / denom
3975 tau = (h * (a - x0) - d * (e - y0)) / denom
3976 zeta = a - x0 + c * tau
3977 eta = b + c * mu + d * tau
3978 theta = d * mu
3979 alph = (-eta + p * np.sqrt(eta**2 - 4 * zeta * theta)) / (2 * theta)
3980 bet = mu * alph + tau
3981 except:
3982 alph = np.nan
3983 bet = np.nan
3984 alpha[m] = alph
3985 beta[m] = bet
3987 return alpha, beta
3990class DiscreteInterp(MetricObject):
3991 """
3992 An interpolator for variables that can only take a discrete set of values.
3994 If the function we wish to interpolate, f(args) can take on the list of
3995 values discrete_vals, this class expects an interpolator for the index of
3996 f's value in discrete_vals.
3997 E.g., if f(a,b,c) = discrete_vals[5], then index_interp(a,b,c) = 5.
3999 Parameters
4000 ----------
4001 index_interp: HARKInterpolator
4002 An interpolator giving an approximation to the index of the value in
4003 discrete_vals that corresponds to a given set of arguments.
4004 discrete_vals: numpy.array
4005 A 1D array containing the values in the range of the discrete function
4006 to be interpolated.
4007 """
4009 distance_criteria = ["index_interp"]
4011 def __init__(self, index_interp, discrete_vals):
4012 self.index_interp = index_interp
4013 self.discrete_vals = discrete_vals
4014 self.n_vals = len(self.discrete_vals)
4016 def __call__(self, *args):
4017 # Interpolate indices and round to integers
4018 inds = np.rint(self.index_interp(*args)).astype(int)
4019 if type(inds) is not np.ndarray:
4020 inds = np.array(inds)
4021 # Deal with out-of range indices
4022 inds[inds < 0] = 0
4023 inds[inds >= self.n_vals] = self.n_vals - 1
4025 # Get values from grid
4026 return self.discrete_vals[inds]
4029class IndexedInterp(MetricObject):
4030 """
4031 An interpolator for functions whose first argument is an integer-valued index.
4032 Constructor takes in a list of functions as its only argument. When evaluated
4033 at f(i,X), interpolator returns f[i](X), where X can be any number of inputs.
4034 This simply provides a different interface for accessing the same functions.
4036 Parameters
4037 ----------
4038 functions : [Callable]
4039 List of one or more functions to be indexed.
4040 """
4042 distance_criteria = ["functions"]
4044 def __init__(self, functions):
4045 self.functions = functions
4046 self.N = len(functions)
4048 def __call__(self, idx, *args):
4049 out = np.empty(idx.shape)
4050 out.fill(np.nan)
4052 for n in range(self.N):
4053 these = idx == n
4054 if not np.any(these):
4055 continue
4056 temp = [arg[these] for arg in args]
4057 out[these] = self.functions[n](*temp)
4058 return out
4061###############################################################################
4062## Functions used in discrete choice models with T1EV taste shocks ############
4063###############################################################################
4066def calc_log_sum_choice_probs(Vals, sigma):
4067 """
4068 Returns the final optimal value and choice probabilities given the choice
4069 specific value functions `Vals`. Probabilities are degenerate if sigma == 0.0.
4070 Parameters
4071 ----------
4072 Vals : [numpy.array]
4073 A numpy.array that holds choice specific values at common grid points.
4074 sigma : float
4075 A number that controls the variance of the taste shocks
4076 Returns
4077 -------
4078 V : [numpy.array]
4079 A numpy.array that holds the integrated value function.
4080 P : [numpy.array]
4081 A numpy.array that holds the discrete choice probabilities
4082 """
4083 # Assumes that NaNs have been replaced by -numpy.inf or similar
4084 if sigma == 0.0:
4085 # We could construct a linear index here and use unravel_index.
4086 Pflat = np.argmax(Vals, axis=0)
4088 V = np.zeros(Vals[0].shape)
4089 Probs = np.zeros(Vals.shape)
4090 for i in range(Vals.shape[0]):
4091 optimalIndices = Pflat == i
4092 V[optimalIndices] = Vals[i][optimalIndices]
4093 Probs[i][optimalIndices] = 1
4094 return V, Probs
4096 # else we have a taste shock
4097 maxV = np.max(Vals, axis=0)
4099 # calculate maxV+sigma*log(sum_i=1^J exp((V[i]-maxV))/sigma)
4100 sumexp = np.sum(np.exp((Vals - maxV) / sigma), axis=0)
4101 LogSumV = np.log(sumexp)
4102 LogSumV = maxV + sigma * LogSumV
4104 Probs = np.exp((Vals - LogSumV) / sigma)
4105 return LogSumV, Probs
4108def calc_choice_probs(Vals, sigma):
4109 """
4110 Returns the choice probabilities given the choice specific value functions
4111 `Vals`. Probabilities are degenerate if sigma == 0.0.
4112 Parameters
4113 ----------
4114 Vals : [numpy.array]
4115 A numpy.array that holds choice specific values at common grid points.
4116 sigma : float
4117 A number that controls the variance of the taste shocks
4118 Returns
4119 -------
4120 Probs : [numpy.array]
4121 A numpy.array that holds the discrete choice probabilities
4122 """
4124 # Assumes that NaNs have been replaced by -numpy.inf or similar
4125 if sigma == 0.0:
4126 # We could construct a linear index here and use unravel_index.
4127 Pflat = np.argmax(Vals, axis=0)
4128 Probs = np.zeros(Vals.shape)
4129 for i in range(Vals.shape[0]):
4130 Probs[i][Pflat == i] = 1
4131 return Probs
4133 maxV = np.max(Vals, axis=0)
4134 Probs = np.divide(
4135 np.exp((Vals - maxV) / sigma), np.sum(np.exp((Vals - maxV) / sigma), axis=0)
4136 )
4137 return Probs
4140def calc_log_sum(Vals, sigma):
4141 """
4142 Returns the optimal value given the choice specific value functions Vals.
4143 Parameters
4144 ----------
4145 Vals : [numpy.array]
4146 A numpy.array that holds choice specific values at common grid points.
4147 sigma : float
4148 A number that controls the variance of the taste shocks
4149 Returns
4150 -------
4151 V : [numpy.array]
4152 A numpy.array that holds the integrated value function.
4153 """
4155 # Assumes that NaNs have been replaced by -numpy.inf or similar
4156 if sigma == 0.0:
4157 # We could construct a linear index here and use unravel_index.
4158 V = np.amax(Vals, axis=0)
4159 return V
4161 # else we have a taste shock
4162 maxV = np.max(Vals, axis=0)
4164 # calculate maxV+sigma*log(sum_i=1^J exp((V[i]-maxV))/sigma)
4165 sumexp = np.sum(np.exp((Vals - maxV) / sigma), axis=0)
4166 LogSumV = np.log(sumexp)
4167 LogSumV = maxV + sigma * LogSumV
4168 return LogSumV
4171###############################################################################
4172# Tools for value and marginal-value functions in models where #
4173# - dvdm = u'(c). #
4174# - u is of the CRRA family. #
4175###############################################################################
4178class ValueFuncCRRA(MetricObject):
4179 """
4180 A class for representing a value function. The underlying interpolation is
4181 in the space of (state,u_inv(v)); this class "re-curves" to the value function.
4183 Parameters
4184 ----------
4185 vFuncNvrs : function
4186 A real function representing the value function composed with the
4187 inverse utility function, defined on the state: u_inv(vFunc(state))
4188 CRRA : float
4189 Coefficient of relative risk aversion.
4190 illegal_value : float, optional
4191 If provided, value to return for "out-of-bounds" inputs that return NaN
4192 from the pseudo-inverse value function. Most common choice is -np.inf,
4193 which makes the outcome infinitely bad.
4194 """
4196 distance_criteria = ["func", "CRRA"]
4198 def __init__(self, vFuncNvrs, CRRA, illegal_value=None):
4199 self.vFuncNvrs = deepcopy(vFuncNvrs)
4200 self.CRRA = CRRA
4201 self.illegal_value = illegal_value
4203 if hasattr(vFuncNvrs, "grid_list"):
4204 self.grid_list = vFuncNvrs.grid_list
4205 else:
4206 self.grid_list = None
4208 def __call__(self, *vFuncArgs):
4209 """
4210 Evaluate the value function at given levels of market resources m.
4212 Parameters
4213 ----------
4214 vFuncArgs : floats or np.arrays, all of the same dimensions.
4215 Values for the state variables. These usually start with 'm',
4216 market resources normalized by the level of permanent income.
4218 Returns
4219 -------
4220 v : float or np.array
4221 Lifetime value of beginning this period with the given states; has
4222 same size as the state inputs.
4223 """
4224 temp = self.vFuncNvrs(*vFuncArgs)
4225 v = CRRAutility(temp, self.CRRA)
4226 if self.illegal_value is not None:
4227 illegal = np.isnan(temp)
4228 v[illegal] = self.illegal_value
4229 return v
4231 def gradient(self, *args):
4232 NvrsGrad = self.vFuncNvrs.gradient(*args)
4233 marg_u = CRRAutilityP(*args, self.CRRA)
4234 grad = [g * marg_u for g in NvrsGrad]
4235 return grad
4237 def _eval_and_grad(self, *args):
4238 return (self.__call__(*args), self.gradient(*args))
4241class MargValueFuncCRRA(MetricObject):
4242 """
4243 A class for representing a marginal value function in models where the
4244 standard envelope condition of dvdm(state) = u'(c(state)) holds (with CRRA utility).
4246 Parameters
4247 ----------
4248 cFunc : function.
4249 Its first argument must be normalized market resources m.
4250 A real function representing the marginal value function composed
4251 with the inverse marginal utility function, defined on the state
4252 variables: uP_inv(dvdmFunc(state)). Called cFunc because when standard
4253 envelope condition applies, uP_inv(dvdm(state)) = cFunc(state).
4254 CRRA : float
4255 Coefficient of relative risk aversion.
4256 """
4258 distance_criteria = ["cFunc", "CRRA"]
4260 def __init__(self, cFunc, CRRA):
4261 self.cFunc = deepcopy(cFunc)
4262 self.CRRA = CRRA
4264 if hasattr(cFunc, "grid_list"):
4265 self.grid_list = cFunc.grid_list
4266 else:
4267 self.grid_list = None
4269 def __call__(self, *cFuncArgs):
4270 """
4271 Evaluate the marginal value function at given levels of market resources m.
4273 Parameters
4274 ----------
4275 cFuncArgs : floats or np.arrays
4276 Values of the state variables at which to evaluate the marginal
4277 value function.
4279 Returns
4280 -------
4281 vP : float or np.array
4282 Marginal lifetime value of beginning this period with state
4283 cFuncArgs
4284 """
4285 return CRRAutilityP(self.cFunc(*cFuncArgs), rho=self.CRRA)
4287 def derivativeX(self, *cFuncArgs):
4288 """
4289 Evaluate the derivative of the marginal value function with respect to
4290 market resources at given state; this is the marginal marginal value
4291 function.
4293 Parameters
4294 ----------
4295 cFuncArgs : floats or np.arrays
4296 State variables.
4298 Returns
4299 -------
4300 vPP : float or np.array
4301 Marginal marginal lifetime value of beginning this period with
4302 state cFuncArgs; has same size as inputs.
4304 """
4306 # The derivative method depends on the dimension of the function
4307 if isinstance(self.cFunc, HARKinterpolator1D):
4308 c, MPC = self.cFunc.eval_with_derivative(*cFuncArgs)
4310 elif hasattr(self.cFunc, "derivativeX"):
4311 c = self.cFunc(*cFuncArgs)
4312 MPC = self.cFunc.derivativeX(*cFuncArgs)
4314 else:
4315 raise Exception(
4316 "cFunc does not have a 'derivativeX' attribute. Can't compute"
4317 + "marginal marginal value."
4318 )
4320 return MPC * CRRAutilityPP(c, rho=self.CRRA)
4323class MargMargValueFuncCRRA(MetricObject):
4324 """
4325 A class for representing a marginal marginal value function in models where
4326 the standard envelope condition of dvdm = u'(c(state)) holds (with CRRA utility).
4328 Parameters
4329 ----------
4330 cFunc : function.
4331 Its first argument must be normalized market resources m.
4332 A real function representing the marginal value function composed
4333 with the inverse marginal utility function, defined on the state
4334 variables: uP_inv(dvdmFunc(state)). Called cFunc because when standard
4335 envelope condition applies, uP_inv(dvdm(state)) = cFunc(state).
4336 CRRA : float
4337 Coefficient of relative risk aversion.
4338 """
4340 distance_criteria = ["cFunc", "CRRA"]
4342 def __init__(self, cFunc, CRRA):
4343 self.cFunc = deepcopy(cFunc)
4344 self.CRRA = CRRA
4346 def __call__(self, *cFuncArgs):
4347 """
4348 Evaluate the marginal marginal value function at given levels of market
4349 resources m.
4351 Parameters
4352 ----------
4353 m : float or np.array
4354 Market resources (normalized by permanent income) whose marginal
4355 marginal value is to be found.
4357 Returns
4358 -------
4359 vPP : float or np.array
4360 Marginal marginal lifetime value of beginning this period with market
4361 resources m; has same size as input m.
4362 """
4364 # The derivative method depends on the dimension of the function
4365 if isinstance(self.cFunc, HARKinterpolator1D):
4366 c, MPC = self.cFunc.eval_with_derivative(*cFuncArgs)
4368 elif hasattr(self.cFunc, "derivativeX"):
4369 c = self.cFunc(*cFuncArgs)
4370 MPC = self.cFunc.derivativeX(*cFuncArgs)
4372 else:
4373 raise Exception(
4374 "cFunc does not have a 'derivativeX' attribute. Can't compute"
4375 + "marginal marginal value."
4376 )
4377 return MPC * CRRAutilityPP(c, rho=self.CRRA)