Internal Architecture

This page describes how pylcm’s source tree is organised and why it is laid out that way. The audience is contributors and advanced users who want to find code, add a feature, or understand which module they should be editing. End-users only ever write from lcm import Model, Regime, ... and never need anything here.

The lcm / _lcm split¶

pylcm’s source is two packages, and the split is a hard binary:

• src/lcm/ — the public surface. Everything a user constructs or consumes lives here, physically: the user-facing classes, the @categorical decorator, the as_leaf helper, the public type aliases, and the exception classes. lcm/__init__.py re-exports the public symbols so users write from lcm import Model.

• src/_lcm/ — the private implementation. The build pipeline, the canonical engine dataclasses, the JAX-traced solve / simulate machinery, validators, I/O plumbing, and the engine-side type aliases. The leading underscore on the package carries the entire “private” signal — modules inside _lcm/ are plainly named.

There is no gradient and no per-module underscore convention: a module is either in lcm/ (public) or in _lcm/ (private). Internal code reaches the user-facing classes through from lcm.regime import Regime as UserRegime etc.; the public lcm/__init__.py imports _lcm first so the jaxtyping patch and the beartype claw are installed before anything else loads.

lcm/
├── __init__.py       ← re-export façade for the public symbols
├── ages.py           ← AgeGrid
├── grids.py          ← LinSpacedGrid, LogSpacedGrid, IrregSpacedGrid, DiscreteGrid,
│                       PiecewiseLinSpacedGrid, PiecewiseLogSpacedGrid,
│                       GridBreakpoint, and the @categorical decorator
├── model.py          ← Model
├── params.py         ← as_leaf + the MappingLeaf / SequenceLeaf re-exports
├── persistence.py    ← SolveSnapshot, SimulateSnapshot, load_snapshot,
│                       complete-result save/load and legacy value reader
├── processes.py      ← the seven *Process classes
├── regime.py         ← Regime (Phased and MarkovTransition re-exported)
├── result.py         ← SimulationResult
├── solver_api.py     ← versioned solver, artifact, replay, and solution contracts
├── solvers.py        ← built-in solvers plus the out-of-tree re-export façade
├── transition.py     ← transition helpers
├── typing.py         ← user-facing type aliases
└── exceptions.py     ← every project-specific exception class

_lcm/
├── __init__.py            ← applies the jaxtyping patch + registers the beartype claw
├── ages.py                ← AgeGrid validators and step parsing
├── beartype_conf.py       ← the beartype configurations
├── config.py              ← build-time configuration constants
├── dtypes.py              ← canonical-dtype resolution
├── engine.py              ← canonical / engine-side dataclasses
├── jaxtyping_patch.py     ← bootstrap patch run before any jaxtyping type
├── model_processing.py    ← Model.__init__ build pipeline
├── pandas_utils.py        ← pd.Series ↔ JAX array bridge
├── reachability.py        ← construction-time solve/simulate regime graphs
├── state_action_space.py  ← state / action space validators
├── transition_checks.py   ← pre-solve regime + state transition prob checks
├── typing.py              ← engine-side type aliases and protocols
├── user_regime_validation.py ← validators for the user-facing Regime
├── variables.py           ← factories that build `Variables` from `Regime`
├── version.py             ← generated version string (hatch-vcs)
├── grids/                 ← grid infrastructure
├── processes/             ← stochastic-process infrastructure
├── persistence/           ← snapshot I/O and versioned solution-archive internals
├── regime_building/       ← per-regime canonicalisation
├── solution/              ← backward induction (solve) + validate_V
├── simulation/            ← forward sampling (simulate) + result helpers
├── params/                ← params templating and processing
└── utils/                 ← small, dependency-free helpers

Names cross the boundary in exactly one direction (user → canonical form) and only twice — once when Model(regimes={...}) is called (which triggers model_processing.build_regimes_and_template → regime_building.processing.process_regimes), and once for flat_params at every solve / simulate call.

The public surface — lcm/¶

lcm/ is the canonical home for every class users construct or consume. Keeping all of them in one package makes the public surface easy to find, review, and keep stable — anything in _lcm/ is fair game for refactoring. lcm/ modules are shallow by design: each file holds class definitions and the small number of public top-level functions that round out the surface. Validators, I/O plumbing, DataFrame assembly, and similar implementation detail live in _lcm/ and are imported back in.

