Partial `fixed_params` into functions upon model creation. Stateless treatment of dynamic `IrregSpacedGrids`, in particular `ShockGrids`. by hmgaudecker · Pull Request #235 · OpenSourceEconomics/pylcm

hmgaudecker · 2026-02-10T12:17:50Z

Ended up a bit of a monster, sorry...

This allows passing any parameter at model creation via fixed_params -> will be partialled into functions.
Strongly-typed shock grids, inheritance diagram below.
Update the IrregSpacedGrids / ShockGrids so they do the same thing: For both IrregSpacedGrids and ShockGrids, we can pass params directly to the instantiation, via fixed_params, or at runtime.
Using terminology "fixed grid" / "points specified at runtime" instead of "static" / "dynamic" (too confusing in a library for dynamic economic models...). Also using mean / std (/ ar1_coeff for
Remove defaults from grid parameters (better be explicit)
There should be a clear, stateless flow now for setting up a Model instance.

classDiagram
    ContinuousGrid <|-- ShockGrid
    ShockGrid <|-- ShockGridIID
    ShockGrid <|-- ShockGridAR1
    ShockGridIID <|-- ShockGridIIDUniform
    ShockGridIID <|-- ShockGridIIDNormal
    ShockGridAR1 <|-- ShockGridAR1Tauchen
    ShockGridAR1 <|-- ShockGridAR1Rouwenhorst

    class ShockGrid {
        +n_points: int
        +params
        +params_to_pass_at_runtime
        +is_fully_specified
        +compute_gridpoints()
        +compute_transition_probs()
        +get_gridpoints()
        +get_transition_probs()
        +to_jax()
        +get_coordinate()
    }

    class ShockGridIID {
        +draw_shock(params, key)
    }

    class ShockGridAR1 {
        +draw_shock(params, key, current_value)
    }

    class ShockGridIIDUniform {
        +start: float
        +stop: float
    }

    class ShockGridIIDNormal {
        +mean: float
        +std: float
        +n_std: float
    }

    class ShockGridAR1Tauchen {
        +ar1_coeff: float
        +std: float
        +mean: float
        +n_std: float
    }

    class ShockGridAR1Rouwenhorst {
        +ar1_coeff: float
        +std: float
        +mean: float
    }

Instead of wrapping each function to accept an internal_regime_params dict and extracting parameters at call time, use dags.signature.rename_arguments() to qualify parameter names with function prefixes (e.g., risk_aversion becomes utility__risk_aversion). This makes the parameter flow explicit and removes the internal_regime_params indirection from all function signatures. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

…at_regime_params. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

…ParamsTemplate. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

…s Callable in result.py. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

…mstances than REGIME_SEPARATOR. Some small improvements.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

…eatment of dynamic IrregSpacedGrids, in particular ShockGrids.

Three bugs in _partial_fixed_params_into_regimes: 1. internal_functions.regime_transition_probs was not updated alongside the top-level attribute, causing simulation to use un-partialled functions for regime transitions. 2. regime_transition_probs was partialled with ALL regime fixed params, but it only accepts a subset — breaking inspect.signature used by dags.concatenate_functions in to_dataframe. 3. to_dataframe additional_targets built DAGs from raw functions that still expected fixed params no longer present in runtime params. Fixes: update internal_functions, filter kwargs per function signature, store resolved_fixed_params on InternalRegime for target computation. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

for more information, see https://pre-commit.ci

…regation

…speed on runs 2 and especially 3 :-)

hmgaudecker

Autoreview.

hmgaudecker · 2026-02-13T13:54:26Z

examples/mahler_yum_2024/model.py

-        initial_regimes=initial_regimes,
-        seed=8295,
-    )
+    for _ in range(3):


Commit msg: Run the MY model three times just for the pleasure of looking at the speed on runs 2 and 3 :-)

Actually, now that I added timings, that is not the only reason. Compile time reduces a lot when passing params directly:

