Coverage for HARK/utilities.py: 78%
243 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-02 05:14 +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, dict):
38 return parent[query]
39 else:
40 return getattr(parent, query)
43# ==============================================================================
44# ============== Some basic function tools ====================================
45# ==============================================================================
48def get_arg_names(function):
49 """
50 Returns a list of strings naming all of the arguments for the passed function.
52 Parameters
53 ----------
54 function : function
55 A function whose argument names are wanted.
57 Returns
58 -------
59 argNames : [string]
60 The names of the arguments of function.
61 """
62 argCount = function.__code__.co_argcount
63 argNames = function.__code__.co_varnames[:argCount]
64 return argNames
67class NullFunc:
68 """
69 A trivial class that acts as a placeholder "do nothing" function.
70 """
72 def __call__(self, *args):
73 """
74 Returns meaningless output no matter what the input(s) is. If no input,
75 returns None. Otherwise, returns an array of NaNs (or a single NaN) of
76 the same size as the first input.
77 """
78 if len(args) == 0:
79 return None
80 else:
81 arg = args[0]
82 if hasattr(arg, "shape"):
83 return np.zeros_like(arg) + np.nan
84 else:
85 return np.nan
87 def distance(self, other):
88 """
89 Trivial distance metric that only cares whether the other object is also
90 an instance of NullFunc. Intentionally does not inherit from HARKobject
91 as this might create dependency problems.
93 Parameters
94 ----------
95 other : any
96 Any object for comparison to this instance of NullFunc.
98 Returns
99 -------
100 (unnamed) : float
101 The distance between self and other. Returns 0 if other is also a
102 NullFunc; otherwise returns an arbitrary high number.
103 """
104 try:
105 if other.__class__ is self.__class__:
106 return 0.0
107 else:
108 return 1000.0
109 except:
110 return 10000.0
113def apply_fun_to_vals(fun, vals):
114 """
115 Applies a function to the arguments defined in `vals`.
116 This is equivalent to `fun(**vals)`, except
117 that `vals` may contain keys that are not named arguments
118 of `fun`.
120 Parameters
121 ----------
122 fun: callable
124 vals: dict
125 """
126 return fun(*[vals[var] for var in signature(fun).parameters])
129# =======================================================
130# ================ Other useful functions ===============
131# =======================================================
134def make_assets_grid(aXtraMin, aXtraMax, aXtraCount, aXtraExtra, aXtraNestFac):
135 """
136 Constructs the base grid of post-decision states, representing end-of-period
137 assets above the absolute minimum. Has three modes, depending on aXtraNestFac:
139 aXtraNestFac = -1 : Uniformly spaced grid.
140 aXtraNestFac = 0 : Ordinary exponentially spaced grid.
141 aXtraNestFac >= 1 : Multi-exponentially nested grid.
143 See :func:`HARK.utilities.make_grid_exp_mult` for more info
145 Parameters
146 ----------
147 aXtraMin: float
148 Minimum value for the assets-above-minimum grid.
149 aXtraMax: float
150 Maximum value for the assets-above-minimum grid.
151 aXtraCount: int
152 Number of nodes in the assets-above-minimum grid, not counting extra values.
153 aXtraExtra: [float]
154 Additional values to insert in the assets-above-minimum grid.
155 aXtraNestFac: int
156 Level of exponential nesting for grid. If -1, the grid is linearly spaced.
158 Returns
159 -------
160 aXtraGrid: np.ndarray
161 Base array of values for the post-decision-state grid.
162 """
163 # Set up post decision state grid:
164 if aXtraNestFac == -1:
165 aXtraGrid = np.linspace(aXtraMin, aXtraMax, aXtraCount)
166 elif aXtraNestFac >= 0:
167 aXtraGrid = make_grid_exp_mult(
168 ming=aXtraMin, maxg=aXtraMax, ng=aXtraCount, timestonest=aXtraNestFac
169 )
170 else:
171 raise ValueError(
172 "Please ensure aXtraNestFac is either -1 or a positive integer."
173 )
175 # Add in additional points for the grid:
176 if aXtraExtra is not None:
177 temp_list = []
178 for i in aXtraExtra:
179 if i is not None:
180 temp_list.append(i)
181 aXtraGrid = np.sort(np.unique(np.concatenate((aXtraGrid, temp_list))))
183 return aXtraGrid
186# ==============================================================================
187# ============== Functions for generating state space grids ===================
188# ==============================================================================
191def make_grid_exp_mult(ming, maxg, ng, timestonest=20):
192 r"""
193 Makes a multi-exponentially spaced grid.
194 If the function :math:`\ln(1+x)` were applied timestonest times,
195 the grid would become linearly spaced.
196 If timestonest is 0, the grid is exponentially spaced.
197 If timestonest is -1, the grid is linearly spaced.
199 Parameters
200 ----------
201 ming : float
202 Minimum value of the grid
203 maxg : float
204 Maximum value of the grid
205 ng : int
206 The number of grid points
207 timestonest : int
208 the number of times to nest the exponentiation
210 Returns
211 -------
212 points : np.array
213 A multi-exponentially spaced grid
215 Notes
216 -----
217 Original Matab code can be found in Chris Carroll's
218 [Solution Methods for Microeconomic Dynamic Optimization Problems]
219 (https://www.econ2.jhu.edu/people/ccarroll/solvingmicrodsops/) toolkit.
220 Latest update: 01 May 2015
221 """
222 if timestonest == -1:
223 grid = np.linspace(ming, maxg, ng)
224 return grid
225 if timestonest > 0:
226 Lming = ming
227 Lmaxg = maxg
228 for j in range(timestonest):
229 Lming = np.log(Lming + 1)
230 Lmaxg = np.log(Lmaxg + 1)
231 Lgrid = np.linspace(Lming, Lmaxg, ng)
232 grid = Lgrid
233 for j in range(timestonest):
234 grid = np.exp(grid) - 1
235 else:
236 Lming = np.log(ming)
237 Lmaxg = np.log(maxg)
238 Lgrid = np.linspace(Lming, Lmaxg, ng)
239 grid = np.exp(Lgrid)
240 return grid
243# ==============================================================================
244# ============== Uncategorized general functions ===================
245# ==============================================================================
248def get_percentiles(data, weights=None, percentiles=None, presorted=False):
249 """
250 Calculates the requested percentiles of (weighted) data. Median by default.
252 Parameters
253 ----------
254 data : numpy.array
255 A 1D array of float data.
256 weights : np.array
257 A weighting vector for the data.
258 percentiles : [float]
259 A list or numpy.array of percentiles to calculate for the data. Each element should
260 be in (0,1).
261 presorted : boolean
262 Indicator for whether data has already been sorted.
264 Returns
265 -------
266 pctl_out : numpy.array
267 The requested percentiles of the data.
268 """
269 if percentiles is None:
270 percentiles = [0.5]
271 else:
272 if (
273 not isinstance(percentiles, (list, np.ndarray))
274 or min(percentiles) <= 0
275 or max(percentiles) >= 1
276 ):
277 raise ValueError(
278 "Percentiles should be a list or numpy array of floats between 0 and 1"
279 )
281 if data.size < 2:
282 return np.zeros(np.array(percentiles).shape) + np.nan
284 if weights is None: # Set equiprobable weights if none were passed
285 weights = np.ones(data.size) / float(data.size)
287 if presorted: # Sort the data if it is not already
288 data_sorted = data
289 weights_sorted = weights
290 else:
291 order = np.argsort(data)
292 data_sorted = data[order]
293 weights_sorted = weights[order]
295 cum_dist = np.cumsum(weights_sorted) / np.sum(
296 weights_sorted
297 ) # cumulative probability distribution
299 # Calculate the requested percentiles by interpolating the data over the
300 # cumulative distribution, then evaluating at the percentile values
301 inv_CDF = interp1d(cum_dist, data_sorted, bounds_error=False, assume_sorted=True)
302 pctl_out = inv_CDF(percentiles)
303 return pctl_out
306def get_lorenz_shares(data, weights=None, percentiles=None, presorted=False):
307 """
308 Calculates the Lorenz curve at the requested percentiles of (weighted) data.
309 Median by default.
311 Parameters
312 ----------
313 data : numpy.array
314 A 1D array of float data.
315 weights : numpy.array
316 A weighting vector for the data.
317 percentiles : [float]
318 A list or numpy.array of percentiles to calculate for the data. Each element should
319 be in (0,1).
320 presorted : boolean
321 Indicator for whether data has already been sorted.
323 Returns
324 -------
325 lorenz_out : numpy.array
326 The requested Lorenz curve points of the data.
327 """
328 if percentiles is None:
329 percentiles = [0.5]
330 else:
331 if (
332 not isinstance(percentiles, (list, np.ndarray))
333 or min(percentiles) <= 0
334 or max(percentiles) >= 1
335 ):
336 raise ValueError(
337 "Percentiles should be a list or numpy array of floats between 0 and 1"
338 )
339 if weights is None: # Set equiprobable weights if none were given
340 weights = np.ones(data.size)
342 if presorted: # Sort the data if it is not already
343 data_sorted = data
344 weights_sorted = weights
345 else:
346 order = np.argsort(data)
347 data_sorted = data[order]
348 weights_sorted = weights[order]
350 cum_dist = np.cumsum(weights_sorted) / np.sum(
351 weights_sorted
352 ) # cumulative probability distribution
353 temp = data_sorted * weights_sorted
354 cum_data = np.cumsum(temp) / sum(temp) # cumulative ownership shares
356 # Calculate the requested Lorenz shares by interpolating the cumulative ownership
357 # shares over the cumulative distribution, then evaluating at requested points
358 lorenzFunc = interp1d(cum_dist, cum_data, bounds_error=False, assume_sorted=True)
359 lorenz_out = lorenzFunc(percentiles)
360 return lorenz_out
363def calc_subpop_avg(data, reference, cutoffs, weights=None):
364 """
365 Calculates the average of (weighted) data between cutoff percentiles of a
366 reference variable.
368 Parameters
369 ----------
370 data : numpy.array
371 A 1D array of float data.
372 reference : numpy.array
373 A 1D array of float data of the same length as data.
374 cutoffs : [(float,float)]
375 A list of doubles with the lower and upper percentile bounds (should be
376 in [0,1]).
377 weights : numpy.array
378 A weighting vector for the data.
380 Returns
381 -------
382 slice_avg
383 The (weighted) average of data that falls within the cutoff percentiles
384 of reference.
386 """
387 if weights is None: # Set equiprobable weights if none were given
388 weights = np.ones(data.size)
390 # Sort the data and generate a cumulative distribution
391 order = np.argsort(reference)
392 data_sorted = data[order]
393 weights_sorted = weights[order]
394 cum_dist = np.cumsum(weights_sorted) / np.sum(weights_sorted)
396 # For each set of cutoffs, calculate the average of data that falls within
397 # the cutoff percentiles of reference
398 slice_avg = []
399 for j in range(len(cutoffs)):
400 bot = np.searchsorted(cum_dist, cutoffs[j][0])
401 top = np.searchsorted(cum_dist, cutoffs[j][1])
402 slice_avg.append(
403 np.sum(data_sorted[bot:top] * weights_sorted[bot:top])
404 / np.sum(weights_sorted[bot:top])
405 )
406 return slice_avg
409def kernel_regression(x, y, bot=None, top=None, N=500, h=None):
410 """
411 Performs a non-parametric Nadaraya-Watson 1D kernel regression on given data
412 with optionally specified range, number of points, and kernel bandwidth.
414 Parameters
415 ----------
416 x : np.array
417 The independent variable in the kernel regression.
418 y : np.array
419 The dependent variable in the kernel regression.
420 bot : float
421 Minimum value of interest in the regression; defaults to min(x).
422 top : float
423 Maximum value of interest in the regression; defaults to max(y).
424 N : int
425 Number of points to compute.
426 h : float
427 The bandwidth of the (Epanechnikov) kernel. To-do: GENERALIZE.
429 Returns
430 -------
431 regression : LinearInterp
432 A piecewise locally linear kernel regression: y = f(x).
433 """
434 # Fix omitted inputs
435 if bot is None:
436 bot = np.min(x)
437 if top is None:
438 top = np.max(x)
439 if h is None:
440 h = 2.0 * (top - bot) / float(N) # This is an arbitrary default
442 # Construct a local linear approximation
443 x_vec = np.linspace(bot, top, num=N)
444 # Evaluate the kernel for all evaluation points at once
445 weights = epanechnikov_kernel(x[:, None], x_vec[None, :], h)
446 weight_sums = np.sum(weights, axis=0)
447 # Avoid division by zero when weights are extremely small
448 weight_sums[weight_sums == 0] = np.nan
449 y_vec = np.dot(weights.T, y) / weight_sums
450 regression = interp1d(x_vec, y_vec, bounds_error=False, assume_sorted=True)
451 return regression
454def epanechnikov_kernel(x, ref_x, h=1.0):
455 """
456 The Epanechnikov kernel, which has been shown to be the most efficient kernel
457 for non-parametric regression.
459 Parameters
460 ----------
461 x : np.array
462 Values at which to evaluate the kernel
463 x_ref : float
464 The reference point
465 h : float
466 Kernel bandwidth
468 Returns
469 -------
470 out : np.array
471 Kernel values at each value of x
472 """
473 u = (x - ref_x) / h # Normalize distance by bandwidth
474 out = 0.75 * (1.0 - u**2)
475 out[np.abs(u) > 1.0] = 0.0 # Kernel = 0 outside [-1,1]
476 return out
479def triangle_kernel(x, ref_x, h=1.0):
480 """
481 The triangle or "hat" kernel.
483 Parameters
484 ----------
485 x : np.array
486 Values at which to evaluate the kernel
487 x_ref : float
488 The reference point
489 h : float
490 Kernel bandwidth
492 Returns
493 -------
494 out : np.array
495 Kernel values at each value of x
496 """
497 u = (x - ref_x) / h # Normalize distance by bandwidth
498 these = np.abs(u) <= 1.0 # Kernel = 0 outside [-1,1]
499 out = np.zeros_like(x) # Initialize kernel output
500 out[these] = 1.0 - np.abs(u[these]) # Evaluate kernel
501 return out
504def make_polynomial_params(coeffs, T, offset=0.0, step=1.0):
505 """
506 Make a T-length array of parameters using polynomial coefficients.
508 Parameters
509 ----------
510 coeffs : [float]
511 Arbitrary length list of polynomial coefficients.
512 T : int
513 Number of elements in the output.
514 offset : float, optional
515 Value at which the X values should start. The default is 0.0.
516 step : float, optional
517 Increment of the X values. The default is 1.0.
519 Returns
520 -------
521 param_vals : np.array
522 T-length array of parameter values calculated using the polynomial coefficients.
523 """
524 X = offset + step * np.arange(T)
525 # np.polyval expects highest power first
526 return np.polyval(coeffs[::-1], X)
529@numba.njit
530def jump_to_grid_1D(m_vals, probs, Dist_mGrid): # pragma: nocover
531 """
532 Distributes values onto a predefined grid, maintaining the means.
535 Parameters
536 ----------
537 m_vals: np.array
538 Market resource values
540 probs: np.array
541 Shock probabilities associated with combinations of m_vals.
542 Can be thought of as the probability mass function of (m_vals).
544 dist_mGrid : np.array
545 Grid over normalized market resources
547 Returns
548 -------
549 probGrid.flatten(): np.array
550 Probabilities of each gridpoint on the combined grid of market resources
552 """
554 probGrid = np.zeros(len(Dist_mGrid))
555 mIndex = np.digitize(m_vals, Dist_mGrid) - 1
556 mIndex[m_vals <= Dist_mGrid[0]] = -1
557 mIndex[m_vals >= Dist_mGrid[-1]] = len(Dist_mGrid) - 1
559 for i in range(len(m_vals)):
560 if mIndex[i] == -1:
561 mlowerIndex = 0
562 mupperIndex = 0
563 mlowerWeight = 1.0
564 mupperWeight = 0.0
565 elif mIndex[i] == len(Dist_mGrid) - 1:
566 mlowerIndex = -1
567 mupperIndex = -1
568 mlowerWeight = 1.0
569 mupperWeight = 0.0
570 else:
571 mlowerIndex = mIndex[i]
572 mupperIndex = mIndex[i] + 1
573 mlowerWeight = (Dist_mGrid[mupperIndex] - m_vals[i]) / (
574 Dist_mGrid[mupperIndex] - Dist_mGrid[mlowerIndex]
575 )
576 mupperWeight = 1.0 - mlowerWeight
578 probGrid[mlowerIndex] += probs[i] * mlowerWeight
579 probGrid[mupperIndex] += probs[i] * mupperWeight
581 return probGrid.flatten()
584@numba.njit
585def jump_to_grid_2D(
586 m_vals, perm_vals, probs, dist_mGrid, dist_pGrid
587): # pragma: nocover
588 """
589 Distributes values onto a predefined grid, maintaining the means. m_vals and perm_vals are realizations of market resources and permanent income while
590 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
591 respective grids. Returns probabilities of each gridpoint on the combined grid of market resources and permanent income.
594 Parameters
595 ----------
596 m_vals: np.array
597 Market resource values
599 perm_vals: np.array
600 Permanent income values
602 probs: np.array
603 Shock probabilities associated with combinations of m_vals and perm_vals.
604 Can be thought of as the probability mass function of (m_vals, perm_vals).
606 dist_mGrid : np.array
607 Grid over normalized market resources
609 dist_pGrid : np.array
610 Grid over permanent income
611 Returns
612 -------
613 probGrid.flatten(): np.array
614 Probabilities of each gridpoint on the combined grid of market resources and permanent income
615 """
617 probGrid = np.zeros((len(dist_mGrid), len(dist_pGrid)))
619 # Maybe use np.searchsorted as opposed to np.digitize
620 mIndex = (
621 np.digitize(m_vals, dist_mGrid) - 1
622 ) # Array indicating in which bin each values of m_vals lies in relative to dist_mGrid. Bins lie between between point of Dist_mGrid.
623 # 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).
624 mIndex[
625 m_vals <= dist_mGrid[0]
626 ] = -1 # if the value is less than the smallest value on dist_mGrid assign it an index of -1
627 mIndex[m_vals >= dist_mGrid[-1]] = (
628 len(dist_mGrid) - 1
629 ) # if value if greater than largest value on dist_mGrid assign it an index of the length of the grid minus 1
631 # the following three lines hold the same intuition as above
632 pIndex = np.digitize(perm_vals, dist_pGrid) - 1
633 pIndex[perm_vals <= dist_pGrid[0]] = -1
634 pIndex[perm_vals >= dist_pGrid[-1]] = len(dist_pGrid) - 1
636 for i in range(len(m_vals)):
637 if (
638 mIndex[i] == -1
639 ): # if mval is below smallest gridpoint, then assign it a weight of 1.0 for lower weight.
640 mlowerIndex = 0
641 mupperIndex = 0
642 mlowerWeight = 1.0
643 mupperWeight = 0.0
644 elif (
645 mIndex[i] == len(dist_mGrid) - 1
646 ): # if mval is greater than maximum gridpoint, then assign the following weights
647 mlowerIndex = -1
648 mupperIndex = -1
649 mlowerWeight = 1.0
650 mupperWeight = 0.0
651 else: # Standard case where mval does not lie past any extremes
652 # identify which two points on the grid the mval is inbetween
653 mlowerIndex = mIndex[i]
654 mupperIndex = mIndex[i] + 1
655 # 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
656 mlowerWeight = (
657 (dist_mGrid[mupperIndex] - m_vals[i])
658 / (dist_mGrid[mupperIndex] - dist_mGrid[mlowerIndex])
659 ) # Metric to determine weight of gridpoint/index below. Intuitively, mvals that are close to gridpoint/index above are assigned a smaller mlowerweight.
660 mupperWeight = 1.0 - mlowerWeight # weight of gridpoint/ index above
662 # Same logic as above except the weights here concern the permanent income grid
663 if pIndex[i] == -1:
664 plowerIndex = 0
665 pupperIndex = 0
666 plowerWeight = 1.0
667 pupperWeight = 0.0
668 elif pIndex[i] == len(dist_pGrid) - 1:
669 plowerIndex = -1
670 pupperIndex = -1
671 plowerWeight = 1.0
672 pupperWeight = 0.0
673 else:
674 plowerIndex = pIndex[i]
675 pupperIndex = pIndex[i] + 1
676 plowerWeight = (dist_pGrid[pupperIndex] - perm_vals[i]) / (
677 dist_pGrid[pupperIndex] - dist_pGrid[plowerIndex]
678 )
679 pupperWeight = 1.0 - plowerWeight
681 # 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,
682 # assigning probabilities to each gridpoint based off the probabilities of the surrounding mvals and pvals and their respective weights placed on the gridpoint.
683 # Note* probs[i] is the probability of mval AND pval occurring
685 probGrid[mlowerIndex][plowerIndex] += (
686 probs[i] * mlowerWeight * plowerWeight
687 ) # probability of gridpoint below mval and pval
688 probGrid[mlowerIndex][pupperIndex] += (
689 probs[i] * mlowerWeight * pupperWeight
690 ) # probability of gridpoint below mval and above pval
691 probGrid[mupperIndex][plowerIndex] += (
692 probs[i] * mupperWeight * plowerWeight
693 ) # probability of gridpoint above mval and below pval
694 probGrid[mupperIndex][pupperIndex] += (
695 probs[i] * mupperWeight * pupperWeight
696 ) # probability of gridpoint above mval and above pval
698 return probGrid.flatten()
701@numba.njit(parallel=True)
702def gen_tran_matrix_1D(
703 dist_mGrid, bNext, shk_prbs, perm_shks, tran_shks, LivPrb, NewBornDist
704): # pragma: nocover
705 """
706 Computes Transition Matrix across normalized market resources.
707 This function is built to non-stochastic simulate the IndShockConsumerType.
708 This function is used exclusively when Harmenberg Neutral Measure is applied and/or if permanent income is not a state variable
709 For more information, see https://econ-ark.org/materials/harmenberg-aggregation?launch
711 Parameters
712 ----------
713 dist_mGrid : np.array
714 Grid over normalized market resources
716 bNext : np.array
717 Grid over bank balances
719 shk_prbs : np.array
720 Array of shock probabilities over combinations of permanent and transitory shocks
722 perm_shks : np.array
723 Array of shocks to permanent income. Shocks should follow Harmenberg neutral measure
725 tran_shks : np.array
726 Array of shocks to transitory
728 LivPrb : float
729 Probability of not dying
731 NewBornDist : np.array
732 array representing distribution of newborns across grid of normalized market resources and grid of permanent income.
734 Returns
735 -------
736 TranMatrix : np.array
737 Transition Matrix over normalized market resources grid.
740 """
742 TranMatrix = np.zeros((len(dist_mGrid), len(dist_mGrid)))
743 for i in numba.prange(len(dist_mGrid)):
744 mNext_ij = (
745 bNext[i] / perm_shks + tran_shks
746 ) # Compute next period's market resources given todays bank balances bnext[i]
747 TranMatrix[:, i] += (
748 LivPrb * jump_to_grid_1D(mNext_ij, shk_prbs, dist_mGrid)
749 + (1.0 - LivPrb) * NewBornDist
750 ) # this is the transition matrix if given you are unemployed today and unemployed tomorrow so you assume the unemployed consumption policy
751 return TranMatrix
754@numba.njit(parallel=True)
755def gen_tran_matrix_2D(
756 dist_mGrid, dist_pGrid, bNext, shk_prbs, perm_shks, tran_shks, LivPrb, NewBornDist
757): # pragma: nocover
758 """
759 Computes Transition Matrix over normalized market resources and permanent income.
760 This function is built to non-stochastic simulate the IndShockConsumerType.
762 Parameters
763 ----------
764 dist_mGrid : np.array
765 Grid over normalized market resources
767 dist_pGrid : np.array
768 Grid over permanent income
770 bNext : np.array
771 Grid over bank balances
773 shk_prbs : np.array
774 Array of shock probabilities over combinations of perm and tran shocks
776 perm_shks : np.array
777 Array of shocks to permanent income
779 tran_shks : np.array
780 Array of shocks to transitory income
782 LivPrb : float
783 Probability of not dying
785 NewBornDist : np.array
786 array representing distribution of newborns across grid of normalized market resources and grid of permanent income.
788 Returns
789 -------
790 TranMatrix : np.array
791 Transition Matrix over normalized market resources grid and permanent income grid
792 """
793 TranMatrix = np.zeros(
794 (len(dist_mGrid) * len(dist_pGrid), len(dist_mGrid) * len(dist_pGrid))
795 )
796 for i in numba.prange(len(dist_mGrid)):
797 for j in numba.prange(len(dist_pGrid)):
798 mNext_ij = (
799 bNext[i] / perm_shks + tran_shks
800 ) # Compute next period's market resources given todays bank balances bnext[i]
801 pNext_ij = (
802 dist_pGrid[j] * perm_shks
803 ) # Computes next period's permanent income level by applying permanent income shock
804 TranMatrix[:, i * len(dist_pGrid) + j] += (
805 LivPrb
806 * jump_to_grid_2D(mNext_ij, pNext_ij, shk_prbs, dist_mGrid, dist_pGrid)
807 + (1.0 - LivPrb) * NewBornDist
808 ) # this is the transition matrix if given you are unemployed today and unemployed tomorrow so you assume the unemployed consumption policy
809 return TranMatrix
812# ==============================================================================
813# ============== Some basic plotting tools ====================================
814# ==============================================================================
817def plot_funcs(functions, bottom, top, N=1000, legend_kwds=None):
818 """
819 Plots 1D function(s) over a given range.
821 Parameters
822 ----------
823 functions : [function] or function
824 A single function, or a list of functions, to be plotted.
825 bottom : float
826 The lower limit of the domain to be plotted.
827 top : float
828 The upper limit of the domain to be plotted.
829 N : int
830 Number of points in the domain to evaluate.
831 legend_kwds: None, or dictionary
832 If not None, the keyword dictionary to pass to plt.legend
834 Returns
835 -------
836 none
837 """
838 import matplotlib.pyplot as plt
840 plt.ion()
842 if type(functions) == list:
843 function_list = functions
844 else:
845 function_list = [functions]
847 for function in function_list:
848 x = np.linspace(bottom, top, N, endpoint=True)
849 y = function(x)
850 plt.plot(x, y)
851 plt.xlim([bottom, top])
852 if legend_kwds is not None:
853 plt.legend(**legend_kwds)
854 plt.show(block=False)
857def plot_funcs_der(functions, bottom, top, N=1000, legend_kwds=None):
858 """
859 Plots the first derivative of 1D function(s) over a given range.
861 Parameters
862 ----------
863 function : function
864 A function or list of functions, the derivatives of which are to be plotted.
865 bottom : float
866 The lower limit of the domain to be plotted.
867 top : float
868 The upper limit of the domain to be plotted.
869 N : int
870 Number of points in the domain to evaluate.
871 legend_kwds: None, or dictionary
872 If not None, the keyword dictionary to pass to plt.legend
874 Returns
875 -------
876 none
877 """
878 import matplotlib.pyplot as plt
880 plt.ion()
882 if type(functions) == list:
883 function_list = functions
884 else:
885 function_list = [functions]
887 step = (top - bottom) / N
888 for function in function_list:
889 x = np.arange(bottom, top, step)
890 y = function.derivative(x)
891 plt.plot(x, y)
892 plt.xlim([bottom, top])
893 if legend_kwds is not None:
894 plt.legend(**legend_kwds)
895 plt.show(block=False)
898###############################################################################
901def determine_platform():
902 """
903 Utility function to return the platform currenlty in use.
905 Returns
906 ---------
907 pf: str
908 'darwin' (MacOS), 'debian'(debian Linux) or 'win' (windows)
909 """
910 import platform
912 pform = platform.system().lower()
913 if "darwin" in pform:
914 pf = "darwin" # MacOS
915 elif "debian" in pform:
916 pf = "debian" # Probably cloud (MyBinder, CoLab, ...)
917 elif "ubuntu" in pform:
918 pf = "debian" # Probably cloud (MyBinder, CoLab, ...)
919 elif "win" in pform:
920 pf = "win"
921 elif "linux" in pform:
922 pf = "linux"
923 else:
924 raise ValueError("Not able to find out the platform.")
925 return pf
928def test_latex_installation(pf): # pragma: no cover
929 """Test to check if latex is installed on the machine.
931 Parameters
932 -----------
933 pf: str (platform)
934 output of determine_platform()
936 Returns
937 --------
938 bool: Boolean
939 True if latex found, else installed in the case of debian
940 otherwise ImportError raised to direct user to install latex manually
941 """
942 # Test whether latex is installed (some of the figures require it)
943 from distutils.spawn import find_executable
945 latexExists = False
947 if find_executable("latex"):
948 latexExists = True
949 return True
951 if not latexExists:
952 print("Some of the figures below require a full installation of LaTeX")
953 # If running on Mac or Win, user can be assumed to be able to install
954 # any missing packages in response to error messages; but not on cloud
955 # so load LaTeX by hand (painfully slowly)
956 if "debian" in pf: # CoLab and MyBinder are both ubuntu
957 print("Installing LaTeX now; please wait 3-5 minutes")
958 from IPython.utils import io
960 with io.capture_output() as captured: # Hide hideously long output
961 os.system("apt-get update")
962 os.system(
963 "apt-get install texlive texlive-latex-extra texlive-xetex dvipng"
964 )
965 latexExists = True
966 return True
967 else:
968 raise ImportError(
969 "Please install a full distribution of LaTeX on your computer then rerun. \n \
970 A full distribution means textlive, texlive-latex-extras, texlive-xetex, dvipng, and ghostscript"
971 )
974def in_ipynb():
975 """If the ipython process contains 'terminal' assume not in a notebook.
977 Returns
978 --------
979 bool: Boolean
980 True if called from a jupyter notebook, else False
981 """
982 try:
983 if "terminal" in str(type(get_ipython())):
984 return False
985 else:
986 return True
987 except NameError:
988 return False
991def setup_latex_env_notebook(pf, latexExists): # pragma: no cover
992 """This is needed for use of the latex_envs notebook extension
993 which allows the use of environments in Markdown.
995 Parameters
996 -----------
997 pf: str (platform)
998 output of determine_platform()
999 """
1000 import os
1002 import matplotlib.pyplot as plt
1004 plt.rc("font", family="serif")
1005 plt.rc("text", usetex=latexExists)
1006 if latexExists:
1007 latex_preamble = (
1008 r"\usepackage{amsmath}\usepackage{amsfonts}"
1009 r"\usepackage[T1]{fontenc}"
1010 r"\providecommand{\Ex}{\mathbb{E}}"
1011 r"\providecommand{\StE}{\check}"
1012 r"\providecommand{\Trg}{\hat}"
1013 r"\providecommand{\PermGroFac}{\Gamma}"
1014 r"\providecommand{\cLev}{\pmb{\mathrm{c}}}"
1015 r"\providecommand{\mLev}{\pmb{\mathrm{m}}}"
1016 r"\providecommand{\Rfree}{\mathsf{R}}"
1017 r"\providecommand{\DiscFac}{\beta}"
1018 r"\providecommand{\CRRA}{\rho}"
1019 r"\providecommand{\MPC}{\kappa}"
1020 r"\providecommand{\UnempPrb}{\wp}"
1021 )
1022 # Latex expects paths to be separated by /. \ might result in pieces
1023 # being interpreted as commands.
1024 latexdefs_path = os.getcwd().replace(os.path.sep, "/") + "/latexdefs.tex"
1025 if os.path.isfile(latexdefs_path):
1026 latex_preamble = latex_preamble + r"\input{" + latexdefs_path + r"}"
1027 else: # the required latex_envs package needs this file to exist even if it is empty
1028 from pathlib import Path
1030 Path(latexdefs_path).touch()
1031 plt.rcParams["text.latex.preamble"] = latex_preamble
1034def make_figs(figure_name, saveFigs, drawFigs, target_dir="Figures"):
1035 """Utility function to save figure in multiple formats and display the image.
1037 Parameters
1038 ----------
1039 figure_name: str
1040 name of the figure
1041 saveFigs: bool
1042 True if the figure needs to be written to disk else False
1043 drawFigs: bool
1044 True if the figure should be displayed using plt.draw()
1045 target_dir: str, default = 'Figures/'
1046 Name of folder to save figures to in the current directory
1048 """
1049 import matplotlib.pyplot as plt
1051 if saveFigs:
1052 import os
1054 # Where to put any figures that the user wants to save
1055 my_file_path = os.getcwd() # Find pathname to this file:
1056 Figures_dir = os.path.join(
1057 my_file_path, f"{target_dir}"
1058 ) # LaTeX document assumes figures will be here
1059 if not os.path.exists(Figures_dir):
1060 os.makedirs(Figures_dir) # If dir does not exist, create it
1061 # Save the figures in several formats
1062 print(f"Saving figure {figure_name} in {target_dir}")
1063 plt.savefig(
1064 os.path.join(target_dir, f"{figure_name}.jpg"),
1065 # metadata is not supported for jpg
1066 ) # For web/html
1067 plt.savefig(
1068 os.path.join(target_dir, f"{figure_name}.png"),
1069 metadata={"CreationDate": None},
1070 ) # For web/html
1071 plt.savefig(
1072 os.path.join(target_dir, f"{figure_name}.pdf"),
1073 metadata={"CreationDate": None},
1074 ) # For LaTeX
1075 plt.savefig(
1076 os.path.join(target_dir, f"{figure_name}.svg"),
1077 metadata={"Date": None},
1078 ) # For html5
1079 # Make sure it's possible to plot it by checking for GUI
1080 if drawFigs and find_gui():
1081 plt.ion() # Counterintuitively, you want interactive mode on if you don't want to interact
1082 plt.draw() # Change to false if you want to pause after the figure
1083 plt.pause(2)
1086def find_gui():
1087 """Quick fix to check if matplotlib is running in a GUI environment.
1089 Returns
1090 -------
1091 bool: Boolean
1092 True if it's a GUI environment, False if not.
1093 """
1094 try:
1095 import matplotlib.pyplot as plt
1096 except:
1097 return False
1098 if plt.get_backend() == "Agg":
1099 return False
1100 return True
1103def benchmark(
1104 agent, sort_by="tottime", max_print=10, filename="restats", return_output=False
1105):
1106 """
1107 Profiling tool for HARK models. Calling `benchmark` on agents calls the solver for
1108 the agents and provides time to solve as well as the top `max_print` function calls
1109 in terms of `sort_by`. Optionally allows for saving a text copy of the profile
1110 as well as returning the `Stats` object for further inspection.
1112 For more details on the python profilers, see
1113 https://docs.python.org/3/library/profile.html#the-stats-class
1115 Parameters
1116 ----------
1117 agent: AgentType
1118 A HARK AgentType with a solve() method.
1119 sort_by: string
1120 A string to sort the stats by.
1121 max_print: int
1122 Number of lines to print
1123 filename: string
1124 Optional filename to save output.
1125 return_output: bool
1126 Boolean to determine whether Stats object should be returned.
1128 Returns
1129 -------
1130 stats: Stats (optional)
1131 Profiling object with call statistics.
1132 """
1133 cProfile.run("agent.solve()", filename)
1134 stats = pstats.Stats(filename)
1135 stats.strip_dirs()
1136 stats.sort_stats(sort_by)
1137 stats.print_stats(max_print)
1138 if return_output:
1139 return stats
1142simpledec = re.compile(r"\d*\.\d{8,20}")
1145def mround(match):
1146 return f"{float(match.group()):.5f}"
1149def round_in_file(filename):
1150 with open(filename, "r+") as file:
1151 filetext = file.read()
1152 filetext = re.sub(simpledec, mround, filetext)
1153 file.seek(0)
1154 file.write(filetext)
1155 file.truncate()
1158def files_in_dir(mypath):
1159 return [
1160 os.path.join(mypath, f)
1161 for f in os.listdir(mypath)
1162 if os.path.isfile(os.path.join(mypath, f))
1163 ]