Coverage for HARK / SSJutils.py: 97%
234 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-03-08 05:31 +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 X = setup["X"]
280 LR_trans = setup["LR_trans"]
281 LR_period = setup["LR_period"]
282 LR_outcomes = setup["LR_outcomes"]
283 outcome_grids = setup["outcome_grids"]
284 SS_dstn = setup["SS_dstn"]
286 SS_outcomes = [np.dot(mat.T, SS_dstn) for mat in LR_outcomes]
288 try:
289 # Solve back one period while perturbing the shock variable
290 t0 = time()
291 base_shock_value, shock_is_list = _perturb_shock(agent, shock)
292 agent.cycles = 1
293 if shock_is_list:
294 temp_value = [base_shock_value + eps]
295 else:
296 temp_value = base_shock_value + eps
297 temp_dict = {shock: temp_value}
298 agent.assign_parameters(**temp_dict)
299 if construct:
300 agent.update()
301 agent.solve(from_solution=LR_soln)
302 agent.initialize_sym()
303 Tm1_soln = deepcopy(agent.solution[0])
304 period_Tm1 = agent._simulator.periods[0]
305 period_T = agent._simulator.periods[-1]
306 t1 = time()
307 if verbose:
308 print(
309 "Solving period T-1 with a perturbed variable took {:.3f}".format(
310 t1 - t0
311 )
312 + " seconds."
313 )
315 # Set up and solve the agent for T_max-1 more periods
316 t0 = time()
317 agent.cycles = T_max - 1
318 if shock_is_list:
319 orig_dict = {shock: [base_shock_value]}
320 else:
321 orig_dict = {shock: base_shock_value}
322 agent.assign_parameters(**orig_dict)
323 if construct:
324 agent.update()
325 agent.solve(from_solution=Tm1_soln)
326 t1 = time()
327 if verbose:
328 print(
329 "Solving the finite horizon model for "
330 + str(T_max - 1)
331 + " more periods took {:.3f}".format(t1 - t0)
332 + " seconds."
333 )
335 # Construct transition and outcome matrices for the "finite horizon"
336 t0 = time()
337 agent.initialize_sym()
338 X = agent._simulator # for easier typing
339 X.periods[-1] = period_Tm1 # substitute period T-1 from above
340 if offset:
341 for name in X.periods[-1].content.keys():
342 if name not in X.solution: # sub in proper T-1 non-solution info
343 X.periods[-1].content[name] = LR_period.content[name]
344 X.periods[-1].distribute_content()
345 X.periods = X.periods[1:] + [period_T]
346 X.make_transition_matrices(grids, norm, fake_news_timing=True)
347 TmX_trans = deepcopy(X.trans_arrays)
348 TmX_outcomes = []
349 for t in range(T_max):
350 Tmt_outcomes = []
351 for var in outcomes:
352 Tmt_outcomes.append(X.periods[t].matrices[var])
353 TmX_outcomes.append(Tmt_outcomes)
354 t1 = time()
355 if verbose:
356 print(
357 "Constructing transition arrays for the finite horizon model took {:.3f}".format(
358 t1 - t0
359 )
360 + " seconds."
361 )
363 # Calculate derivatives of transition and outcome matrices by first differences
364 t0 = time()
365 J = len(outcomes)
366 K = SS_dstn.size
367 D_dstn_array = calc_derivs_of_state_dstns(
368 T_max, J, np.array(TmX_trans), LR_trans, SS_dstn
369 )
370 dY_news_array = np.empty((T_max, J))
371 for j in range(J):
372 temp_outcomes = np.array([TmX_outcomes[t][j] for t in range(T_max)])
373 dY_news_array[:, j] = calc_derivs_of_policy_funcs(
374 T_max, temp_outcomes, LR_outcomes[j], outcome_grids[j], SS_dstn
375 )
376 t1 = time()
377 if verbose:
378 print(
379 "Calculating derivatives by first differences took {:.3f}".format(
380 t1 - t0
381 )
382 + " seconds."
383 )
385 # Construct the "fake news" matrices, one for each outcome variable
386 t0 = time()
387 expectation_vectors = np.empty((J, K)) # Initialize expectation vectors
388 for j in range(J):
389 expectation_vectors[j, :] = np.dot(LR_outcomes[j], outcome_grids[j])
390 FN = make_fake_news_matrices(
391 T_max,
392 J,
393 dY_news_array,
394 D_dstn_array,
395 LR_trans.T,
396 expectation_vectors.copy(),
397 )
398 t1 = time()
399 if verbose:
400 print(
401 "Constructing the fake news matrices took {:.3f}".format(t1 - t0)
402 + " seconds."
403 )
405 # Construct the SSJ matrices, one for each outcome variable
406 t0 = time()
407 SSJ_array = calc_ssj_from_fake_news_matrices(T_max, J, FN, eps)
408 SSJ = [SSJ_array[j, :, :] for j in range(J)] # unpack into a list of arrays
409 t1 = time()
410 if verbose:
411 print(
412 "Constructing the sequence space Jacobians took {:.3f}".format(t1 - t0)
413 + " seconds."
414 )
416 if no_list:
417 return SSJ[0]
418 else:
419 return SSJ
420 finally:
421 _restore_agent(agent, LR_soln, simulator_backup)
424def calc_shock_response_manually(
425 agent,
426 shock,
427 outcomes,
428 grids,
429 s=0,
430 eps=1e-4,
431 T_max=300,
432 norm=None,
433 solved=False,
434 construct=[],
435 offset=False,
436 verbose=False,
437):
438 """
439 Compute an AgentType instance's timepath of outcome responses to learning at
440 t=0 that the named shock variable will be perturbed at t=s. This is equivalent
441 to calculating only the s-th column of the SSJs *manually*, rather than using
442 the fake news algorithm. This function can be used to verify and/or debug the
443 output of the fake news SSJ algorithm.
445 Important: Mortality (or death and replacement generally) should be turned
446 off in the model (via parameter values) for this to work properly. Or does it?
448 Parameters
449 ----------
450 agent : AgentType
451 Agent for which the response(s) should be calculated. Must have T_cycle=1
452 and cycles=0, or the function will throw an error. Must have a model
453 file defined or this won't work at all.
454 shock : str
455 Name of the variable that the response will be computed with respect to.
456 It does not need to be a "shock" in a modeling sense, but it must be a
457 single-valued parameter (possibly a singleton list) that can be changed.
458 outcomes : str or [str]
459 Names of outcome variables of interest; an SSJ matrix will be constructed
460 for each variable named here. If a single string is passed, the output
461 will be a single np.array. If a list of strings are passed, the output
462 will be a list of dYdX vectors in the order specified here.
463 grids : dict
464 Dictionary of dictionaries with discretizing grid information. The grids
465 should include all arrival variables other than those that are normalized
466 out. They should also include all variables named in outcomes, except
467 outcomes that are continuation variables that remap to arrival variables.
468 Grid specification must include number of nodes N, should also include
469 min and max if the variable is continuous.
470 s : int
471 Period in which the shock variable is perturbed, relative to current t=0.
472 The default is 0.
473 eps : float
474 Amount by which to perturb the shock variable. The default is 1e-4.
475 T_max : int
476 The length of the simulation for this exercise. The default is 300.
477 norm : str or None
478 Name of the model variable to normalize by for Harmenberg aggregation,
479 if any. For many HARK models, this should be 'PermShk', which enables
480 the grid over permanent income to be omitted as an explicit state.
481 solved : bool
482 Whether the agent's model has already been solved. If False (default),
483 it will be solved as the very first step.
484 construct : [str]
485 List of constructed objects that will be changed by perturbing shock.
486 These should all share an "offset status" (True or False). Default is [].
487 offset : bool
488 Whether the shock variable is "offset in time" for the solver, with a
489 default of False. This should be set to True if the named shock variable
490 (or the constructed model input that it affects) is indexed by t+1 from
491 the perspective of the solver. For example, the period t solver for the
492 ConsIndShock model takes in risk free interest factor Rfree as an argument,
493 but it represents the value of R that will occur at the start of t+1.
494 verbose : bool
495 Whether to display timing/progress to screen. The default is False.
497 Returns
498 -------
499 dYdX : np.array or [np.array]
500 One or more vectors of length T_max.
501 """
502 setup = _prepare_ssj_computation(agent, outcomes, grids, norm, solved, verbose)
503 outcomes = setup["outcomes"]
504 no_list = setup["no_list"]
505 simulator_backup = setup["simulator_backup"]
506 LR_soln = setup["LR_soln"]
507 X = setup["X"]
508 LR_outcomes = setup["LR_outcomes"]
509 outcome_grids = setup["outcome_grids"]
510 SS_dstn = setup["SS_dstn"]
512 SS_outcomes = [np.dot(mat.T, SS_dstn) for mat in LR_outcomes]
513 SS_avgs = [np.dot(ss, grid) for ss, grid in zip(SS_outcomes, outcome_grids)]
515 try:
516 # Make a temporary agent to construct the perturbed constructed objects
517 t0 = time()
518 temp_agent = deepcopy(agent)
519 base_shock_value, shock_is_list = _perturb_shock(agent, shock)
520 if shock_is_list:
521 temp_value = [base_shock_value + eps]
522 else:
523 temp_value = base_shock_value + eps
524 temp_dict = {shock: temp_value}
525 temp_agent.assign_parameters(**temp_dict)
526 if len(construct) > 0:
527 temp_agent.update()
528 for var in construct:
529 temp_dict[var] = getattr(temp_agent, var)
531 # Build the finite horizon version of this agent
532 FH_agent = deepcopy(agent)
533 FH_agent.del_param("solution")
534 FH_agent.del_param("_simulator")
535 FH_agent.del_from_time_vary("solution")
536 FH_agent.del_from_time_inv(shock)
537 FH_agent.add_to_time_vary(shock)
538 FH_agent.del_from_time_inv(*construct)
539 FH_agent.add_to_time_vary(*construct)
540 finite_dict = {"T_cycle": T_max, "cycles": 1}
541 for var in FH_agent.time_vary:
542 if var in construct:
543 sequence = [deepcopy(getattr(agent, var)[0]) for t in range(T_max)]
544 sequence[s] = deepcopy(getattr(temp_agent, var)[0])
545 else:
546 sequence = T_max * [deepcopy(getattr(agent, var)[0])]
547 finite_dict[var] = sequence
548 shock_seq = T_max * [base_shock_value]
549 shock_seq[s] = base_shock_value + eps
550 finite_dict[shock] = shock_seq
551 FH_agent.assign_parameters(**finite_dict)
552 del temp_agent
553 t1 = time()
554 if verbose:
555 print(
556 "Building the finite horizon agent took {:.3f}".format(t1 - t0)
557 + " seconds."
558 )
560 # Solve the finite horizon agent
561 t0 = time()
562 FH_agent.solve(from_solution=LR_soln)
563 t1 = time()
564 if verbose:
565 print(
566 "Solving the "
567 + str(T_max)
568 + " period problem took {:.3f}".format(t1 - t0)
569 + " seconds."
570 )
572 # Build transition matrices for the finite horizon problem
573 t0 = time()
574 FH_agent.initialize_sym()
575 FH_agent._simulator.make_transition_matrices(
576 grids, norm=norm, fake_news_timing=True
577 )
578 t1 = time()
579 if verbose:
580 print(
581 "Constructing transition matrices took {:.3f}".format(t1 - t0)
582 + " seconds."
583 )
585 # Use grid simulation to find the timepath of requested variables, and compute
586 # the derivative with respect to baseline outcomes
587 t0 = time()
588 FH_agent._simulator.simulate_cohort_by_grids(outcomes, from_dstn=SS_dstn)
589 dYdX = []
590 for j, var in enumerate(outcomes):
591 diff_path = (FH_agent._simulator.history_avg[var] - SS_avgs[j]) / eps
592 if offset:
593 dYdX.append(diff_path[1:])
594 else:
595 dYdX.append(diff_path[:-1])
596 t1 = time()
597 if verbose:
598 print(
599 "Calculating impulse responses by grid simulation took {:.3f}".format(
600 t1 - t0
601 )
602 + " seconds."
603 )
605 del FH_agent
606 if no_list:
607 return dYdX[0]
608 else:
609 return dYdX
610 finally:
611 _restore_agent(agent, LR_soln, simulator_backup)
614@njit
615def calc_derivs_of_state_dstns(T, J, trans_by_t, trans_LR, SS_dstn): # pragma: no cover
616 """
617 Numba-compatible helper function to calculate the derivative of the state
618 distribution by period.
620 Parameters
621 ----------
622 T : int
623 Maximum time horizon for the fake news algorithm.
624 J : int
625 Number of outcomes of interest.
626 trans_by_t : np.array
627 Array of shape (T,K,K) representing the transition matrix in each period.
628 trans_LR : np.array
629 Array of shape (K,K) representing the long run transition matrix.
630 SS_dstn : np.array
631 Array of size K representing the long run steady state distribution.
633 Returns
634 -------
635 D_dstn_news : np.array
636 Array of shape (T,K) representing dD_1^s from the SSJ paper, where K
637 is the number of arrival state space nodes.
639 """
640 K = SS_dstn.size
641 D_dstn_news = np.empty((T, K)) # this is dD_1^s in the SSJ paper (equation 24)
642 for t in range(T - 1, -1, -1):
643 D_dstn_news[T - t - 1, :] = np.dot((trans_by_t[t, :, :] - trans_LR).T, SS_dstn)
644 return D_dstn_news
647@njit
648def calc_derivs_of_policy_funcs(T, Y_by_t, Y_LR, Y_grid, SS_dstn): # pragma: no cover
649 """
650 Numba-compatible helper function to calculate the derivative of an outcome
651 function in each period.
653 Parameters
654 ----------
655 T : int
656 Maximum time horizon for the fake news algorithm.
657 Y_by_t : np.array
658 Array of shape (T,K,N) with the stochastic outcome, mapping from K arrival
659 state space nodes to N outcome space nodes, for each of the T periods.
660 Y_LR : np.array
661 Array of shape (K,N) representing the stochastic outcome in the long run.
662 Y_grid : np.array
663 Array of size N representing outcome space gridpoints.
664 SS_dstn : np.array
665 Array of size K representing the long run steady state distribution.
667 Returns
668 -------
669 dY_news : np.array
670 Array of size T representing the change in average outcome in each period
671 when the shock arrives unexpectedly in that period.
672 """
673 dY_news = np.empty(T) # this is dY_0^s in the SSJ paper (equation 24)
674 for t in range(T - 1, -1, -1):
675 temp = (Y_by_t[t, :, :] - Y_LR).T
676 dY_news[T - t - 1] = np.dot(np.dot(temp, SS_dstn), Y_grid)
677 return dY_news
680@njit
681def make_fake_news_matrices(T, J, dY, D_dstn, trans_LR, E): # pragma: no cover
682 """
683 Numba-compatible function to calculate the fake news array from first order
684 perturbation information.
686 Parameters
687 ----------
688 T : int
689 Maximum time horizon for the fake news algorithm.
690 J : int
691 Number of outcomes of interest.
692 dY : int
693 Array shape (T,J) representing dY_0 from the SSJ paper.
694 D_dstn : np.array
695 Array of shape (T,K) representing dD_1^s from the SSJ paper, where K
696 is the number of arrival state space nodes.
697 trans_LR : np.array
698 Array of shape (K,K) representing the transpose of the long run transition matrix.
699 E : np.array
700 Initial expectation vectors combined into a single array of shape (J,K).
702 Returns
703 -------
704 FN : np.array
705 Fake news array of shape (J,T,T).
706 """
707 FN = np.empty((J, T, T))
708 FN[:, 0, :] = dY.T # Fill in row zero
709 for t in range(1, T): # Loop over other rows
710 for s in range(T):
711 FN[:, t, s] = np.dot(E, D_dstn[s, :])
712 E = np.dot(E, trans_LR)
713 return FN
716@njit
717def calc_ssj_from_fake_news_matrices(T, J, FN, dx): # pragma: no cover
718 """
719 Numba-compatible function to calculate the HA-SSJ from fake news matrices.
721 Parameters
722 ----------
723 T : int
724 Maximum time horizon for the fake news algorithm.
725 J : int
726 Number of outcomes of interest.
727 FN : np.array
728 Fake news array of shape (J,T,T).
729 dx : float
730 Size of the perturbation of the shock variables (epsilon).
732 Returns
733 -------
734 SSJ : np.array
735 HA-SSJ array of shape (J,T,T).
736 """
737 SSJ = np.empty((J, T, T))
738 SSJ[:, 0, :] = FN[:, 0, :] # Fill in row zero
739 SSJ[:, :, 0] = FN[:, :, 0] # Fill in column zero
740 for t in range(1, T): # Loop over other rows
741 for s in range(1, T): # Loop over other columns
742 SSJ[:, t, s] = SSJ[:, t - 1, s - 1] + FN[:, t, s]
743 SSJ *= dx**-1.0 # Scale by dx
744 return SSJ