A Tiny Example

A three-period consumption-savings model with two regimes:

Working life (ages 25 and 45): The agent chooses whether to work and how much to consume. A simple tax-and-transfer system guarantees a consumption floor. Savings earn interest.
Retirement (age 65): Terminal regime. The agent consumes out of remaining wealth.

This notebook teaches the common GridSearch workflow, not a universal model template. If your own model needs EGM, nested EGM, or declared institutional kinks, start from the matching specialized regime and margin declarations. Make the solver-choice checkpoint before adapting the code below.

Model¶

An agent lives for three periods (ages 25, 45, and 65). In the first two periods (working life), the agent chooses whether to work $d_t \in \{0, 1\}$ and how much to consume $c_t$ . In the final period (retirement), the agent consumes out of remaining wealth.

Working life (ages 25 and 45):

V_t(w_t) = \max_{d_t,\, c_t} \left\{ \frac{c_t^{1-\sigma}}{1-\sigma} - \phi \, d_t + \beta \, V_{t+1}(w_{t+1}) \right\}

(1)

subject to

\begin{align} e_t &= d_t \cdot \bar{w} \\[4pt] \tau(e_t, w_t) &= \begin{cases} \theta\,(e_t - \underline{c}) & \text{if } e_t \geq \underline{c} \\ \min(0,\; w_t + e_t - \underline{c}) & \text{otherwise} \end{cases} \\[4pt] a_t &= w_t + e_t - \tau(e_t, w_t) - c_t \\[4pt] w_{t+1} &= (1 + r)\, a_t \\[4pt] a_t &\geq 0 \end{align}

(2)

where $w_t$ is wealth, $e_t$ earnings, $\bar{w}$ the wage, $\underline{c}$ a consumption floor guaranteed by transfers, $\theta$ the tax rate, and $a_t$ end-of-period wealth. The transfer only kicks in when the agent’s resources ( $w_t + e_t$ ) fall below the consumption floor.

Retirement (age 65, terminal):

V_2(w_2) = \max_{c_2 \leq w_2} \frac{c_2^{1-\sigma}}{1-\sigma}

(3)

from pprint import pprint

import jax.numpy as jnp
import numpy as np
import pandas as pd
import plotly.express as px

from lcm import (
    AgeGrid,
    DeterministicTransition,
    DiscreteGrid,
    LinSpacedGrid,
    LogSpacedGrid,
    Model,
    Regime,
    Transition,
    categorical,
)
from lcm.typing import (
    BoolND,
    ContinuousAction,
    ContinuousState,
    DiscreteAction,
    FloatND,
    ScalarInt,
)

Categorical Variables¶

@categorical(ordered=True)
class LaborSupply:
    do_not_work: ScalarInt
    work: ScalarInt


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

Model Functions¶

# Utility


def utility(
    *,
    consumption: ContinuousAction,
    labor_supply: DiscreteAction,
    disutility_of_work: float,
    risk_aversion: float,
) -> FloatND:
    return consumption ** (1 - risk_aversion) / (
        1 - risk_aversion
    ) - disutility_of_work * (labor_supply == LaborSupply.work)


def utility_retirement(*, wealth: ContinuousState, risk_aversion: float) -> FloatND:
    return wealth ** (1 - risk_aversion) / (1 - risk_aversion)


# Auxiliary functions


def earnings(*, labor_supply: DiscreteAction, wage: float) -> FloatND:
    return jnp.where(labor_supply == LaborSupply.work, wage, 0.0)


def taxes_transfers(
    *,
    earnings: FloatND,
    wealth: ContinuousState,
    consumption_floor: float,
    tax_rate: float,
) -> FloatND:
    return jnp.where(
        earnings >= consumption_floor,
        tax_rate * (earnings - consumption_floor),
        jnp.minimum(0.0, wealth + earnings - consumption_floor),
    )


def end_of_period_wealth(
    *,
    wealth: ContinuousState,
    earnings: FloatND,
    taxes_transfers: FloatND,
    consumption: ContinuousAction,
) -> FloatND:
    return wealth + earnings - taxes_transfers - consumption


# State transition


def next_wealth(
    *, end_of_period_wealth: FloatND, interest_rate: float
) -> ContinuousState:
    return (1 + interest_rate) * end_of_period_wealth


