The Phase Grammar

Every pylcm model lives in two phases: solve (backward induction over the state-action grid) and simulate (forward sampling of subjects). Most regime slots mean the same thing in both phases — but some quantities genuinely differ between them, and the phase grammar lets you say so in one place.

The grammar is one idea: phase is a broadcast dimension of the regime specification.

A bare value broadcasts to both phases — write a function or a grid once and both phases use it.
Phased(solve=..., simulate=...) specifies each phase explicitly.

Phased is accepted where a per-phase variant makes sense:

functions — per-phase implementations.
state_transitions — per-phase laws of motion.
transition — per-phase regime transitions (matching forms; for per-target dicts, identical key sets).
states — only the combination Phased(solve=callable, simulate=Grid), the carried state described below.

constraints, actions, active, and derived_categoricals are phase-invariant — solve and simulate must agree on what is feasible and what can be chosen, otherwise simulated agents would face a different problem than the one their policy was computed for. Phased is rejected there with an explanation, is outermost-only (never inside a per-target dict), and never nests.

Carried States¶

A carried state is a state that the policy does not condition on, but whose true value the simulation must track — e.g. actual wealth in private pensions when the policy was solved on a value imputed from public pensions etc. in order to keep the state space manageable. It is spelled Phased(solve=callable, simulate=Grid) in states, giving one quantity two roles:

solve: a derived function — the quantity is computed from other states by the callable and never becomes a grid dimension, so the value-function grid does not grow;
simulate: a genuine state — seeded from the initial conditions and evolved each period by its ordinary state_transitions law, with the Grid as its domain.

Decisions are evaluated at the solve-phase imputation — the value the solved policy was computed for; every other simulate consumer reads the carried true value.

This pays off when a state matters for subjects’ histories but is well approximated by a function of other states for decision-making: you keep the per-subject dynamics in simulation without paying for another axis of the value function.

The example below tracks pension wealth. During solve it is imputed from average earnings (aime); during simulation it is a real state that compounds at a fixed rate.

import pprint

import jax.numpy as jnp

from lcm import AgeGrid, LinSpacedGrid, Model, Phased, Regime, categorical
from lcm.typing import FloatND, ScalarInt

RETIREMENT_AGE = 62
MAX_AGE = 63


@categorical(ordered=False)
class RegimeId:
    working: ScalarInt
    dead: ScalarInt


def next_regime(age: float) -> ScalarInt:
    return jnp.where(age >= RETIREMENT_AGE, RegimeId.dead, RegimeId.working)


def impute_pension_wealth(aime: float) -> float:
    """Solve-phase pension wealth: imputed from average earnings."""
    return 0.1 * aime


def evolve_pension_wealth(pension_wealth: float) -> float:
    """Simulate-phase law of motion: compounds at a fixed rate."""
    return 1.03 * pension_wealth


def utility(consumption: float) -> FloatND:
    return jnp.log(consumption)


def next_wealth(wealth: float, consumption: float, pension_wealth: float) -> float:
    return wealth - consumption + pension_wealth


def next_aime(aime: float) -> float:
    return aime


def consumption_feasible(consumption: float, wealth: float) -> bool:
    return consumption <= wealth

working = Regime(
    transition=next_regime,
    active=lambda age: age < MAX_AGE,
    states={
        "wealth": LinSpacedGrid(start=1.0, stop=100.0, n_points=10),
        "aime": LinSpacedGrid(start=1.0, stop=50.0, n_points=5),
        # The carried state: derived during solve, a real state in simulation.
        "pension_wealth": Phased(
            solve=impute_pension_wealth,
            simulate=LinSpacedGrid(start=0.0, stop=20.0, n_points=4),
        ),
    },
    state_transitions={
        "wealth": next_wealth,
        "aime": next_aime,
        # The carried state's law of motion is an ordinary entry.
        "pension_wealth": evolve_pension_wealth,
    },
    actions={"consumption": LinSpacedGrid(start=1.0, stop=10.0, n_points=5)},
    constraints={"consumption_feasible": consumption_feasible},
    functions={"utility": utility},
)

dead = Regime(transition=None, functions={"utility": lambda: 0.0})

model = Model(
    regimes={"working": working, "dead": dead},
    ages=AgeGrid(start=60, stop=63, step="Y"),
    regime_id_class=RegimeId,
)

