Transitions

Transitions define how agents move between regimes and how state variables evolve over time. Both are specified on the Regime:

Regime transitions (transition field) — which regime next period?
State transitions (state_transitions field) — how do states evolve?

import jax.numpy as jnp

from lcm import (
    AgeGrid,
    DiscreteGrid,
    LinSpacedGrid,
    MarkovTransition,
    Model,
    Regime,
    categorical,
    fixed_transition,
)
from lcm.typing import DiscreteState, FloatND, ScalarInt

Coarse and Granular Spellings¶

Two words recur throughout this page, describing how precisely a spelling names its target regimes:

A granular spelling names each target regime explicitly — {"retired": law_a, "dead": law_b} — so the engine knows exactly which target gets what.
A coarse spelling is target-blind — a bare law, a bare MarkovTransition — and means “the same thing toward every (reachable) target”.

The canonical form the engine stores is always granular; coarse user spellings are sugar that broadcasts into it. This one rule governs state transitions, regime transitions, and parameters.

Regime Transitions¶

The regime transition (the transition field) determines which regime an agent enters in the next period. It takes one of three forms:

Bare callable (coarse) — deterministic: the function returns the target regime ID. Every regime is reachable.
MarkovTransition (coarse) — stochastic: the function returns a probability vector over all regimes. Every regime is reachable.
Per-target dict (granular) — stochastic: {target_regime: MarkovTransition(prob_func)}. Each cell returns that target’s probability, and the key set declares the regime’s reachable targets — omitted regimes are structurally unreachable.

transition=None marks a terminal regime.

Deterministic¶

The function returns an integer regime ID (from the @categorical RegimeId class). Use this for transitions that depend deterministically on states and actions — for example, mandatory retirement at a certain age, or a retirement decision triggered by a discrete action.

Pass the function directly to Regime(transition=...):

@categorical(ordered=False)
class RegimeId:
    working_life: ScalarInt
    retirement: ScalarInt


def next_regime(age: float, retirement_age: float) -> ScalarInt:
    return jnp.where(age >= retirement_age, RegimeId.retirement, RegimeId.working_life)


print(next_regime)  # just a plain function — no wrapper needed

<function next_regime at 0x7ce2fd8f6560>

Stochastic¶

The function returns a probability array over all regimes. Use this when the regime transition is uncertain — for example, a mortality risk that determines whether the agent survives to the next period.

Wrap the transition function in MarkovTransition:

@categorical(ordered=False)
class RegimeIdMortality:
    alive: ScalarInt
    dead: ScalarInt


def survival_transition(survival_prob: float) -> FloatND:
    """Return [P(alive), P(dead)]."""
    return jnp.array([survival_prob, 1 - survival_prob])


example_regime = Regime(
    transition=MarkovTransition(survival_transition),
    functions={"utility": lambda: 0.0},
)

Internally, deterministic transitions are converted to one-hot probability arrays, so both forms end up in the same format during the solve step.

Per-Target Dict¶

When not every regime should be reachable, key the transition by target regime name. Each cell is a MarkovTransition-wrapped function returning that target’s probability; across cells the probabilities must sum to one. Regimes not listed are structurally unreachable from this regime:

transition={
    "retired": MarkovTransition(prob_retired),  # P(retired next period)
    "dead": MarkovTransition(prob_dead),  # P(dead next period)
    # "work" omitted — returning to work is structurally impossible
}

Narrowing reachability is not just documentation: state values are handed across the regime boundary only toward reachable targets, so granular keys shrink the carries at the boundary and make structural intent explicit. Cell parameters nest under the target in the params template — prob_retired’s parameters live at template[<regime>]["retired"]["next_regime"] (see Per-target parameters below). transition={} is rejected — terminality is spelled transition=None.

State Transitions¶

The state_transitions dict on a Regime defines how each state variable evolves over time. Every state in a non-terminal regime must have an entry — except continuous stochastic processes, which carry their own transitions (see below).

Deterministic Transitions¶

Pass a callable that returns the next-period value:

state_transitions={
    "wealth": next_wealth,  # callable: (wealth, consumption, ...) -> next_wealth
}

Fixed States (`fixed_transition`)¶

States that don’t change over time get the identity law via fixed_transition(state_name) — the argument must match the dict key:

from lcm import fixed_transition

state_transitions={
    "education": fixed_transition("education"),  # education stays fixed
}

Stochastic Transitions (`MarkovTransition`)¶

For states with stochastic transitions (e.g., a Markov chain over health states), wrap the transition function in MarkovTransition. The function returns a probability vector over grid points:

from lcm import MarkovTransition

state_transitions={
    "health": MarkovTransition(health_transition_probs),
}