# Constraints


def borrowing_constraint_working(end_of_period_wealth: FloatND) -> BoolND:
    return end_of_period_wealth >= 0


# Regime transition


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

Regimes and Model¶

age_grid = AgeGrid(start=25, inclusive_stop=65, step="20Y")
retirement_age = age_grid.exact_values[-1]

working_life = Regime(
    states={
        "wealth": LinSpacedGrid(start=0, stop=50, n_points=25),
    },
    state_transitions={
        "wealth": next_wealth,
    },
    actions={
        "labor_supply": DiscreteGrid(category_class=LaborSupply),
        "consumption": LogSpacedGrid(start=4, stop=50, n_points=100),
    },
    functions={
        "utility": utility,
        "earnings": earnings,
        "taxes_transfers": taxes_transfers,
        "end_of_period_wealth": end_of_period_wealth,
    },
    constraints={
        "borrowing_constraint_working": borrowing_constraint_working,
    },
)

retirement = Regime(
    states={
        "wealth": LinSpacedGrid(start=0, stop=50, n_points=25),
    },
    functions={"utility": utility_retirement},
)

# At 25 the graph permits both working life and retirement, and the age-based
# `next_regime` selects one; at 45 retirement is the only edge, and `next_regime`
# returns it there. `retirement` has no outgoing edges, which makes it terminal.
edges = {
    "working_life": Transition(
        targets={"working_life": 25, "retirement": (25, 45)},
        law=DeterministicTransition(func=next_regime),
    )
}

model = Model(
    regimes={
        "working_life": working_life,
        "retirement": retirement,
    },
    ages=age_grid,
    edges=edges,
    initial_nodes=((25, "working_life"),),
    regime_id_class=RegimeId,
    description="A tiny three-period consumption-savings model.",
)

Parameters¶

Use model.get_params_template() to see what parameters the model expects, organized by regime and function. The regime-transition law belongs to the edges, so its parameters sit under the "edges" root, keyed by the source regime.

pprint(dict(model.get_params_template()))

{'edges': {'working_life': {'last_working_age': 'float'}},
 'retirement': {'utility': {'risk_aversion': 'float'}},
 'working_life': {'borrowing_constraint_working': {},
                  'certainty_equivalent': {},
                  'earnings': {'wage': 'float'},
                  'end_of_period_wealth': {},
                  'koopmans_aggregator': {'discount_factor': 'FloatND'},
                  'next_wealth': {'interest_rate': 'float'},
                  'taxes_transfers': {'consumption_floor': 'float',
                                      'tax_rate': 'float'},
                  'utility': {'disutility_of_work': 'float',
                              'risk_aversion': 'float'}}}

Parameters shared across regimes (risk_aversion, discount_factor, interest_rate) can be specified at the model level. Parameters unique to one regime go under the regime name. last_working_age, read by the law next_regime, goes under "edges" and the source regime working_life.

params = {
    "discount_factor": 0.95,
    "risk_aversion": 1.5,
    "interest_rate": 0.03,
    "working_life": {
        "utility": {"disutility_of_work": 1.0},
        "earnings": {"wage": 20.0},
        "taxes_transfers": {"consumption_floor": 2.0, "tax_rate": 0.2},
    },
    "edges": {"working_life": {"last_working_age": age_grid.exact_values[-2]}},
}

Simulate¶

n_agents = 100

initial_df = pd.DataFrame(
    {
        "regime_name": "working_life",
        "age": float(age_grid.exact_values[0]),
        "wealth": np.linspace(1, 20, n_agents),
    }
)

result = model.simulate(
    params=params,
    initial_conditions=initial_df,
    log_level="debug",
)

solve call da1d0bf1d0e2 phase public_simulate begin

INFO:lcm:solve call da1d0bf1d0e2 phase public_simulate begin

solve call da1d0bf1d0e2 phase params_validation begin

INFO:lcm:solve call da1d0bf1d0e2 phase params_validation begin

solve call da1d0bf1d0e2 phase params_validation end status=ok seconds=0.107506

INFO:lcm:solve call da1d0bf1d0e2 phase params_validation end status=ok seconds=0.107506

solve call da1d0bf1d0e2 phase simulation_inputs begin

