Coverage for HARK / interpolation.py: 97%
1081 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-05-10 06:19 +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 _coerce_1d_grid(arr):
55 """Return ``arr`` as a 1D numpy array, flattening if necessary."""
56 a = np.asarray(arr)
57 if a.ndim != 1:
58 warnings.warn("input not of the size (n, ), attempting to flatten")
59 return a.flatten()
60 return a
63def _broadcast_eval(inner, *args):
64 """Broadcast ``args`` to a common shape, call ``inner`` on the flattened
65 arrays, and reshape the result.
67 Shared by the ``__call__``/``derivativeX``/``derivativeY``/... methods of
68 :class:`HARKinterpolator2D`, :class:`HARKinterpolator3D`, and
69 :class:`HARKinterpolator4D`.
70 """
71 arrs = list(np.broadcast_arrays(*[np.asarray(a) for a in args]))
72 return inner(*[a.flatten() for a in arrs]).reshape(arrs[0].shape)
75def _locate_clipped(grid, values, n):
76 """Return ``np.searchsorted(grid, values)`` clipped into ``[1, n - 1]``.
78 Shared by every interpolator that brackets queries with ``a_list[idx - 1]``
79 and ``a_list[idx]``: a single clipped index per axis is enough for
80 1D/2D/3D/4D evaluation and partial-derivative loops.
81 """
82 return np.clip(np.searchsorted(grid, values), 1, n - 1)
85def _cell_fraction(grid, idx, queries):
86 """Linear-cell fractional position of ``queries`` within ``[grid[idx-1], grid[idx]]``.
88 Returns ``(queries - grid[idx - 1]) / (grid[idx] - grid[idx - 1])``. Works
89 with ``idx`` as a scalar (interp-on-interp loops) or an integer array
90 (tensor-grid interpolators).
91 """
92 lower = grid[idx - 1]
93 return (queries - lower) / (grid[idx] - lower)
96def _iter_unique_pairs(*positions):
97 """Yield ``(*indices, mask)`` for each unique observed combination of axis ``positions``.
99 Accepts any number of equal-length 1D position arrays. Cells that no
100 query falls into are silently skipped. Yielded indices are Python
101 ``int``s, safe for list indexing. No-op when no positions are passed
102 or when each axis is empty.
103 """
104 if not positions or positions[0].size == 0:
105 return
106 stacked = np.column_stack(positions)
107 combos, inverse = np.unique(stacked, axis=0, return_inverse=True)
108 for k, combo in enumerate(combos):
109 yield (*(int(v) for v in combo), inverse == k)
112def _envelope_partial(envelope, args, deriv_attr):
113 """Compute an envelope partial derivative.
115 Evaluates each member function on the broadcast ``args`` to identify the
116 active function per point (via ``envelope.argcompare``), then takes the
117 requested derivative (``deriv_attr``) of the active function on its slice.
118 Shared by ``LowerEnvelope2D`` and ``LowerEnvelope3D`` partial derivatives.
119 """
120 primary = args[0]
121 temp = np.column_stack([f(*args) for f in envelope.functions])
122 active = envelope.argcompare(temp, axis=1)
123 out = np.zeros_like(primary)
124 for j in np.unique(active):
125 c = active == j
126 out[c] = getattr(envelope.functions[j], deriv_attr)(*[a[c] for a in args])
127 return out
130class HARKinterpolator1D(MetricObject):
131 """
132 A wrapper class for 1D interpolation methods in HARK.
133 """
135 distance_criteria = []
137 def __call__(self, x):
138 """
139 Evaluates the interpolated function at the given input.
141 Parameters
142 ----------
143 x : np.array or float
144 Real values to be evaluated in the interpolated function.
146 Returns
147 -------
148 y : np.array or float
149 The interpolated function evaluated at x: y = f(x), with the same
150 shape as x.
151 """
152 z = np.asarray(x)
153 return (self._evaluate(z.flatten())).reshape(z.shape)
155 def derivative(self, x):
156 """
157 Evaluates the derivative of the interpolated function at the given input.
159 Parameters
160 ----------
161 x : np.array or float
162 Real values to be evaluated in the interpolated function.
164 Returns
165 -------
166 dydx : np.array or float
167 The interpolated function's first derivative evaluated at x:
168 dydx = f'(x), with the same shape as x.
169 """
170 z = np.asarray(x)
171 return (self._der(z.flatten())).reshape(z.shape)
173 derivativeX = derivative # alias
175 def eval_with_derivative(self, x):
176 """
177 Evaluates the interpolated function and its derivative at the given input.
179 Parameters
180 ----------
181 x : np.array or float
182 Real values to be evaluated in the interpolated function.
184 Returns
185 -------
186 y : np.array or float
187 The interpolated function evaluated at x: y = f(x), with the same
188 shape as x.
189 dydx : np.array or float
190 The interpolated function's first derivative evaluated at x:
191 dydx = f'(x), with the same shape as x.
192 """
193 z = np.asarray(x)
194 y, dydx = self._evalAndDer(z.flatten())
195 return y.reshape(z.shape), dydx.reshape(z.shape)
197 def _evaluate(self, x):
198 """
199 Interpolated function evaluator, to be defined in subclasses.
200 """
201 raise NotImplementedError()
203 def _der(self, x):
204 """
205 Default or fallback derivative method using finite difference approximation.
206 Subclasses of HARKinterpolator1D should define their own more specific method.
207 """
208 eps = 1e-8
209 f0 = self.__call__(x)
210 f1 = self.__call__(x + eps)
211 dydx = (f1 - f0) / eps
212 return dydx
214 def _evalAndDer(self, x):
215 """
216 Interpolated function and derivative evaluator, to be defined in subclasses.
217 Default implementation separately calls the _evaluate and _der methods, which
218 might be inefficient relative to interpolator-specific implementation.
219 """
220 y = self._evaluate(x)
221 dydx = self._der(x)
222 return y, dydx
224 def _init_cubic_grids(self, x_list, y_list, dydx_list):
225 """
226 Coerce ``x_list``, ``y_list``, ``dydx_list`` to validated 1D arrays.
228 Stores them as ``self.x_list``, ``self.y_list``, ``self.dydx_list``,
229 sets ``self.n``, and runs ``_check_grid_dimensions`` against ``x_list``.
230 Shared between :class:`CubicInterp` and :class:`CubicHermiteInterp`.
231 """
232 self.x_list = _coerce_1d_grid(x_list)
233 self.y_list = _coerce_1d_grid(y_list)
234 self.dydx_list = _coerce_1d_grid(dydx_list)
235 _check_grid_dimensions(1, self.y_list, self.x_list)
236 _check_grid_dimensions(1, self.dydx_list, self.x_list)
237 self.n = self.x_list.size
240class HARKinterpolator2D(MetricObject):
241 """
242 A wrapper class for 2D interpolation methods in HARK.
243 """
245 distance_criteria = []
247 def __call__(self, x, y):
248 """
249 Evaluates the interpolated function at the given input.
251 Parameters
252 ----------
253 x : np.array or float
254 Real values to be evaluated in the interpolated function.
255 y : np.array or float
256 Real values to be evaluated in the interpolated function. If both
257 are arrays, they must be broadcastable to the same shape.
258 Scalar inputs will be broadcast to match array inputs.
260 Returns
261 -------
262 fxy : np.array or float
263 The interpolated function evaluated at x,y: fxy = f(x,y), with the
264 same shape as x and y.
265 """
266 return _broadcast_eval(self._evaluate, x, y)
268 def derivativeX(self, x, y):
269 """
270 Evaluates the partial derivative of interpolated function with respect
271 to x (the first argument) at the given input.
273 Parameters
274 ----------
275 x : np.array or float
276 Real values to be evaluated in the interpolated function.
277 y : np.array or float
278 Real values to be evaluated in the interpolated function. If both
279 are arrays, they must be broadcastable to the same shape.
280 Scalar inputs will be broadcast to match array inputs.
282 Returns
283 -------
284 dfdx : np.array or float
285 The derivative of the interpolated function with respect to x, eval-
286 uated at x,y: dfdx = f_x(x,y), with the same shape as x and y.
287 """
288 return _broadcast_eval(self._derX, x, y)
290 def derivativeY(self, x, y):
291 """
292 Evaluates the partial derivative of interpolated function with respect
293 to y (the second argument) at the given input.
295 Parameters
296 ----------
297 x : np.array or float
298 Real values to be evaluated in the interpolated function.
299 y : np.array or float
300 Real values to be evaluated in the interpolated function. If both
301 are arrays, they must be broadcastable to the same shape.
302 Scalar inputs will be broadcast to match array inputs.
304 Returns
305 -------
306 dfdy : np.array or float
307 The derivative of the interpolated function with respect to y, eval-
308 uated at x,y: dfdx = f_y(x,y), with the same shape as x and y.
309 """
310 return _broadcast_eval(self._derY, x, y)
312 def _evaluate(self, x, y):
313 """
314 Interpolated function evaluator, to be defined in subclasses.
315 """
316 raise NotImplementedError()
318 def _derX(self, x, y):
319 """
320 Default or fallback derivative with respect to x, using finite difference approximation.
321 Subclasses of HARKinterpolator2D should define their own more specific method.
322 """
323 eps = 1e-8
324 f0 = self.__call__(x, y)
325 f1 = self.__call__(x + eps, y)
326 dfdx = (f1 - f0) / eps
327 return dfdx
329 def _derY(self, x, y):
330 """
331 Default or fallback derivative with respect to y, using finite difference approximation.
332 Subclasses of HARKinterpolator2D should define their own more specific method.
333 """
334 eps = 1e-8
335 f0 = self.__call__(x, y)
336 f1 = self.__call__(x, y + eps)
337 dfdy = (f1 - f0) / eps
338 return dfdy
341class HARKinterpolator3D(MetricObject):
342 """
343 A wrapper class for 3D interpolation methods in HARK.
344 """
346 distance_criteria = []
348 def __call__(self, x, y, z):
349 """
350 Evaluates the interpolated function at the given input.
352 Parameters
353 ----------
354 x : np.array or float
355 Real values to be evaluated in the interpolated function.
356 y : np.array or float
357 Real values to be evaluated in the interpolated function. If multiple
358 inputs are arrays, they must be broadcastable to the same shape.
359 Scalar inputs will be broadcast to match array inputs.
360 z : np.array or float
361 Real values to be evaluated in the interpolated function. If multiple
362 inputs are arrays, they must be broadcastable to the same shape.
363 Scalar inputs will be broadcast to match array inputs.
365 Returns
366 -------
367 fxyz : np.array or float
368 The interpolated function evaluated at x,y,z: fxyz = f(x,y,z), with
369 the same shape as x, y, and z.
370 """
371 return _broadcast_eval(self._evaluate, x, y, z)
373 def derivativeX(self, x, y, z):
374 """
375 Evaluates the partial derivative of the interpolated function with respect
376 to x (the first argument) at the given input.
378 Parameters
379 ----------
380 x : np.array or float
381 Real values to be evaluated in the interpolated function.
382 y : np.array or float
383 Real values to be evaluated in the interpolated function. If multiple
384 inputs are arrays, they must be broadcastable to the same shape.
385 Scalar inputs will be broadcast to match array inputs.
386 z : np.array or float
387 Real values to be evaluated in the interpolated function. If multiple
388 inputs are arrays, they must be broadcastable to the same shape.
389 Scalar inputs will be broadcast to match array inputs.
391 Returns
392 -------
393 dfdx : np.array or float
394 The derivative with respect to x of the interpolated function evaluated
395 at x,y,z: dfdx = f_x(x,y,z), with the same shape as x, y, and z.
396 """
397 return _broadcast_eval(self._derX, x, y, z)
399 def derivativeY(self, x, y, z):
400 """
401 Evaluates the partial derivative of the interpolated function with respect
402 to y (the second argument) at the given input.
404 Parameters
405 ----------
406 x : np.array or float
407 Real values to be evaluated in the interpolated function.
408 y : np.array or float
409 Real values to be evaluated in the interpolated function. If multiple
410 inputs are arrays, they must be broadcastable to the same shape.
411 Scalar inputs will be broadcast to match array inputs.
412 z : np.array or float
413 Real values to be evaluated in the interpolated function. If multiple
414 inputs are arrays, they must be broadcastable to the same shape.
415 Scalar inputs will be broadcast to match array inputs.
417 Returns
418 -------
419 dfdy : np.array or float
420 The derivative with respect to y of the interpolated function evaluated
421 at x,y,z: dfdy = f_y(x,y,z), with the same shape as x, y, and z.
422 """
423 return _broadcast_eval(self._derY, x, y, z)
425 def derivativeZ(self, x, y, z):
426 """
427 Evaluates the partial derivative of the interpolated function with respect
428 to z (the third argument) at the given input.
430 Parameters
431 ----------
432 x : np.array or float
433 Real values to be evaluated in the interpolated function.
434 y : np.array or float
435 Real values to be evaluated in the interpolated function. If multiple
436 inputs are arrays, they must be broadcastable to the same shape.
437 Scalar inputs will be broadcast to match array inputs.
438 z : np.array or float
439 Real values to be evaluated in the interpolated function. If multiple
440 inputs are arrays, they must be broadcastable to the same shape.
441 Scalar inputs will be broadcast to match array inputs.
443 Returns
444 -------
445 dfdz : np.array or float
446 The derivative with respect to z of the interpolated function evaluated
447 at x,y,z: dfdz = f_z(x,y,z), with the same shape as x, y, and z.
448 """
449 return _broadcast_eval(self._derZ, x, y, z)
451 def _evaluate(self, x, y, z):
452 """
453 Interpolated function evaluator, to be defined in subclasses.
454 """
455 raise NotImplementedError()
457 def _derX(self, x, y, z):
458 """
459 Default or fallback derivative with respect to x, using finite difference approximation.
460 Subclasses of HARKinterpolator3D should define their own more specific method.
461 """
462 eps = 1e-8
463 f0 = self.__call__(x, y, z)
464 f1 = self.__call__(x + eps, y, z)
465 dfdx = (f1 - f0) / eps
466 return dfdx
468 def _derY(self, x, y, z):
469 """
470 Default or fallback derivative with respect to y, using finite difference approximation.
471 Subclasses of HARKinterpolator3D should define their own more specific method.
472 """
473 eps = 1e-8
474 f0 = self.__call__(x, y, z)
475 f1 = self.__call__(x, y + eps, z)
476 dfdy = (f1 - f0) / eps
477 return dfdy
479 def _derZ(self, x, y, z):
480 """
481 Default or fallback derivative with respect to z, using finite difference approximation.
482 Subclasses of HARKinterpolator3D should define their own more specific method.
483 """
484 eps = 1e-8
485 f0 = self.__call__(x, y, z)
486 f1 = self.__call__(x, y, z + eps)
487 dfdz = (f1 - f0) / eps
488 return dfdz
491class HARKinterpolator4D(MetricObject):
492 """
493 A wrapper class for 4D interpolation methods in HARK.
494 """
496 distance_criteria = []
498 def __call__(self, w, x, y, z):
499 """
500 Evaluates the interpolated function at the given input.
502 Parameters
503 ----------
504 w : np.array or float
505 Real values to be evaluated in the interpolated function.
506 x : np.array or float
507 Real values to be evaluated in the interpolated function. If multiple
508 inputs are arrays, they must be broadcastable to the same shape.
509 Scalar inputs will be broadcast to match array inputs.
510 y : np.array or float
511 Real values to be evaluated in the interpolated function. If multiple
512 inputs are arrays, they must be broadcastable to the same shape.
513 Scalar inputs will be broadcast to match array inputs.
514 z : np.array or float
515 Real values to be evaluated in the interpolated function. If multiple
516 inputs are arrays, they must be broadcastable to the same shape.
517 Scalar inputs will be broadcast to match array inputs.
519 Returns
520 -------
521 fwxyz : np.array or float
522 The interpolated function evaluated at w,x,y,z: fwxyz = f(w,x,y,z),
523 with the same shape as w, x, y, and z.
524 """
525 return _broadcast_eval(self._evaluate, w, x, y, z)
527 def derivativeW(self, w, x, y, z):
528 """
529 Evaluates the partial derivative with respect to w (the first argument)
530 of the interpolated function at the given input.
532 Parameters
533 ----------
534 w : np.array or float
535 Real values to be evaluated in the interpolated function.
536 x : np.array or float
537 Real values to be evaluated in the interpolated function. If multiple
538 inputs are arrays, they must be broadcastable to the same shape.
539 Scalar inputs will be broadcast to match array inputs.
540 y : np.array or float
541 Real values to be evaluated in the interpolated function. If multiple
542 inputs are arrays, they must be broadcastable to the same shape.
543 Scalar inputs will be broadcast to match array inputs.
544 z : np.array or float
545 Real values to be evaluated in the interpolated function. If multiple
546 inputs are arrays, they must be broadcastable to the same shape.
547 Scalar inputs will be broadcast to match array inputs.
549 Returns
550 -------
551 dfdw : np.array or float
552 The derivative with respect to w of the interpolated function eval-
553 uated at w,x,y,z: dfdw = f_w(w,x,y,z), with the same shape as inputs.
554 """
555 return _broadcast_eval(self._derW, w, x, y, z)
557 def derivativeX(self, w, x, y, z):
558 """
559 Evaluates the partial derivative with respect to x (the second argument)
560 of the interpolated function at the given input.
562 Parameters
563 ----------
564 w : np.array or float
565 Real values to be evaluated in the interpolated function.
566 x : np.array or float
567 Real values to be evaluated in the interpolated function. If multiple
568 inputs are arrays, they must be broadcastable to the same shape.
569 Scalar inputs will be broadcast to match array inputs.
570 y : np.array or float
571 Real values to be evaluated in the interpolated function. If multiple
572 inputs are arrays, they must be broadcastable to the same shape.
573 Scalar inputs will be broadcast to match array inputs.
574 z : np.array or float
575 Real values to be evaluated in the interpolated function. If multiple
576 inputs are arrays, they must be broadcastable to the same shape.
577 Scalar inputs will be broadcast to match array inputs.
579 Returns
580 -------
581 dfdx : np.array or float
582 The derivative with respect to x of the interpolated function eval-
583 uated at w,x,y,z: dfdx = f_x(w,x,y,z), with the same shape as inputs.
584 """
585 return _broadcast_eval(self._derX, w, x, y, z)
587 def derivativeY(self, w, x, y, z):
588 """
589 Evaluates the partial derivative with respect to y (the third argument)
590 of the interpolated function at the given input.
592 Parameters
593 ----------
594 w : np.array or float
595 Real values to be evaluated in the interpolated function.
596 x : np.array or float
597 Real values to be evaluated in the interpolated function. If multiple
598 inputs are arrays, they must be broadcastable to the same shape.
599 Scalar inputs will be broadcast to match array inputs.
600 y : np.array or float
601 Real values to be evaluated in the interpolated function. If multiple
602 inputs are arrays, they must be broadcastable to the same shape.
603 Scalar inputs will be broadcast to match array inputs.
604 z : np.array or float
605 Real values to be evaluated in the interpolated function. If multiple
606 inputs are arrays, they must be broadcastable to the same shape.
607 Scalar inputs will be broadcast to match array inputs.
609 Returns
610 -------
611 dfdy : np.array or float
612 The derivative with respect to y of the interpolated function eval-
613 uated at w,x,y,z: dfdy = f_y(w,x,y,z), with the same shape as inputs.
614 """
615 return _broadcast_eval(self._derY, w, x, y, z)
617 def derivativeZ(self, w, x, y, z):
618 """
619 Evaluates the partial derivative with respect to z (the fourth argument)
620 of the interpolated function at the given input.
622 Parameters
623 ----------
624 w : np.array or float
625 Real values to be evaluated in the interpolated function.
626 x : np.array or float
627 Real values to be evaluated in the interpolated function. If multiple
628 inputs are arrays, they must be broadcastable to the same shape.
629 Scalar inputs will be broadcast to match array inputs.
630 y : np.array or float
631 Real values to be evaluated in the interpolated function. If multiple
632 inputs are arrays, they must be broadcastable to the same shape.
633 Scalar inputs will be broadcast to match array inputs.
634 z : np.array or float
635 Real values to be evaluated in the interpolated function. If multiple
636 inputs are arrays, they must be broadcastable to the same shape.
637 Scalar inputs will be broadcast to match array inputs.
639 Returns
640 -------
641 dfdz : np.array or float
642 The derivative with respect to z of the interpolated function eval-
643 uated at w,x,y,z: dfdz = f_z(w,x,y,z), with the same shape as inputs.
644 """
645 return _broadcast_eval(self._derZ, w, x, y, z)
647 def _evaluate(self, w, x, y, z):
648 """
649 Interpolated function evaluator, to be defined in subclasses.
650 """
651 raise NotImplementedError()
653 def _derW(self, w, x, y, z):
654 """
655 Default or fallback derivative with respect to w, using finite difference approximation.
656 Subclasses of HARKinterpolator4D should define their own more specific method.
657 """
658 eps = 1e-8
659 f0 = self.__call__(w, x, y, z)
660 f1 = self.__call__(w + eps, x, y, z)
661 dfdw = (f1 - f0) / eps
662 return dfdw
664 def _derX(self, w, x, y, z):
665 """
666 Default or fallback derivative with respect to x, using finite difference approximation.
667 Subclasses of HARKinterpolator4D should define their own more specific method.
668 """
669 eps = 1e-8
670 f0 = self.__call__(w, x, y, z)
671 f1 = self.__call__(w, x + eps, y, z)
672 dfdx = (f1 - f0) / eps
673 return dfdx
675 def _derY(self, w, x, y, z):
676 """
677 Default or fallback derivative with respect to y, using finite difference approximation.
678 Subclasses of HARKinterpolator4D should define their own more specific method.
679 """
680 eps = 1e-8
681 f0 = self.__call__(w, x, y, z)
682 f1 = self.__call__(w, x, y + eps, z)
683 dfdy = (f1 - f0) / eps
684 return dfdy
686 def _derZ(self, w, x, y, z):
687 """
688 Default or fallback derivative with respect to z, using finite difference approximation.
689 Subclasses of HARKinterpolator4D should define their own more specific method.
690 """
691 eps = 1e-8
692 f0 = self.__call__(w, x, y, z)
693 f1 = self.__call__(w, x, y, z + eps)
694 dfdz = (f1 - f0) / eps
695 return dfdz
698class IdentityFunction(MetricObject):
699 """
700 A fairly trivial interpolator that simply returns one of its arguments. Useful for avoiding
701 numeric error in extreme cases.
703 Parameters
704 ----------
705 i_dim : int
706 Index of the dimension on which the identity is defined. ``f(*x) = x[i]``
707 n_dims : int
708 Total number of input dimensions for this function.
709 """
711 distance_criteria = ["i_dim"]
713 def __init__(self, i_dim=0, n_dims=1):
714 self.i_dim = i_dim
715 self.n_dims = n_dims
717 def __call__(self, *args):
718 """
719 Evaluate the identity function.
720 """
721 return args[self.i_dim]
723 def derivative(self, *args):
724 """
725 Returns the derivative of the function with respect to the first dimension.
726 """
727 if self.i_dim == 0:
728 return np.ones_like(args[0])
729 else:
730 return np.zeros_like(args[0])
732 def derivativeX(self, *args):
733 """
734 Returns the derivative of the function with respect to the X dimension.
735 This is the first input whenever n_dims < 4 and the second input otherwise.
736 """
737 if self.n_dims >= 4:
738 j = 1
739 else:
740 j = 0
741 if self.i_dim == j:
742 return np.ones_like(args[0])
743 else:
744 return np.zeros_like(args[0])
746 def derivativeY(self, *args):
747 """
748 Returns the derivative of the function with respect to the Y dimension.
749 This is the second input whenever n_dims < 4 and the third input otherwise.
750 """
751 if self.n_dims >= 4:
752 j = 2
753 else:
754 j = 1
755 if self.i_dim == j:
756 return np.ones_like(args[0])
757 else:
758 return np.zeros_like(args[0])
760 def derivativeZ(self, *args):
761 """
762 Returns the derivative of the function with respect to the Z dimension.
763 This is the third input whenever n_dims < 4 and the fourth input otherwise.
764 """
765 if self.n_dims >= 4:
766 j = 3
767 else:
768 j = 2
769 if self.i_dim == j:
770 return np.ones_like(args[0])
771 else:
772 return np.zeros_like(args[0])
774 def derivativeW(self, *args):
775 """
776 Returns the derivative of the function with respect to the W dimension.
777 This should only exist when n_dims >= 4.
778 """
779 if self.n_dims < 4:
780 raise RuntimeError(
781 "Derivative with respect to W can't be called when n_dims < 4!"
782 )
783 j = 0
784 if self.i_dim == j:
785 return np.ones_like(args[0])
786 else:
787 return np.zeros_like(args[0])
790class ConstantFunction(MetricObject):
791 """
792 A class for representing trivial functions that return the same real output for any input. This
793 is convenient for models where an object might be a (non-trivial) function, but in some variations
794 that object is just a constant number. Rather than needing to make a (Bi/Tri/Quad)-
795 LinearInterpolation with trivial state grids and the same f_value in every entry, ConstantFunction
796 allows the user to quickly make a constant/trivial function. This comes up, e.g., in models
797 with endogenous pricing of insurance contracts; a contract's premium might depend on some state
798 variables of the individual, but in some variations the premium of a contract is just a number.
800 Parameters
801 ----------
802 value : float
803 The constant value that the function returns.
804 """
806 distance_criteria = ["value"]
808 def __init__(self, value):
809 self.value = float(value)
811 def __call__(self, *args):
812 """
813 Evaluate the constant function. The first input must exist and should be an array.
814 Returns an array of identical shape to args[0] (if it exists).
815 """
816 if (
817 len(args) > 0
818 ): # If there is at least one argument, return appropriately sized array
819 if _isscalar(args[0]):
820 return self.value
821 else:
822 shape = args[0].shape
823 return self.value * np.ones(shape)
824 else: # Otherwise, return a single instance of the constant value
825 return self.value
827 def _der(self, *args):
828 """
829 Evaluate the derivative of the function. The first input must exist and should be an array.
830 Returns an array of identical shape to args[0] (if it exists). This is an array of zeros.
831 """
832 if len(args) > 0:
833 if _isscalar(args[0]):
834 return 0.0
835 else:
836 shape = args[0].shape
837 return np.zeros(shape)
838 else:
839 return 0.0
841 def eval_with_derivative(self, x):
842 val = self(x)
843 der = self._der(x)
844 return val, der
846 # All other derivatives are also zero everywhere, so these methods just point to derivative
847 derivative = _der
848 derivativeX = derivative
849 derivativeY = derivative
850 derivativeZ = derivative
851 derivativeW = derivative
852 derivativeXX = derivative
855class LinearInterp(HARKinterpolator1D):
856 """
857 A "from scratch" 1D linear interpolation class. Allows for linear or decay
858 extrapolation (approaching a limiting linear function from below).
860 NOTE: When no input is given for the limiting linear function, linear
861 extrapolation is used above the highest gridpoint.
863 Parameters
864 ----------
865 x_list : np.array
866 List of x values composing the grid.
867 y_list : np.array
868 List of y values, representing f(x) at the points in x_list.
869 intercept_limit : float
870 Intercept of limiting linear function.
871 slope_limit : float
872 Slope of limiting linear function.
873 lower_extrap : bool
874 Indicator for whether lower extrapolation is allowed. False means
875 f(x) = NaN for x < min(x_list); True means linear extrapolation.
876 pre_compute : bool
877 Indicator for whether interpolation coefficients should be pre-computed
878 and stored as attributes of self (default False). More memory will be used,
879 and instantiation will take slightly longer, but later evaluation will
880 be faster due to less arithmetic.
881 indexer : function or None (default)
882 If provided, a custom function that identifies the index of the interpolant
883 segment for each query point. Should return results identically to the
884 default behavior of np.maximum(np.searchsorted(self.x_list[:-1], x), 1).
885 WARNING: User is responsible for verifying that their custom indexer is
886 actually correct versus default behavior.
887 """
889 distance_criteria = ["x_list", "y_list"]
891 def __init__(
892 self,
893 x_list,
894 y_list,
895 intercept_limit=None,
896 slope_limit=None,
897 lower_extrap=False,
898 pre_compute=False,
899 indexer=None,
900 ):
901 # Make the basic linear spline interpolation
902 self.x_list = _coerce_1d_grid(x_list)
903 self.y_list = _coerce_1d_grid(y_list)
904 _check_grid_dimensions(1, self.y_list, self.x_list)
905 self.lower_extrap = lower_extrap
906 self.x_n = self.x_list.size
907 self.indexer = indexer
909 # Make a decay extrapolation
910 if intercept_limit is not None and slope_limit is not None:
911 slope_at_top = (y_list[-1] - y_list[-2]) / (x_list[-1] - x_list[-2])
912 level_diff = intercept_limit + slope_limit * x_list[-1] - y_list[-1]
913 slope_diff = slope_limit - slope_at_top
914 # If the model that can handle uncertainty has been calibrated with
915 # with uncertainty set to zero, the 'extrapolation' will blow up
916 # Guard against that and nearby problems by testing slope equality
917 if not np.isclose(slope_limit, slope_at_top, atol=1e-15):
918 self.decay_extrap_A = level_diff
919 self.decay_extrap_B = -slope_diff / level_diff
920 self.intercept_limit = intercept_limit
921 self.slope_limit = slope_limit
922 self.decay_extrap = True
923 else:
924 self.decay_extrap = False
925 else:
926 self.decay_extrap = False
928 # Calculate interpolation coefficients now rather than at evaluation time
929 if pre_compute:
930 self.slopes = (self.y_list[1:] - self.y_list[:-1]) / (
931 self.x_list[1:] - self.x_list[:-1]
932 )
933 self.intercepts = self.y_list[:-1] - self.slopes * self.x_list[:-1]
935 def _segment_index(self, x):
936 """Return the bracketing right-endpoint index for each query in ``x``."""
937 if self.indexer is None:
938 return np.maximum(np.searchsorted(self.x_list[:-1], x), 1)
939 return self.indexer(x)
941 def _segment_values(self, x, i, want_y, want_dydx):
942 """Compute ``(y, dydx)`` on the linear segment to the right of ``i - 1``.
944 Skipped outputs return ``None``. Returned arrays are fresh allocations
945 safe for in-place patching by ``_apply_lower_bound`` /
946 ``_apply_upper_decay``: in the pre-computed branch ``self.slopes[j]``
947 is itself a fancy-index copy, so mutating ``dydx`` does not touch
948 ``self.slopes``.
949 """
950 if hasattr(self, "slopes"):
951 j = i - 1
952 slopes_j = self.slopes[j]
953 y = self.intercepts[j] + slopes_j * x if want_y else None
954 dydx = slopes_j if want_dydx else None
955 return y, dydx
956 x_lo = self.x_list[i - 1]
957 x_hi = self.x_list[i]
958 y_lo = self.y_list[i - 1]
959 y_hi = self.y_list[i]
960 if want_y:
961 alpha = (x - x_lo) / (x_hi - x_lo)
962 y = (1.0 - alpha) * y_lo + alpha * y_hi
963 else:
964 y = None
965 dydx = (y_hi - y_lo) / (x_hi - x_lo) if want_dydx else None
966 return y, dydx
968 def _apply_lower_bound(self, x, y, dydx):
969 """In-place: mark queries below ``x_list[0]`` as NaN. ``y`` and ``dydx``
970 may each be ``None`` to skip; no-op when ``self.lower_extrap`` is True."""
971 if self.lower_extrap or (y is None and dydx is None):
972 return
973 below = x < self.x_list[0]
974 if y is not None:
975 y[below] = np.nan
976 if dydx is not None:
977 dydx[below] = np.nan
979 def _apply_upper_decay(self, x, y, dydx):
980 """In-place: replace queries above ``x_list[-1]`` with the limiting linear
981 plus exponential-decay envelope. ``y`` and ``dydx`` may each be ``None``
982 to skip; no-op when ``self.decay_extrap`` is False."""
983 if not self.decay_extrap or (y is None and dydx is None):
984 return
985 above = x > self.x_list[-1]
986 if not np.any(above):
987 return
988 x_temp = x[above] - self.x_list[-1]
989 decay = self.decay_extrap_A * np.exp(-self.decay_extrap_B * x_temp)
990 if y is not None:
991 y[above] = self.intercept_limit + self.slope_limit * x[above] - decay
992 if dydx is not None:
993 dydx[above] = self.slope_limit + self.decay_extrap_B * decay
995 def _evalOrDer(self, x, _eval, _Der):
996 """
997 Returns the level and/or first derivative of the function at each value in
998 x. Only called internally by HARKinterpolator1D.eval_and_der (etc).
1000 Parameters
1001 ----------
1002 x : scalar or np.array
1003 Set of points where we want to evaluate the interpolated function and/or its derivative.
1004 _eval : boolean
1005 Indicator for whether to evaluate the level of the interpolated function.
1006 _Der : boolean
1007 Indicator for whether to evaluate the derivative of the interpolated function.
1009 Returns
1010 -------
1011 A list including the level and/or derivative of the interpolated function where requested.
1012 """
1013 i = self._segment_index(x)
1014 y, dydx = self._segment_values(x, i, want_y=_eval, want_dydx=_Der)
1015 self._apply_lower_bound(x, y, dydx)
1016 self._apply_upper_decay(x, y, dydx)
1017 output = []
1018 if _eval:
1019 output.append(y)
1020 if _Der:
1021 output.append(dydx)
1022 return output
1024 def _evaluate(self, x, return_indices=False):
1025 """
1026 Returns the level of the interpolated function at each value in x. Only
1027 called internally by HARKinterpolator1D.__call__ (etc).
1028 """
1029 return self._evalOrDer(x, True, False)[0]
1031 def _der(self, x):
1032 """
1033 Returns the first derivative of the interpolated function at each value
1034 in x. Only called internally by HARKinterpolator1D.derivative (etc).
1035 """
1036 return self._evalOrDer(x, False, True)[0]
1038 def _evalAndDer(self, x):
1039 """
1040 Returns the level and first derivative of the function at each value in
1041 x. Only called internally by HARKinterpolator1D.eval_and_der (etc).
1042 """
1043 y, dydx = self._evalOrDer(x, True, True)
1045 return y, dydx
1048class CubicInterp(HARKinterpolator1D):
1049 """
1050 An interpolating function using piecewise cubic splines. Matches level and
1051 slope of 1D function at gridpoints, smoothly interpolating in between.
1052 Extrapolation above highest gridpoint approaches a limiting linear function
1053 if desired (linear extrapolation also enabled.)
1055 NOTE: When no input is given for the limiting linear function, linear
1056 extrapolation is used above the highest gridpoint.
1058 Parameters
1059 ----------
1060 x_list : np.array
1061 List of x values composing the grid.
1062 y_list : np.array
1063 List of y values, representing f(x) at the points in x_list.
1064 dydx_list : np.array
1065 List of dydx values, representing f'(x) at the points in x_list
1066 intercept_limit : float
1067 Intercept of limiting linear function.
1068 slope_limit : float
1069 Slope of limiting linear function.
1070 lower_extrap : boolean
1071 Indicator for whether lower extrapolation is allowed. False means
1072 f(x) = NaN for x < min(x_list); True means linear extrapolation.
1073 """
1075 distance_criteria = ["x_list", "y_list", "dydx_list"]
1077 def __init__(
1078 self,
1079 x_list,
1080 y_list,
1081 dydx_list,
1082 intercept_limit=None,
1083 slope_limit=None,
1084 lower_extrap=False,
1085 ):
1086 self._init_cubic_grids(x_list, y_list, dydx_list)
1088 # Define lower extrapolation as linear function (or just NaN)
1089 if lower_extrap:
1090 lower_row = [y_list[0], dydx_list[0], 0.0, 0.0]
1091 else:
1092 lower_row = [np.nan, np.nan, np.nan, np.nan]
1094 # Per-segment cubic coefficients on segments mapped to [0,1] (vectorized)
1095 xL = self.x_list[:-1]
1096 xR = self.x_list[1:]
1097 yL = self.y_list[:-1]
1098 yR = self.y_list[1:]
1099 Span = xR - xL
1100 dydxL = self.dydx_list[:-1] * Span
1101 dydxR = self.dydx_list[1:] * Span
1102 seg = np.column_stack(
1103 [
1104 yL,
1105 dydxL,
1106 3 * (yR - yL) - 2 * dydxL - dydxR,
1107 2 * (yL - yR) + dydxL + dydxR,
1108 ]
1109 )
1111 # Calculate extrapolation coefficients as a decay toward limiting function y = mx+b
1112 x_top = self.x_list[-1]
1113 y_top = self.y_list[-1]
1114 if slope_limit is None and intercept_limit is None:
1115 slope_limit = self.dydx_list[-1]
1116 intercept_limit = y_top - slope_limit * x_top
1117 gap = slope_limit * x_top + intercept_limit - y_top
1118 slope = slope_limit - self.dydx_list[-1]
1119 if (gap != 0) and (slope <= 0):
1120 upper_row = [intercept_limit, slope_limit, gap, slope / gap]
1121 elif slope > 0:
1122 # fixing a problem when slope is positive
1123 upper_row = [intercept_limit, slope_limit, 0, 0]
1124 else:
1125 upper_row = [intercept_limit, slope_limit, gap, 0]
1127 self.coeffs = np.vstack([lower_row, seg, upper_row])
1129 def _classify_segments(self, x):
1130 """Bucket ``x`` into below-grid, above-grid, and in-bounds positions and
1131 precompute in-bounds coefficient slices and the local segment ``alpha``.
1132 Returns ``(m, out_bot, out_top, in_bnds, i, coeffs_in, alpha)``."""
1133 m = len(x)
1134 pos = np.searchsorted(self.x_list, x, side="right")
1135 out_bot = pos == 0
1136 out_top = pos == self.n
1137 in_bnds = np.logical_not(np.logical_or(out_bot, out_top))
1138 i = pos[in_bnds]
1139 coeffs_in = self.coeffs[i, :]
1140 alpha = (x[in_bnds] - self.x_list[i - 1]) / (
1141 self.x_list[i] - self.x_list[i - 1]
1142 )
1143 return m, out_bot, out_top, in_bnds, i, coeffs_in, alpha
1145 def _eval_y_outbounds(self, y, out_bot, out_top, x):
1146 """Apply lower/upper extrapolation values to ``y`` at out-of-bounds points."""
1147 y[out_bot] = self.coeffs[0, 0] + self.coeffs[0, 1] * (
1148 x[out_bot] - self.x_list[0]
1149 )
1150 alpha_top = x[out_top] - self.x_list[self.n - 1]
1151 y[out_top] = (
1152 self.coeffs[self.n, 0]
1153 + x[out_top] * self.coeffs[self.n, 1]
1154 - self.coeffs[self.n, 2] * np.exp(alpha_top * self.coeffs[self.n, 3])
1155 )
1156 return alpha_top
1158 def _eval_dydx_outbounds(self, dydx, out_bot, out_top, alpha_top):
1159 """Apply lower/upper extrapolation derivatives to ``dydx``."""
1160 dydx[out_bot] = self.coeffs[0, 1]
1161 dydx[out_top] = self.coeffs[self.n, 1] - self.coeffs[self.n, 2] * self.coeffs[
1162 self.n, 3
1163 ] * np.exp(alpha_top * self.coeffs[self.n, 3])
1165 def _evaluate(self, x):
1166 """
1167 Returns the level of the interpolated function at each value in x. Only
1168 called internally by HARKinterpolator1D.__call__ (etc).
1169 """
1170 m, out_bot, out_top, in_bnds, _i, coeffs_in, alpha = self._classify_segments(x)
1171 y = np.zeros(m)
1172 if y.size > 0:
1173 y[in_bnds] = coeffs_in[:, 0] + alpha * (
1174 coeffs_in[:, 1] + alpha * (coeffs_in[:, 2] + alpha * coeffs_in[:, 3])
1175 )
1176 self._eval_y_outbounds(y, out_bot, out_top, x)
1177 return y
1179 def _der(self, x):
1180 """
1181 Returns the first derivative of the interpolated function at each value
1182 in x. Only called internally by HARKinterpolator1D.derivative (etc).
1183 """
1184 m, out_bot, out_top, in_bnds, i, coeffs_in, alpha = self._classify_segments(x)
1185 dydx = np.zeros(m)
1186 if dydx.size > 0:
1187 dydx[in_bnds] = (
1188 coeffs_in[:, 1]
1189 + alpha * (2 * coeffs_in[:, 2] + alpha * 3 * coeffs_in[:, 3])
1190 ) / (self.x_list[i] - self.x_list[i - 1])
1191 alpha_top = x[out_top] - self.x_list[self.n - 1]
1192 self._eval_dydx_outbounds(dydx, out_bot, out_top, alpha_top)
1193 return dydx
1195 def _evalAndDer(self, x):
1196 """
1197 Returns the level and first derivative of the function at each value in
1198 x. Only called internally by HARKinterpolator1D.eval_and_der (etc).
1199 """
1200 m, out_bot, out_top, in_bnds, i, coeffs_in, alpha = self._classify_segments(x)
1201 y = np.zeros(m)
1202 dydx = np.zeros(m)
1203 if y.size > 0:
1204 y[in_bnds] = coeffs_in[:, 0] + alpha * (
1205 coeffs_in[:, 1] + alpha * (coeffs_in[:, 2] + alpha * coeffs_in[:, 3])
1206 )
1207 dydx[in_bnds] = (
1208 coeffs_in[:, 1]
1209 + alpha * (2 * coeffs_in[:, 2] + alpha * 3 * coeffs_in[:, 3])
1210 ) / (self.x_list[i] - self.x_list[i - 1])
1211 alpha_top = self._eval_y_outbounds(y, out_bot, out_top, x)
1212 self._eval_dydx_outbounds(dydx, out_bot, out_top, alpha_top)
1213 return y, dydx
1216class CubicHermiteInterp(HARKinterpolator1D):
1217 """
1218 An interpolating function using piecewise cubic splines. Matches level and
1219 slope of 1D function at gridpoints, smoothly interpolating in between.
1220 Extrapolation above highest gridpoint approaches a limiting linear function
1221 if desired (linear extrapolation also enabled.)
1223 NOTE: When no input is given for the limiting linear function, linear
1224 extrapolation is used above the highest gridpoint.
1226 Parameters
1227 ----------
1228 x_list : np.array
1229 List of x values composing the grid.
1230 y_list : np.array
1231 List of y values, representing f(x) at the points in x_list.
1232 dydx_list : np.array
1233 List of dydx values, representing f'(x) at the points in x_list
1234 intercept_limit : float
1235 Intercept of limiting linear function.
1236 slope_limit : float
1237 Slope of limiting linear function.
1238 lower_extrap : boolean
1239 Indicator for whether lower extrapolation is allowed. False means
1240 f(x) = NaN for x < min(x_list); True means linear extrapolation.
1241 """
1243 distance_criteria = ["x_list", "y_list", "dydx_list"]
1245 def __init__(
1246 self,
1247 x_list,
1248 y_list,
1249 dydx_list,
1250 intercept_limit=None,
1251 slope_limit=None,
1252 lower_extrap=False,
1253 ):
1254 self._init_cubic_grids(x_list, y_list, dydx_list)
1256 self._chs = CubicHermiteSpline(
1257 self.x_list, self.y_list, self.dydx_list, extrapolate=None
1258 )
1259 self.coeffs = np.flip(self._chs.c.T, 1)
1261 # Define lower extrapolation as linear function (or just NaN)
1262 if lower_extrap:
1263 temp = np.array([self.y_list[0], self.dydx_list[0], 0, 0])
1264 else:
1265 temp = np.array([np.nan, np.nan, np.nan, np.nan])
1267 self.coeffs = np.vstack((temp, self.coeffs))
1269 x1 = self.x_list[self.n - 1]
1270 y1 = self.y_list[self.n - 1]
1272 # Calculate extrapolation coefficients as a decay toward limiting function y = mx+b
1273 if slope_limit is None and intercept_limit is None:
1274 slope_limit = self.dydx_list[-1]
1275 intercept_limit = self.y_list[-1] - slope_limit * self.x_list[-1]
1276 gap = slope_limit * x1 + intercept_limit - y1
1277 slope = slope_limit - self.dydx_list[self.n - 1]
1278 if (gap != 0) and (slope <= 0):
1279 temp = np.array([intercept_limit, slope_limit, gap, slope / gap])
1280 elif slope > 0:
1281 # fixing a problem when slope is positive
1282 temp = np.array([intercept_limit, slope_limit, 0, 0])
1283 else:
1284 temp = np.array([intercept_limit, slope_limit, gap, 0])
1285 self.coeffs = np.vstack((self.coeffs, temp))
1287 def out_of_bounds(self, x):
1288 out_bot = x < self.x_list[0]
1289 out_top = x > self.x_list[-1]
1290 return out_bot, out_top
1292 def _evaluate(self, x):
1293 """
1294 Returns the level of the interpolated function at each value in x. Only
1295 called internally by HARKinterpolator1D.__call__ (etc).
1296 """
1297 out_bot, out_top = self.out_of_bounds(x)
1299 return self._eval_helper(x, out_bot, out_top)
1301 def _eval_helper(self, x, out_bot, out_top):
1302 y = self._chs(x)
1304 # Do the "out of bounds" evaluation points
1305 if any(out_bot):
1306 y[out_bot] = self.coeffs[0, 0] + self.coeffs[0, 1] * (
1307 x[out_bot] - self.x_list[0]
1308 )
1310 if any(out_top):
1311 alpha = x[out_top] - self.x_list[self.n - 1]
1312 y[out_top] = (
1313 self.coeffs[self.n, 0]
1314 + x[out_top] * self.coeffs[self.n, 1]
1315 - self.coeffs[self.n, 2] * np.exp(alpha * self.coeffs[self.n, 3])
1316 )
1318 return y
1320 def _der(self, x):
1321 """
1322 Returns the first derivative of the interpolated function at each value
1323 in x. Only called internally by HARKinterpolator1D.derivative (etc).
1324 """
1325 out_bot, out_top = self.out_of_bounds(x)
1327 return self._der_helper(x, out_bot, out_top)
1329 def _der_helper(self, x, out_bot, out_top):
1330 dydx = self._chs(x, nu=1)
1332 # Do the "out of bounds" evaluation points
1333 if any(out_bot):
1334 dydx[out_bot] = self.coeffs[0, 1]
1335 if any(out_top):
1336 alpha = x[out_top] - self.x_list[self.n - 1]
1337 dydx[out_top] = self.coeffs[self.n, 1] - self.coeffs[
1338 self.n, 2
1339 ] * self.coeffs[self.n, 3] * np.exp(alpha * self.coeffs[self.n, 3])
1341 return dydx
1343 def _evalAndDer(self, x):
1344 """
1345 Returns the level and first derivative of the function at each value in
1346 x. Only called internally by HARKinterpolator1D.eval_and_der (etc).
1347 """
1348 out_bot, out_top = self.out_of_bounds(x)
1349 y = self._eval_helper(x, out_bot, out_top)
1350 dydx = self._der_helper(x, out_bot, out_top)
1351 return y, dydx
1353 def der_interp(self, nu=1):
1354 """
1355 Construct a new piecewise polynomial representing the derivative.
1356 See `scipy` for additional documentation:
1357 https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.CubicHermiteSpline.html
1358 """
1359 return self._chs.derivative(nu)
1361 def antider_interp(self, nu=1):
1362 """
1363 Construct a new piecewise polynomial representing the antiderivative.
1365 Antiderivative is also the indefinite integral of the function,
1366 and derivative is its inverse operation.
1368 See `scipy` for additional documentation:
1369 https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.CubicHermiteSpline.html
1370 """
1371 return self._chs.antiderivative(nu)
1373 def integrate(self, a, b, extrapolate=False):
1374 """
1375 Compute a definite integral over a piecewise polynomial.
1377 See `scipy` for additional documentation:
1378 https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.CubicHermiteSpline.html
1379 """
1380 return self._chs.integrate(a, b, extrapolate)
1382 def roots(self, discontinuity=True, extrapolate=False):
1383 """
1384 Find real roots of the the piecewise polynomial.
1386 See `scipy` for additional documentation:
1387 https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.CubicHermiteSpline.html
1388 """
1389 return self._chs.roots(discontinuity, extrapolate)
1391 def solve(self, y=0.0, discontinuity=True, extrapolate=False):
1392 """
1393 Find real solutions of the the equation ``pp(x) == y``.
1395 See `scipy` for additional documentation:
1396 https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.CubicHermiteSpline.html
1397 """
1398 return self._chs.solve(y, discontinuity, extrapolate)
1401class BilinearInterp(HARKinterpolator2D):
1402 """
1403 Bilinear full (or tensor) grid interpolation of a function f(x,y).
1405 Parameters
1406 ----------
1407 f_values : numpy.array
1408 An array of size (x_n,y_n) such that f_values[i,j] = f(x_list[i],y_list[j])
1409 x_list : numpy.array
1410 An array of x values, with length designated x_n.
1411 y_list : numpy.array
1412 An array of y values, with length designated y_n.
1413 """
1415 distance_criteria = ["x_list", "y_list", "f_values"]
1417 def __init__(self, f_values, x_list, y_list):
1418 self.f_values = f_values
1419 self.x_list = _coerce_1d_grid(x_list)
1420 self.y_list = _coerce_1d_grid(y_list)
1421 _check_grid_dimensions(2, self.f_values, self.x_list, self.y_list)
1422 self.x_n = self.x_list.size
1423 self.y_n = self.y_list.size
1425 def _locate_xy_indices(self, x, y):
1426 """Return clamped search indices for ``x`` and ``y`` shared by ``_evaluate``,
1427 ``_derX``, and ``_derY``."""
1428 return (
1429 _locate_clipped(self.x_list, x, self.x_n),
1430 _locate_clipped(self.y_list, y, self.y_n),
1431 )
1433 def _evaluate(self, x, y):
1434 """
1435 Returns the level of the interpolated function at each value in x,y.
1436 Only called internally by HARKinterpolator2D.__call__ (etc).
1437 """
1438 x_pos, y_pos = self._locate_xy_indices(x, y)
1439 alpha = _cell_fraction(self.x_list, x_pos, x)
1440 beta = _cell_fraction(self.y_list, y_pos, y)
1441 f = (
1442 (1 - alpha) * (1 - beta) * self.f_values[x_pos - 1, y_pos - 1]
1443 + (1 - alpha) * beta * self.f_values[x_pos - 1, y_pos]
1444 + alpha * (1 - beta) * self.f_values[x_pos, y_pos - 1]
1445 + alpha * beta * self.f_values[x_pos, y_pos]
1446 )
1447 return f
1449 def _derX(self, x, y):
1450 """
1451 Returns the derivative with respect to x of the interpolated function
1452 at each value in x,y. Only called internally by HARKinterpolator2D.derivativeX.
1453 """
1454 x_pos, y_pos = self._locate_xy_indices(x, y)
1455 beta = _cell_fraction(self.y_list, y_pos, y)
1456 dfdx = (
1457 (
1458 (1 - beta) * self.f_values[x_pos, y_pos - 1]
1459 + beta * self.f_values[x_pos, y_pos]
1460 )
1461 - (
1462 (1 - beta) * self.f_values[x_pos - 1, y_pos - 1]
1463 + beta * self.f_values[x_pos - 1, y_pos]
1464 )
1465 ) / (self.x_list[x_pos] - self.x_list[x_pos - 1])
1466 return dfdx
1468 def _derY(self, x, y):
1469 """
1470 Returns the derivative with respect to y of the interpolated function
1471 at each value in x,y. Only called internally by HARKinterpolator2D.derivativeY.
1472 """
1473 x_pos, y_pos = self._locate_xy_indices(x, y)
1474 alpha = _cell_fraction(self.x_list, x_pos, x)
1475 dfdy = (
1476 (
1477 (1 - alpha) * self.f_values[x_pos - 1, y_pos]
1478 + alpha * self.f_values[x_pos, y_pos]
1479 )
1480 - (
1481 (1 - alpha) * self.f_values[x_pos - 1, y_pos - 1]
1482 + alpha * self.f_values[x_pos, y_pos - 1]
1483 )
1484 ) / (self.y_list[y_pos] - self.y_list[y_pos - 1])
1485 return dfdy
1488class TrilinearInterp(HARKinterpolator3D):
1489 """
1490 Trilinear full (or tensor) grid interpolation of a function f(x,y,z).
1492 Parameters
1493 ----------
1494 f_values : numpy.array
1495 An array of size (x_n,y_n,z_n) such that f_values[i,j,k] =
1496 f(x_list[i],y_list[j],z_list[k])
1497 x_list : numpy.array
1498 An array of x values, with length designated x_n.
1499 y_list : numpy.array
1500 An array of y values, with length designated y_n.
1501 z_list : numpy.array
1502 An array of z values, with length designated z_n.
1503 """
1505 distance_criteria = ["f_values", "x_list", "y_list", "z_list"]
1507 def __init__(self, f_values, x_list, y_list, z_list):
1508 self.f_values = f_values
1509 self.x_list = _coerce_1d_grid(x_list)
1510 self.y_list = _coerce_1d_grid(y_list)
1511 self.z_list = _coerce_1d_grid(z_list)
1512 _check_grid_dimensions(3, self.f_values, self.x_list, self.y_list, self.z_list)
1513 self.x_n = self.x_list.size
1514 self.y_n = self.y_list.size
1515 self.z_n = self.z_list.size
1517 def _locate_xyz_indices(self, x, y, z):
1518 """Return clamped search indices for ``x``, ``y``, ``z`` shared by
1519 ``_evaluate`` and the three derivative methods."""
1520 return (
1521 _locate_clipped(self.x_list, x, self.x_n),
1522 _locate_clipped(self.y_list, y, self.y_n),
1523 _locate_clipped(self.z_list, z, self.z_n),
1524 )
1526 def _evaluate(self, x, y, z):
1527 """
1528 Returns the level of the interpolated function at each value in x,y,z.
1529 Only called internally by HARKinterpolator3D.__call__ (etc).
1530 """
1531 x_pos, y_pos, z_pos = self._locate_xyz_indices(x, y, z)
1532 alpha = _cell_fraction(self.x_list, x_pos, x)
1533 beta = _cell_fraction(self.y_list, y_pos, y)
1534 gamma = _cell_fraction(self.z_list, z_pos, z)
1535 f = (
1536 (1 - alpha)
1537 * (1 - beta)
1538 * (1 - gamma)
1539 * self.f_values[x_pos - 1, y_pos - 1, z_pos - 1]
1540 + (1 - alpha)
1541 * (1 - beta)
1542 * gamma
1543 * self.f_values[x_pos - 1, y_pos - 1, z_pos]
1544 + (1 - alpha)
1545 * beta
1546 * (1 - gamma)
1547 * self.f_values[x_pos - 1, y_pos, z_pos - 1]
1548 + (1 - alpha) * beta * gamma * self.f_values[x_pos - 1, y_pos, z_pos]
1549 + alpha
1550 * (1 - beta)
1551 * (1 - gamma)
1552 * self.f_values[x_pos, y_pos - 1, z_pos - 1]
1553 + alpha * (1 - beta) * gamma * self.f_values[x_pos, y_pos - 1, z_pos]
1554 + alpha * beta * (1 - gamma) * self.f_values[x_pos, y_pos, z_pos - 1]
1555 + alpha * beta * gamma * self.f_values[x_pos, y_pos, z_pos]
1556 )
1557 return f
1559 def _derX(self, x, y, z):
1560 """
1561 Returns the derivative with respect to x of the interpolated function
1562 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeX.
1563 """
1564 x_pos, y_pos, z_pos = self._locate_xyz_indices(x, y, z)
1565 beta = _cell_fraction(self.y_list, y_pos, y)
1566 gamma = _cell_fraction(self.z_list, z_pos, z)
1567 dfdx = (
1568 (
1569 (1 - beta) * (1 - gamma) * self.f_values[x_pos, y_pos - 1, z_pos - 1]
1570 + (1 - beta) * gamma * self.f_values[x_pos, y_pos - 1, z_pos]
1571 + beta * (1 - gamma) * self.f_values[x_pos, y_pos, z_pos - 1]
1572 + beta * gamma * self.f_values[x_pos, y_pos, z_pos]
1573 )
1574 - (
1575 (1 - beta)
1576 * (1 - gamma)
1577 * self.f_values[x_pos - 1, y_pos - 1, z_pos - 1]
1578 + (1 - beta) * gamma * self.f_values[x_pos - 1, y_pos - 1, z_pos]
1579 + beta * (1 - gamma) * self.f_values[x_pos - 1, y_pos, z_pos - 1]
1580 + beta * gamma * self.f_values[x_pos - 1, y_pos, z_pos]
1581 )
1582 ) / (self.x_list[x_pos] - self.x_list[x_pos - 1])
1583 return dfdx
1585 def _derY(self, x, y, z):
1586 """
1587 Returns the derivative with respect to y of the interpolated function
1588 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeY.
1589 """
1590 x_pos, y_pos, z_pos = self._locate_xyz_indices(x, y, z)
1591 alpha = _cell_fraction(self.x_list, x_pos, x)
1592 gamma = _cell_fraction(self.z_list, z_pos, z)
1593 dfdy = (
1594 (
1595 (1 - alpha) * (1 - gamma) * self.f_values[x_pos - 1, y_pos, z_pos - 1]
1596 + (1 - alpha) * gamma * self.f_values[x_pos - 1, y_pos, z_pos]
1597 + alpha * (1 - gamma) * self.f_values[x_pos, y_pos, z_pos - 1]
1598 + alpha * gamma * self.f_values[x_pos, y_pos, z_pos]
1599 )
1600 - (
1601 (1 - alpha)
1602 * (1 - gamma)
1603 * self.f_values[x_pos - 1, y_pos - 1, z_pos - 1]
1604 + (1 - alpha) * gamma * self.f_values[x_pos - 1, y_pos - 1, z_pos]
1605 + alpha * (1 - gamma) * self.f_values[x_pos, y_pos - 1, z_pos - 1]
1606 + alpha * gamma * self.f_values[x_pos, y_pos - 1, z_pos]
1607 )
1608 ) / (self.y_list[y_pos] - self.y_list[y_pos - 1])
1609 return dfdy
1611 def _derZ(self, x, y, z):
1612 """
1613 Returns the derivative with respect to z of the interpolated function
1614 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeZ.
1615 """
1616 x_pos, y_pos, z_pos = self._locate_xyz_indices(x, y, z)
1617 alpha = _cell_fraction(self.x_list, x_pos, x)
1618 beta = _cell_fraction(self.y_list, y_pos, y)
1619 dfdz = (
1620 (
1621 (1 - alpha) * (1 - beta) * self.f_values[x_pos - 1, y_pos - 1, z_pos]
1622 + (1 - alpha) * beta * self.f_values[x_pos - 1, y_pos, z_pos]
1623 + alpha * (1 - beta) * self.f_values[x_pos, y_pos - 1, z_pos]
1624 + alpha * beta * self.f_values[x_pos, y_pos, z_pos]
1625 )
1626 - (
1627 (1 - alpha)
1628 * (1 - beta)
1629 * self.f_values[x_pos - 1, y_pos - 1, z_pos - 1]
1630 + (1 - alpha) * beta * self.f_values[x_pos - 1, y_pos, z_pos - 1]
1631 + alpha * (1 - beta) * self.f_values[x_pos, y_pos - 1, z_pos - 1]
1632 + alpha * beta * self.f_values[x_pos, y_pos, z_pos - 1]
1633 )
1634 ) / (self.z_list[z_pos] - self.z_list[z_pos - 1])
1635 return dfdz
1638class QuadlinearInterp(HARKinterpolator4D):
1639 """
1640 Quadlinear full (or tensor) grid interpolation of a function f(w,x,y,z).
1642 Parameters
1643 ----------
1644 f_values : numpy.array
1645 An array of size (w_n,x_n,y_n,z_n) such that f_values[i,j,k,l] =
1646 f(w_list[i],x_list[j],y_list[k],z_list[l])
1647 w_list : numpy.array
1648 An array of w values, with length designated w_n.
1649 x_list : numpy.array
1650 An array of x values, with length designated x_n.
1651 y_list : numpy.array
1652 An array of y values, with length designated y_n.
1653 z_list : numpy.array
1654 An array of z values, with length designated z_n.
1655 """
1657 distance_criteria = ["f_values", "w_list", "x_list", "y_list", "z_list"]
1659 def __init__(self, f_values, w_list, x_list, y_list, z_list):
1660 self.f_values = f_values
1661 self.w_list = _coerce_1d_grid(w_list)
1662 self.x_list = _coerce_1d_grid(x_list)
1663 self.y_list = _coerce_1d_grid(y_list)
1664 self.z_list = _coerce_1d_grid(z_list)
1665 _check_grid_dimensions(
1666 4, self.f_values, self.w_list, self.x_list, self.y_list, self.z_list
1667 )
1668 self.w_n = self.w_list.size
1669 self.x_n = self.x_list.size
1670 self.y_n = self.y_list.size
1671 self.z_n = self.z_list.size
1673 def _locate_quad_indices(self, w, x, y, z):
1674 """
1675 Return clipped lookup indices ``(i, j, k, l)`` for ``(w, x, y, z)``.
1677 Each axis runs ``np.searchsorted`` and clips the result into
1678 ``[1, n - 1]`` so that ``a_list[idx - 1]`` and ``a_list[idx]``
1679 always bracket the query point.
1680 """
1681 return (
1682 _locate_clipped(self.w_list, w, self.w_n),
1683 _locate_clipped(self.x_list, x, self.x_n),
1684 _locate_clipped(self.y_list, y, self.y_n),
1685 _locate_clipped(self.z_list, z, self.z_n),
1686 )
1688 def _evaluate(self, w, x, y, z):
1689 """
1690 Returns the level of the interpolated function at each value in x,y,z.
1691 Only called internally by HARKinterpolator4D.__call__ (etc).
1692 """
1693 i, j, k, l = self._locate_quad_indices(w, x, y, z)
1694 alpha = _cell_fraction(self.w_list, i, w)
1695 beta = _cell_fraction(self.x_list, j, x)
1696 gamma = _cell_fraction(self.y_list, k, y)
1697 delta = _cell_fraction(self.z_list, l, z)
1698 f = (1 - alpha) * (
1699 (1 - beta)
1700 * (
1701 (1 - gamma) * (1 - delta) * self.f_values[i - 1, j - 1, k - 1, l - 1]
1702 + (1 - gamma) * delta * self.f_values[i - 1, j - 1, k - 1, l]
1703 + gamma * (1 - delta) * self.f_values[i - 1, j - 1, k, l - 1]
1704 + gamma * delta * self.f_values[i - 1, j - 1, k, l]
1705 )
1706 + beta
1707 * (
1708 (1 - gamma) * (1 - delta) * self.f_values[i - 1, j, k - 1, l - 1]
1709 + (1 - gamma) * delta * self.f_values[i - 1, j, k - 1, l]
1710 + gamma * (1 - delta) * self.f_values[i - 1, j, k, l - 1]
1711 + gamma * delta * self.f_values[i - 1, j, k, l]
1712 )
1713 ) + alpha * (
1714 (1 - beta)
1715 * (
1716 (1 - gamma) * (1 - delta) * self.f_values[i, j - 1, k - 1, l - 1]
1717 + (1 - gamma) * delta * self.f_values[i, j - 1, k - 1, l]
1718 + gamma * (1 - delta) * self.f_values[i, j - 1, k, l - 1]
1719 + gamma * delta * self.f_values[i, j - 1, k, l]
1720 )
1721 + beta
1722 * (
1723 (1 - gamma) * (1 - delta) * self.f_values[i, j, k - 1, l - 1]
1724 + (1 - gamma) * delta * self.f_values[i, j, k - 1, l]
1725 + gamma * (1 - delta) * self.f_values[i, j, k, l - 1]
1726 + gamma * delta * self.f_values[i, j, k, l]
1727 )
1728 )
1729 return f
1731 def _derW(self, w, x, y, z):
1732 """
1733 Returns the derivative with respect to w of the interpolated function
1734 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeW.
1735 """
1736 i, j, k, l = self._locate_quad_indices(w, x, y, z)
1737 beta = _cell_fraction(self.x_list, j, x)
1738 gamma = _cell_fraction(self.y_list, k, y)
1739 delta = _cell_fraction(self.z_list, l, z)
1740 dfdw = (
1741 (
1742 (1 - beta)
1743 * (1 - gamma)
1744 * (1 - delta)
1745 * self.f_values[i, j - 1, k - 1, l - 1]
1746 + (1 - beta) * (1 - gamma) * delta * self.f_values[i, j - 1, k - 1, l]
1747 + (1 - beta) * gamma * (1 - delta) * self.f_values[i, j - 1, k, l - 1]
1748 + (1 - beta) * gamma * delta * self.f_values[i, j - 1, k, l]
1749 + beta * (1 - gamma) * (1 - delta) * self.f_values[i, j, k - 1, l - 1]
1750 + beta * (1 - gamma) * delta * self.f_values[i, j, k - 1, l]
1751 + beta * gamma * (1 - delta) * self.f_values[i, j, k, l - 1]
1752 + beta * gamma * delta * self.f_values[i, j, k, l]
1753 )
1754 - (
1755 (1 - beta)
1756 * (1 - gamma)
1757 * (1 - delta)
1758 * self.f_values[i - 1, j - 1, k - 1, l - 1]
1759 + (1 - beta)
1760 * (1 - gamma)
1761 * delta
1762 * self.f_values[i - 1, j - 1, k - 1, l]
1763 + (1 - beta)
1764 * gamma
1765 * (1 - delta)
1766 * self.f_values[i - 1, j - 1, k, l - 1]
1767 + (1 - beta) * gamma * delta * self.f_values[i - 1, j - 1, k, l]
1768 + beta
1769 * (1 - gamma)
1770 * (1 - delta)
1771 * self.f_values[i - 1, j, k - 1, l - 1]
1772 + beta * (1 - gamma) * delta * self.f_values[i - 1, j, k - 1, l]
1773 + beta * gamma * (1 - delta) * self.f_values[i - 1, j, k, l - 1]
1774 + beta * gamma * delta * self.f_values[i - 1, j, k, l]
1775 )
1776 ) / (self.w_list[i] - self.w_list[i - 1])
1777 return dfdw
1779 def _derX(self, w, x, y, z):
1780 """
1781 Returns the derivative with respect to x of the interpolated function
1782 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeX.
1783 """
1784 i, j, k, l = self._locate_quad_indices(w, x, y, z)
1785 alpha = _cell_fraction(self.w_list, i, w)
1786 gamma = _cell_fraction(self.y_list, k, y)
1787 delta = _cell_fraction(self.z_list, l, z)
1788 dfdx = (
1789 (
1790 (1 - alpha)
1791 * (1 - gamma)
1792 * (1 - delta)
1793 * self.f_values[i - 1, j, k - 1, l - 1]
1794 + (1 - alpha) * (1 - gamma) * delta * self.f_values[i - 1, j, k - 1, l]
1795 + (1 - alpha) * gamma * (1 - delta) * self.f_values[i - 1, j, k, l - 1]
1796 + (1 - alpha) * gamma * delta * self.f_values[i - 1, j, k, l]
1797 + alpha * (1 - gamma) * (1 - delta) * self.f_values[i, j, k - 1, l - 1]
1798 + alpha * (1 - gamma) * delta * self.f_values[i, j, k - 1, l]
1799 + alpha * gamma * (1 - delta) * self.f_values[i, j, k, l - 1]
1800 + alpha * gamma * delta * self.f_values[i, j, k, l]
1801 )
1802 - (
1803 (1 - alpha)
1804 * (1 - gamma)
1805 * (1 - delta)
1806 * self.f_values[i - 1, j - 1, k - 1, l - 1]
1807 + (1 - alpha)
1808 * (1 - gamma)
1809 * delta
1810 * self.f_values[i - 1, j - 1, k - 1, l]
1811 + (1 - alpha)
1812 * gamma
1813 * (1 - delta)
1814 * self.f_values[i - 1, j - 1, k, l - 1]
1815 + (1 - alpha) * gamma * delta * self.f_values[i - 1, j - 1, k, l]
1816 + alpha
1817 * (1 - gamma)
1818 * (1 - delta)
1819 * self.f_values[i, j - 1, k - 1, l - 1]
1820 + alpha * (1 - gamma) * delta * self.f_values[i, j - 1, k - 1, l]
1821 + alpha * gamma * (1 - delta) * self.f_values[i, j - 1, k, l - 1]
1822 + alpha * gamma * delta * self.f_values[i, j - 1, k, l]
1823 )
1824 ) / (self.x_list[j] - self.x_list[j - 1])
1825 return dfdx
1827 def _derY(self, w, x, y, z):
1828 """
1829 Returns the derivative with respect to y of the interpolated function
1830 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeY.
1831 """
1832 i, j, k, l = self._locate_quad_indices(w, x, y, z)
1833 alpha = _cell_fraction(self.w_list, i, w)
1834 beta = _cell_fraction(self.x_list, j, x)
1835 delta = _cell_fraction(self.z_list, l, z)
1836 dfdy = (
1837 (
1838 (1 - alpha)
1839 * (1 - beta)
1840 * (1 - delta)
1841 * self.f_values[i - 1, j - 1, k, l - 1]
1842 + (1 - alpha) * (1 - beta) * delta * self.f_values[i - 1, j - 1, k, l]
1843 + (1 - alpha) * beta * (1 - delta) * self.f_values[i - 1, j, k, l - 1]
1844 + (1 - alpha) * beta * delta * self.f_values[i - 1, j, k, l]
1845 + alpha * (1 - beta) * (1 - delta) * self.f_values[i, j - 1, k, l - 1]
1846 + alpha * (1 - beta) * delta * self.f_values[i, j - 1, k, l]
1847 + alpha * beta * (1 - delta) * self.f_values[i, j, k, l - 1]
1848 + alpha * beta * delta * self.f_values[i, j, k, l]
1849 )
1850 - (
1851 (1 - alpha)
1852 * (1 - beta)
1853 * (1 - delta)
1854 * self.f_values[i - 1, j - 1, k - 1, l - 1]
1855 + (1 - alpha)
1856 * (1 - beta)
1857 * delta
1858 * self.f_values[i - 1, j - 1, k - 1, l]
1859 + (1 - alpha)
1860 * beta
1861 * (1 - delta)
1862 * self.f_values[i - 1, j, k - 1, l - 1]
1863 + (1 - alpha) * beta * delta * self.f_values[i - 1, j, k - 1, l]
1864 + alpha
1865 * (1 - beta)
1866 * (1 - delta)
1867 * self.f_values[i, j - 1, k - 1, l - 1]
1868 + alpha * (1 - beta) * delta * self.f_values[i, j - 1, k - 1, l]
1869 + alpha * beta * (1 - delta) * self.f_values[i, j, k - 1, l - 1]
1870 + alpha * beta * delta * self.f_values[i, j, k - 1, l]
1871 )
1872 ) / (self.y_list[k] - self.y_list[k - 1])
1873 return dfdy
1875 def _derZ(self, w, x, y, z):
1876 """
1877 Returns the derivative with respect to z of the interpolated function
1878 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeZ.
1879 """
1880 i, j, k, l = self._locate_quad_indices(w, x, y, z)
1881 alpha = _cell_fraction(self.w_list, i, w)
1882 beta = _cell_fraction(self.x_list, j, x)
1883 gamma = _cell_fraction(self.y_list, k, y)
1884 dfdz = (
1885 (
1886 (1 - alpha)
1887 * (1 - beta)
1888 * (1 - gamma)
1889 * self.f_values[i - 1, j - 1, k - 1, l]
1890 + (1 - alpha) * (1 - beta) * gamma * self.f_values[i - 1, j - 1, k, l]
1891 + (1 - alpha) * beta * (1 - gamma) * self.f_values[i - 1, j, k - 1, l]
1892 + (1 - alpha) * beta * gamma * self.f_values[i - 1, j, k, l]
1893 + alpha * (1 - beta) * (1 - gamma) * self.f_values[i, j - 1, k - 1, l]
1894 + alpha * (1 - beta) * gamma * self.f_values[i, j - 1, k, l]
1895 + alpha * beta * (1 - gamma) * self.f_values[i, j, k - 1, l]
1896 + alpha * beta * gamma * self.f_values[i, j, k, l]
1897 )
1898 - (
1899 (1 - alpha)
1900 * (1 - beta)
1901 * (1 - gamma)
1902 * self.f_values[i - 1, j - 1, k - 1, l - 1]
1903 + (1 - alpha)
1904 * (1 - beta)
1905 * gamma
1906 * self.f_values[i - 1, j - 1, k, l - 1]
1907 + (1 - alpha)
1908 * beta
1909 * (1 - gamma)
1910 * self.f_values[i - 1, j, k - 1, l - 1]
1911 + (1 - alpha) * beta * gamma * self.f_values[i - 1, j, k, l - 1]
1912 + alpha
1913 * (1 - beta)
1914 * (1 - gamma)
1915 * self.f_values[i, j - 1, k - 1, l - 1]
1916 + alpha * (1 - beta) * gamma * self.f_values[i, j - 1, k, l - 1]
1917 + alpha * beta * (1 - gamma) * self.f_values[i, j, k - 1, l - 1]
1918 + alpha * beta * gamma * self.f_values[i, j, k, l - 1]
1919 )
1920 ) / (self.z_list[l] - self.z_list[l - 1])
1921 return dfdz
1924def _init_envelope_state(obj, functions, nan_bool, lower=True):
1925 """Set ``compare``/``argcompare``/``functions``/``funcCount`` for an envelope."""
1926 if lower:
1927 obj.compare = np.nanmin if nan_bool else np.min
1928 obj.argcompare = np.nanargmin if nan_bool else np.argmin
1929 else:
1930 obj.compare = np.nanmax if nan_bool else np.max
1931 obj.argcompare = np.nanargmax if nan_bool else np.argmax
1932 obj.functions = list(functions)
1933 obj.funcCount = len(obj.functions)
1936class _Envelope1D(HARKinterpolator1D):
1937 """
1938 Base class for the lower/upper envelope of a finite set of 1D functions.
1940 Concrete subclasses set ``self.compare`` and ``self.argcompare`` in
1941 ``__init__`` (e.g. ``np.nanmin``/``np.nanargmin`` for the lower envelope,
1942 ``np.nanmax``/``np.nanargmax`` for the upper envelope). All evaluation
1943 logic is shared via ``self.compare`` and ``self.argcompare``.
1944 """
1946 distance_criteria = ["functions"]
1948 def _evaluate(self, x):
1949 """
1950 Returns the level of the envelope at each value in x. Only called
1951 internally by HARKinterpolator1D.__call__.
1952 """
1953 fx = np.column_stack([f(x) for f in self.functions])
1954 return self.compare(fx, axis=1)
1956 def _der(self, x):
1957 """
1958 Returns the first derivative of the envelope at each value in x. Only
1959 called internally by HARKinterpolator1D.derivative.
1960 """
1961 y, dydx = self._evalAndDer(x)
1962 return dydx # Sadly, this is the fastest / most convenient way...
1964 def _evalAndDer(self, x):
1965 """
1966 Returns the level and first derivative of the envelope at each value
1967 in x. Only called internally by HARKinterpolator1D.eval_and_der.
1968 """
1969 fx = np.column_stack([f(x) for f in self.functions])
1970 i = self.argcompare(fx, axis=1)
1971 y = fx[np.arange(len(x)), i]
1972 dydx = np.zeros_like(y)
1973 for j in np.unique(i):
1974 c = i == j
1975 dydx[c] = self.functions[j].derivative(x[c])
1976 return y, dydx
1979class LowerEnvelope(_Envelope1D):
1980 """
1981 The lower envelope of a finite set of 1D functions, each of which can be of
1982 any class that has the methods __call__, derivative, and eval_with_derivative.
1983 Generally: it combines HARKinterpolator1Ds.
1985 Parameters
1986 ----------
1987 *functions : function
1988 Any number of real functions; often instances of HARKinterpolator1D
1989 nan_bool : boolean
1990 An indicator for whether the solver should exclude NA's when
1991 forming the lower envelope
1992 """
1994 def __init__(self, *functions, nan_bool=True):
1995 _init_envelope_state(self, functions, nan_bool, lower=True)
1998class UpperEnvelope(_Envelope1D):
1999 """
2000 The upper envelope of a finite set of 1D functions, each of which can be of
2001 any class that has the methods __call__, derivative, and eval_with_derivative.
2002 Generally: it combines HARKinterpolator1Ds.
2004 Parameters
2005 ----------
2006 *functions : function
2007 Any number of real functions; often instances of HARKinterpolator1D
2008 nan_bool : boolean
2009 An indicator for whether the solver should exclude NA's when forming
2010 the upper envelope.
2011 """
2013 def __init__(self, *functions, nan_bool=True):
2014 _init_envelope_state(self, functions, nan_bool, lower=False)
2017class LowerEnvelope2D(HARKinterpolator2D):
2018 """
2019 The lower envelope of a finite set of 2D functions, each of which can be of
2020 any class that has the methods __call__, derivativeX, and derivativeY.
2021 Generally: it combines HARKinterpolator2Ds.
2023 Parameters
2024 ----------
2025 *functions : function
2026 Any number of real functions; often instances of HARKinterpolator2D
2027 nan_bool : boolean
2028 An indicator for whether the solver should exclude NA's when forming
2029 the lower envelope.
2030 """
2032 distance_criteria = ["functions"]
2034 def __init__(self, *functions, nan_bool=True):
2035 _init_envelope_state(self, functions, nan_bool, lower=True)
2037 def _evaluate(self, x, y):
2038 """
2039 Returns the level of the function at each value in (x,y) as the minimum
2040 among all of the functions. Only called internally by
2041 HARKinterpolator2D.__call__.
2042 """
2043 temp = np.column_stack([f(x, y) for f in self.functions])
2044 return self.compare(temp, axis=1)
2046 def _derX(self, x, y):
2047 """
2048 Returns the first derivative of the function with respect to X at each
2049 value in (x,y). Only called internally by HARKinterpolator2D._derX.
2050 """
2051 return _envelope_partial(self, (x, y), "derivativeX")
2053 def _derY(self, x, y):
2054 """
2055 Returns the first derivative of the function with respect to Y at each
2056 value in (x,y). Only called internally by HARKinterpolator2D._derY.
2057 """
2058 return _envelope_partial(self, (x, y), "derivativeY")
2061class LowerEnvelope3D(HARKinterpolator3D):
2062 """
2063 The lower envelope of a finite set of 3D functions, each of which can be of
2064 any class that has the methods __call__, derivativeX, derivativeY, and
2065 derivativeZ. Generally: it combines HARKinterpolator2Ds.
2067 Parameters
2068 ----------
2069 *functions : function
2070 Any number of real functions; often instances of HARKinterpolator3D
2071 nan_bool : boolean
2072 An indicator for whether the solver should exclude NA's when forming
2073 the lower envelope.
2074 """
2076 distance_criteria = ["functions"]
2078 def __init__(self, *functions, nan_bool=True):
2079 _init_envelope_state(self, functions, nan_bool, lower=True)
2081 def _evaluate(self, x, y, z):
2082 """
2083 Returns the level of the function at each value in (x,y,z) as the minimum
2084 among all of the functions. Only called internally by
2085 HARKinterpolator3D.__call__.
2086 """
2087 temp = np.column_stack([f(x, y, z) for f in self.functions])
2088 return self.compare(temp, axis=1)
2090 def _derX(self, x, y, z):
2091 """
2092 Returns the first derivative of the function with respect to X at each
2093 value in (x,y,z). Only called internally by HARKinterpolator3D._derX.
2094 """
2095 return _envelope_partial(self, (x, y, z), "derivativeX")
2097 def _derY(self, x, y, z):
2098 """
2099 Returns the first derivative of the function with respect to Y at each
2100 value in (x,y,z). Only called internally by HARKinterpolator3D._derY.
2101 """
2102 return _envelope_partial(self, (x, y, z), "derivativeY")
2104 def _derZ(self, x, y, z):
2105 """
2106 Returns the first derivative of the function with respect to Z at each
2107 value in (x,y,z). Only called internally by HARKinterpolator3D._derZ.
2108 """
2109 return _envelope_partial(self, (x, y, z), "derivativeZ")
2112class VariableLowerBoundFunc2D(HARKinterpolator2D):
2113 """
2114 A class for representing a function with two real inputs whose lower bound
2115 in the first input depends on the second input. Useful for managing curved
2116 natural borrowing constraints, as occurs in the persistent shocks model.
2118 Parameters
2119 ----------
2120 func : function
2121 A function f: (R_+ x R) --> R representing the function of interest
2122 shifted by its lower bound in the first input.
2123 lowerBound : function
2124 The lower bound in the first input of the function of interest, as
2125 a function of the second input.
2126 """
2128 distance_criteria = ["func", "lowerBound"]
2130 def __init__(self, func, lowerBound):
2131 self.func = func
2132 self.lowerBound = lowerBound
2134 def __call__(self, x, y):
2135 """
2136 Evaluate the function at given state space points.
2138 Parameters
2139 ----------
2140 x : np.array
2141 First input values.
2142 y : np.array
2143 Second input values; should be of same shape as x.
2145 Returns
2146 -------
2147 f_out : np.array
2148 Function evaluated at (x,y), of same shape as inputs.
2149 """
2150 xShift = self.lowerBound(y)
2151 f_out = self.func(x - xShift, y)
2152 return f_out
2154 def _derX(self, x, y):
2155 """
2156 Evaluate the first derivative with respect to x of the function at given
2157 state space points.
2159 Parameters
2160 ----------
2161 x : np.array
2162 First input values.
2163 y : np.array
2164 Second input values; should be of same shape as x.
2166 Returns
2167 -------
2168 dfdx_out : np.array
2169 First derivative of function with respect to the first input,
2170 evaluated at (x,y), of same shape as inputs.
2171 """
2172 xShift = self.lowerBound(y)
2173 dfdx_out = self.func.derivativeX(x - xShift, y)
2174 return dfdx_out
2176 def _derY(self, x, y):
2177 """
2178 Evaluate the first derivative with respect to y of the function at given
2179 state space points.
2181 Parameters
2182 ----------
2183 x : np.array
2184 First input values.
2185 y : np.array
2186 Second input values; should be of same shape as x.
2188 Returns
2189 -------
2190 dfdy_out : np.array
2191 First derivative of function with respect to the second input,
2192 evaluated at (x,y), of same shape as inputs.
2193 """
2194 xShift, xShiftDer = self.lowerBound.eval_with_derivative(y)
2195 dfdy_out = self.func.derivativeY(
2196 x - xShift, y
2197 ) - xShiftDer * self.func.derivativeX(x - xShift, y)
2198 return dfdy_out
2201class VariableLowerBoundFunc3D(HARKinterpolator3D):
2202 """
2203 A class for representing a function with three real inputs whose lower bound
2204 in the first input depends on the second input. Useful for managing curved
2205 natural borrowing constraints.
2207 Parameters
2208 ----------
2209 func : function
2210 A function f: (R_+ x R^2) --> R representing the function of interest
2211 shifted by its lower bound in the first input.
2212 lowerBound : function
2213 The lower bound in the first input of the function of interest, as
2214 a function of the second input.
2215 """
2217 distance_criteria = ["func", "lowerBound"]
2219 def __init__(self, func, lowerBound):
2220 self.func = func
2221 self.lowerBound = lowerBound
2223 def __call__(self, x, y, z):
2224 """
2225 Evaluate the function at given state space points.
2227 Parameters
2228 ----------
2229 x : np.array
2230 First input values.
2231 y : np.array
2232 Second input values; should be of same shape as x.
2233 z : np.array
2234 Third input values; should be of same shape as x.
2236 Returns
2237 -------
2238 f_out : np.array
2239 Function evaluated at (x,y,z), of same shape as inputs.
2240 """
2241 xShift = self.lowerBound(y)
2242 f_out = self.func(x - xShift, y, z)
2243 return f_out
2245 def _derX(self, x, y, z):
2246 """
2247 Evaluate the first derivative with respect to x of the function at given
2248 state space points.
2250 Parameters
2251 ----------
2252 x : np.array
2253 First input values.
2254 y : np.array
2255 Second input values; should be of same shape as x.
2256 z : np.array
2257 Third input values; should be of same shape as x.
2259 Returns
2260 -------
2261 dfdx_out : np.array
2262 First derivative of function with respect to the first input,
2263 evaluated at (x,y,z), of same shape as inputs.
2264 """
2265 xShift = self.lowerBound(y)
2266 dfdx_out = self.func.derivativeX(x - xShift, y, z)
2267 return dfdx_out
2269 def _derY(self, x, y, z):
2270 """
2271 Evaluate the first derivative with respect to y of the function at given
2272 state space points.
2274 Parameters
2275 ----------
2276 x : np.array
2277 First input values.
2278 y : np.array
2279 Second input values; should be of same shape as x.
2280 z : np.array
2281 Third input values; should be of same shape as x.
2283 Returns
2284 -------
2285 dfdy_out : np.array
2286 First derivative of function with respect to the second input,
2287 evaluated at (x,y,z), of same shape as inputs.
2288 """
2289 xShift, xShiftDer = self.lowerBound.eval_with_derivative(y)
2290 dfdy_out = self.func.derivativeY(
2291 x - xShift, y, z
2292 ) - xShiftDer * self.func.derivativeX(x - xShift, y, z)
2293 return dfdy_out
2295 def _derZ(self, x, y, z):
2296 """
2297 Evaluate the first derivative with respect to z of the function at given
2298 state space points.
2300 Parameters
2301 ----------
2302 x : np.array
2303 First input values.
2304 y : np.array
2305 Second input values; should be of same shape as x.
2306 z : np.array
2307 Third input values; should be of same shape as x.
2309 Returns
2310 -------
2311 dfdz_out : np.array
2312 First derivative of function with respect to the third input,
2313 evaluated at (x,y,z), of same shape as inputs.
2314 """
2315 xShift = self.lowerBound(y)
2316 dfdz_out = self.func.derivativeZ(x - xShift, y, z)
2317 return dfdz_out
2320class LinearInterpOnInterp1D(HARKinterpolator2D):
2321 """
2322 A 2D interpolator that linearly interpolates among a list of 1D interpolators.
2324 Parameters
2325 ----------
2326 xInterpolators : [HARKinterpolator1D]
2327 A list of 1D interpolations over the x variable. The nth element of
2328 xInterpolators represents f(x,y_values[n]).
2329 y_values: numpy.array
2330 An array of y values equal in length to xInterpolators.
2331 """
2333 distance_criteria = ["xInterpolators", "y_list"]
2335 def __init__(self, xInterpolators, y_values):
2336 self.xInterpolators = xInterpolators
2337 self.y_list = y_values
2338 self.y_n = y_values.size
2340 def _linear_y_blend(self, x, y, eval_func):
2341 """Evaluate ``eval_func`` on each cell's bracketing 1D interpolators
2342 and combine with the y-direction linear weights ``(1 - alpha)`` and
2343 ``alpha``. Shared by ``_evaluate`` and ``_derX``.
2344 """
2345 m = len(x)
2346 y_pos = _locate_clipped(self.y_list, y, self.y_n)
2347 out = np.full(m, np.nan)
2348 for i, c in _iter_unique_pairs(y_pos):
2349 alpha = _cell_fraction(self.y_list, i, y[c])
2350 out[c] = (1 - alpha) * eval_func(
2351 self.xInterpolators[i - 1], x[c]
2352 ) + alpha * eval_func(self.xInterpolators[i], x[c])
2353 return out
2355 def _evaluate(self, x, y):
2356 """
2357 Returns the level of the interpolated function at each value in x,y.
2358 Only called internally by HARKinterpolator2D.__call__ (etc).
2359 """
2360 return self._linear_y_blend(x, y, lambda interp, xs: interp(xs))
2362 def _derX(self, x, y):
2363 """
2364 Returns the derivative with respect to x of the interpolated function
2365 at each value in x,y. Only called internally by HARKinterpolator2D.derivativeX.
2366 """
2367 return self._linear_y_blend(x, y, lambda interp, xs: interp._der(xs))
2369 def _derY(self, x, y):
2370 """
2371 Returns the derivative with respect to y of the interpolated function
2372 at each value in x,y. Only called internally by HARKinterpolator2D.derivativeY.
2373 """
2374 m = len(x)
2375 y_pos = _locate_clipped(self.y_list, y, self.y_n)
2376 dfdy = np.full(m, np.nan)
2377 for i, c in _iter_unique_pairs(y_pos):
2378 dfdy[c] = (
2379 self.xInterpolators[i](x[c]) - self.xInterpolators[i - 1](x[c])
2380 ) / (self.y_list[i] - self.y_list[i - 1])
2381 return dfdy
2384class BilinearInterpOnInterp1D(HARKinterpolator3D):
2385 """
2386 A 3D interpolator that bilinearly interpolates among a list of lists of 1D
2387 interpolators.
2389 Constructor for the class, generating an approximation to a function of
2390 the form f(x,y,z) using interpolations over f(x,y_0,z_0) for a fixed grid
2391 of y_0 and z_0 values.
2393 Parameters
2394 ----------
2395 xInterpolators : [[HARKinterpolator1D]]
2396 A list of lists of 1D interpolations over the x variable. The i,j-th
2397 element of xInterpolators represents f(x,y_values[i],z_values[j]).
2398 y_values: numpy.array
2399 An array of y values equal in length to xInterpolators.
2400 z_values: numpy.array
2401 An array of z values equal in length to xInterpolators[0].
2402 """
2404 distance_criteria = ["xInterpolators", "y_list", "z_list"]
2406 def __init__(self, xInterpolators, y_values, z_values):
2407 self.xInterpolators = xInterpolators
2408 self.y_list = y_values
2409 self.y_n = y_values.size
2410 self.z_list = z_values
2411 self.z_n = z_values.size
2413 def _locate_yz_indices(self, y, z):
2414 """Return clipped ``searchsorted`` indices for ``y`` and ``z`` shared
2415 by ``_evaluate`` and the three derivative methods."""
2416 return (
2417 _locate_clipped(self.y_list, y, self.y_n),
2418 _locate_clipped(self.z_list, z, self.z_n),
2419 )
2421 def _bilinear_loop(self, x, y, z, eval_func):
2422 """Bilinear blend of ``eval_func`` across ``xInterpolators`` corners.
2424 Shared by ``_evaluate`` (``f(x)``) and ``_derX`` (``f._der(x)``).
2425 """
2426 m = len(x)
2427 y_pos, z_pos = self._locate_yz_indices(y, z)
2428 out = np.full(m, np.nan)
2429 for i, j, c in _iter_unique_pairs(y_pos, z_pos):
2430 alpha = _cell_fraction(self.y_list, i, y[c])
2431 beta = _cell_fraction(self.z_list, j, z[c])
2432 xc = x[c]
2433 out[c] = (
2434 (1 - alpha)
2435 * (1 - beta)
2436 * eval_func(self.xInterpolators[i - 1][j - 1], xc)
2437 + (1 - alpha) * beta * eval_func(self.xInterpolators[i - 1][j], xc)
2438 + alpha * (1 - beta) * eval_func(self.xInterpolators[i][j - 1], xc)
2439 + alpha * beta * eval_func(self.xInterpolators[i][j], xc)
2440 )
2441 return out
2443 def _evaluate(self, x, y, z):
2444 """
2445 Returns the level of the interpolated function at each value in x,y,z.
2446 Only called internally by HARKinterpolator3D.__call__ (etc).
2447 """
2448 return self._bilinear_loop(x, y, z, lambda f, xc: f(xc))
2450 def _derX(self, x, y, z):
2451 """
2452 Returns the derivative with respect to x of the interpolated function
2453 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeX.
2454 """
2455 return self._bilinear_loop(x, y, z, lambda f, xc: f._der(xc))
2457 def _derY(self, x, y, z):
2458 """
2459 Returns the derivative with respect to y of the interpolated function
2460 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeY.
2461 """
2462 m = len(x)
2463 y_pos, z_pos = self._locate_yz_indices(y, z)
2464 dfdy = np.full(m, np.nan)
2465 for i, j, c in _iter_unique_pairs(y_pos, z_pos):
2466 beta = _cell_fraction(self.z_list, j, z[c])
2467 dfdy[c] = (
2468 (
2469 (1 - beta) * self.xInterpolators[i][j - 1](x[c])
2470 + beta * self.xInterpolators[i][j](x[c])
2471 )
2472 - (
2473 (1 - beta) * self.xInterpolators[i - 1][j - 1](x[c])
2474 + beta * self.xInterpolators[i - 1][j](x[c])
2475 )
2476 ) / (self.y_list[i] - self.y_list[i - 1])
2477 return dfdy
2479 def _derZ(self, x, y, z):
2480 """
2481 Returns the derivative with respect to z of the interpolated function
2482 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeZ.
2483 """
2484 m = len(x)
2485 y_pos, z_pos = self._locate_yz_indices(y, z)
2486 dfdz = np.full(m, np.nan)
2487 for i, j, c in _iter_unique_pairs(y_pos, z_pos):
2488 alpha = _cell_fraction(self.y_list, i, y[c])
2489 dfdz[c] = (
2490 (
2491 (1 - alpha) * self.xInterpolators[i - 1][j](x[c])
2492 + alpha * self.xInterpolators[i][j](x[c])
2493 )
2494 - (
2495 (1 - alpha) * self.xInterpolators[i - 1][j - 1](x[c])
2496 + alpha * self.xInterpolators[i][j - 1](x[c])
2497 )
2498 ) / (self.z_list[j] - self.z_list[j - 1])
2499 return dfdz
2502class TrilinearInterpOnInterp1D(HARKinterpolator4D):
2503 """
2504 A 4D interpolator that trilinearly interpolates among a list of lists of 1D interpolators.
2506 Constructor for the class, generating an approximation to a function of
2507 the form f(w,x,y,z) using interpolations over f(w,x_0,y_0,z_0) for a fixed
2508 grid of y_0 and z_0 values.
2510 Parameters
2511 ----------
2512 wInterpolators : [[[HARKinterpolator1D]]]
2513 A list of lists of lists of 1D interpolations over the x variable.
2514 The i,j,k-th element of wInterpolators represents f(w,x_values[i],y_values[j],z_values[k]).
2515 x_values: numpy.array
2516 An array of x values equal in length to wInterpolators.
2517 y_values: numpy.array
2518 An array of y values equal in length to wInterpolators[0].
2519 z_values: numpy.array
2520 An array of z values equal in length to wInterpolators[0][0]
2521 """
2523 distance_criteria = ["wInterpolators", "x_list", "y_list", "z_list"]
2525 def __init__(self, wInterpolators, x_values, y_values, z_values):
2526 self.wInterpolators = wInterpolators
2527 self.x_list = x_values
2528 self.x_n = x_values.size
2529 self.y_list = y_values
2530 self.y_n = y_values.size
2531 self.z_list = z_values
2532 self.z_n = z_values.size
2534 def _locate_xyz_indices(self, x, y, z):
2535 """Return clamped ``searchsorted`` indices for ``x``, ``y``, ``z`` shared
2536 by ``_evaluate`` and the four derivative methods."""
2537 return (
2538 _locate_clipped(self.x_list, x, self.x_n),
2539 _locate_clipped(self.y_list, y, self.y_n),
2540 _locate_clipped(self.z_list, z, self.z_n),
2541 )
2543 def _iter_xyz_cells(self, x, y, z, x_pos, y_pos, z_pos):
2544 """Yield ``(i, j, k, c, alpha, beta, gamma)`` for each non-empty cell of
2545 the (x, y, z) grid. Shared by ``_trilinear_loop`` and the partial-derivative
2546 methods, all of which use a subset of these values."""
2547 for i, j, k, c in _iter_unique_pairs(x_pos, y_pos, z_pos):
2548 alpha = _cell_fraction(self.x_list, i, x[c])
2549 beta = _cell_fraction(self.y_list, j, y[c])
2550 gamma = _cell_fraction(self.z_list, k, z[c])
2551 yield i, j, k, c, alpha, beta, gamma
2553 def _trilinear_loop(self, w, x, y, z, eval_func):
2554 """Trilinear interpolation over ``wInterpolators[i,j,k]`` evaluated by
2555 ``eval_func``. Shared by ``_evaluate`` (``f(w)``) and ``_derW`` (``f._der(w)``)."""
2556 m = len(x)
2557 x_pos, y_pos, z_pos = self._locate_xyz_indices(x, y, z)
2558 out = np.full(m, np.nan)
2559 for i, j, k, c, alpha, beta, gamma in self._iter_xyz_cells(
2560 x, y, z, x_pos, y_pos, z_pos
2561 ):
2562 wc = w[c]
2563 out[c] = (
2564 (1 - alpha)
2565 * (1 - beta)
2566 * (1 - gamma)
2567 * eval_func(self.wInterpolators[i - 1][j - 1][k - 1], wc)
2568 + (1 - alpha)
2569 * (1 - beta)
2570 * gamma
2571 * eval_func(self.wInterpolators[i - 1][j - 1][k], wc)
2572 + (1 - alpha)
2573 * beta
2574 * (1 - gamma)
2575 * eval_func(self.wInterpolators[i - 1][j][k - 1], wc)
2576 + (1 - alpha)
2577 * beta
2578 * gamma
2579 * eval_func(self.wInterpolators[i - 1][j][k], wc)
2580 + alpha
2581 * (1 - beta)
2582 * (1 - gamma)
2583 * eval_func(self.wInterpolators[i][j - 1][k - 1], wc)
2584 + alpha
2585 * (1 - beta)
2586 * gamma
2587 * eval_func(self.wInterpolators[i][j - 1][k], wc)
2588 + alpha
2589 * beta
2590 * (1 - gamma)
2591 * eval_func(self.wInterpolators[i][j][k - 1], wc)
2592 + alpha * beta * gamma * eval_func(self.wInterpolators[i][j][k], wc)
2593 )
2594 return out
2596 def _evaluate(self, w, x, y, z):
2597 """
2598 Returns the level of the interpolated function at each value in w,x,y,z.
2599 Only called internally by HARKinterpolator4D.__call__ (etc).
2600 """
2601 return self._trilinear_loop(w, x, y, z, lambda f, ww: f(ww))
2603 def _derW(self, w, x, y, z):
2604 """
2605 Returns the derivative with respect to w of the interpolated function
2606 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeW.
2607 """
2608 return self._trilinear_loop(w, x, y, z, lambda f, ww: f._der(ww))
2610 def _derX(self, w, x, y, z):
2611 """
2612 Returns the derivative with respect to x of the interpolated function
2613 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeX.
2614 """
2615 m = len(x)
2616 x_pos, y_pos, z_pos = self._locate_xyz_indices(x, y, z)
2617 dfdx = np.full(m, np.nan)
2618 for i, j, k, c, _alpha, beta, gamma in self._iter_xyz_cells(
2619 x, y, z, x_pos, y_pos, z_pos
2620 ):
2621 wc = w[c]
2622 dfdx[c] = (
2623 (
2624 (1 - beta) * (1 - gamma) * self.wInterpolators[i][j - 1][k - 1](wc)
2625 + (1 - beta) * gamma * self.wInterpolators[i][j - 1][k](wc)
2626 + beta * (1 - gamma) * self.wInterpolators[i][j][k - 1](wc)
2627 + beta * gamma * self.wInterpolators[i][j][k](wc)
2628 )
2629 - (
2630 (1 - beta)
2631 * (1 - gamma)
2632 * self.wInterpolators[i - 1][j - 1][k - 1](wc)
2633 + (1 - beta) * gamma * self.wInterpolators[i - 1][j - 1][k](wc)
2634 + beta * (1 - gamma) * self.wInterpolators[i - 1][j][k - 1](wc)
2635 + beta * gamma * self.wInterpolators[i - 1][j][k](wc)
2636 )
2637 ) / (self.x_list[i] - self.x_list[i - 1])
2638 return dfdx
2640 def _derY(self, w, x, y, z):
2641 """
2642 Returns the derivative with respect to y of the interpolated function
2643 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeY.
2644 """
2645 m = len(x)
2646 x_pos, y_pos, z_pos = self._locate_xyz_indices(x, y, z)
2647 dfdy = np.full(m, np.nan)
2648 for i, j, k, c, alpha, _beta, gamma in self._iter_xyz_cells(
2649 x, y, z, x_pos, y_pos, z_pos
2650 ):
2651 wc = w[c]
2652 dfdy[c] = (
2653 (
2654 (1 - alpha) * (1 - gamma) * self.wInterpolators[i - 1][j][k - 1](wc)
2655 + (1 - alpha) * gamma * self.wInterpolators[i - 1][j][k](wc)
2656 + alpha * (1 - gamma) * self.wInterpolators[i][j][k - 1](wc)
2657 + alpha * gamma * self.wInterpolators[i][j][k](wc)
2658 )
2659 - (
2660 (1 - alpha)
2661 * (1 - gamma)
2662 * self.wInterpolators[i - 1][j - 1][k - 1](wc)
2663 + (1 - alpha) * gamma * self.wInterpolators[i - 1][j - 1][k](wc)
2664 + alpha * (1 - gamma) * self.wInterpolators[i][j - 1][k - 1](wc)
2665 + alpha * gamma * self.wInterpolators[i][j - 1][k](wc)
2666 )
2667 ) / (self.y_list[j] - self.y_list[j - 1])
2668 return dfdy
2670 def _derZ(self, w, x, y, z):
2671 """
2672 Returns the derivative with respect to z of the interpolated function
2673 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeZ.
2674 """
2675 m = len(x)
2676 x_pos, y_pos, z_pos = self._locate_xyz_indices(x, y, z)
2677 dfdz = np.full(m, np.nan)
2678 for i, j, k, c, alpha, beta, _gamma in self._iter_xyz_cells(
2679 x, y, z, x_pos, y_pos, z_pos
2680 ):
2681 wc = w[c]
2682 dfdz[c] = (
2683 (
2684 (1 - alpha) * (1 - beta) * self.wInterpolators[i - 1][j - 1][k](wc)
2685 + (1 - alpha) * beta * self.wInterpolators[i - 1][j][k](wc)
2686 + alpha * (1 - beta) * self.wInterpolators[i][j - 1][k](wc)
2687 + alpha * beta * self.wInterpolators[i][j][k](wc)
2688 )
2689 - (
2690 (1 - alpha)
2691 * (1 - beta)
2692 * self.wInterpolators[i - 1][j - 1][k - 1](wc)
2693 + (1 - alpha) * beta * self.wInterpolators[i - 1][j][k - 1](wc)
2694 + alpha * (1 - beta) * self.wInterpolators[i][j - 1][k - 1](wc)
2695 + alpha * beta * self.wInterpolators[i][j][k - 1](wc)
2696 )
2697 ) / (self.z_list[k] - self.z_list[k - 1])
2698 return dfdz
2701class LinearInterpOnInterp2D(HARKinterpolator3D):
2702 """
2703 A 3D interpolation method that linearly interpolates between "layers" of
2704 arbitrary 2D interpolations. Useful for models with two endogenous state
2705 variables and one exogenous state variable when solving with the endogenous
2706 grid method. NOTE: should not be used if an exogenous 3D grid is used, will
2707 be significantly slower than TrilinearInterp.
2709 Constructor for the class, generating an approximation to a function of
2710 the form f(x,y,z) using interpolations over f(x,y,z_0) for a fixed grid
2711 of z_0 values.
2713 Parameters
2714 ----------
2715 xyInterpolators : [HARKinterpolator2D]
2716 A list of 2D interpolations over the x and y variables. The nth
2717 element of xyInterpolators represents f(x,y,z_values[n]).
2718 z_values: numpy.array
2719 An array of z values equal in length to xyInterpolators.
2720 """
2722 distance_criteria = ["xyInterpolators", "z_list"]
2724 def __init__(self, xyInterpolators, z_values):
2725 self.xyInterpolators = xyInterpolators
2726 self.z_list = z_values
2727 self.z_n = z_values.size
2729 def _linear_z_blend(self, x, y, z, eval_func):
2730 """Linear blend of ``eval_func`` between consecutive ``xyInterpolators``
2731 layers along ``z``. Shared by ``_evaluate``, ``_derX``, ``_derY``."""
2732 m = len(x)
2733 z_pos = _locate_clipped(self.z_list, z, self.z_n)
2734 out = np.full(m, np.nan)
2735 for i, c in _iter_unique_pairs(z_pos):
2736 alpha = _cell_fraction(self.z_list, i, z[c])
2737 lower = eval_func(self.xyInterpolators[i - 1], x[c], y[c])
2738 upper = eval_func(self.xyInterpolators[i], x[c], y[c])
2739 out[c] = (1 - alpha) * lower + alpha * upper
2740 return out
2742 def _evaluate(self, x, y, z):
2743 """
2744 Returns the level of the interpolated function at each value in x,y,z.
2745 Only called internally by HARKinterpolator3D.__call__ (etc).
2746 """
2747 return self._linear_z_blend(x, y, z, lambda f, xv, yv: f(xv, yv))
2749 def _derX(self, x, y, z):
2750 """
2751 Returns the derivative with respect to x of the interpolated function
2752 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeX.
2753 """
2754 return self._linear_z_blend(x, y, z, lambda f, xv, yv: f.derivativeX(xv, yv))
2756 def _derY(self, x, y, z):
2757 """
2758 Returns the derivative with respect to y of the interpolated function
2759 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeY.
2760 """
2761 return self._linear_z_blend(x, y, z, lambda f, xv, yv: f.derivativeY(xv, yv))
2763 def _derZ(self, x, y, z):
2764 """
2765 Returns the derivative with respect to z of the interpolated function
2766 at each value in x,y,z. Only called internally by HARKinterpolator3D.derivativeZ.
2767 """
2768 m = len(x)
2769 z_pos = _locate_clipped(self.z_list, z, self.z_n)
2770 dfdz = np.full(m, np.nan)
2771 for i, c in _iter_unique_pairs(z_pos):
2772 dfdz[c] = (
2773 self.xyInterpolators[i](x[c], y[c])
2774 - self.xyInterpolators[i - 1](x[c], y[c])
2775 ) / (self.z_list[i] - self.z_list[i - 1])
2776 return dfdz
2779class BilinearInterpOnInterp2D(HARKinterpolator4D):
2780 """
2781 A 4D interpolation method that bilinearly interpolates among "layers" of
2782 arbitrary 2D interpolations. Useful for models with two endogenous state
2783 variables and two exogenous state variables when solving with the endogenous
2784 grid method. NOTE: should not be used if an exogenous 4D grid is used, will
2785 be significantly slower than QuadlinearInterp.
2787 Constructor for the class, generating an approximation to a function of
2788 the form f(w,x,y,z) using interpolations over f(w,x,y_0,z_0) for a fixed
2789 grid of y_0 and z_0 values.
2791 Parameters
2792 ----------
2793 wxInterpolators : [[HARKinterpolator2D]]
2794 A list of lists of 2D interpolations over the w and x variables.
2795 The i,j-th element of wxInterpolators represents
2796 f(w,x,y_values[i],z_values[j]).
2797 y_values: numpy.array
2798 An array of y values equal in length to wxInterpolators.
2799 z_values: numpy.array
2800 An array of z values equal in length to wxInterpolators[0].
2801 """
2803 distance_criteria = ["wxInterpolators", "y_list", "z_list"]
2805 def __init__(self, wxInterpolators, y_values, z_values):
2806 self.wxInterpolators = wxInterpolators
2807 self.y_list = y_values
2808 self.y_n = y_values.size
2809 self.z_list = z_values
2810 self.z_n = z_values.size
2812 def _locate_yz_indices(self, y, z):
2813 """Return clamped ``searchsorted`` indices for ``y`` and ``z`` shared
2814 by ``_evaluate`` and the four derivative methods."""
2815 return (
2816 _locate_clipped(self.y_list, y, self.y_n),
2817 _locate_clipped(self.z_list, z, self.z_n),
2818 )
2820 def _bilinear_loop(self, w, x, y, z, eval_func):
2821 """Bilinear interpolation across (y, z) layers of ``wxInterpolators``,
2822 with each corner evaluated by ``eval_func``. Shared by ``_evaluate``,
2823 ``_derW``, ``_derX`` (the latter two pick a derivative method)."""
2824 m = len(x)
2825 y_pos, z_pos = self._locate_yz_indices(y, z)
2826 out = np.full(m, np.nan)
2827 for i, j, c in _iter_unique_pairs(y_pos, z_pos):
2828 alpha = _cell_fraction(self.y_list, i, y[c])
2829 beta = _cell_fraction(self.z_list, j, z[c])
2830 wc, xc = w[c], x[c]
2831 out[c] = (
2832 (1 - alpha)
2833 * (1 - beta)
2834 * eval_func(self.wxInterpolators[i - 1][j - 1], wc, xc)
2835 + (1 - alpha) * beta * eval_func(self.wxInterpolators[i - 1][j], wc, xc)
2836 + alpha * (1 - beta) * eval_func(self.wxInterpolators[i][j - 1], wc, xc)
2837 + alpha * beta * eval_func(self.wxInterpolators[i][j], wc, xc)
2838 )
2839 return out
2841 def _evaluate(self, w, x, y, z):
2842 """
2843 Returns the level of the interpolated function at each value in x,y,z.
2844 Only called internally by HARKinterpolator4D.__call__ (etc).
2845 """
2846 return self._bilinear_loop(w, x, y, z, lambda f, wc, xc: f(wc, xc))
2848 def _derW(self, w, x, y, z):
2849 """
2850 Returns the derivative with respect to w of the interpolated function
2851 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeW.
2852 """
2853 # This may look strange, as we call the derivativeX() method to get the
2854 # derivative with respect to w, but that's just a quirk of 4D interpolations
2855 # beginning with w rather than x. The derivative wrt the first dimension
2856 # of an element of wxInterpolators is the w-derivative of the main function.
2857 return self._bilinear_loop(w, x, y, z, lambda f, wc, xc: f.derivativeX(wc, xc))
2859 def _derX(self, w, x, y, z):
2860 """
2861 Returns the derivative with respect to x of the interpolated function
2862 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeX.
2863 """
2864 # This may look strange, as we call the derivativeY() method to get the
2865 # derivative with respect to x, but that's just a quirk of 4D interpolations
2866 # beginning with w rather than x. The derivative wrt the second dimension
2867 # of an element of wxInterpolators is the x-derivative of the main function.
2868 return self._bilinear_loop(w, x, y, z, lambda f, wc, xc: f.derivativeY(wc, xc))
2870 def _derY(self, w, x, y, z):
2871 """
2872 Returns the derivative with respect to y of the interpolated function
2873 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeY.
2874 """
2875 m = len(x)
2876 y_pos, z_pos = self._locate_yz_indices(y, z)
2877 dfdy = np.full(m, np.nan)
2878 for i, j, c in _iter_unique_pairs(y_pos, z_pos):
2879 beta = _cell_fraction(self.z_list, j, z[c])
2880 dfdy[c] = (
2881 (
2882 (1 - beta) * self.wxInterpolators[i][j - 1](w[c], x[c])
2883 + beta * self.wxInterpolators[i][j](w[c], x[c])
2884 )
2885 - (
2886 (1 - beta) * self.wxInterpolators[i - 1][j - 1](w[c], x[c])
2887 + beta * self.wxInterpolators[i - 1][j](w[c], x[c])
2888 )
2889 ) / (self.y_list[i] - self.y_list[i - 1])
2890 return dfdy
2892 def _derZ(self, w, x, y, z):
2893 """
2894 Returns the derivative with respect to z of the interpolated function
2895 at each value in w,x,y,z. Only called internally by HARKinterpolator4D.derivativeZ.
2896 """
2897 m = len(x)
2898 y_pos, z_pos = self._locate_yz_indices(y, z)
2899 dfdz = np.full(m, np.nan)
2900 for i, j, c in _iter_unique_pairs(y_pos, z_pos):
2901 alpha = _cell_fraction(self.y_list, i, y[c])
2902 dfdz[c] = (
2903 (
2904 (1 - alpha) * self.wxInterpolators[i - 1][j](w[c], x[c])
2905 + alpha * self.wxInterpolators[i][j](w[c], x[c])
2906 )
2907 - (
2908 (1 - alpha) * self.wxInterpolators[i - 1][j - 1](w[c], x[c])
2909 + alpha * self.wxInterpolators[i][j - 1](w[c], x[c])
2910 )
2911 ) / (self.z_list[j] - self.z_list[j - 1])
2912 return dfdz
2915class Curvilinear2DInterp(HARKinterpolator2D):
2916 """
2917 A 2D interpolation method for curvilinear or "warped grid" interpolation, as
2918 in White (2015). Used for models with two endogenous states that are solved
2919 with the endogenous grid method. Allows multiple function outputs, but all of
2920 the interpolated functions must share a common curvilinear grid.
2922 Parameters
2923 ----------
2924 f_values: numpy.array or [numpy.array]
2925 One or more 2D arrays of function values such that f_values[i,j] =
2926 f(x_values[i,j],y_values[i,j]).
2927 x_values: numpy.array
2928 A 2D array of x values of the same shape as f_values.
2929 y_values: numpy.array
2930 A 2D array of y values of the same shape as f_values.
2931 """
2933 distance_criteria = ["f_values", "x_values", "y_values"]
2935 def __init__(self, f_values, x_values, y_values):
2936 self.multi = isinstance(f_values, list)
2937 f_list = f_values if self.multi else [f_values]
2938 my_shape = x_values.shape
2939 if my_shape != y_values.shape:
2940 raise ValueError("y_values must have the same shape as x_values!")
2941 prefix = "Each element of f_values" if self.multi else "f_values"
2942 for arr in f_list:
2943 if my_shape != arr.shape:
2944 raise ValueError(f"{prefix} must have the same shape as x_values!")
2946 # Stack as (N_funcs, x_n, y_n) so per-corner indexing vectorizes
2947 # across functions: ``self.f_values[:, x_pos, y_pos]`` returns the
2948 # (N_funcs, M) corner table in one numpy call.
2949 self.f_values = np.stack(f_list)
2950 self.x_values = x_values
2951 self.y_values = y_values
2952 self.x_n, self.y_n = my_shape
2953 self.N_funcs = len(f_list)
2954 self.update_polarity()
2956 def _dispatch(self, x, y, inner):
2957 """Run ``inner`` on flattened ``(x, y)``, reshape per-function results,
2958 and unwrap the single-function case. Shared by ``__call__``,
2959 ``derivativeX``, and ``derivativeY``."""
2960 xa = np.asarray(x)
2961 ya = np.asarray(y)
2962 S = xa.shape
2963 result = inner(xa.flatten(), ya.flatten())
2964 output = [r.reshape(S) for r in result]
2965 return output if self.multi else output[0]
2967 def __call__(self, x, y):
2968 """
2969 Modification of HARKinterpolator2D.__call__ to account for multiple outputs.
2970 """
2971 return self._dispatch(x, y, self._evaluate)
2973 def derivativeX(self, x, y):
2974 """
2975 Modification of HARKinterpolator2D.derivativeX to account for multiple outputs.
2976 """
2977 return self._dispatch(x, y, self._derX)
2979 def derivativeY(self, x, y):
2980 """
2981 Modification of HARKinterpolator2D.derivativeY to account for multiple outputs.
2982 """
2983 return self._dispatch(x, y, self._derY)
2985 def update_polarity(self):
2986 """
2987 Fills in the polarity attribute of the interpolation, determining whether
2988 the "plus" (True) or "minus" (False) solution of the system of equations
2989 should be used for each sector. Needs to be called in __init__.
2991 Parameters
2992 ----------
2993 none
2995 Returns
2996 -------
2997 none
2998 """
2999 # Grab a point known to be inside each sector: the midway point between
3000 # the lower left and upper right vertex of each sector
3001 x_temp = 0.5 * (
3002 self.x_values[0 : (self.x_n - 1), 0 : (self.y_n - 1)]
3003 + self.x_values[1 : self.x_n, 1 : self.y_n]
3004 )
3005 y_temp = 0.5 * (
3006 self.y_values[0 : (self.x_n - 1), 0 : (self.y_n - 1)]
3007 + self.y_values[1 : self.x_n, 1 : self.y_n]
3008 )
3009 size = (self.x_n - 1) * (self.y_n - 1)
3010 x_temp = np.reshape(x_temp, size)
3011 y_temp = np.reshape(y_temp, size)
3012 y_pos = np.tile(np.arange(0, self.y_n - 1), self.x_n - 1)
3013 x_pos = np.reshape(
3014 np.tile(np.arange(0, self.x_n - 1), (self.y_n - 1, 1)).transpose(), size
3015 )
3017 # Set the polarity of all sectors to "plus", then test each sector
3018 self.polarity = np.ones((self.x_n - 1, self.y_n - 1), dtype=bool)
3019 alpha, beta = self.find_coords(x_temp, y_temp, x_pos, y_pos)
3020 polarity = np.logical_and(
3021 np.logical_and(alpha > 0, alpha < 1), np.logical_and(beta > 0, beta < 1)
3022 )
3024 # Update polarity: if (alpha,beta) not in the unit square, then that
3025 # sector must use the "minus" solution instead
3026 self.polarity = np.reshape(polarity, (self.x_n - 1, self.y_n - 1))
3028 def find_sector(self, x, y):
3029 """
3030 Finds the quadrilateral "sector" for each (x,y) point in the input.
3031 Only called as a subroutine of _evaluate(), etc. Uses a numba helper
3032 function below to accelerate computation.
3034 Parameters
3035 ----------
3036 x : np.array
3037 Values whose sector should be found.
3038 y : np.array
3039 Values whose sector should be found. Should be same size as x.
3041 Returns
3042 -------
3043 x_pos : np.array
3044 Sector x-coordinates for each point of the input, of the same size.
3045 y_pos : np.array
3046 Sector y-coordinates for each point of the input, of the same size.
3047 """
3048 x_pos, y_pos = find_sector_numba(x, y, self.x_values, self.y_values)
3049 return x_pos, y_pos
3051 def find_coords(self, x, y, x_pos, y_pos):
3052 """
3053 Calculates the relative coordinates (alpha,beta) for each point (x,y),
3054 given the sectors (x_pos,y_pos) in which they reside. Only called as
3055 a subroutine of _evaluate(), etc. Uses a numba helper function to acc-
3056 elerate computation, and has a "backup method" for when the math fails.
3058 Parameters
3059 ----------
3060 x : np.array
3061 Values whose sector should be found.
3062 y : np.array
3063 Values whose sector should be found. Should be same size as x.
3064 x_pos : np.array
3065 Sector x-coordinates for each point in (x,y), of the same size.
3066 y_pos : np.array
3067 Sector y-coordinates for each point in (x,y), of the same size.
3069 Returns
3070 -------
3071 alpha : np.array
3072 Relative "horizontal" position of the input in their respective sectors.
3073 beta : np.array
3074 Relative "vertical" position of the input in their respective sectors.
3075 """
3076 alpha, beta = find_coords_numba(
3077 x, y, x_pos, y_pos, self.x_values, self.y_values, self.polarity
3078 )
3080 # Alternate method if there are sectors that are "too regular"
3081 # These points weren't able to identify coordinates
3082 z = np.logical_or(np.isnan(alpha), np.isnan(beta))
3083 if np.any(z):
3084 ii = x_pos[z]
3085 jj = y_pos[z]
3086 xA = self.x_values[ii, jj]
3087 xB = self.x_values[ii + 1, jj]
3088 xC = self.x_values[ii, jj + 1]
3089 xD = self.x_values[ii + 1, jj + 1]
3090 yA = self.y_values[ii, jj]
3091 yB = self.y_values[ii + 1, jj]
3092 yC = self.y_values[ii, jj + 1]
3093 # yD = self.y_values[ii + 1, jj + 1]
3094 b = xB - xA
3095 f = yB - yA
3096 kappa = f / b
3097 int_bot = yA - kappa * xA
3098 int_top = yC - kappa * xC
3099 int_these = y[z] - kappa * x[z]
3100 beta_temp = (int_these - int_bot) / (int_top - int_bot)
3101 x_left = beta_temp * xC + (1.0 - beta_temp) * xA
3102 x_right = beta_temp * xD + (1.0 - beta_temp) * xB
3103 alpha_temp = (x[z] - x_left) / (x_right - x_left)
3104 beta[z] = beta_temp
3105 alpha[z] = alpha_temp
3107 return alpha, beta
3109 def _evaluate(self, x, y):
3110 """
3111 Returns the level of the interpolated function at each value in x,y.
3112 Only called internally by __call__ (etc).
3114 Returns an ``(N_funcs, M)`` array of bilinearly interpolated values.
3115 """
3116 x_pos, y_pos = self.find_sector(x, y)
3117 alpha, beta = self.find_coords(x, y, x_pos, y_pos)
3119 alpha_C = 1.0 - alpha
3120 beta_C = 1.0 - beta
3121 wA = alpha_C * beta_C
3122 wB = alpha * beta_C
3123 wC = alpha_C * beta
3124 wD = alpha * beta
3126 # Bilinear interpolation, vectorized over both queries and N_funcs.
3127 return (
3128 wA * self.f_values[:, x_pos, y_pos]
3129 + wB * self.f_values[:, x_pos + 1, y_pos]
3130 + wC * self.f_values[:, x_pos, y_pos + 1]
3131 + wD * self.f_values[:, x_pos + 1, y_pos + 1]
3132 )
3134 def _curvilinear_partials(self, x, y):
3135 """
3136 Compute the inverse Jacobian of the (alpha, beta) -> (x, y) curvilinear
3137 map at each sample point, plus the function-level (dfda, dfdb) arrays.
3139 Returns a 5-tuple ``(x_alpha, x_beta, y_alpha, y_beta, (dfda, dfdb))``
3140 where ``dfda`` and ``dfdb`` are ``(N_funcs, M)`` arrays. Used by both
3141 ``_derX`` and ``_derY``.
3142 """
3143 x_pos, y_pos = self.find_sector(x, y)
3144 alpha, beta = self.find_coords(x, y, x_pos, y_pos)
3146 # Get four corners data for each point
3147 xA = self.x_values[x_pos, y_pos]
3148 xB = self.x_values[x_pos + 1, y_pos]
3149 xC = self.x_values[x_pos, y_pos + 1]
3150 xD = self.x_values[x_pos + 1, y_pos + 1]
3151 yA = self.y_values[x_pos, y_pos]
3152 yB = self.y_values[x_pos + 1, y_pos]
3153 yC = self.y_values[x_pos, y_pos + 1]
3154 yD = self.y_values[x_pos + 1, y_pos + 1]
3156 # Components of the alpha,beta --> x,y delta translation matrix.
3157 alpha_C = 1 - alpha
3158 beta_C = 1 - beta
3159 alpha_x = beta_C * (xB - xA) + beta * (xD - xC)
3160 alpha_y = beta_C * (yB - yA) + beta * (yD - yC)
3161 beta_x = alpha_C * (xC - xA) + alpha * (xD - xB)
3162 beta_y = alpha_C * (yC - yA) + alpha * (yD - yB)
3164 # Invert the delta translation matrix into x,y --> alpha,beta.
3165 det = alpha_x * beta_y - beta_x * alpha_y
3166 x_alpha = beta_y / det
3167 x_beta = -alpha_y / det
3168 y_alpha = -beta_x / det
3169 y_beta = alpha_x / det
3171 # Function corners, vectorized over (N_funcs, M).
3172 fA = self.f_values[:, x_pos, y_pos]
3173 fB = self.f_values[:, x_pos + 1, y_pos]
3174 fC = self.f_values[:, x_pos, y_pos + 1]
3175 fD = self.f_values[:, x_pos + 1, y_pos + 1]
3176 dfda = beta_C * (fB - fA) + beta * (fD - fC)
3177 dfdb = alpha_C * (fC - fA) + alpha * (fD - fB)
3178 return x_alpha, x_beta, y_alpha, y_beta, (dfda, dfdb)
3180 def _derX(self, x, y):
3181 """
3182 Returns the derivative with respect to x of the interpolated function
3183 at each value in x,y. Only called internally by derivativeX.
3184 """
3185 x_alpha, x_beta, _, _, (dfda, dfdb) = self._curvilinear_partials(x, y)
3186 return x_alpha * dfda + x_beta * dfdb
3188 def _derY(self, x, y):
3189 """
3190 Returns the derivative with respect to y of the interpolated function
3191 at each value in x,y. Only called internally by derivativeY.
3192 """
3193 _, _, y_alpha, y_beta, (dfda, dfdb) = self._curvilinear_partials(x, y)
3194 return y_alpha * dfda + y_beta * dfdb
3197# Define a function that checks whether a set of points violates a linear boundary
3198# defined by (x1,y1) and (x2,y2), where the latter is *COUNTER CLOCKWISE* from the
3199# former. Returns 1 if the point is outside the boundary and 0 otherwise.
3200@njit
3201def boundary_check(xq, yq, x1, y1, x2, y2): # pragma: no cover
3202 return int((y2 - y1) * xq - (x2 - x1) * yq > x1 * y2 - y1 * x2)
3205# Define a numba helper function for finding the sector in the irregular grid
3206@njit
3207def find_sector_numba(X_query, Y_query, X_values, Y_values): # pragma: no cover
3208 # Initialize the sector guess
3209 M = X_query.size
3210 x_n = X_values.shape[0]
3211 y_n = X_values.shape[1]
3212 ii = int(x_n / 2)
3213 jj = int(y_n / 2)
3214 top_ii = x_n - 2
3215 top_jj = y_n - 2
3217 # Initialize the output arrays
3218 X_pos = np.empty(M, dtype=np.int32)
3219 Y_pos = np.empty(M, dtype=np.int32)
3221 # Identify the correct sector for each point to be evaluated
3222 max_loops = x_n + y_n
3223 for m in range(M):
3224 found = False
3225 loops = 0
3226 while not found and loops < max_loops:
3227 # Get coordinates for the four vertices: (xA,yA),...,(xD,yD)
3228 x0 = X_query[m]
3229 y0 = Y_query[m]
3230 xA = X_values[ii, jj]
3231 xB = X_values[ii + 1, jj]
3232 xC = X_values[ii, jj + 1]
3233 xD = X_values[ii + 1, jj + 1]
3234 yA = Y_values[ii, jj]
3235 yB = Y_values[ii + 1, jj]
3236 yC = Y_values[ii, jj + 1]
3237 yD = Y_values[ii + 1, jj + 1]
3239 # Check the "bounding box" for the sector: is this guess plausible?
3240 D = int(y0 < np.minimum(yA, yB))
3241 R = int(x0 > np.maximum(xB, xD))
3242 U = int(y0 > np.maximum(yC, yD))
3243 L = int(x0 < np.minimum(xA, xC))
3245 # Check which boundaries are violated (and thus where to look next)
3246 in_box = np.all(np.logical_not(np.array([D, R, U, L])))
3247 if in_box:
3248 D = boundary_check(x0, y0, xA, yA, xB, yB)
3249 R = boundary_check(x0, y0, xB, yB, xD, yD)
3250 U = boundary_check(x0, y0, xD, yD, xC, yC)
3251 L = boundary_check(x0, y0, xC, yC, xA, yA)
3253 # Update the sector guess based on the violations
3254 ii_next = np.maximum(np.minimum(ii - L + R, top_ii), 0)
3255 jj_next = np.maximum(np.minimum(jj - D + U, top_jj), 0)
3257 # Check whether sector guess changed and go to next iteration
3258 found = (ii == ii_next) and (jj == jj_next)
3259 ii = ii_next
3260 jj = jj_next
3261 loops += 1
3263 # Put the final sector guess into the output array
3264 X_pos[m] = ii
3265 Y_pos[m] = jj
3267 # Return the output
3268 return X_pos, Y_pos
3271# Define a numba helper function for finding relative coordinates within sector
3272@njit
3273def find_coords_numba(
3274 X_query, Y_query, X_pos, Y_pos, X_values, Y_values, polarity
3275): # pragma: no cover
3276 M = X_query.size
3277 alpha = np.empty(M)
3278 beta = np.empty(M)
3280 # Calculate relative coordinates in the sector for each point
3281 for m in range(M):
3282 try:
3283 x0 = X_query[m]
3284 y0 = Y_query[m]
3285 ii = X_pos[m]
3286 jj = Y_pos[m]
3287 xA = X_values[ii, jj]
3288 xB = X_values[ii + 1, jj]
3289 xC = X_values[ii, jj + 1]
3290 xD = X_values[ii + 1, jj + 1]
3291 yA = Y_values[ii, jj]
3292 yB = Y_values[ii + 1, jj]
3293 yC = Y_values[ii, jj + 1]
3294 yD = Y_values[ii + 1, jj + 1]
3295 p = 2.0 * polarity[ii, jj] - 1.0
3296 a = xA
3297 b = xB - xA
3298 c = xC - xA
3299 d = xA - xB - xC + xD
3300 e = yA
3301 f = yB - yA
3302 g = yC - yA
3303 h = yA - yB - yC + yD
3304 denom = d * g - h * c
3305 mu = (h * b - d * f) / denom
3306 tau = (h * (a - x0) - d * (e - y0)) / denom
3307 zeta = a - x0 + c * tau
3308 eta = b + c * mu + d * tau
3309 theta = d * mu
3310 alph = (-eta + p * np.sqrt(eta**2 - 4 * zeta * theta)) / (2 * theta)
3311 bet = mu * alph + tau
3312 except Exception:
3313 alph = np.nan
3314 bet = np.nan
3315 alpha[m] = alph
3316 beta[m] = bet
3318 return alpha, beta
3321class DiscreteInterp(MetricObject):
3322 """
3323 An interpolator for variables that can only take a discrete set of values.
3325 If the function we wish to interpolate, f(args) can take on the list of
3326 values discrete_vals, this class expects an interpolator for the index of
3327 f's value in discrete_vals.
3328 E.g., if f(a,b,c) = discrete_vals[5], then index_interp(a,b,c) = 5.
3330 Parameters
3331 ----------
3332 index_interp: HARKInterpolator
3333 An interpolator giving an approximation to the index of the value in
3334 discrete_vals that corresponds to a given set of arguments.
3335 discrete_vals: numpy.array
3336 A 1D array containing the values in the range of the discrete function
3337 to be interpolated.
3338 """
3340 distance_criteria = ["index_interp"]
3342 def __init__(self, index_interp, discrete_vals):
3343 self.index_interp = index_interp
3344 self.discrete_vals = discrete_vals
3345 self.n_vals = len(self.discrete_vals)
3347 def __call__(self, *args):
3348 # Interpolate indices and round to integers
3349 inds = np.rint(self.index_interp(*args)).astype(int)
3350 if type(inds) is not np.ndarray:
3351 inds = np.array(inds)
3352 # Deal with out-of range indices
3353 inds[inds < 0] = 0
3354 inds[inds >= self.n_vals] = self.n_vals - 1
3356 # Get values from grid
3357 return self.discrete_vals[inds]
3360class IndexedInterp(MetricObject):
3361 """
3362 An interpolator for functions whose first argument is an integer-valued index.
3363 Constructor takes in a list of functions as its only argument. When evaluated
3364 at f(i,X), interpolator returns f[i](X), where X can be any number of inputs.
3365 This simply provides a different interface for accessing the same functions.
3367 Parameters
3368 ----------
3369 functions : [Callable]
3370 List of one or more functions to be indexed.
3371 """
3373 distance_criteria = ["functions"]
3375 def __init__(self, functions):
3376 self.functions = functions
3377 self.N = len(functions)
3379 def __call__(self, idx, *args):
3380 out = np.empty(idx.shape)
3381 out.fill(np.nan)
3383 for n in range(self.N):
3384 these = idx == n
3385 if not np.any(these):
3386 continue
3387 temp = [arg[these] for arg in args]
3388 out[these] = self.functions[n](*temp)
3389 return out
3392###############################################################################
3393## Functions used in discrete choice models with T1EV taste shocks ############
3394###############################################################################
3397def _log_sum_taste_shock(Vals, sigma):
3398 """Stabilized log-sum-exp under a T1EV taste shock with scale ``sigma``.
3400 Returns ``maxV + sigma * log(sum_j exp((V_j - maxV) / sigma))``.
3401 Caller must ensure ``sigma != 0``.
3402 """
3403 maxV = np.max(Vals, axis=0)
3404 sumexp = np.sum(np.exp((Vals - maxV) / sigma), axis=0)
3405 return maxV + sigma * np.log(sumexp)
3408def calc_log_sum_choice_probs(Vals, sigma):
3409 """
3410 Returns the final optimal value and choice probabilities given the choice
3411 specific value functions `Vals`. Probabilities are degenerate if sigma == 0.0.
3412 Parameters
3413 ----------
3414 Vals : [numpy.array]
3415 A numpy.array that holds choice specific values at common grid points.
3416 sigma : float
3417 A number that controls the variance of the taste shocks
3418 Returns
3419 -------
3420 V : [numpy.array]
3421 A numpy.array that holds the integrated value function.
3422 P : [numpy.array]
3423 A numpy.array that holds the discrete choice probabilities
3424 """
3425 # Assumes that NaNs have been replaced by -numpy.inf or similar
3426 if sigma == 0.0:
3427 Pflat = np.argmax(Vals, axis=0)
3428 V = np.max(Vals, axis=0)
3429 Probs = np.zeros(Vals.shape)
3430 np.put_along_axis(Probs, Pflat[None, ...], 1, axis=0)
3431 return V, Probs
3433 LogSumV = _log_sum_taste_shock(Vals, sigma)
3434 Probs = np.exp((Vals - LogSumV) / sigma)
3435 return LogSumV, Probs
3438def calc_choice_probs(Vals, sigma):
3439 """
3440 Returns the choice probabilities given the choice specific value functions
3441 `Vals`. Probabilities are degenerate if sigma == 0.0.
3442 Parameters
3443 ----------
3444 Vals : [numpy.array]
3445 A numpy.array that holds choice specific values at common grid points.
3446 sigma : float
3447 A number that controls the variance of the taste shocks
3448 Returns
3449 -------
3450 Probs : [numpy.array]
3451 A numpy.array that holds the discrete choice probabilities
3452 """
3454 # Assumes that NaNs have been replaced by -numpy.inf or similar
3455 if sigma == 0.0:
3456 Pflat = np.argmax(Vals, axis=0)
3457 Probs = np.zeros(Vals.shape)
3458 np.put_along_axis(Probs, Pflat[None, ...], 1, axis=0)
3459 return Probs
3461 maxV = np.max(Vals, axis=0)
3462 weights = np.exp((Vals - maxV) / sigma)
3463 return weights / np.sum(weights, axis=0)
3466def calc_log_sum(Vals, sigma):
3467 """
3468 Returns the optimal value given the choice specific value functions Vals.
3469 Parameters
3470 ----------
3471 Vals : [numpy.array]
3472 A numpy.array that holds choice specific values at common grid points.
3473 sigma : float
3474 A number that controls the variance of the taste shocks
3475 Returns
3476 -------
3477 V : [numpy.array]
3478 A numpy.array that holds the integrated value function.
3479 """
3480 # Assumes that NaNs have been replaced by -numpy.inf or similar
3481 if sigma == 0.0:
3482 return np.amax(Vals, axis=0)
3483 return _log_sum_taste_shock(Vals, sigma)
3486###############################################################################
3487# Tools for value and marginal-value functions in models where #
3488# - dvdm = u'(c). #
3489# - u is of the CRRA family. #
3490###############################################################################
3493class ValueFuncCRRA(MetricObject):
3494 """
3495 A class for representing a value function. The underlying interpolation is
3496 in the space of (state,u_inv(v)); this class "re-curves" to the value function.
3498 Parameters
3499 ----------
3500 vFuncNvrs : function
3501 A real function representing the value function composed with the
3502 inverse utility function, defined on the state: u_inv(vFunc(state))
3503 CRRA : float
3504 Coefficient of relative risk aversion.
3505 illegal_value : float, optional
3506 If provided, value to return for "out-of-bounds" inputs that return NaN
3507 from the pseudo-inverse value function. Most common choice is -np.inf,
3508 which makes the outcome infinitely bad.
3509 """
3511 distance_criteria = ["func", "CRRA"]
3513 def __init__(self, vFuncNvrs, CRRA, illegal_value=None):
3514 self.vFuncNvrs = deepcopy(vFuncNvrs)
3515 self.CRRA = CRRA
3516 self.illegal_value = illegal_value
3518 if hasattr(vFuncNvrs, "grid_list"):
3519 self.grid_list = vFuncNvrs.grid_list
3520 else:
3521 self.grid_list = None
3523 def __call__(self, *vFuncArgs):
3524 """
3525 Evaluate the value function at given levels of market resources m.
3527 Parameters
3528 ----------
3529 vFuncArgs : floats or np.arrays, all of the same dimensions.
3530 Values for the state variables. These usually start with 'm',
3531 market resources normalized by the level of permanent income.
3533 Returns
3534 -------
3535 v : float or np.array
3536 Lifetime value of beginning this period with the given states; has
3537 same size as the state inputs.
3538 """
3539 temp = self.vFuncNvrs(*vFuncArgs)
3540 v = CRRAutility(temp, self.CRRA)
3541 if self.illegal_value is not None:
3542 illegal = np.isnan(temp)
3543 v[illegal] = self.illegal_value
3544 return v
3546 def gradient(self, *args):
3547 # V(s) = u(vFuncNvrs(s)), so by the chain rule
3548 # dV/ds_i = u'(vFuncNvrs(s)) * d vFuncNvrs / ds_i.
3549 NvrsGrad = self.vFuncNvrs.gradient(*args)
3550 marg_u = CRRAutilityP(self.vFuncNvrs(*args), self.CRRA)
3551 grad = [g * marg_u for g in NvrsGrad]
3552 return grad
3554 def _eval_and_grad(self, *args):
3555 return (self.__call__(*args), self.gradient(*args))
3558def _eval_c_and_mpc(cFunc, *cFuncArgs):
3559 """Return ``(c, MPC)`` from ``cFunc`` regardless of whether it exposes
3560 ``eval_with_derivative`` (1D) or a ``derivativeX`` attribute (multi-D)."""
3561 if isinstance(cFunc, HARKinterpolator1D):
3562 return cFunc.eval_with_derivative(*cFuncArgs)
3563 if hasattr(cFunc, "derivativeX"):
3564 return cFunc(*cFuncArgs), cFunc.derivativeX(*cFuncArgs)
3565 raise TypeError(
3566 "cFunc does not have a 'derivativeX' attribute. Can't compute "
3567 "marginal marginal value."
3568 )
3571class MargValueFuncCRRA(MetricObject):
3572 """
3573 A class for representing a marginal value function in models where the
3574 standard envelope condition of dvdm(state) = u'(c(state)) holds (with CRRA utility).
3576 Parameters
3577 ----------
3578 cFunc : function.
3579 Its first argument must be normalized market resources m.
3580 A real function representing the marginal value function composed
3581 with the inverse marginal utility function, defined on the state
3582 variables: uP_inv(dvdmFunc(state)). Called cFunc because when standard
3583 envelope condition applies, uP_inv(dvdm(state)) = cFunc(state).
3584 CRRA : float
3585 Coefficient of relative risk aversion.
3586 """
3588 distance_criteria = ["cFunc", "CRRA"]
3590 def __init__(self, cFunc, CRRA):
3591 self.cFunc = deepcopy(cFunc)
3592 self.CRRA = CRRA
3594 if hasattr(cFunc, "grid_list"):
3595 self.grid_list = cFunc.grid_list
3596 else:
3597 self.grid_list = None
3599 def __call__(self, *cFuncArgs):
3600 """
3601 Evaluate the marginal value function at given levels of market resources m.
3603 Parameters
3604 ----------
3605 cFuncArgs : floats or np.arrays
3606 Values of the state variables at which to evaluate the marginal
3607 value function.
3609 Returns
3610 -------
3611 vP : float or np.array
3612 Marginal lifetime value of beginning this period with state
3613 cFuncArgs
3614 """
3615 return CRRAutilityP(self.cFunc(*cFuncArgs), rho=self.CRRA)
3617 def derivativeX(self, *cFuncArgs):
3618 """
3619 Evaluate the derivative of the marginal value function with respect to
3620 market resources at given state; this is the marginal marginal value
3621 function.
3623 Parameters
3624 ----------
3625 cFuncArgs : floats or np.arrays
3626 State variables.
3628 Returns
3629 -------
3630 vPP : float or np.array
3631 Marginal marginal lifetime value of beginning this period with
3632 state cFuncArgs; has same size as inputs.
3634 """
3636 c, MPC = _eval_c_and_mpc(self.cFunc, *cFuncArgs)
3637 return MPC * CRRAutilityPP(c, rho=self.CRRA)
3640class MargMargValueFuncCRRA(MetricObject):
3641 """
3642 A class for representing a marginal marginal value function in models where
3643 the standard envelope condition of dvdm = u'(c(state)) holds (with CRRA utility).
3645 Parameters
3646 ----------
3647 cFunc : function.
3648 Its first argument must be normalized market resources m.
3649 A real function representing the marginal value function composed
3650 with the inverse marginal utility function, defined on the state
3651 variables: uP_inv(dvdmFunc(state)). Called cFunc because when standard
3652 envelope condition applies, uP_inv(dvdm(state)) = cFunc(state).
3653 CRRA : float
3654 Coefficient of relative risk aversion.
3655 """
3657 distance_criteria = ["cFunc", "CRRA"]
3659 def __init__(self, cFunc, CRRA):
3660 self.cFunc = deepcopy(cFunc)
3661 self.CRRA = CRRA
3663 def __call__(self, *cFuncArgs):
3664 """
3665 Evaluate the marginal marginal value function at given levels of market
3666 resources m.
3668 Parameters
3669 ----------
3670 m : float or np.array
3671 Market resources (normalized by permanent income) whose marginal
3672 marginal value is to be found.
3674 Returns
3675 -------
3676 vPP : float or np.array
3677 Marginal marginal lifetime value of beginning this period with market
3678 resources m; has same size as input m.
3679 """
3681 c, MPC = _eval_c_and_mpc(self.cFunc, *cFuncArgs)
3682 return MPC * CRRAutilityPP(c, rho=self.CRRA)