Parameter Specification Workflow

This notebook explains how to specify parameters for PyLCM models. It covers:

Understanding model.params_template - what parameters your model needs
Specifying parameters at different levels (function, regime, model)
Parameter propagation - how higher-level specifications flow down
Error handling - what’s not allowed and why

Setup: A Simple Two-Regime Model¶

Let’s create a simple consumption-savings model with two regimes: working and retired. This will help us understand how parameters are organized.

from pprint import pprint

import jax.numpy as jnp

from lcm import (
    AgeGrid,
    DiscreteGrid,
    LinSpacedGrid,
    Model,
    Regime,
    RegimeTransition,
    categorical,
)
from lcm.typing import ContinuousAction, ContinuousState, DiscreteAction, FloatND

@categorical
class WorkingStatus:
    retired: int
    working: int


@categorical
class RegimeId:
    working: int
    retired: int


def next_regime() -> int:
    return RegimeId.retired


def utility_working(
    consumption: ContinuousAction,
    working: DiscreteAction,
    risk_aversion: float,
    wage: float,
    disutility_of_work: float,
) -> FloatND:
    return (
        consumption ** (1 - risk_aversion) / (1 - risk_aversion)
        - disutility_of_work * jnp.log(wage) * working
    )


def utility_retired(consumption: ContinuousAction, risk_aversion: float) -> FloatND:
    return consumption ** (1 - risk_aversion) / (1 - risk_aversion)


def labor_income(wage: float, working: DiscreteAction) -> FloatND:
    return wage * working


def next_wealth(
    wealth: ContinuousState,
    consumption: ContinuousAction,
    labor_income: FloatND,
    interest_rate: float,
) -> ContinuousState:
    return (1 + interest_rate) * (wealth + labor_income - consumption)


def borrowing_constraint(
    consumption: ContinuousAction, wealth: ContinuousState
) -> FloatND:
    return consumption <= wealth

consumption_grid = LinSpacedGrid(start=1, stop=100, n_points=50)

working_regime = Regime(
    transition=RegimeTransition(next_regime),
    constraints={"borrowing_constraint": borrowing_constraint},
    functions={"utility": utility_working, "labor_income": labor_income},
    actions={
        "working": DiscreteGrid(WorkingStatus),
        "consumption": consumption_grid,
    },
    states={
        "wealth": LinSpacedGrid(start=1, stop=100, n_points=50, transition=next_wealth)
    },
)

retired_regime = Regime(
    transition=None,
    functions={"utility": utility_retired},
    constraints={"borrowing_constraint": borrowing_constraint},
    actions={"consumption": consumption_grid},
    states={"wealth": LinSpacedGrid(start=1, stop=100, n_points=50, transition=None)},
)

model = Model(
    description="Simple two-regime consumption-savings model",
    ages=AgeGrid(start=60, stop=62, step="Y"),
    regimes={
        "working": working_regime,
        "retired": retired_regime,
    },
    regime_id_class=RegimeId,
)

Understanding `model.params_template`¶

After creating a model, you can inspect model.params_template to see what parameters are required. The template is organized hierarchically:

params_template
├── regime_name
│   ├── function_name
│   │   ├── param_1: type
│   │   └── param_2: type
│   └── another_function
│       └── param_3: type
└── another_regime
    └── ...

pprint(dict(model.params_template))

{'retired': mappingproxy({'borrowing_constraint': mappingproxy({}),
                          'next_wealth': mappingproxy({}),
                          'utility': mappingproxy({'risk_aversion': <class 'float'>})}),
 'working': mappingproxy({'H': mappingproxy({'discount_factor': <class 'float'>}),
                          'borrowing_constraint': mappingproxy({}),
                          'labor_income': mappingproxy({'wage': <class 'float'>}),
                          'next_regime': mappingproxy({}),
                          'next_wealth': mappingproxy({'interest_rate': <class 'float'>}),
                          'utility': mappingproxy({'disutility_of_work': <class 'float'>,
                                                   'risk_aversion': <class 'float'>,
                                                   'wage': <class 'float'>})})}

Let’s break this down:

Working regime:

H function needs: discount_factor (for discounting future utility)
labor_income function needs: wage
next_wealth function needs: interest_rate
utility function needs: risk_aversion, disutility_of_work, and wage

Retired regime (terminal):

utility function needs: risk_aversion

Notice that:

borrowing_constraint and next_regime appear but need no parameters
wage is needed by both labor_income and utility in the working regime
risk_aversion is needed in both regimes (in each regime’s utility)
discount_factor only appears in the working regime (retired is terminal)
interest_rate only appears in the working regime (retired has no state transitions)

Specifying Parameters: Three Levels¶

PyLCM allows you to specify parameters at three levels:

Function level: Most specific - directly in the function’s dict
Regime level: Parameters propagate to all functions in that regime
Model level: Parameters propagate to all regimes and functions

This flexibility reduces repetition when the same parameter value is used across multiple functions or regimes. The following examples progressively move parameters up from function level to model level.

Level 1: Function-Level Specification (Most Explicit)¶

Every parameter is specified directly in its corresponding function dict. Note that wage and risk_aversion must be specified separately for each function that uses them.

