Skip to article frontmatterSkip to article content
Site not loading correctly?

This may be due to an incorrect BASE_URL configuration. See the MyST Documentation for reference.

The NB-EGM solver

NBEGM is an endogenous-grid solver for a one-dimensional consumption–saving regime whose budget is split by declared institutional breakpoints — asset tests, subsidy brackets, benefit notches, consumption floors. The model author exposes each boundary as metadata (a case piece); within each piece the budget is smooth, so NBEGM runs EGM per piece and merges the pieces on the liquid grid with a branch-aware upper envelope, resolving the kinks and jumps exactly at their declared locations.

Reach for NBEGM when a regime carries institutional discontinuities. A dense brute grid (GridSearch) can only approximate the optimum near a cliff — unless you can specify your grid statically as described in PiecewiseLinSpacedGrid, it places no candidate exactly at the threshold and averages across it — so brute force is a diagnostic here, not the correctness reference (see Validating an NB-EGM regime). For the full theory, correctness results, and the conditions under which NB-EGM beats brute force, see the NB-EGM methods paper; this page is the how-to.

When it applies

NBEGM solves the sub-class where:

Declaring feasibility constraints

First decide whether the restriction itself needs retained structure. Use an ordinary callable for arbitrary executable logic that does not fit the named-comparison algebra:

import jax.numpy as jnp


def finite_positive_consumption(consumption, resources):
    return jnp.isfinite(consumption) & (consumption > 0.0) & (consumption <= resources)

Use lcm.ref when the restriction is a named comparison whose structure pylcm must retain. The borrowing limit of a consumption-saving problem is the central example:

import lcm

nonnegative_savings = lcm.ref("savings") >= 0.0

An NBEGM solve does not call an arbitrary feasibility predicate along its endogenous savings path. It can nevertheless prove savings >= 0 because its savings grid already enforces that exact lower bound. Prefer the specialized spelling when the bound belongs to the regime’s declared liquid margin:

from lcm.consumption_savings_regime import (
    LiquidMargin,
    post_decision_lower_bound,
)

liquid = LiquidMargin(
    state="liquid",
    action="consumption",
    resources="liquid",
    post_decision_state="savings",
)
borrowing_limit = post_decision_lower_bound(margin=liquid, lower=0.0)

The specialized declaration and lcm.ref("savings") >= 0.0 mean the same thing. The specialized form prevents the condition’s post-decision name from drifting away from the margin declaration. A general callable remains opaque; NBEGM accepts it only if a route can evaluate every required name and otherwise refuses the model before lowering. Condition itself grants no extra solver capability.

NBEGM does compile one useful retained shape: a conjunction of ordered comparisons between the current liquid state and either a literal or a flat parameter. For example:

asset_band = (lcm.ref("liquid") >= lcm.ref("minimum_assets")) & (
    lcm.ref("liquid") < 10.0
)

The operators determine exact boundary ownership. >= and <= include equality on the feasible side; > and < exclude it. During the solve, NBEGM turns these comparisons into an owner-aware liquid-axis partition and masks every value, policy, and marginal channel together. During simulation, pylcm evaluates the original condition on the simulated candidates.

This compiled route supports:

NNBEGM applies the same compiled subset independently inside its adjuster and keeper routes, so both outer candidates enforce the identical liquid feasibility geometry.

It requires jump_read="one_sided", because a parent must not interpolate from a finite value across an infeasible child region. NBEGM refuses a union, complement, implication, equality comparison, opaque callable, or a boundary that depends on an action, post-decision value, computed output, or another state. It also refuses this feasibility topology combined with a finite-value schedule jump or a binary case piece. Use GridSearch when the restriction falls outside the compiled subset and can be evaluated on complete state-action candidates.

See Writing a constraint: callable or Condition for named-to-literal and named-to-named comparisons, intersections, unions, complements, implications, and the cross-solver decision table.

Declaring breakpoints

Institutional boundaries are declared, not discovered — no solver can recover a threshold’s exact location from finitely many black-box evaluations, so the model author exposes each boundary as metadata. The decorators only attach metadata and return the function unchanged, so the same model still solves identically under GridSearch.

There are two declaration routes, and which one a boundary belongs to is not a matter of taste: case pieces describe a means test that hands two different formulas to the two sides, and piecewise-affine schedules describe a bracket structure — a tax, a taper, a floor — that one formula already spans. The case-piece route is deliberately narrow; everything else goes through a schedule.