INFO:lcm:solve call da1d0bf1d0e2 phase simulation_inputs begin

solve call da1d0bf1d0e2 phase simulation_inputs end status=ok seconds=0.390899

INFO:lcm:solve call da1d0bf1d0e2 phase simulation_inputs end status=ok seconds=0.390899

solve call da1d0bf1d0e2 phase authority_fingerprint begin

INFO:lcm:solve call da1d0bf1d0e2 phase authority_fingerprint begin

solve call da1d0bf1d0e2 phase authority_fingerprint end status=ok seconds=0.004405

INFO:lcm:solve call da1d0bf1d0e2 phase authority_fingerprint end status=ok seconds=0.004405

solve call da1d0bf1d0e2 phase solver_param_checks begin

INFO:lcm:solve call da1d0bf1d0e2 phase solver_param_checks begin

solve call da1d0bf1d0e2 phase solver_param_checks end status=ok seconds=0.000043

INFO:lcm:solve call da1d0bf1d0e2 phase solver_param_checks end status=ok seconds=0.000043

solve call da1d0bf1d0e2 phase state_action_spaces begin

INFO:lcm:solve call da1d0bf1d0e2 phase state_action_spaces begin

solve call da1d0bf1d0e2 phase state_action_spaces end status=ok seconds=0.000248

INFO:lcm:solve call da1d0bf1d0e2 phase state_action_spaces end status=ok seconds=0.000248

solve call da1d0bf1d0e2 phase continuation_templates begin

INFO:lcm:solve call da1d0bf1d0e2 phase continuation_templates begin

solve call da1d0bf1d0e2 phase continuation_templates end status=ok seconds=0.001134

INFO:lcm:solve call da1d0bf1d0e2 phase continuation_templates end status=ok seconds=0.001134

solve call da1d0bf1d0e2 phase program_graphs begin

INFO:lcm:solve call da1d0bf1d0e2 phase program_graphs begin

solve call da1d0bf1d0e2 phase program_graphs end status=ok seconds=0.000652

INFO:lcm:solve call da1d0bf1d0e2 phase program_graphs end status=ok seconds=0.000652

solve call da1d0bf1d0e2 phase structural_resolution begin

INFO:lcm:solve call da1d0bf1d0e2 phase structural_resolution begin

structural blueprint: miss (hits=0, misses=1, entries=1)

DEBUG:lcm:structural blueprint: miss (hits=0, misses=1, entries=1)

solve call da1d0bf1d0e2 phase structural_resolution end status=ok seconds=0.033701

INFO:lcm:solve call da1d0bf1d0e2 phase structural_resolution end status=ok seconds=0.033701

solve call da1d0bf1d0e2 phase residency_inventory begin

INFO:lcm:solve call da1d0bf1d0e2 phase residency_inventory begin

solve call da1d0bf1d0e2 phase residency_inventory end status=ok seconds=0.000007

INFO:lcm:solve call da1d0bf1d0e2 phase residency_inventory end status=ok seconds=0.000007

solve call da1d0bf1d0e2 phase compilation_waves begin

INFO:lcm:solve call da1d0bf1d0e2 phase compilation_waves begin

candidate evaluation ('working_life', 0, 'main') {'action_product': 64, 'cell': 16}

DEBUG:lcm:candidate evaluation ('working_life', 0, 'main') {'action_product': 64, 'cell': 16}

candidate evaluation ('working_life', 1, 'main') {'action_product': 64, 'cell': 16}

DEBUG:lcm:candidate evaluation ('working_life', 1, 'main') {'action_product': 64, 'cell': 16}

candidate evaluation ('retirement', 1, 'main') {'cell': 16}

DEBUG:lcm:candidate evaluation ('retirement', 1, 'main') {'cell': 16}

candidate evaluation ('retirement', 2, 'main') {'cell': 16}

DEBUG:lcm:candidate evaluation ('retirement', 2, 'main') {'cell': 16}

AOT compilation wave 0: 3 unique lowerings for 4 regime-period-core triples (2 workers)

INFO:lcm:AOT compilation wave 0: 3 unique lowerings for 4 regime-period-core triples (2 workers)

1/3  working_life main (age 25, widths={'action_product': 64, 'cell': 16})