params_function_level = {
    "working": {
        "H": {"discount_factor": 0.95},
        "labor_income": {"wage": 10.0},
        "next_wealth": {"interest_rate": 0.04},
        "utility": {"risk_aversion": 1.5, "disutility_of_work": 0.1, "wage": 10.0},
    },
    "retired": {
        "utility": {"risk_aversion": 1.5},
    },
}

V = model.solve(params_function_level, debug_mode=False)
print("Solved successfully with function-level params!")

Solved successfully with function-level params!

Level 2: Regime-Level Specification¶

wage has the same value in both labor_income and utility. By specifying it at the regime level, it propagates to all functions that need it — no duplication:

params_regime_level = {
    "working": {
        "H": {"discount_factor": 0.95},
        "wage": 10.0,
        "next_wealth": {"interest_rate": 0.04},
        "utility": {"risk_aversion": 1.5, "disutility_of_work": 0.1},
    },
    "retired": {
        "utility": {"risk_aversion": 1.5},
    },
}

V = model.solve(params_regime_level, debug_mode=False)
print("Solved successfully with regime-level params!")

Solved successfully with regime-level params!

Level 3: Model-Level Specification¶

risk_aversion has the same value in both regimes. Specifying it at the model level removes the need for a separate retired entry entirely:

params_model_level = {
    "risk_aversion": 1.5,
    "working": {
        "H": {"discount_factor": 0.95},
        "wage": 10.0,
        "next_wealth": {"interest_rate": 0.04},
        "utility": {"disutility_of_work": 0.1},
    },
}

V = model.solve(params_model_level, debug_mode=False)
print("Solved successfully with model-level params!")

Solved successfully with model-level params!

All Parameters at Model Level¶

If we do not need different parameter values depending on where they appear in the tree, we can just specify all of them at the model level. They will just propagate to the places where they are needed.

params_all_model_level = {
    "discount_factor": 0.95,
    "risk_aversion": 1.5,
    "wage": 10.0,
    "interest_rate": 0.04,
    "disutility_of_work": 0.1,
}

V = model.solve(params_all_model_level, debug_mode=False)
print("Solved successfully with all params at model level!")

Solved successfully with all params at model level!

What’s NOT Allowed: Ambiguous Specifications¶

The key rule is: you cannot specify the same parameter at multiple levels within the same subtree. This would be ambiguous - which value should be used?

PyLCM will raise an InvalidNameError if you try to do this.

Error 1: Same parameter at model AND regime level¶

from lcm.exceptions import InvalidNameError

params_ambiguous_model_regime = {
    "discount_factor": 0.95,
    "risk_aversion": 1.5,  # Model level
    "working": {
        "risk_aversion": 2.0,  # Also at regime level - AMBIGUOUS!
        "wage": 10.0,
        "interest_rate": 0.04,
        "disutility_of_work": 0.1,
    },
}

try:
    model.solve(params_ambiguous_model_regime)
except InvalidNameError as e:
    print(f"Error: {e}")

Error: Ambiguous parameter specification for 'working__utility__risk_aversion'. Found values at: ['working__risk_aversion', 'risk_aversion']

Error 2: Same parameter at regime AND function level¶

params_ambiguous_regime_function = {
    "discount_factor": 0.95,
    "risk_aversion": 1.5,
    "working": {
        "wage": 10.0,  # Regime level
        "labor_income": {
            "wage": 12.0,  # Also at function level - AMBIGUOUS!
        },
        "interest_rate": 0.04,
        "disutility_of_work": 0.1,
    },
}

try:
    model.solve(params_ambiguous_regime_function)
except InvalidNameError as e:
    print(f"Error: {e}")

Error: Ambiguous parameter specification for 'working__labor_income__wage'. Found values at: ['working__labor_income__wage', 'working__wage']

Error 3: Same parameter at model AND function level¶

params_ambiguous_model_function = {
    "wage": 15.0,  # Model level
    "discount_factor": 0.95,
    "risk_aversion": 1.5,
    "working": {
        "labor_income": {
            "wage": 10.0,  # Also at function level - AMBIGUOUS!
        },
        "interest_rate": 0.04,
        "disutility_of_work": 0.1,
    },
}

try:
    model.solve(params_ambiguous_model_function)
except InvalidNameError as e:
    print(f"Error: {e}")

Error: Ambiguous parameter specification for 'working__labor_income__wage'. Found values at: ['working__labor_income__wage', 'wage']

Name Requirements¶

To enable unambiguous parameter propagation, PyLCM enforces naming rules:

Regime names, function names, and argument names must not overlap with each other (within the same category across regimes)
Names cannot contain the separator __ (double underscore)

These rules are checked when the model is created. If violated, you’ll get an error during Model() initialization.

Summary¶

Level	Syntax	Propagation
Model	`{"param": value, ...}`	To all regimes and functions
Regime	`{"regime": {"param": value, ...}}`	To all functions in that regime
Function	`{"regime": {"func": {"param": value}}}`	Direct specification

Key Rules:

You can specify parameters at any level
Higher-level specifications propagate down automatically
The same parameter cannot be specified at multiple levels within a subtree
Use model.params_template to see what parameters are needed