Coverage for HARK / SSJutils.py: 91%
377 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-05-10 06:19 +0000
1"""
2Functions for building heterogeneous agent sequence space Jacobian matrices from
3HARK AgentType instances. The top-level functions are accessible as methods on
4AgentType itself.
5"""
7from time import time
8from copy import deepcopy
9import numpy as np
10from numba import njit
13def _prepare_ssj_computation(agent, outcomes, grids, norm, solved, verbose):
14 """
15 Shared setup for make_basic_SSJ_matrices and calc_shock_response_manually.
16 Validates the agent, normalizes outcomes, optionally solves the long run model,
17 builds transition matrices, and finds the steady state distribution.
19 Parameters
20 ----------
21 agent : AgentType
22 Agent for which setup should be performed.
23 outcomes : str or [str]
24 Names of outcome variables of interest (will be normalized to a list).
25 grids : dict
26 Dictionary of dictionaries with discretizing grid information.
27 norm : str or None
28 Name of the model variable to normalize by for Harmenberg aggregation.
29 solved : bool
30 Whether the agent's model has already been solved.
31 verbose : bool
32 Whether to display timing/progress to screen.
34 Returns
35 -------
36 setup : dict
37 Dictionary containing:
38 - outcomes : [str] normalized list of outcome variable names
39 - no_list : bool True if outcomes was passed as a single string
40 - simulator_backup : object agent._simulator backup, or None if not present
41 - LR_soln : object deepcopy of the long run solution
42 - X : object reference to agent._simulator
43 - LR_trans : np.array long run transition matrix
44 - LR_period : object the long run period object from the simulator
45 - LR_outcomes : [np.array] outcome matrices in the long run model
46 - outcome_grids : [np.array] outcome grids in the long run model
47 - SS_dstn : np.array steady state distribution
48 """
49 if (agent.cycles > 0) or (agent.T_cycle != 1):
50 raise ValueError(
51 "This function is only compatible with one period infinite horizon models!"
52 )
53 if not isinstance(outcomes, list):
54 outcomes = [outcomes]
55 no_list = True
56 else:
57 no_list = False
59 # Store the simulator if it exists
60 if hasattr(agent, "_simulator"):
61 simulator_backup = agent._simulator
62 else:
63 simulator_backup = None
65 # Solve the long run model if it wasn't already
66 if not solved:
67 t0 = time()
68 agent.solve()
69 t1 = time()
70 if verbose:
71 print(
72 "Solving the long run model took {:.3f}".format(t1 - t0) + " seconds."
73 )
74 LR_soln = deepcopy(agent.solution[0])
76 # Construct the transition matrix for the long run model
77 t0 = time()
78 agent.initialize_sym()
79 X = agent._simulator # for easier referencing
80 X.make_transition_matrices(grids, norm)
81 LR_trans = X.trans_arrays[0].copy() # the transition matrix in LR model
82 LR_period = X.periods[0]
83 LR_outcomes = []
84 outcome_grids = []
85 for var in outcomes:
86 if var not in X.periods[0].matrices:
87 raise ValueError(
88 "Outcome " + var + " was requested but has no transition matrix."
89 )
90 if var not in X.periods[0].grids:
91 raise ValueError(
92 "Outcome " + var + " was requested but no grid was provided!"
93 )
94 LR_outcomes.append(X.periods[0].matrices[var])
95 outcome_grids.append(X.periods[0].grids[var])
96 t1 = time()
97 if verbose:
98 print(
99 "Making the transition matrix for the long run model took {:.3f}".format(
100 t1 - t0
101 )
102 + " seconds."
103 )
105 # Find the steady state for the long run model
106 t0 = time()
107 X.find_steady_state()
108 SS_dstn = X.steady_state_dstn.copy()
109 t1 = time()
110 if verbose:
111 print(
112 "Finding the long run steady state took {:.3f}".format(t1 - t0)
113 + " seconds."
114 )
116 return {
117 "outcomes": outcomes,
118 "no_list": no_list,
119 "simulator_backup": simulator_backup,
120 "LR_soln": LR_soln,
121 "X": X,
122 "LR_trans": LR_trans,
123 "LR_period": LR_period,
124 "LR_outcomes": LR_outcomes,
125 "outcome_grids": outcome_grids,
126 "SS_dstn": SS_dstn,
127 }
130def _perturb_shock(agent, shock):
131 """
132 Retrieve and validate the named shock attribute on an agent, returning the
133 scalar base value and a flag indicating whether it is stored as a singleton list.
135 Parameters
136 ----------
137 agent : AgentType
138 Agent whose shock attribute will be read.
139 shock : str
140 Name of the shock attribute to perturb.
142 Returns
143 -------
144 base_shock_value : float or np.floating
145 The scalar value of the shock attribute (unwrapped from a list if needed).
146 shock_is_list : bool
147 True if the shock attribute is stored as a list, False otherwise.
148 """
149 if not hasattr(agent, shock):
150 raise ValueError(
151 "The agent doesn't have anything called " + shock + " to perturb!"
152 )
153 base = getattr(agent, shock)
154 if isinstance(base, list):
155 base_shock_value = base[0]
156 shock_is_list = True
157 else:
158 base_shock_value = base
159 shock_is_list = False
160 if isinstance(base_shock_value, bool) or not isinstance(
161 base_shock_value, (float, np.floating)
162 ):
163 raise TypeError(
164 "The shock attribute '"
165 + shock
166 + "' must be a scalar float (Python float or np.floating), but got "
167 + str(type(base_shock_value))
168 + "."
169 )
170 return base_shock_value, shock_is_list
173def _restore_agent(agent, LR_soln, simulator_backup):
174 """
175 Reset the agent to its original long run state after SSJ or impulse response
176 computation, restoring the simulator backup if one was saved.
178 Parameters
179 ----------
180 agent : AgentType
181 Agent to restore.
182 LR_soln : object
183 Long run solution to reinstall as the agent's solution.
184 simulator_backup : object or None
185 The original simulator to restore, or None if no backup was stored.
186 """
187 agent.solution = [LR_soln]
188 agent.cycles = 0
189 agent._simulator.reset()
190 if simulator_backup is not None:
191 agent._simulator = simulator_backup
192 else:
193 del agent._simulator
196def make_basic_SSJ_matrices(
197 agent,
198 shock,
199 outcomes,
200 grids,
201 eps=1e-4,
202 T_max=300,
203 norm=None,
204 solved=False,
205 construct=True,
206 offset=False,
207 verbose=False,
208):
209 """
210 Constructs one or more sequence space Jacobian (SSJ) matrices for specified
211 outcomes over one shock variable. It is "basic" in the sense that it only
212 works for "one period infinite horizon" models, as in the original SSJ paper.
214 Parameters
215 ----------
216 agent : AgentType
217 Agent for which the SSJ(s) should be constructed. Must have T_cycle=1
218 and cycles=0, or the function will throw an error. Must have a model
219 file defined or this won't work at all.
220 shock : str
221 Name of the variable that Jacobians will be computed with respect to.
222 It does not need to be a "shock" in a modeling sense, but it must be a
223 single-valued parameter (possibly a singleton list) that can be changed.
224 outcomes : str or [str]
225 Names of outcome variables of interest; an SSJ matrix will be constructed
226 for each variable named here. If a single string is passed, the output
227 will be a single np.array. If a list of strings are passed, the output
228 will be a list of SSJ matrices in the order specified here.
229 grids : dict
230 Dictionary of dictionaries with discretizing grid information. The grids
231 should include all arrival variables other than those that are normalized
232 out. They should also include all variables named in outcomes, except
233 outcomes that are continuation variables that remap to arrival variables.
234 Grid specification must include number of nodes N, should also include
235 min and max if the variable is continuous.
236 eps : float
237 Amount by which to perturb the shock variable. The default is 1e-4.
238 T_max : int
239 Size of the SSJ matrices: the maximum number of periods to consider.
240 The default is 300.
241 norm : str or None
242 Name of the model variable to normalize by for Harmenberg aggregation,
243 if any. For many HARK models, this should be 'PermShk', which enables
244 the grid over permanent income to be omitted as an explicit state.
245 solved : bool
246 Whether the agent's model has already been solved. If False (default),
247 it will be solved as the very first step. Solving the agent's long run
248 model before constructing SSJ matrices has the advantage of not needing
249 to re-solve the long run model for each shock variable.
250 construct : bool
251 Whether the construct (update) method should be run after the shock is
252 updated. The default is True, which is the "safe" option. If the shock
253 variable is a parameter that enters the model only *directly*, rather
254 than being used to build a more complex model input, then this can be
255 set to False to save a (very) small amount of time during computation.
256 If it is set to False improperly, the SSJs will be very wrong, potentially
257 just zero everywhere.
258 offset : bool
259 Whether the shock variable is "offset in time" for the solver, with a
260 default of False. This should be set to True if the named shock variable
261 (or the constructed model input that it affects) is indexed by t+1 from
262 the perspective of the solver. For example, the period t solver for the
263 ConsIndShock model takes in risk free interest factor Rfree as an argument,
264 but it represents the value of R that will occur at the start of t+1.
265 verbose : bool
266 Whether to display timing/progress to screen. The default is False.
268 Returns
269 -------
270 SSJ : np.array or [np.array]
271 One or more sequence space Jacobian arrays over the outcome variables
272 with respect to the named shock variable.
273 """
274 setup = _prepare_ssj_computation(agent, outcomes, grids, norm, solved, verbose)
275 outcomes = setup["outcomes"]
276 no_list = setup["no_list"]
277 simulator_backup = setup["simulator_backup"]
278 LR_soln = setup["LR_soln"]
279 LR_trans = setup["LR_trans"]
280 LR_period = setup["LR_period"]
281 LR_outcomes = setup["LR_outcomes"]
282 outcome_grids = setup["outcome_grids"]
283 SS_dstn = setup["SS_dstn"]
285 try:
286 base_shock_value, shock_is_list = _perturb_shock(agent, shock)
287 Tm1_soln, period_Tm1, period_T = _solve_perturbed_Tm1(
288 agent,
289 shock,
290 base_shock_value,
291 shock_is_list,
292 eps,
293 construct,
294 LR_soln,
295 verbose,
296 )
297 _solve_finite_horizon(
298 agent,
299 shock,
300 base_shock_value,
301 shock_is_list,
302 T_max,
303 construct,
304 Tm1_soln,
305 verbose,
306 )
307 TmX_trans, TmX_outcomes = _build_finite_horizon_matrices(
308 agent,
309 period_Tm1,
310 period_T,
311 LR_period,
312 grids,
313 norm,
314 offset,
315 T_max,
316 outcomes,
317 verbose,
318 )
320 J = len(outcomes)
321 K = SS_dstn.size
322 D_dstn_array, dY_news_array = _compute_finite_horizon_derivatives(
323 T_max,
324 J,
325 outcomes,
326 TmX_trans,
327 TmX_outcomes,
328 LR_trans,
329 LR_outcomes,
330 outcome_grids,
331 SS_dstn,
332 verbose,
333 )
335 FN = _build_fake_news(
336 T_max,
337 J,
338 K,
339 dY_news_array,
340 D_dstn_array,
341 LR_trans,
342 LR_outcomes,
343 outcome_grids,
344 verbose,
345 )
347 # Construct the SSJ matrices, one for each outcome variable
348 t0 = time()
349 SSJ_array = FN.copy()
350 for t in range(1, T_max):
351 SSJ_array[:, 1:, t] += SSJ_array[:, :-1, t - 1]
352 SSJ_array /= eps
353 SSJ = [SSJ_array[j, :, :] for j in range(J)] # unpack into a list of arrays
354 _log_timing(verbose, "Constructing the sequence space Jacobians", t0, time())
356 if no_list:
357 return SSJ[0]
358 else:
359 return SSJ
361 finally:
362 _restore_agent(agent, LR_soln, simulator_backup)
365def _log_timing(verbose, label, t0, t1):
366 if verbose:
367 print(label + " took {:.3f}".format(t1 - t0) + " seconds.")
370def _shock_value(base_value, shock_is_list, eps=0.0):
371 return [base_value + eps] if shock_is_list else base_value + eps
374def _solve_perturbed_Tm1(
375 agent, shock, base_shock_value, shock_is_list, eps, construct, LR_soln, verbose
376):
377 t0 = time()
378 agent.cycles = 1
379 agent.assign_parameters(
380 **{shock: _shock_value(base_shock_value, shock_is_list, eps)}
381 )
382 if construct:
383 agent.update()
384 agent.solve(from_solution=LR_soln)
385 agent.initialize_sym()
386 Tm1_soln = deepcopy(agent.solution[0])
387 period_Tm1 = agent._simulator.periods[0]
388 period_T = agent._simulator.periods[-1]
389 _log_timing(verbose, "Solving period T-1 with a perturbed variable", t0, time())
390 return Tm1_soln, period_Tm1, period_T
393def _solve_finite_horizon(
394 agent, shock, base_shock_value, shock_is_list, T_max, construct, Tm1_soln, verbose
395):
396 t0 = time()
397 agent.cycles = T_max - 1
398 agent.assign_parameters(**{shock: _shock_value(base_shock_value, shock_is_list)})
399 if construct:
400 agent.update()
401 agent.solve(from_solution=Tm1_soln)
402 _log_timing(
403 verbose,
404 "Solving the finite horizon model for " + str(T_max - 1) + " more periods",
405 t0,
406 time(),
407 )
410def _build_finite_horizon_matrices(
411 agent,
412 period_Tm1,
413 period_T,
414 LR_period,
415 grids,
416 norm,
417 offset,
418 T_max,
419 outcomes,
420 verbose,
421):
422 t0 = time()
423 agent.initialize_sym()
424 X = agent._simulator
425 X.periods[-1] = period_Tm1
426 if offset:
427 for name in X.periods[-1].content.keys():
428 if name not in X.solution:
429 X.periods[-1].content[name] = LR_period.content[name]
430 X.periods[-1].distribute_content()
431 X.periods = X.periods[1:] + [period_T]
432 X.make_transition_matrices(grids, norm, fake_news_timing=True)
433 TmX_trans = deepcopy(X.trans_arrays)
434 TmX_outcomes = [
435 [X.periods[t].matrices[var] for var in outcomes] for t in range(T_max)
436 ]
437 _log_timing(
438 verbose,
439 "Constructing transition arrays for the finite horizon model",
440 t0,
441 time(),
442 )
443 return TmX_trans, TmX_outcomes
446def _compute_finite_horizon_derivatives(
447 T_max,
448 J,
449 outcomes,
450 TmX_trans,
451 TmX_outcomes,
452 LR_trans,
453 LR_outcomes,
454 outcome_grids,
455 SS_dstn,
456 verbose,
457):
458 t0 = time()
459 D_dstn_array = calc_derivs_of_state_dstns(
460 T_max, J, np.array(TmX_trans), LR_trans, SS_dstn
461 )
462 dY_news_array = np.empty((T_max, J))
463 for j in range(J):
464 temp_outcomes = np.array([TmX_outcomes[t][j] for t in range(T_max)])
465 dY_news_array[:, j] = calc_derivs_of_policy_funcs(
466 T_max, temp_outcomes, LR_outcomes[j], outcome_grids[j], SS_dstn
467 )
468 _log_timing(verbose, "Calculating derivatives by first differences", t0, time())
469 return D_dstn_array, dY_news_array
472def _build_fake_news(
473 T_max,
474 J,
475 K,
476 dY_news_array,
477 D_dstn_array,
478 LR_trans,
479 LR_outcomes,
480 outcome_grids,
481 verbose,
482):
483 t0 = time()
484 expectation_vectors = np.empty((J, K))
485 for j in range(J):
486 expectation_vectors[j, :] = np.dot(LR_outcomes[j], outcome_grids[j])
487 FN = make_fake_news_matrices(
488 T_max,
489 J,
490 dY_news_array,
491 D_dstn_array,
492 LR_trans.T,
493 expectation_vectors.copy(),
494 )
495 _log_timing(verbose, "Constructing the fake news matrices", t0, time())
496 return FN
499def make_flat_LC_SSJ_matrices(
500 agent,
501 shock,
502 outcomes,
503 grids,
504 eps=1e-4,
505 T_max=100,
506 norm=None,
507 trend=None,
508 pop_gro=1.0,
509 prod_gro=1.0,
510 solved=False,
511 age_agg=True,
512 construct=True,
513 offset=False,
514 verbose=False,
515):
516 """
517 Constructs one or more sequence space Jacobian (SSJ) matrices for specified
518 outcomes over one shock variable. This version of the function is for life-
519 cycle models with "flat" demographic dynamics: the long run distribution of
520 ages is stable. This requires that survival probability is not endogenous to
521 agent actions and thus cannot be affected by shocks.
523 "Flat" demographic dynamics permit two very specific growth trends: constant
524 population growth and constant aggregate productivity growth.
526 This algorithm (and some of the code) are directly taken from Mateo Velasquez
527 and Bence Bardoczy's paper on life-cycle Jacobians, and its accompanying repo.
529 Parameters
530 ----------
531 agent : AgentType
532 Agent for which the SSJ(s) should be constructed. Must have cycles = 1
533 or the function will throw an error. Must have a model file defined or
534 this won't work at all.
535 shock : str
536 Name of the variable that Jacobians will be computed with respect to.
537 It does not need to be a "shock" in a modeling sense, but it must be a
538 single-valued parameter (possibly a singleton list) that can be changed.
539 outcomes : str or [str]
540 Names of outcome variables of interest; an SSJ matrix will be constructed
541 for each variable named here. If a single string is passed, the output
542 will be a single np.array. If a list of strings are passed, the output
543 will be a list of SSJ matrices in the order specified here.
544 grids : dict
545 Dictionary of dictionaries with discretizing grid information. The grids
546 should include all arrival variables other than those that are normalized
547 out. They should also include all variables named in outcomes, except
548 outcomes that are continuation variables that remap to arrival variables.
549 Grid specification must include number of nodes N, should also include
550 min and max if the variable is continuous.
551 eps : float
552 Amount by which to perturb the shock variable. The default is 1e-4.
553 T_max : int
554 Size of the SSJ matrices: the maximum number of periods to consider.
555 The default is 100.
556 norm : str or None
557 Name of the model variable to normalize by for Harmenberg aggregation,
558 if any. For many HARK models, this should be 'PermShk', which enables
559 the grid over permanent income to be omitted as an explicit state.
560 trend : str or None
561 Name of the model variable that represents the "normalization trend factor"
562 for the outcomes. For example, most consumption-saving models in HARK are
563 normalized by permanent income, which grows by factor `PermGroFac` each
564 period of the life-cycle; `PermGroFac` should be named as the `trend` for
565 any model outputs that are normalized by permanent income (i.e. they have
566 `Nrm` in their name). In contrast, if you wanted the fraction of agents
567 that have `Lbr > 0.0` for `LaborIntMargConsumerType`, a binary indicator
568 for this outcome should *not* have `PermGroFac` named as the `trend`--
569 you don't want to upweight people who have accumulated more income growth
570 more when calculating the employment rate!
571 pop_gro : float
572 Constant population growth factor, defaulting to 1. Each successive
573 birth cohort is this factor bigger than the prior birth cohort. With flat
574 demographic dynamics, this results in the entire population growing by
575 this factor each period as well. NOT YET IMPLEMENTED
576 prod_gro : float
577 Constant aggregate productivity growth factor, defaulting to 1. Each
578 successive birth cohort has permanent labor productivity that is this
579 factor bigger than the prior cohort. With flat demographic dynamics,
580 this results in aggregate productivity growing by this factor as well.
581 NOT YET IMPLEMENTED
582 solved : bool
583 Whether the agent's model has already been solved. If False (default),
584 it will be solved as the very first step. Solving the agent's long run
585 model before constructing SSJ matrices has the advantage of not needing
586 to re-solve the long run model for each shock variable.
587 age_agg : bool
588 Whether the returned SSJs should combine effects across ages (default True)
589 or leave effects disaggregated by age (False). When False, each SSJ will
590 be shape (T_age, T_max, T_max), and the overall SSJ matrix can be found by
591 doing np.sum(SSJ, axis=0).
592 construct : bool
593 Whether the construct (update) method should be run after the shock is
594 updated. The default is True, which is the "safe" option. If the shock
595 variable is a parameter that enters the model only *directly*, rather
596 than being used to build a more complex model input, then this can be
597 set to False to save a (very) small amount of time during computation.
598 If it is set to False improperly, the SSJs will be very wrong, potentially
599 just zero everywhere.
600 offset : bool
601 Whether the shock variable is "offset in time" for the solver, with a
602 default of False. This should be set to True if the named shock variable
603 (or the constructed model input that it affects) is indexed by t+1 from
604 the perspective of the solver. For example, the period t solver for the
605 ConsIndShock model takes in risk free interest factor Rfree as an argument,
606 but it represents the value of R that will occur at the start of t+1.
607 verbose : bool
608 Whether to display timing/progress to screen. The default is False.
610 Returns
611 -------
612 SSJ : np.array or [np.array]
613 One or more sequence space Jacobian arrays over the outcome variables
614 with respect to the named shock variable. Each is shape (T_max, T_max).
615 """
616 if agent.cycles != 1:
617 raise ValueError("This function is only compatible with life-cycle models!")
618 if not isinstance(outcomes, list):
619 outcomes = [outcomes]
620 no_list = True
621 else:
622 no_list = False
623 J = len(outcomes)
625 # Check for attempts to use future functionality
626 if pop_gro != 1.0:
627 raise ValueError(
628 "Population growth is not yet implemented for make_flat_LC_SSJ_matrices!"
629 )
630 if prod_gro != 1.0:
631 raise ValueError(
632 "Productivity growth is not yet implemented for make_flat_LC_SSJ_matrices!"
633 )
635 # Store the simulator if it exists
636 simulator_backup = agent._simulator if hasattr(agent, "_simulator") else None
638 # Make sure the shock variable is age-varying
639 if shock in agent.time_inv:
640 temp = getattr(agent, shock)
641 original_shock_value = temp
642 setattr(agent, shock, agent.T_cycle * [temp])
643 agent.del_from_time_inv(shock)
644 agent.add_to_time_vary(shock)
645 shock_was_time_inv = True
646 else:
647 shock_was_time_inv = False
649 # Solve the long run model if it wasn't already
650 if not solved:
651 t0 = time()
652 agent.solve()
653 t1 = time()
654 if verbose:
655 print(
656 "Solving the long run model took {:.3f}".format(t1 - t0) + " seconds."
657 )
658 LR_soln = deepcopy(agent.solution)
660 try:
661 t0 = time()
662 agent.initialize_sym()
663 X = agent._simulator # for easier referencing
665 # Construct the transition matrices for the long run model
666 X.make_transition_matrices(grids, norm)
667 LR_trans = deepcopy(X.trans_arrays) # the transition matrices in LR model
668 T_age = len(LR_trans)
669 if T_max < T_age:
670 raise ValueError(
671 "T_max must be greater than or equal to T_age in order to pad "
672 "fake_news_array without truncation."
673 )
674 LR_outcomes = []
675 outcome_grids = []
676 for var in outcomes:
677 try:
678 LR_outcomes.append([X.periods[t].matrices[var] for t in range(T_age)])
679 outcome_grids.append([X.periods[t].grids[var] for t in range(T_age)])
680 except KeyError as exc:
681 raise KeyError(
682 "Outcome " + var + " was requested, but no grid was provided!"
683 ) from exc
685 # Extract the normalizing trend
686 if trend is not None:
687 trend_adj_fac = np.array(
688 [X.periods[t].content[trend] for t in range(T_age)]
689 )
690 trend_adj_fac[0] = 1.0
691 trend_adj_cum = np.cumprod(trend_adj_fac)
692 else:
693 trend_adj_cum = np.ones(T_age)
695 t1 = time()
696 if verbose:
697 print(
698 "Making the transition matrix for the long run model took {:.3f}".format(
699 t1 - t0
700 )
701 + " seconds."
702 )
704 # Find the steady state for the long run model
705 t0 = time()
706 X.simulate_cohort_by_grids(outcomes=["dead"] + outcomes, calc_dstn=True)
707 SS_dstn = deepcopy(X.state_dstn_by_age)
708 SS_outcomes = {}
709 for j in range(J):
710 name = outcomes[j]
711 SS_outcomes[name] = [
712 np.dot(LR_outcomes[j][t], outcome_grids[j][t]) for t in range(T_age)
713 ]
715 # Re-apply mortality to downweight older ages
716 survival_by_age = 1.0 - X.history_avg["dead"]
717 survival_by_age[-1] = 0.0 # Force automatic death
718 cum_liv_prb = 1.0
719 pop_sum = 0.0
720 for a in range(T_age):
721 SS_dstn[a] *= cum_liv_prb
722 pop_sum += cum_liv_prb
723 cum_liv_prb *= survival_by_age[a]
725 t1 = time()
726 if verbose:
727 print(
728 "Finding the long run steady state took {:.3f}".format(t1 - t0)
729 + " seconds."
730 )
732 # Construct the "expectation vectors" for all outcomes at all ages
733 t0 = time()
734 E_vecs = {}
735 for j in range(J):
736 name = outcomes[j]
737 E_temp = [[SS_outcomes[name][a].copy()] for a in range(T_age)]
738 for t in range(1, T_age):
739 for a in range(T_age - t):
740 S = survival_by_age[a]
741 E_temp[a].append(np.dot(S * LR_trans[a], E_temp[a + 1][-1]))
742 E_vecs[name] = E_temp
743 t1 = time()
745 # Rearrange the expectation vectors for better access later
746 E_curly = [
747 np.stack(
748 [
749 np.stack([E_vecs[name][a][t] for name in outcomes])
750 for t in range(T_age - a)
751 ]
752 )
753 for a in range(T_age)
754 ]
756 if verbose:
757 print(
758 "Constructing expectation vectors took {:.3f}".format(t1 - t0)
759 + " seconds."
760 )
762 # Each entry of the E_vecs dictionary is a nested list. The outer index of the
763 # list is a, the age at t=0, and the inner index is time period t. The elements
764 # in the nested list are expectation vectors: the expected value of the outcome
765 # in period t conditional on being age a and at state space gridpoint n at t=0.
767 # Initialize the fake news matrices for each output
768 fake_news_array = np.zeros((J, T_age, T_age, T_age))
769 # Dimensions of fake news arrays:
770 # dim 0 --> j: index of outcome variable
771 # dim 1 --> a: age in period t
772 # dim 2 --> t: periods since news arrived
773 # dim 3 --> s: periods ahead about which the news arrived
775 # Loop over ages of the model and have the news shock apply at each one;
776 # k is the age index at which the shock arrives
777 t0 = time()
778 for k in reversed(range(T_age)):
779 # Adjust the timing for "offset" shocks
780 l = k - int(offset)
781 shock_val_orig = getattr(agent, shock)[l]
782 shock_val_new = shock_val_orig + eps
784 # Perturb the shock variable at age k, which corresponds to "solver period" l
785 if l >= 0:
786 getattr(agent, shock)[l] = shock_val_new
788 # Solve the model backwards from age l
789 if construct:
790 agent.construct()
791 agent.solve(from_solution=LR_soln[l + 1], from_t=l + 1)
792 else:
793 agent.solution = LR_soln
795 # Build transitions and outcomes up to age k. Don't use "fake news timing" option!
796 agent.initialize_sym()
797 X = agent._simulator # for easier typing
798 if l < 0:
799 setattr(X.periods[0], shock, shock_val_new)
800 X.make_transition_matrices(grids, norm, for_t=range(k + 1))
801 shocked_trans = deepcopy(X.trans_arrays)
802 shocked_outcomes = []
803 for var in outcomes:
804 temp_outcomes = []
805 for a in range(k + 1):
806 temp_outcomes.append(X.periods[a].matrices[var])
807 shocked_outcomes.append(temp_outcomes)
809 # Update the t=0 row of the fake news matrices
810 for j in range(J):
811 for a in range(k + 1):
812 temp = np.dot(
813 SS_dstn[a], shocked_outcomes[j][a] - LR_outcomes[j][a]
814 )
815 fake_news_array[j, a, 0, k - a] += np.dot(temp, outcome_grids[j][a])
817 # Update the other t rows of the fake news matrices
818 for a in range(k + 1):
819 if a >= T_age - 1:
820 continue
821 S = survival_by_age[a]
822 D_dstn_news = (
823 np.dot(S * shocked_trans[a].T, SS_dstn[a]) - SS_dstn[a + 1]
824 )
825 update_FN_mats(
826 fake_news_array, E_curly[a + 1], D_dstn_news, T_age, a, k
827 )
829 # Reset the shock variable at age l
830 if l >= 0:
831 getattr(agent, shock)[l] = shock_val_orig
833 t1 = time()
834 if verbose:
835 print(
836 "Making fake news arrays for each period of the problem took {:.3f}".format(
837 t1 - t0
838 )
839 + " seconds."
840 )
842 t0 = time()
843 # Pad out the fake news array with zeros
844 FN_pad = np.zeros((J, T_age, T_max, T_max))
845 FN_pad[:, :, :T_age, :T_age] = fake_news_array
847 # Construct age-specific Jacobian matrices
848 SSJ_by_age = FN_pad.copy()
849 for t in range(1, T_max):
850 SSJ_by_age[:, :, 1:, t] += SSJ_by_age[:, :, :-1, t - 1]
852 # Apply normalization factors
853 SSJ_by_age *= np.reshape(trend_adj_cum, (1, T_age, 1, 1))
854 SSJ_by_age /= pop_sum * eps
856 t1 = time()
857 if verbose:
858 print(
859 "Constructing the sequence space Jacobians took {:.3f}".format(t1 - t0)
860 + " seconds."
861 )
863 # Structure and return outputs, aggregating by age if requested
864 SSJ = [SSJ_by_age[j, :, :, :] for j in range(J)]
865 if age_agg:
866 for j in range(J):
867 SSJ[j] = np.sum(SSJ[j], axis=0)
868 if no_list:
869 return SSJ[0]
870 else:
871 return SSJ
873 finally:
874 # Make sure the agent wasn't unexpectedly mutated in this method
875 _restore_agent(agent, LR_soln, simulator_backup)
876 if shock_was_time_inv:
877 setattr(agent, shock, original_shock_value)
878 agent.del_from_time_vary(shock)
879 agent.add_to_time_inv(shock)
882def calc_shock_response_manually(
883 agent,
884 shock,
885 outcomes,
886 grids,
887 s=0,
888 eps=1e-4,
889 T_max=300,
890 norm=None,
891 solved=False,
892 construct=[],
893 offset=False,
894 verbose=False,
895):
896 """
897 Compute an AgentType instance's timepath of outcome responses to learning at
898 t=0 that the named shock variable will be perturbed at t=s. This is equivalent
899 to calculating only the s-th column of the SSJs *manually*, rather than using
900 the fake news algorithm. This function can be used to verify and/or debug the
901 output of the fake news SSJ algorithm.
903 Important: Mortality (or death and replacement generally) should be turned
904 off in the model (via parameter values) for this to work properly. Or does it?
906 Parameters
907 ----------
908 agent : AgentType
909 Agent for which the response(s) should be calculated. Must have T_cycle=1
910 and cycles=0, or the function will throw an error. Must have a model
911 file defined or this won't work at all.
912 shock : str
913 Name of the variable that the response will be computed with respect to.
914 It does not need to be a "shock" in a modeling sense, but it must be a
915 single-valued parameter (possibly a singleton list) that can be changed.
916 outcomes : str or [str]
917 Names of outcome variables of interest; an SSJ matrix will be constructed
918 for each variable named here. If a single string is passed, the output
919 will be a single np.array. If a list of strings are passed, the output
920 will be a list of dYdX vectors in the order specified here.
921 grids : dict
922 Dictionary of dictionaries with discretizing grid information. The grids
923 should include all arrival variables other than those that are normalized
924 out. They should also include all variables named in outcomes, except
925 outcomes that are continuation variables that remap to arrival variables.
926 Grid specification must include number of nodes N, should also include
927 min and max if the variable is continuous.
928 s : int
929 Period in which the shock variable is perturbed, relative to current t=0.
930 The default is 0.
931 eps : float
932 Amount by which to perturb the shock variable. The default is 1e-4.
933 T_max : int
934 The length of the simulation for this exercise. The default is 300.
935 norm : str or None
936 Name of the model variable to normalize by for Harmenberg aggregation,
937 if any. For many HARK models, this should be 'PermShk', which enables
938 the grid over permanent income to be omitted as an explicit state.
939 solved : bool
940 Whether the agent's model has already been solved. If False (default),
941 it will be solved as the very first step.
942 construct : [str]
943 List of constructed objects that will be changed by perturbing shock.
944 These should all share an "offset status" (True or False). Default is [].
945 offset : bool
946 Whether the shock variable is "offset in time" for the solver, with a
947 default of False. This should be set to True if the named shock variable
948 (or the constructed model input that it affects) is indexed by t+1 from
949 the perspective of the solver. For example, the period t solver for the
950 ConsIndShock model takes in risk free interest factor Rfree as an argument,
951 but it represents the value of R that will occur at the start of t+1.
952 verbose : bool
953 Whether to display timing/progress to screen. The default is False.
955 Returns
956 -------
957 dYdX : np.array or [np.array]
958 One or more vectors of length T_max.
959 """
960 setup = _prepare_ssj_computation(agent, outcomes, grids, norm, solved, verbose)
961 outcomes = setup["outcomes"]
962 no_list = setup["no_list"]
963 simulator_backup = setup["simulator_backup"]
964 LR_soln = setup["LR_soln"]
965 X = setup["X"]
966 LR_outcomes = setup["LR_outcomes"]
967 outcome_grids = setup["outcome_grids"]
968 SS_dstn = setup["SS_dstn"]
970 SS_outcomes = [np.dot(mat.T, SS_dstn) for mat in LR_outcomes]
971 SS_avgs = [np.dot(ss, grid) for ss, grid in zip(SS_outcomes, outcome_grids)]
973 try:
974 # Make a temporary agent to construct the perturbed constructed objects
975 t0 = time()
976 temp_agent = deepcopy(agent)
977 base_shock_value, shock_is_list = _perturb_shock(agent, shock)
978 if shock_is_list:
979 temp_value = [base_shock_value + eps]
980 else:
981 temp_value = base_shock_value + eps
982 temp_dict = {shock: temp_value}
983 temp_agent.assign_parameters(**temp_dict)
984 if len(construct) > 0:
985 temp_agent.update()
986 for var in construct:
987 temp_dict[var] = getattr(temp_agent, var)
989 # Build the finite horizon version of this agent
990 FH_agent = deepcopy(agent)
991 FH_agent.del_param("solution")
992 FH_agent.del_param("_simulator")
993 FH_agent.del_from_time_vary("solution")
994 FH_agent.del_from_time_inv(shock)
995 FH_agent.add_to_time_vary(shock)
996 FH_agent.del_from_time_inv(*construct)
997 FH_agent.add_to_time_vary(*construct)
998 finite_dict = {"T_cycle": T_max, "cycles": 1}
999 for var in FH_agent.time_vary:
1000 if var in construct:
1001 sequence = [deepcopy(getattr(agent, var)[0]) for t in range(T_max)]
1002 sequence[s] = deepcopy(getattr(temp_agent, var)[0])
1003 else:
1004 sequence = T_max * [deepcopy(getattr(agent, var)[0])]
1005 finite_dict[var] = sequence
1006 shock_seq = T_max * [base_shock_value]
1007 shock_seq[s] = base_shock_value + eps
1008 finite_dict[shock] = shock_seq
1009 FH_agent.assign_parameters(**finite_dict)
1010 del temp_agent
1011 t1 = time()
1012 if verbose:
1013 print(
1014 "Building the finite horizon agent took {:.3f}".format(t1 - t0)
1015 + " seconds."
1016 )
1018 # Solve the finite horizon agent
1019 t0 = time()
1020 FH_agent.solve(from_solution=LR_soln)
1021 t1 = time()
1022 if verbose:
1023 print(
1024 "Solving the "
1025 + str(T_max)
1026 + " period problem took {:.3f}".format(t1 - t0)
1027 + " seconds."
1028 )
1030 # Build transition matrices for the finite horizon problem
1031 t0 = time()
1032 FH_agent.initialize_sym()
1033 FH_agent._simulator.make_transition_matrices(
1034 grids, norm=norm, fake_news_timing=True
1035 )
1036 t1 = time()
1037 if verbose:
1038 print(
1039 "Constructing transition matrices took {:.3f}".format(t1 - t0)
1040 + " seconds."
1041 )
1043 # Use grid simulation to find the timepath of requested variables, and compute
1044 # the derivative with respect to baseline outcomes
1045 t0 = time()
1046 FH_agent._simulator.simulate_cohort_by_grids(outcomes, from_dstn=SS_dstn)
1047 dYdX = []
1048 for j, var in enumerate(outcomes):
1049 diff_path = (FH_agent._simulator.history_avg[var] - SS_avgs[j]) / eps
1050 if offset:
1051 dYdX.append(diff_path[1:])
1052 else:
1053 dYdX.append(diff_path[:-1])
1054 t1 = time()
1055 if verbose:
1056 print(
1057 "Calculating impulse responses by grid simulation took {:.3f}".format(
1058 t1 - t0
1059 )
1060 + " seconds."
1061 )
1063 del FH_agent
1064 if no_list:
1065 return dYdX[0]
1066 else:
1067 return dYdX
1068 finally:
1069 _restore_agent(agent, LR_soln, simulator_backup)
1072@njit
1073def calc_derivs_of_state_dstns(T, J, trans_by_t, trans_LR, SS_dstn): # pragma: no cover
1074 """
1075 Numba-compatible helper function to calculate the derivative of the state
1076 distribution by period.
1078 Parameters
1079 ----------
1080 T : int
1081 Maximum time horizon for the fake news algorithm.
1082 J : int
1083 Number of outcomes of interest.
1084 trans_by_t : np.array
1085 Array of shape (T,K,K) representing the transition matrix in each period.
1086 trans_LR : np.array
1087 Array of shape (K,K) representing the long run transition matrix.
1088 SS_dstn : np.array
1089 Array of size K representing the long run steady state distribution.
1091 Returns
1092 -------
1093 D_dstn_news : np.array
1094 Array of shape (T,K) representing dD_1^s from the SSJ paper, where K
1095 is the number of arrival state space nodes.
1097 """
1098 K = SS_dstn.size
1099 D_dstn_news = np.empty((T, K)) # this is dD_1^s in the SSJ paper (equation 24)
1100 for t in range(T - 1, -1, -1):
1101 D_dstn_news[T - t - 1, :] = np.dot((trans_by_t[t, :, :] - trans_LR).T, SS_dstn)
1102 return D_dstn_news
1105@njit
1106def calc_derivs_of_policy_funcs(T, Y_by_t, Y_LR, Y_grid, SS_dstn): # pragma: no cover
1107 """
1108 Numba-compatible helper function to calculate the derivative of an outcome
1109 function in each period.
1111 Parameters
1112 ----------
1113 T : int
1114 Maximum time horizon for the fake news algorithm.
1115 Y_by_t : np.array
1116 Array of shape (T,K,N) with the stochastic outcome, mapping from K arrival
1117 state space nodes to N outcome space nodes, for each of the T periods.
1118 Y_LR : np.array
1119 Array of shape (K,N) representing the stochastic outcome in the long run.
1120 Y_grid : np.array
1121 Array of size N representing outcome space gridpoints.
1122 SS_dstn : np.array
1123 Array of size K representing the long run steady state distribution.
1125 Returns
1126 -------
1127 dY_news : np.array
1128 Array of size T representing the change in average outcome in each period
1129 when the shock arrives unexpectedly in that period.
1130 """
1131 dY_news = np.empty(T) # this is dY_0^s in the SSJ paper (equation 24)
1132 for t in range(T - 1, -1, -1):
1133 temp = (Y_by_t[t, :, :] - Y_LR).T
1134 dY_news[T - t - 1] = np.dot(np.dot(temp, SS_dstn), Y_grid)
1135 return dY_news
1138@njit
1139def make_fake_news_matrices(T, J, dY, D_dstn, trans_LR, E): # pragma: no cover
1140 """
1141 Numba-compatible function to calculate the fake news array from first order
1142 perturbation information.
1144 Parameters
1145 ----------
1146 T : int
1147 Maximum time horizon for the fake news algorithm.
1148 J : int
1149 Number of outcomes of interest.
1150 dY : int
1151 Array shape (T,J) representing dY_0 from the SSJ paper.
1152 D_dstn : np.array
1153 Array of shape (T,K) representing dD_1^s from the SSJ paper, where K
1154 is the number of arrival state space nodes.
1155 trans_LR : np.array
1156 Array of shape (K,K) representing the transpose of the long run transition matrix.
1157 E : np.array
1158 Initial expectation vectors combined into a single array of shape (J,K).
1160 Returns
1161 -------
1162 FN : np.array
1163 Fake news array of shape (J,T,T).
1164 """
1165 FN = np.empty((J, T, T))
1166 FN[:, 0, :] = dY.T # Fill in row zero
1167 for t in range(1, T): # Loop over other rows
1168 for s in range(T):
1169 FN[:, t, s] = np.dot(E, D_dstn[s, :])
1170 E = np.dot(E, trans_LR)
1171 return FN
1174@njit
1175def update_FN_mats(FN_mats, evecs, dD1, A, a, k): # pragma: no cover
1176 """
1177 This is adapted from Mateo's code.
1179 FN_mats: (J, A, A, A)
1180 evecs : (T, J, G)
1181 dD1 : (G,)
1182 """
1183 J = FN_mats.shape[0]
1184 G = dD1.shape[0]
1186 for j in range(J):
1187 for t in range(1, A - a):
1188 # compute dot(evecs[t-1, j, :], dD1) by hand
1189 v = 0.0
1190 for g in range(G):
1191 v += evecs[t - 1, j, g] * dD1[g]
1192 FN_mats[j, a + t, t, k - a] += v