The mapping of public names to files:

Why a package boundary, not just naming?¶

A naming convention (a _private_ prefix on every internal module) tells readers what is internal. A package boundary makes it visible to tools: code-search, auto-import, public-API audits, and the linter can all key off _lcm/. The boundary is enforced by the absence of imports — lcm/ modules import from _lcm/ only at well-defined wiring points, and _lcm/ reaches the user-facing classes through aliased imports (from lcm.regime import Regime as UserRegime).

Grid and process infrastructure: _lcm/grids/ and _lcm/processes/¶

_lcm/grids/
├── base.py            ← Grid, ContinuousGrid, UniformContinuousGrid (ABCs)
├── continuous.py      ← LinSpacedGrid, LogSpacedGrid, IrregSpacedGrid
├── discrete.py        ← DiscreteGrid
├── piecewise.py       ← PiecewiseLinSpacedGrid, PiecewiseLogSpacedGrid,
│                         GridBreakpoint
├── categorical.py     ← @categorical decorator + validators
└── coordinates.py     ← coordinate lookup helpers used by interpolation

_lcm/processes/
├── base.py            ← _ContinuousStochasticProcess + Gauss-Hermite / mixture helpers
├── iid.py             ← UniformIIDProcess, NormalIIDProcess, LogNormalIIDProcess,
│                        NormalMixtureIIDProcess
└── ar1.py             ← TauchenAR1Process, RouwenhorstAR1Process,
                         TauchenNormalMixtureAR1Process

The leaf classes are surfaced through lcm/grids.py and lcm/processes.py; the ABCs (Grid, _ContinuousStochasticProcess, etc.) are used by internal code but are not part of the documented public API.

Two design points worth knowing:

• Process classes bundle both a discretization grid AND a transition mechanism, unlike ordinary grids which are pure outcome-space. Users place UniformIIDProcess(...) in Regime(states=...) directly — the transition is invoked automatically. Putting a process class in state_transitions is a bug.

• All vocabulary in the engine speaks of process, not shock. Use is_process (on VariableInfo), process_names (on Variables), and ProcessName (typing alias). shock is reserved for the colloquial meaning and never appears as an identifier.

Private siblings of the public modules¶

Several lcm/ modules have a private counterpart in _lcm/ that holds their implementation detail. The pattern is the same throughout: lcm/ keeps the class definitions and the public top-level functions; the _lcm/ counterpart holds validators, helpers, and I/O plumbing that internal code is free to refactor.

_lcm/ages.py            ← STEP_UNITS, _parse_step,
                          _validate_age_grid / _validate_range / _validate_values
_lcm/user_regime_validation.py  ← the validators called from Regime.__post_init__

_lcm/simulation/initial_conditions.py
                        ← MISSING_CAT_CODE, PSEUDO_STATE_NAMES, and the
                          build / validate helpers for initial conditions

_lcm/persistence/
├── io.py               ← _atomic_dump, _save_pkl, _save_h5, _load_h5,
│                          _get_platform, _next_counter, _enforce_retention,
│                          _write_metadata, _write_environment_files
└── snapshots.py        ← _save_solve_snapshot, _save_simulate_snapshot,
                          _strip_V_arr_from_result, _bind_forward_refs

_lcm/simulation/result_metadata.py
                        ← ResultMetadata + _compute_metadata, _get_output_dtypes
_lcm/simulation/result_dataframe.py
                        ← _create_flat_dataframe and the per-regime / per-period
                          assembly helpers, plus categorical conversion
_lcm/simulation/additional_targets.py
                        ← _resolve_targets, _compute_targets, and DAG helpers
                          for to_dataframe(additional_targets=...)

Why split these out? Two reasons:

• The public surface is easier to audit. regime.py, persistence.py, and result.py each contain only the dozen-or-so symbols users actually touch. A reader looking for “what is the public contract of a Regime?” sees that contract directly, without scrolling past validator bodies.

• Internal helpers can move freely. Anything in _lcm/ is internal — its location, signature, and existence can change without bumping the user surface.

