Age-specialized functions and grids

Some model functions do not just take the agent’s age as an argument — their whole body changes with it. The canonical case is a tax-transfer system that changes as a cohort ages through calendar time. With each age, the laws that apply to the cohort change, so the function computing net income at age 58 is a different closure from the one at age 63, not the same closure with a different age argument.

AgeSpecializedFunction marks such a function. At model build, pylcm resolves the marker per period: each period’s age gets its own concrete function, compiled into that period’s programs. There is no runtime dispatch on age or calendar year.

from lcm import AgeSpecializedFunction


def make_net_income(age: float):
    """Build the net-income closure for the policy year this age falls in."""
    policy_env = load_policy_environment(year_for(age))

    def net_income(gross_income):
        return policy_env.apply(gross_income)

    return net_income


def policy_key(age: float):
    """Identify the closure: ages in the same policy year share a program."""
    return year_for(age)


functions = {
    "net_income": AgeSpecializedFunction(build=make_net_income, signature=policy_key),
    # ... other regime functions ...
}

The two contracts¶

AgeSpecializedFunction(build, signature) places two obligations on the user; pylcm cannot verify either automatically.

Invariant call signature. Every concrete function build(age) returns must expose the same argument list at every age. Only the constants the closure binds may differ. pylcm’s static passes (the parameter template, used-variable validation) read the function at one representative age and rely on this.
signature is a correctness precondition, not a performance hint. Periods whose signature(age) values are equal share a single compiled program. An equal signature must therefore imply identical closure behavior — same policy date, same price level, same overrides, same everything the closure binds. If two ages collide to one signature but build returns different closures, the solve is silently wrong. When in doubt, put more into the signature (a tuple of every varying ingredient), never less.

Deduplication is what keeps this affordable: an age-invariant signature collapses to one program for the whole horizon (identical to not using the wrapper), and a policy-date signature compiles one program per distinct policy year, not per period.

Where it is allowed¶

AgeSpecializedFunction is accepted in functions and constraints of non-terminal regimes. Everything else is rejected at Regime construction, loudly:

the regime transition, bare or wrapped in MarkovTransition — periodizing the regime-transition probability path is not supported;
a regime transition that reads an AgeSpecializedFunction, directly or through plain helper functions — regime-transition probabilities are built once, so a policy-specialized value flowing into them would reuse one age’s closure everywhere;
MarkovTransition(AgeSpecializedFunction(...)) — policy-specialized stochastic transitions are not supported;
a direct state_transitions value — see the pattern below instead;
anything in a terminal regime — the terminal value program is built once and shared across all periods.

Policy-dependent laws of motion¶

A state transition whose content depends on the policy year is written as a plain transition reading an age-specialized helper function:

functions = {
    "new_pension_points": AgeSpecializedFunction(
        build=make_points, signature=policy_key
    ),
    # ... other regime functions ...
}

state_transitions = {
    # A plain function — the policy content lives in the helper it reads.
    "pension_points": lambda pension_points, new_pension_points: (
        pension_points + new_pension_points
    ),
}

The simulate-phase next-state programs are built per period, so the helper resolves to the current period’s closure inside every transition — the law of motion tracks the policy year without the transition itself carrying a marker.

What is resolved when¶

Per period (exact): the solve Q_and_F programs, the simulate decision programs, and the simulate next-state programs. These are what determine model behavior, and they always use each period’s own closure.
At one representative age (the regime’s first active age): the parameter template, used-variable validation, and the published function sets (regime.solution.functions, regime.simulation.functions). These consumers only need the (age-invariant) call signature — except one:

to_dataframe(additional_targets=...) computes targets from the published simulation functions. A target that depends on an age-specialized function would therefore be evaluated under the representative age’s policy, not each row’s period — so pylcm rejects such targets with InvalidAdditionalTargetsError. Quantities you need per period from a specialized function should be carried inside the model (for example as a state fed by a plain transition reading the helper).

Age-specialized functions and grids are fully resolved during model creation; solve and simulation only select the prebuilt period-specific objects — no build(age) factory is ever called from backward induction, simulation, AOT compilation, or diagnostics.

Age-specialized continuous-state grids¶

AgeSpecializedGrid is the grid counterpart of AgeSpecializedFunction. Use it when a continuous state’s grid bounds or nodes are known from age at model creation, for example an age-dependent borrowing limit.

from lcm import AgeSpecializedGrid, LinSpacedGrid

states = {
    "assets": AgeSpecializedGrid(
        build=lambda age: LinSpacedGrid(
            start=borrowing_floor(age),
            stop=ASSET_MAX,
            n_points=40,
        ),
        signature=lambda age: borrowing_floor(age),
    ),
}

pylcm builds the concrete grid for every active age during model creation. The grid class, number of points, node-array shape, and dtype must stay fixed; only bounds or node values may change. Equal signature(age) values must identify identical grids, because those periods may share a compiled program; pylcm cross-checks a shared group’s resolved nodes at build time and raises if an equal signature turns out to hide genuinely different grids.

AgeSpecializedGrid is currently supported for continuous states, not actions, discrete states, or stochastic-process grids. A grid whose points depend on model parameters should use the ordinary runtime-points IrregSpacedGrid mechanism instead; an AgeSpecializedGrid that resolves to a runtime-points grid is rejected at model creation.