Write Economics, Not Glue Code

pylcm lets you write economic equations that look like the math in your papers. You define small functions — one for utility, one for earnings, one for the budget constraint — and pylcm wires them together automatically. This page shows why that matters and how it works.

Named functions are also the seam used by specialized solvers. Keep economically meaningful resources and post-decision quantities as named DAG nodes. When a solver needs a boundary or comparison itself—not only its computed value—use the corresponding structured declaration. See Choosing a solver.

The Model on Paper¶

Consider a consumption-savings model where the agent chooses labor supply $d_t \in \{\text{work}, \text{retire}\}$ and consumption $c_t$ each period. An economist would write it like this:

V_t(w_t) = \max_{d_t,\, c_t} \left\{ u(c_t, d_t) + \beta\, V_{t+1}(w_{t+1}) \right\}

(1)

\begin{align} u(c_t, d_t) &= \ln c_t - \phi \cdot \mathbf{1}[d_t = \text{work}] \\ y_t &= \mathbf{1}[d_t = \text{work}] \cdot \bar{w} \\ w_{t+1} &= (1 + r)(w_t - c_t) + y_t \\ c_t &\leq w_t \end{align}

(2)

Notice the dependency structure implicit in the math:

Utility depends on whether the agent works
Next-period wealth depends on labor income
Labor income depends on whether the agent works

When you write these equations in a paper, the evaluation order is obvious to the reader. In code, someone has to enforce it.

The Dependency Graph¶

Those implicit dependencies form a directed acyclic graph (DAG):

Reading the DAG: an arrow from A to B means “B needs the value of A”. So next_wealth needs labor_income, which needs is_working, which needs the action labor_supply. The question is: who manages this wiring?

The Traditional Approach¶

In a traditional implementation, you manage the call order and pass intermediate results by hand:

import jax.numpy as jnp


def evaluate_period_traditional(
    *, wealth, consumption, labor_supply, wage, interest_rate, disutility_of_work
):
    # Step 1: compute intermediate values — YOU manage the order
    _is_working = labor_supply == 1
    _labor_income = jnp.where(_is_working, wage, 0.0)

    # Step 2: compute utility — YOU pass the intermediate
    _work_disutility = jnp.where(_is_working, disutility_of_work, 0.0)
    _utility = jnp.log(consumption) - _work_disutility

    # Step 3: compute next-period state — YOU pass the intermediate
    _next_wealth = (1 + interest_rate) * (wealth - consumption) + _labor_income

    # Step 4: check constraints — YOU pass the right variables
    _feasible = consumption <= wealth

    return _utility, _next_wealth, _feasible

This works, but has real costs:

You manage the call order. Swap steps 1 and 2 and you get a bug.
You pass intermediate results. Add a new auxiliary variable and you rewire the entire function.
The economics is buried in plumbing. The actual equations are the short lines; most of the code is scaffolding.
It doesn’t scale. With five states, three auxiliary variables, and multiple regimes, this monolith becomes unreadable.

The pylcm Way¶

In pylcm, each equation from the paper becomes its own function. The function’s name is the variable it computes, and its parameters are the variables it needs:

from lcm import (
    AgeGrid,
    DeterministicTransition,
    DiscreteGrid,
    LinSpacedGrid,
    Regime,
    Transition,
)
from lcm.typing import (
    BoolND,
    ContinuousAction,
    ContinuousState,
    DiscreteAction,
    FloatND,
    ScalarInt,
)
from lcm_examples.mortality import LaborSupply, RegimeId

def is_working(labor_supply: DiscreteAction) -> BoolND:
    """Convert the discrete labor-supply action to a boolean."""
    return labor_supply == LaborSupply.work


def labor_income(*, is_working: BoolND, wage: float | FloatND) -> FloatND:
    """Compute labor income (zero if not working)."""
    return jnp.where(is_working, wage, 0.0)


def utility_working(
    *, consumption: ContinuousAction, is_working: BoolND, disutility_of_work: float
) -> FloatND:
    """Log utility with a work-disutility term."""
    work_disutility = jnp.where(is_working, disutility_of_work, 0.0)
    return jnp.log(consumption) - work_disutility


def next_wealth(
    *,
    wealth: ContinuousState,
    consumption: ContinuousAction,
    labor_income: FloatND,
    interest_rate: float,
) -> ContinuousState:
    """Budget constraint: next-period wealth."""
    return (1 + interest_rate) * (wealth - consumption) + labor_income


def borrowing_constraint(
    *, consumption: ContinuousAction, wealth: ContinuousState
) -> BoolND:
    """The agent cannot consume more than current wealth."""
    return consumption <= wealth

Compare next_wealth above with the paper equation:

w_{t+1} = (1 + r)(w_t - c_t) + y_t

(3)

The code reads almost identically. And notice that labor_income appears as both a function name and a parameter of next_wealth. pylcm (via dags) resolves this: it calls labor_income(...) first and feeds the result into next_wealth. You never write labor_income = labor_income(is_working, wage) yourself.