A note on shadowing: the canonical Regime lives in _lcm/engine.py. The validators in _lcm/user_regime_validation.py operate on the user-facing lcm.regime.Regime and reach it through TYPE_CHECKING-guarded imports to break the circular dependency at import time; beartype resolves the forward references at first call.

Engine-side: _lcm/engine.py¶

engine.py holds the canonical post-processing dataclasses — the form the DP machinery operates on:

• Regime — the canonical regime (distinct from the user-facing lcm.regime.Regime; in source files that import both we alias the user-facing one as UserRegime).

• StateActionSpace — pre-built state and action grids for a regime, with a state_action_space(params) method that fills in runtime-supplied grid points.

• SolveFunctions / SimulateFunctions — the compiled function bundles consumed by solve and simulate.

• Variables / VariableInfo — name + kind + topology metadata for every state and action in a regime.

• PeriodRegimeSimulationData — raw simulation output for one (regime, period) pair, before SimulationResult materialises a DataFrame.

The file name engine.py reflects what’s inside: the engine’s view of a model.

Build pipeline: model_processing.py and regime_building/¶

A user Regime is finalized at model build — model-level slots merged, broadcast variables pruned, the Koopmans aggregator and certainty equivalent injected, completeness validated — into the plain, complete Regimes exposed as model.user_regimes. The params template reads this user-vocabulary form, while process_regimes internally splits each regime into canonical per-phase slices and compiles the engine Regime.

_lcm/model_processing.py  ← top-level pipeline:
                            user regimes + params → canonical Model

_lcm/regime_building/
├── broadcast.py          ← model-level slot merge (exactly-one-level rule,
│                            `None` masking) + DAG-reachability pruning of
│                            broadcast states and actions
├── finalize.py           ← finalize_regimes: derived-categorical merge,
│                            Koopmans-aggregator and certainty-equivalent
│                            injection, completeness validation; output
│                            stays a plain lcm.regime.Regime
├── phases.py             ← normalize_regime_phases: expand every regime
│                            slot into per-phase RegimePhaseSpec slices
│                            (the Phased grammar boundary)
├── age_normalization.py  ← model-level normalization of age specialization
├── age_specialization.py ← per-age-specialized node resolution and
│                            grid-shape validation
├── canonicalize.py       ← canonicalize_regimes: rewrite every phase
│                            slice's laws and regime transition into the
│                            canonical target-granular form over exactly
│                            the reachable targets
├── processing.py         ← per-regime canonicalisation:
│                            UserRegime → engine.Regime
├── transitions.py        ← collect_state_transitions: walk user-supplied
│                            state_transitions into per-target callables
├── stochastic_state_transitions.py
│                         ← process-time AST + n_outcomes derivation for
│                            stochastic state transitions (raises
│                            InvalidStateTransitionProbabilitiesError on
│                            subscript-order mismatches)
├── Q_and_F.py            ← build (Q, F) closure for solve / simulate;
│                            also resolves the utility/feasibility DAG,
│                            whose two targets share one upstream chain
├── argmax.py             ← argmax helpers over action grids
├── max_Q_over_a.py       ← argmax / max over action grids
├── V.py                  ← value-function interpolation info
├── w_dag.py              ← user-DAG resolution for the Koopmans
│                            aggregator's *extra* params — those beyond
│                            utility and CE, which the Bellman step wires
│                            directly
├── next_state.py         ← compose per-state transitions into a single
│                            next_state function for simulation
├── ndimage.py            ← map-coordinates wrapper for continuous interp
└── diagnostics.py        ← cold-path machinery invoked by validate_V to
                            pinpoint *which* intermediate produced a NaN

The two-step name (model_processing at the model level, regime_building per regime) reflects what each layer actually does — the top level merges regimes and resolves fixed params; each regime is then canonicalised independently.

The numerical checks fired at solve / simulate time live outside regime_building/:

• regime_building/stochastic_state_transitions.py runs at Model(...) construction time and can fail the build before any params are involved. It catches malformed user functions (e.g., probs_array[health, age] where the signature is (age, health)) via AST analysis. Always on, never gated.

