Redesign the solver seam as an ABC (GridSearch + DCEGM stub) by hmgaudecker · Pull Request #388 · OpenSourceEconomics/pylcm

hmgaudecker · 2026-06-17T15:15:42Z

Stacked on #389 (feat/target-batch-eviction).

Behavior-preserving refactor that introduces a per-regime solver-selection
seam as a polymorphic Solver ABC, so the upcoming DC-EGM PR is a pure
addition whose review can concentrate on the DC-EGM logic alone.

What this does

Solver(ABC) (_lcm/solution/contract.py): an abstract
build_period_kernels(*, context) -> SolverKernels plus a default no-op
validate(*, context). The engine dispatches polymorphically on the solver
instance — solver.validate(context) then
solver.build_period_kernels(context) — with no switch on solver type and no
BruteForce | DCEGM union. contract.py is a true engine leaf (no import
cycle: import lcm resolves cleanly).
GridSearch(Solver) (the renamed BruteForce) and DCEGM(Solver) live in
_lcm/solution/solvers.py, with validate + build_period_kernels methods
and deferred numerical imports; lcm/solvers.py is a thin re-export façade.
The type(solver)-keyed SOLVER_KERNEL_BUILDERS registry and registry.py
are deleted.
Regime.solver: Solver (default GridSearch()), exported from lcm
alongside the DCEGM configuration class.
DCEGM is published as the final configuration surface (fields + field
validation), but its engine is not yet wired in: a regime requesting it is
rejected at model build with a clear NotImplementedError.

Not in scope (the DC-EGM PR, stacked on top of this)

