Coverage for HARK / utilities.py: 91%
265 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-05-10 06:19 +0000
1"""
2General purpose / miscellaneous functions. Includes functions to approximate
3continuous distributions with discrete ones, utility functions (and their
4derivatives), manipulation of discrete distributions, and basic plotting tools.
5"""
7import cProfile
8import os
9import pstats
10import re
12import numba
13import numpy as np # Python's numeric library, abbreviated "np"
14from scipy.interpolate import interp1d
16from inspect import signature
19class get_it_from:
20 """
21 Class whose instances act as a special case trivial constructor that merely
22 grabs an attribute or entry from the named attribute. This is useful when
23 there are constructed model inputs that are "built together". Simply have a
24 constructor that makes a dictionary (or object) containing the several inputs,
25 then use get_it_from(that_dict_name) as the constructor for each of them.
27 Parameters
28 ----------
29 name : str
30 Name of the parent dictionary or object from which to take the object.
31 """
33 def __init__(self, name):
34 self.name = name
36 def __call__(self, parent, query):
37 if isinstance(parent, (int, float, bool, complex, str)):
38 return parent # the desired result is the thing itself
39 elif isinstance(parent, dict):
40 return parent[query]
41 else:
42 return getattr(parent, query)
45# ==============================================================================
46# ============== Some basic function tools ====================================
47# ==============================================================================
50def get_arg_names(function):
51 """
52 Returns a list of strings naming all of the arguments for the passed function.
54 Parameters
55 ----------
56 function : function
57 A function whose argument names are wanted.
59 Returns
60 -------
61 argNames : [string]
62 The names of the arguments of function.
63 """
64 argCount = function.__code__.co_argcount
65 argNames = function.__code__.co_varnames[:argCount]
66 return argNames
69class NullFunc:
70 """
71 A trivial class that acts as a placeholder "do nothing" function.
72 """
74 def __call__(self, *args):
75 """
76 Returns meaningless output no matter what the input(s) is. If no input,
77 returns None. Otherwise, returns an array of NaNs (or a single NaN) of
78 the same size as the first input.
79 """
80 if len(args) == 0:
81 return None
82 else:
83 arg = args[0]
84 if hasattr(arg, "shape"):
85 return np.zeros_like(arg) + np.nan
86 else:
87 return np.nan
89 def distance(self, other):
90 """
91 Trivial distance metric that only cares whether the other object is also
92 an instance of NullFunc. Intentionally does not inherit from HARKobject
93 as this might create dependency problems.
95 Parameters
96 ----------
97 other : any
98 Any object for comparison to this instance of NullFunc.
100 Returns
101 -------
102 (unnamed) : float
103 The distance between self and other. Returns 0 if other is also a
104 NullFunc; otherwise returns an arbitrary high number.
105 """
106 if other.__class__ is self.__class__:
107 return 0.0
108 else:
109 return 1000.0
112def apply_fun_to_vals(fun, vals):
113 """
114 Applies a function to the arguments defined in `vals`.
115 This is equivalent to `fun(**vals)`, except
116 that `vals` may contain keys that are not named arguments
117 of `fun`.
119 Parameters
120 ----------
121 fun: callable
123 vals: dict
124 """
125 return fun(*[vals[var] for var in signature(fun).parameters])
128# =======================================================
129# ================ Other useful functions ===============
130# =======================================================
133def make_assets_grid(aXtraMin, aXtraMax, aXtraCount, aXtraExtra, aXtraNestFac):
134 """
135 Constructs the base grid of post-decision states, representing end-of-period
136 assets above the absolute minimum. Has three modes, depending on aXtraNestFac:
138 aXtraNestFac = -1 : Uniformly spaced grid.
139 aXtraNestFac = 0 : Ordinary exponentially spaced grid.
140 aXtraNestFac >= 1 : Multi-exponentially nested grid.
142 See :func:`HARK.utilities.make_grid_exp_mult` for more info
144 Parameters
145 ----------
146 aXtraMin: float
147 Minimum value for the assets-above-minimum grid.
148 aXtraMax: float
149 Maximum value for the assets-above-minimum grid.
150 aXtraCount: int
151 Number of nodes in the assets-above-minimum grid, not counting extra values.
152 aXtraExtra: [float]
153 Additional values to insert in the assets-above-minimum grid.
154 aXtraNestFac: int
155 Level of exponential nesting for grid. If -1, the grid is linearly spaced.
157 Returns
158 -------
159 aXtraGrid: np.ndarray
160 Base array of values for the post-decision-state grid.
161 """
162 # Set up post decision state grid:
163 if aXtraNestFac == -1:
164 aXtraGrid = np.linspace(aXtraMin, aXtraMax, aXtraCount)
165 elif (aXtraNestFac >= 0) and type(aXtraNestFac) is int:
166 aXtraGrid = make_grid_exp_mult(
167 ming=aXtraMin, maxg=aXtraMax, ng=aXtraCount, timestonest=aXtraNestFac
168 )
169 else:
170 raise ValueError(
171 "Please ensure aXtraNestFac is either -1 or a positive integer."
172 )
174 # Add in additional points for the grid:
175 if aXtraExtra is not None:
176 temp_list = []
177 for i in aXtraExtra:
178 if i is not None:
179 temp_list.append(i)
180 aXtraGrid = np.sort(np.unique(np.concatenate((aXtraGrid, temp_list))))
182 return aXtraGrid
185# ==============================================================================
186# ============== Functions for generating state space grids ===================
187# ==============================================================================
190def make_grid_exp_mult(ming, maxg, ng, timestonest=20, offset=0.0):
191 r"""
192 Makes a multi-exponentially spaced grid.
193 If the function :math:`\ln(1+x)` were applied timestonest times, the grid would
194 become linearly spaced. If timestonest is 0, the grid is exponentially spaced.
195 If timestonest is -1, the grid is linearly spaced.
197 NOTE: The bounds of the grid must be non-negative, else this function will
198 return an invalid grid with NaNs in it. If you want a non-linearly spaced
199 grid that spans negative numbers, use make_polynomial_grid or specify an offset.
201 Parameters
202 ----------
203 ming : float
204 Minimum value of the grid, which must be non-negative.
205 maxg : float
206 Maximum value of the grid, which must be greater than ming.
207 ng : int
208 The number of gridpoints
209 timestonest : int, optional
210 The number of times to nest the exponentiation; the default is 20.
211 offset : float
212 Offset added to the final grid, so it spans [ming+offset, maxg+offset].
213 The default is zero.
215 Returns
216 -------
217 points : np.array
218 A multi-exponentially spaced grid.
220 Notes
221 -----
222 Original Matab code can be found in Chris Carroll's
223 [Solution Methods for Microeconomic Dynamic Optimization Problems]
224 (https://www.econ2.jhu.edu/people/ccarroll/solvingmicrodsops/) toolkit.
225 """
226 if timestonest == -1:
227 grid = np.linspace(ming, maxg, ng) + offset
228 return grid
229 if timestonest > 0:
230 Lming = ming
231 Lmaxg = maxg
232 for j in range(timestonest):
233 Lming = np.log(Lming + 1)
234 Lmaxg = np.log(Lmaxg + 1)
235 Lgrid = np.linspace(Lming, Lmaxg, ng)
236 grid = Lgrid
237 for j in range(timestonest):
238 grid = np.exp(grid) - 1
239 else:
240 Lming = np.log(ming)
241 Lmaxg = np.log(maxg)
242 Lgrid = np.linspace(Lming, Lmaxg, ng)
243 grid = np.exp(Lgrid)
244 grid += offset
245 return grid
248def make_polynomial_grid(ming, maxg, ng, order=1.0):
249 """
250 Construct a polynomially spaced grid with chosen polynomial order.
251 A uniformly spaced grid on [0,1] is raised to the chosen order, then linearly
252 remapped to the specified interval. Supports any real valued grid bounds.
254 Parameters
255 ----------
256 ming : float
257 Lower bound of grid.
258 maxg : float
259 Upper bound of grid.
260 ng : int
261 Number of points in the grid.
262 order : float, optional
263 Polynomial spacing order for the grid. The default is 1.0, or linear.
265 Returns
266 -------
267 grid : np.array
268 Polynomial spaced grid on [ming, maxg] with ng points.
269 """
270 grid = np.linspace(0.0, 1.0, ng) ** order * (maxg - ming) + ming
271 return grid
274# ==============================================================================
275# ============== Uncategorized general functions ===================
276# ==============================================================================
279def _validate_percentiles(percentiles):
280 """Default to ``[0.5]`` if ``None``; otherwise validate and return."""
281 if percentiles is None:
282 return [0.5]
283 if (
284 not isinstance(percentiles, (list, np.ndarray))
285 or min(percentiles) <= 0
286 or max(percentiles) >= 1
287 ):
288 raise ValueError(
289 "Percentiles should be a list or numpy array of floats between 0 and 1"
290 )
291 return percentiles
294def _sort_and_cum_dist(data, weights, presorted):
295 """Sort ``data``/``weights`` if needed and return ``(data_sorted,
296 weights_sorted, cum_dist)`` where ``cum_dist`` is the normalized cumulative
297 weight."""
298 if weights is None:
299 weights = np.ones(data.size)
300 if presorted:
301 data_sorted = data
302 weights_sorted = weights
303 else:
304 order = np.argsort(data)
305 data_sorted = data[order]
306 weights_sorted = weights[order]
307 cum_dist = np.cumsum(weights_sorted) / np.sum(weights_sorted)
308 return data_sorted, weights_sorted, cum_dist
311def get_percentiles(data, weights=None, percentiles=None, presorted=False):
312 """
313 Calculates the requested percentiles of (weighted) data. Median by default.
315 Parameters
316 ----------
317 data : numpy.array
318 A 1D array of float data.
319 weights : np.array
320 A weighting vector for the data.
321 percentiles : [float]
322 A list or numpy.array of percentiles to calculate for the data. Each element should
323 be in (0,1).
324 presorted : boolean
325 Indicator for whether data has already been sorted.
327 Returns
328 -------
329 pctl_out : numpy.array
330 The requested percentiles of the data.
331 """
332 percentiles = _validate_percentiles(percentiles)
333 if data.size < 2:
334 return np.zeros(np.array(percentiles).shape) + np.nan
335 data_sorted, _, cum_dist = _sort_and_cum_dist(data, weights, presorted)
336 inv_CDF = interp1d(cum_dist, data_sorted, bounds_error=False, assume_sorted=True)
337 return inv_CDF(percentiles)
340def get_lorenz_shares(data, weights=None, percentiles=None, presorted=False):
341 """
342 Calculates the Lorenz curve at the requested percentiles of (weighted) data.
343 Median by default.
345 Parameters
346 ----------
347 data : numpy.array
348 A 1D array of float data.
349 weights : numpy.array
350 A weighting vector for the data.
351 percentiles : [float]
352 A list or numpy.array of percentiles to calculate for the data. Each element should
353 be in (0,1).
354 presorted : boolean
355 Indicator for whether data has already been sorted.
357 Returns
358 -------
359 lorenz_out : numpy.array
360 The requested Lorenz curve points of the data.
361 """
362 percentiles = _validate_percentiles(percentiles)
363 data_sorted, weights_sorted, cum_dist = _sort_and_cum_dist(data, weights, presorted)
364 temp = data_sorted * weights_sorted
365 cum_data = np.cumsum(temp) / sum(temp)
366 lorenzFunc = interp1d(cum_dist, cum_data, bounds_error=False, assume_sorted=True)
367 return lorenzFunc(percentiles)
370def calc_subpop_avg(data, reference, cutoffs, weights=None):
371 """
372 Calculates the average of (weighted) data between cutoff percentiles of a
373 reference variable.
375 Parameters
376 ----------
377 data : numpy.array
378 A 1D array of float data.
379 reference : numpy.array
380 A 1D array of float data of the same length as data.
381 cutoffs : [(float,float)]
382 A list of doubles with the lower and upper percentile bounds (should be
383 in [0,1]).
384 weights : numpy.array
385 A weighting vector for the data.
387 Returns
388 -------
389 slice_avg
390 The (weighted) average of data that falls within the cutoff percentiles
391 of reference.
393 """
394 if weights is None: # Set equiprobable weights if none were given
395 weights = np.ones(data.size)
397 # Sort the data and generate a cumulative distribution
398 order = np.argsort(reference)
399 data_sorted = data[order]
400 weights_sorted = weights[order]
401 cum_dist = np.cumsum(weights_sorted) / np.sum(weights_sorted)
403 # For each set of cutoffs, calculate the average of data that falls within
404 # the cutoff percentiles of reference
405 slice_avg = []
406 for j in range(len(cutoffs)):
407 bot = np.searchsorted(cum_dist, cutoffs[j][0])
408 top = np.searchsorted(cum_dist, cutoffs[j][1])
409 slice_avg.append(
410 np.sum(data_sorted[bot:top] * weights_sorted[bot:top])
411 / np.sum(weights_sorted[bot:top])
412 )
413 return slice_avg
416def epanechnikov_kernel(x, ref_x, h=1.0):
417 """
418 The Epanechnikov kernel, which has been shown to be the most efficient kernel
419 for non-parametric regression.
421 Parameters
422 ----------
423 x : np.array
424 Values at which to evaluate the kernel
425 x_ref : float
426 The reference point
427 h : float
428 Kernel bandwidth
430 Returns
431 -------
432 out : np.array
433 Kernel values at each value of x
434 """
435 u = (x - ref_x) / h # Normalize distance by bandwidth
436 out = 0.75 * (1.0 - u**2)
437 out[np.abs(u) > 1.0] = 0.0 # Kernel = 0 outside [-1,1]
438 return out
441def triangle_kernel(x, ref_x, h=1.0):
442 """
443 The triangle or "hat" kernel.
445 Parameters
446 ----------
447 x : np.array
448 Values at which to evaluate the kernel
449 x_ref : float
450 The reference point
451 h : float
452 Kernel bandwidth
454 Returns
455 -------
456 out : np.array
457 Kernel values at each value of x
458 """
459 u = (x - ref_x) / h # Normalize distance by bandwidth
460 these = np.abs(u) <= 1.0 # Kernel = 0 outside [-1,1]
461 out = np.zeros_like(u) # Initialize kernel output
462 out[these] = 1.0 - np.abs(u[these]) # Evaluate kernel
463 return out
466kernel_dict = {
467 "epanechnikov": epanechnikov_kernel,
468 "triangle": triangle_kernel,
469 "hat": triangle_kernel,
470}
473def kernel_regression(x, y, bot=None, top=None, N=500, h=None, kernel="epanechnikov"):
474 """
475 Performs a non-parametric Nadaraya-Watson 1D kernel regression on given data
476 with optionally specified range, number of points, and kernel bandwidth.
478 Parameters
479 ----------
480 x : np.array
481 The independent variable in the kernel regression.
482 y : np.array
483 The dependent variable in the kernel regression.
484 bot : float
485 Minimum value of interest in the regression; defaults to min(x).
486 top : float
487 Maximum value of interest in the regression; defaults to max(y).
488 N : int
489 Number of points to compute.
490 h : float
491 The bandwidth of the (Epanechnikov) kernel. To-do: GENERALIZE.
493 Returns
494 -------
495 regression : LinearInterp
496 A piecewise locally linear kernel regression: y = f(x).
497 """
498 # Fix omitted inputs
499 if bot is None:
500 bot = np.min(x)
501 if top is None:
502 top = np.max(x)
503 if h is None:
504 h = 2.0 * (top - bot) / float(N) # This is an arbitrary default
506 # Get kernel if possible
507 try:
508 kern = kernel_dict[kernel]
509 except KeyError:
510 raise ValueError(f"Can't find a kernel named '{kernel}'!") from None
512 # Construct a local linear approximation
513 x_vec = np.linspace(bot, top, num=N)
514 # Evaluate the kernel for all evaluation points at once
515 weights = kern(x[:, None], x_vec[None, :], h)
516 weight_sums = np.sum(weights, axis=0)
517 # Avoid division by zero when weights are extremely small
518 weight_sums[weight_sums == 0] = np.nan
519 y_vec = np.dot(weights.T, y) / weight_sums
520 regression = interp1d(x_vec, y_vec, bounds_error=False, assume_sorted=True)
521 return regression
524def make_polynomial_params(coeffs, T, offset=0.0, step=1.0):
525 """
526 Make a T-length array of parameters using polynomial coefficients.
528 Parameters
529 ----------
530 coeffs : [float]
531 Arbitrary length list of polynomial coefficients.
532 T : int
533 Number of elements in the output.
534 offset : float, optional
535 Value at which the X values should start. The default is 0.0.
536 step : float, optional
537 Increment of the X values. The default is 1.0.
539 Returns
540 -------
541 param_vals : np.array
542 T-length array of parameter values calculated using the polynomial coefficients.
543 """
544 X = offset + step * np.arange(T)
545 # np.polyval expects highest power first
546 return np.polyval(coeffs[::-1], X)
549@numba.njit
550def jump_to_grid_1D(m_vals, probs, Dist_mGrid): # pragma: nocover
551 """
552 Distributes values onto a predefined grid, maintaining the means.
555 Parameters
556 ----------
557 m_vals: np.array
558 Market resource values
560 probs: np.array
561 Shock probabilities associated with combinations of m_vals.
562 Can be thought of as the probability mass function of (m_vals).
564 dist_mGrid : np.array
565 Grid over normalized market resources
567 Returns
568 -------
569 probGrid.flatten(): np.array
570 Probabilities of each gridpoint on the combined grid of market resources
572 """
574 probGrid = np.zeros(len(Dist_mGrid))
575 mIndex = np.digitize(m_vals, Dist_mGrid) - 1
576 mIndex[m_vals <= Dist_mGrid[0]] = -1
577 mIndex[m_vals >= Dist_mGrid[-1]] = len(Dist_mGrid) - 1
579 for i in range(len(m_vals)):
580 if mIndex[i] == -1:
581 mlowerIndex = 0
582 mupperIndex = 0
583 mlowerWeight = 1.0
584 mupperWeight = 0.0
585 elif mIndex[i] == len(Dist_mGrid) - 1:
586 mlowerIndex = -1
587 mupperIndex = -1
588 mlowerWeight = 1.0
589 mupperWeight = 0.0
590 else:
591 mlowerIndex = mIndex[i]
592 mupperIndex = mIndex[i] + 1
593 mlowerWeight = (Dist_mGrid[mupperIndex] - m_vals[i]) / (
594 Dist_mGrid[mupperIndex] - Dist_mGrid[mlowerIndex]
595 )
596 mupperWeight = 1.0 - mlowerWeight
598 probGrid[mlowerIndex] += probs[i] * mlowerWeight
599 probGrid[mupperIndex] += probs[i] * mupperWeight
601 return probGrid.flatten()
604@numba.njit
605def jump_to_grid_2D(
606 m_vals, perm_vals, probs, dist_mGrid, dist_pGrid
607): # pragma: nocover
608 """
609 Distributes values onto a predefined grid, maintaining the means. m_vals and perm_vals are realizations of market resources and permanent income while
610 dist_mGrid and dist_pGrid are the predefined grids of market resources and permanent income, respectively. That is, m_vals and perm_vals do not necesarily lie on their
611 respective grids. Returns probabilities of each gridpoint on the combined grid of market resources and permanent income.
614 Parameters
615 ----------
616 m_vals: np.array
617 Market resource values
619 perm_vals: np.array
620 Permanent income values
622 probs: np.array
623 Shock probabilities associated with combinations of m_vals and perm_vals.
624 Can be thought of as the probability mass function of (m_vals, perm_vals).
626 dist_mGrid : np.array
627 Grid over normalized market resources
629 dist_pGrid : np.array
630 Grid over permanent income
631 Returns
632 -------
633 probGrid.flatten(): np.array
634 Probabilities of each gridpoint on the combined grid of market resources and permanent income
635 """
637 probGrid = np.zeros((len(dist_mGrid), len(dist_pGrid)))
639 # Maybe use np.searchsorted as opposed to np.digitize
640 mIndex = (
641 np.digitize(m_vals, dist_mGrid) - 1
642 ) # Array indicating in which bin each values of m_vals lies in relative to dist_mGrid. Bins lie between between point of Dist_mGrid.
643 # For instance, if mval lies between dist_mGrid[4] and dist_mGrid[5] it is in bin 4 (would be 5 if 1 was not subtracted in the previous line).
644 mIndex[
645 m_vals <= dist_mGrid[0]
646 ] = -1 # if the value is less than the smallest value on dist_mGrid assign it an index of -1
647 mIndex[m_vals >= dist_mGrid[-1]] = (
648 len(dist_mGrid) - 1
649 ) # if value if greater than largest value on dist_mGrid assign it an index of the length of the grid minus 1
651 # the following three lines hold the same intuition as above
652 pIndex = np.digitize(perm_vals, dist_pGrid) - 1
653 pIndex[perm_vals <= dist_pGrid[0]] = -1
654 pIndex[perm_vals >= dist_pGrid[-1]] = len(dist_pGrid) - 1
656 for i in range(len(m_vals)):
657 if (
658 mIndex[i] == -1
659 ): # if mval is below smallest gridpoint, then assign it a weight of 1.0 for lower weight.
660 mlowerIndex = 0
661 mupperIndex = 0
662 mlowerWeight = 1.0
663 mupperWeight = 0.0
664 elif (
665 mIndex[i] == len(dist_mGrid) - 1
666 ): # if mval is greater than maximum gridpoint, then assign the following weights
667 mlowerIndex = -1
668 mupperIndex = -1
669 mlowerWeight = 1.0
670 mupperWeight = 0.0
671 else: # Standard case where mval does not lie past any extremes
672 # identify which two points on the grid the mval is inbetween
673 mlowerIndex = mIndex[i]
674 mupperIndex = mIndex[i] + 1
675 # Assign weight to the indices that bound the m_vals point. Intuitively, an mval perfectly between two points on the mgrid will assign a weight of .5 to the gridpoint above and below
676 mlowerWeight = (
677 (dist_mGrid[mupperIndex] - m_vals[i])
678 / (dist_mGrid[mupperIndex] - dist_mGrid[mlowerIndex])
679 ) # Metric to determine weight of gridpoint/index below. Intuitively, mvals that are close to gridpoint/index above are assigned a smaller mlowerweight.
680 mupperWeight = 1.0 - mlowerWeight # weight of gridpoint/ index above
682 # Same logic as above except the weights here concern the permanent income grid
683 if pIndex[i] == -1:
684 plowerIndex = 0
685 pupperIndex = 0
686 plowerWeight = 1.0
687 pupperWeight = 0.0
688 elif pIndex[i] == len(dist_pGrid) - 1:
689 plowerIndex = -1
690 pupperIndex = -1
691 plowerWeight = 1.0
692 pupperWeight = 0.0
693 else:
694 plowerIndex = pIndex[i]
695 pupperIndex = pIndex[i] + 1
696 plowerWeight = (dist_pGrid[pupperIndex] - perm_vals[i]) / (
697 dist_pGrid[pupperIndex] - dist_pGrid[plowerIndex]
698 )
699 pupperWeight = 1.0 - plowerWeight
701 # Compute probabilities of each gridpoint on the combined market resources and permanent income grid by looping through each point on the combined market resources and permanent income grid,
702 # assigning probabilities to each gridpoint based off the probabilities of the surrounding mvals and pvals and their respective weights placed on the gridpoint.
703 # Note* probs[i] is the probability of mval AND pval occurring
705 probGrid[mlowerIndex][plowerIndex] += (
706 probs[i] * mlowerWeight * plowerWeight
707 ) # probability of gridpoint below mval and pval
708 probGrid[mlowerIndex][pupperIndex] += (
709 probs[i] * mlowerWeight * pupperWeight
710 ) # probability of gridpoint below mval and above pval
711 probGrid[mupperIndex][plowerIndex] += (
712 probs[i] * mupperWeight * plowerWeight
713 ) # probability of gridpoint above mval and below pval
714 probGrid[mupperIndex][pupperIndex] += (
715 probs[i] * mupperWeight * pupperWeight
716 ) # probability of gridpoint above mval and above pval
718 return probGrid.flatten()
721@numba.njit(parallel=True)
722def gen_tran_matrix_1D(
723 dist_mGrid, bNext, shk_prbs, perm_shks, tran_shks, LivPrb, NewBornDist
724): # pragma: nocover
725 """
726 Computes Transition Matrix across normalized market resources.
727 This function is built to non-stochastic simulate the IndShockConsumerType.
728 This function is used exclusively when Harmenberg Neutral Measure is applied and/or if permanent income is not a state variable
729 For more information, see https://econ-ark.org/materials/harmenberg-aggregation?launch
731 Parameters
732 ----------
733 dist_mGrid : np.array
734 Grid over normalized market resources
736 bNext : np.array
737 Grid over bank balances
739 shk_prbs : np.array
740 Array of shock probabilities over combinations of permanent and transitory shocks
742 perm_shks : np.array
743 Array of shocks to permanent income. Shocks should follow Harmenberg neutral measure
745 tran_shks : np.array
746 Array of shocks to transitory
748 LivPrb : float
749 Probability of not dying
751 NewBornDist : np.array
752 array representing distribution of newborns across grid of normalized market resources and grid of permanent income.
754 Returns
755 -------
756 TranMatrix : np.array
757 Transition Matrix over normalized market resources grid.
760 """
762 TranMatrix = np.zeros((len(dist_mGrid), len(dist_mGrid)))
763 for i in numba.prange(len(dist_mGrid)):
764 mNext_ij = (
765 bNext[i] / perm_shks + tran_shks
766 ) # Compute next period's market resources given todays bank balances bnext[i]
767 TranMatrix[:, i] += (
768 LivPrb * jump_to_grid_1D(mNext_ij, shk_prbs, dist_mGrid)
769 + (1.0 - LivPrb) * NewBornDist
770 ) # this is the transition matrix if given you are unemployed today and unemployed tomorrow so you assume the unemployed consumption policy
771 return TranMatrix
774@numba.njit(parallel=True)
775def gen_tran_matrix_2D(
776 dist_mGrid, dist_pGrid, bNext, shk_prbs, perm_shks, tran_shks, LivPrb, NewBornDist
777): # pragma: nocover
778 """
779 Computes Transition Matrix over normalized market resources and permanent income.
780 This function is built to non-stochastic simulate the IndShockConsumerType.
782 Parameters
783 ----------
784 dist_mGrid : np.array
785 Grid over normalized market resources
787 dist_pGrid : np.array
788 Grid over permanent income
790 bNext : np.array
791 Grid over bank balances
793 shk_prbs : np.array
794 Array of shock probabilities over combinations of perm and tran shocks
796 perm_shks : np.array
797 Array of shocks to permanent income
799 tran_shks : np.array
800 Array of shocks to transitory income
802 LivPrb : float
803 Probability of not dying
805 NewBornDist : np.array
806 array representing distribution of newborns across grid of normalized market resources and grid of permanent income.
808 Returns
809 -------
810 TranMatrix : np.array
811 Transition Matrix over normalized market resources grid and permanent income grid
812 """
813 TranMatrix = np.zeros(
814 (len(dist_mGrid) * len(dist_pGrid), len(dist_mGrid) * len(dist_pGrid))
815 )
816 for i in numba.prange(len(dist_mGrid)):
817 for j in numba.prange(len(dist_pGrid)):
818 mNext_ij = (
819 bNext[i] / perm_shks + tran_shks
820 ) # Compute next period's market resources given todays bank balances bnext[i]
821 pNext_ij = (
822 dist_pGrid[j] * perm_shks
823 ) # Computes next period's permanent income level by applying permanent income shock
824 TranMatrix[:, i * len(dist_pGrid) + j] += (
825 LivPrb
826 * jump_to_grid_2D(mNext_ij, pNext_ij, shk_prbs, dist_mGrid, dist_pGrid)
827 + (1.0 - LivPrb) * NewBornDist
828 ) # this is the transition matrix if given you are unemployed today and unemployed tomorrow so you assume the unemployed consumption policy
829 return TranMatrix
832# ==============================================================================
833# ============== Some basic plotting tools ====================================
834# ==============================================================================
837def _plot_functions_grid(
838 functions, bottom, top, N, legend_kwds, xlabel, ylabel, x_grid, evaluate
839):
840 """Plot one or more 1D ``functions`` over ``[bottom, top]``.
842 ``x_grid(bottom, top, N)`` produces the abscissa, and
843 ``evaluate(function, x)`` returns the ordinate values.
844 """
845 import matplotlib.pyplot as plt
847 plt.ion()
848 function_list = functions if type(functions) == list else [functions]
849 for function in function_list:
850 x = x_grid(bottom, top, N)
851 plt.plot(x, evaluate(function, x))
852 plt.xlim([bottom, top])
853 if xlabel is not None:
854 plt.xlabel(xlabel)
855 if ylabel is not None:
856 plt.ylabel(ylabel)
857 if legend_kwds is not None:
858 plt.legend(**legend_kwds)
859 plt.show(block=False)
862def plot_funcs(
863 functions, bottom, top, N=1000, legend_kwds=None, xlabel=None, ylabel=None
864):
865 """
866 Plots 1D function(s) over a given range.
868 Parameters
869 ----------
870 functions : [function] or function
871 A single function, or a list of functions, to be plotted.
872 bottom : float
873 The lower limit of the domain to be plotted.
874 top : float
875 The upper limit of the domain to be plotted.
876 N : int
877 Number of points in the domain to evaluate.
878 legend_kwds: None or dictionary
879 If not None, the keyword dictionary to pass to plt.legend
880 xlabel : None or str
881 Optional horizontal axis label.
882 ylabel : None or str
883 Optional vertical axis label.
885 Returns
886 -------
887 none
888 """
889 _plot_functions_grid(
890 functions,
891 bottom,
892 top,
893 N,
894 legend_kwds,
895 xlabel,
896 ylabel,
897 x_grid=lambda b, t, n: np.linspace(b, t, n, endpoint=True),
898 evaluate=lambda f, x: f(x),
899 )
902def plot_funcs_der(
903 functions, bottom, top, N=1000, legend_kwds=None, xlabel=None, ylabel=None
904):
905 """
906 Plots the first derivative of 1D function(s) over a given range.
908 Parameters
909 ----------
910 function : function
911 A function or list of functions, the derivatives of which are to be plotted.
912 bottom : float
913 The lower limit of the domain to be plotted.
914 top : float
915 The upper limit of the domain to be plotted.
916 N : int
917 Number of points in the domain to evaluate.
918 legend_kwds: None or dictionary
919 If not None, the keyword dictionary to pass to plt.legend
920 xlabel : None or str
921 Optional horizontal axis label.
922 ylabel : None or str
923 Optional vertical axis label.
925 Returns
926 -------
927 none
928 """
929 _plot_functions_grid(
930 functions,
931 bottom,
932 top,
933 N,
934 legend_kwds,
935 xlabel,
936 ylabel,
937 x_grid=lambda b, t, n: np.arange(b, t, (t - b) / n),
938 evaluate=lambda f, x: f.derivative(x),
939 )
942def plot_func_slices(
943 func,
944 bot,
945 top,
946 N=1000,
947 xdim=0,
948 zdim=1,
949 zmin=None,
950 zmax=None,
951 zn=11,
952 zorder=1.0,
953 Z=None,
954 const=None,
955 legend_kwds=None,
956 xlabel=None,
957 ylabel=None,
958):
959 """
960 Plots "slices" of a function with more than one argument. User specifies
961 range of the "x" (horizontal) dimension and selects which other dimension
962 is "z". Can specify set of z values explicitly with list Z or describe an
963 exponentially spaced grid with zmin, zmax, (zn, zorder).
965 Parameters
966 ----------
967 func : callable
968 The function whose slices are to be plotted. Must take more than one argument.
969 bot : float
970 Lowest value in the xdim to plot.
971 top : float
972 Highest value in the xdim to plot.
973 N : int
974 Number of x values to plot for each slice (default 1000).
975 xdim : int
976 Index of the input that serves as "x", the horizontal plot dimension (default 0).
977 zdim : int
978 Index of the input that serves as "z", which is varied for each slice plotted (default 1).
979 zmin : None or float
980 If specified, the lowest value of z that will be plotted.
981 zmax : None or float
982 If specified, the highest value of z that will be plotted.
983 zn : int
984 The number of slices to plot if zmin and zmax are specified (default 11).
985 zorder : float
986 The exponential order of the set of z values, if zmin and zmax are specified (default 1.0).
987 Z : None or [float]
988 A user-specified list (or array) of z values. Cannot be used of zmin and zmax are provided.
989 const : None or [float]
990 Constant values at which to hold fixed function arguments *other* than x and z.
991 E.g. if user wants to plot f(x, 3.0, z, 5.0), they specify xdim=0, zdim=2, const=[3.0, 5.0].
992 legend_kwds: None or dictionary
993 If not None, the keyword dictionary to pass to plt.legend
994 xlabel : None or str
995 Optional horizontal axis label.
996 ylabel : None or str
997 Optional vertical axis label.
999 Returns
1000 -------
1001 None
1002 """
1003 import matplotlib.pyplot as plt
1005 plt.ion()
1007 if xdim == zdim:
1008 raise ValueError("xdim and zdim cannot refer to the same argument!")
1010 # Check whether the set for z has been correctly specified
1011 if (zmin is not None) and (zmax is not None):
1012 if Z is None:
1013 if zmax > zmin:
1014 Z = make_polynomial_grid(zmin, zmax, zn, order=zorder)
1015 else:
1016 raise ValueError("zmax must be greater than zmin!")
1017 else:
1018 raise ValueError(
1019 "Cannot provide set Z if zmin and zmax are also specified!"
1020 )
1021 if Z is None:
1022 raise ValueError("Must specify set Z or grid of z with zmin and zmax!")
1024 # Build the vectors of x values and constant values
1025 X = np.linspace(bot, top, num=N)
1026 if const is not None:
1027 Q = [const[j] * np.ones(N) for j in range(len(const))]
1028 else:
1029 Q = []
1030 D = 2 + len(Q)
1032 # Assemble a list of function arguments in the right order, leaving z blank
1033 args = []
1034 i = 0
1035 for d in range(D):
1036 if d == xdim:
1037 args.append(X)
1038 elif d == zdim:
1039 args.append(None)
1040 else:
1041 args.append(Q[i])
1042 i += 1
1044 # Plot a slice for each z in Z
1045 for j in range(len(Z)):
1046 z = Z[j]
1047 args[zdim] = z * np.ones(N)
1048 plt.plot(X, func(*args))
1050 # Format and display the figure
1051 plt.xlim(bot, top)
1052 if xlabel is not None:
1053 plt.xlabel(xlabel)
1054 if ylabel is not None:
1055 plt.ylabel(ylabel)
1056 if legend_kwds is not None:
1057 plt.legend(**legend_kwds)
1058 plt.show(block=False)
1061def plot_SSJ(jac, S, outcome=None, shock=None, t_max=None):
1062 """
1063 Plots selected columns of an HA-SSJ matrix.
1065 Parameters
1066 ----------
1067 jac : np.array
1068 T x T array representing an HA-SSJ matrix.
1069 S : int | Sequence[int]
1070 Which columns of the SSJ should be plotted, representing how many periods
1071 ahead the shock happens after announcement at t=0.
1072 outcome : str, optional
1073 The name or description of the outcome to be plotted.
1074 shock : str, optional
1075 The name or description of the variable that is shocked at t=s.
1076 t_max : int, optional
1077 Optional last period t to plot, truncating the graph to the right.
1079 Returns
1080 -------
1081 None
1082 """
1083 import matplotlib.pyplot as plt
1085 plt.ion()
1087 top = jac.shape[0] + 1 if t_max is None else t_max + 1
1088 if isinstance(S, (int, np.integer)) and not isinstance(S, bool):
1089 S = [S]
1090 for s in S:
1091 plt.plot(jac[:, s], "-", label="s=" + str(s))
1092 plt.legend()
1093 plt.xlabel(r"time $t$")
1094 if outcome is None:
1095 plt.ylabel("rate of change")
1096 else:
1097 plt.ylabel("rate of change of " + outcome)
1098 if outcome is not None and shock is not None:
1099 plt.title("SSJ for " + outcome + " with respect to " + shock + r" at time $s$")
1100 elif shock is not None:
1101 plt.title("SSJ with respect to " + shock + r" at time $s$")
1102 elif outcome is not None:
1103 plt.title("SSJ for " + outcome + r" for a shock at time $s$")
1104 plt.tight_layout()
1105 plt.xlim(-1, top)
1106 plt.show(block=False)
1109###############################################################################
1112def determine_platform(): # pragma: nocover
1113 """
1114 Utility function to return the platform currenlty in use.
1116 Returns
1117 ---------
1118 pf: str
1119 'darwin' (MacOS), 'debian'(debian Linux) or 'win' (windows)
1120 """
1121 import platform
1123 pform = platform.system().lower()
1124 if "darwin" in pform:
1125 pf = "darwin" # MacOS
1126 elif "debian" in pform:
1127 pf = "debian" # Probably cloud (MyBinder, CoLab, ...)
1128 elif "ubuntu" in pform:
1129 pf = "debian" # Probably cloud (MyBinder, CoLab, ...)
1130 elif "win" in pform:
1131 pf = "win"
1132 elif "linux" in pform:
1133 pf = "linux"
1134 else:
1135 raise ValueError("Not able to find out the platform.")
1136 return pf
1139def test_latex_installation(pf): # pragma: no cover
1140 """Test to check if latex is installed on the machine.
1142 Parameters
1143 -----------
1144 pf: str (platform)
1145 output of determine_platform()
1147 Returns
1148 --------
1149 bool: Boolean
1150 True if latex found, else installed in the case of debian
1151 otherwise ImportError raised to direct user to install latex manually
1152 """
1153 # Test whether latex is installed (some of the figures require it)
1154 from distutils.spawn import find_executable
1156 latexExists = False
1158 if find_executable("latex"):
1159 latexExists = True
1160 return True
1162 if not latexExists:
1163 print("Some of the figures below require a full installation of LaTeX")
1164 # If running on Mac or Win, user can be assumed to be able to install
1165 # any missing packages in response to error messages; but not on cloud
1166 # so load LaTeX by hand (painfully slowly)
1167 if "debian" in pf: # CoLab and MyBinder are both ubuntu
1168 print("Installing LaTeX now; please wait 3-5 minutes")
1169 from IPython.utils import io
1171 with io.capture_output() as captured: # Hide hideously long output
1172 os.system("apt-get update")
1173 os.system(
1174 "apt-get install texlive texlive-latex-extra texlive-xetex dvipng"
1175 )
1176 latexExists = True
1177 return True
1178 else:
1179 raise ImportError(
1180 "Please install a full distribution of LaTeX on your computer then rerun. \n \
1181 A full distribution means textlive, texlive-latex-extras, texlive-xetex, dvipng, and ghostscript"
1182 )
1185def in_ipynb():
1186 """If the ipython process contains 'terminal' assume not in a notebook.
1188 Returns
1189 --------
1190 bool: Boolean
1191 True if called from a jupyter notebook, else False
1192 """
1193 try:
1194 if "terminal" in str(type(get_ipython())):
1195 return False
1196 else:
1197 return True
1198 except NameError:
1199 return False
1202def setup_latex_env_notebook(pf, latexExists): # pragma: nocover
1203 """This is needed for use of the latex_envs notebook extension
1204 which allows the use of environments in Markdown.
1206 Parameters
1207 -----------
1208 pf: str (platform)
1209 output of determine_platform()
1210 """
1211 import os
1213 import matplotlib.pyplot as plt
1215 plt.rc("font", family="serif")
1216 plt.rc("text", usetex=latexExists)
1217 if latexExists:
1218 latex_preamble = (
1219 r"\usepackage{amsmath}\usepackage{amsfonts}"
1220 r"\usepackage[T1]{fontenc}"
1221 r"\providecommand{\Ex}{\mathbb{E}}"
1222 r"\providecommand{\StE}{\check}"
1223 r"\providecommand{\Trg}{\hat}"
1224 r"\providecommand{\PermGroFac}{\Gamma}"
1225 r"\providecommand{\cLev}{\pmb{\mathrm{c}}}"
1226 r"\providecommand{\mLev}{\pmb{\mathrm{m}}}"
1227 r"\providecommand{\Rfree}{\mathsf{R}}"
1228 r"\providecommand{\DiscFac}{\beta}"
1229 r"\providecommand{\CRRA}{\rho}"
1230 r"\providecommand{\MPC}{\kappa}"
1231 r"\providecommand{\UnempPrb}{\wp}"
1232 )
1233 # Latex expects paths to be separated by /. \ might result in pieces
1234 # being interpreted as commands.
1235 latexdefs_path = os.getcwd().replace(os.path.sep, "/") + "/latexdefs.tex"
1236 if os.path.isfile(latexdefs_path):
1237 latex_preamble = latex_preamble + r"\input{" + latexdefs_path + r"}"
1238 else: # the required latex_envs package needs this file to exist even if it is empty
1239 from pathlib import Path
1241 Path(latexdefs_path).touch()
1242 plt.rcParams["text.latex.preamble"] = latex_preamble
1245def make_figs(figure_name, saveFigs, drawFigs, target_dir="Figures"):
1246 """Utility function to save figure in multiple formats and display the image.
1248 Parameters
1249 ----------
1250 figure_name: str
1251 name of the figure
1252 saveFigs: bool
1253 True if the figure needs to be written to disk else False
1254 drawFigs: bool
1255 True if the figure should be displayed using plt.draw()
1256 target_dir: str, default = 'Figures/'
1257 Name of folder to save figures to in the current directory
1259 """
1260 import matplotlib.pyplot as plt
1262 if saveFigs:
1263 import os
1265 # Where to put any figures that the user wants to save
1266 my_file_path = os.getcwd() # Find pathname to this file:
1267 Figures_dir = os.path.join(
1268 my_file_path, f"{target_dir}"
1269 ) # LaTeX document assumes figures will be here
1270 if not os.path.exists(Figures_dir):
1271 os.makedirs(Figures_dir) # If dir does not exist, create it
1272 # Save the figures in several formats
1273 print(f"Saving figure {figure_name} in {target_dir}")
1274 plt.savefig(
1275 os.path.join(target_dir, f"{figure_name}.jpg"),
1276 # metadata is not supported for jpg
1277 ) # For web/html
1278 plt.savefig(
1279 os.path.join(target_dir, f"{figure_name}.png"),
1280 metadata={"CreationDate": None},
1281 ) # For web/html
1282 plt.savefig(
1283 os.path.join(target_dir, f"{figure_name}.pdf"),
1284 metadata={"CreationDate": None},
1285 ) # For LaTeX
1286 plt.savefig(
1287 os.path.join(target_dir, f"{figure_name}.svg"),
1288 metadata={"Date": None},
1289 ) # For html5
1290 # Make sure it's possible to plot it by checking for GUI
1291 if drawFigs and find_gui():
1292 plt.ion() # Counterintuitively, you want interactive mode on if you don't want to interact
1293 plt.draw() # Change to false if you want to pause after the figure
1294 plt.pause(2)
1297def find_gui():
1298 """Quick fix to check if matplotlib is running in a GUI environment.
1300 Returns
1301 -------
1302 bool: Boolean
1303 True if it's a GUI environment, False if not.
1304 """
1305 try:
1306 import matplotlib.pyplot as plt
1307 except ImportError:
1308 return False
1309 if plt.get_backend() == "Agg":
1310 return False
1311 return True
1314def benchmark(
1315 agent, sort_by="tottime", max_print=10, filename="restats", return_output=False
1316): # pragma: nocover
1317 """
1318 Profiling tool for HARK models. Calling `benchmark` on agents calls the solver for
1319 the agents and provides time to solve as well as the top `max_print` function calls
1320 in terms of `sort_by`. Optionally allows for saving a text copy of the profile
1321 as well as returning the `Stats` object for further inspection.
1323 For more details on the python profilers, see
1324 https://docs.python.org/3/library/profile.html#the-stats-class
1326 Parameters
1327 ----------
1328 agent: AgentType
1329 A HARK AgentType with a solve() method.
1330 sort_by: string
1331 A string to sort the stats by.
1332 max_print: int
1333 Number of lines to print
1334 filename: string
1335 Optional filename to save output.
1336 return_output: bool
1337 Boolean to determine whether Stats object should be returned.
1339 Returns
1340 -------
1341 stats: Stats (optional)
1342 Profiling object with call statistics.
1343 """
1344 cProfile.run("agent.solve()", filename)
1345 stats = pstats.Stats(filename)
1346 stats.strip_dirs()
1347 stats.sort_stats(sort_by)
1348 stats.print_stats(max_print)
1349 if return_output:
1350 return stats
1353simpledec = re.compile(r"\d*\.\d{8,20}")
1356def mround(match):
1357 return f"{float(match.group()):.5f}"
1360def round_in_file(filename): # pragma: nocover
1361 with open(filename, "r+") as file:
1362 filetext = file.read()
1363 filetext = re.sub(simpledec, mround, filetext)
1364 file.seek(0)
1365 file.write(filetext)
1366 file.truncate()
1369def files_in_dir(mypath):
1370 return [
1371 os.path.join(mypath, f)
1372 for f in os.listdir(mypath)
1373 if os.path.isfile(os.path.join(mypath, f))
1374 ]