prod_shock_grid = ShockGridAR1Rouwenhorst(n_points=5, ar1_coeff=rho, mean=0, std=1) ... "adjustment_cost": ShockGridIIDUniform(n_points=5, start=0, stop=1)

INFO:lcm:--- Timing summary --- INFO:lcm: Run 1: 84.696s INFO:lcm: Run 2: 7.181s INFO:lcm: Run 3: 7.132s

vs. via fixed_params:

fixed_params={ "alive": { "productivity_shock": {"ar1_coeff": rho, "mean": 0, "std": 1}, "adjustment_cost": {"start": 0, "stop": 1}, } },

INFO:lcm:--- Timing summary --- INFO:lcm: Run 1: 117.166s INFO:lcm: Run 2: 7.301s INFO:lcm: Run 3: 7.331s

whether the runtime is truly affected or my machine was doing more other stuff, I cannot say.

hmgaudecker · 2026-02-13T17:43:20Z

src/lcm/shocks.py


 def _discretized_uniform_distribution_gridpoints(
-    n_points: int, start: float = 0, stop: float = 1
+    n_points: int, start: float, stop: float


Note: I prefer require setting all parameters.

hmgaudecker · 2026-02-14T07:54:47Z

@mj023: Not asking for a full review, but it would be great if you could take a look at the shocks code!

mj023 · 2026-02-14T15:19:56Z

src/lcm/shock_grids.py

+    r"""AR(1) shock discretized via Tauchen (1986).
+
+    The process is
+    :math:`y_t = \mu_\varepsilon + \rho \, (y_{t-1} - \mu_\varepsilon) + \varepsilon_t`,


This AR1 Process is a bit odd. The function that draws the shock is $\rho Y_{t-1} - \rho \mu_{\epsilon} + \epsilon_t$. Which is not what is stated here, but also generally not what I would expect. The first $\mu_{\epsilon}$ probably needs to be new parameter if we want something that just shifts the mean.

Oops, will have to look into that again. Sorry, this was a very late addition (since it is now ready to support that) and I did not have time to check it well. Should be very standard...

Sorry for the noise. We are now using the form

$$y_{t+1} = μ + ρ y_{t} + ϵ_{t+1}$$

for AR(1)-processes throughout, which is the standard in econ. It matches what, e.g., QuantEcon does. Sorry again, I had thought this would be so standard that I did not check at all -- but there are different representations around!

src/lcm/shock_grids.py

…rocesses throughout. Stick to rho - mu - sigma for all shock parameters.

…ecture-dependent.

timmens

Very nice PR! 🚀

I have a few comments before full approval, but nothing serious. In addition, I have two comments here:

I think this PR is missing an economically sensible test for the shocks. A test where we predict (some part of) the result of the solution or simulation given the chosen shock. Currently there is no end-to-end test like this. Only the regression test, but that does not test a behavior.
At some point I believe we must think about a clearer separation of code/objects that are used for the solution and code/objects used for the simulation. Definitely not in this PR -- was just reminded of that while reading the changes.

timmens · 2026-02-14T13:34:10Z

examples/mahler_yum_2024/model.py

 from lcm import AgeGrid, DiscreteGrid, LinSpacedGrid, Model, Regime, categorical
 from lcm.dispatchers import _base_productmap
-from lcm.grids import ShockGrid
+from lcm.shock_grids import ShockGridAR1Rouwenhorst, ShockGridIIDUniform


For a more pleasant user-interface, I would leave out the "ShockGrid" prefix. Potentially, I would even rename the shock_grids module to shocks.

I can also see the argument that you want to make clear to the user that they are describing grids. Personally, I still prefer the minimalistic UI (without prefix), which focuses on the actual object they are trying to model (the shock) and not the implementation (the grid).

We could even go one step further and export a shocks object that behaves as follows:

uniform_shock = lcm.shocks.IID.Uniform(start=0, stop=1) rouwenhorst_shock = lcm.shocks.AR1.Rouwenhorst(...)

timmens · 2026-02-14T13:52:03Z

src/lcm/grids.py

    def get_coordinate(self, value: ScalarFloat) -> ScalarFloat:
        """Return the generalized coordinate of a value in the grid."""
+        if self.points is None:
+            raise GridInitializationError(
+                "Cannot compute coordinate without points. Pass points at "
+                "initialization or use IrregSpacedGrid(n_points=...) and "
+                "supply points at runtime via params."
+            )
        return grid_helpers.get_irreg_coordinate(value, self.to_jax())


Is this tested? If the error is raised in jitted code the user might see a strange error.

timmens · 2026-02-14T15:00:25Z

tests/solution/test_solve_brute.py

    - state_action_spaces: StateActionSpace object
    - max_Q_over_a_functions: dict mapping period to max_Q_over_a function
    - active: list of periods the regime is active
+    - gridspecs: grid specifications (needed by _replace_dynamic_states)


Read somewhere that we do not want to use the terminology "dynamic". Is the comment outdated or the function name?

timmens · 2026-02-14T15:08:17Z

tests/test_runtime_params.py

Nice tests in general. But the style is slightly off, probably because of Claude. The model functions don't adhere to our standard terminology (_simple_utility instead of just utility). We also don't usually use Test-classes with pytest. Maybe have a quick look at it again to align the style? On top, I would expect _simple_model to be named _create_model, or similar.

timmens · 2026-02-14T16:47:35Z

CLAUDE.md


+### Testing Style
+
+- Use plain pytest functions, never test classes (`class TestFoo`)


Ah, very nice! Must have been added in a later commit then 😅

timmens · 2026-02-15T18:06:31Z

src/lcm/shock_grids.py

+    The class *is* the distribution type — no ``distribution_type`` string needed.
+


Suggested change

The class *is* the distribution type — no ``distribution_type`` string needed.

timmens · 2026-02-15T18:08:11Z

src/lcm/shock_grids.py

+    def get_transition_probs(self) -> FloatND:
+        """Get the transition probabilities at the gridpoints.
+
+        Returns uniform probabilities when required params are missing.


Do we actually want that, or do we want to throw an error if its not fully specified and we request the transition probs?

timmens · 2026-02-15T18:13:23Z

src/lcm/shock_grids.py

For the mathematical formulas / implementation of Tauchen and Rouwenhorst, is there another reference than the original papers, or did we choose the specific implementation ourselves given the original papers?

timmens · 2026-02-15T18:17:17Z

tests/test_shock_grids.py

+@pytest.mark.parametrize(
+    "distribution_type", ["uniform", "normal", "tauchen", "rouwenhorst"]
+)
+def test_model_with_shock(distribution_type):


This is just a regression test, right? We don't test any simulation result changes based on the different shock types?

timmens · 2026-02-15T18:31:48Z

src/lcm/solution/solve_brute.py

    return MappingProxyType(solution)
+
+
+def _replace_runtime_states(


Shouldn't this be a method of InternalRegime? Something like:

state_action_space = internal_regime.build_state_action_space(internal_params[name])

Potentially with a different prefix than build_.

I think we could even make the internal_regime.state_action_space attribute private then.

hmgaudecker and others added 22 commits February 10, 2026 06:29

Rename InternalRegimeParams to FlatRegimeParams and flat_params to fl…

1669e9d

…at_regime_params. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Rename params_template to regime_params_template where type is Regime…

f0b5cbe

…ParamsTemplate. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Rename internal_fixed_params to flat_regime_fixed_params; type func a…

66b31c4

…s Callable in result.py. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Use QNAME_DELIMITER throughout since we need it in more general circu…

44a7627

…mstances than REGIME_SEPARATOR. Some small improvements.

Use pyproject-fmt and update hooks.

4ff6b86

Simplify.

c90ee02

Fix MY model.

60927f1

Improve naming.

b560602

Remove unnecessary defaults.

aaeabe4

Factor out time aggregation function H.

be2b079

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Include utility among user functions (same as H).

0f235c5

Run code-simplifier and some manual fixes in that vein.

1ef6108

Further simplifications.

2eb5b92

Partial fixed_params into functions upon model creation. Stateless tr…

2250a85

…eatment of dynamic IrregSpacedGrids, in particular ShockGrids.

[pre-commit.ci] auto fixes from pre-commit.com hooks

f41eae8

for more information, see https://pre-commit.ci

Back to dags main; tiny harmonisation efforts in pyproject.toml

24f07b5

Improve typing.

8f1e84d

Merge branch 'remove-internal_regime_params' into factor-out-time-agg…

0b2e5d3

…regation

Update lockfile, remove now-irrelevant ty:ignore

876b8af

Remove a couple of future imports

187304f

hmgaudecker mentioned this pull request Feb 13, 2026

Replace internal_regime_params with flat kwargs via dags.signature.rename_arguments. #233

Merged

hmgaudecker added 6 commits February 13, 2026 11:13

prek autoupdate and include [tool.pyproject-fmt] in pyproject.toml

7ce65c9

Merge branch 'remove-internal_regime_params' into factor-out-time-agg…

4183875

…regation

Merge branch 'factor-out-time-aggregation' into partial-fixed_params

c706b26

Address review comments from #233.

abdc30b

Merge branch 'main' into factor-out-time-aggregation

607fd53

Merge branch 'factor-out-time-aggregation' into partial-fixed_params

e5340f8

Base automatically changed from factor-out-time-aggregation to main February 13, 2026 12:36

hmgaudecker added 10 commits February 13, 2026 13:36

Merge branch 'main' into partial-fixed_params

ebffbaf

Simplify.

f7084ea

Improve naming of grids that may change at runtime.

844c87d

Simplify grid logic further.

3642954

Run the MY model three times just for the pleasure of looking at the …

cbce3b4

…speed on runs 2 and especially 3 :-)

Improve naming.

167f4b7

Simplify fixed_params mechanism.

cb17244

Stateless setting up of model.

63c52cb

Split up test file in a now-sensible way.

00c9ed0

Fix params flow.

fabdae9

hmgaudecker commented Feb 13, 2026

View reviewed changes

hmgaudecker marked this pull request as ready for review February 13, 2026 17:54

hmgaudecker requested a review from timmens February 13, 2026 17:54

Refactor ShockGrids to be strongly typed.

b02d435

hmgaudecker requested a review from mj023 February 14, 2026 07:54

hmgaudecker mentioned this pull request Feb 14, 2026

Add the NestedMappingParams type. #237

Open

2 tasks

mj023 reviewed Feb 14, 2026

View reviewed changes

Review comments: Use form y_{t+1} = μ + ρ y_{t} + ϵ_{t+1} for AR(1)-p…

a52684d

…rocesses throughout. Stick to rho - mu - sigma for all shock parameters.

hmgaudecker force-pushed the partial-fixed_params branch from 65b0c67 to a52684d Compare February 15, 2026 06:57

hmgaudecker added 4 commits February 15, 2026 08:15

Stateless Rouwenhorst.

d5452bd

More readable formatting.

a424870

Normalise probabilities explicitly to get rid of failure on 32bit.

365754d

Revert previous change; make rounding precision in tests to be archit…

bc9d0f3

…ecture-dependent.

timmens reviewed Feb 15, 2026

View reviewed changes


		### Testing Style

		- Use plain pytest functions, never test classes (`class TestFoo`)

		The class is the distribution type — no ``distribution_type`` string needed.

		return MappingProxyType(solution)


		def _replace_runtime_states(

Conversation

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

Uh oh!

Uh oh!

hmgaudecker left a comment

Choose a reason for hiding this comment

Uh oh!

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

hmgaudecker commented Feb 14, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

timmens left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

hmgaudecker commented Feb 10, 2026 •

edited

Loading

hmgaudecker Feb 13, 2026 •

edited

Loading