Case pieces — one means-tested cliff

The case-piece kernels solve a single means-tested cliff on the liquid margin, so NBEGM refuses anything wider at model build. A case-piece regime must:

The "when" owner, the other two kinds, and a boundary on any other variable are accepted by lcm.boundary because a schedule uses them — they are rejected on the case-piece route.

import jax.numpy as jnp

import lcm
from lcm.typing import BoolND, ContinuousState, FloatND

# The kernels form cash-on-hand themselves, so the regime declares pylcm's own
# node rather than a local spelling of the same arithmetic.
resources = lcm.cash_on_hand_with_subsidy


@lcm.case_boundary(
    lcm.boundary(
        variable="liquid",
        threshold="medicaid_asset_limit",
        equality="otherwise",
        kind="jump",
    )
)
def medicaid_eligible(liquid: ContinuousState, medicaid_asset_limit: float) -> BoolND:
    """Medicaid asset test: eligible while liquid wealth is below the limit."""
    return liquid < medicaid_asset_limit


@lcm.piece(output="subsidy", when=medicaid_eligible)
def subsidy_medicaid(subsidy_high: float) -> FloatND:
    """Subsidy into market resources for the Medicaid-eligible (low-asset) case."""
    return jnp.asarray(subsidy_high)


@lcm.piece(output="subsidy", otherwise=medicaid_eligible)
def subsidy_private(subsidy_low: float) -> FloatND:
    """Subsidy into market resources for the private (high-asset) case."""
    return jnp.asarray(subsidy_low)

The Medicaid-eligible subsidy exceeds the private one, so market resources — and hence the value function — jump down as liquid wealth crosses the limit upward.

The budget node and the liquid law

The case-piece kernels do not call the regime’s budget node: they form cash-on-hand as liquid + subsidy themselves. No finite check establishes that an arbitrary callable computes that same thing — a global rescaling agrees at every sampled point and still moves every state’s value — so pylcm exports the declaration itself and the route accepts it by identity:

A budget node those kernels cannot form does not belong on the case-piece route. Declare a lcm.piecewise_affine schedule with a post_decision_function — which composes the budget from the DAG and reads whatever it declares — or solve the regime with GridSearch.

The liquid law carries no such restriction. NBEGM reads the law the regime declares — where each level of savings lands next period, and how that landing point moves when savings move — so a per-period fixed cost, a rescaled income, or a return compounded over sub-periods is part of the problem solved rather than structure dropped from it.

What the law may not do is depend on the consumption choice by any route other than post-decision savings. The Euler inversion runs on a grid of savings and reads the continuation off the landing points that grid reaches, so a law stated as next_liquid(resources, consumption, …) has no single continuation to read and is refused at build — even when it happens to depend on the difference alone. The regime therefore declares the node the law reads, named savings by default and otherwise named to NBEGM(post_decision_function=…):

from lcm.typing import ContinuousAction, FloatND


def savings(resources: FloatND, consumption: ContinuousAction) -> FloatND:
    """Post-decision savings: cash-on-hand net of consumption."""
    return resources - consumption

pylcm exports the conventional affine law in both forms, as ordinary executable functions; a regime is free to write its own instead:

Piecewise-affine schedules — brackets, tapers, and floors

A tax schedule, a benefit taper, or a consumption floor is one formula that is affine between its thresholds, so there is nothing to split into pieces — the model author declares where the thresholds are and what kind of discontinuity each carries.

import jax.numpy as jnp

import lcm
from lcm.typing import ContinuousState, FloatND


@lcm.piecewise_affine(
    output="tax",
    variable="liquid",
    breakpoints=(
        lcm.affine_breakpoint(threshold="tax_exemption", kind="continuous_kink"),
    ),
)
def tax(liquid: ContinuousState, tax_rate: float, tax_exemption: float) -> FloatND:
    """Continuous tax: zero below the exemption, `tax_rate` on the excess above."""
    return tax_rate * jnp.maximum(liquid - tax_exemption, 0.0)


def resources(liquid: ContinuousState, tax: FloatND, base_income: float) -> FloatND:
    """Cash-on-hand: liquid wealth plus base income, net of the tax."""
    return liquid + base_income - tax

The thresholds are ordinary parameters, so an estimator moves them — but the step’s case structure is fixed at build time from the declared order. Declare a schedule’s breakpoints ascending in value, and keep every parameter draw in that order; a mixed-kind schedule whose thresholds arrive reordered is refused rather than solved.

