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