Assembling the Regime¶

The functions are collected into functions and constraints dicts. The dict key is the name pylcm uses for wiring — so utility_working is registered under the key "utility" (which every regime must have):

working_functions = {
    "utility": utility_working,
    "labor_income": labor_income,
    "is_working": is_working,
}

working_constraints = {
    "borrowing_constraint": borrowing_constraint,
}

State transitions live on the Regime via the state_transitions dict. This separates outcome-space definitions (grids) from dynamics (transitions). Regime transitions live on the model’s edges: a source with several destinations at some age carries a Transition whose law picks one, and a regime without outgoing edges, like dead, is terminal:

def next_regime_from_working(
    *, labor_supply: DiscreteAction, age: float, final_age_alive: float
) -> ScalarInt:
    return jnp.where(
        age >= final_age_alive,
        RegimeId.dead,
        jnp.where(
            labor_supply == LaborSupply.retire,
            RegimeId.retirement,
            RegimeId.working_life,
        ),
    )


def next_regime_from_retirement(*, age: float, final_age_alive: float) -> ScalarInt:
    return jnp.where(age >= final_age_alive, RegimeId.dead, RegimeId.retirement)


def utility_retirement(consumption: ContinuousAction) -> FloatND:
    return jnp.log(consumption)


ages = AgeGrid(start=40, inclusive_stop=90, step="10Y")

working_life = Regime(
    states={
        "wealth": LinSpacedGrid(start=1, stop=400, n_points=100),
    },
    state_transitions={
        "wealth": next_wealth,
    },
    actions={
        "labor_supply": DiscreteGrid(category_class=LaborSupply),
        "consumption": LinSpacedGrid(start=1, stop=400, n_points=500),
    },
    functions=working_functions,
    constraints=working_constraints,
)

retirement = Regime(
    states={
        "wealth": LinSpacedGrid(start=1, stop=400, n_points=100),
    },
    state_transitions={
        "wealth": next_wealth,
    },
    actions={"consumption": LinSpacedGrid(start=1, stop=400, n_points=500)},
    functions={"utility": utility_retirement},
    constraints={"borrowing_constraint": borrowing_constraint},
)

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

# Pass this topology alongside the regimes when constructing Model.
edges = {
    # Work, retire or die until the age before the last one, then die.
    "working_life": Transition(
        targets={
            "working_life": (40, 50, 60, 70),
            "retirement": (40, 50, 60, 70),
            "dead": (40, 50, 60, 70, 80),
        },
        law=DeterministicTransition(func=next_regime_from_working),
    ),
    "retirement": Transition(
        targets={
            "retirement": (40, 50, 60, 70),
            "dead": (40, 50, 60, 70, 80),
        },
        law=DeterministicTransition(func=next_regime_from_retirement),
    ),
}

pylcm inspects the signatures of every registered function and constraint, builds the dependency graph, and determines that is_working must run before labor_income, which must run before next_wealth. You describe the economics; pylcm handles the plumbing.

Each function above is the same closure for every period the regime is solved. See Age-specialized functions and grids to let a function’s closure vary with age instead — for example, a tax-transfer system pinned to a policy date that changes as a cohort ages through calendar time.

For more on regimes — terminal vs non-terminal, stochastic transitions — see Regimes; for age-dependent laws with ByAge, see Age-indexed regimes.

What You Don’t Have to Do¶

Compare this with the traditional approach from above:

No manual call ordering. pylcm topologically sorts the DAG for you.
No passing intermediate results. The output of is_working is automatically fed to labor_income and utility_working.
No shared state dict. Each function declares exactly what it needs via its signature.
Adding a new auxiliary function is trivial. Add it to the functions dict and use its name as a parameter in other functions — done.

The One Rule: One Function, One Output¶

The key mental-model difference from traditional code: each pylcm function returns exactly one thing. In traditional code you might compute several values in one block; in pylcm you split them into separate functions, each returning a single result. This is what makes the DAG work — each node produces one named value that other nodes can reference by name.

Example contrast:

# Traditional: one function computes multiple things
def compute_income_and_tax(*, wage, hours, tax_rate):
    gross = wage * hours
    tax = gross * tax_rate
    net = gross - tax
    return gross, tax, net  # tuple — caller must unpack


# pylcm: each equation is its own function
def gross_income(*, wage, hours):
    return wage * hours


def tax(*, gross_income, tax_rate):
    return gross_income * tax_rate


def net_income(*, gross_income, tax):
    return gross_income - tax

This is actually closer to how economists write equations in papers — each equation defines one variable.

To see how to parametrise and solve this model, continue to Parameters and Solving and Simulating. For a complete end-to-end walkthrough, see A Tiny Example. To let a function’s closure or a continuous state’s grid vary with age, see Age-specialized functions and grids.