The Params Template Unions Both Phases¶

The params template reads the regime in user vocabulary, before the phase split. Where a slot differs by phase, the parameters of both variants appear in the template — a parameter needed by only one phase is still a parameter of the model. Below, pension_wealth (the solve-phase imputation) and next_pension_wealth (the simulate-phase law) both surface:

pprint.pprint(model.get_params_template())

{'dead': {'utility': {}},
 'working': {'H': {'discount_factor': 'FloatND'},
             'consumption_feasible': {},
             'next_aime': {},
             'next_pension_wealth': {},
             'next_regime': {},
             'next_wealth': {},
             'pension_wealth': {},
             'utility': {}}}

Both Phases in Action¶

Solving uses the imputation (no pension_wealth axis in the value function); simulation seeds pension wealth from the initial conditions and compounds it at 3% per period:

result = model.simulate(
    params={"discount_factor": 0.95},
    period_to_regime_to_V_arr=None,
    initial_conditions={
        "age": jnp.array([60.0, 60.0]),
        "wealth": jnp.array([20.0, 70.0]),
        "aime": jnp.array([10.0, 40.0]),
        "pension_wealth": jnp.array([2.0, 8.0]),
        "regime_id": jnp.array([RegimeId.working] * 2),
    },
    log_level="warning",
)
result.to_dataframe()[
    ["period", "subject_id", "regime_name", "wealth", "pension_wealth"]
]

The pension_wealth column starts at the seeded values (2.0 and 8.0) and grows by the factor 1.03 each period — the simulate-phase law — while the solve phase never saw a pension-wealth grid axis at all.

Wrong Beliefs: Different Transitions in Solve and Simulate¶

The motivating use case for a Phased law of motion: agents believe a state evolves one way, while the data-generating process differs — the policy is solved under the belief, the simulation evolves the truth. Beliefs and truth are distinct parameters, so name them apart with dags.rename_arguments; everything that keeps one name stays one shared parameter. The params template unions both sides:

rho_belief binds only in the solve variant,
rho_true binds only in the simulate variant,
sigma keeps one name, so both phases share its value.

The rational-expectations counterfactual is one line: set rho_belief equal to rho_true in the params dict.

from dags import rename_arguments


def next_income(income: float, rho: float, sigma: float) -> float:
    return rho * income + sigma


believer = Regime(
    transition=next_regime,
    active=lambda age: age < MAX_AGE,
    states={"income": LinSpacedGrid(start=0.0, stop=10.0, n_points=11)},
    state_transitions={
        "income": Phased(
            solve=rename_arguments(next_income, mapper={"rho": "rho_belief"}),
            simulate=rename_arguments(next_income, mapper={"rho": "rho_true"}),
        ),
    },
    actions={"consumption": LinSpacedGrid(start=0.1, stop=1.0, n_points=3)},
    functions={
        "utility": lambda consumption, income: jnp.log(consumption + 0.1 * income)
    },
)

beliefs_model = Model(
    regimes={"working": believer, "dead": dead},
    ages=AgeGrid(start=60, stop=63, step="Y"),
    regime_id_class=RegimeId,
)
pprint.pprint(beliefs_model.get_params_template()["working"]["next_income"])

{'rho_belief': 'float', 'rho_true': 'float', 'sigma': 'float'}

beliefs_result = beliefs_model.simulate(
    params={
        "discount_factor": 0.95,
        "working": {
            "next_income": {"rho_belief": 0.95, "rho_true": 0.8, "sigma": 0.5},
        },
    },
    period_to_regime_to_V_arr=None,
    initial_conditions={
        "age": jnp.array([60.0, 60.0]),
        "income": jnp.array([2.0, 4.0]),
        "regime_id": jnp.array([RegimeId.working] * 2),
    },
    log_level="warning",
)
beliefs_result.to_dataframe()[["period", "subject_id", "regime_name", "income"]]

The realized income paths follow the true law (0.8 * income + 0.5: 2.0 → 2.1 and 4.0 → 3.7), while the policy was solved under the believed persistence of 0.95 — the simulated agents act on beliefs and live in the truth.

The Phase Grammar

Carried States¶

The Params Template Unions Both Phases¶

Both Phases in Action¶

Wrong Beliefs: Different Transitions in Solve and Simulate¶

See Also¶