Coverage for HARK / utilities.py: 93%
267 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-03-08 05:31 +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):
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_exponential_grid; see below.
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 grid points
209 timestonest : int
210 the number of times to nest the exponentiation
212 Returns
213 -------
214 points : np.array
215 A multi-exponentially spaced grid
217 Notes
218 -----
219 Original Matab code can be found in Chris Carroll's
220 [Solution Methods for Microeconomic Dynamic Optimization Problems]
221 (https://www.econ2.jhu.edu/people/ccarroll/solvingmicrodsops/) toolkit.
222 Latest update: 01 May 2015
223 """
224 if timestonest == -1:
225 grid = np.linspace(ming, maxg, ng)
226 return grid
227 if timestonest > 0:
228 Lming = ming
229 Lmaxg = maxg
230 for j in range(timestonest):
231 Lming = np.log(Lming + 1)
232 Lmaxg = np.log(Lmaxg + 1)
233 Lgrid = np.linspace(Lming, Lmaxg, ng)
234 grid = Lgrid
235 for j in range(timestonest):
236 grid = np.exp(grid) - 1
237 else:
238 Lming = np.log(ming)
239 Lmaxg = np.log(maxg)
240 Lgrid = np.linspace(Lming, Lmaxg, ng)
241 grid = np.exp(Lgrid)
242 return grid
245def make_exponential_grid(ming, maxg, ng, order=1.0):
246 """
247 Construct an exponentially spaced grid with chosen exponential order.
248 A uniformly spaced grid on [0,1] is raised to the chosen order, then linearly
249 remapped to the specified interval. Supports any real valued grid bounds.
251 Parameters
252 ----------
253 ming : float
254 Lower bound of grid.
255 maxg : float
256 Upper bound of grid.
257 ng : int
258 Number of points in the grid.
259 order : float, optional
260 Exponential spacing order for the grid. The default is 1.0, or linear.
262 Returns
263 -------
264 grid : np.array
265 Exponentially spaced grid on [ming, maxg] with ng points.
266 """
267 grid = np.linspace(0.0, 1.0, ng) ** order * (maxg - ming) + ming
268 return grid
271# ==============================================================================
272# ============== Uncategorized general functions ===================
273# ==============================================================================
276def get_percentiles(data, weights=None, percentiles=None, presorted=False):
277 """
278 Calculates the requested percentiles of (weighted) data. Median by default.
280 Parameters
281 ----------
282 data : numpy.array
283 A 1D array of float data.
284 weights : np.array
285 A weighting vector for the data.
286 percentiles : [float]
287 A list or numpy.array of percentiles to calculate for the data. Each element should
288 be in (0,1).
289 presorted : boolean
290 Indicator for whether data has already been sorted.
292 Returns
293 -------
294 pctl_out : numpy.array
295 The requested percentiles of the data.
296 """
297 if percentiles is None:
298 percentiles = [0.5]
299 else:
300 if (
301 not isinstance(percentiles, (list, np.ndarray))
302 or min(percentiles) <= 0
303 or max(percentiles) >= 1
304 ):
305 raise ValueError(
306 "Percentiles should be a list or numpy array of floats between 0 and 1"
307 )
309 if data.size < 2:
310 return np.zeros(np.array(percentiles).shape) + np.nan
312 if weights is None: # Set equiprobable weights if none were passed
313 weights = np.ones(data.size) / float(data.size)
315 if presorted: # Sort the data if it is not already
316 data_sorted = data
317 weights_sorted = weights
318 else:
319 order = np.argsort(data)
320 data_sorted = data[order]
321 weights_sorted = weights[order]
323 cum_dist = np.cumsum(weights_sorted) / np.sum(
324 weights_sorted
325 ) # cumulative probability distribution
327 # Calculate the requested percentiles by interpolating the data over the
328 # cumulative distribution, then evaluating at the percentile values
329 inv_CDF = interp1d(cum_dist, data_sorted, bounds_error=False, assume_sorted=True)
330 pctl_out = inv_CDF(percentiles)
331 return pctl_out
334def get_lorenz_shares(data, weights=None, percentiles=None, presorted=False):
335 """
336 Calculates the Lorenz curve at the requested percentiles of (weighted) data.
337 Median by default.
339 Parameters
340 ----------
341 data : numpy.array
342 A 1D array of float data.
343 weights : numpy.array
344 A weighting vector for the data.
345 percentiles : [float]
346 A list or numpy.array of percentiles to calculate for the data. Each element should
347 be in (0,1).
348 presorted : boolean
349 Indicator for whether data has already been sorted.
351 Returns
352 -------
353 lorenz_out : numpy.array
354 The requested Lorenz curve points of the data.
355 """
356 if percentiles is None:
357 percentiles = [0.5]
358 else:
359 if (
360 not isinstance(percentiles, (list, np.ndarray))
361 or min(percentiles) <= 0
362 or max(percentiles) >= 1
363 ):
364 raise ValueError(
365 "Percentiles should be a list or numpy array of floats between 0 and 1"
366 )
367 if weights is None: # Set equiprobable weights if none were given
368 weights = np.ones(data.size)
370 if presorted: # Sort the data if it is not already
371 data_sorted = data
372 weights_sorted = weights
373 else:
374 order = np.argsort(data)
375 data_sorted = data[order]
376 weights_sorted = weights[order]
378 cum_dist = np.cumsum(weights_sorted) / np.sum(
379 weights_sorted
380 ) # cumulative probability distribution
381 temp = data_sorted * weights_sorted
382 cum_data = np.cumsum(temp) / sum(temp) # cumulative ownership shares
384 # Calculate the requested Lorenz shares by interpolating the cumulative ownership
385 # shares over the cumulative distribution, then evaluating at requested points
386 lorenzFunc = interp1d(cum_dist, cum_data, bounds_error=False, assume_sorted=True)
387 lorenz_out = lorenzFunc(percentiles)
388 return lorenz_out
391def calc_subpop_avg(data, reference, cutoffs, weights=None):
392 """
393 Calculates the average of (weighted) data between cutoff percentiles of a
394 reference variable.
396 Parameters
397 ----------
398 data : numpy.array
399 A 1D array of float data.
400 reference : numpy.array
401 A 1D array of float data of the same length as data.
402 cutoffs : [(float,float)]
403 A list of doubles with the lower and upper percentile bounds (should be
404 in [0,1]).
405 weights : numpy.array
406 A weighting vector for the data.
408 Returns
409 -------
410 slice_avg
411 The (weighted) average of data that falls within the cutoff percentiles
412 of reference.
414 """
415 if weights is None: # Set equiprobable weights if none were given
416 weights = np.ones(data.size)
418 # Sort the data and generate a cumulative distribution
419 order = np.argsort(reference)
420 data_sorted = data[order]
421 weights_sorted = weights[order]
422 cum_dist = np.cumsum(weights_sorted) / np.sum(weights_sorted)
424 # For each set of cutoffs, calculate the average of data that falls within
425 # the cutoff percentiles of reference
426 slice_avg = []
427 for j in range(len(cutoffs)):
428 bot = np.searchsorted(cum_dist, cutoffs[j][0])
429 top = np.searchsorted(cum_dist, cutoffs[j][1])
430 slice_avg.append(
431 np.sum(data_sorted[bot:top] * weights_sorted[bot:top])
432 / np.sum(weights_sorted[bot:top])
433 )
434 return slice_avg
437def epanechnikov_kernel(x, ref_x, h=1.0):
438 """
439 The Epanechnikov kernel, which has been shown to be the most efficient kernel
440 for non-parametric regression.
442 Parameters
443 ----------
444 x : np.array
445 Values at which to evaluate the kernel
446 x_ref : float
447 The reference point
448 h : float
449 Kernel bandwidth
451 Returns
452 -------
453 out : np.array
454 Kernel values at each value of x
455 """
456 u = (x - ref_x) / h # Normalize distance by bandwidth
457 out = 0.75 * (1.0 - u**2)
458 out[np.abs(u) > 1.0] = 0.0 # Kernel = 0 outside [-1,1]
459 return out
462def triangle_kernel(x, ref_x, h=1.0):
463 """
464 The triangle or "hat" kernel.
466 Parameters
467 ----------
468 x : np.array
469 Values at which to evaluate the kernel
470 x_ref : float
471 The reference point
472 h : float
473 Kernel bandwidth
475 Returns
476 -------
477 out : np.array
478 Kernel values at each value of x
479 """
480 u = (x - ref_x) / h # Normalize distance by bandwidth
481 these = np.abs(u) <= 1.0 # Kernel = 0 outside [-1,1]
482 out = np.zeros_like(u) # Initialize kernel output
483 out[these] = 1.0 - np.abs(u[these]) # Evaluate kernel
484 return out
487kernel_dict = {
488 "epanechnikov": epanechnikov_kernel,
489 "triangle": triangle_kernel,
490 "hat": triangle_kernel,
491}
494def kernel_regression(x, y, bot=None, top=None, N=500, h=None, kernel="epanechnikov"):
495 """
496 Performs a non-parametric Nadaraya-Watson 1D kernel regression on given data
497 with optionally specified range, number of points, and kernel bandwidth.
499 Parameters
500 ----------
501 x : np.array
502 The independent variable in the kernel regression.
503 y : np.array
504 The dependent variable in the kernel regression.
505 bot : float
506 Minimum value of interest in the regression; defaults to min(x).
507 top : float
508 Maximum value of interest in the regression; defaults to max(y).
509 N : int
510 Number of points to compute.
511 h : float
512 The bandwidth of the (Epanechnikov) kernel. To-do: GENERALIZE.
514 Returns
515 -------
516 regression : LinearInterp
517 A piecewise locally linear kernel regression: y = f(x).
518 """
519 # Fix omitted inputs
520 if bot is None:
521 bot = np.min(x)
522 if top is None:
523 top = np.max(x)
524 if h is None:
525 h = 2.0 * (top - bot) / float(N) # This is an arbitrary default
527 # Get kernel if possible
528 try:
529 kern = kernel_dict[kernel]
530 except:
531 raise ValueError("Can't find a kernel named '" + kernel + "'!")
533 # Construct a local linear approximation
534 x_vec = np.linspace(bot, top, num=N)
535 # Evaluate the kernel for all evaluation points at once
536 weights = kern(x[:, None], x_vec[None, :], h)
537 weight_sums = np.sum(weights, axis=0)
538 # Avoid division by zero when weights are extremely small
539 weight_sums[weight_sums == 0] = np.nan
540 y_vec = np.dot(weights.T, y) / weight_sums
541 regression = interp1d(x_vec, y_vec, bounds_error=False, assume_sorted=True)
542 return regression
545def make_polynomial_params(coeffs, T, offset=0.0, step=1.0):
546 """
547 Make a T-length array of parameters using polynomial coefficients.
549 Parameters
550 ----------
551 coeffs : [float]
552 Arbitrary length list of polynomial coefficients.
553 T : int
554 Number of elements in the output.
555 offset : float, optional
556 Value at which the X values should start. The default is 0.0.
557 step : float, optional
558 Increment of the X values. The default is 1.0.
560 Returns
561 -------
562 param_vals : np.array
563 T-length array of parameter values calculated using the polynomial coefficients.
564 """
565 X = offset + step * np.arange(T)
566 # np.polyval expects highest power first
567 return np.polyval(coeffs[::-1], X)
570@numba.njit
571def jump_to_grid_1D(m_vals, probs, Dist_mGrid): # pragma: nocover
572 """
573 Distributes values onto a predefined grid, maintaining the means.
576 Parameters
577 ----------
578 m_vals: np.array
579 Market resource values
581 probs: np.array
582 Shock probabilities associated with combinations of m_vals.
583 Can be thought of as the probability mass function of (m_vals).
585 dist_mGrid : np.array
586 Grid over normalized market resources
588 Returns
589 -------
590 probGrid.flatten(): np.array
591 Probabilities of each gridpoint on the combined grid of market resources
593 """
595 probGrid = np.zeros(len(Dist_mGrid))
596 mIndex = np.digitize(m_vals, Dist_mGrid) - 1
597 mIndex[m_vals <= Dist_mGrid[0]] = -1
598 mIndex[m_vals >= Dist_mGrid[-1]] = len(Dist_mGrid) - 1
600 for i in range(len(m_vals)):
601 if mIndex[i] == -1:
602 mlowerIndex = 0
603 mupperIndex = 0
604 mlowerWeight = 1.0
605 mupperWeight = 0.0
606 elif mIndex[i] == len(Dist_mGrid) - 1:
607 mlowerIndex = -1
608 mupperIndex = -1
609 mlowerWeight = 1.0
610 mupperWeight = 0.0
611 else:
612 mlowerIndex = mIndex[i]
613 mupperIndex = mIndex[i] + 1
614 mlowerWeight = (Dist_mGrid[mupperIndex] - m_vals[i]) / (
615 Dist_mGrid[mupperIndex] - Dist_mGrid[mlowerIndex]
616 )
617 mupperWeight = 1.0 - mlowerWeight
619 probGrid[mlowerIndex] += probs[i] * mlowerWeight
620 probGrid[mupperIndex] += probs[i] * mupperWeight
622 return probGrid.flatten()
625@numba.njit
626def jump_to_grid_2D(
627 m_vals, perm_vals, probs, dist_mGrid, dist_pGrid
628): # pragma: nocover
629 """
630 Distributes values onto a predefined grid, maintaining the means. m_vals and perm_vals are realizations of market resources and permanent income while
631 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
632 respective grids. Returns probabilities of each gridpoint on the combined grid of market resources and permanent income.
635 Parameters
636 ----------
637 m_vals: np.array
638 Market resource values
640 perm_vals: np.array
641 Permanent income values
643 probs: np.array
644 Shock probabilities associated with combinations of m_vals and perm_vals.
645 Can be thought of as the probability mass function of (m_vals, perm_vals).
647 dist_mGrid : np.array
648 Grid over normalized market resources
650 dist_pGrid : np.array
651 Grid over permanent income
652 Returns
653 -------
654 probGrid.flatten(): np.array
655 Probabilities of each gridpoint on the combined grid of market resources and permanent income
656 """
658 probGrid = np.zeros((len(dist_mGrid), len(dist_pGrid)))
660 # Maybe use np.searchsorted as opposed to np.digitize
661 mIndex = (
662 np.digitize(m_vals, dist_mGrid) - 1
663 ) # Array indicating in which bin each values of m_vals lies in relative to dist_mGrid. Bins lie between between point of Dist_mGrid.
664 # 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).
665 mIndex[
666 m_vals <= dist_mGrid[0]
667 ] = -1 # if the value is less than the smallest value on dist_mGrid assign it an index of -1
668 mIndex[m_vals >= dist_mGrid[-1]] = (
669 len(dist_mGrid) - 1
670 ) # if value if greater than largest value on dist_mGrid assign it an index of the length of the grid minus 1
672 # the following three lines hold the same intuition as above
673 pIndex = np.digitize(perm_vals, dist_pGrid) - 1
674 pIndex[perm_vals <= dist_pGrid[0]] = -1
675 pIndex[perm_vals >= dist_pGrid[-1]] = len(dist_pGrid) - 1
677 for i in range(len(m_vals)):
678 if (
679 mIndex[i] == -1
680 ): # if mval is below smallest gridpoint, then assign it a weight of 1.0 for lower weight.
681 mlowerIndex = 0
682 mupperIndex = 0
683 mlowerWeight = 1.0
684 mupperWeight = 0.0
685 elif (
686 mIndex[i] == len(dist_mGrid) - 1
687 ): # if mval is greater than maximum gridpoint, then assign the following weights
688 mlowerIndex = -1
689 mupperIndex = -1
690 mlowerWeight = 1.0
691 mupperWeight = 0.0
692 else: # Standard case where mval does not lie past any extremes
693 # identify which two points on the grid the mval is inbetween
694 mlowerIndex = mIndex[i]
695 mupperIndex = mIndex[i] + 1
696 # 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
697 mlowerWeight = (
698 (dist_mGrid[mupperIndex] - m_vals[i])
699 / (dist_mGrid[mupperIndex] - dist_mGrid[mlowerIndex])
700 ) # Metric to determine weight of gridpoint/index below. Intuitively, mvals that are close to gridpoint/index above are assigned a smaller mlowerweight.
701 mupperWeight = 1.0 - mlowerWeight # weight of gridpoint/ index above
703 # Same logic as above except the weights here concern the permanent income grid
704 if pIndex[i] == -1:
705 plowerIndex = 0
706 pupperIndex = 0
707 plowerWeight = 1.0
708 pupperWeight = 0.0
709 elif pIndex[i] == len(dist_pGrid) - 1:
710 plowerIndex = -1
711 pupperIndex = -1
712 plowerWeight = 1.0
713 pupperWeight = 0.0
714 else:
715 plowerIndex = pIndex[i]
716 pupperIndex = pIndex[i] + 1
717 plowerWeight = (dist_pGrid[pupperIndex] - perm_vals[i]) / (
718 dist_pGrid[pupperIndex] - dist_pGrid[plowerIndex]
719 )
720 pupperWeight = 1.0 - plowerWeight
722 # 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,
723 # assigning probabilities to each gridpoint based off the probabilities of the surrounding mvals and pvals and their respective weights placed on the gridpoint.
724 # Note* probs[i] is the probability of mval AND pval occurring
726 probGrid[mlowerIndex][plowerIndex] += (
727 probs[i] * mlowerWeight * plowerWeight
728 ) # probability of gridpoint below mval and pval
729 probGrid[mlowerIndex][pupperIndex] += (
730 probs[i] * mlowerWeight * pupperWeight
731 ) # probability of gridpoint below mval and above pval
732 probGrid[mupperIndex][plowerIndex] += (
733 probs[i] * mupperWeight * plowerWeight
734 ) # probability of gridpoint above mval and below pval
735 probGrid[mupperIndex][pupperIndex] += (
736 probs[i] * mupperWeight * pupperWeight
737 ) # probability of gridpoint above mval and above pval
739 return probGrid.flatten()
742@numba.njit(parallel=True)
743def gen_tran_matrix_1D(
744 dist_mGrid, bNext, shk_prbs, perm_shks, tran_shks, LivPrb, NewBornDist
745): # pragma: nocover
746 """
747 Computes Transition Matrix across normalized market resources.
748 This function is built to non-stochastic simulate the IndShockConsumerType.
749 This function is used exclusively when Harmenberg Neutral Measure is applied and/or if permanent income is not a state variable
750 For more information, see https://econ-ark.org/materials/harmenberg-aggregation?launch
752 Parameters
753 ----------
754 dist_mGrid : np.array
755 Grid over normalized market resources
757 bNext : np.array
758 Grid over bank balances
760 shk_prbs : np.array
761 Array of shock probabilities over combinations of permanent and transitory shocks
763 perm_shks : np.array
764 Array of shocks to permanent income. Shocks should follow Harmenberg neutral measure
766 tran_shks : np.array
767 Array of shocks to transitory
769 LivPrb : float
770 Probability of not dying
772 NewBornDist : np.array
773 array representing distribution of newborns across grid of normalized market resources and grid of permanent income.
775 Returns
776 -------
777 TranMatrix : np.array
778 Transition Matrix over normalized market resources grid.
781 """
783 TranMatrix = np.zeros((len(dist_mGrid), len(dist_mGrid)))
784 for i in numba.prange(len(dist_mGrid)):
785 mNext_ij = (
786 bNext[i] / perm_shks + tran_shks
787 ) # Compute next period's market resources given todays bank balances bnext[i]
788 TranMatrix[:, i] += (
789 LivPrb * jump_to_grid_1D(mNext_ij, shk_prbs, dist_mGrid)
790 + (1.0 - LivPrb) * NewBornDist
791 ) # this is the transition matrix if given you are unemployed today and unemployed tomorrow so you assume the unemployed consumption policy
792 return TranMatrix
795@numba.njit(parallel=True)
796def gen_tran_matrix_2D(
797 dist_mGrid, dist_pGrid, bNext, shk_prbs, perm_shks, tran_shks, LivPrb, NewBornDist
798): # pragma: nocover
799 """
800 Computes Transition Matrix over normalized market resources and permanent income.
801 This function is built to non-stochastic simulate the IndShockConsumerType.
803 Parameters
804 ----------
805 dist_mGrid : np.array
806 Grid over normalized market resources
808 dist_pGrid : np.array
809 Grid over permanent income
811 bNext : np.array
812 Grid over bank balances
814 shk_prbs : np.array
815 Array of shock probabilities over combinations of perm and tran shocks
817 perm_shks : np.array
818 Array of shocks to permanent income
820 tran_shks : np.array
821 Array of shocks to transitory income
823 LivPrb : float
824 Probability of not dying
826 NewBornDist : np.array
827 array representing distribution of newborns across grid of normalized market resources and grid of permanent income.
829 Returns
830 -------
831 TranMatrix : np.array
832 Transition Matrix over normalized market resources grid and permanent income grid
833 """
834 TranMatrix = np.zeros(
835 (len(dist_mGrid) * len(dist_pGrid), len(dist_mGrid) * len(dist_pGrid))
836 )
837 for i in numba.prange(len(dist_mGrid)):
838 for j in numba.prange(len(dist_pGrid)):
839 mNext_ij = (
840 bNext[i] / perm_shks + tran_shks
841 ) # Compute next period's market resources given todays bank balances bnext[i]
842 pNext_ij = (
843 dist_pGrid[j] * perm_shks
844 ) # Computes next period's permanent income level by applying permanent income shock
845 TranMatrix[:, i * len(dist_pGrid) + j] += (
846 LivPrb
847 * jump_to_grid_2D(mNext_ij, pNext_ij, shk_prbs, dist_mGrid, dist_pGrid)
848 + (1.0 - LivPrb) * NewBornDist
849 ) # this is the transition matrix if given you are unemployed today and unemployed tomorrow so you assume the unemployed consumption policy
850 return TranMatrix
853# ==============================================================================
854# ============== Some basic plotting tools ====================================
855# ==============================================================================
858def plot_funcs(
859 functions, bottom, top, N=1000, legend_kwds=None, xlabel=None, ylabel=None
860):
861 """
862 Plots 1D function(s) over a given range.
864 Parameters
865 ----------
866 functions : [function] or function
867 A single function, or a list of functions, to be plotted.
868 bottom : float
869 The lower limit of the domain to be plotted.
870 top : float
871 The upper limit of the domain to be plotted.
872 N : int
873 Number of points in the domain to evaluate.
874 legend_kwds: None or dictionary
875 If not None, the keyword dictionary to pass to plt.legend
876 xlabel : None or str
877 Optional horizontal axis label.
878 ylabel : None or str
879 Optional vertical axis label.
881 Returns
882 -------
883 none
884 """
885 import matplotlib.pyplot as plt
887 plt.ion()
889 if type(functions) == list:
890 function_list = functions
891 else:
892 function_list = [functions]
894 for function in function_list:
895 x = np.linspace(bottom, top, N, endpoint=True)
896 y = function(x)
897 plt.plot(x, y)
898 plt.xlim([bottom, top])
899 if xlabel is not None:
900 plt.xlabel(xlabel)
901 if ylabel is not None:
902 plt.ylabel(ylabel)
903 if legend_kwds is not None:
904 plt.legend(**legend_kwds)
905 plt.show(block=False)
908def plot_funcs_der(
909 functions, bottom, top, N=1000, legend_kwds=None, xlabel=None, ylabel=None
910):
911 """
912 Plots the first derivative of 1D function(s) over a given range.
914 Parameters
915 ----------
916 function : function
917 A function or list of functions, the derivatives of which are to be plotted.
918 bottom : float
919 The lower limit of the domain to be plotted.
920 top : float
921 The upper limit of the domain to be plotted.
922 N : int
923 Number of points in the domain to evaluate.
924 legend_kwds: None or dictionary
925 If not None, the keyword dictionary to pass to plt.legend
926 xlabel : None or str
927 Optional horizontal axis label.
928 ylabel : None or str
929 Optional vertical axis label.
931 Returns
932 -------
933 none
934 """
935 import matplotlib.pyplot as plt
937 plt.ion()
939 if type(functions) == list:
940 function_list = functions
941 else:
942 function_list = [functions]
944 step = (top - bottom) / N
945 for function in function_list:
946 x = np.arange(bottom, top, step)
947 y = function.derivative(x)
948 plt.plot(x, y)
949 plt.xlim([bottom, top])
950 if xlabel is not None:
951 plt.xlabel(xlabel)
952 if ylabel is not None:
953 plt.ylabel(ylabel)
954 if legend_kwds is not None:
955 plt.legend(**legend_kwds)
956 plt.show(block=False)
959def plot_func_slices(
960 func,
961 bot,
962 top,
963 N=1000,
964 xdim=0,
965 zdim=1,
966 zmin=None,
967 zmax=None,
968 zn=11,
969 zorder=1.0,
970 Z=None,
971 const=None,
972 legend_kwds=None,
973 xlabel=None,
974 ylabel=None,
975):
976 """
977 Plots "slices" of a function with more than one argument. User specifies
978 range of the "x" (horizontal) dimension and selects which other dimension
979 is "z". Can specify set of z values explicitly with list Z or describe an
980 exponentially spaced grid with zmin, zmax, (zn, zorder).
982 Parameters
983 ----------
984 func : callable
985 The function whose slices are to be plotted. Must take more than one argument.
986 bot : float
987 Lowest value in the xdim to plot.
988 top : float
989 Highest value in the xdim to plot.
990 N : int
991 Number of x values to plot for each slice (default 1000).
992 xdim : int
993 Index of the input that serves as "x", the horizontal plot dimension (default 0).
994 zdim : int
995 Index of the input that serves as "z", which is varied for each slice plotted (default 1).
996 zmin : None or float
997 If specified, the lowest value of z that will be plotted.
998 zmax : None or float
999 If specified, the highest value of z that will be plotted.
1000 zn : int
1001 The number of slices to plot if zmin and zmax are specified (default 11).
1002 zorder : float
1003 The exponential order of the set of z values, if zmin and zmax are specified (default 1.0).
1004 Z : None or [float]
1005 A user-specified list (or array) of z values. Cannot be used of zmin and zmax are provided.
1006 const : None or [float]
1007 Constant values at which to hold fixed function arguments *other* than x and z.
1008 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].
1009 legend_kwds: None or dictionary
1010 If not None, the keyword dictionary to pass to plt.legend
1011 xlabel : None or str
1012 Optional horizontal axis label.
1013 ylabel : None or str
1014 Optional vertical axis label.
1016 Returns
1017 -------
1018 None
1019 """
1020 import matplotlib.pyplot as plt
1022 plt.ion()
1024 if xdim == zdim:
1025 raise ValueError("xdim and zdim cannot refer to the same argument!")
1027 # Check whether the set for z has been correctly specified
1028 if (zmin is not None) and (zmax is not None):
1029 if Z is None:
1030 if zmax > zmin:
1031 Z = make_exponential_grid(zmin, zmax, zn, order=zorder)
1032 else:
1033 raise ValueError("zmax must be greater than zmin!")
1034 else:
1035 raise ValueError(
1036 "Cannot provide set Z if zmin and zmax are also specified!"
1037 )
1038 if Z is None:
1039 raise ValueError("Must specify set Z or grid of z with zmin and zmax!")
1041 # Build the vectors of x values and constant values
1042 X = np.linspace(bot, top, num=N)
1043 if const is not None:
1044 Q = [const[j] * np.ones(N) for j in range(len(const))]
1045 else:
1046 Q = []
1047 D = 2 + len(Q)
1049 # Assemble a list of function arguments in the right order, leaving z blank
1050 args = []
1051 i = 0
1052 for d in range(D):
1053 if d == xdim:
1054 args.append(X)
1055 elif d == zdim:
1056 args.append(None)
1057 else:
1058 args.append(Q[i])
1059 i += 1
1061 # Plot a slice for each z in Z
1062 for j in range(len(Z)):
1063 z = Z[j]
1064 args[zdim] = z * np.ones(N)
1065 plt.plot(X, func(*args))
1067 # Format and display the figure
1068 plt.xlim(bot, top)
1069 if xlabel is not None:
1070 plt.xlabel(xlabel)
1071 if ylabel is not None:
1072 plt.ylabel(ylabel)
1073 if legend_kwds is not None:
1074 plt.legend(**legend_kwds)
1075 plt.show(block=False)
1078###############################################################################
1081def determine_platform(): # pragma: nocover
1082 """
1083 Utility function to return the platform currenlty in use.
1085 Returns
1086 ---------
1087 pf: str
1088 'darwin' (MacOS), 'debian'(debian Linux) or 'win' (windows)
1089 """
1090 import platform
1092 pform = platform.system().lower()
1093 if "darwin" in pform:
1094 pf = "darwin" # MacOS
1095 elif "debian" in pform:
1096 pf = "debian" # Probably cloud (MyBinder, CoLab, ...)
1097 elif "ubuntu" in pform:
1098 pf = "debian" # Probably cloud (MyBinder, CoLab, ...)
1099 elif "win" in pform:
1100 pf = "win"
1101 elif "linux" in pform:
1102 pf = "linux"
1103 else:
1104 raise ValueError("Not able to find out the platform.")
1105 return pf
1108def test_latex_installation(pf): # pragma: no cover
1109 """Test to check if latex is installed on the machine.
1111 Parameters
1112 -----------
1113 pf: str (platform)
1114 output of determine_platform()
1116 Returns
1117 --------
1118 bool: Boolean
1119 True if latex found, else installed in the case of debian
1120 otherwise ImportError raised to direct user to install latex manually
1121 """
1122 # Test whether latex is installed (some of the figures require it)
1123 from distutils.spawn import find_executable
1125 latexExists = False
1127 if find_executable("latex"):
1128 latexExists = True
1129 return True
1131 if not latexExists:
1132 print("Some of the figures below require a full installation of LaTeX")
1133 # If running on Mac or Win, user can be assumed to be able to install
1134 # any missing packages in response to error messages; but not on cloud
1135 # so load LaTeX by hand (painfully slowly)
1136 if "debian" in pf: # CoLab and MyBinder are both ubuntu
1137 print("Installing LaTeX now; please wait 3-5 minutes")
1138 from IPython.utils import io
1140 with io.capture_output() as captured: # Hide hideously long output
1141 os.system("apt-get update")
1142 os.system(
1143 "apt-get install texlive texlive-latex-extra texlive-xetex dvipng"
1144 )
1145 latexExists = True
1146 return True
1147 else:
1148 raise ImportError(
1149 "Please install a full distribution of LaTeX on your computer then rerun. \n \
1150 A full distribution means textlive, texlive-latex-extras, texlive-xetex, dvipng, and ghostscript"
1151 )
1154def in_ipynb():
1155 """If the ipython process contains 'terminal' assume not in a notebook.
1157 Returns
1158 --------
1159 bool: Boolean
1160 True if called from a jupyter notebook, else False
1161 """
1162 try:
1163 if "terminal" in str(type(get_ipython())):
1164 return False
1165 else:
1166 return True
1167 except NameError:
1168 return False
1171def setup_latex_env_notebook(pf, latexExists): # pragma: nocover
1172 """This is needed for use of the latex_envs notebook extension
1173 which allows the use of environments in Markdown.
1175 Parameters
1176 -----------
1177 pf: str (platform)
1178 output of determine_platform()
1179 """
1180 import os
1182 import matplotlib.pyplot as plt
1184 plt.rc("font", family="serif")
1185 plt.rc("text", usetex=latexExists)
1186 if latexExists:
1187 latex_preamble = (
1188 r"\usepackage{amsmath}\usepackage{amsfonts}"
1189 r"\usepackage[T1]{fontenc}"
1190 r"\providecommand{\Ex}{\mathbb{E}}"
1191 r"\providecommand{\StE}{\check}"
1192 r"\providecommand{\Trg}{\hat}"
1193 r"\providecommand{\PermGroFac}{\Gamma}"
1194 r"\providecommand{\cLev}{\pmb{\mathrm{c}}}"
1195 r"\providecommand{\mLev}{\pmb{\mathrm{m}}}"
1196 r"\providecommand{\Rfree}{\mathsf{R}}"
1197 r"\providecommand{\DiscFac}{\beta}"
1198 r"\providecommand{\CRRA}{\rho}"
1199 r"\providecommand{\MPC}{\kappa}"
1200 r"\providecommand{\UnempPrb}{\wp}"
1201 )
1202 # Latex expects paths to be separated by /. \ might result in pieces
1203 # being interpreted as commands.
1204 latexdefs_path = os.getcwd().replace(os.path.sep, "/") + "/latexdefs.tex"
1205 if os.path.isfile(latexdefs_path):
1206 latex_preamble = latex_preamble + r"\input{" + latexdefs_path + r"}"
1207 else: # the required latex_envs package needs this file to exist even if it is empty
1208 from pathlib import Path
1210 Path(latexdefs_path).touch()
1211 plt.rcParams["text.latex.preamble"] = latex_preamble
1214def make_figs(figure_name, saveFigs, drawFigs, target_dir="Figures"):
1215 """Utility function to save figure in multiple formats and display the image.
1217 Parameters
1218 ----------
1219 figure_name: str
1220 name of the figure
1221 saveFigs: bool
1222 True if the figure needs to be written to disk else False
1223 drawFigs: bool
1224 True if the figure should be displayed using plt.draw()
1225 target_dir: str, default = 'Figures/'
1226 Name of folder to save figures to in the current directory
1228 """
1229 import matplotlib.pyplot as plt
1231 if saveFigs:
1232 import os
1234 # Where to put any figures that the user wants to save
1235 my_file_path = os.getcwd() # Find pathname to this file:
1236 Figures_dir = os.path.join(
1237 my_file_path, f"{target_dir}"
1238 ) # LaTeX document assumes figures will be here
1239 if not os.path.exists(Figures_dir):
1240 os.makedirs(Figures_dir) # If dir does not exist, create it
1241 # Save the figures in several formats
1242 print(f"Saving figure {figure_name} in {target_dir}")
1243 plt.savefig(
1244 os.path.join(target_dir, f"{figure_name}.jpg"),
1245 # metadata is not supported for jpg
1246 ) # For web/html
1247 plt.savefig(
1248 os.path.join(target_dir, f"{figure_name}.png"),
1249 metadata={"CreationDate": None},
1250 ) # For web/html
1251 plt.savefig(
1252 os.path.join(target_dir, f"{figure_name}.pdf"),
1253 metadata={"CreationDate": None},
1254 ) # For LaTeX
1255 plt.savefig(
1256 os.path.join(target_dir, f"{figure_name}.svg"),
1257 metadata={"Date": None},
1258 ) # For html5
1259 # Make sure it's possible to plot it by checking for GUI
1260 if drawFigs and find_gui():
1261 plt.ion() # Counterintuitively, you want interactive mode on if you don't want to interact
1262 plt.draw() # Change to false if you want to pause after the figure
1263 plt.pause(2)
1266def find_gui():
1267 """Quick fix to check if matplotlib is running in a GUI environment.
1269 Returns
1270 -------
1271 bool: Boolean
1272 True if it's a GUI environment, False if not.
1273 """
1274 try:
1275 import matplotlib.pyplot as plt
1276 except:
1277 return False
1278 if plt.get_backend() == "Agg":
1279 return False
1280 return True
1283def benchmark(
1284 agent, sort_by="tottime", max_print=10, filename="restats", return_output=False
1285): # pragma: nocover
1286 """
1287 Profiling tool for HARK models. Calling `benchmark` on agents calls the solver for
1288 the agents and provides time to solve as well as the top `max_print` function calls
1289 in terms of `sort_by`. Optionally allows for saving a text copy of the profile
1290 as well as returning the `Stats` object for further inspection.
1292 For more details on the python profilers, see
1293 https://docs.python.org/3/library/profile.html#the-stats-class
1295 Parameters
1296 ----------
1297 agent: AgentType
1298 A HARK AgentType with a solve() method.
1299 sort_by: string
1300 A string to sort the stats by.
1301 max_print: int
1302 Number of lines to print
1303 filename: string
1304 Optional filename to save output.
1305 return_output: bool
1306 Boolean to determine whether Stats object should be returned.
1308 Returns
1309 -------
1310 stats: Stats (optional)
1311 Profiling object with call statistics.
1312 """
1313 cProfile.run("agent.solve()", filename)
1314 stats = pstats.Stats(filename)
1315 stats.strip_dirs()
1316 stats.sort_stats(sort_by)
1317 stats.print_stats(max_print)
1318 if return_output:
1319 return stats
1322simpledec = re.compile(r"\d*\.\d{8,20}")
1325def mround(match):
1326 return f"{float(match.group()):.5f}"
1329def round_in_file(filename): # pragma: nocover
1330 with open(filename, "r+") as file:
1331 filetext = file.read()
1332 filetext = re.sub(simpledec, mround, filetext)
1333 file.seek(0)
1334 file.write(filetext)
1335 file.truncate()
1338def files_in_dir(mypath):
1339 return [
1340 os.path.join(mypath, f)
1341 for f in os.listdir(mypath)
1342 if os.path.isfile(os.path.join(mypath, f))
1343 ]