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¶

The flagship use of the grammar is the carried state: Phased(solve=callable, simulate=Grid) in states gives one quantity two roles:

solve: a derived function — the quantity is imputed from other states and never becomes a grid dimension, so the solve 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.

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.

The Phase Grammar

Carried States¶

The Params Template Unions Both Phases¶

Both Phases in Action¶

See Also¶