Replace `internal_regime_params` with flat kwargs via `dags.signature.rename_arguments`. by hmgaudecker · Pull Request #233 · OpenSourceEconomics/pylcm

hmgaudecker · 2026-02-10T05:17:10Z

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.

It is a crucial prerequisite for solving #174 and the partialling of fixed_params described in #219, which will be done in subsequent PRs.

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.

hmgaudecker

Autoreview.

hmgaudecker · 2026-02-10T05:28:30Z

src/lcm/input_processing/regime_processing.py

-
-    return cast("InternalUserFunction", processed_func)
+    """
+    param_names = list(params_template[param_key])  # ty: ignore[invalid-argument-type]


ty: ignore will be gone once in PR that follows immediately.

hmgaudecker · 2026-02-10T07:57:35Z

src/lcm/input_processing/params_processing.py

+    """Get all flat parameter names from a regime params template.
+
+    Converts nested template entries like {"utility": {"risk_aversion": type}} to
+    flat names like "utility__risk_aversion". Top-level params like


Will be simplified in subsequent PR.

hmgaudecker · 2026-02-10T09:06:43Z

src/lcm/input_processing/regime_processing.py

-
-    return cast("InternalUserFunction", processed_func)
+    """
+    param_names = list(regime_params_template[fn_key])  # ty: ignore[invalid-argument-type]


Note: ty: ignore will disappear in subsequent PR.

timmens

While starting the review I was a bit skeptical about the change, but was quickly convinced that this is a big improvement. Very nice PR! Only minor comments, but already approved.

timmens · 2026-02-12T22:36:26Z

tests/input_processing/test_process_params.py

    assert set(internal_params.keys()) == set(params_template.keys())
    for regime in params_template:
-        assert set(internal_params[regime].keys()) == set(
-            params_template[regime].keys()
-        )
+        expected_flat_keys = set()
+        for func, func_params in params_template[regime].items():
+            for arg in func_params:
+                expected_flat_keys.add(f"{func}__{arg}")
+        assert set(internal_params[regime].keys()) == expected_flat_keys


Seeing this check now for the third time. Make a function out of it?

timmens · 2026-02-13T07:38:18Z

tests/test_next_state.py

-            },
-        }
-    )
+    flat_regime_params = {


Is there a reason why this dictionary does not need to be a MappingProxyType anymore?

src/lcm/input_processing/params_processing.py

timmens · 2026-02-13T07:41:17Z

src/lcm/input_processing/params_processing.py

+    {"working": {"discount_factor": value}}.
+
+    """
+    result: dict[str, dict[str, Any]] = {}


Suggested change

result: dict[str, dict[str, Any]] = {}

result: dict[RegimeName, dict[str, Any]] = {}

timmens · 2026-02-13T07:46:44Z

src/lcm/input_processing/params_processing.py

+            for param_name in value:
+                result.add(f"{key}{QNAME_DELIMITER}{param_name}")
+        else:
+            # Top-level param (e.g., "discount_factor": float)


This exact comment was renamed above to

Scalar param (currently unused - all params are under namespaces)

timmens · 2026-02-13T07:53:21Z

src/lcm/input_processing/regime_components.py

+    params = tuple(inspect.signature(next_state).parameters)

    next_state_vmapped = vmap_1d(
        func=next_state,
        variables=tuple(
-            parameter
-            for parameter in parameters
-            if parameter not in ("period", "age", "internal_regime_params")
+            p for p in params if p not in _get_non_vmap_params(regime_params_template)


I think params and/or the p in the loop should be replaced by something like argnames and var. I was slightly confused for which reason we would vmap over parameter, which we don't, we vmap over states and actions that are arguments of the next_state function.

timmens · 2026-02-13T07:54:44Z

src/lcm/input_processing/regime_components.py

+    params = tuple(inspect.signature(next_regime_accepting_all).parameters)

    next_regime_vmapped = vmap_1d(
        func=next_regime_accepting_all,
        variables=tuple(
-            parameter
-            for parameter in parameters
-            if parameter not in ("period", "age", "internal_regime_params")
+            p for p in params if p not in _get_non_vmap_params(regime_params_template)


Same comment as above. It confuses me a bit that we use params to denote function arguments but also to denote the model parameters. I think it should be clear from the naming that these are different concepts

timmens · 2026-02-13T08:05:24Z

src/lcm/max_Q_over_a.py

-            internal_regime_params=internal_regime_params,
            next_V_arr=next_V_arr,
-            **states_and_actions,
+            **kwargs,


In the **states_and_actions case I liked that it was clear that the unpacked dictionary contains the states and actions. If I understand it correctly, kwargs now also contains the flat parameters, which is why we need the more general name. Could we still add a comment for the reader on what is expected to be found inside kwargs?

timmens · 2026-02-13T08:05:42Z

src/lcm/max_Q_over_c.py

-        internal_regime_params: InternalRegimeParams,
        period: Period,
-        **states_and_actions: Array,
+        **kwargs: Array,


Same comment on kwargs as above.

hmgaudecker · 2026-02-13T09:13:56Z

These are great comments, thanks!

Could you live with me applying them to #235 ? Would make life much easier...

timmens · 2026-02-13T09:29:21Z

Yes, lets do it in #235 !!

hmgaudecker and others added 2 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>

hmgaudecker force-pushed the remove-internal_regime_params branch from 8aef1ab to 1669e9d Compare February 10, 2026 06:23

hmgaudecker and others added 3 commits February 10, 2026 07:30

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.

hmgaudecker force-pushed the remove-internal_regime_params branch from 883fa5f to 44a7627 Compare February 10, 2026 08:01

hmgaudecker added 5 commits February 10, 2026 09:24

Use pyproject-fmt and update hooks.

4ff6b86

Simplify.

c90ee02

Fix MY model.

60927f1

Improve naming.

b560602

Remove unnecessary defaults.

aaeabe4

hmgaudecker commented Feb 10, 2026

View reviewed changes

hmgaudecker requested a review from timmens February 10, 2026 09:17

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

24f07b5

timmens approved these changes Feb 13, 2026

View reviewed changes

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

7ce65c9

hmgaudecker added a commit that referenced this pull request Feb 13, 2026

Address review comments from #233.

abdc30b

hmgaudecker merged commit c65ad31 into main Feb 13, 2026
8 checks passed

hmgaudecker deleted the remove-internal_regime_params branch February 13, 2026 12:22

hmgaudecker mentioned this pull request Feb 14, 2026

Allow more deeply nested mappings in params while retaining tight typing. #237

Merged

2 tasks

	result: dict[str, dict[str, Any]] = {}
	result: dict[RegimeName, dict[str, Any]] = {}

Conversation

hmgaudecker commented Feb 10, 2026

Uh oh!

hmgaudecker 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!

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!

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!

hmgaudecker commented Feb 13, 2026

Uh oh!

timmens commented Feb 13, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants