Reorder dimensions in productmap (#335)

Co-authored-by: Hans-Martin von Gaudecker <hmgaudecker@gmail.com>
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 4 people

src/lcm/exceptions.py

@@ -50,6 +50,10 @@ class GridInitializationError(PyLCMError):
     """Raised when there is an error in the grid initialization."""
 
 
+class FunctionDispatchError(PyLCMError):
+    """Raised when there is an error during the function dispatch."""
+
+
 def format_messages(errors: str | list[str]) -> str:
     """Convert message or list of messages into a single string."""
     if isinstance(errors, str):

src/lcm/regime_building/variable_info.py

@@ -1,3 +1,4 @@
+import math
 from types import MappingProxyType
 
 import pandas as pd
@@ -34,15 +35,32 @@ def get_variable_info(regime: Regime) -> pd.DataFrame:
     ]
     info["is_discrete"] = ~info["is_continuous"]
 
-    order = info.query("is_discrete & is_state").index.tolist()
-    order += info.query("is_discrete & is_action").index.tolist()
-    order += info.query("is_continuous & is_state").index.tolist()
-    order += info.query("is_continuous & is_action").index.tolist()
+    ordered_discrete_states = sorted(
+        info.query("is_discrete & is_state").index.tolist(),
+        key=lambda x: (
+            regime.states[x].batch_size
+            if regime.states[x].batch_size != 0
+            else math.inf
+        ),
+    )
+    ordered_continuous_states = sorted(
+        info.query("is_continuous & is_state").index.tolist(),
+        key=lambda x: (
+            regime.states[x].batch_size
+            if regime.states[x].batch_size != 0
+            else math.inf
+        ),
+    )
+    ordered_states_and_actions = [
+        *ordered_discrete_states,
+        *ordered_continuous_states,
+        *info.query("is_action").index.tolist(),
+    ]
 
-    if set(order) != set(info.index):
+    if set(ordered_states_and_actions) != set(info.index):
         raise ValueError("Order and index do not match.")
 