INFO:lcm:1/3  working_life main (age 25, widths={'action_product': 64, 'cell': 16})

  lowering working_life main (age 25, widths={'action_product': 64, 'cell': 16}) ...

INFO:lcm:  lowering working_life main (age 25, widths={'action_product': 64, 'cell': 16}) ...

  lowered in 332.6ms

INFO:lcm:  lowered in 332.6ms

  compiling working_life main (age 25, widths={'action_product': 64, 'cell': 16}) ...

INFO:lcm:  compiling working_life main (age 25, widths={'action_product': 64, 'cell': 16}) ...

2/3  working_life main (age 45, widths={'action_product': 64, 'cell': 16})

INFO:lcm:2/3  working_life main (age 45, widths={'action_product': 64, 'cell': 16})

  lowering working_life main (age 45, widths={'action_product': 64, 'cell': 16}) ...

INFO:lcm:  lowering working_life main (age 45, widths={'action_product': 64, 'cell': 16}) ...

  lowered in 261.0ms

INFO:lcm:  lowered in 261.0ms

  compiling working_life main (age 45, widths={'action_product': 64, 'cell': 16}) ...

INFO:lcm:  compiling working_life main (age 45, widths={'action_product': 64, 'cell': 16}) ...

3/3  retirement main (age 45, widths={'cell': 16})

INFO:lcm:3/3  retirement main (age 45, widths={'cell': 16})

  lowering retirement main (age 45, widths={'cell': 16}) ...

INFO:lcm:  lowering retirement main (age 45, widths={'cell': 16}) ...

  lowered in 105.3ms

INFO:lcm:  lowered in 105.3ms

  compiled  working_life main (age 25, widths={'action_product': 64, 'cell': 16})  804.0ms

INFO:lcm:  compiled  working_life main (age 25, widths={'action_product': 64, 'cell': 16})  804.0ms

  compiling retirement main (age 45, widths={'cell': 16}) ...

INFO:lcm:  compiling retirement main (age 45, widths={'cell': 16}) ...

  compiled  retirement main (age 45, widths={'cell': 16})  138.2ms

INFO:lcm:  compiled  retirement main (age 45, widths={'cell': 16})  138.2ms

  compiled  working_life main (age 45, widths={'action_product': 64, 'cell': 16})  705.7ms

INFO:lcm:  compiled  working_life main (age 45, widths={'action_product': 64, 'cell': 16})  705.7ms

AOT compilation wave 0: 3 unique variants compiled, 0 fallback evaluations skipped, 0 fallback requests avoided

INFO:lcm:AOT compilation wave 0: 3 unique variants compiled, 0 fallback evaluations skipped, 0 fallback requests avoided

solve call da1d0bf1d0e2 phase compilation_waves end status=ok seconds=1.333464

INFO:lcm:solve call da1d0bf1d0e2 phase compilation_waves end status=ok seconds=1.333464

solve call da1d0bf1d0e2 phase workspace_selection begin

INFO:lcm:solve call da1d0bf1d0e2 phase workspace_selection begin

core plan record {"active_replica_bytes":0,"axis_extents":{"action_product":200,"cell":25},"compiled_collectives":{},"compiler_peak_bytes":939,"compiler_reservation_bytes":null,"core":"main","device_ids":[0],"dispatch":{"action_product":"streamed","cell":"streamed"},"logical_gathered_bytes":0,"logical_value_shape":[25],"period":0,"period_transfer_scratch_bytes":{},"physical_value_shape":[25],"regime":"working_life","resident_bytes":null,"selected_block":null,"state_extents":{"wealth":25},"stored_owner_bytes":{},"transfer_workspace_bytes":0,"widths":{"action_product":64,"cell":16}}

DEBUG:lcm:core plan record {"active_replica_bytes":0,"axis_extents":{"action_product":200,"cell":25},"compiled_collectives":{},"compiler_peak_bytes":939,"compiler_reservation_bytes":null,"core":"main","device_ids":[0],"dispatch":{"action_product":"streamed","cell":"streamed"},"logical_gathered_bytes":0,"logical_value_shape":[25],"period":0,"period_transfer_scratch_bytes":{},"physical_value_shape":[25],"regime":"working_life","resident_bytes":null,"selected_block":null,"state_extents":{"wealth":25},"stored_owner_bytes":{},"transfer_workspace_bytes":0,"widths":{"action_product":64,"cell":16}}