Per-Target Dicts¶

When a state’s transition depends on which regime the agent enters next period, use a dict keyed by target regime name:

state_transitions={
    "health": {
        "working": MarkovTransition(health_probs_working),
        "retired": MarkovTransition(health_probs_retired),
    },
}

Within a per-target dict, stochasticity must be consistent — all entries must be MarkovTransition or all must be plain callables (fixed_transition entries count as plain callables). Per-target dicts are the building block of multi-regime models; see Cross-regime transitions below for the full semantics.

Terminal Regimes¶

Terminal regimes (transition=None) must have empty state_transitions — there is no next period, so no transitions are needed.

Continuous Stochastic Processes¶

Continuous stochastic processes (e.g., NormalIIDProcess, TauchenAR1Process) derive their own transition probabilities from the discretization scheme, so you do not specify them in state_transitions — placing one there is an error. See Continuous stochastic processes for the available process classes.

Where transitions live — and why¶

State transitions come in two kinds, and the kind decides where you specify the transition:

Functional — a next_<state> function you write yourself. Because you supply the dynamics, the function goes in state_transitions. Deterministic transitions, fixed_transition entries, MarkovTransition-wrapped stochastic transitions, and per-target dicts are all functional.
Parametric — the dynamics follow from a distribution and its parameters. A continuous stochastic process (NormalIIDProcess, TauchenAR1Process, and the other *Process classes) derives both its grid points and its transition matrix from that single spec, so the grid and the transition travel together as one object placed in states.

This is why a continuous stochastic process lives in states and never in state_transitions: its transition is not something you write, it falls out of the discretization. It is also why there is no DiscreteProcess — a DiscreteGrid is just a set of categories, and its Markov transition is an arbitrary function you supply. The two share no common parametrization, so they stay separate: the grid in states, the MarkovTransition-wrapped function in state_transitions.

Cross-Regime Transitions¶

A multi-regime model needs five pieces of boundary semantics, and one worked example covers them all:

Reachability comes from the regime transition. A per-target dict’s key set narrows which regimes are reachable; the coarse forms (bare callable, MarkovTransition) make every regime reachable. Narrowing buys smaller carries at the boundary and makes structural intent explicit.
Coverage rule. Every reachable target that carries the state needs a law of motion for it. A bare law broadcasts to all of them; a per-target dict lists them explicitly.
Handing a value into a state the source regime does not carry. An entering law keyed by the target regime computes the target’s state from the source’s variables.
States the target does not carry are simply dropped — no law needed toward that target.
The same state name can change grids across the boundary (a cross-grid law maps source values into the target’s grid), and fixed_transition works inside per-target cells.

The example: agents work, retire (possibly involuntarily), and die. Retirement is absorbing — there is no way back to work. Health is tracked at three levels while working but only two in retirement; job tenure exists only while working; pension points exist only in retirement.

LAST_WORKING_AGE = 61
FORCED_RETIREMENT_AGE = 62
MAX_AGE = 63


@categorical(ordered=False)
class RegimeIdRetirement:
    work: ScalarInt
    retired: ScalarInt
    dead: ScalarInt


@categorical(ordered=True)
class HealthWork:
    bad: ScalarInt
    fair: ScalarInt
    good: ScalarInt


@categorical(ordered=True)
class HealthRetired:
    bad: ScalarInt
    good: ScalarInt


def utility_work(consumption: float, health: DiscreteState) -> FloatND:
    return jnp.log(consumption) + 0.1 * health


def utility_retired(
    consumption: float, health: DiscreteState, pension_points: float
) -> FloatND:
    return jnp.log(consumption) + 0.05 * health + 0.1 * pension_points


def next_wealth(wealth: float, consumption: float) -> float:
    return (wealth - consumption) * 1.05


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


def coarsen_health(health: DiscreteState) -> ScalarInt:
    """Map three working-life health levels onto the two retirement levels."""
    return jnp.where(health == HealthWork.good, HealthRetired.good, HealthRetired.bad)


def accrue_pension_points(job_tenure: float) -> float:
    """Convert job tenure into pension points on entering retirement."""
    return jnp.clip(0.1 * job_tenure, 0.0, 3.0)


def grow_job_tenure(job_tenure: float) -> float:
    return job_tenure + 1.0