-    return info.loc[order]
+    return info.loc[ordered_states_and_actions]
 
 
 def get_grids(

src/lcm/utils/dispatchers.py

@@ -8,6 +8,7 @@
 import jax.numpy as jnp
 from jax import Array, vmap
 
+from lcm.exceptions import FunctionDispatchError
 from lcm.typing import Float1D, FloatND
 from lcm.utils.containers import find_duplicates
 from lcm.utils.functools import allow_args, allow_only_kwargs
@@ -224,7 +225,7 @@ def _base_productmap_batched(
     Like `jax.lax.map`, this function does not preserve the function signature.
 
     Args:
-        func: The function to be dispatched. Cannot have keyword-only arguments.
+        func: The function to be dispatched. Cannot have positional-only parameters.
         product_axes: Tuple with names of arguments over which we apply
             `jax.lax.map`.
         batch_sizes: Dict with the batch sizes for each product_axis.
@@ -234,6 +235,13 @@ def _base_productmap_batched(
 
     """
     parameters = inspect.signature(func).parameters
+    for name, param in parameters.items():
+        if param.kind == inspect.Parameter.POSITIONAL_ONLY:
+            raise FunctionDispatchError(
+                "Positional-only parameters are not allowed in dispatched functions. "
+                f"The parameter '{name}' to the function {func.__name__} "
+                "is POSITIONAL_ONLY."
+            )
 
     def batched_vmap(**kwargs: FloatND) -> FloatND:
         non_array_kwargs = {
@@ -250,14 +258,6 @@ def map_one_more(
             def func_mapped_over_one_more_axis(
                 *already_mapped_args: Float1D, **already_mapped_kwargs: Float1D
             ) -> FloatND:
-                if parameters[axis].kind == inspect.Parameter.POSITIONAL_ONLY:
-                    return jax.lax.map(
-                        lambda axis_i: loop_func(
-                            axis_i, *already_mapped_args, **already_mapped_kwargs
-                        ),
-                        jnp.atleast_1d(kwargs[axis]),
-                        batch_size=batch_sizes[axis],
-                    )
                 return jax.lax.map(
                     lambda axis_i: loop_func(
                         *already_mapped_args, **{axis: axis_i}, **already_mapped_kwargs
@@ -271,6 +271,7 @@ def func_mapped_over_one_more_axis(
         # Loop over all product axes
         for axis in reversed(product_axes):
             func_with_partialled_args = map_one_more(func_with_partialled_args, axis)
+
         return cast("FloatND", func_with_partialled_args())
 
     return cast("FunctionWithArrayReturn", batched_vmap)

src/lcm_examples/mahler_yum_2024/_model.py

@@ -326,7 +326,7 @@ def dead_is_active(age: int, initial_age: float) -> bool:
         "education": DiscreteGrid(Education),
         "productivity": DiscreteGrid(ProductivityType),
         "health_type": DiscreteGrid(HealthType),
-        "discount_type": DiscreteGrid(DiscountType, batch_size=1),
+        "discount_type": DiscreteGrid(DiscountType),
     },
     state_transitions={
         "wealth": next_wealth,
@@ -380,7 +380,7 @@ def dead_utility(discount_type: DiscreteState) -> FloatND:  # noqa: ARG001
     transition=None,
     active=partial(dead_is_active, initial_age=ages.values[0]),
     states={
-        "discount_type": DiscreteGrid(DiscountType, batch_size=1),
+        "discount_type": DiscreteGrid(DiscountType),
     },
     functions={"utility": dead_utility},
 )

tests/regime_building/test_regime_processing.py

@@ -7,8 +7,8 @@
 from numpy.testing import assert_array_equal
 from pandas.testing import assert_frame_equal
 
-from lcm import DiscreteGrid
 from lcm.ages import AgeGrid
+from lcm.grids import DiscreteGrid, LinSpacedGrid
 from lcm.regime_building.processing import (
     _rename_params_to_qnames,
     process_regimes,
@@ -76,6 +76,35 @@ def next_c(a, b):
     assert got["c"].codes == (0, 1)
 
 
+def test_get_grids_reorder(binary_category_class):
+    def next_state(a, b):
+        pass
+
+    regime_mock = RegimeMock(
+        actions={
+            "a": DiscreteGrid(binary_category_class),
+        },
+        states={
+            "b": DiscreteGrid(binary_category_class),
+            "c": DiscreteGrid(binary_category_class, batch_size=1),
+            "d": LinSpacedGrid(start=0, stop=1, n_points=5, batch_size=3),
+            "e": LinSpacedGrid(start=0, stop=1, n_points=5, batch_size=1),
+            "f": LinSpacedGrid(start=0, stop=1, n_points=5),
+        },
+        state_transitions={
+            "b": next_state,
+            "c": next_state,
+            "d": next_state,
+            "e": next_state,
+            "f": next_state,
+        },
+        functions={"utility": lambda _c: None},
+    )
+
+    got = get_grids(regime_mock)  # ty: ignore[invalid-argument-type]
+    assert list(got.keys()) == ["c", "b", "e", "d", "f", "a"]
+
+
 def test_process_regimes():
     ages = AgeGrid(start=0, stop=4, step="Y")
     regimes = {"working_life": working_life, "dead": dead}
@@ -93,12 +122,12 @@ def test_process_regimes():
     # Variable Info
     assert (
         internal_working_regime.variable_info["is_state"].to_numpy()
-        == np.array([False, True, False])
+        == np.array([True, False, False])
     ).all()
 
     assert (
         internal_working_regime.variable_info["is_continuous"].to_numpy()
-        == np.array([False, True, True])
+        == np.array([True, False, True])
     ).all()
 
     # Grids — compare the grid objects (which now include transition attributes)

tests/test_dispatchers.py

@@ -4,6 +4,7 @@
 import pytest
 from numpy.testing import assert_array_almost_equal as aaae
 
+from lcm.exceptions import FunctionDispatchError
 from lcm.utils.dispatchers import (
     productmap,
     simulation_spacemap,
@@ -12,30 +13,14 @@
 from lcm.utils.functools import allow_args
 
 
-def f(a, /, *, b, c):
-    """Tests that dispatchers can handle positional-only and keyword-only arguments.
+def f(a, *, b, c):
+    """Tests that dispatchers can handle standard arguments and keyword-only arguments.
 
-    a is positional-only, b and c are keyword-only
+    a is positional-or-keyword, b and c are keyword-only
     """
     return jnp.sin(a) + jnp.cos(b) + jnp.tan(c)
 
 
-def f2(b, a, /, *, c):
-    """Tests that dispatchers can handle positional-only and keyword-only arguments.
-
-    b and a are positional-only, c is keyword-only
-    """
-    return jnp.sin(a) + jnp.cos(b) + jnp.tan(c)
-
-
-def g(a, /, b, *, c, d):
-    """Tests that dispatchers can handle positional-only and keyword-only arguments.
-
-    a is positional-only, b is positional-or-keyword, c and d are keyword-only
-    """
-    return f(a, b=b, c=c) + jnp.log(d)
-
-
 @pytest.fixture
 def setup_productmap_f():
     return {
@@ -57,34 +42,10 @@ def expected_productmap_f():
     return allow_args(f)(*helper).reshape(10, 7, 5)
 
 
-@pytest.fixture
-def setup_productmap_g():
-    return {
-        "a": jnp.linspace(-5, 5, 10),
-        "b": jnp.linspace(0, 3, 7),
-        "c": jnp.linspace(1, 5, 5),
-        "d": jnp.linspace(1, 3, 4),
-    }
-
-
-@pytest.fixture
-def expected_productmap_g():
-    grids = {
-        "a": jnp.linspace(-5, 5, 10),
-        "b": jnp.linspace(0, 3, 7),
-        "c": jnp.linspace(1, 5, 5),
-        "d": jnp.linspace(1, 3, 4),
-    }
-
-    helper = jnp.array(list(itertools.product(*grids.values()))).T
-    return allow_args(g)(*helper).reshape(10, 7, 5, 4)
-
-
 @pytest.mark.parametrize(
     ("func", "args", "grids", "expected"),
     [
         (f, ["a", "b", "c"], "setup_productmap_f", "expected_productmap_f"),
-        (g, ["a", "b", "c", "d"], "setup_productmap_g", "expected_productmap_g"),
     ],
 )
 def test_productmap_with_all_arguments_mapped(func, args, grids, expected, request):
@@ -112,24 +73,13 @@ def test_productmap_with_positional_args(setup_productmap_f):
         decorated(*setup_productmap_f.values())  # ty: ignore[missing-argument]
 
 
-def test_productmap_different_func_order(setup_productmap_f):
-    _bs = dict.fromkeys(("a", "b", "c"), 0)
-    decorated_f = productmap(func=f, variables=("a", "b", "c"), batch_sizes=_bs)
-    expected = decorated_f(**setup_productmap_f)  # ty: ignore[missing-argument]
-
-    decorated_f2 = productmap(func=f2, variables=("a", "b", "c"), batch_sizes=_bs)
-    calculated_f2 = decorated_f2(**setup_productmap_f)  # ty: ignore[missing-argument]
-
-    aaae(calculated_f2, expected)
-
-
 def test_productmap_change_arg_order(setup_productmap_f, expected_productmap_f):
     expected = jnp.transpose(expected_productmap_f, (1, 0, 2))
 
     decorated = productmap(
         func=f, variables=("b", "a", "c"), batch_sizes=dict.fromkeys(("b", "a", "c"), 0)
     )
-    calculated = decorated(**setup_productmap_f)  # ty: ignore[missing-argument]
+    calculated = decorated(**setup_productmap_f)
 
     aaae(calculated, expected)
 
@@ -148,7 +98,7 @@ def test_productmap_with_all_arguments_mapped_some_len_one():
     decorated = productmap(
         func=f, variables=("a", "b", "c"), batch_sizes=dict.fromkeys(("a", "b", "c"), 0)
     )
-    calculated = decorated(**grids)  # ty: ignore[missing-argument]
+    calculated = decorated(**grids)
     aaae(calculated, expected)
 
 
@@ -166,7 +116,7 @@ def test_productmap_with_some_arguments_mapped():
     decorated = productmap(
         func=f, variables=("a", "c"), batch_sizes=dict.fromkeys(("a", "c"), 0)
     )
-    calculated = decorated(**grids)  # ty: ignore[missing-argument]
+    calculated = decorated(**grids)
     aaae(calculated, expected)
 
 
@@ -201,6 +151,14 @@ def test_productmap_with_some_argument_mapped_twice():
         )
 
 
+def test_productmap_rejects_positional_only():
+    def h(a, /, *, b):
+        return a + b
+
+    with pytest.raises(FunctionDispatchError, match="POSITIONAL_ONLY"):
+        productmap(func=h, variables=("a", "b"), batch_sizes={"a": 0, "b": 0})
+
+
 @pytest.fixture
 def setup_spacemap():
     value_grid = {
@@ -210,14 +168,12 @@ def setup_spacemap():
 
     combination_values = {
         "c": jnp.array([7.0, 8, 9, 10]),
-        "d": jnp.array([9.0, 10, 11, 12, 13]),
     }
 
     helper = jnp.array(list(itertools.product(*combination_values.values()))).T
 
     combination_grid = {
         "c": helper[0],
-        "d": helper[1],
     }
     return value_grid, combination_grid
 
@@ -231,13 +187,12 @@ def expected_spacemap():
 
     combination_grid = {
         "c": jnp.array([7.0, 8, 9, 10]),
-        "d": jnp.array([9.0, 10, 11, 12, 13]),
     }
 
     all_grids = {**value_grid, **combination_grid}
     helper = jnp.array(list(itertools.product(*all_grids.values()))).T
 
-    return allow_args(g)(*helper).reshape(3, 2, 4 * 5)
+    return allow_args(f)(*helper).reshape(3, 2, 4)
 
 
 def test_spacemap_all_arguments_mapped(
@@ -247,11 +202,11 @@ def test_spacemap_all_arguments_mapped(
     product_vars, combination_vars = setup_spacemap
 
     decorated = simulation_spacemap(
-        func=g,
+        func=f,
         action_names=tuple(product_vars),
         state_names=tuple(combination_vars),
     )
-    calculated = decorated(**product_vars, **combination_vars)  # ty: ignore[missing-argument]
+    calculated = decorated(**product_vars, **combination_vars)
 
     aaae(calculated, jnp.transpose(expected_spacemap, axes=(2, 0, 1)))
 
@@ -274,7 +229,7 @@ def test_spacemap_all_arguments_mapped(
 def test_spacemap_arguments_overlap(error_msg, product_vars, combination_vars):
     with pytest.raises(ValueError, match=error_msg):
         simulation_spacemap(
-            func=g, action_names=product_vars, state_names=combination_vars
+            func=f, action_names=product_vars, state_names=combination_vars
         )