core plan record {"active_replica_bytes":0,"axis_extents":{"action_product":200,"cell":25},"compiled_collectives":{},"compiler_peak_bytes":835,"compiler_reservation_bytes":null,"core":"main","device_ids":[0],"dispatch":{"action_product":"streamed","cell":"streamed"},"logical_gathered_bytes":0,"logical_value_shape":[25],"period":1,"period_transfer_scratch_bytes":{},"physical_value_shape":[25],"regime":"working_life","resident_bytes":null,"selected_block":null,"state_extents":{"wealth":25},"stored_owner_bytes":{},"transfer_workspace_bytes":0,"widths":{"action_product":64,"cell":16}}

DEBUG:lcm:core plan record {"active_replica_bytes":0,"axis_extents":{"action_product":200,"cell":25},"compiled_collectives":{},"compiler_peak_bytes":835,"compiler_reservation_bytes":null,"core":"main","device_ids":[0],"dispatch":{"action_product":"streamed","cell":"streamed"},"logical_gathered_bytes":0,"logical_value_shape":[25],"period":1,"period_transfer_scratch_bytes":{},"physical_value_shape":[25],"regime":"working_life","resident_bytes":null,"selected_block":null,"state_extents":{"wealth":25},"stored_owner_bytes":{},"transfer_workspace_bytes":0,"widths":{"action_product":64,"cell":16}}

core plan record {"active_replica_bytes":0,"axis_extents":{"cell":25},"compiled_collectives":{},"compiler_peak_bytes":236,"compiler_reservation_bytes":null,"core":"main","device_ids":[0],"dispatch":{"cell":"streamed"},"logical_gathered_bytes":0,"logical_value_shape":[25],"period":1,"period_transfer_scratch_bytes":{},"physical_value_shape":[25],"regime":"retirement","resident_bytes":null,"selected_block":null,"state_extents":{"wealth":25},"stored_owner_bytes":{},"transfer_workspace_bytes":0,"widths":{"cell":16}}

DEBUG:lcm:core plan record {"active_replica_bytes":0,"axis_extents":{"cell":25},"compiled_collectives":{},"compiler_peak_bytes":236,"compiler_reservation_bytes":null,"core":"main","device_ids":[0],"dispatch":{"cell":"streamed"},"logical_gathered_bytes":0,"logical_value_shape":[25],"period":1,"period_transfer_scratch_bytes":{},"physical_value_shape":[25],"regime":"retirement","resident_bytes":null,"selected_block":null,"state_extents":{"wealth":25},"stored_owner_bytes":{},"transfer_workspace_bytes":0,"widths":{"cell":16}}

core plan record {"active_replica_bytes":0,"axis_extents":{"cell":25},"compiled_collectives":{},"compiler_peak_bytes":236,"compiler_reservation_bytes":null,"core":"main","device_ids":[0],"dispatch":{"cell":"streamed"},"logical_gathered_bytes":0,"logical_value_shape":[25],"period":2,"period_transfer_scratch_bytes":{},"physical_value_shape":[25],"regime":"retirement","resident_bytes":null,"selected_block":null,"state_extents":{"wealth":25},"stored_owner_bytes":{},"transfer_workspace_bytes":0,"widths":{"cell":16}}

DEBUG:lcm:core plan record {"active_replica_bytes":0,"axis_extents":{"cell":25},"compiled_collectives":{},"compiler_peak_bytes":236,"compiler_reservation_bytes":null,"core":"main","device_ids":[0],"dispatch":{"cell":"streamed"},"logical_gathered_bytes":0,"logical_value_shape":[25],"period":2,"period_transfer_scratch_bytes":{},"physical_value_shape":[25],"regime":"retirement","resident_bytes":null,"selected_block":null,"state_extents":{"wealth":25},"stored_owner_bytes":{},"transfer_workspace_bytes":0,"widths":{"cell":16}}

solve call da1d0bf1d0e2 phase workspace_selection end status=ok seconds=0.023469

INFO:lcm:solve call da1d0bf1d0e2 phase workspace_selection end status=ok seconds=0.023469

