Coverage for HARK/SSJutils.py: 88%
262 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-02 05:14 +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 make_basic_SSJ_matrices(
14 agent,
15 shock,
16 outcomes,
17 grids,
18 eps=1e-4,
19 T_max=300,
20 norm=None,
21 solved=False,
22 construct=True,
23 offset=False,
24 verbose=False,
25):
26 """
27 Constructs one or more sequence space Jacobian (SSJ) matrices for specified
28 outcomes over one shock variable. It is "basic" in the sense that it only
29 works for "one period infinite horizon" models, as in the original SSJ paper.
31 Parameters
32 ----------
33 agent : AgentType
34 Agent for which the SSJ(s) should be constructed. Must have T_cycle=1
35 and cycles=0, or the function will throw an error. Must have a model
36 file defined or this won't work at all.
37 shock : str
38 Name of the variable that Jacobians will be computed with respect to.
39 It does not need to be a "shock" in a modeling sense, but it must be a
40 single-valued parameter (possibly a singleton list) that can be changed.
41 outcomes : str or [str]
42 Names of outcome variables of interest; an SSJ matrix will be constructed
43 for each variable named here. If a single string is passed, the output
44 will be a single np.array. If a list of strings are passed, the output
45 will be a list of SSJ matrices in the order specified here.
46 grids : dict
47 Dictionary of dictionaries with discretizing grid information. The grids
48 should include all arrival variables other than those that are normalized
49 out. They should also include all variables named in outcomes, except
50 outcomes that are continuation variables that remap to arrival variables.
51 Grid specification must include number of nodes N, should also include
52 min and max if the variable is continuous.
53 eps : float
54 Amount by which to perturb the shock variable. The default is 1e-4.
55 T_max : int
56 Size of the SSJ matrices: the maximum number of periods to consider.
57 The default is 300.
58 norm : str or None
59 Name of the model variable to normalize by for Harmenberg aggregation,
60 if any. For many HARK models, this should be 'PermShk', which enables
61 the grid over permanent income to be omitted as an explicit state.
62 solved : bool
63 Whether the agent's model has already been solved. If False (default),
64 it will be solved as the very first step. Solving the agent's long run
65 model before constructing SSJ matrices has the advantage of not needing
66 to re-solve the long run model for each shock variable.
67 construct : bool
68 Whether the construct (update) method should be run after the shock is
69 updated. The default is True, which is the "safe" option. If the shock
70 variable is a parameter that enters the model only *directly*, rather
71 than being used to build a more complex model input, then this can be
72 set to False to save a (very) small amount of time during computation.
73 If it is set to False improperly, the SSJs will be very wrong, potentially
74 just zero everywhere.
75 offset : bool
76 Whether the shock variable is "offset in time" for the solver, with a
77 default of False. This should be set to True if the named shock variable
78 (or the constructed model input that it affects) is indexed by t+1 from
79 the perspective of the solver. For example, the period t solver for the
80 ConsIndShock model takes in risk free interest factor Rfree as an argument,
81 but it represents the value of R that will occur at the start of t+1.
82 verbose : bool
83 Whether to display timing/progress to screen. The default is False.
85 Returns
86 -------
87 SSJ : np.array or [np.array]
88 One or more sequence space Jacobian arrays over the outcome variables
89 with respect to the named shock variable.
90 """
91 if (agent.cycles > 0) or (agent.T_cycle != 1):
92 raise ValueError(
93 "This function is only compatible with one period infinite horizon models!"
94 )
95 if not isinstance(outcomes, list):
96 outcomes = [outcomes]
97 no_list = True
98 else:
99 no_list = False
101 # Store the simulator if it exists
102 if hasattr(agent, "_simulator"):
103 simulator_backup = agent._simulator
105 # Solve the long run model if it wasn't already
106 if not solved:
107 t0 = time()
108 agent.solve()
109 t1 = time()
110 if verbose:
111 print(
112 "Solving the long run model took {:.3f}".format(t1 - t0) + " seconds."
113 )
114 LR_soln = deepcopy(agent.solution[0])
116 # Construct the transition matrix for the long run model
117 t0 = time()
118 agent.initialize_sym()
119 X = agent._simulator # for easier referencing
120 X.make_transition_matrices(grids, norm)
121 LR_trans = X.trans_arrays[0].copy() # the transition matrix in LR model
122 LR_period = X.periods[0]
123 LR_outcomes = []
124 outcome_grids = []
125 for var in outcomes:
126 try:
127 LR_outcomes.append(X.periods[0].matrices[var])
128 outcome_grids.append(X.periods[0].grids[var])
129 except:
130 raise ValueError(
131 "Outcome " + var + " was requested, but no grid was provided!"
132 )
133 t1 = time()
134 if verbose:
135 print(
136 "Making the transition matrix for the long run model took {:.3f}".format(
137 t1 - t0
138 )
139 + " seconds."
140 )
142 # Find the steady state for the long run model
143 t0 = time()
144 X.find_steady_state()
145 SS_dstn = X.steady_state_dstn.copy()
146 SS_outcomes = []
147 for j in range(len(outcomes)):
148 SS_outcomes.append(np.dot(LR_outcomes[j].transpose(), SS_dstn))
149 t1 = time()
150 if verbose:
151 print(
152 "Finding the long run steady state took {:.3f}".format(t1 - t0)
153 + " seconds."
154 )
156 # Solve back one period while perturbing the shock variable
157 t0 = time()
158 try:
159 base = getattr(agent, shock)
160 except:
161 raise ValueError(
162 "The agent doesn't have anything called " + shock + " to perturb!"
163 )
164 if isinstance(base, list):
165 base_shock_value = base[0]
166 shock_is_list = True
167 else:
168 base_shock_value = base
169 shock_is_list = False
170 if not isinstance(base_shock_value, float):
171 raise TypeError(
172 "Only a single real-valued object can be perturbed in this way!"
173 )
174 agent.cycles = 1
175 if shock_is_list:
176 temp_value = [base_shock_value + eps]
177 else:
178 temp_value = base_shock_value + eps
179 temp_dict = {shock: temp_value}
180 agent.assign_parameters(**temp_dict)
181 if construct:
182 agent.update()
183 agent.solve(from_solution=LR_soln)
184 agent.initialize_sym()
185 Tm1_soln = deepcopy(agent.solution[0])
186 period_Tm1 = agent._simulator.periods[0]
187 period_T = agent._simulator.periods[-1]
188 t1 = time()
189 if verbose:
190 print(
191 "Solving period T-1 with a perturbed variable took {:.3f}".format(t1 - t0)
192 + " seconds."
193 )
195 # Set up and solve the agent for T_max-1 more periods
196 t0 = time()
197 agent.cycles = T_max - 1
198 if shock_is_list:
199 orig_dict = {shock: [base_shock_value]}
200 else:
201 orig_dict = {shock: base_shock_value}
202 agent.assign_parameters(**orig_dict)
203 if construct:
204 agent.update()
205 agent.solve(from_solution=Tm1_soln)
206 t1 = time()
207 if verbose:
208 print(
209 "Solving the finite horizon model for "
210 + str(T_max - 1)
211 + " more periods took {:.3f}".format(t1 - t0)
212 + " seconds."
213 )
215 # Construct transition and outcome matrices for the "finite horizon"
216 t0 = time()
217 agent.initialize_sym()
218 X = agent._simulator # for easier typing
219 X.periods[-1] = period_Tm1 # substitute period T-1 from above
220 if offset:
221 for name in X.periods[-1].content.keys():
222 if name not in X.solution: # sub in proper T-1 non-solution info
223 X.periods[-1].content[name] = LR_period.content[name]
224 X.periods[-1].distribute_content()
225 X.periods = X.periods[1:] + [period_T]
226 X.make_transition_matrices(grids, norm, fake_news_timing=True)
227 TmX_trans = deepcopy(X.trans_arrays)
228 TmX_outcomes = []
229 for t in range(T_max):
230 Tmt_outcomes = []
231 for var in outcomes:
232 Tmt_outcomes.append(X.periods[t].matrices[var])
233 TmX_outcomes.append(Tmt_outcomes)
234 t1 = time()
235 if verbose:
236 print(
237 "Constructing transition arrays for the finite horizon model took {:.3f}".format(
238 t1 - t0
239 )
240 + " seconds."
241 )
243 # Calculate derivatives of transition and outcome matrices by first differences
244 t0 = time()
245 J = len(outcomes)
246 K = SS_dstn.size
247 D_dstn_array = calc_derivs_of_state_dstns(
248 T_max, J, np.array(TmX_trans), LR_trans, SS_dstn
249 )
250 dY_news_array = np.empty((T_max, J))
251 for j in range(J):
252 temp_outcomes = np.array([TmX_outcomes[t][j] for t in range(T_max)])
253 dY_news_array[:, j] = calc_derivs_of_policy_funcs(
254 T_max, temp_outcomes, LR_outcomes[j], outcome_grids[j], SS_dstn
255 )
256 t1 = time()
257 if verbose:
258 print(
259 "Calculating derivatives by first differences took {:.3f}".format(t1 - t0)
260 + " seconds."
261 )
263 # Construct the "fake news" matrices, one for each outcome variable
264 t0 = time()
265 expectation_vectors = np.empty((J, K)) # Initialize expectation vectors
266 for j in range(J):
267 expectation_vectors[j, :] = np.dot(LR_outcomes[j], outcome_grids[j])
268 FN = make_fake_news_matrices(
269 T_max, J, dY_news_array, D_dstn_array, LR_trans.T, expectation_vectors.copy()
270 )
271 t1 = time()
272 if verbose:
273 print(
274 "Constructing the fake news matrices took {:.3f}".format(t1 - t0)
275 + " seconds."
276 )
278 # Construct the SSJ matrices, one for each outcome variable
279 t0 = time()
280 SSJ_array = calc_ssj_from_fake_news_matrices(T_max, J, FN, eps)
281 SSJ = [SSJ_array[j, :, :] for j in range(J)] # unpack into a list of arrays
282 t1 = time()
283 if verbose:
284 print(
285 "Constructing the sequence space Jacobians took {:.3f}".format(t1 - t0)
286 + " seconds."
287 )
289 # Reset the agent to its original state and return the output
290 agent.solution = [LR_soln]
291 agent.cycles = 0
292 agent._simulator.reset()
293 try:
294 agent._simulator = simulator_backup
295 except:
296 del agent._simulator
297 if no_list:
298 return SSJ[0]
299 else:
300 return SSJ
303def calc_shock_response_manually(
304 agent,
305 shock,
306 outcomes,
307 grids,
308 s=0,
309 eps=1e-4,
310 T_max=300,
311 norm=None,
312 solved=False,
313 construct=[],
314 offset=False,
315 verbose=False,
316):
317 """
318 Compute an AgentType instance's timepath of outcome responses to learning at
319 t=0 that the named shock variable will be perturbed at t=s. This is equivalent
320 to calculating only the s-th column of the SSJs *manually*, rather than using
321 the fake news algorithm. This function can be used to verify and/or debug the
322 output of the fake news SSJ algorithm.
324 Important: Mortality (or death and replacement generally) should be turned
325 off in the model (via parameter values) for this to work properly. Or does it?
327 Parameters
328 ----------
329 agent : AgentType
330 Agent for which the response(s) should be calculated. Must have T_cycle=1
331 and cycles=0, or the function will throw an error. Must have a model
332 file defined or this won't work at all.
333 shock : str
334 Name of the variable that the response will be computed with respect to.
335 It does not need to be a "shock" in a modeling sense, but it must be a
336 single-valued parameter (possibly a singleton list) that can be changed.
337 outcomes : str or [str]
338 Names of outcome variables of interest; an SSJ matrix will be constructed
339 for each variable named here. If a single string is passed, the output
340 will be a single np.array. If a list of strings are passed, the output
341 will be a list of dYdX vectors in the order specified here.
342 grids : dict
343 Dictionary of dictionaries with discretizing grid information. The grids
344 should include all arrival variables other than those that are normalized
345 out. They should also include all variables named in outcomes, except
346 outcomes that are continuation variables that remap to arrival variables.
347 Grid specification must include number of nodes N, should also include
348 min and max if the variable is continuous.
349 s : int
350 Period in which the shock variable is perturbed, relative to current t=0.
351 The default is 0.
352 eps : float
353 Amount by which to perturb the shock variable. The default is 1e-4.
354 T_max : int
355 The length of the simulation for this exercise. The default is 300.
356 norm : str or None
357 Name of the model variable to normalize by for Harmenberg aggregation,
358 if any. For many HARK models, this should be 'PermShk', which enables
359 the grid over permanent income to be omitted as an explicit state.
360 solved : bool
361 Whether the agent's model has already been solved. If False (default),
362 it will be solved as the very first step.
363 construct : [str]
364 List of constructed objects that will be changed by perturbing shock.
365 These should all share an "offset status" (True or False). Default is [].
366 offset : bool
367 Whether the shock variable is "offset in time" for the solver, with a
368 default of False. This should be set to True if the named shock variable
369 (or the constructed model input that it affects) is indexed by t+1 from
370 the perspective of the solver. For example, the period t solver for the
371 ConsIndShock model takes in risk free interest factor Rfree as an argument,
372 but it represents the value of R that will occur at the start of t+1.
373 verbose : bool
374 Whether to display timing/progress to screen. The default is False.
376 Returns
377 -------
378 dYdX : np.array or [np.array]
379 One or more vectors of length T_max.
380 """
381 if (agent.cycles > 0) or (agent.T_cycle != 1):
382 raise ValueError(
383 "This function is only compatible with one period infinite horizon models!"
384 )
385 if not isinstance(outcomes, list):
386 outcomes = [outcomes]
387 no_list = True
388 else:
389 no_list = False
391 # Store the simulator if it exists
392 if hasattr(agent, "_simulator"):
393 simulator_backup = agent._simulator
395 # Solve the long run model if it wasn't already
396 if not solved:
397 t0 = time()
398 agent.solve()
399 t1 = time()
400 if verbose:
401 print(
402 "Solving the long run model took {:.3f}".format(t1 - t0) + " seconds."
403 )
404 LR_soln = deepcopy(agent.solution[0])
406 # Construct the transition matrix for the long run model
407 t0 = time()
408 agent.initialize_sym()
409 X = agent._simulator # for easier referencing
410 X.make_transition_matrices(grids, norm)
411 LR_outcomes = []
412 outcome_grids = []
413 for var in outcomes:
414 try:
415 LR_outcomes.append(X.periods[0].matrices[var])
416 outcome_grids.append(X.periods[0].grids[var])
417 except:
418 raise ValueError(
419 "Outcome " + var + " was requested, but no grid was provided!"
420 )
421 t1 = time()
422 if verbose:
423 print(
424 "Making the transition matrix for the long run model took {:.3f}".format(
425 t1 - t0
426 )
427 + " seconds."
428 )
430 # Find the steady state for the long run model
431 t0 = time()
432 X.find_steady_state()
433 SS_dstn = X.steady_state_dstn.copy()
434 SS_outcomes = []
435 SS_avgs = []
436 for j in range(len(outcomes)):
437 SS_outcomes.append(np.dot(LR_outcomes[j].transpose(), SS_dstn))
438 SS_avgs.append(np.dot(SS_outcomes[j], outcome_grids[j]))
439 t1 = time()
440 if verbose:
441 print(
442 "Finding the long run steady state took {:.3f}".format(t1 - t0)
443 + " seconds."
444 )
446 # Make a temporary agent to construct the perturbed constructed objects
447 t0 = time()
448 temp_agent = deepcopy(agent)
449 try:
450 base = getattr(agent, shock)
451 except:
452 raise ValueError(
453 "The agent doesn't have anything called " + shock + " to perturb!"
454 )
455 if isinstance(base, list):
456 base_shock_value = base[0]
457 shock_is_list = True
458 else:
459 base_shock_value = base
460 shock_is_list = False
461 if not isinstance(base_shock_value, float):
462 raise TypeError(
463 "Only a single real-valued object can be perturbed in this way!"
464 )
465 if shock_is_list:
466 temp_value = [base_shock_value + eps]
467 else:
468 temp_value = base_shock_value + eps
469 temp_dict = {shock: temp_value}
470 temp_agent.assign_parameters(**temp_dict)
471 if len(construct) > 0:
472 temp_agent.update()
473 for var in construct:
474 temp_dict[var] = getattr(temp_agent, var)
476 # Build the finite horizon version of this agent
477 FH_agent = deepcopy(agent)
478 FH_agent.del_param("solution")
479 FH_agent.del_param("_simulator")
480 FH_agent.del_from_time_vary("solution")
481 FH_agent.del_from_time_inv(shock)
482 FH_agent.add_to_time_vary(shock)
483 FH_agent.del_from_time_inv(*construct)
484 FH_agent.add_to_time_vary(*construct)
485 finite_dict = {"T_cycle": T_max, "cycles": 1}
486 for var in FH_agent.time_vary:
487 if var in construct:
488 sequence = [deepcopy(getattr(agent, var)[0]) for t in range(T_max)]
489 sequence[s] = deepcopy(getattr(temp_agent, var)[0])
490 else:
491 sequence = T_max * [deepcopy(getattr(agent, var)[0])]
492 finite_dict[var] = sequence
493 shock_seq = T_max * [base_shock_value]
494 shock_seq[s] = base_shock_value + eps
495 finite_dict[shock] = shock_seq
496 FH_agent.assign_parameters(**finite_dict)
497 del temp_agent
498 t1 = time()
499 if verbose:
500 print(
501 "Building the finite horizon agent took {:.3f}".format(t1 - t0)
502 + " seconds."
503 )
505 # Solve the finite horizon agent
506 t0 = time()
507 FH_agent.solve(from_solution=LR_soln)
508 t1 = time()
509 if verbose:
510 print(
511 "Solving the "
512 + str(T_max)
513 + " period problem took {:.3f}".format(t1 - t0)
514 + " seconds."
515 )
517 # Build transition matrices for the finite horizon problem
518 t0 = time()
519 FH_agent.initialize_sym()
520 FH_agent._simulator.make_transition_matrices(
521 grids, norm=norm, fake_news_timing=True
522 )
523 t1 = time()
524 if verbose:
525 print(
526 "Constructing transition matrices took {:.3f}".format(t1 - t0) + " seconds."
527 )
529 # Use grid simulation to find the timepath of requested variables, and compute
530 # the derivative with respect to baseline outcomes
531 t0 = time()
532 FH_agent._simulator.simulate_cohort_by_grids(outcomes, from_dstn=SS_dstn)
533 dYdX = []
534 for j, var in enumerate(outcomes):
535 diff_path = (FH_agent._simulator.history_avg[var] - SS_avgs[j]) / eps
536 if offset:
537 dYdX.append(diff_path[1:])
538 else:
539 dYdX.append(diff_path[:-1])
540 t1 = time()
541 if verbose:
542 print(
543 "Calculating impulse responses by grid simulation took {:.3f}".format(
544 t1 - t0
545 )
546 + " seconds."
547 )
549 # Reset the agent to its original state and return the output
550 del FH_agent
551 agent.solution = [LR_soln]
552 agent.cycles = 0
553 agent._simulator.reset()
554 try:
555 agent._simulator = simulator_backup
556 except:
557 del agent._simulator
558 if no_list:
559 return dYdX[0]
560 else:
561 return dYdX
564@njit
565def calc_derivs_of_state_dstns(T, J, trans_by_t, trans_LR, SS_dstn): # pragma: no cover
566 """
567 Numba-compatible helper function to calculate the derivative of the state
568 distribution by period.
570 Parameters
571 ----------
572 T : int
573 Maximum time horizon for the fake news algorithm.
574 J : int
575 Number of outcomes of interest.
576 trans_by_t : np.array
577 Array of shape (T,K,K) representing the transition matrix in each period.
578 trans_LR : np.array
579 Array of shape (K,K) representing the long run transition matrix.
580 SS_dstn : np.array
581 Array of size K representing the long run steady state distribution.
583 Returns
584 -------
585 D_dstn_news : np.array
586 Array of shape (T,K) representing dD_1^s from the SSJ paper, where K
587 is the number of arrival state space nodes.
589 """
590 K = SS_dstn.size
591 D_dstn_news = np.empty((T, K)) # this is dD_1^s in the SSJ paper (equation 24)
592 for t in range(T - 1, -1, -1):
593 D_dstn_news[T - t - 1, :] = np.dot((trans_by_t[t, :, :] - trans_LR).T, SS_dstn)
594 return D_dstn_news
597@njit
598def calc_derivs_of_policy_funcs(T, Y_by_t, Y_LR, Y_grid, SS_dstn): # pragma: no cover
599 """
600 Numba-compatible helper function to calculate the derivative of an outcome
601 function in each period.
603 Parameters
604 ----------
605 T : int
606 Maximum time horizon for the fake news algorithm.
607 Y_by_t : np.array
608 Array of shape (T,K,N) with the stochastic outcome, mapping from K arrival
609 state space nodes to N outcome space nodes, for each of the T periods.
610 Y_LR : np.array
611 Array of shape (K,N) representing the stochastic outcome in the long run.
612 Y_grid : np.array
613 Array of size N representing outcome space gridpoints.
614 SS_dstn : np.array
615 Array of size K representing the long run steady state distribution.
617 Returns
618 -------
619 dY_news : np.array
620 Array of size T representing the change in average outcome in each period
621 when the shock arrives unexpectedly in that period.
622 """
623 dY_news = np.empty(T) # this is dY_0^s in the SSJ paper (equation 24)
624 for t in range(T - 1, -1, -1):
625 temp = (Y_by_t[t, :, :] - Y_LR).T
626 dY_news[T - t - 1] = np.dot(np.dot(temp, SS_dstn), Y_grid)
627 return dY_news
630@njit
631def make_fake_news_matrices(T, J, dY, D_dstn, trans_LR, E): # pragma: no cover
632 """
633 Numba-compatible function to calculate the fake news array from first order
634 perturbation information.
636 Parameters
637 ----------
638 T : int
639 Maximum time horizon for the fake news algorithm.
640 J : int
641 Number of outcomes of interest.
642 dY : int
643 Array shape (T,J) representing dY_0 from the SSJ paper.
644 D_dstn : np.array
645 Array of shape (T,K) representing dD_1^s from the SSJ paper, where K
646 is the number of arrival state space nodes.
647 trans_LR : np.array
648 Array of shape (K,K) representing the transpose of the long run transition matrix.
649 E : np.array
650 Initial expectation vectors combined into a single array of shape (J,K).
652 Returns
653 -------
654 FN : np.array
655 Fake news array of shape (J,T,T).
656 """
657 FN = np.empty((J, T, T))
658 FN[:, 0, :] = dY.T # Fill in row zero
659 for t in range(1, T): # Loop over other rows
660 for s in range(T):
661 FN[:, t, s] = np.dot(E, D_dstn[s, :])
662 E = np.dot(E, trans_LR)
663 return FN
666@njit
667def calc_ssj_from_fake_news_matrices(T, J, FN, dx): # pragma: no cover
668 """
669 Numba-compatible function to calculate the HA-SSJ from fake news matrices.
671 Parameters
672 ----------
673 T : int
674 Maximum time horizon for the fake news algorithm.
675 J : int
676 Number of outcomes of interest.
677 FN : np.array
678 Fake news array of shape (J,T,T).
679 dx : float
680 Size of the perturbation of the shock variables (epsilon).
682 Returns
683 -------
684 SSJ : np.array
685 HA-SSJ array of shape (J,T,T).
686 """
687 SSJ = np.empty((J, T, T))
688 SSJ[:, 0, :] = FN[:, 0, :] # Fill in row zero
689 SSJ[:, :, 0] = FN[:, :, 0] # Fill in column zero
690 for t in range(1, T): # Loop over other rows
691 for s in range(1, T): # Loop over other columns
692 SSJ[:, t, s] = SSJ[:, t - 1, s - 1] + FN[:, t, s]
693 SSJ *= dx**-1.0 # Scale by dx
694 return SSJ