Sequence-Space Jacobian (SSJ) Toolkit

The best way to learn SSJ is by working through the provided Jupyter notebooks. The notebooks cover a progression of complexity:

1. RBC: Learn to represent macro models as collections of blocks (DAG), write SimpleBlocks and CombinedBlocks, and compute linearized and non-linear (perfect-foresight) impulse responses.
2. Krusell-Smith: Learn to write HetBlocks for heterogeneous agents, construct general-equilibrium Jacobians manually, and compute model log-likelihood given time-series data.
3. One-asset HANK: Learn to adapt off-the-shelf HetBlock to macro environments using helper functions and advanced calibration.
4. Two-asset HANK: Learn to write SolvedBlocks for implicit aggregate equilibrium conditions, re-use saved Jacobians, and fine-tune block method options.
5. Labor search: Learn to handle multiple exogenous states and shocks to transition matrices.

You can download all notebooks as a zip file from the repository's notebooks/ directory.

You can adapt off-the-shelf Heterogeneous Agent (HA) blocks (like hh_labor or hh_sim) to specific macro models by adding custom input and output logic. This avoids rewriting the core backward iteration (Bellman equation) solver.

• hetinputs: Functions evaluated before the core HetBlock backward iteration. They can handle acyclic dependencies (one hetinput using the output of another) and can take scalar-valued, time-varying inputs (like aggregate dividends or taxes) to compute micro-level variables.
• hetoutputs: Functions evaluated after the backward iteration has converged. They can take both the multidimensional outputs of the backward iteration and the outputs of the hetinputs as arguments (e.g., to compute aggregate effective labor supply from individual labor supply).

Use HetBlock.add_hetinputs([...]) and HetBlock.add_hetoutputs([...]) to attach these functions.

# Example: Adding hetinputs and hetoutputs to a household block
hh1 = hh.add_hetinputs([make_grid, transfers, wages])
hh_ext = hh1.add_hetoutputs([labor_supply])

The Block class is the fundamental building block in sequence-jacobian. It represents a component of a model with defined inputs and outputs. Users typically subclass Block to implement specific economic equations or model segments.

Key capabilities of a Block include:

• Evaluating steady states (steady_state)
• Calculating non-linear impulse responses (impulse_nonlinear)
• Calculating linear impulse responses (impulse_linear)
• Computing Jacobians (jacobian and partial_jacobians)
• Solving for general equilibrium steady states (solve_steady_state)
• Solving for general equilibrium impulse responses (solve_impulse_nonlinear, solve_impulse_linear, and solve_jacobian)

Note that Block is an abstract base class. For a block to be functional, you must implement the underlying _steady_state, _impulse_nonlinear, _impulse_linear, and _jacobian methods (prefixed with underscores).

A CombinedBlock (or Model) acts as a container for a sequence of blocks. It inherits from Block, Parent, and DAG, allowing it to:

• Evaluate partial or general equilibrium steady states via steady_state().
• Compute impulse responses (nonlinear or linear) via impulse_nonlinear() or impulse_linear().
• Calculate Jacobians along the dependency graph via jacobian() or partial_jacobians().

Important: Do not manually change the attributes of a CombinedBlock via assignment after instantiation. Automated features and topological sorting are inferred during initialization and may not update correctly if attributes are reassigned.

from sequence_jacobian.blocks.combined_block import create_model

# 1. Define blocks
# 2. Combine them into a model
model = create_model(blocks, name="macro_model")

# 3. Use model methods (inherited from Block/DAG/Parent)
ss = model.steady_state(calibration, dissolve=some_dict)
js = model.jacobian(ss, inputs, outputs, T, Js, options)

If you need to separate the calculation of interpolation weights from the application of those weights to a different dataset, use the combination of interpolate_coord and apply_coord. This is useful when you want to find the interpolation representation for a set of query points once, and then apply that same representation to multiple different $y$ arrays.

1. Get coordinates with interpolate_coord

Calculates the indices and weights for query points xq against data points x.

• Input: x (ascending data points), xq (ascending query points).
• Output: xqi (indices of lower bracketing gridpoints), xqpi (weights on lower bracketing gridpoints).
• Relationship: xq = xqpi * x[xqi] + (1-xqpi) * x[xqi+1]

2. Apply coordinates with apply_coord

Uses the previously calculated coordinates to find interpolated values for a new data array y.

• Input: xqi (indices), xqpi (weights), y (data points).
• Output: yq (interpolated points).
• Relationship: yq = xqpi * y[xqi] + (1-xqpi) * y[xqi+1]

from sequence_jacobian.utilities.interpolate import interpolate_coord, apply_coord
import numpy as np