Starting solution

INFO:lcm:Starting solution

solve call da1d0bf1d0e2 phase backward_induction begin

INFO:lcm:solve call da1d0bf1d0e2 phase backward_induction begin

Age 65 (1 regimes):

INFO:lcm:Age 65 (1 regimes):

  finished in 141.0ms

INFO:lcm:  finished in 141.0ms

Age 45 (2 regimes):

INFO:lcm:Age 45 (2 regimes):

  finished in 26.5ms

INFO:lcm:  finished in 26.5ms

Age 25 (1 regimes):

INFO:lcm:Age 25 (1 regimes):

  finished in 4.2ms

INFO:lcm:  finished in 4.2ms

Inf in V_arr for 1 other state(s) for regime 'retirement', period 2 (age 65.0)

WARNING:lcm:Inf in V_arr for 1 other state(s) for regime 'retirement', period 2 (age 65.0)

Inf in V_arr for 1 other state(s) for regime 'retirement', period 1 (age 45.0)

WARNING:lcm:Inf in V_arr for 1 other state(s) for regime 'retirement', period 1 (age 45.0)

  retirement  age 65.0   V min=-inf  max=-0.283  mean=-inf

DEBUG:lcm:  retirement  age 65.0   V min=-inf  max=-0.283  mean=-inf

  working_life  age 45.0   V min=-2.35  max=-0.774  mean=-1.29

DEBUG:lcm:  working_life  age 45.0   V min=-2.35  max=-0.774  mean=-1.29

  retirement  age 45.0   V min=-inf  max=-0.283  mean=-inf

DEBUG:lcm:  retirement  age 45.0   V min=-inf  max=-0.283  mean=-inf

  working_life  age 25.0   V min=-3.42  max=-1.38  mean=-2.15

DEBUG:lcm:  working_life  age 25.0   V min=-3.42  max=-1.38  mean=-2.15

solve call da1d0bf1d0e2 phase backward_induction end status=ok seconds=0.199210

INFO:lcm:solve call da1d0bf1d0e2 phase backward_induction end status=ok seconds=0.199210

Solution complete  (200.7ms)

INFO:lcm:Solution complete  (200.7ms)

solve call da1d0bf1d0e2 phase result_assembly begin

INFO:lcm:solve call da1d0bf1d0e2 phase result_assembly begin

solve call da1d0bf1d0e2 phase result_assembly end status=ok seconds=0.004465

INFO:lcm:solve call da1d0bf1d0e2 phase result_assembly end status=ok seconds=0.004465

solve call da1d0bf1d0e2 phase solution_resolution begin

INFO:lcm:solve call da1d0bf1d0e2 phase solution_resolution begin

solve call da1d0bf1d0e2 phase solution_resolution end status=ok seconds=0.001790

INFO:lcm:solve call da1d0bf1d0e2 phase solution_resolution end status=ok seconds=0.001790

solve call da1d0bf1d0e2 phase solution_handover begin

INFO:lcm:solve call da1d0bf1d0e2 phase solution_handover begin

solve call da1d0bf1d0e2 phase solution_handover end status=ok seconds=0.000004

INFO:lcm:solve call da1d0bf1d0e2 phase solution_handover end status=ok seconds=0.000004

solve call da1d0bf1d0e2 phase chunk_planning begin

INFO:lcm:solve call da1d0bf1d0e2 phase chunk_planning begin

solve call da1d0bf1d0e2 phase chunk_planning end status=ok seconds=0.000290

INFO:lcm:solve call da1d0bf1d0e2 phase chunk_planning end status=ok seconds=0.000290

solve call da1d0bf1d0e2 phase simulation_setup begin

INFO:lcm:solve call da1d0bf1d0e2 phase simulation_setup begin

Starting simulation

INFO:lcm:Starting simulation

Simulation plan: route=legacy; 1 cpu device(s) (0,); outer chunks=1 (widths=(100,)); unbudgeted.

INFO:lcm:Simulation plan: route=legacy; 1 cpu device(s) (0,); outer chunks=1 (widths=(100,)); unbudgeted.