• _lcm/transition_checks.py runs from Model.solve() / Model.simulate() before backward induction starts. It evaluates the regime and state transition functions on the regime’s grid Cartesian product and verifies output shape, [0, 1] range, and sum-to-1. State checks are gated by log_level != "off" because the Cartesian product can blow up on models with many continuous-grid-dependent stochastic states.

• _lcm/solution/validate_V.py runs during backward induction (after each period in backward_induction.py, and once on the V handed to simulate.py). On NaN it invokes the diagnostic-intermediates closure built in regime_building/diagnostics.py to pinpoint which intermediate (U, F, E[V], Q) produced the NaN.

Reachability: _lcm/reachability.py¶

build_model_reachability builds the model’s static solve and simulate graphs once, at model construction, from the single canonical active_periods_by_regime mapping (regime_building.processing.compute_active_periods_by_regime) and the declared regime transitions. There is no runtime topology pass — the graph never changes after construction, and no runtime probability value narrows or widens it.

Every retained edge is EdgeStatus.CONDITIONAL; there is no TRUE status, because no declaration form (not even a per-target dict with one key) proves unconditional positive probability independently of state, action, and free runtime parameters. A coarse (bare callable / bare MarkovTransition) regime transition is therefore conservative: it retains an edge to every regime active in the next period, and every such edge’s state handoff is checked at model build — a carried state, a deterministic/stochastic law, or an explicit target-local/entry law must supply each target state’s next-period value. A per-target dict narrows support to its declared key set instead.

The solve and simulate phases build independent graphs (ModelReachability.solution / .simulation), because a regime transition’s Phased sides can differ between them — so the two graphs may retain different edges for the same source period.

Solver and simulation runtime code (_lcm/solution/, _lcm/simulation/) consume this graph — PhaseReachability.targets, .union_targets, .edge_status, ... — but never infers reachability itself: no runtime module calls an activity predicate, inspects a declared transition’s raw mapping keys, or derives continuation-target membership from state-law-bundle keys. regime_building/processing.py and diagnostics.py read the graph to decide which targets a period’s Q_and_F (or diagnostic) closure needs to build; they do not re-derive it.

Solve and simulate¶

_lcm/solution/
├── backward_induction.py  ← backward induction loop:
│                             V[T], V[T-1], ..., V[0] via max_Q_over_a
└── validate_V.py          ← per-period NaN / Inf validation

_lcm/simulation/
├── simulate.py         ← forward sampling loop with state-action draws
├── compile.py          ← compiled-function assembly for the simulate phase
├── random.py           ← PRNG-key handling for the sampling draws
├── transitions.py      ← per-state transition composition for simulation
└── initial_conditions.py
                        ← canonicalize / validate the user's
                          initial_conditions kwarg

These are the JAX-traced hot paths. The DP and sampling logic is the only thing here; everything that constructs the inputs (parameters, grids, transitions, compiled callables) lives in regime_building/ and is read out of the canonical Regime instances.

The solver seam: keys and routes¶

Two declarations connect a solver to the engine without either reading the other’s concrete types.

A continuation is keyed, not concrete. _lcm/continuation.py defines ContinuationSpec, pairing an all-finite template with the ArtifactKey under which its kernels publish it; EGMContinuationSpec is the EGM family’s specialization, adding the layout a reading parent needs. A solver declares what it reads as Solver.required_continuation_keys, and process_regimes matches every declared key against what each reachable target publishes before anything compiles, so a version mismatch is a build error naming both regimes rather than a failure inside the first rolled period. ContinuationPayload is the ContinuationArtifact protocol — one property, artifact_key — so backward induction stores and rolls a payload of any type. The shipped EGM family still exchanges concrete EGMCarry fields between its own producers and readers; the transition ledger records what closing that gap requires.

A replay route is declared, not discovered. Every canonical regime answers simulation.replay_route with one object carrying a ReplayMode, the exact payload class it retains, whether a solve owes one, and the reader that consumes it. EGMPolicyRead and NNBEGMPolicyRead are the built-in routes for the EGM and nested NB-EGM families; a regime that retains nothing declares grid recomputation. An external solver instead returns an ExecutableReplayRoute from SolutionKernels. Its stable plugin and route identities, period-specific artifact requirements, per-cell artifact authorities, mathematical validator, and JAX-transformable reader all use types re-exported by lcm.solvers.