The _lcm/egm/* engine, build-time DC-EGM validation, the SolutionPhase EGM
fields, the EGM-specific SolverBuildContext/SolverKernels fields, and the
DC-EGM kernel builder. solve_brute.py and the simulation phase are unchanged.

Verification

Numerics unchanged: routing the existing grid search through the ABC contract
is a faithful move; solver=GridSearch() explicitly ≡ the default value
function.
Full gpu-01 suite green (1208 passed, 15 skipped). tests/test_solvers.py
asserts the default solver, the polymorphic-dispatch parity, DCEGM
construction + field validation, and the build-time guard.
ty and prek clean.

🤖 Generated with Claude Code

Route the existing brute-force grid search through a per-regime solver configuration and a builder registry, with no change to the numerics. - `Regime.solver: BruteForce | DCEGM` (default `BruteForce()`), exported from `lcm` alongside the `DCEGM` configuration class. - `_lcm.solution.registry`: `SolverBuildContext`, `SolverKernels`, the `SolverKernelBuilder` protocol, and `_build_brute_force_kernels` (the former `_build_max_Q_over_a_per_period`), dispatched on `type(regime.solver)` via `SOLVER_KERNEL_BUILDERS`. - `DCEGM` is published as the final configuration surface (fields + field validation), but its engine is not yet wired in; a regime requesting it is rejected at model build with `NotImplementedError`. Behavior-preserving: the full test suite passes unchanged and an explicit `BruteForce()` yields the same value function as the default. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

The full discrete-continuous endogenous grid method, re-rooted as a single commit on top of the solver-selection seam (PR #388). The tree reproduces the reviewed feat/dcegm tip exactly, minus two root-level audit scratch files. On top of the seam this adds the `_lcm.egm` engine, DC-EGM build-time validation, the `_build_dcegm_kernels` builder and the EGM-specific SolverBuildContext / SolverKernels / SolutionPhase fields, and replaces the seam's build-time `NotImplementedError` guard with the real solver. To be split into reviewable sub-PRs at review time. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

read-the-docs-community · 2026-06-17T15:20:16Z

Documentation build overview

📚 pylcm | 🛠️ Build #33237115 | 📁 Comparing d903f83 against latest (dc59124)

🔍 Preview build

37 files changed · ± 37 modified

± Modified

github-actions · 2026-06-17T16:27:22Z

Benchmark comparison (main → HEAD)

Comparing dc591241 (main) → d903f83a (HEAD)

Benchmark	Statistic	before	after	Ratio	Alert
aca-baseline	execution time	15.217 s	13.419 s	0.88
	peak GPU mem	581 MB	588 MB	1.01
	compilation time	338.71 s	373.51 s	1.10	❌
	peak CPU mem	6.80 GB	7.87 GB	1.16	❌
aca-baseline-debug	execution time	80.285 s	88.834 s	1.11	❌
	peak GPU mem	581 MB	587 MB	1.01
	compilation time	450.65 s	567.53 s	1.26	❌
	peak CPU mem	7.89 GB	8.19 GB	1.04
Mahler-Yum	execution time	4.487 s	4.839 s	1.08
	peak GPU mem	520 MB	520 MB	1.00
	compilation time	11.22 s	11.40 s	1.02
	peak CPU mem	1.58 GB	1.58 GB	1.00
Precautionary Savings - Solve	execution time	24.9 ms	23.0 ms	0.92
	peak GPU mem	8 MB	8 MB	1.00
	compilation time	1.59 s	1.83 s	1.15	❌
	peak CPU mem	1.16 GB	1.16 GB	1.00
Precautionary Savings - Simulate	execution time	62.4 ms	63.7 ms	1.02
	peak GPU mem	157 MB	157 MB	1.00
	compilation time	3.61 s	3.79 s	1.05
	peak CPU mem	1.33 GB	1.33 GB	1.00
Precautionary Savings - Solve & Simulate	execution time	101.4 ms	94.7 ms	0.93
	peak GPU mem	566 MB	566 MB	1.00
	compilation time	4.78 s	4.98 s	1.04
	peak CPU mem	1.31 GB	1.30 GB	1.00
Precautionary Savings - Solve & Simulate (irreg)	execution time	202.7 ms	207.7 ms	1.02
	peak GPU mem	2.18 GB	2.18 GB	1.00
	compilation time	5.03 s	5.11 s	1.02
	peak CPU mem	1.36 GB	1.37 GB	1.00
IskhakovEtAl2017Simulate	execution time	194.1 ms	220.7 ms	1.14	❌
	compilation time	4.19 s	4.43 s	1.06
	peak CPU mem	1.30 GB	1.29 GB	1.00
IskhakovEtAl2017Solve	execution time	44.3 ms	47.0 ms	1.06
	compilation time	0.66 s	0.80 s	1.20	❌
	peak CPU mem	1.15 GB	1.15 GB	0.99
IskhakovEtAl2017SimulateGpuPeakMem	peak GPU mem	281 MB	281 MB	1.00
IskhakovEtAl2017SolveGpuPeakMem	peak GPU mem	67 MB	67 MB	1.00

Points the benchmark feature's aca-model at feat/dcegm-solver's tip — main's audit fixes + the DC-EGM solver + smooth-share eligibility — the version the solver-seam work is developed against. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

timmens

I really like the idea of making the backward induction solver a replaceable component.

I am not a huge fan of the current implementation, but I could be convinced otherwise. Some things that I find odd:

To me a Solver abstraction is the perfect use-case for inheritance, or at least a Protocol interface. I have looked for a single interface / abstraction in the PR that tells me what the "solver contract" is. Having a clear abstraction should also make it simple for users to implement custom solvers (thinking of optimagic's algorithm interface here).
The tests don't seem to test anything interesting again.
I actually don't see why we need a DCEGM stub in this PR. I would expect the first PR to design a solver interface / abstraction, for which we then implement the GridSearch(Solver) and only check that this works. Then in the next PR we add DCEGM.

Replace the `type(solver)`-keyed builder registry with a polymorphic `Solver` ABC. The engine now calls `solver.validate(context)` then `solver.build_period_kernels(context)` — no `SOLVER_KERNEL_BUILDERS` dict, no `BruteForce | DCEGM` union, no standalone DC-EGM guard. - `_lcm/solution/contract.py` (new): the `Solver` ABC (abstract `build_period_kernels`, default no-op `validate`), `SolverBuildContext`, and `SolverKernels`. An engine leaf — imports nothing that reaches `lcm.solvers`, so the façade can re-export it without an import cycle. - `_lcm/solution/solvers.py` (new): `GridSearch(Solver)` (the relocated grid-search builder, with function-local `jax`/`get_max_Q_over_a` imports) and `DCEGM(Solver)` (the published config; `validate` raises the not-yet-available guard, so a regime requesting it is rejected at model build). - `lcm/solvers.py` → thin re-export façade; `registry.py` deleted; the `processing` dispatch and `Regime.solver` field updated. - Rename the default solver `BruteForce` → `GridSearch` (more descriptive; alpha permits the break). Faithful to dcegm-solver-seam-abc-design.md (ABC, not Protocol). Layer 2 (generic KernelResult) does not apply here — the stub seam has no EGM fork. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

mj023

I also like the object-oriented approach more. I think the stub is fine, the PR is small enough for this. The interface is a bit limiting as I can imagine that some solvers (Reinforcement learning) would need different aspects of the Regime, like the transition functions, or would return something else then the max_Q_over_A functions. But the interface is not complicated, so I would think we can cahnge this if we ever want to implement such a solver.

The full discrete-continuous endogenous grid method, re-rooted as a single commit on top of the solver-selection seam (PR #388). The tree reproduces the reviewed feat/dcegm tip exactly, minus two root-level audit scratch files. On top of the seam this adds the `_lcm.egm` engine, DC-EGM build-time validation, the `_build_dcegm_kernels` builder and the EGM-specific SolverBuildContext / SolverKernels / SolutionPhase fields, and replaces the seam's build-time `NotImplementedError` guard with the real solver. To be split into reviewable sub-PRs at review time. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

The solve module's job is backward induction; name it for that. Renames src/_lcm/solution/solve_brute.py → backward_induction.py and tests/solution/test_solve_brute.py → test_backward_induction.py, updates every import (model.py, simulation/compile.py, test_beartype_claw.py), the two test-function names, and the doc/docstring references. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

A module-level jnp.array constant triggers the PREALLOCATE=true pool reservation on device 0 at import time — before any solve, and even in processes that only import the model code to schedule work. A second such process then OOMs on device 0 at startup. Document the host-array (numpy) pattern for module-level constants in the tuning guide's memory section. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

hmgaudecker requested review from mj023 and timmens June 17, 2026 17:31

timmens reviewed Jun 18, 2026

View reviewed changes

hmgaudecker force-pushed the feat/solver-seam branch from 44e8522 to e0fecdd Compare June 19, 2026 14:00

hmgaudecker changed the title ~~Introduce the solver-selection seam (BruteForce dispatch + DCEGM stub)~~ Redesign the solver seam as an ABC (GridSearch + DCEGM stub) Jun 19, 2026

hmgaudecker changed the base branch from main to feat/target-batch-eviction June 19, 2026 14:00

mj023 approved these changes Jun 19, 2026

View reviewed changes

hmgaudecker mentioned this pull request Jun 19, 2026

Add the endogenous-grid solver family: EGM, DC-EGM, and NEGM #390

Merged

hmgaudecker changed the base branch from feat/target-batch-eviction to main June 21, 2026 15:23

hmgaudecker force-pushed the feat/solver-seam branch from e0fecdd to 44e8522 Compare June 21, 2026 15:29

hmgaudecker and others added 2 commits June 21, 2026 18:48

hmgaudecker merged commit 3e29f91 into main Jun 21, 2026
3 checks passed

hmgaudecker deleted the feat/solver-seam branch June 21, 2026 16:57

Navigation Menu

Redesign the solver seam as an ABC (GridSearch + DCEGM stub) - #388

Redesign the solver seam as an ABC (GridSearch + DCEGM stub)#388
hmgaudecker merged 5 commits into
mainfrom
feat/solver-seam

hmgaudecker commented Jun 17, 2026 •

edited

Loading

Uh oh!

read-the-docs-community Bot commented Jun 17, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Jun 17, 2026 •

edited

Loading

Uh oh!

timmens left a comment

Uh oh!

mj023 left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

hmgaudecker commented Jun 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What this does

Not in scope (the DC-EGM PR, stacked on top of this)

Verification

Uh oh!

read-the-docs-community Bot commented Jun 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Documentation build overview

Uh oh!

github-actions Bot commented Jun 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Benchmark comparison (main → HEAD)

Uh oh!

timmens left a comment

Choose a reason for hiding this comment

Uh oh!

mj023 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

hmgaudecker commented Jun 17, 2026 •

edited

Loading

read-the-docs-community Bot commented Jun 17, 2026 •

edited

Loading

github-actions Bot commented Jun 17, 2026 •

edited

Loading