Simulation plan detail: route=legacy; subject_devices=(cpu) (0,); axis_widths_by_regime=[working_life: {}; retirement: {}]; outer_chunk_count=1; admitted_chunk_widths=(100,); budget_mode=unbudgeted; effective_device_memory_bytes=None; subject_grouping=None.

DEBUG:lcm:Simulation plan detail: route=legacy; subject_devices=(cpu) (0,); axis_widths_by_regime=[working_life: {}; retirement: {}]; outer_chunk_count=1; admitted_chunk_widths=(100,); budget_mode=unbudgeted; effective_device_memory_bytes=None; subject_grouping=None.

solve call da1d0bf1d0e2 phase simulation_setup end status=ok seconds=0.262583

INFO:lcm:solve call da1d0bf1d0e2 phase simulation_setup end status=ok seconds=0.262583

solve call da1d0bf1d0e2 phase simulation_chunk begin

INFO:lcm:solve call da1d0bf1d0e2 phase simulation_chunk begin

Age 25 (1 regimes):

INFO:lcm:Age 25 (1 regimes):

solve call da1d0bf1d0e2 phase simulation_compilation begin

INFO:lcm:solve call da1d0bf1d0e2 phase simulation_compilation begin

solve call da1d0bf1d0e2 phase simulation_compilation end status=ok seconds=0.481220

INFO:lcm:solve call da1d0bf1d0e2 phase simulation_compilation end status=ok seconds=0.481220

solve call da1d0bf1d0e2 phase simulation_compilation begin

INFO:lcm:solve call da1d0bf1d0e2 phase simulation_compilation begin

solve call da1d0bf1d0e2 phase simulation_compilation end status=ok seconds=0.045990

INFO:lcm:solve call da1d0bf1d0e2 phase simulation_compilation end status=ok seconds=0.045990

solve call da1d0bf1d0e2 phase simulation_compilation begin

INFO:lcm:solve call da1d0bf1d0e2 phase simulation_compilation begin

solve call da1d0bf1d0e2 phase simulation_compilation end status=ok seconds=0.035131

INFO:lcm:solve call da1d0bf1d0e2 phase simulation_compilation end status=ok seconds=0.035131

Off-grid next states: regime='working_life', target='retirement', period=0, state='wealth', phase='realized', count=0

INFO:lcm:Off-grid next states: regime='working_life', target='retirement', period=0, state='wealth', phase='realized', count=0

Off-grid next states: regime='working_life', target='working_life', period=0, state='wealth', phase='realized', count=0

INFO:lcm:Off-grid next states: regime='working_life', target='working_life', period=0, state='wealth', phase='realized', count=0

  transitions:
  - working_life → working_life = 100

DEBUG:lcm:  transitions:
  - working_life → working_life = 100

  finished in 2.1s

INFO:lcm:  finished in 2.1s

Age 45 (2 regimes):

INFO:lcm:Age 45 (2 regimes):

solve call da1d0bf1d0e2 phase simulation_compilation begin

INFO:lcm:solve call da1d0bf1d0e2 phase simulation_compilation begin

solve call da1d0bf1d0e2 phase simulation_compilation end status=ok seconds=0.313000

INFO:lcm:solve call da1d0bf1d0e2 phase simulation_compilation end status=ok seconds=0.313000

solve call da1d0bf1d0e2 phase simulation_compilation begin

INFO:lcm:solve call da1d0bf1d0e2 phase simulation_compilation begin

solve call da1d0bf1d0e2 phase simulation_compilation end status=ok seconds=0.037715

INFO:lcm:solve call da1d0bf1d0e2 phase simulation_compilation end status=ok seconds=0.037715

Off-grid next states: regime='working_life', target='retirement', period=1, state='wealth', phase='realized', count=0

INFO:lcm:Off-grid next states: regime='working_life', target='retirement', period=1, state='wealth', phase='realized', count=0

solve call da1d0bf1d0e2 phase simulation_compilation begin

INFO:lcm:solve call da1d0bf1d0e2 phase simulation_compilation begin

solve call da1d0bf1d0e2 phase simulation_compilation end status=ok seconds=0.043507

INFO:lcm:solve call da1d0bf1d0e2 phase simulation_compilation end status=ok seconds=0.043507

  transitions:
  - working_life → retirement = 100