model_authority.py builds authority from the canonical model and consuming route. A restored or caller-supplied result is canonicalized once; required lazy entries are materialized and checksum checked; then pylcm validation, plugin validation, reader construction, and the forward loop consume that same immutable snapshot. Descriptors transported in SolutionMetadata remain descriptive and cannot authenticate their own payloads.

Solution identity and persistence¶

Model.solve() returns one SolutionResult containing a ValueStore, addressed artifact stores, descriptive metadata, and explicit omission reasons. In-memory entries are loaded; entries restored by load_solution are independently lazy. Loading one period/regime value or artifact does not project or materialize its siblings, and an unloaded entry is present rather than omitted.

The model fingerprint in SolutionMetadata is a deterministic digest of the canonical mathematical declarations and parameters needed to interpret a result. It includes grid support, category order, solver/replay/artifact identities, and callable semantics while excluding device, compiler, tiling, sharding, JIT, and other execution-only choices. An in-memory result additionally carries a process-local model-instance guard; a restored archive is accepted by a separately constructed compatible model through the durable fingerprint.

_lcm/persistence/solution.py writes one HDF5 archive through an atomic sibling file. Its manifest is JSON; every value or artifact leaf is a separately addressed numerical dataset whose checksum binds its logical address, shape, dtype, and bytes. The archive serializes no Python implementation. load_solution checks the manifest and exact format/solution/solver-interface versions, then returns lazy handles that reopen the archive and verify a leaf before caching it. A whole-archive checksum pass deliberately does not change load state. Plugin PyTrees are reconstructed only from an installed route’s model-authoritative template.

Persistence is a per-artifact decision. ArtifactDescriptor.persistence is MODEL_VERIFIABLE only when another model process can independently rebuild the corresponding ArtifactAuthority; otherwise saving records NOT_PERSISTED and omits the payload. This keeps adaptive solve-generated coordinates out of the trust root. The same artifact identities drive computation: REPLAY programs are value-producing alternatives and ARTIFACT programs are additive, while each names its exact CoreProgram.retained_artifact_keys and the exact final KernelOutput payload type for every retained key. Every producer of one key must agree, including programs republished through a composite kernel; every replay program also names the exact values-only program it replaces. The engine selects them independently for every period/regime cell from model-authoritative ArtifactRef values, so requesting one persistable auxiliary or replay artifact does not suppress an unrelated values program.

Params: boundary form vs. canonical form¶

_lcm/params/
├── processing.py       ← cast_params_to_canonical_dtypes:
│                          User-supplied dicts (with int/float/np.array
│                          leaves) → flat MappingProxyType keyed by
│                          qualified names with JAX-array leaves.
├── regime_template.py  ← per-regime template construction:
│                          inspect Regime functions to derive what
│                          parameters they need.
├── mapping_leaf.py     ← UserMappingLeaf / MappingLeaf — wrapper that
│                          carries an immutable dict through a JAX pytree
│                          without becoming a Mapping itself.
└── sequence_leaf.py    ← UserSequenceLeaf / SequenceLeaf — same for
                           sequences.

The public lcm/params.py module exposes as_leaf and re-exports the four leaf classes; their definitions and the engine params machinery live in _lcm/params/.

Two leaf types exist because params dicts can contain heterogeneous leaves (scalars, arrays, named tuples, etc.). Wrapping them in MappingLeaf / SequenceLeaf lets the pytree machinery treat them as opaque leaves rather than walking into them — important when a “leaf” is itself a dict mapping named arguments to JAX arrays.

The User* types accept the wide boundary form (int, float, np.ndarray, pd.Series, etc.). After cast_params_to_canonical_dtypes runs, only canonical JAX-array leaves and canonical-narrow MappingLeaf / SequenceLeaf instances survive. The downstream solve / simulate code only ever sees the canonical form.

Pandas bridge: _lcm/pandas_utils.py¶

