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 MarkovTransition, Regime, categorical
from lcm.typing import FloatND, ScalarInt

Detecting C-based callable <method-wrapper '__call__' of jaxlib._jax.PjitFunction object at 0x7c9cee452d50> isomorphism...

Regime Transitions¶

The regime transition function determines which regime an agent enters in the next period. There are two kinds:

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 0x7c9d2117a6c0>

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 types end up in the same format during the solve step.

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 (`None`)¶

States that don’t change over time. An identity transition is auto-generated:

state_transitions={
    "education": None,  # fixed — identity auto-generated
}

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),
}

Target-Regime-Dependent Transitions¶

When a state’s transition depends on which regime the agent enters next period, use a dict keyed by target regime name. Every reachable target must be listed:

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.

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, 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.