DEBUG:lcm:  transitions:
  - working_life → retirement = 100

  finished in 755.3ms

INFO:lcm:  finished in 755.3ms

Age 65 (1 regimes):

INFO:lcm:Age 65 (1 regimes):

  transitions:
  - retirement → retirement = 100

DEBUG:lcm:  transitions:
  - retirement → retirement = 100

  finished in 61.4ms

INFO:lcm:  finished in 61.4ms

solve call da1d0bf1d0e2 phase simulation_chunk end status=ok seconds=2.906751

INFO:lcm:solve call da1d0bf1d0e2 phase simulation_chunk end status=ok seconds=2.906751

solve call da1d0bf1d0e2 phase simulation_completion begin

INFO:lcm:solve call da1d0bf1d0e2 phase simulation_completion begin

Simulation complete  (3.2s)

INFO:lcm:Simulation complete  (3.2s)

solve call da1d0bf1d0e2 phase simulation_completion end status=ok seconds=0.001010

INFO:lcm:solve call da1d0bf1d0e2 phase simulation_completion end status=ok seconds=0.001010

solve call da1d0bf1d0e2 phase simulation_result begin

INFO:lcm:solve call da1d0bf1d0e2 phase simulation_result begin

solve call da1d0bf1d0e2 phase simulation_result end status=ok seconds=0.005306

INFO:lcm:solve call da1d0bf1d0e2 phase simulation_result end status=ok seconds=0.005306

solve call da1d0bf1d0e2 phase result_finalization begin

INFO:lcm:solve call da1d0bf1d0e2 phase result_finalization begin

solve call da1d0bf1d0e2 phase result_finalization end status=ok seconds=0.000008

INFO:lcm:solve call da1d0bf1d0e2 phase result_finalization end status=ok seconds=0.000008

solve call da1d0bf1d0e2 phase public_simulate end status=ok seconds=5.334448

INFO:lcm:solve call da1d0bf1d0e2 phase public_simulate end status=ok seconds=5.334448

df = result.to_dataframe(additional_targets="all")
df["age"] = df["age"].astype(int)
df.loc[df["age"] == retirement_age, "consumption"] = df.loc[
    df["age"] == retirement_age, "wealth"
]
columns = [
    "regime_name",
    "labor_supply",
    "consumption",
    "wealth",
    "earnings",
    "taxes_transfers",
    "end_of_period_wealth",
    "value",
]
df.set_index(["subject_id", "age"])[columns].head(20).style.format(
    precision=1,
    na_rep="",
)

# Classify agents by work pattern across the two working-life periods
first_working_age = age_grid.exact_values[0]
last_working_age = age_grid.exact_values[-2]

df_working = df[df["regime_name"] == "working_life"]
work_by_age = df_working.pivot_table(
    index="subject_id",
    columns="age",
    values="labor_supply",
    aggfunc="first",
)
work_pattern = (
    work_by_age[first_working_age].astype(str)
    + ", "
    + work_by_age[last_working_age].astype(str)
)
assert "work, work" not in work_pattern.to_numpy(), (
    "Plotting assumes that no agent works in both periods of working life."
)

label_map = {
    "work, do_not_work": "low",  # work early, not later
    "do_not_work, work": "medium",  # coast early, work later
    "do_not_work, do_not_work": "high",  # never work
}
groups = work_pattern.map(label_map).rename("initial_wealth")

# Combined descriptives and work decisions table
initial_wealth = df[df["age"] == first_working_age].set_index("subject_id")["wealth"]
group_desc = initial_wealth.groupby(groups).agg(["min", "max"]).round(1)

df_groups = df.copy()
df_groups["initial_wealth"] = df_groups["subject_id"].map(groups)
df_mean = df_groups.groupby(["initial_wealth", "age"], as_index=False).mean(
    numeric_only=True,
)
work_table = df_mean[df_mean["age"] < retirement_age].pivot_table(
    index="initial_wealth",
    columns="age",
    values="earnings",
)
work_table = (work_table > 0).astype(int)
work_table.columns = [f"works {c}" for c in work_table.columns]

summary = pd.concat([group_desc, work_table], axis=1)
summary.index.name = "initial_wealth"
summary.loc[["low", "medium", "high"]].style.format(precision=1, na_rep="")