A single module for converting between user-friendly pd.Series / pd.DataFrame representations and the JAX arrays the engine expects. array_from_series is the workhorse: it inspects the function source via AST helpers (in utils/ast_inspection.py) to determine the expected multi-index order, then materialises a properly-shaped JAX array.

This file gets used both at params processing (for any pd.Series leaves in user params) and at simulation output (for building the result DataFrame).

Utilities: _lcm/utils/¶

Small, dependency-light helpers grouped by topic:

• ast_inspection.py — Parse a function body to find probs_array[a, b] subscript patterns. Used by the static AST check and by pandas_utils.

• containers.py — ensure_containers_are_immutable, first_non_none, invert_regime_ids.

• dispatchers.py — productmap, vmap_1d, simulation_spacemap. See Dispatchers.

• error_messages.py — format_messages, which collapses a list of validation errors into a single string.

• functools.py — all_as_kwargs, get_union_of_args.

• logging.py — get_logger, format_duration, log-formatting helpers.

• namespace.py — flatten_regime_namespace / unflatten_regime_namespace for the qualified-name pytree keys.

Type aliases: lcm/typing.py vs _lcm/typing.py¶

lcm/typing.py    ← user-facing aliases: jaxtyping array shapes (FloatND,
                   ScalarInt, ...), Period, Age, and the User* boundary
                   aliases (UserParams, UserInitialConditions, ...)
_lcm/typing.py   ← engine-side aliases and protocols: string labels
                   (RegimeName, StateName, ...), compound mapping
                   aliases, canonical post-processing forms (Params,
                   InitialConditions, ...), and the structural Protocol
                   classes (EconFunction, TransitionFunction, ...)

The split mirrors the public / private package boundary. lcm/typing.py holds the aliases a user needs to annotate model functions and the User* aliases that accept wide boundary types; it imports nothing from _lcm. _lcm/typing.py holds the engine-internal aliases — including the post-canonicalization forms (Params, InitialConditions) — and builds on the public aliases it imports from lcm.typing.

Exceptions: lcm/exceptions.py¶

Every project-specific exception class lives here, all inheriting from PyLCMError. They split into two categories:

• Initialization errors — raised at Model(...) / Regime(...) / LinSpacedGrid(...) construction time. These map beartype violations on user-facing constructors to a project-typed error (so users see e.g. ModelInitializationError, not a BeartypeCallHintViolation).

• Runtime errors — InvalidValueFunctionError, InvalidRegimeTransitionProbabilitiesError, InvalidStateTransitionProbabilitiesError, InvalidParamsError, InvalidInitialConditionsError, SolutionIntegrityError, and IncompatibleSolutionError. These fire from transition/value checks, restored-result preflight, or the solution-archive reader during solve / simulate / load.

The exception classes are public — both from lcm.exceptions import InvalidParamsError and except lcm.InvalidParamsError work. format_messages, the helper that assembles a list of validation errors into one string, is internal validation plumbing and lives in _lcm/utils/error_messages.py.

Bootstrap modules¶

A few _lcm/ modules exist for ordering reasons rather than for any conceptual grouping:

• jaxtyping_patch.py — Bootstrap patch that has to run before any jaxtyping-annotated type is created. _lcm/__init__.py applies it as its first statement.

• beartype_conf.py — Holds the beartype configurations used in the package (the internal-claw conf + the user-facing constructor-decorator confs).

• config.py — Build-time configuration constants (paths to test data, etc.).

• dtypes.py — Canonical-dtype resolution (canonical_float_dtype()), which depends on the JAX x64 setting.

Reading order for new contributors¶

If you’re reading the codebase for the first time, the path of least confusion is:

1. lcm/regime.py to see what users supply.

2. lcm/model.py to see what Model.__init__ triggers.

3. _lcm/model_processing.py for the top-level pipeline.

4. _lcm/regime_building/processing.py for per-regime canonicalisation — the longest single file and the heart of the build.

5. _lcm/engine.py for the canonical dataclasses the DP machinery consumes.

6. _lcm/reachability.py for the static solve/simulate graphs solve and simulate consume but never infer.

7. _lcm/solution/backward_induction.py and _lcm/simulation/simulate.py for the actual DP and sampling.

By the time you reach (7), the canonical form should feel familiar and the JAX-traced code becomes easy to read.