x = np.array([1.0, 2.0, 3.0])
xq = np.array([1.5, 2.5])
y1 = np.array([10.0, 20.0, 30.0])
y2 = np.array([100.0, 200.0, 300.0])

# Step 1: Get representation
xqi, xqpi = interpolate_coord(x, xq)

# Step 2: Apply to different y arrays
yq1 = apply_coord(xqi, xqpi, y1)
# yq2 = apply_coord(xqi, xqpi, y2)

A key assumption of the SSJ method is that blocks are connected via scalar sequences. To handle multidimensional inputs (like grids) or outputs (like specific policy functions) within a HetBlock, use add_hetinputs and add_hetoutputs.

• hetinputs: Map scalars from outside the block into the multidimensional inputs the block requires (e.g., creating an asset grid from a scalar parameter).
• hetoutputs: Map multidimensional outputs of the block into new multidimensional outputs of interest.

A HetBlock evaluation follows this lifecycle:

1. Evaluate hetinputs to obtain multidimensional inputs.
2. Evaluate the core block (backward iteration) to obtain policy functions.
3. Evaluate hetoutputs to obtain desired outputs.

def make_grid(rho_e, sd_e, nE, amin, amax, nA):
    e_grid, _, Pi = grids.markov_rouwenhorst(rho=rho_e, sigma=sd_e, N=nE)
    a_grid = grids.agrid(amin=amin, amax=amax, n=nA)
    return e_grid, Pi, a_grid

# Attach the grid creation as a hetinput
household_ext = household.add_hetinputs([make_grid])

A Solved Block is an extension of a simple block designed to handle implicit relationships that require numerical evaluation rather than analytical mapping. While simple blocks map aggregate inputs to aggregate outputs analytically, solved blocks are used when the relationship between variables is defined by an equation that must be solved (e.g., a Phillips curve or an arbitrage condition).

When the SSJ toolkit encounters a solved block, it computes its Jacobian via the implicit function theorem. This allows you to embed complex, numerically-solved miniature models within a larger Directed Acyclic Graph (DAG) without changing the rest of your model's structure.

@solved(unknowns={'pi': (-0.1, 0.1)}, targets=['nkpc'], solver="brentq")
def pricing_solved(pi, mc, r, Y, kappap, mup):
    nkpc = kappap * (mc - 1/mup) + Y(+1) / Y * (1 + pi(+1)).apply(np.log) / \
           (1 + r(+1)) - (1 + pi).apply(np.log)
    return nkpc

To represent heterogeneous agents at both micro and macro levels, you must define a backward step function that represents a single step of an iterative solution method (like value function iteration) for the Bellman equation.

Once defined, use the @het decorator to transform this function into a HetBlock. You must specify:

• exogenous: The name of the transition matrix for discrete exogenous state(s).
• policy: The name of the policy corresponding to continuous endogenous state(s).
• backward: The name of the backward variable(s) used for iteration (e.g., the derivative of the value function).
• backward_init: A function to initialize the guess for the backward variable.

HA blocks are unique because they produce both macro outputs (aggregates) and micro outputs (distributions, policy functions, or internals).

@het(exogenous='Pi', policy='a', backward='Va', backward_init=household_init)
def household(Va_p, a_grid, e_grid, r, w, beta, eis):
    # ... implementation of a single backward iteration step ...
    return Va, a, c

Starting in version 1.0.0, HetBlocks support an arbitrary number of discrete exogenous state variables. This allows for modeling scenarios where transition probabilities (e.g., job finding or separation rates) vary over time.

To implement this, follow these requirements:

1. State Space Dimensionality: Ensure your NumPy expressions support the additional dimensions. For example, if adding employment status $s$, your state space becomes 3-dimensional (s, e, a) instead of 2-dimensional (e, a).
2. Chronological Order: Provide the names of your exogenous Markov matrices in the order they occur chronologically.
3. Hetinput Function: Write a hetinput function that maps your time-varying parameters (like job-finding rates) into the required Markov matrices.

Once these steps are completed, you can compute Jacobians and impulse responses with respect to these time-varying exogenous parameters.

# Example of defining a time-varying Markov matrix via hetinputs
def search_frictions(f, s):
    # f is the job-finding rate, s is the separation rate
    Pi_s = np.vstack(([1 - s, s], [f, 1 - f]))
    return Pi_s

# Applying the decorator with exogenous matrices
@het(exogenous=['Pi_s', 'Pi_e'], policy='a', backward='Va', backward_init=household_init)
def household(Va_p, a_grid, y, r, beta, eis):
    # ... implementation ...
    return Va, a, c