work = Regime(
    # Reachability (1): all three regimes are reachable from work.
    transition={
        "work": MarkovTransition(
            lambda age: jnp.where(age >= LAST_WORKING_AGE, 0.0, 0.6)
        ),
        "retired": MarkovTransition(
            lambda age: jnp.where(age >= LAST_WORKING_AGE, 0.8, 0.3)
        ),
        "dead": MarkovTransition(
            lambda age: jnp.where(age >= LAST_WORKING_AGE, 0.2, 0.1)
        ),
    },
    active=lambda age: age < FORCED_RETIREMENT_AGE,
    states={
        "wealth": LinSpacedGrid(start=1.0, stop=100.0, n_points=10),
        "health": DiscreteGrid(HealthWork),
        "job_tenure": LinSpacedGrid(start=0.0, stop=40.0, n_points=5),
    },
    state_transitions={
        # Coverage (2): a bare law broadcasts to every reachable target that
        # carries the state — here `work` and `retired` (dead carries nothing).
        "wealth": next_wealth,
        # Cross-grid + fixed_transition in cells (5): identity toward work,
        # a 3-to-2 category map toward retirement.
        "health": {
            "work": fixed_transition("health"),
            "retired": coarsen_health,
        },
        # Dropping (4): only work carries job_tenure — no law toward `retired`
        # or `dead` is needed; the state is dropped at the boundary.
        "job_tenure": {"work": grow_job_tenure},
        # Entering law (3): work does not carry pension_points; the law keyed
        # by `retired` hands a value into the target's state.
        "pension_points": {"retired": accrue_pension_points},
    },
    actions={"consumption": LinSpacedGrid(start=1.0, stop=10.0, n_points=5)},
    constraints={"consumption_feasible": consumption_feasible},
    functions={"utility": utility_work},
)

retired = Regime(
    # Reachability (1): `work` is omitted — returning to work is structurally
    # impossible, so no laws of motion toward work are ever required here.
    transition={
        "retired": MarkovTransition(
            lambda age: jnp.where(age >= FORCED_RETIREMENT_AGE, 0.0, 0.9)
        ),
        "dead": MarkovTransition(
            lambda age: jnp.where(age >= FORCED_RETIREMENT_AGE, 1.0, 0.1)
        ),
    },
    active=lambda age: age < MAX_AGE,
    states={
        "wealth": LinSpacedGrid(start=1.0, stop=100.0, n_points=10),
        "health": DiscreteGrid(HealthRetired),
        "pension_points": LinSpacedGrid(start=0.0, stop=3.0, n_points=4),
    },
    state_transitions={
        "wealth": next_wealth,
        "health": fixed_transition("health"),
        "pension_points": fixed_transition("pension_points"),
    },
    actions={"consumption": LinSpacedGrid(start=1.0, stop=10.0, n_points=5)},
    constraints={"consumption_feasible": consumption_feasible},
    functions={"utility": utility_retired},
)

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

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

Solving and simulating shows the boundary semantics at work: health is coarsened from three categories to two on retirement, pension points appear there computed from job tenure, and job tenure disappears.

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([30.0, 80.0]),
        "health": jnp.array([HealthWork.fair, HealthWork.good]),
        "job_tenure": jnp.array([10.0, 30.0]),
        "regime_id": jnp.array([RegimeIdRetirement.work] * 2),
    },
    log_level="warning",
)
df = result.to_dataframe()
df[
    [
        "period",
        "subject_id",
        "regime_name",
        "wealth",
        "health",
        "job_tenure",
        "pension_points",
    ]
]

Reading the output:

health takes the working-life labels (fair, good) while in work and the coarsened retirement labels (bad, good) after the cross-grid law fires at the boundary.
pension_points is empty during work (the regime does not carry it) and is seeded by the entering law accrue_pension_points on retirement.
job_tenure grows while working and is dropped on retirement — no law toward retired exists, and none is required.

Per-Target Parameters¶

The params template mirrors how you spelled each transition. A bare law keeps a single entry under its own name; a per-target dict (and each cell of a per-target regime transition) nests its parameters under the target regime’s name — the target is a genuine level of the params tree:

template = model.get_params_template()
# A bare law's params:           template["work"]["next_wealth"]
# A per-target law's params:     template["work"]["retired"]["next_wealth"]
# A regime-transition cell's:    template["work"]["retired"]["next_regime"]

Values can be supplied at four levels, most to least specific:

Target level — {"work": {"retired": {"next_wealth": {"exit_tax": 0.1}}}}: a value for that target’s law only.
Function level — {"work": {"next_wealth": {"exit_tax": 0.1}}}: one value broadcasts over all targets of the law.
Regime level — {"work": {"exit_tax": 0.1}}: propagates to every function in the regime needing exit_tax.
Model level — {"exit_tax": 0.1}: propagates to all regimes.

Each parameter must be supplied at exactly one level — multi-level specifications are rejected as ambiguous. A coarse regime transition (bare callable or MarkovTransition) is evaluated once and shared across targets, so it takes no per-target parameters — its params keep the single next_regime entry.