[WIP] Add life-cycle HA-SSJ capability based on Mateo's paper#1718
[WIP] Add life-cycle HA-SSJ capability based on Mateo's paper#1718
Conversation
All of this code has been sitting in a temp folder on my desktop since about June 2025. I began a HARK implementation of Mateo's lifecycle SSJ method around then, but cut it out so that it didn't hold up the release of 0.16.0. I'm now going to return and finish it up.
There was a problem hiding this comment.
Pull request overview
Adds an initial (WIP) implementation of lifecycle (“flat demographics”) HA sequence-space Jacobians (SSJs) to complement the existing infinite-horizon make_basic_SSJ_matrices utilities.
Changes:
- Introduces
make_flat_LC_SSJ_matricesas a new lifecycle SSJ constructor (currently incomplete). - Adds a Numba-jitted helper
update_FN_matsintended to update lifecycle “fake news” matrices.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| # Construct the (aggregate) SSJ object | ||
|
|
||
| # Structure and return outputs | ||
| return None |
There was a problem hiding this comment.
T_max, PopGroFac, and ProdGroFac are currently unused, and the function returns None despite the docstring promising SSJ matrices. If this is intentionally incomplete/WIP, consider raising NotImplementedError (or removing the public function) to avoid silent misuse.
| return None | |
| raise NotImplementedError( | |
| "make_basic_SSJ_matrices is not yet implemented: it currently does not " | |
| "construct or return SSJ matrices. Please implement or remove this " | |
| "function before using it." | |
| ) |
| T_max : int | ||
| Size of the SSJ matrices: the maximum number of periods to consider. | ||
| The default is 300. |
There was a problem hiding this comment.
The docstring says T_max default is 300, but the signature sets T_max=100. Please reconcile the documented default with the actual default value to avoid confusing callers.
| for a in range(k + 1): | ||
| D_dstn_news = ( | ||
| np.dot(shocked_trans[a], SS_dstn[a]) - SS_dstn[a + 1] | ||
| ).flatten() | ||
| update_FN_mats(fake_news_array, E_vecs[a + 1], D_dstn_news, T_age - 1, a, k) |
There was a problem hiding this comment.
update_FN_mats is @njit, so the evecs argument must be a NumPy array with a Numba-compatible dtype/shape. In the current caller, evecs comes from E_vecs[...], which is built as nested Python lists/dicts above; this will fail compilation when executed. Consider constructing E_vecs as a single NumPy array (e.g., shape (T, J, K)), and pass a slice of that array here.
| # Make sure the shock variable is age-varying | ||
| if shock in agent.time_inv: | ||
| temp = getattr(agent, shock) | ||
| setattr(agent, shock, agent.T_cycle * [temp]) | ||
| agent.del_from_time_inv(shock) | ||
| agent.add_to_time_vary(shock) | ||
|
|
There was a problem hiding this comment.
This function converts a time-invariant shock parameter into a time-varying list (and updates time_inv/time_vary), but there is no corresponding restoration step before returning. That permanently changes the passed-in agent’s parameter structure for callers. Consider snapshotting and restoring the original attribute value and time-inv/time-vary membership on exit (similar to make_basic_SSJ_matrices).
| for j in len(outcomes): | ||
| name = outcomes[j] | ||
| SS_outcomes[name] = [ | ||
| np.dot(LR_outcomes[j][t], outcome_grids[j][t]) for t in range(T_age) | ||
| ] |
There was a problem hiding this comment.
for j in len(outcomes): will raise TypeError: 'int' object is not iterable. Use for j in range(len(outcomes)): (or enumerate(outcomes)) when building SS_outcomes.
| for k in reversed(range(T_age - 1)): | ||
| # Perturb the shock variable at age k | ||
| shock_val_orig = getattr(agent, shock)[k] | ||
| shock_val_new = shock_val_orig + eps | ||
| getattr(agent, shock)[k] = shock_val_new | ||
|
|
||
| # Solve the model starting from age k | ||
| if construct: | ||
| agent.update() | ||
| agent.solve(from_solution=LR_soln[k + 1], from_t=k + 1) | ||
|
|
There was a problem hiding this comment.
Inside the age-perturbation loop, the shocked parameter value is never restored (and neither is agent.solution). This means later iterations will compound earlier perturbations and produce incorrect Jacobians; also leaves the passed-in agent mutated after this function finishes.
| # Update the t=0 row of the fake news matrices | ||
| for j in range(J): | ||
| for a in range(k + 1): | ||
| fake_news_array[j, a, 0, k - a] += np.sum( | ||
| (shocked_outcomes[j][a] - LR_outcomes[j][a]) * SS_dstn[a] | ||
| ) |
There was a problem hiding this comment.
The t=0 fake-news update is aggregating (shocked_outcomes - LR_outcomes) * SS_dstn but does not apply the outcome grid. Since X.periods[a].matrices[var] is an outcome-mapping matrix, this computes a sum of weighted matrix entries rather than the change in the average outcome. The aggregation should incorporate outcome_grids similarly to calc_derivs_of_policy_funcs.
| raise ValueError("This function is only compatible with life-cycle models!") | ||
| if not isinstance(outcomes, list): | ||
| outcomes = [outcomes] | ||
| no_list = True |
There was a problem hiding this comment.
Variable no_list is not used.
| outcomes = [outcomes] | ||
| no_list = True | ||
| else: | ||
| no_list = False |
There was a problem hiding this comment.
Variable no_list is not used.
|
|
||
| # Store the simulator if it exists | ||
| if hasattr(agent, "_simulator"): | ||
| simulator_backup = agent._simulator |
There was a problem hiding this comment.
Variable simulator_backup is not used.
|
This is WIP.
…On Wed, Feb 25, 2026 at 8:48 PM Alan Lujan ***@***.***> wrote:
***@***.**** approved this pull request.
—
Reply to this email directly, view it on GitHub
<#1718 (review)>,
or unsubscribe
<https://github.com/notifications/unsubscribe-auth/ADKRAFKBNSB6LSEEZUYARZT4NZGGTAVCNFSM6AAAAACUHZXLS6VHI2DSMVQWIX3LMV43YUDVNRWFEZLROVSXG5CSMV3GSZLXHMZTQNJXHEZTAMZVGQ>
.
You are receiving this because you authored the thread.Message ID:
***@***.***>
|
I began a HARK implementation of Mateo's lifecycle SSJ method around June 2025, but cut it out so that it didn't hold up the release of 0.16.0. It is incomplete and non-functional right now, but I'm going to return and finish it up.
When complete, it will work just like the infinite horizon
make_basic_SSJmethod, and I might unify theAgentTypeinterface for them (i.e. the same top-level method is used whether it's infinite horizon or lifecycle; which function gets called by the method happens automatically under the hood).The code in the first commit represents what I had done back then; it's been sitting in a temp file on my desktop since July 1, 2025.