Two further limits apply to a regime with no ride-along co-state: it may declare only one schedule (a second schedule’s thresholds would not enter the interval partition), and it may not combine a "hard_constraint" with a "jump" in the same schedule.

Selecting the solver

The solver is a per-regime slot. Pass an NBEGM instance where you would otherwise leave the default GridSearch:

from lcm import LinSpacedGrid
from lcm.consumption_savings_regime import (
    ConsumptionSavingsRegime,
    LiquidMargin,
    post_decision_lower_bound,
)
from lcm.solvers import NBEGM

liquid = LiquidMargin(
    state="liquid",
    action="consumption",
    resources="liquid",
    post_decision_state="savings",
)

alive_regime = ConsumptionSavingsRegime(
    transition=next_regime,
    states={"liquid": LinSpacedGrid(start=0.0, stop=20.0, n_points=80)},
    actions={"consumption": LinSpacedGrid(start=0.0, stop=20.0, n_points=80)},
    state_transitions={"liquid": next_liquid},
    functions={
        "utility": utility,
        "medicaid_eligible": medicaid_eligible,
        "subsidy_medicaid": subsidy_medicaid,
        "subsidy_private": subsidy_private,
        "savings": savings,
    },
    constraints={
        "borrowing_limit": post_decision_lower_bound(
            margin=liquid,
            lower=0.0,
        )
    },
    solver=NBEGM(savings_grid=LinSpacedGrid(start=0.0, stop=20.0, n_points=100)),
    liquid=liquid,
)

NBEGM requires a savings_grid (the post-decision savings nodes). Key optional arguments:

The two cliff-read modes

A child regime’s value cliffs cannot be represented by a single continuous interpolant. jump_read selects how the continuation carry is published to parents:

The smoothness gate

Declared breakpoints are only as good as the smoothness of what lies between them. At model build, NBEGM.validate runs two validators over the user economic functions reachable in each case:

Mark a reviewed numerical clip/max/abs guard with @lcm.smooth_helper to exempt it, stating the domain on which it is smooth.

The piecewise-constancy probe

When the continuation reads the current liquid state (a co-state’s next-state law or a regime-transition probability switched at a declared threshold), NBEGM solves one continuation row per declared interval. This is exact only if that liquid dependence is piecewise-constant on the declared partition. A probe screens for it and refuses by default (probe_failure="reject") when it detects smooth dependence or cannot differentiate the model’s DAG. Passing probe_failure="assume_declared" asserts the precondition explicitly (emitting a warning); every exactness claim is then conditional on that assertion, which must be discharged by independent validation.

The probe — and the affine-budget probe alongside it — runs on the first solve, not at model build, because it needs parameter values: a budget reading tax schedules or a law reading an interpolation table cannot be differentiated until those are supplied. Running it there is also what makes it meaningful, since the real bracket structure is what gets differentiated. It runs once per model: what it tests is a property of the model’s functional structure, so an estimation loop pays for it on its first iteration only.

The probe is a finite diagnostic, not a certificate — a dependence whose derivative vanishes at every probed point passes undetected.

Performance knobs

Every batching knob streams one axis and changes peak memory and schedule only, never the result (up to floating-point reassociation):

KnobAxis streamed
cell_block_sizeride-along cells
branch_batch_sizediscrete-action branches (lax.map, body compiled once)
interval_batch_sizeper-interval continuation reads
stochastic_node_batch_sizechild stochastic-node mesh
envelope_segment_block_sizeenvelope segment blocks (two-pass scan)

All default to 0 (the whole axis in one vectorized pass). Raise a knob when the corresponding buffer is the memory wall.

Relation to other EGM methods

NB-EGM sits alongside pylcm’s other endogenous-grid solvers rather than replacing them (see Choosing a solver for the full map):

Validating an NB-EGM regime

GridSearch is a diagnostic, not the correctness oracle. Brute force evaluates the combined (jnp.where) budget on a finite action grid, so it smooths across every breakpoint. The optimal policy is often to save to just below an eligibility cliff (to keep the benefit), so the optimum sits one step inside the eligible side — a point no finite grid holds exactly, though a PiecewiseLinSpacedGrid aligned to the breakpoint may come close enough for practical purposes. Asserting exact agreement with brute is therefore the wrong acceptance test.

Euler residuals are a useful report but not the acceptance criterion — they are blind to the corner and boundary candidates that carry the economics of interest.

pylcm
Choosing a solver
pylcm
Performance and Memory Tuning