Coverage for HARK/interpolation.py: 96%
1554 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-02 05:14 +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)
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 y[x == self.x_list[0]] = self.y_list[0]
1074 return y
1076 def _der(self, x):
1077 """
1078 Returns the first derivative of the interpolated function at each value
1079 in x. Only called internally by HARKinterpolator1D.derivative (etc).
1080 """
1082 m = len(x)
1083 pos = np.searchsorted(self.x_list, x)
1084 dydx = np.zeros(m)
1085 if dydx.size > 0:
1086 out_bot = pos == 0
1087 out_top = pos == self.n
1088 in_bnds = np.logical_not(np.logical_or(out_bot, out_top))
1090 # Do the "in bounds" evaluation points
1091 i = pos[in_bnds]
1092 coeffs_in = self.coeffs[i, :]
1093 alpha = (x[in_bnds] - self.x_list[i - 1]) / (
1094 self.x_list[i] - self.x_list[i - 1]
1095 )
1096 dydx[in_bnds] = (
1097 coeffs_in[:, 1]
1098 + alpha * (2 * coeffs_in[:, 2] + alpha * 3 * coeffs_in[:, 3])
1099 ) / (self.x_list[i] - self.x_list[i - 1])
1101 # Do the "out of bounds" evaluation points
1102 dydx[out_bot] = self.coeffs[0, 1]
1103 alpha = x[out_top] - self.x_list[self.n - 1]
1104 dydx[out_top] = self.coeffs[self.n, 1] - self.coeffs[
1105 self.n, 2
1106 ] * self.coeffs[self.n, 3] * np.exp(alpha * self.coeffs[self.n, 3])
1107 return dydx
1109 def _evalAndDer(self, x):
1110 """
1111 Returns the level and first derivative of the function at each value in
1112 x. Only called internally by HARKinterpolator1D.eval_and_der (etc).
1113 """
1114 m = len(x)
1115 pos = np.searchsorted(self.x_list, x)
1116 y = np.zeros(m)
1117 dydx = np.zeros(m)
1118 if y.size > 0:
1119 out_bot = pos == 0
1120 out_top = pos == self.n
1121 in_bnds = np.logical_not(np.logical_or(out_bot, out_top))
1123 # Do the "in bounds" evaluation points
1124 i = pos[in_bnds]
1125 coeffs_in = self.coeffs[i, :]
1126 alpha = (x[in_bnds] - self.x_list[i - 1]) / (
1127 self.x_list[i] - self.x_list[i - 1]
1128 )
1129 y[in_bnds] = coeffs_in[:, 0] + alpha * (
1130 coeffs_in[:, 1] + alpha * (coeffs_in[:, 2] + alpha * coeffs_in[:, 3])
1131 )
1132 dydx[in_bnds] = (
1133 coeffs_in[:, 1]
1134 + alpha * (2 * coeffs_in[:, 2] + alpha * 3 * coeffs_in[:, 3])
1135 ) / (self.x_list[i] - self.x_list[i - 1])
1137 # Do the "out of bounds" evaluation points
1138 y[out_bot] = self.coeffs[0, 0] + self.coeffs[0, 1] * (
1139 x[out_bot] - self.x_list[0]
1140 )
1141 dydx[out_bot] = self.coeffs[0, 1]
1142 alpha = x[out_top] - self.x_list[self.n - 1]
1143 y[out_top] = (
1144 self.coeffs[self.n, 0]
1145 + x[out_top] * self.coeffs[self.n, 1]
1146 - self.coeffs[self.n, 2] * np.exp(alpha * self.coeffs[self.n, 3])
1147 )
1148 dydx[out_top] = self.coeffs[self.n, 1] - self.coeffs[
1149 self.n, 2
1150 ] * self.coeffs[self.n, 3] * np.exp(alpha * self.coeffs[self.n, 3])
1151 return y, dydx
1154class CubicHermiteInterp(HARKinterpolator1D):
1155 """
1156 An interpolating function using piecewise cubic splines. Matches level and
1157 slope of 1D function at gridpoints, smoothly interpolating in between.
1158 Extrapolation above highest gridpoint approaches a limiting linear function
1159 if desired (linear extrapolation also enabled.)
1161 NOTE: When no input is given for the limiting linear function, linear
1162 extrapolation is used above the highest gridpoint.
1164 Parameters
1165 ----------
1166 x_list : np.array
1167 List of x values composing the grid.
1168 y_list : np.array
1169 List of y values, representing f(x) at the points in x_list.
1170 dydx_list : np.array
1171 List of dydx values, representing f'(x) at the points in x_list
1172 intercept_limit : float
1173 Intercept of limiting linear function.
1174 slope_limit : float
1175 Slope of limiting linear function.
1176 lower_extrap : boolean
1177 Indicator for whether lower extrapolation is allowed. False means
1178 f(x) = NaN for x < min(x_list); True means linear extrapolation.
1179 """
1181 distance_criteria = ["x_list", "y_list", "dydx_list"]
1183 def __init__(
1184 self,
1185 x_list,
1186 y_list,
1187 dydx_list,
1188 intercept_limit=None,
1189 slope_limit=None,
1190 lower_extrap=False,
1191 ):
1192 self.x_list = (
1193 np.asarray(x_list)
1194 if _check_flatten(1, x_list)
1195 else np.array(x_list).flatten()
1196 )
1197 self.y_list = (
1198 np.asarray(y_list)
1199 if _check_flatten(1, y_list)
1200 else np.array(y_list).flatten()
1201 )
1202 self.dydx_list = (
1203 np.asarray(dydx_list)
1204 if _check_flatten(1, dydx_list)
1205 else np.array(dydx_list).flatten()
1206 )
1207 _check_grid_dimensions(1, self.y_list, self.x_list)
1208 _check_grid_dimensions(1, self.dydx_list, self.x_list)
1210 self.n = len(x_list)
1212 self._chs = CubicHermiteSpline(
1213 self.x_list, self.y_list, self.dydx_list, extrapolate=None
1214 )
1215 self.coeffs = np.flip(self._chs.c.T, 1)
1217 # Define lower extrapolation as linear function (or just NaN)
1218 if lower_extrap:
1219 temp = np.array([self.y_list[0], self.dydx_list[0], 0, 0])
1220 else:
1221 temp = np.array([np.nan, np.nan, np.nan, np.nan])
1223 self.coeffs = np.vstack((temp, self.coeffs))
1225 x1 = self.x_list[self.n - 1]
1226 y1 = self.y_list[self.n - 1]
1228 # Calculate extrapolation coefficients as a decay toward limiting function y = mx+b
1229 if slope_limit is None and intercept_limit is None:
1230 slope_limit = self.dydx_list[-1]
1231 intercept_limit = self.y_list[-1] - slope_limit * self.x_list[-1]
1232 gap = slope_limit * x1 + intercept_limit - y1
1233 slope = slope_limit - self.dydx_list[self.n - 1]
1234 if (gap != 0) and (slope <= 0):
1235 temp = np.array([intercept_limit, slope_limit, gap, slope / gap])
1236 elif slope > 0:
1237 # fixing a problem when slope is positive
1238 temp = np.array([intercept_limit, slope_limit, 0, 0])
1239 else:
1240 temp = np.array([intercept_limit, slope_limit, gap, 0])
1241 self.coeffs = np.vstack((self.coeffs, temp))
1243 def out_of_bounds(self, x):
1244 out_bot = x < self.x_list[0]
1245 out_top = x > self.x_list[-1]
1246 return out_bot, out_top
1248 def _evaluate(self, x):
1249 """
1250 Returns the level of the interpolated function at each value in x. Only
1251 called internally by HARKinterpolator1D.__call__ (etc).
1252 """
1253 out_bot, out_top = self.out_of_bounds(x)
1255 return self._eval_helper(x, out_bot, out_top)
1257 def _eval_helper(self, x, out_bot, out_top):
1258 y = self._chs(x)
1260 # Do the "out of bounds" evaluation points
1261 if any(out_bot):
1262 y[out_bot] = self.coeffs[0, 0] + self.coeffs[0, 1] * (
1263 x[out_bot] - self.x_list[0]
1264 )
1266 if any(out_top):
1267 alpha = x[out_top] - self.x_list[self.n - 1]
1268 y[out_top] = (
1269 self.coeffs[self.n, 0]
1270 + x[out_top] * self.coeffs[self.n, 1]
1271 - self.coeffs[self.n, 2] * np.exp(alpha * self.coeffs[self.n, 3])
1272 )
1274 return y
1276 def _der(self, x):
1277 """
1278 Returns the first derivative of the interpolated function at each value
1279 in x. Only called internally by HARKinterpolator1D.derivative (etc).
1280 """
1281 out_bot, out_top = self.out_of_bounds(x)
1283 return self._der_helper(x, out_bot, out_top)
1285 def _der_helper(self, x, out_bot, out_top):
1286 dydx = self._chs(x, nu=1)
1288 # Do the "out of bounds" evaluation points
1289 if any(out_bot):
1290 dydx[out_bot] = self.coeffs[0, 1]
1291 if any(out_top):
1292 alpha = x[out_top] - self.x_list[self.n - 1]
1293 dydx[out_top] = self.coeffs[self.n, 1] - self.coeffs[
1294 self.n, 2
1295 ] * self.coeffs[self.n, 3] * np.exp(alpha * self.coeffs[self.n, 3])
1297 return dydx
1299 def _evalAndDer(self, x):
1300 """
1301 Returns the level and first derivative of the function at each value in
1302 x. Only called internally by HARKinterpolator1D.eval_and_der (etc).
1303 """
1304 out_bot, out_top = self.out_of_bounds(x)
1305 y = self._eval_helper(x, out_bot, out_top)
1306 dydx = self._der_helper(x, out_bot, out_top)
1307 return y, dydx
1309 def der_interp(self, nu=1):
1310 """
1311 Construct a new piecewise polynomial representing the derivative.
1312 See `scipy` for additional documentation:
1313 https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.CubicHermiteSpline.html
1314 """
1315 return self._chs.derivative(nu)
1317 def antider_interp(self, nu=1):
1318 """
1319 Construct a new piecewise polynomial representing the antiderivative.
1321 Antiderivative is also the indefinite integral of the function,
1322 and derivative is its inverse operation.
1324 See `scipy` for additional documentation:
1325 https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.CubicHermiteSpline.html
1326 """
1327 return self._chs.antiderivative(nu)
1329 def integrate(self, a, b, extrapolate=False):
1330 """
1331 Compute a definite integral over a piecewise polynomial.
1333 See `scipy` for additional documentation:
1334 https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.CubicHermiteSpline.html
1335 """
1336 return self._chs.integrate(a, b, extrapolate)
1338 def roots(self, discontinuity=True, extrapolate=False):
1339 """
1340 Find real roots of the the piecewise polynomial.
1342 See `scipy` for additional documentation:
1343 https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.CubicHermiteSpline.html
1344 """
1345 return self._chs.roots(discontinuity, extrapolate)
1347 def solve(self, y=0.0, discontinuity=True, extrapolate=False):
1348 """
1349 Find real solutions of the the equation ``pp(x) == y``.
1351 See `scipy` for additional documentation:
1352 https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.CubicHermiteSpline.html
1353 """
1354 return self._chs.solve(y, discontinuity, extrapolate)
1357class BilinearInterp(HARKinterpolator2D):
1358 """
1359 Bilinear full (or tensor) grid interpolation of a function f(x,y).
1361 Parameters
1362 ----------
1363 f_values : numpy.array
1364 An array of size (x_n,y_n) such that f_values[i,j] = f(x_list[i],y_list[j])
1365 x_list : numpy.array
1366 An array of x values, with length designated x_n.
1367 y_list : numpy.array
1368 An array of y values, with length designated y_n.
1369 xSearchFunc : function
1370 An optional function that returns the reference location for x values:
1371 indices = xSearchFunc(x_list,x). Default is np.searchsorted
1372 ySearchFunc : function
1373 An optional function that returns the reference location for y values:
1374 indices = ySearchFunc(y_list,y). Default is np.searchsorted
1375 """
1377 distance_criteria = ["x_list", "y_list", "f_values"]
1379 def __init__(self, f_values, x_list, y_list, xSearchFunc=None, ySearchFunc=None):
1380 self.f_values = f_values
1381 self.x_list = (
1382 np.array(x_list)
1383 if _check_flatten(1, x_list)
1384 else np.array(x_list).flatten()
1385 )
1386 self.y_list = (
1387 np.array(y_list)
1388 if _check_flatten(1, y_list)
1389 else np.array(y_list).flatten()
1390 )
1391 _check_grid_dimensions(2, self.f_values, self.x_list, self.y_list)
1392 self.x_n = x_list.size
1393 self.y_n = y_list.size
1394 if xSearchFunc is None:
1395 xSearchFunc = np.searchsorted
1396 if ySearchFunc is None:
1397 ySearchFunc = np.searchsorted
1398 self.xSearchFunc = xSearchFunc
1399 self.ySearchFunc = ySearchFunc
1401 def _evaluate(self, x, y):
1402 """
1403 Returns the level of the interpolated function at each value in x,y.
1404 Only called internally by HARKinterpolator2D.__call__ (etc).
1405 """
1406 x_pos = self.xSearchFunc(self.x_list, x)
1407 x_pos[x_pos < 1] = 1
1408 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
1409 y_pos = self.ySearchFunc(self.y_list, y)
1410 y_pos[y_pos < 1] = 1
1411 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
1412 alpha = (x - self.x_list[x_pos - 1]) / (
1413 self.x_list[x_pos] - self.x_list[x_pos - 1]
1414 )
1415 beta = (y - self.y_list[y_pos - 1]) / (
1416 self.y_list[y_pos] - self.y_list[y_pos - 1]
1417 )
1418 f = (
1419 (1 - alpha) * (1 - beta) * self.f_values[x_pos - 1, y_pos - 1]
1420 + (1 - alpha) * beta * self.f_values[x_pos - 1, y_pos]
1421 + alpha * (1 - beta) * self.f_values[x_pos, y_pos - 1]
1422 + alpha * beta * self.f_values[x_pos, y_pos]
1423 )
1424 return f
1426 def _derX(self, x, y):
1427 """
1428 Returns the derivative with respect to x of the interpolated function
1429 at each value in x,y. Only called internally by HARKinterpolator2D.derivativeX.
1430 """
1431 x_pos = self.xSearchFunc(self.x_list, x)
1432 x_pos[x_pos < 1] = 1
1433 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
1434 y_pos = self.ySearchFunc(self.y_list, y)
1435 y_pos[y_pos < 1] = 1
1436 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
1437 beta = (y - self.y_list[y_pos - 1]) / (
1438 self.y_list[y_pos] - self.y_list[y_pos - 1]
1439 )
1440 dfdx = (
1441 (
1442 (1 - beta) * self.f_values[x_pos, y_pos - 1]
1443 + beta * self.f_values[x_pos, y_pos]
1444 )
1445 - (
1446 (1 - beta) * self.f_values[x_pos - 1, y_pos - 1]
1447 + beta * self.f_values[x_pos - 1, y_pos]
1448 )
1449 ) / (self.x_list[x_pos] - self.x_list[x_pos - 1])
1450 return dfdx
1452 def _derY(self, x, y):
1453 """
1454 Returns the derivative with respect to y of the interpolated function
1455 at each value in x,y. Only called internally by HARKinterpolator2D.derivativeY.
1456 """
1457 x_pos = self.xSearchFunc(self.x_list, x)
1458 x_pos[x_pos < 1] = 1
1459 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
1460 y_pos = self.ySearchFunc(self.y_list, y)
1461 y_pos[y_pos < 1] = 1
1462 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
1463 alpha = (x - self.x_list[x_pos - 1]) / (
1464 self.x_list[x_pos] - self.x_list[x_pos - 1]
1465 )
1466 dfdy = (
1467 (
1468 (1 - alpha) * self.f_values[x_pos - 1, y_pos]
1469 + alpha * self.f_values[x_pos, y_pos]
1470 )
1471 - (
1472 (1 - alpha) * self.f_values[x_pos - 1, y_pos - 1]
1473 + alpha * self.f_values[x_pos, y_pos - 1]
1474 )
1475 ) / (self.y_list[y_pos] - self.y_list[y_pos - 1])
1476 return dfdy
1479class TrilinearInterp(HARKinterpolator3D):
1480 """
1481 Trilinear full (or tensor) grid interpolation of a function f(x,y,z).
1483 Parameters
1484 ----------
1485 f_values : numpy.array
1486 An array of size (x_n,y_n,z_n) such that f_values[i,j,k] =
1487 f(x_list[i],y_list[j],z_list[k])
1488 x_list : numpy.array
1489 An array of x values, with length designated x_n.
1490 y_list : numpy.array
1491 An array of y values, with length designated y_n.
1492 z_list : numpy.array
1493 An array of z values, with length designated z_n.
1494 xSearchFunc : function
1495 An optional function that returns the reference location for x values:
1496 indices = xSearchFunc(x_list,x). Default is np.searchsorted
1497 ySearchFunc : function
1498 An optional function that returns the reference location for y values:
1499 indices = ySearchFunc(y_list,y). Default is np.searchsorted
1500 zSearchFunc : function
1501 An optional function that returns the reference location for z values:
1502 indices = zSearchFunc(z_list,z). Default is np.searchsorted
1503 """
1505 distance_criteria = ["f_values", "x_list", "y_list", "z_list"]
1507 def __init__(
1508 self,
1509 f_values,
1510 x_list,
1511 y_list,
1512 z_list,
1513 xSearchFunc=None,
1514 ySearchFunc=None,
1515 zSearchFunc=None,
1516 ):
1517 self.f_values = f_values
1518 self.x_list = (
1519 np.array(x_list)
1520 if _check_flatten(1, x_list)
1521 else np.array(x_list).flatten()
1522 )
1523 self.y_list = (
1524 np.array(y_list)
1525 if _check_flatten(1, y_list)
1526 else np.array(y_list).flatten()
1527 )
1528 self.z_list = (
1529 np.array(z_list)
1530 if _check_flatten(1, z_list)
1531 else np.array(z_list).flatten()
1532 )
1533 _check_grid_dimensions(3, self.f_values, self.x_list, self.y_list, self.z_list)
1534 self.x_n = x_list.size
1535 self.y_n = y_list.size
1536 self.z_n = z_list.size
1537 if xSearchFunc is None:
1538 xSearchFunc = np.searchsorted
1539 if ySearchFunc is None:
1540 ySearchFunc = np.searchsorted
1541 if zSearchFunc is None:
1542 zSearchFunc = np.searchsorted
1543 self.xSearchFunc = xSearchFunc
1544 self.ySearchFunc = ySearchFunc
1545 self.zSearchFunc = zSearchFunc
1547 def _evaluate(self, x, y, z):
1548 """
1549 Returns the level of the interpolated function at each value in x,y,z.
1550 Only called internally by HARKinterpolator3D.__call__ (etc).
1551 """
1552 x_pos = self.xSearchFunc(self.x_list, x)
1553 x_pos[x_pos < 1] = 1
1554 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
1555 y_pos = self.ySearchFunc(self.y_list, y)
1556 y_pos[y_pos < 1] = 1
1557 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
1558 z_pos = self.zSearchFunc(self.z_list, z)
1559 z_pos[z_pos < 1] = 1
1560 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
1561 alpha = (x - self.x_list[x_pos - 1]) / (
1562 self.x_list[x_pos] - self.x_list[x_pos - 1]
1563 )
1564 beta = (y - self.y_list[y_pos - 1]) / (
1565 self.y_list[y_pos] - self.y_list[y_pos - 1]
1566 )
1567 gamma = (z - self.z_list[z_pos - 1]) / (
1568 self.z_list[z_pos] - self.z_list[z_pos - 1]
1569 )
1570 f = (
1571 (1 - alpha)
1572 * (1 - beta)
1573 * (1 - gamma)
1574 * self.f_values[x_pos - 1, y_pos - 1, z_pos - 1]
1575 + (1 - alpha)
1576 * (1 - beta)
1577 * gamma
1578 * self.f_values[x_pos - 1, y_pos - 1, z_pos]
1579 + (1 - alpha)
1580 * beta
1581 * (1 - gamma)
1582 * self.f_values[x_pos - 1, y_pos, z_pos - 1]
1583 + (1 - alpha) * beta * gamma * self.f_values[x_pos - 1, y_pos, z_pos]
1584 + alpha
1585 * (1 - beta)
1586 * (1 - gamma)
1587 * self.f_values[x_pos, y_pos - 1, z_pos - 1]
1588 + alpha * (1 - beta) * gamma * self.f_values[x_pos, y_pos - 1, z_pos]
1589 + alpha * beta * (1 - gamma) * self.f_values[x_pos, y_pos, z_pos - 1]
1590 + alpha * beta * gamma * self.f_values[x_pos, y_pos, z_pos]
1591 )
1592 return f
1594 def _derX(self, x, y, z):
1595 """
1596 Returns the derivative with respect to x of the interpolated function
1597 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeX.
1598 """
1599 x_pos = self.xSearchFunc(self.x_list, x)
1600 x_pos[x_pos < 1] = 1
1601 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
1602 y_pos = self.ySearchFunc(self.y_list, y)
1603 y_pos[y_pos < 1] = 1
1604 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
1605 z_pos = self.zSearchFunc(self.z_list, z)
1606 z_pos[z_pos < 1] = 1
1607 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
1608 beta = (y - self.y_list[y_pos - 1]) / (
1609 self.y_list[y_pos] - self.y_list[y_pos - 1]
1610 )
1611 gamma = (z - self.z_list[z_pos - 1]) / (
1612 self.z_list[z_pos] - self.z_list[z_pos - 1]
1613 )
1614 dfdx = (
1615 (
1616 (1 - beta) * (1 - gamma) * self.f_values[x_pos, y_pos - 1, z_pos - 1]
1617 + (1 - beta) * gamma * self.f_values[x_pos, y_pos - 1, z_pos]
1618 + beta * (1 - gamma) * self.f_values[x_pos, y_pos, z_pos - 1]
1619 + beta * gamma * self.f_values[x_pos, y_pos, z_pos]
1620 )
1621 - (
1622 (1 - beta)
1623 * (1 - gamma)
1624 * self.f_values[x_pos - 1, y_pos - 1, z_pos - 1]
1625 + (1 - beta) * gamma * self.f_values[x_pos - 1, y_pos - 1, z_pos]
1626 + beta * (1 - gamma) * self.f_values[x_pos - 1, y_pos, z_pos - 1]
1627 + beta * gamma * self.f_values[x_pos - 1, y_pos, z_pos]
1628 )
1629 ) / (self.x_list[x_pos] - self.x_list[x_pos - 1])
1630 return dfdx
1632 def _derY(self, x, y, z):
1633 """
1634 Returns the derivative with respect to y of the interpolated function
1635 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeY.
1636 """
1637 x_pos = self.xSearchFunc(self.x_list, x)
1638 x_pos[x_pos < 1] = 1
1639 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
1640 y_pos = self.ySearchFunc(self.y_list, y)
1641 y_pos[y_pos < 1] = 1
1642 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
1643 z_pos = self.zSearchFunc(self.z_list, z)
1644 z_pos[z_pos < 1] = 1
1645 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
1646 alpha = (x - self.x_list[x_pos - 1]) / (
1647 self.x_list[x_pos] - self.x_list[x_pos - 1]
1648 )
1649 gamma = (z - self.z_list[z_pos - 1]) / (
1650 self.z_list[z_pos] - self.z_list[z_pos - 1]
1651 )
1652 dfdy = (
1653 (
1654 (1 - alpha) * (1 - gamma) * self.f_values[x_pos - 1, y_pos, z_pos - 1]
1655 + (1 - alpha) * gamma * self.f_values[x_pos - 1, y_pos, z_pos]
1656 + alpha * (1 - gamma) * self.f_values[x_pos, y_pos, z_pos - 1]
1657 + alpha * gamma * self.f_values[x_pos, y_pos, z_pos]
1658 )
1659 - (
1660 (1 - alpha)
1661 * (1 - gamma)
1662 * self.f_values[x_pos - 1, y_pos - 1, z_pos - 1]
1663 + (1 - alpha) * gamma * self.f_values[x_pos - 1, y_pos - 1, z_pos]
1664 + alpha * (1 - gamma) * self.f_values[x_pos, y_pos - 1, z_pos - 1]
1665 + alpha * gamma * self.f_values[x_pos, y_pos - 1, z_pos]
1666 )
1667 ) / (self.y_list[y_pos] - self.y_list[y_pos - 1])
1668 return dfdy
1670 def _derZ(self, x, y, z):
1671 """
1672 Returns the derivative with respect to z of the interpolated function
1673 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeZ.
1674 """
1675 x_pos = self.xSearchFunc(self.x_list, x)
1676 x_pos[x_pos < 1] = 1
1677 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
1678 y_pos = self.ySearchFunc(self.y_list, y)
1679 y_pos[y_pos < 1] = 1
1680 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
1681 z_pos = self.zSearchFunc(self.z_list, z)
1682 z_pos[z_pos < 1] = 1
1683 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
1684 alpha = (x - self.x_list[x_pos - 1]) / (
1685 self.x_list[x_pos] - self.x_list[x_pos - 1]
1686 )
1687 beta = (y - self.y_list[y_pos - 1]) / (
1688 self.y_list[y_pos] - self.y_list[y_pos - 1]
1689 )
1690 dfdz = (
1691 (
1692 (1 - alpha) * (1 - beta) * self.f_values[x_pos - 1, y_pos - 1, z_pos]
1693 + (1 - alpha) * beta * self.f_values[x_pos - 1, y_pos, z_pos]
1694 + alpha * (1 - beta) * self.f_values[x_pos, y_pos - 1, z_pos]
1695 + alpha * beta * self.f_values[x_pos, y_pos, z_pos]
1696 )
1697 - (
1698 (1 - alpha)
1699 * (1 - beta)
1700 * self.f_values[x_pos - 1, y_pos - 1, z_pos - 1]
1701 + (1 - alpha) * beta * self.f_values[x_pos - 1, y_pos, z_pos - 1]
1702 + alpha * (1 - beta) * self.f_values[x_pos, y_pos - 1, z_pos - 1]
1703 + alpha * beta * self.f_values[x_pos, y_pos, z_pos - 1]
1704 )
1705 ) / (self.z_list[z_pos] - self.z_list[z_pos - 1])
1706 return dfdz
1709class QuadlinearInterp(HARKinterpolator4D):
1710 """
1711 Quadlinear full (or tensor) grid interpolation of a function f(w,x,y,z).
1713 Parameters
1714 ----------
1715 f_values : numpy.array
1716 An array of size (w_n,x_n,y_n,z_n) such that f_values[i,j,k,l] =
1717 f(w_list[i],x_list[j],y_list[k],z_list[l])
1718 w_list : numpy.array
1719 An array of x values, with length designated w_n.
1720 x_list : numpy.array
1721 An array of x values, with length designated x_n.
1722 y_list : numpy.array
1723 An array of y values, with length designated y_n.
1724 z_list : numpy.array
1725 An array of z values, with length designated z_n.
1726 wSearchFunc : function
1727 An optional function that returns the reference location for w values:
1728 indices = wSearchFunc(w_list,w). Default is np.searchsorted
1729 xSearchFunc : function
1730 An optional function that returns the reference location for x values:
1731 indices = xSearchFunc(x_list,x). Default is np.searchsorted
1732 ySearchFunc : function
1733 An optional function that returns the reference location for y values:
1734 indices = ySearchFunc(y_list,y). Default is np.searchsorted
1735 zSearchFunc : function
1736 An optional function that returns the reference location for z values:
1737 indices = zSearchFunc(z_list,z). Default is np.searchsorted
1738 """
1740 distance_criteria = ["f_values", "w_list", "x_list", "y_list", "z_list"]
1742 def __init__(
1743 self,
1744 f_values,
1745 w_list,
1746 x_list,
1747 y_list,
1748 z_list,
1749 wSearchFunc=None,
1750 xSearchFunc=None,
1751 ySearchFunc=None,
1752 zSearchFunc=None,
1753 ):
1754 self.f_values = f_values
1755 self.w_list = (
1756 np.array(w_list)
1757 if _check_flatten(1, w_list)
1758 else np.array(w_list).flatten()
1759 )
1760 self.x_list = (
1761 np.array(x_list)
1762 if _check_flatten(1, x_list)
1763 else np.array(x_list).flatten()
1764 )
1765 self.y_list = (
1766 np.array(y_list)
1767 if _check_flatten(1, y_list)
1768 else np.array(y_list).flatten()
1769 )
1770 self.z_list = (
1771 np.array(z_list)
1772 if _check_flatten(1, z_list)
1773 else np.array(z_list).flatten()
1774 )
1775 _check_grid_dimensions(
1776 4, self.f_values, self.w_list, self.x_list, self.y_list, self.z_list
1777 )
1778 self.w_n = w_list.size
1779 self.x_n = x_list.size
1780 self.y_n = y_list.size
1781 self.z_n = z_list.size
1782 if wSearchFunc is None:
1783 wSearchFunc = np.searchsorted
1784 if xSearchFunc is None:
1785 xSearchFunc = np.searchsorted
1786 if ySearchFunc is None:
1787 ySearchFunc = np.searchsorted
1788 if zSearchFunc is None:
1789 zSearchFunc = np.searchsorted
1790 self.wSearchFunc = wSearchFunc
1791 self.xSearchFunc = xSearchFunc
1792 self.ySearchFunc = ySearchFunc
1793 self.zSearchFunc = zSearchFunc
1795 def _evaluate(self, w, x, y, z):
1796 """
1797 Returns the level of the interpolated function at each value in x,y,z.
1798 Only called internally by HARKinterpolator4D.__call__ (etc).
1799 """
1800 w_pos = self.wSearchFunc(self.w_list, w)
1801 w_pos[w_pos < 1] = 1
1802 w_pos[w_pos > self.w_n - 1] = self.w_n - 1
1803 x_pos = self.xSearchFunc(self.x_list, x)
1804 x_pos[x_pos < 1] = 1
1805 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
1806 y_pos = self.ySearchFunc(self.y_list, y)
1807 y_pos[y_pos < 1] = 1
1808 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
1809 z_pos = self.zSearchFunc(self.z_list, z)
1810 z_pos[z_pos < 1] = 1
1811 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
1812 i = w_pos # for convenience
1813 j = x_pos
1814 k = y_pos
1815 l = z_pos
1816 alpha = (w - self.w_list[i - 1]) / (self.w_list[i] - self.w_list[i - 1])
1817 beta = (x - self.x_list[j - 1]) / (self.x_list[j] - self.x_list[j - 1])
1818 gamma = (y - self.y_list[k - 1]) / (self.y_list[k] - self.y_list[k - 1])
1819 delta = (z - self.z_list[l - 1]) / (self.z_list[l] - self.z_list[l - 1])
1820 f = (1 - alpha) * (
1821 (1 - beta)
1822 * (
1823 (1 - gamma) * (1 - delta) * self.f_values[i - 1, j - 1, k - 1, l - 1]
1824 + (1 - gamma) * delta * self.f_values[i - 1, j - 1, k - 1, l]
1825 + gamma * (1 - delta) * self.f_values[i - 1, j - 1, k, l - 1]
1826 + gamma * delta * self.f_values[i - 1, j - 1, k, l]
1827 )
1828 + beta
1829 * (
1830 (1 - gamma) * (1 - delta) * self.f_values[i - 1, j, k - 1, l - 1]
1831 + (1 - gamma) * delta * self.f_values[i - 1, j, k - 1, l]
1832 + gamma * (1 - delta) * self.f_values[i - 1, j, k, l - 1]
1833 + gamma * delta * self.f_values[i - 1, j, k, l]
1834 )
1835 ) + alpha * (
1836 (1 - beta)
1837 * (
1838 (1 - gamma) * (1 - delta) * self.f_values[i, j - 1, k - 1, l - 1]
1839 + (1 - gamma) * delta * self.f_values[i, j - 1, k - 1, l]
1840 + gamma * (1 - delta) * self.f_values[i, j - 1, k, l - 1]
1841 + gamma * delta * self.f_values[i, j - 1, k, l]
1842 )
1843 + beta
1844 * (
1845 (1 - gamma) * (1 - delta) * self.f_values[i, j, k - 1, l - 1]
1846 + (1 - gamma) * delta * self.f_values[i, j, k - 1, l]
1847 + gamma * (1 - delta) * self.f_values[i, j, k, l - 1]
1848 + gamma * delta * self.f_values[i, j, k, l]
1849 )
1850 )
1851 return f
1853 def _derW(self, w, x, y, z):
1854 """
1855 Returns the derivative with respect to w of the interpolated function
1856 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeW.
1857 """
1858 w_pos = self.wSearchFunc(self.w_list, w)
1859 w_pos[w_pos < 1] = 1
1860 w_pos[w_pos > self.w_n - 1] = self.w_n - 1
1861 x_pos = self.xSearchFunc(self.x_list, x)
1862 x_pos[x_pos < 1] = 1
1863 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
1864 y_pos = self.ySearchFunc(self.y_list, y)
1865 y_pos[y_pos < 1] = 1
1866 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
1867 z_pos = self.zSearchFunc(self.z_list, z)
1868 z_pos[z_pos < 1] = 1
1869 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
1870 i = w_pos # for convenience
1871 j = x_pos
1872 k = y_pos
1873 l = z_pos
1874 beta = (x - self.x_list[j - 1]) / (self.x_list[j] - self.x_list[j - 1])
1875 gamma = (y - self.y_list[k - 1]) / (self.y_list[k] - self.y_list[k - 1])
1876 delta = (z - self.z_list[l - 1]) / (self.z_list[l] - self.z_list[l - 1])
1877 dfdw = (
1878 (
1879 (1 - beta)
1880 * (1 - gamma)
1881 * (1 - delta)
1882 * self.f_values[i, j - 1, k - 1, l - 1]
1883 + (1 - beta) * (1 - gamma) * delta * self.f_values[i, j - 1, k - 1, l]
1884 + (1 - beta) * gamma * (1 - delta) * self.f_values[i, j - 1, k, l - 1]
1885 + (1 - beta) * gamma * delta * self.f_values[i, j - 1, k, l]
1886 + beta * (1 - gamma) * (1 - delta) * self.f_values[i, j, k - 1, l - 1]
1887 + beta * (1 - gamma) * delta * self.f_values[i, j, k - 1, l]
1888 + beta * gamma * (1 - delta) * self.f_values[i, j, k, l - 1]
1889 + beta * gamma * delta * self.f_values[i, j, k, l]
1890 )
1891 - (
1892 (1 - beta)
1893 * (1 - gamma)
1894 * (1 - delta)
1895 * self.f_values[i - 1, j - 1, k - 1, l - 1]
1896 + (1 - beta)
1897 * (1 - gamma)
1898 * delta
1899 * self.f_values[i - 1, j - 1, k - 1, l]
1900 + (1 - beta)
1901 * gamma
1902 * (1 - delta)
1903 * self.f_values[i - 1, j - 1, k, l - 1]
1904 + (1 - beta) * gamma * delta * self.f_values[i - 1, j - 1, k, l]
1905 + beta
1906 * (1 - gamma)
1907 * (1 - delta)
1908 * self.f_values[i - 1, j, k - 1, l - 1]
1909 + beta * (1 - gamma) * delta * self.f_values[i - 1, j, k - 1, l]
1910 + beta * gamma * (1 - delta) * self.f_values[i - 1, j, k, l - 1]
1911 + beta * gamma * delta * self.f_values[i - 1, j, k, l]
1912 )
1913 ) / (self.w_list[i] - self.w_list[i - 1])
1914 return dfdw
1916 def _derX(self, w, x, y, z):
1917 """
1918 Returns the derivative with respect to x of the interpolated function
1919 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeX.
1920 """
1921 w_pos = self.wSearchFunc(self.w_list, w)
1922 w_pos[w_pos < 1] = 1
1923 w_pos[w_pos > self.w_n - 1] = self.w_n - 1
1924 x_pos = self.xSearchFunc(self.x_list, x)
1925 x_pos[x_pos < 1] = 1
1926 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
1927 y_pos = self.ySearchFunc(self.y_list, y)
1928 y_pos[y_pos < 1] = 1
1929 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
1930 z_pos = self.zSearchFunc(self.z_list, z)
1931 z_pos[z_pos < 1] = 1
1932 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
1933 i = w_pos # for convenience
1934 j = x_pos
1935 k = y_pos
1936 l = z_pos
1937 alpha = (w - self.w_list[i - 1]) / (self.w_list[i] - self.w_list[i - 1])
1938 gamma = (y - self.y_list[k - 1]) / (self.y_list[k] - self.y_list[k - 1])
1939 delta = (z - self.z_list[l - 1]) / (self.z_list[l] - self.z_list[l - 1])
1940 dfdx = (
1941 (
1942 (1 - alpha)
1943 * (1 - gamma)
1944 * (1 - delta)
1945 * self.f_values[i - 1, j, k - 1, l - 1]
1946 + (1 - alpha) * (1 - gamma) * delta * self.f_values[i - 1, j, k - 1, l]
1947 + (1 - alpha) * gamma * (1 - delta) * self.f_values[i - 1, j, k, l - 1]
1948 + (1 - alpha) * gamma * delta * self.f_values[i - 1, j, k, l]
1949 + alpha * (1 - gamma) * (1 - delta) * self.f_values[i, j, k - 1, l - 1]
1950 + alpha * (1 - gamma) * delta * self.f_values[i, j, k - 1, l]
1951 + alpha * gamma * (1 - delta) * self.f_values[i, j, k, l - 1]
1952 + alpha * gamma * delta * self.f_values[i, j, k, l]
1953 )
1954 - (
1955 (1 - alpha)
1956 * (1 - gamma)
1957 * (1 - delta)
1958 * self.f_values[i - 1, j - 1, k - 1, l - 1]
1959 + (1 - alpha)
1960 * (1 - gamma)
1961 * delta
1962 * self.f_values[i - 1, j - 1, k - 1, l]
1963 + (1 - alpha)
1964 * gamma
1965 * (1 - delta)
1966 * self.f_values[i - 1, j - 1, k, l - 1]
1967 + (1 - alpha) * gamma * delta * self.f_values[i - 1, j - 1, k, l]
1968 + alpha
1969 * (1 - gamma)
1970 * (1 - delta)
1971 * self.f_values[i, j - 1, k - 1, l - 1]
1972 + alpha * (1 - gamma) * delta * self.f_values[i, j - 1, k - 1, l]
1973 + alpha * gamma * (1 - delta) * self.f_values[i, j - 1, k, l - 1]
1974 + alpha * gamma * delta * self.f_values[i, j - 1, k, l]
1975 )
1976 ) / (self.x_list[j] - self.x_list[j - 1])
1977 return dfdx
1979 def _derY(self, w, x, y, z):
1980 """
1981 Returns the derivative with respect to y of the interpolated function
1982 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeY.
1983 """
1984 w_pos = self.wSearchFunc(self.w_list, w)
1985 w_pos[w_pos < 1] = 1
1986 w_pos[w_pos > self.w_n - 1] = self.w_n - 1
1987 x_pos = self.xSearchFunc(self.x_list, x)
1988 x_pos[x_pos < 1] = 1
1989 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
1990 y_pos = self.ySearchFunc(self.y_list, y)
1991 y_pos[y_pos < 1] = 1
1992 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
1993 z_pos = self.zSearchFunc(self.z_list, z)
1994 z_pos[z_pos < 1] = 1
1995 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
1996 i = w_pos # for convenience
1997 j = x_pos
1998 k = y_pos
1999 l = z_pos
2000 alpha = (w - self.w_list[i - 1]) / (self.w_list[i] - self.w_list[i - 1])
2001 beta = (x - self.x_list[j - 1]) / (self.x_list[j] - self.x_list[j - 1])
2002 delta = (z - self.z_list[l - 1]) / (self.z_list[l] - self.z_list[l - 1])
2003 dfdy = (
2004 (
2005 (1 - alpha)
2006 * (1 - beta)
2007 * (1 - delta)
2008 * self.f_values[i - 1, j - 1, k, l - 1]
2009 + (1 - alpha) * (1 - beta) * delta * self.f_values[i - 1, j - 1, k, l]
2010 + (1 - alpha) * beta * (1 - delta) * self.f_values[i - 1, j, k, l - 1]
2011 + (1 - alpha) * beta * delta * self.f_values[i - 1, j, k, l]
2012 + alpha * (1 - beta) * (1 - delta) * self.f_values[i, j - 1, k, l - 1]
2013 + alpha * (1 - beta) * delta * self.f_values[i, j - 1, k, l]
2014 + alpha * beta * (1 - delta) * self.f_values[i, j, k, l - 1]
2015 + alpha * beta * delta * self.f_values[i, j, k, l]
2016 )
2017 - (
2018 (1 - alpha)
2019 * (1 - beta)
2020 * (1 - delta)
2021 * self.f_values[i - 1, j - 1, k - 1, l - 1]
2022 + (1 - alpha)
2023 * (1 - beta)
2024 * delta
2025 * self.f_values[i - 1, j - 1, k - 1, l]
2026 + (1 - alpha)
2027 * beta
2028 * (1 - delta)
2029 * self.f_values[i - 1, j, k - 1, l - 1]
2030 + (1 - alpha) * beta * delta * self.f_values[i - 1, j, k - 1, l]
2031 + alpha
2032 * (1 - beta)
2033 * (1 - delta)
2034 * self.f_values[i, j - 1, k - 1, l - 1]
2035 + alpha * (1 - beta) * delta * self.f_values[i, j - 1, k - 1, l]
2036 + alpha * beta * (1 - delta) * self.f_values[i, j, k - 1, l - 1]
2037 + alpha * beta * delta * self.f_values[i, j, k - 1, l]
2038 )
2039 ) / (self.y_list[k] - self.y_list[k - 1])
2040 return dfdy
2042 def _derZ(self, w, x, y, z):
2043 """
2044 Returns the derivative with respect to z of the interpolated function
2045 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeZ.
2046 """
2047 w_pos = self.wSearchFunc(self.w_list, w)
2048 w_pos[w_pos < 1] = 1
2049 w_pos[w_pos > self.w_n - 1] = self.w_n - 1
2050 x_pos = self.xSearchFunc(self.x_list, x)
2051 x_pos[x_pos < 1] = 1
2052 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
2053 y_pos = self.ySearchFunc(self.y_list, y)
2054 y_pos[y_pos < 1] = 1
2055 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
2056 z_pos = self.zSearchFunc(self.z_list, z)
2057 z_pos[z_pos < 1] = 1
2058 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
2059 i = w_pos # for convenience
2060 j = x_pos
2061 k = y_pos
2062 l = z_pos
2063 alpha = (w - self.w_list[i - 1]) / (self.w_list[i] - self.w_list[i - 1])
2064 beta = (x - self.x_list[j - 1]) / (self.x_list[j] - self.x_list[j - 1])
2065 gamma = (y - self.y_list[k - 1]) / (self.y_list[k] - self.y_list[k - 1])
2066 dfdz = (
2067 (
2068 (1 - alpha)
2069 * (1 - beta)
2070 * (1 - gamma)
2071 * self.f_values[i - 1, j - 1, k - 1, l]
2072 + (1 - alpha) * (1 - beta) * gamma * self.f_values[i - 1, j - 1, k, l]
2073 + (1 - alpha) * beta * (1 - gamma) * self.f_values[i - 1, j, k - 1, l]
2074 + (1 - alpha) * beta * gamma * self.f_values[i - 1, j, k, l]
2075 + alpha * (1 - beta) * (1 - gamma) * self.f_values[i, j - 1, k - 1, l]
2076 + alpha * (1 - beta) * gamma * self.f_values[i, j - 1, k, l]
2077 + alpha * beta * (1 - gamma) * self.f_values[i, j, k - 1, l]
2078 + alpha * beta * gamma * self.f_values[i, j, k, l]
2079 )
2080 - (
2081 (1 - alpha)
2082 * (1 - beta)
2083 * (1 - gamma)
2084 * self.f_values[i - 1, j - 1, k - 1, l - 1]
2085 + (1 - alpha)
2086 * (1 - beta)
2087 * gamma
2088 * self.f_values[i - 1, j - 1, k, l - 1]
2089 + (1 - alpha)
2090 * beta
2091 * (1 - gamma)
2092 * self.f_values[i - 1, j, k - 1, l - 1]
2093 + (1 - alpha) * beta * gamma * self.f_values[i - 1, j, k, l - 1]
2094 + alpha
2095 * (1 - beta)
2096 * (1 - gamma)
2097 * self.f_values[i, j - 1, k - 1, l - 1]
2098 + alpha * (1 - beta) * gamma * self.f_values[i, j - 1, k, l - 1]
2099 + alpha * beta * (1 - gamma) * self.f_values[i, j, k - 1, l - 1]
2100 + alpha * beta * gamma * self.f_values[i, j, k, l - 1]
2101 )
2102 ) / (self.z_list[l] - self.z_list[l - 1])
2103 return dfdz
2106class LowerEnvelope(HARKinterpolator1D):
2107 """
2108 The lower envelope of a finite set of 1D functions, each of which can be of
2109 any class that has the methods __call__, derivative, and eval_with_derivative.
2110 Generally: it combines HARKinterpolator1Ds.
2112 Parameters
2113 ----------
2114 *functions : function
2115 Any number of real functions; often instances of HARKinterpolator1D
2116 nan_bool : boolean
2117 An indicator for whether the solver should exclude NA's when
2118 forming the lower envelope
2119 """
2121 distance_criteria = ["functions"]
2123 def __init__(self, *functions, nan_bool=True):
2124 if nan_bool:
2125 self.compare = np.nanmin
2126 self.argcompare = np.nanargmin
2127 else:
2128 self.compare = np.min
2129 self.argcompare = np.argmin
2131 self.functions = []
2132 for function in functions:
2133 self.functions.append(function)
2134 self.funcCount = len(self.functions)
2136 def _evaluate(self, x):
2137 """
2138 Returns the level of the function at each value in x as the minimum among
2139 all of the functions. Only called internally by HARKinterpolator1D.__call__.
2140 """
2141 m = len(x)
2142 fx = np.zeros((m, self.funcCount))
2143 for j in range(self.funcCount):
2144 fx[:, j] = self.functions[j](x)
2145 y = self.compare(fx, axis=1)
2146 return y
2148 def _der(self, x):
2149 """
2150 Returns the first derivative of the function at each value in x. Only
2151 called internally by HARKinterpolator1D.derivative.
2152 """
2153 y, dydx = self._evalAndDer(x)
2154 return dydx # Sadly, this is the fastest / most convenient way...
2156 def _evalAndDer(self, x):
2157 """
2158 Returns the level and first derivative of the function at each value in
2159 x. Only called internally by HARKinterpolator1D.eval_and_der.
2160 """
2161 m = len(x)
2162 fx = np.zeros((m, self.funcCount))
2163 for j in range(self.funcCount):
2164 fx[:, j] = self.functions[j](x)
2165 i = self.argcompare(fx, axis=1)
2166 y = fx[np.arange(m), i]
2167 dydx = np.zeros_like(y)
2168 for j in np.unique(i):
2169 c = i == j
2170 dydx[c] = self.functions[j].derivative(x[c])
2171 return y, dydx
2174class UpperEnvelope(HARKinterpolator1D):
2175 """
2176 The upper envelope of a finite set of 1D functions, each of which can be of
2177 any class that has the methods __call__, derivative, and eval_with_derivative.
2178 Generally: it combines HARKinterpolator1Ds.
2180 Parameters
2181 ----------
2182 *functions : function
2183 Any number of real functions; often instances of HARKinterpolator1D
2184 nan_bool : boolean
2185 An indicator for whether the solver should exclude NA's when forming
2186 the lower envelope.
2187 """
2189 distance_criteria = ["functions"]
2191 def __init__(self, *functions, nan_bool=True):
2192 if nan_bool:
2193 self.compare = np.nanmax
2194 self.argcompare = np.nanargmax
2195 else:
2196 self.compare = np.max
2197 self.argcompare = np.argmax
2198 self.functions = []
2199 for function in functions:
2200 self.functions.append(function)
2201 self.funcCount = len(self.functions)
2203 def _evaluate(self, x):
2204 """
2205 Returns the level of the function at each value in x as the maximum among
2206 all of the functions. Only called internally by HARKinterpolator1D.__call__.
2207 """
2208 m = len(x)
2209 fx = np.zeros((m, self.funcCount))
2210 for j in range(self.funcCount):
2211 fx[:, j] = self.functions[j](x)
2212 y = self.compare(fx, axis=1)
2213 return y
2215 def _der(self, x):
2216 """
2217 Returns the first derivative of the function at each value in x. Only
2218 called internally by HARKinterpolator1D.derivative.
2219 """
2220 y, dydx = self._evalAndDer(x)
2221 return dydx # Sadly, this is the fastest / most convenient way...
2223 def _evalAndDer(self, x):
2224 """
2225 Returns the level and first derivative of the function at each value in
2226 x. Only called internally by HARKinterpolator1D.eval_and_der.
2227 """
2228 m = len(x)
2229 fx = np.zeros((m, self.funcCount))
2230 for j in range(self.funcCount):
2231 fx[:, j] = self.functions[j](x)
2232 i = self.argcompare(fx, axis=1)
2233 y = fx[np.arange(m), i]
2234 dydx = np.zeros_like(y)
2235 for j in np.unique(i):
2236 c = i == j
2237 dydx[c] = self.functions[j].derivative(x[c])
2238 return y, dydx
2241class LowerEnvelope2D(HARKinterpolator2D):
2242 """
2243 The lower envelope of a finite set of 2D functions, each of which can be of
2244 any class that has the methods __call__, derivativeX, and derivativeY.
2245 Generally: it combines HARKinterpolator2Ds.
2247 Parameters
2248 ----------
2249 *functions : function
2250 Any number of real functions; often instances of HARKinterpolator2D
2251 nan_bool : boolean
2252 An indicator for whether the solver should exclude NA's when forming
2253 the lower envelope.
2254 """
2256 distance_criteria = ["functions"]
2258 def __init__(self, *functions, nan_bool=True):
2259 if nan_bool:
2260 self.compare = np.nanmin
2261 self.argcompare = np.nanargmin
2262 else:
2263 self.compare = np.min
2264 self.argcompare = np.argmin
2265 self.functions = []
2266 for function in functions:
2267 self.functions.append(function)
2268 self.funcCount = len(self.functions)
2270 def _evaluate(self, x, y):
2271 """
2272 Returns the level of the function at each value in (x,y) as the minimum
2273 among all of the functions. Only called internally by
2274 HARKinterpolator2D.__call__.
2275 """
2276 m = len(x)
2277 temp = np.zeros((m, self.funcCount))
2278 for j in range(self.funcCount):
2279 temp[:, j] = self.functions[j](x, y)
2280 f = self.compare(temp, axis=1)
2281 return f
2283 def _derX(self, x, y):
2284 """
2285 Returns the first derivative of the function with respect to X at each
2286 value in (x,y). Only called internally by HARKinterpolator2D._derX.
2287 """
2288 m = len(x)
2289 temp = np.zeros((m, self.funcCount))
2290 for j in range(self.funcCount):
2291 temp[:, j] = self.functions[j](x, y)
2292 i = self.argcompare(temp, axis=1)
2293 dfdx = np.zeros_like(x)
2294 for j in np.unique(i):
2295 c = i == j
2296 dfdx[c] = self.functions[j].derivativeX(x[c], y[c])
2297 return dfdx
2299 def _derY(self, x, y):
2300 """
2301 Returns the first derivative of the function with respect to Y at each
2302 value in (x,y). Only called internally by HARKinterpolator2D._derY.
2303 """
2304 m = len(x)
2305 temp = np.zeros((m, self.funcCount))
2306 for j in range(self.funcCount):
2307 temp[:, j] = self.functions[j](x, y)
2308 i = self.argcompare(temp, axis=1)
2309 y = temp[np.arange(m), i]
2310 dfdy = np.zeros_like(x)
2311 for j in np.unique(i):
2312 c = i == j
2313 dfdy[c] = self.functions[j].derivativeY(x[c], y[c])
2314 return dfdy
2317class LowerEnvelope3D(HARKinterpolator3D):
2318 """
2319 The lower envelope of a finite set of 3D functions, each of which can be of
2320 any class that has the methods __call__, derivativeX, derivativeY, and
2321 derivativeZ. Generally: it combines HARKinterpolator2Ds.
2323 Parameters
2324 ----------
2325 *functions : function
2326 Any number of real functions; often instances of HARKinterpolator3D
2327 nan_bool : boolean
2328 An indicator for whether the solver should exclude NA's when forming
2329 the lower envelope.
2330 """
2332 distance_criteria = ["functions"]
2334 def __init__(self, *functions, nan_bool=True):
2335 if nan_bool:
2336 self.compare = np.nanmin
2337 self.argcompare = np.nanargmin
2338 else:
2339 self.compare = np.min
2340 self.argcompare = np.argmin
2341 self.functions = []
2342 for function in functions:
2343 self.functions.append(function)
2344 self.funcCount = len(self.functions)
2346 def _evaluate(self, x, y, z):
2347 """
2348 Returns the level of the function at each value in (x,y,z) as the minimum
2349 among all of the functions. Only called internally by
2350 HARKinterpolator3D.__call__.
2351 """
2352 m = len(x)
2353 temp = np.zeros((m, self.funcCount))
2354 for j in range(self.funcCount):
2355 temp[:, j] = self.functions[j](x, y, z)
2356 f = self.compare(temp, axis=1)
2357 return f
2359 def _derX(self, x, y, z):
2360 """
2361 Returns the first derivative of the function with respect to X at each
2362 value in (x,y,z). Only called internally by HARKinterpolator3D._derX.
2363 """
2364 m = len(x)
2365 temp = np.zeros((m, self.funcCount))
2366 for j in range(self.funcCount):
2367 temp[:, j] = self.functions[j](x, y, z)
2368 i = self.argcompare(temp, axis=1)
2369 dfdx = np.zeros_like(x)
2370 for j in np.unique(i):
2371 c = i == j
2372 dfdx[c] = self.functions[j].derivativeX(x[c], y[c], z[c])
2373 return dfdx
2375 def _derY(self, x, y, z):
2376 """
2377 Returns the first derivative of the function with respect to Y at each
2378 value in (x,y,z). Only called internally by HARKinterpolator3D._derY.
2379 """
2380 m = len(x)
2381 temp = np.zeros((m, self.funcCount))
2382 for j in range(self.funcCount):
2383 temp[:, j] = self.functions[j](x, y, z)
2384 i = self.argcompare(temp, axis=1)
2385 y = temp[np.arange(m), i]
2386 dfdy = np.zeros_like(x)
2387 for j in np.unique(i):
2388 c = i == j
2389 dfdy[c] = self.functions[j].derivativeY(x[c], y[c], z[c])
2390 return dfdy
2392 def _derZ(self, x, y, z):
2393 """
2394 Returns the first derivative of the function with respect to Z at each
2395 value in (x,y,z). Only called internally by HARKinterpolator3D._derZ.
2396 """
2397 m = len(x)
2398 temp = np.zeros((m, self.funcCount))
2399 for j in range(self.funcCount):
2400 temp[:, j] = self.functions[j](x, y, z)
2401 i = self.argcompare(temp, axis=1)
2402 y = temp[np.arange(m), i]
2403 dfdz = np.zeros_like(x)
2404 for j in np.unique(i):
2405 c = i == j
2406 dfdz[c] = self.functions[j].derivativeZ(x[c], y[c], z[c])
2407 return dfdz
2410class VariableLowerBoundFunc2D(MetricObject):
2411 """
2412 A class for representing a function with two real inputs whose lower bound
2413 in the first input depends on the second input. Useful for managing curved
2414 natural borrowing constraints, as occurs in the persistent shocks model.
2416 Parameters
2417 ----------
2418 func : function
2419 A function f: (R_+ x R) --> R representing the function of interest
2420 shifted by its lower bound in the first input.
2421 lowerBound : function
2422 The lower bound in the first input of the function of interest, as
2423 a function of the second input.
2424 """
2426 distance_criteria = ["func", "lowerBound"]
2428 def __init__(self, func, lowerBound):
2429 self.func = func
2430 self.lowerBound = lowerBound
2432 def __call__(self, x, y):
2433 """
2434 Evaluate the function at given state space points.
2436 Parameters
2437 ----------
2438 x : np.array
2439 First input values.
2440 y : np.array
2441 Second input values; should be of same shape as x.
2443 Returns
2444 -------
2445 f_out : np.array
2446 Function evaluated at (x,y), of same shape as inputs.
2447 """
2448 xShift = self.lowerBound(y)
2449 f_out = self.func(x - xShift, y)
2450 return f_out
2452 def derivativeX(self, x, y):
2453 """
2454 Evaluate the first derivative with respect to x of the function at given
2455 state space points.
2457 Parameters
2458 ----------
2459 x : np.array
2460 First input values.
2461 y : np.array
2462 Second input values; should be of same shape as x.
2464 Returns
2465 -------
2466 dfdx_out : np.array
2467 First derivative of function with respect to the first input,
2468 evaluated at (x,y), of same shape as inputs.
2469 """
2470 xShift = self.lowerBound(y)
2471 dfdx_out = self.func.derivativeX(x - xShift, y)
2472 return dfdx_out
2474 def derivativeY(self, x, y):
2475 """
2476 Evaluate the first derivative with respect to y of the function at given
2477 state space points.
2479 Parameters
2480 ----------
2481 x : np.array
2482 First input values.
2483 y : np.array
2484 Second input values; should be of same shape as x.
2486 Returns
2487 -------
2488 dfdy_out : np.array
2489 First derivative of function with respect to the second input,
2490 evaluated at (x,y), of same shape as inputs.
2491 """
2492 xShift, xShiftDer = self.lowerBound.eval_with_derivative(y)
2493 dfdy_out = self.func.derivativeY(
2494 x - xShift, y
2495 ) - xShiftDer * self.func.derivativeX(x - xShift, y)
2496 return dfdy_out
2499class VariableLowerBoundFunc3D(MetricObject):
2500 """
2501 A class for representing a function with three real inputs whose lower bound
2502 in the first input depends on the second input. Useful for managing curved
2503 natural borrowing constraints.
2505 Parameters
2506 ----------
2507 func : function
2508 A function f: (R_+ x R^2) --> R representing the function of interest
2509 shifted by its lower bound in the first input.
2510 lowerBound : function
2511 The lower bound in the first input of the function of interest, as
2512 a function of the second input.
2513 """
2515 distance_criteria = ["func", "lowerBound"]
2517 def __init__(self, func, lowerBound):
2518 self.func = func
2519 self.lowerBound = lowerBound
2521 def __call__(self, x, y, z):
2522 """
2523 Evaluate the function at given state space points.
2525 Parameters
2526 ----------
2527 x : np.array
2528 First input values.
2529 y : np.array
2530 Second input values; should be of same shape as x.
2531 z : np.array
2532 Third input values; should be of same shape as x.
2534 Returns
2535 -------
2536 f_out : np.array
2537 Function evaluated at (x,y,z), of same shape as inputs.
2538 """
2539 xShift = self.lowerBound(y)
2540 f_out = self.func(x - xShift, y, z)
2541 return f_out
2543 def derivativeX(self, x, y, z):
2544 """
2545 Evaluate the first derivative with respect to x of the function at given
2546 state space points.
2548 Parameters
2549 ----------
2550 x : np.array
2551 First input values.
2552 y : np.array
2553 Second input values; should be of same shape as x.
2554 z : np.array
2555 Third input values; should be of same shape as x.
2557 Returns
2558 -------
2559 dfdx_out : np.array
2560 First derivative of function with respect to the first input,
2561 evaluated at (x,y,z), of same shape as inputs.
2562 """
2563 xShift = self.lowerBound(y)
2564 dfdx_out = self.func.derivativeX(x - xShift, y, z)
2565 return dfdx_out
2567 def derivativeY(self, x, y, z):
2568 """
2569 Evaluate the first derivative with respect to y of the function at given
2570 state space points.
2572 Parameters
2573 ----------
2574 x : np.array
2575 First input values.
2576 y : np.array
2577 Second input values; should be of same shape as x.
2578 z : np.array
2579 Third input values; should be of same shape as x.
2581 Returns
2582 -------
2583 dfdy_out : np.array
2584 First derivative of function with respect to the second input,
2585 evaluated at (x,y,z), of same shape as inputs.
2586 """
2587 xShift, xShiftDer = self.lowerBound.eval_with_derivative(y)
2588 dfdy_out = self.func.derivativeY(
2589 x - xShift, y, z
2590 ) - xShiftDer * self.func.derivativeX(x - xShift, y, z)
2591 return dfdy_out
2593 def derivativeZ(self, x, y, z):
2594 """
2595 Evaluate the first derivative with respect to z of the function at given
2596 state space points.
2598 Parameters
2599 ----------
2600 x : np.array
2601 First input values.
2602 y : np.array
2603 Second input values; should be of same shape as x.
2604 z : np.array
2605 Third input values; should be of same shape as x.
2607 Returns
2608 -------
2609 dfdz_out : np.array
2610 First derivative of function with respect to the third input,
2611 evaluated at (x,y,z), of same shape as inputs.
2612 """
2613 xShift = self.lowerBound(y)
2614 dfdz_out = self.func.derivativeZ(x - xShift, y, z)
2615 return dfdz_out
2618class LinearInterpOnInterp1D(HARKinterpolator2D):
2619 """
2620 A 2D interpolator that linearly interpolates among a list of 1D interpolators.
2622 Parameters
2623 ----------
2624 xInterpolators : [HARKinterpolator1D]
2625 A list of 1D interpolations over the x variable. The nth element of
2626 xInterpolators represents f(x,y_values[n]).
2627 y_values: numpy.array
2628 An array of y values equal in length to xInterpolators.
2629 """
2631 distance_criteria = ["xInterpolators", "y_list"]
2633 def __init__(self, xInterpolators, y_values):
2634 self.xInterpolators = xInterpolators
2635 self.y_list = y_values
2636 self.y_n = y_values.size
2638 def _evaluate(self, x, y):
2639 """
2640 Returns the level of the interpolated function at each value in x,y.
2641 Only called internally by HARKinterpolator2D.__call__ (etc).
2642 """
2643 m = len(x)
2644 y_pos = np.searchsorted(self.y_list, y)
2645 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
2646 y_pos[y_pos < 1] = 1
2647 f = np.zeros(m) + np.nan
2648 if y.size > 0:
2649 for i in range(1, self.y_n):
2650 c = y_pos == i
2651 if np.any(c):
2652 alpha = (y[c] - self.y_list[i - 1]) / (
2653 self.y_list[i] - self.y_list[i - 1]
2654 )
2655 f[c] = (1 - alpha) * self.xInterpolators[i - 1](
2656 x[c]
2657 ) + alpha * self.xInterpolators[i](x[c])
2658 return f
2660 def _derX(self, x, y):
2661 """
2662 Returns the derivative with respect to x of the interpolated function
2663 at each value in x,y. Only called internally by HARKinterpolator2D.derivativeX.
2664 """
2665 m = len(x)
2666 y_pos = np.searchsorted(self.y_list, y)
2667 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
2668 y_pos[y_pos < 1] = 1
2669 dfdx = np.zeros(m) + np.nan
2670 if y.size > 0:
2671 for i in range(1, self.y_n):
2672 c = y_pos == i
2673 if np.any(c):
2674 alpha = (y[c] - self.y_list[i - 1]) / (
2675 self.y_list[i] - self.y_list[i - 1]
2676 )
2677 dfdx[c] = (1 - alpha) * self.xInterpolators[i - 1]._der(
2678 x[c]
2679 ) + alpha * self.xInterpolators[i]._der(x[c])
2680 return dfdx
2682 def _derY(self, x, y):
2683 """
2684 Returns the derivative with respect to y of the interpolated function
2685 at each value in x,y. Only called internally by HARKinterpolator2D.derivativeY.
2686 """
2687 m = len(x)
2688 y_pos = np.searchsorted(self.y_list, y)
2689 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
2690 y_pos[y_pos < 1] = 1
2691 dfdy = np.zeros(m) + np.nan
2692 if y.size > 0:
2693 for i in range(1, self.y_n):
2694 c = y_pos == i
2695 if np.any(c):
2696 dfdy[c] = (
2697 self.xInterpolators[i](x[c]) - self.xInterpolators[i - 1](x[c])
2698 ) / (self.y_list[i] - self.y_list[i - 1])
2699 return dfdy
2702class BilinearInterpOnInterp1D(HARKinterpolator3D):
2703 """
2704 A 3D interpolator that bilinearly interpolates among a list of lists of 1D
2705 interpolators.
2707 Constructor for the class, generating an approximation to a function of
2708 the form f(x,y,z) using interpolations over f(x,y_0,z_0) for a fixed grid
2709 of y_0 and z_0 values.
2711 Parameters
2712 ----------
2713 xInterpolators : [[HARKinterpolator1D]]
2714 A list of lists of 1D interpolations over the x variable. The i,j-th
2715 element of xInterpolators represents f(x,y_values[i],z_values[j]).
2716 y_values: numpy.array
2717 An array of y values equal in length to xInterpolators.
2718 z_values: numpy.array
2719 An array of z values equal in length to xInterpolators[0].
2720 """
2722 distance_criteria = ["xInterpolators", "y_list", "z_list"]
2724 def __init__(self, xInterpolators, y_values, z_values):
2725 self.xInterpolators = xInterpolators
2726 self.y_list = y_values
2727 self.y_n = y_values.size
2728 self.z_list = z_values
2729 self.z_n = z_values.size
2731 def _evaluate(self, x, y, z):
2732 """
2733 Returns the level of the interpolated function at each value in x,y,z.
2734 Only called internally by HARKinterpolator3D.__call__ (etc).
2735 """
2736 m = len(x)
2737 y_pos = np.searchsorted(self.y_list, y)
2738 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
2739 y_pos[y_pos < 1] = 1
2740 z_pos = np.searchsorted(self.z_list, z)
2741 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
2742 z_pos[z_pos < 1] = 1
2743 f = np.zeros(m) + np.nan
2744 for i in range(1, self.y_n):
2745 for j in range(1, self.z_n):
2746 c = np.logical_and(i == y_pos, j == z_pos)
2747 if np.any(c):
2748 alpha = (y[c] - self.y_list[i - 1]) / (
2749 self.y_list[i] - self.y_list[i - 1]
2750 )
2751 beta = (z[c] - self.z_list[j - 1]) / (
2752 self.z_list[j] - self.z_list[j - 1]
2753 )
2754 f[c] = (
2755 (1 - alpha)
2756 * (1 - beta)
2757 * self.xInterpolators[i - 1][j - 1](x[c])
2758 + (1 - alpha) * beta * self.xInterpolators[i - 1][j](x[c])
2759 + alpha * (1 - beta) * self.xInterpolators[i][j - 1](x[c])
2760 + alpha * beta * self.xInterpolators[i][j](x[c])
2761 )
2762 return f
2764 def _derX(self, x, y, z):
2765 """
2766 Returns the derivative with respect to x of the interpolated function
2767 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeX.
2768 """
2769 m = len(x)
2770 y_pos = np.searchsorted(self.y_list, y)
2771 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
2772 y_pos[y_pos < 1] = 1
2773 z_pos = np.searchsorted(self.z_list, z)
2774 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
2775 z_pos[z_pos < 1] = 1
2776 dfdx = np.zeros(m) + np.nan
2777 for i in range(1, self.y_n):
2778 for j in range(1, self.z_n):
2779 c = np.logical_and(i == y_pos, j == z_pos)
2780 if np.any(c):
2781 alpha = (y[c] - self.y_list[i - 1]) / (
2782 self.y_list[i] - self.y_list[i - 1]
2783 )
2784 beta = (z[c] - self.z_list[j - 1]) / (
2785 self.z_list[j] - self.z_list[j - 1]
2786 )
2787 dfdx[c] = (
2788 (1 - alpha)
2789 * (1 - beta)
2790 * self.xInterpolators[i - 1][j - 1]._der(x[c])
2791 + (1 - alpha) * beta * self.xInterpolators[i - 1][j]._der(x[c])
2792 + alpha * (1 - beta) * self.xInterpolators[i][j - 1]._der(x[c])
2793 + alpha * beta * self.xInterpolators[i][j]._der(x[c])
2794 )
2795 return dfdx
2797 def _derY(self, x, y, z):
2798 """
2799 Returns the derivative with respect to y of the interpolated function
2800 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeY.
2801 """
2802 m = len(x)
2803 y_pos = np.searchsorted(self.y_list, y)
2804 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
2805 y_pos[y_pos < 1] = 1
2806 z_pos = np.searchsorted(self.z_list, z)
2807 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
2808 z_pos[z_pos < 1] = 1
2809 dfdy = np.zeros(m) + np.nan
2810 for i in range(1, self.y_n):
2811 for j in range(1, self.z_n):
2812 c = np.logical_and(i == y_pos, j == z_pos)
2813 if np.any(c):
2814 beta = (z[c] - self.z_list[j - 1]) / (
2815 self.z_list[j] - self.z_list[j - 1]
2816 )
2817 dfdy[c] = (
2818 (
2819 (1 - beta) * self.xInterpolators[i][j - 1](x[c])
2820 + beta * self.xInterpolators[i][j](x[c])
2821 )
2822 - (
2823 (1 - beta) * self.xInterpolators[i - 1][j - 1](x[c])
2824 + beta * self.xInterpolators[i - 1][j](x[c])
2825 )
2826 ) / (self.y_list[i] - self.y_list[i - 1])
2827 return dfdy
2829 def _derZ(self, x, y, z):
2830 """
2831 Returns the derivative with respect to z of the interpolated function
2832 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeZ.
2833 """
2834 m = len(x)
2835 y_pos = np.searchsorted(self.y_list, y)
2836 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
2837 y_pos[y_pos < 1] = 1
2838 z_pos = np.searchsorted(self.z_list, z)
2839 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
2840 z_pos[z_pos < 1] = 1
2841 dfdz = np.zeros(m) + np.nan
2842 for i in range(1, self.y_n):
2843 for j in range(1, self.z_n):
2844 c = np.logical_and(i == y_pos, j == z_pos)
2845 if np.any(c):
2846 alpha = (y[c] - self.y_list[i - 1]) / (
2847 self.y_list[i] - self.y_list[i - 1]
2848 )
2849 dfdz[c] = (
2850 (
2851 (1 - alpha) * self.xInterpolators[i - 1][j](x[c])
2852 + alpha * self.xInterpolators[i][j](x[c])
2853 )
2854 - (
2855 (1 - alpha) * self.xInterpolators[i - 1][j - 1](x[c])
2856 + alpha * self.xInterpolators[i][j - 1](x[c])
2857 )
2858 ) / (self.z_list[j] - self.z_list[j - 1])
2859 return dfdz
2862class TrilinearInterpOnInterp1D(HARKinterpolator4D):
2863 """
2864 A 4D interpolator that trilinearly interpolates among a list of lists of 1D interpolators.
2866 Constructor for the class, generating an approximation to a function of
2867 the form f(w,x,y,z) using interpolations over f(w,x_0,y_0,z_0) for a fixed
2868 grid of y_0 and z_0 values.
2870 Parameters
2871 ----------
2872 wInterpolators : [[[HARKinterpolator1D]]]
2873 A list of lists of lists of 1D interpolations over the x variable.
2874 The i,j,k-th element of wInterpolators represents f(w,x_values[i],y_values[j],z_values[k]).
2875 x_values: numpy.array
2876 An array of x values equal in length to wInterpolators.
2877 y_values: numpy.array
2878 An array of y values equal in length to wInterpolators[0].
2879 z_values: numpy.array
2880 An array of z values equal in length to wInterpolators[0][0]
2881 """
2883 distance_criteria = ["wInterpolators", "x_list", "y_list", "z_list"]
2885 def __init__(self, wInterpolators, x_values, y_values, z_values):
2886 self.wInterpolators = wInterpolators
2887 self.x_list = x_values
2888 self.x_n = x_values.size
2889 self.y_list = y_values
2890 self.y_n = y_values.size
2891 self.z_list = z_values
2892 self.z_n = z_values.size
2894 def _evaluate(self, w, x, y, z):
2895 """
2896 Returns the level of the interpolated function at each value in w,x,y,z.
2897 Only called internally by HARKinterpolator4D.__call__ (etc).
2898 """
2899 m = len(x)
2900 x_pos = np.searchsorted(self.x_list, x)
2901 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
2902 y_pos = np.searchsorted(self.y_list, y)
2903 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
2904 y_pos[y_pos < 1] = 1
2905 z_pos = np.searchsorted(self.z_list, z)
2906 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
2907 z_pos[z_pos < 1] = 1
2908 f = np.zeros(m) + np.nan
2909 for i in range(1, self.x_n):
2910 for j in range(1, self.y_n):
2911 for k in range(1, self.z_n):
2912 c = np.logical_and(
2913 np.logical_and(i == x_pos, j == y_pos), k == z_pos
2914 )
2915 if np.any(c):
2916 alpha = (x[c] - self.x_list[i - 1]) / (
2917 self.x_list[i] - self.x_list[i - 1]
2918 )
2919 beta = (y[c] - self.y_list[j - 1]) / (
2920 self.y_list[j] - self.y_list[j - 1]
2921 )
2922 gamma = (z[c] - self.z_list[k - 1]) / (
2923 self.z_list[k] - self.z_list[k - 1]
2924 )
2925 f[c] = (
2926 (1 - alpha)
2927 * (1 - beta)
2928 * (1 - gamma)
2929 * self.wInterpolators[i - 1][j - 1][k - 1](w[c])
2930 + (1 - alpha)
2931 * (1 - beta)
2932 * gamma
2933 * self.wInterpolators[i - 1][j - 1][k](w[c])
2934 + (1 - alpha)
2935 * beta
2936 * (1 - gamma)
2937 * self.wInterpolators[i - 1][j][k - 1](w[c])
2938 + (1 - alpha)
2939 * beta
2940 * gamma
2941 * self.wInterpolators[i - 1][j][k](w[c])
2942 + alpha
2943 * (1 - beta)
2944 * (1 - gamma)
2945 * self.wInterpolators[i][j - 1][k - 1](w[c])
2946 + alpha
2947 * (1 - beta)
2948 * gamma
2949 * self.wInterpolators[i][j - 1][k](w[c])
2950 + alpha
2951 * beta
2952 * (1 - gamma)
2953 * self.wInterpolators[i][j][k - 1](w[c])
2954 + alpha * beta * gamma * self.wInterpolators[i][j][k](w[c])
2955 )
2956 return f
2958 def _derW(self, w, x, y, z):
2959 """
2960 Returns the derivative with respect to w of the interpolated function
2961 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeW.
2962 """
2963 m = len(x)
2964 x_pos = np.searchsorted(self.x_list, x)
2965 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
2966 y_pos = np.searchsorted(self.y_list, y)
2967 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
2968 y_pos[y_pos < 1] = 1
2969 z_pos = np.searchsorted(self.z_list, z)
2970 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
2971 z_pos[z_pos < 1] = 1
2972 dfdw = np.zeros(m) + np.nan
2973 for i in range(1, self.x_n):
2974 for j in range(1, self.y_n):
2975 for k in range(1, self.z_n):
2976 c = np.logical_and(
2977 np.logical_and(i == x_pos, j == y_pos), k == z_pos
2978 )
2979 if np.any(c):
2980 alpha = (x[c] - self.x_list[i - 1]) / (
2981 self.x_list[i] - self.x_list[i - 1]
2982 )
2983 beta = (y[c] - self.y_list[j - 1]) / (
2984 self.y_list[j] - self.y_list[j - 1]
2985 )
2986 gamma = (z[c] - self.z_list[k - 1]) / (
2987 self.z_list[k] - self.z_list[k - 1]
2988 )
2989 dfdw[c] = (
2990 (1 - alpha)
2991 * (1 - beta)
2992 * (1 - gamma)
2993 * self.wInterpolators[i - 1][j - 1][k - 1]._der(w[c])
2994 + (1 - alpha)
2995 * (1 - beta)
2996 * gamma
2997 * self.wInterpolators[i - 1][j - 1][k]._der(w[c])
2998 + (1 - alpha)
2999 * beta
3000 * (1 - gamma)
3001 * self.wInterpolators[i - 1][j][k - 1]._der(w[c])
3002 + (1 - alpha)
3003 * beta
3004 * gamma
3005 * self.wInterpolators[i - 1][j][k]._der(w[c])
3006 + alpha
3007 * (1 - beta)
3008 * (1 - gamma)
3009 * self.wInterpolators[i][j - 1][k - 1]._der(w[c])
3010 + alpha
3011 * (1 - beta)
3012 * gamma
3013 * self.wInterpolators[i][j - 1][k]._der(w[c])
3014 + alpha
3015 * beta
3016 * (1 - gamma)
3017 * self.wInterpolators[i][j][k - 1]._der(w[c])
3018 + alpha
3019 * beta
3020 * gamma
3021 * self.wInterpolators[i][j][k]._der(w[c])
3022 )
3023 return dfdw
3025 def _derX(self, w, x, y, z):
3026 """
3027 Returns the derivative with respect to x of the interpolated function
3028 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeX.
3029 """
3030 m = len(x)
3031 x_pos = np.searchsorted(self.x_list, x)
3032 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
3033 y_pos = np.searchsorted(self.y_list, y)
3034 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
3035 y_pos[y_pos < 1] = 1
3036 z_pos = np.searchsorted(self.z_list, z)
3037 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
3038 z_pos[z_pos < 1] = 1
3039 dfdx = np.zeros(m) + np.nan
3040 for i in range(1, self.x_n):
3041 for j in range(1, self.y_n):
3042 for k in range(1, self.z_n):
3043 c = np.logical_and(
3044 np.logical_and(i == x_pos, j == y_pos), k == z_pos
3045 )
3046 if np.any(c):
3047 beta = (y[c] - self.y_list[j - 1]) / (
3048 self.y_list[j] - self.y_list[j - 1]
3049 )
3050 gamma = (z[c] - self.z_list[k - 1]) / (
3051 self.z_list[k] - self.z_list[k - 1]
3052 )
3053 dfdx[c] = (
3054 (
3055 (1 - beta)
3056 * (1 - gamma)
3057 * self.wInterpolators[i][j - 1][k - 1](w[c])
3058 + (1 - beta)
3059 * gamma
3060 * self.wInterpolators[i][j - 1][k](w[c])
3061 + beta
3062 * (1 - gamma)
3063 * self.wInterpolators[i][j][k - 1](w[c])
3064 + beta * gamma * self.wInterpolators[i][j][k](w[c])
3065 )
3066 - (
3067 (1 - beta)
3068 * (1 - gamma)
3069 * self.wInterpolators[i - 1][j - 1][k - 1](w[c])
3070 + (1 - beta)
3071 * gamma
3072 * self.wInterpolators[i - 1][j - 1][k](w[c])
3073 + beta
3074 * (1 - gamma)
3075 * self.wInterpolators[i - 1][j][k - 1](w[c])
3076 + beta * gamma * self.wInterpolators[i - 1][j][k](w[c])
3077 )
3078 ) / (self.x_list[i] - self.x_list[i - 1])
3079 return dfdx
3081 def _derY(self, w, x, y, z):
3082 """
3083 Returns the derivative with respect to y of the interpolated function
3084 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeY.
3085 """
3086 m = len(x)
3087 x_pos = np.searchsorted(self.x_list, x)
3088 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
3089 y_pos = np.searchsorted(self.y_list, y)
3090 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
3091 y_pos[y_pos < 1] = 1
3092 z_pos = np.searchsorted(self.z_list, z)
3093 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
3094 z_pos[z_pos < 1] = 1
3095 dfdy = np.zeros(m) + np.nan
3096 for i in range(1, self.x_n):
3097 for j in range(1, self.y_n):
3098 for k in range(1, self.z_n):
3099 c = np.logical_and(
3100 np.logical_and(i == x_pos, j == y_pos), k == z_pos
3101 )
3102 if np.any(c):
3103 alpha = (x[c] - self.x_list[i - 1]) / (
3104 self.x_list[i] - self.x_list[i - 1]
3105 )
3106 gamma = (z[c] - self.z_list[k - 1]) / (
3107 self.z_list[k] - self.z_list[k - 1]
3108 )
3109 dfdy[c] = (
3110 (
3111 (1 - alpha)
3112 * (1 - gamma)
3113 * self.wInterpolators[i - 1][j][k - 1](w[c])
3114 + (1 - alpha)
3115 * gamma
3116 * self.wInterpolators[i - 1][j][k](w[c])
3117 + alpha
3118 * (1 - gamma)
3119 * self.wInterpolators[i][j][k - 1](w[c])
3120 + alpha * gamma * self.wInterpolators[i][j][k](w[c])
3121 )
3122 - (
3123 (1 - alpha)
3124 * (1 - gamma)
3125 * self.wInterpolators[i - 1][j - 1][k - 1](w[c])
3126 + (1 - alpha)
3127 * gamma
3128 * self.wInterpolators[i - 1][j - 1][k](w[c])
3129 + alpha
3130 * (1 - gamma)
3131 * self.wInterpolators[i][j - 1][k - 1](w[c])
3132 + alpha * gamma * self.wInterpolators[i][j - 1][k](w[c])
3133 )
3134 ) / (self.y_list[j] - self.y_list[j - 1])
3135 return dfdy
3137 def _derZ(self, w, x, y, z):
3138 """
3139 Returns the derivative with respect to z of the interpolated function
3140 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeZ.
3141 """
3142 m = len(x)
3143 x_pos = np.searchsorted(self.x_list, x)
3144 x_pos[x_pos > self.x_n - 1] = self.x_n - 1
3145 y_pos = np.searchsorted(self.y_list, y)
3146 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
3147 y_pos[y_pos < 1] = 1
3148 z_pos = np.searchsorted(self.z_list, z)
3149 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
3150 z_pos[z_pos < 1] = 1
3151 dfdz = np.zeros(m) + np.nan
3152 for i in range(1, self.x_n):
3153 for j in range(1, self.y_n):
3154 for k in range(1, self.z_n):
3155 c = np.logical_and(
3156 np.logical_and(i == x_pos, j == y_pos), k == z_pos
3157 )
3158 if np.any(c):
3159 alpha = (x[c] - self.x_list[i - 1]) / (
3160 self.x_list[i] - self.x_list[i - 1]
3161 )
3162 beta = (y[c] - self.y_list[j - 1]) / (
3163 self.y_list[j] - self.y_list[j - 1]
3164 )
3165 dfdz[c] = (
3166 (
3167 (1 - alpha)
3168 * (1 - beta)
3169 * self.wInterpolators[i - 1][j - 1][k](w[c])
3170 + (1 - alpha)
3171 * beta
3172 * self.wInterpolators[i - 1][j][k](w[c])
3173 + alpha
3174 * (1 - beta)
3175 * self.wInterpolators[i][j - 1][k](w[c])
3176 + alpha * beta * self.wInterpolators[i][j][k](w[c])
3177 )
3178 - (
3179 (1 - alpha)
3180 * (1 - beta)
3181 * self.wInterpolators[i - 1][j - 1][k - 1](w[c])
3182 + (1 - alpha)
3183 * beta
3184 * self.wInterpolators[i - 1][j][k - 1](w[c])
3185 + alpha
3186 * (1 - beta)
3187 * self.wInterpolators[i][j - 1][k - 1](w[c])
3188 + alpha * beta * self.wInterpolators[i][j][k - 1](w[c])
3189 )
3190 ) / (self.z_list[k] - self.z_list[k - 1])
3191 return dfdz
3194class LinearInterpOnInterp2D(HARKinterpolator3D):
3195 """
3196 A 3D interpolation method that linearly interpolates between "layers" of
3197 arbitrary 2D interpolations. Useful for models with two endogenous state
3198 variables and one exogenous state variable when solving with the endogenous
3199 grid method. NOTE: should not be used if an exogenous 3D grid is used, will
3200 be significantly slower than TrilinearInterp.
3202 Constructor for the class, generating an approximation to a function of
3203 the form f(x,y,z) using interpolations over f(x,y,z_0) for a fixed grid
3204 of z_0 values.
3206 Parameters
3207 ----------
3208 xyInterpolators : [HARKinterpolator2D]
3209 A list of 2D interpolations over the x and y variables. The nth
3210 element of xyInterpolators represents f(x,y,z_values[n]).
3211 z_values: numpy.array
3212 An array of z values equal in length to xyInterpolators.
3213 """
3215 distance_criteria = ["xyInterpolators", "z_list"]
3217 def __init__(self, xyInterpolators, z_values):
3218 self.xyInterpolators = xyInterpolators
3219 self.z_list = z_values
3220 self.z_n = z_values.size
3222 def _evaluate(self, x, y, z):
3223 """
3224 Returns the level of the interpolated function at each value in x,y,z.
3225 Only called internally by HARKinterpolator3D.__call__ (etc).
3226 """
3227 m = len(x)
3228 z_pos = np.searchsorted(self.z_list, z)
3229 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
3230 z_pos[z_pos < 1] = 1
3231 f = np.zeros(m) + np.nan
3232 if x.size > 0:
3233 for i in range(1, self.z_n):
3234 c = z_pos == i
3235 if np.any(c):
3236 alpha = (z[c] - self.z_list[i - 1]) / (
3237 self.z_list[i] - self.z_list[i - 1]
3238 )
3239 f[c] = (1 - alpha) * self.xyInterpolators[i - 1](
3240 x[c], y[c]
3241 ) + alpha * self.xyInterpolators[i](x[c], y[c])
3242 return f
3244 def _derX(self, x, y, z):
3245 """
3246 Returns the derivative with respect to x of the interpolated function
3247 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeX.
3248 """
3249 m = len(x)
3250 z_pos = np.searchsorted(self.z_list, z)
3251 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
3252 z_pos[z_pos < 1] = 1
3253 dfdx = np.zeros(m) + np.nan
3254 if x.size > 0:
3255 for i in range(1, self.z_n):
3256 c = z_pos == i
3257 if np.any(c):
3258 alpha = (z[c] - self.z_list[i - 1]) / (
3259 self.z_list[i] - self.z_list[i - 1]
3260 )
3261 dfdx[c] = (1 - alpha) * self.xyInterpolators[i - 1].derivativeX(
3262 x[c], y[c]
3263 ) + alpha * self.xyInterpolators[i].derivativeX(x[c], y[c])
3264 return dfdx
3266 def _derY(self, x, y, z):
3267 """
3268 Returns the derivative with respect to y of the interpolated function
3269 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeY.
3270 """
3271 m = len(x)
3272 z_pos = np.searchsorted(self.z_list, z)
3273 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
3274 z_pos[z_pos < 1] = 1
3275 dfdy = np.zeros(m) + np.nan
3276 if x.size > 0:
3277 for i in range(1, self.z_n):
3278 c = z_pos == i
3279 if np.any(c):
3280 alpha = (z[c] - self.z_list[i - 1]) / (
3281 self.z_list[i] - self.z_list[i - 1]
3282 )
3283 dfdy[c] = (1 - alpha) * self.xyInterpolators[i - 1].derivativeY(
3284 x[c], y[c]
3285 ) + alpha * self.xyInterpolators[i].derivativeY(x[c], y[c])
3286 return dfdy
3288 def _derZ(self, x, y, z):
3289 """
3290 Returns the derivative with respect to z of the interpolated function
3291 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeZ.
3292 """
3293 m = len(x)
3294 z_pos = np.searchsorted(self.z_list, z)
3295 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
3296 z_pos[z_pos < 1] = 1
3297 dfdz = np.zeros(m) + np.nan
3298 if x.size > 0:
3299 for i in range(1, self.z_n):
3300 c = z_pos == i
3301 if np.any(c):
3302 dfdz[c] = (
3303 self.xyInterpolators[i](x[c], y[c])
3304 - self.xyInterpolators[i - 1](x[c], y[c])
3305 ) / (self.z_list[i] - self.z_list[i - 1])
3306 return dfdz
3309class BilinearInterpOnInterp2D(HARKinterpolator4D):
3310 """
3311 A 4D interpolation method that bilinearly interpolates among "layers" of
3312 arbitrary 2D interpolations. Useful for models with two endogenous state
3313 variables and two exogenous state variables when solving with the endogenous
3314 grid method. NOTE: should not be used if an exogenous 4D grid is used, will
3315 be significantly slower than QuadlinearInterp.
3317 Constructor for the class, generating an approximation to a function of
3318 the form f(w,x,y,z) using interpolations over f(w,x,y_0,z_0) for a fixed
3319 grid of y_0 and z_0 values.
3321 Parameters
3322 ----------
3323 wxInterpolators : [[HARKinterpolator2D]]
3324 A list of lists of 2D interpolations over the w and x variables.
3325 The i,j-th element of wxInterpolators represents
3326 f(w,x,y_values[i],z_values[j]).
3327 y_values: numpy.array
3328 An array of y values equal in length to wxInterpolators.
3329 z_values: numpy.array
3330 An array of z values equal in length to wxInterpolators[0].
3331 """
3333 distance_criteria = ["wxInterpolators", "y_list", "z_list"]
3335 def __init__(self, wxInterpolators, y_values, z_values):
3336 self.wxInterpolators = wxInterpolators
3337 self.y_list = y_values
3338 self.y_n = y_values.size
3339 self.z_list = z_values
3340 self.z_n = z_values.size
3342 def _evaluate(self, w, x, y, z):
3343 """
3344 Returns the level of the interpolated function at each value in x,y,z.
3345 Only called internally by HARKinterpolator4D.__call__ (etc).
3346 """
3347 m = len(x)
3348 y_pos = np.searchsorted(self.y_list, y)
3349 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
3350 y_pos[y_pos < 1] = 1
3351 z_pos = np.searchsorted(self.z_list, z)
3352 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
3353 z_pos[z_pos < 1] = 1
3354 f = np.zeros(m) + np.nan
3355 for i in range(1, self.y_n):
3356 for j in range(1, self.z_n):
3357 c = np.logical_and(i == y_pos, j == z_pos)
3358 if np.any(c):
3359 alpha = (y[c] - self.y_list[i - 1]) / (
3360 self.y_list[i] - self.y_list[i - 1]
3361 )
3362 beta = (z[c] - self.z_list[j - 1]) / (
3363 self.z_list[j] - self.z_list[j - 1]
3364 )
3365 f[c] = (
3366 (1 - alpha)
3367 * (1 - beta)
3368 * self.wxInterpolators[i - 1][j - 1](w[c], x[c])
3369 + (1 - alpha)
3370 * beta
3371 * self.wxInterpolators[i - 1][j](w[c], x[c])
3372 + alpha
3373 * (1 - beta)
3374 * self.wxInterpolators[i][j - 1](w[c], x[c])
3375 + alpha * beta * self.wxInterpolators[i][j](w[c], x[c])
3376 )
3377 return f
3379 def _derW(self, w, x, y, z):
3380 """
3381 Returns the derivative with respect to w of the interpolated function
3382 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeW.
3383 """
3384 # This may look strange, as we call the derivativeX() method to get the
3385 # derivative with respect to w, but that's just a quirk of 4D interpolations
3386 # beginning with w rather than x. The derivative wrt the first dimension
3387 # of an element of wxInterpolators is the w-derivative of the main function.
3388 m = len(x)
3389 y_pos = np.searchsorted(self.y_list, y)
3390 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
3391 y_pos[y_pos < 1] = 1
3392 z_pos = np.searchsorted(self.z_list, z)
3393 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
3394 z_pos[z_pos < 1] = 1
3395 dfdw = np.zeros(m) + np.nan
3396 for i in range(1, self.y_n):
3397 for j in range(1, self.z_n):
3398 c = np.logical_and(i == y_pos, j == z_pos)
3399 if np.any(c):
3400 alpha = (y[c] - self.y_list[i - 1]) / (
3401 self.y_list[i] - self.y_list[i - 1]
3402 )
3403 beta = (z[c] - self.z_list[j - 1]) / (
3404 self.z_list[j] - self.z_list[j - 1]
3405 )
3406 dfdw[c] = (
3407 (1 - alpha)
3408 * (1 - beta)
3409 * self.wxInterpolators[i - 1][j - 1].derivativeX(w[c], x[c])
3410 + (1 - alpha)
3411 * beta
3412 * self.wxInterpolators[i - 1][j].derivativeX(w[c], x[c])
3413 + alpha
3414 * (1 - beta)
3415 * self.wxInterpolators[i][j - 1].derivativeX(w[c], x[c])
3416 + alpha
3417 * beta
3418 * self.wxInterpolators[i][j].derivativeX(w[c], x[c])
3419 )
3420 return dfdw
3422 def _derX(self, w, x, y, z):
3423 """
3424 Returns the derivative with respect to x of the interpolated function
3425 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeX.
3426 """
3427 # This may look strange, as we call the derivativeY() method to get the
3428 # derivative with respect to x, but that's just a quirk of 4D interpolations
3429 # beginning with w rather than x. The derivative wrt the second dimension
3430 # of an element of wxInterpolators is the x-derivative of the main function.
3431 m = len(x)
3432 y_pos = np.searchsorted(self.y_list, y)
3433 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
3434 y_pos[y_pos < 1] = 1
3435 z_pos = np.searchsorted(self.z_list, z)
3436 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
3437 z_pos[z_pos < 1] = 1
3438 dfdx = np.zeros(m) + np.nan
3439 for i in range(1, self.y_n):
3440 for j in range(1, self.z_n):
3441 c = np.logical_and(i == y_pos, j == z_pos)
3442 if np.any(c):
3443 alpha = (y[c] - self.y_list[i - 1]) / (
3444 self.y_list[i] - self.y_list[i - 1]
3445 )
3446 beta = (z[c] - self.z_list[j - 1]) / (
3447 self.z_list[j] - self.z_list[j - 1]
3448 )
3449 dfdx[c] = (
3450 (1 - alpha)
3451 * (1 - beta)
3452 * self.wxInterpolators[i - 1][j - 1].derivativeY(w[c], x[c])
3453 + (1 - alpha)
3454 * beta
3455 * self.wxInterpolators[i - 1][j].derivativeY(w[c], x[c])
3456 + alpha
3457 * (1 - beta)
3458 * self.wxInterpolators[i][j - 1].derivativeY(w[c], x[c])
3459 + alpha
3460 * beta
3461 * self.wxInterpolators[i][j].derivativeY(w[c], x[c])
3462 )
3463 return dfdx
3465 def _derY(self, w, x, y, z):
3466 """
3467 Returns the derivative with respect to y of the interpolated function
3468 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeY.
3469 """
3470 m = len(x)
3471 y_pos = np.searchsorted(self.y_list, y)
3472 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
3473 y_pos[y_pos < 1] = 1
3474 z_pos = np.searchsorted(self.z_list, z)
3475 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
3476 z_pos[z_pos < 1] = 1
3477 dfdy = np.zeros(m) + np.nan
3478 for i in range(1, self.y_n):
3479 for j in range(1, self.z_n):
3480 c = np.logical_and(i == y_pos, j == z_pos)
3481 if np.any(c):
3482 beta = (z[c] - self.z_list[j - 1]) / (
3483 self.z_list[j] - self.z_list[j - 1]
3484 )
3485 dfdy[c] = (
3486 (
3487 (1 - beta) * self.wxInterpolators[i][j - 1](w[c], x[c])
3488 + beta * self.wxInterpolators[i][j](w[c], x[c])
3489 )
3490 - (
3491 (1 - beta) * self.wxInterpolators[i - 1][j - 1](w[c], x[c])
3492 + beta * self.wxInterpolators[i - 1][j](w[c], x[c])
3493 )
3494 ) / (self.y_list[i] - self.y_list[i - 1])
3495 return dfdy
3497 def _derZ(self, w, x, y, z):
3498 """
3499 Returns the derivative with respect to z of the interpolated function
3500 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeZ.
3501 """
3502 m = len(x)
3503 y_pos = np.searchsorted(self.y_list, y)
3504 y_pos[y_pos > self.y_n - 1] = self.y_n - 1
3505 y_pos[y_pos < 1] = 1
3506 z_pos = np.searchsorted(self.z_list, z)
3507 z_pos[z_pos > self.z_n - 1] = self.z_n - 1
3508 z_pos[z_pos < 1] = 1
3509 dfdz = np.zeros(m) + np.nan
3510 for i in range(1, self.y_n):
3511 for j in range(1, self.z_n):
3512 c = np.logical_and(i == y_pos, j == z_pos)
3513 if np.any(c):
3514 alpha = (y[c] - self.y_list[i - 1]) / (
3515 self.y_list[i] - self.y_list[i - 1]
3516 )
3517 dfdz[c] = (
3518 (
3519 (1 - alpha) * self.wxInterpolators[i - 1][j](w[c], x[c])
3520 + alpha * self.wxInterpolators[i][j](w[c], x[c])
3521 )
3522 - (
3523 (1 - alpha) * self.wxInterpolators[i - 1][j - 1](w[c], x[c])
3524 + alpha * self.wxInterpolators[i][j - 1](w[c], x[c])
3525 )
3526 ) / (self.z_list[j] - self.z_list[j - 1])
3527 return dfdz
3530class Curvilinear2DInterp(HARKinterpolator2D):
3531 """
3532 A 2D interpolation method for curvilinear or "warped grid" interpolation, as
3533 in White (2015). Used for models with two endogenous states that are solved
3534 with the endogenous grid method. Allows multiple function outputs, but all of
3535 the interpolated functions must share a common curvilinear grid.
3537 Parameters
3538 ----------
3539 f_values: numpy.array or [numpy.array]
3540 One or more 2D arrays of function values such that f_values[i,j] =
3541 f(x_values[i,j],y_values[i,j]).
3542 x_values: numpy.array
3543 A 2D array of x values of the same shape as f_values.
3544 y_values: numpy.array
3545 A 2D array of y values of the same shape as f_values.
3546 """
3548 distance_criteria = ["f_values", "x_values", "y_values"]
3550 def __init__(self, f_values, x_values, y_values):
3551 if isinstance(f_values, list):
3552 N_funcs = len(f_values)
3553 multi = True
3554 else:
3555 N_funcs = 1
3556 multi = False
3557 my_shape = x_values.shape
3558 if not (my_shape == y_values.shape):
3559 raise ValueError("y_values must have the same shape as x_values!")
3560 if multi:
3561 for n in range(N_funcs):
3562 if not (my_shape == f_values[n].shape):
3563 raise ValueError(
3564 "Each element of f_values must have the same shape as x_values!"
3565 )
3566 else:
3567 if not (my_shape == f_values.shape):
3568 raise ValueError("f_values must have the same shape as x_values!")
3570 if multi:
3571 self.f_values = f_values
3572 else:
3573 self.f_values = [f_values]
3574 self.x_values = x_values
3575 self.y_values = y_values
3576 self.x_n = my_shape[0]
3577 self.y_n = my_shape[1]
3578 self.N_funcs = N_funcs
3579 self.multi = multi
3580 self.update_polarity()
3582 def __call__(self, x, y):
3583 """
3584 Modification of HARKinterpolator2D.__call__ to account for multiple outputs.
3585 """
3586 xa = np.asarray(x)
3587 ya = np.asarray(y)
3588 S = xa.shape
3589 fa = self._evaluate(xa.flatten(), ya.flatten())
3590 output = [fa[n].reshape(S) for n in range(self.N_funcs)]
3591 if self.multi:
3592 return output
3593 else:
3594 return output[0]
3596 def derivativeX(self, x, y):
3597 """
3598 Modification of HARKinterpolator2D.derivativeX to account for multiple outputs.
3599 """
3600 xa = np.asarray(x)
3601 ya = np.asarray(y)
3602 S = xa.shape
3603 dfdxa = self._derX(xa.flatten(), ya.flatten())
3604 output = [dfdxa[n].reshape(S) for n in range(self.N_funcs)]
3605 if self.multi:
3606 return output
3607 else:
3608 return output[0]
3610 def derivativeY(self, x, y):
3611 """
3612 Modification of HARKinterpolator2D.derivativeY to account for multiple outputs.
3613 """
3614 xa = np.asarray(x)
3615 ya = np.asarray(y)
3616 S = xa.shape
3617 dfdya = self._derY(xa.flatten(), ya.flatten())
3618 output = [dfdya[n].reshape(S) for n in range(self.N_funcs)]
3619 if self.multi:
3620 return output
3621 else:
3622 return output[0]
3624 def update_polarity(self):
3625 """
3626 Fills in the polarity attribute of the interpolation, determining whether
3627 the "plus" (True) or "minus" (False) solution of the system of equations
3628 should be used for each sector. Needs to be called in __init__.
3630 Parameters
3631 ----------
3632 none
3634 Returns
3635 -------
3636 none
3637 """
3638 # Grab a point known to be inside each sector: the midway point between
3639 # the lower left and upper right vertex of each sector
3640 x_temp = 0.5 * (
3641 self.x_values[0 : (self.x_n - 1), 0 : (self.y_n - 1)]
3642 + self.x_values[1 : self.x_n, 1 : self.y_n]
3643 )
3644 y_temp = 0.5 * (
3645 self.y_values[0 : (self.x_n - 1), 0 : (self.y_n - 1)]
3646 + self.y_values[1 : self.x_n, 1 : self.y_n]
3647 )
3648 size = (self.x_n - 1) * (self.y_n - 1)
3649 x_temp = np.reshape(x_temp, size)
3650 y_temp = np.reshape(y_temp, size)
3651 y_pos = np.tile(np.arange(0, self.y_n - 1), self.x_n - 1)
3652 x_pos = np.reshape(
3653 np.tile(np.arange(0, self.x_n - 1), (self.y_n - 1, 1)).transpose(), size
3654 )
3656 # Set the polarity of all sectors to "plus", then test each sector
3657 self.polarity = np.ones((self.x_n - 1, self.y_n - 1), dtype=bool)
3658 alpha, beta = self.find_coords(x_temp, y_temp, x_pos, y_pos)
3659 polarity = np.logical_and(
3660 np.logical_and(alpha > 0, alpha < 1), np.logical_and(beta > 0, beta < 1)
3661 )
3663 # Update polarity: if (alpha,beta) not in the unit square, then that
3664 # sector must use the "minus" solution instead
3665 self.polarity = np.reshape(polarity, (self.x_n - 1, self.y_n - 1))
3667 def find_sector(self, x, y):
3668 """
3669 Finds the quadrilateral "sector" for each (x,y) point in the input.
3670 Only called as a subroutine of _evaluate(), etc. Uses a numba helper
3671 function below to accelerate computation.
3673 Parameters
3674 ----------
3675 x : np.array
3676 Values whose sector should be found.
3677 y : np.array
3678 Values whose sector should be found. Should be same size as x.
3680 Returns
3681 -------
3682 x_pos : np.array
3683 Sector x-coordinates for each point of the input, of the same size.
3684 y_pos : np.array
3685 Sector y-coordinates for each point of the input, of the same size.
3686 """
3687 x_pos, y_pos = find_sector_numba(x, y, self.x_values, self.y_values)
3688 return x_pos, y_pos
3690 def find_coords(self, x, y, x_pos, y_pos):
3691 """
3692 Calculates the relative coordinates (alpha,beta) for each point (x,y),
3693 given the sectors (x_pos,y_pos) in which they reside. Only called as
3694 a subroutine of _evaluate(), etc. Uses a numba helper function to acc-
3695 elerate computation, and has a "backup method" for when the math fails.
3697 Parameters
3698 ----------
3699 x : np.array
3700 Values whose sector should be found.
3701 y : np.array
3702 Values whose sector should be found. Should be same size as x.
3703 x_pos : np.array
3704 Sector x-coordinates for each point in (x,y), of the same size.
3705 y_pos : np.array
3706 Sector y-coordinates for each point in (x,y), of the same size.
3708 Returns
3709 -------
3710 alpha : np.array
3711 Relative "horizontal" position of the input in their respective sectors.
3712 beta : np.array
3713 Relative "vertical" position of the input in their respective sectors.
3714 """
3715 alpha, beta = find_coords_numba(
3716 x, y, x_pos, y_pos, self.x_values, self.y_values, self.polarity
3717 )
3719 # Alternate method if there are sectors that are "too regular"
3720 # These points weren't able to identify coordinates
3721 z = np.logical_or(np.isnan(alpha), np.isnan(beta))
3722 if np.any(z):
3723 ii = x_pos[z]
3724 jj = y_pos[z]
3725 xA = self.x_values[ii, jj]
3726 xB = self.x_values[ii + 1, jj]
3727 xC = self.x_values[ii, jj + 1]
3728 xD = self.x_values[ii + 1, jj + 1]
3729 yA = self.y_values[ii, jj]
3730 yB = self.y_values[ii + 1, jj]
3731 yC = self.y_values[ii, jj + 1]
3732 # yD = self.y_values[ii + 1, jj + 1]
3733 b = xB - xA
3734 f = yB - yA
3735 kappa = f / b
3736 int_bot = yA - kappa * xA
3737 int_top = yC - kappa * xC
3738 int_these = y[z] - kappa * x[z]
3739 beta_temp = (int_these - int_bot) / (int_top - int_bot)
3740 x_left = beta_temp * xC + (1.0 - beta_temp) * xA
3741 x_right = beta_temp * xD + (1.0 - beta_temp) * xB
3742 alpha_temp = (x[z] - x_left) / (x_right - x_left)
3743 beta[z] = beta_temp
3744 alpha[z] = alpha_temp
3746 return alpha, beta
3748 def _evaluate(self, x, y):
3749 """
3750 Returns the level of the interpolated function at each value in x,y.
3751 Only called internally by __call__ (etc).
3752 """
3753 x_pos, y_pos = self.find_sector(x, y)
3754 alpha, beta = self.find_coords(x, y, x_pos, y_pos)
3756 # Get weights on each vertex
3757 alpha_C = 1.0 - alpha
3758 beta_C = 1.0 - beta
3759 wA = alpha_C * beta_C
3760 wB = alpha * beta_C
3761 wC = alpha_C * beta
3762 wD = alpha * beta
3764 # Evaluate each function by bilinear interpolation
3765 f = []
3766 for n in range(self.N_funcs):
3767 f_n = (
3768 0.0
3769 + wA * self.f_values[n][x_pos, y_pos]
3770 + wB * self.f_values[n][x_pos + 1, y_pos]
3771 + wC * self.f_values[n][x_pos, y_pos + 1]
3772 + wD * self.f_values[n][x_pos + 1, y_pos + 1]
3773 )
3774 f.append(f_n)
3775 return f
3777 def _derX(self, x, y):
3778 """
3779 Returns the derivative with respect to x of the interpolated function
3780 at each value in x,y. Only called internally by derivativeX.
3781 """
3782 x_pos, y_pos = self.find_sector(x, y)
3783 alpha, beta = self.find_coords(x, y, x_pos, y_pos)
3785 # Get four corners data for each point
3786 xA = self.x_values[x_pos, y_pos]
3787 xB = self.x_values[x_pos + 1, y_pos]
3788 xC = self.x_values[x_pos, y_pos + 1]
3789 xD = self.x_values[x_pos + 1, y_pos + 1]
3790 yA = self.y_values[x_pos, y_pos]
3791 yB = self.y_values[x_pos + 1, y_pos]
3792 yC = self.y_values[x_pos, y_pos + 1]
3793 yD = self.y_values[x_pos + 1, y_pos + 1]
3795 # Calculate components of the alpha,beta --> x,y delta translation matrix
3796 alpha_C = 1 - alpha
3797 beta_C = 1 - beta
3798 alpha_x = beta_C * (xB - xA) + beta * (xD - xC)
3799 alpha_y = beta_C * (yB - yA) + beta * (yD - yC)
3800 beta_x = alpha_C * (xC - xA) + alpha * (xD - xB)
3801 beta_y = alpha_C * (yC - yA) + alpha * (yD - yB)
3803 # Invert the delta translation matrix into x,y --> alpha,beta
3804 det = alpha_x * beta_y - beta_x * alpha_y
3805 x_alpha = beta_y / det
3806 x_beta = -alpha_y / det
3808 # Calculate the derivative of f w.r.t. alpha and beta for each function
3809 dfdx = []
3810 for n in range(self.N_funcs):
3811 fA = self.f_values[n][x_pos, y_pos]
3812 fB = self.f_values[n][x_pos + 1, y_pos]
3813 fC = self.f_values[n][x_pos, y_pos + 1]
3814 fD = self.f_values[n][x_pos + 1, y_pos + 1]
3815 dfda = beta_C * (fB - fA) + beta * (fD - fC)
3816 dfdb = alpha_C * (fC - fA) + alpha * (fD - fB)
3818 # Calculate the derivative with respect to x
3819 dfdx_n = x_alpha * dfda + x_beta * dfdb
3820 dfdx.append(dfdx_n)
3821 return dfdx
3823 def _derY(self, x, y):
3824 """
3825 Returns the derivative with respect to y of the interpolated function
3826 at each value in x,y. Only called internally by derivativeY.
3827 """
3828 x_pos, y_pos = self.find_sector(x, y)
3829 alpha, beta = self.find_coords(x, y, x_pos, y_pos)
3831 # Get four corners data for each point
3832 xA = self.x_values[x_pos, y_pos]
3833 xB = self.x_values[x_pos + 1, y_pos]
3834 xC = self.x_values[x_pos, y_pos + 1]
3835 xD = self.x_values[x_pos + 1, y_pos + 1]
3836 yA = self.y_values[x_pos, y_pos]
3837 yB = self.y_values[x_pos + 1, y_pos]
3838 yC = self.y_values[x_pos, y_pos + 1]
3839 yD = self.y_values[x_pos + 1, y_pos + 1]
3841 # Calculate components of the alpha,beta --> x,y delta translation matrix
3842 alpha_C = 1 - alpha
3843 beta_C = 1 - beta
3844 alpha_x = beta_C * (xB - xA) + beta * (xD - xC)
3845 alpha_y = beta_C * (yB - yA) + beta * (yD - yC)
3846 beta_x = alpha_C * (xC - xA) + alpha * (xD - xB)
3847 beta_y = alpha_C * (yC - yA) + alpha * (yD - yB)
3849 # Invert the delta translation matrix into x,y --> alpha,beta
3850 det = alpha_x * beta_y - beta_x * alpha_y
3851 y_alpha = -beta_x / det
3852 y_beta = alpha_x / det
3854 # Calculate the derivative of f w.r.t. alpha and beta for each function
3855 dfdy = []
3856 for n in range(self.N_funcs):
3857 fA = self.f_values[n][x_pos, y_pos]
3858 fB = self.f_values[n][x_pos + 1, y_pos]
3859 fC = self.f_values[n][x_pos, y_pos + 1]
3860 fD = self.f_values[n][x_pos + 1, y_pos + 1]
3861 dfda = beta_C * (fB - fA) + beta * (fD - fC)
3862 dfdb = alpha_C * (fC - fA) + alpha * (fD - fB)
3864 # Calculate the derivative with respect to y
3865 dfdy_n = y_alpha * dfda + y_beta * dfdb
3866 dfdy.append(dfdy_n)
3867 return dfdy
3870# Define a function that checks whether a set of points violates a linear boundary
3871# defined by (x1,y1) and (x2,y2), where the latter is *COUNTER CLOCKWISE* from the
3872# former. Returns 1 if the point is outside the boundary and 0 otherwise.
3873@njit
3874def boundary_check(xq, yq, x1, y1, x2, y2): # pragma: no cover
3875 return int((y2 - y1) * xq - (x2 - x1) * yq > x1 * y2 - y1 * x2)
3878# Define a numba helper function for finding the sector in the irregular grid
3879@njit
3880def find_sector_numba(X_query, Y_query, X_values, Y_values): # pragma: no cover
3881 # Initialize the sector guess
3882 M = X_query.size
3883 x_n = X_values.shape[0]
3884 y_n = X_values.shape[1]
3885 ii = int(x_n / 2)
3886 jj = int(y_n / 2)
3887 top_ii = x_n - 2
3888 top_jj = y_n - 2
3890 # Initialize the output arrays
3891 X_pos = np.empty(M, dtype=np.int32)
3892 Y_pos = np.empty(M, dtype=np.int32)
3894 # Identify the correct sector for each point to be evaluated
3895 max_loops = x_n + y_n
3896 for m in range(M):
3897 found = False
3898 loops = 0
3899 while not found and loops < max_loops:
3900 # Get coordinates for the four vertices: (xA,yA),...,(xD,yD)
3901 x0 = X_query[m]
3902 y0 = Y_query[m]
3903 xA = X_values[ii, jj]
3904 xB = X_values[ii + 1, jj]
3905 xC = X_values[ii, jj + 1]
3906 xD = X_values[ii + 1, jj + 1]
3907 yA = Y_values[ii, jj]
3908 yB = Y_values[ii + 1, jj]
3909 yC = Y_values[ii, jj + 1]
3910 yD = Y_values[ii + 1, jj + 1]
3912 # Check the "bounding box" for the sector: is this guess plausible?
3913 D = int(y0 < np.minimum(yA, yB))
3914 R = int(x0 > np.maximum(xB, xD))
3915 U = int(y0 > np.maximum(yC, yD))
3916 L = int(x0 < np.minimum(xA, xC))
3918 # Check which boundaries are violated (and thus where to look next)
3919 in_box = np.all(np.logical_not(np.array([D, R, U, L])))
3920 if in_box:
3921 D = boundary_check(x0, y0, xA, yA, xB, yB)
3922 R = boundary_check(x0, y0, xB, yB, xD, yD)
3923 U = boundary_check(x0, y0, xD, yD, xC, yC)
3924 L = boundary_check(x0, y0, xC, yC, xA, yA)
3926 # Update the sector guess based on the violations
3927 ii_next = np.maximum(np.minimum(ii - L + R, top_ii), 0)
3928 jj_next = np.maximum(np.minimum(jj - D + U, top_jj), 0)
3930 # Check whether sector guess changed and go to next iteration
3931 found = (ii == ii_next) and (jj == jj_next)
3932 ii = ii_next
3933 jj = jj_next
3934 loops += 1
3936 # Put the final sector guess into the output array
3937 X_pos[m] = ii
3938 Y_pos[m] = jj
3940 # Return the output
3941 return X_pos, Y_pos
3944# Define a numba helper function for finding relative coordinates within sector
3945@njit
3946def find_coords_numba(
3947 X_query, Y_query, X_pos, Y_pos, X_values, Y_values, polarity
3948): # pragma: no cover
3949 M = X_query.size
3950 alpha = np.empty(M)
3951 beta = np.empty(M)
3953 # Calculate relative coordinates in the sector for each point
3954 for m in range(M):
3955 try:
3956 x0 = X_query[m]
3957 y0 = Y_query[m]
3958 ii = X_pos[m]
3959 jj = Y_pos[m]
3960 xA = X_values[ii, jj]
3961 xB = X_values[ii + 1, jj]
3962 xC = X_values[ii, jj + 1]
3963 xD = X_values[ii + 1, jj + 1]
3964 yA = Y_values[ii, jj]
3965 yB = Y_values[ii + 1, jj]
3966 yC = Y_values[ii, jj + 1]
3967 yD = Y_values[ii + 1, jj + 1]
3968 p = 2.0 * polarity[ii, jj] - 1.0
3969 a = xA
3970 b = xB - xA
3971 c = xC - xA
3972 d = xA - xB - xC + xD
3973 e = yA
3974 f = yB - yA
3975 g = yC - yA
3976 h = yA - yB - yC + yD
3977 denom = d * g - h * c
3978 mu = (h * b - d * f) / denom
3979 tau = (h * (a - x0) - d * (e - y0)) / denom
3980 zeta = a - x0 + c * tau
3981 eta = b + c * mu + d * tau
3982 theta = d * mu
3983 alph = (-eta + p * np.sqrt(eta**2 - 4 * zeta * theta)) / (2 * theta)
3984 bet = mu * alph + tau
3985 except:
3986 alph = np.nan
3987 bet = np.nan
3988 alpha[m] = alph
3989 beta[m] = bet
3991 return alpha, beta
3994class DiscreteInterp(MetricObject):
3995 """
3996 An interpolator for variables that can only take a discrete set of values.
3998 If the function we wish to interpolate, f(args) can take on the list of
3999 values discrete_vals, this class expects an interpolator for the index of
4000 f's value in discrete_vals.
4001 E.g., if f(a,b,c) = discrete_vals[5], then index_interp(a,b,c) = 5.
4003 Parameters
4004 ----------
4005 index_interp: HARKInterpolator
4006 An interpolator giving an approximation to the index of the value in
4007 discrete_vals that corresponds to a given set of arguments.
4008 discrete_vals: numpy.array
4009 A 1D array containing the values in the range of the discrete function
4010 to be interpolated.
4011 """
4013 distance_criteria = ["index_interp"]
4015 def __init__(self, index_interp, discrete_vals):
4016 self.index_interp = index_interp
4017 self.discrete_vals = discrete_vals
4018 self.n_vals = len(self.discrete_vals)
4020 def __call__(self, *args):
4021 # Interpolate indices and round to integers
4022 inds = np.rint(self.index_interp(*args)).astype(int)
4023 if type(inds) is not np.ndarray:
4024 inds = np.array(inds)
4025 # Deal with out-of range indices
4026 inds[inds < 0] = 0
4027 inds[inds >= self.n_vals] = self.n_vals - 1
4029 # Get values from grid
4030 return self.discrete_vals[inds]
4033class IndexedInterp(MetricObject):
4034 """
4035 An interpolator for functions whose first argument is an integer-valued index.
4036 Constructor takes in a list of functions as its only argument. When evaluated
4037 at f(i,X), interpolator returns f[i](X), where X can be any number of inputs.
4038 This simply provides a different interface for accessing the same functions.
4040 Parameters
4041 ----------
4042 functions : [Callable]
4043 List of one or more functions to be indexed.
4044 """
4046 distance_criteria = ["functions"]
4048 def __init__(self, functions):
4049 self.functions = functions
4050 self.N = len(functions)
4052 def __call__(self, idx, *args):
4053 out = np.empty(idx.shape)
4054 out.fill(np.nan)
4056 for n in range(self.N):
4057 these = idx == n
4058 if not np.any(these):
4059 continue
4060 temp = [arg[these] for arg in args]
4061 out[these] = self.functions[n](*temp)
4062 return out
4065###############################################################################
4066## Functions used in discrete choice models with T1EV taste shocks ############
4067###############################################################################
4070def calc_log_sum_choice_probs(Vals, sigma):
4071 """
4072 Returns the final optimal value and choice probabilities given the choice
4073 specific value functions `Vals`. Probabilities are degenerate if sigma == 0.0.
4074 Parameters
4075 ----------
4076 Vals : [numpy.array]
4077 A numpy.array that holds choice specific values at common grid points.
4078 sigma : float
4079 A number that controls the variance of the taste shocks
4080 Returns
4081 -------
4082 V : [numpy.array]
4083 A numpy.array that holds the integrated value function.
4084 P : [numpy.array]
4085 A numpy.array that holds the discrete choice probabilities
4086 """
4087 # Assumes that NaNs have been replaced by -numpy.inf or similar
4088 if sigma == 0.0:
4089 # We could construct a linear index here and use unravel_index.
4090 Pflat = np.argmax(Vals, axis=0)
4092 V = np.zeros(Vals[0].shape)
4093 Probs = np.zeros(Vals.shape)
4094 for i in range(Vals.shape[0]):
4095 optimalIndices = Pflat == i
4096 V[optimalIndices] = Vals[i][optimalIndices]
4097 Probs[i][optimalIndices] = 1
4098 return V, Probs
4100 # else we have a taste shock
4101 maxV = np.max(Vals, axis=0)
4103 # calculate maxV+sigma*log(sum_i=1^J exp((V[i]-maxV))/sigma)
4104 sumexp = np.sum(np.exp((Vals - maxV) / sigma), axis=0)
4105 LogSumV = np.log(sumexp)
4106 LogSumV = maxV + sigma * LogSumV
4108 Probs = np.exp((Vals - LogSumV) / sigma)
4109 return LogSumV, Probs
4112def calc_choice_probs(Vals, sigma):
4113 """
4114 Returns the choice probabilities given the choice specific value functions
4115 `Vals`. Probabilities are degenerate if sigma == 0.0.
4116 Parameters
4117 ----------
4118 Vals : [numpy.array]
4119 A numpy.array that holds choice specific values at common grid points.
4120 sigma : float
4121 A number that controls the variance of the taste shocks
4122 Returns
4123 -------
4124 Probs : [numpy.array]
4125 A numpy.array that holds the discrete choice probabilities
4126 """
4128 # Assumes that NaNs have been replaced by -numpy.inf or similar
4129 if sigma == 0.0:
4130 # We could construct a linear index here and use unravel_index.
4131 Pflat = np.argmax(Vals, axis=0)
4132 Probs = np.zeros(Vals.shape)
4133 for i in range(Vals.shape[0]):
4134 Probs[i][Pflat == i] = 1
4135 return Probs
4137 maxV = np.max(Vals, axis=0)
4138 Probs = np.divide(
4139 np.exp((Vals - maxV) / sigma), np.sum(np.exp((Vals - maxV) / sigma), axis=0)
4140 )
4141 return Probs
4144def calc_log_sum(Vals, sigma):
4145 """
4146 Returns the optimal value given the choice specific value functions Vals.
4147 Parameters
4148 ----------
4149 Vals : [numpy.array]
4150 A numpy.array that holds choice specific values at common grid points.
4151 sigma : float
4152 A number that controls the variance of the taste shocks
4153 Returns
4154 -------
4155 V : [numpy.array]
4156 A numpy.array that holds the integrated value function.
4157 """
4159 # Assumes that NaNs have been replaced by -numpy.inf or similar
4160 if sigma == 0.0:
4161 # We could construct a linear index here and use unravel_index.
4162 V = np.amax(Vals, axis=0)
4163 return V
4165 # else we have a taste shock
4166 maxV = np.max(Vals, axis=0)
4168 # calculate maxV+sigma*log(sum_i=1^J exp((V[i]-maxV))/sigma)
4169 sumexp = np.sum(np.exp((Vals - maxV) / sigma), axis=0)
4170 LogSumV = np.log(sumexp)
4171 LogSumV = maxV + sigma * LogSumV
4172 return LogSumV
4175###############################################################################
4176# Tools for value and marginal-value functions in models where #
4177# - dvdm = u'(c). #
4178# - u is of the CRRA family. #
4179###############################################################################
4182class ValueFuncCRRA(MetricObject):
4183 """
4184 A class for representing a value function. The underlying interpolation is
4185 in the space of (state,u_inv(v)); this class "re-curves" to the value function.
4187 Parameters
4188 ----------
4189 vFuncNvrs : function
4190 A real function representing the value function composed with the
4191 inverse utility function, defined on the state: u_inv(vFunc(state))
4192 CRRA : float
4193 Coefficient of relative risk aversion.
4194 illegal_value : float, optional
4195 If provided, value to return for "out-of-bounds" inputs that return NaN
4196 from the pseudo-inverse value function. Most common choice is -np.inf,
4197 which makes the outcome infinitely bad.
4198 """
4200 distance_criteria = ["func", "CRRA"]
4202 def __init__(self, vFuncNvrs, CRRA, illegal_value=None):
4203 self.vFuncNvrs = deepcopy(vFuncNvrs)
4204 self.CRRA = CRRA
4205 self.illegal_value = illegal_value
4207 if hasattr(vFuncNvrs, "grid_list"):
4208 self.grid_list = vFuncNvrs.grid_list
4209 else:
4210 self.grid_list = None
4212 def __call__(self, *vFuncArgs):
4213 """
4214 Evaluate the value function at given levels of market resources m.
4216 Parameters
4217 ----------
4218 vFuncArgs : floats or np.arrays, all of the same dimensions.
4219 Values for the state variables. These usually start with 'm',
4220 market resources normalized by the level of permanent income.
4222 Returns
4223 -------
4224 v : float or np.array
4225 Lifetime value of beginning this period with the given states; has
4226 same size as the state inputs.
4227 """
4228 temp = self.vFuncNvrs(*vFuncArgs)
4229 v = CRRAutility(temp, self.CRRA)
4230 if self.illegal_value is not None:
4231 illegal = np.isnan(temp)
4232 v[illegal] = self.illegal_value
4233 return v
4235 def gradient(self, *args):
4236 NvrsGrad = self.vFuncNvrs.gradient(*args)
4237 marg_u = CRRAutilityP(*args, self.CRRA)
4238 grad = [g * marg_u for g in NvrsGrad]
4239 return grad
4241 def _eval_and_grad(self, *args):
4242 return (self.__call__(*args), self.gradient(*args))
4245class MargValueFuncCRRA(MetricObject):
4246 """
4247 A class for representing a marginal value function in models where the
4248 standard envelope condition of dvdm(state) = u'(c(state)) holds (with CRRA utility).
4250 Parameters
4251 ----------
4252 cFunc : function.
4253 Its first argument must be normalized market resources m.
4254 A real function representing the marginal value function composed
4255 with the inverse marginal utility function, defined on the state
4256 variables: uP_inv(dvdmFunc(state)). Called cFunc because when standard
4257 envelope condition applies, uP_inv(dvdm(state)) = cFunc(state).
4258 CRRA : float
4259 Coefficient of relative risk aversion.
4260 """
4262 distance_criteria = ["cFunc", "CRRA"]
4264 def __init__(self, cFunc, CRRA):
4265 self.cFunc = deepcopy(cFunc)
4266 self.CRRA = CRRA
4268 if hasattr(cFunc, "grid_list"):
4269 self.grid_list = cFunc.grid_list
4270 else:
4271 self.grid_list = None
4273 def __call__(self, *cFuncArgs):
4274 """
4275 Evaluate the marginal value function at given levels of market resources m.
4277 Parameters
4278 ----------
4279 cFuncArgs : floats or np.arrays
4280 Values of the state variables at which to evaluate the marginal
4281 value function.
4283 Returns
4284 -------
4285 vP : float or np.array
4286 Marginal lifetime value of beginning this period with state
4287 cFuncArgs
4288 """
4289 return CRRAutilityP(self.cFunc(*cFuncArgs), rho=self.CRRA)
4291 def derivativeX(self, *cFuncArgs):
4292 """
4293 Evaluate the derivative of the marginal value function with respect to
4294 market resources at given state; this is the marginal marginal value
4295 function.
4297 Parameters
4298 ----------
4299 cFuncArgs : floats or np.arrays
4300 State variables.
4302 Returns
4303 -------
4304 vPP : float or np.array
4305 Marginal marginal lifetime value of beginning this period with
4306 state cFuncArgs; has same size as inputs.
4308 """
4310 # The derivative method depends on the dimension of the function
4311 if isinstance(self.cFunc, HARKinterpolator1D):
4312 c, MPC = self.cFunc.eval_with_derivative(*cFuncArgs)
4314 elif hasattr(self.cFunc, "derivativeX"):
4315 c = self.cFunc(*cFuncArgs)
4316 MPC = self.cFunc.derivativeX(*cFuncArgs)
4318 else:
4319 raise Exception(
4320 "cFunc does not have a 'derivativeX' attribute. Can't compute"
4321 + "marginal marginal value."
4322 )
4324 return MPC * CRRAutilityPP(c, rho=self.CRRA)
4327class MargMargValueFuncCRRA(MetricObject):
4328 """
4329 A class for representing a marginal marginal value function in models where
4330 the standard envelope condition of dvdm = u'(c(state)) holds (with CRRA utility).
4332 Parameters
4333 ----------
4334 cFunc : function.
4335 Its first argument must be normalized market resources m.
4336 A real function representing the marginal value function composed
4337 with the inverse marginal utility function, defined on the state
4338 variables: uP_inv(dvdmFunc(state)). Called cFunc because when standard
4339 envelope condition applies, uP_inv(dvdm(state)) = cFunc(state).
4340 CRRA : float
4341 Coefficient of relative risk aversion.
4342 """
4344 distance_criteria = ["cFunc", "CRRA"]
4346 def __init__(self, cFunc, CRRA):
4347 self.cFunc = deepcopy(cFunc)
4348 self.CRRA = CRRA
4350 def __call__(self, *cFuncArgs):
4351 """
4352 Evaluate the marginal marginal value function at given levels of market
4353 resources m.
4355 Parameters
4356 ----------
4357 m : float or np.array
4358 Market resources (normalized by permanent income) whose marginal
4359 marginal value is to be found.
4361 Returns
4362 -------
4363 vPP : float or np.array
4364 Marginal marginal lifetime value of beginning this period with market
4365 resources m; has same size as input m.
4366 """
4368 # The derivative method depends on the dimension of the function
4369 if isinstance(self.cFunc, HARKinterpolator1D):
4370 c, MPC = self.cFunc.eval_with_derivative(*cFuncArgs)
4372 elif hasattr(self.cFunc, "derivativeX"):
4373 c = self.cFunc(*cFuncArgs)
4374 MPC = self.cFunc.derivativeX(*cFuncArgs)
4376 else:
4377 raise Exception(
4378 "cFunc does not have a 'derivativeX' attribute. Can't compute"
4379 + "marginal marginal value."
4380 )
4381 return MPC * CRRAutilityPP(c, rho=self.CRRA)