Add batched productmap (#280)

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

Author: 4 people

pixi.lock

@@ -6,6 +6,16 @@
 class Grid(ABC):
     """LCM Grid base class."""
 
+    @property
+    @abstractmethod
+    def batch_size(self) -> int:
+        """Size of the batches looped over during the solution.
+
+        `ContinuousGrid` overrides this via its dataclass field.
+        `DiscreteGrid` overrides this via its own property.
+
+        """
+
     @abstractmethod
     def to_jax(self) -> Int1D | Float1D:
         """Convert the grid to a Jax array."""

src/lcm/grids/base.py

@@ -16,6 +16,7 @@
 )
 
 
+@dataclass(frozen=True, kw_only=True)
 class ContinuousGrid(Grid):
     """Base class for grids representing continuous values with coordinate lookup.
 
@@ -24,6 +25,9 @@ class ContinuousGrid(Grid):
 
     """
 
+    batch_size: int = 0
+    """Size of the batches that are looped over during the solution."""
+
     @overload
     def get_coordinate(self, value: ScalarFloat) -> ScalarFloat: ...
     @overload

src/lcm/grids/continuous.py

@@ -6,15 +6,26 @@
 from lcm.utils.containers import get_field_names_and_values
 
 
-class _DiscreteGridBase(Grid):
-    """Base class for discrete grids: categories, codes, and JAX conversion."""
+class DiscreteGrid(Grid):
+    """A discrete grid defining the outcome space of a categorical variable.
+
+    Args:
+        category_class: The category class representing the grid categories. Must
+            be a dataclass with fields that have unique int values.
+
+    Raises:
+        GridInitializationError: If the `category_class` is not a dataclass with int
+            fields.
+
+    """
 
-    def __init__(self, category_class: type) -> None:
+    def __init__(self, category_class: type, batch_size: int = 0) -> None:
         _validate_discrete_grid(category_class)
         names_and_values = get_field_names_and_values(category_class)
         self.__categories = tuple(names_and_values.keys())
         self.__codes = tuple(names_and_values.values())
         self.__ordered: bool = getattr(category_class, "_ordered", False)
+        self.__batch_size: int = batch_size
 
     @property
     def categories(self) -> tuple[str, ...]:
@@ -31,23 +42,11 @@ def ordered(self) -> bool:
         """Return whether the categories have a meaningful ordering."""
         return self.__ordered
 
+    @property
+    def batch_size(self) -> int:
+        """Return batch size during solution."""
+        return self.__batch_size
+
     def to_jax(self) -> Int1D:
         """Convert the grid to a Jax array."""
         return jnp.array(self.codes)
-
-
-class DiscreteGrid(_DiscreteGridBase):
-    """A discrete grid defining the outcome space of a categorical variable.
-
-    Args:
-        category_class: The category class representing the grid categories. Must
-            be a dataclass with fields that have unique int values.
-
-    Raises:
-        GridInitializationError: If the `category_class` is not a dataclass with int
-            fields.
-
-    """
-
-    def __init__(self, category_class: type) -> None:
-        super().__init__(category_class)

src/lcm/grids/discrete.py

@@ -385,7 +385,7 @@ def _check_leaf(value: object, path: str) -> None:
         type_name = type(value).__module__ + "." + type(value).__name__
         msg = (
             f"Parameter '{path}' is a {type_name} (shape {value.shape}). "
-            f"Use jax.numpy.array() or pass a pd.Series with a named index."
+            f"Use jnp.array() or pass a pd.Series with a named index."
         )
         raise InvalidParamsError(msg)
     type_name = type(value).__module__ + "." + type(value).__name__

src/lcm/model.py

@@ -357,7 +357,7 @@ def array_from_series(
     display_params = ["age" if p == "period" else p for p in indexing_params]
 
     level_mappings = _build_level_mappings_for_param(
-        indexing_params=display_params, all_grids=grids, ages=model.ages
+        indexing_params=display_params, grids=grids, ages=model.ages
     )
 
     # Append outcome axis for transition probability arrays (next_* functions
@@ -368,7 +368,7 @@ def array_from_series(
         ]
         if next_levels:
             outcome_mapping = _build_outcome_mapping(
-                func_name=func_name, all_grids=grids, model=model
+                func_name=func_name, grids=grids, model=model
             )
             level_mappings = (*level_mappings, outcome_mapping)
 
@@ -566,8 +566,8 @@ class _LevelMapping:
     size: int
     """Number of positions along this axis."""
 
-    label_to_index: Callable[[object], int]
-    """Map a single label value to its integer index."""
+    get_code_from_label: Callable[[str], int]
+    """Return the integer code for a label."""
 
     valid_labels: tuple[str, ...] = ()
     """Valid label names, for error messages. Empty for age levels."""
@@ -578,7 +578,7 @@ def _age_level_mapping(ages: AgeGrid) -> _LevelMapping:
     return _LevelMapping(
         name="age",
         size=ages.n_periods,
-        label_to_index=ages.age_to_period,  # ty: ignore[invalid-argument-type]
+        get_code_from_label=ages.age_to_period,  # ty: ignore[invalid-argument-type]
     )
 
 
@@ -597,23 +597,23 @@ def _grid_level_mapping(*, name: str, grid: DiscreteGrid) -> _LevelMapping:
     return _LevelMapping(
         name=name,
         size=len(grid.categories),
-        label_to_index=label_to_code.__getitem__,
+        get_code_from_label=label_to_code.__getitem__,
         valid_labels=grid.categories,
     )
 
 
 def _build_level_mappings_for_param(
     *,
     indexing_params: list[str],
-    all_grids: dict[str, DiscreteGrid],
+    grids: dict[str, DiscreteGrid],
     ages: AgeGrid,
 ) -> tuple[_LevelMapping, ...]:
     """Build level mappings for `array_from_series` from indexing params.
 
     Args:
         indexing_params: Parameter names in output axis order, with
             `"period"` already replaced by `"age"`.
-        all_grids: Categorical grid lookup.
+        grids: Categorical grid lookup.
         ages: The model's `AgeGrid`.
 
     Returns:
@@ -624,12 +624,12 @@ def _build_level_mappings_for_param(
     for param in indexing_params:
         if param == "age":
             mappings.append(_age_level_mapping(ages))
-        elif param in all_grids:
-            mappings.append(_grid_level_mapping(name=param, grid=all_grids[param]))
+        elif param in grids:
+            mappings.append(_grid_level_mapping(name=param, grid=grids[param]))
         else:
             msg = (
                 f"Unrecognised indexing parameter '{param}'. Expected 'age' "
-                f"or a discrete grid name ({sorted(all_grids)}). If "
+                f"or a discrete grid name ({sorted(grids)}). If "
                 f"'{param}' is a DAG function output, pass "
                 f'derived_categoricals={{"{param}": DiscreteGrid(...)}} '
                 f"to solve() / simulate()."
@@ -641,7 +641,7 @@ def _build_level_mappings_for_param(
 def _build_outcome_mapping(
     *,
     func_name: str,
-    all_grids: dict[str, DiscreteGrid],
+    grids: dict[str, DiscreteGrid],
     model: Model,
 ) -> _LevelMapping:
     """Build a `_LevelMapping` for the outcome axis of a `next_*` function.
@@ -651,7 +651,7 @@ def _build_outcome_mapping(
 
     Args:
         func_name: Function name starting with `"next_"`.
-        all_grids: Categorical grid lookup.
+        grids: Categorical grid lookup.
         model: The LCM Model instance.
 
     Returns:
@@ -663,13 +663,13 @@ def _build_outcome_mapping(
         return _LevelMapping(
             name="next_regime",
             size=len(regime_ids),
-            label_to_index=regime_ids.__getitem__,  # ty: ignore[invalid-argument-type]
+            get_code_from_label=regime_ids.__getitem__,
             valid_labels=tuple(regime_ids),
         )
 
     path = tree_path_from_qname(func_name)
     state_name = path[0].removeprefix("next_")
-    return _grid_level_mapping(name=f"next_{state_name}", grid=all_grids[state_name])
+    return _grid_level_mapping(name=f"next_{state_name}", grid=grids[state_name])
 
 
 def _scatter_series(
@@ -744,7 +744,7 @@ def _map_level(*, mapping: _LevelMapping, level_values: pd.Index) -> np.ndarray:
         raise ValueError(msg)
 
     try:
-        return np.array([mapping.label_to_index(v) for v in level_values])
+        return np.array([mapping.get_code_from_label(v) for v in level_values])
     except ValueError:
         # Age levels: age_to_period raises ValueError with a good message
         raise

src/lcm/pandas_utils.py

@@ -109,11 +109,13 @@ def get_Q_and_F(
             - set(target_transitions)
             - {V_arr_name}
         )
+        stochastic_variables = tuple(
+            key for key in target_transitions if key in stochastic_transition_names
+        )
         next_V[target_regime_name] = productmap(
             func=next_V_interpolator,
-            variables=tuple(
-                key for key in target_transitions if key in stochastic_transition_names
-            ),
+            variables=stochastic_variables,
+            batch_sizes=dict.fromkeys(stochastic_variables, 0),
         )
 
     # ----------------------------------------------------------------------------------
@@ -342,7 +344,10 @@ def _outer(**kwargs: Float1D) -> FloatND:
         weights = jnp.array(list(kwargs.values()))
         return jnp.prod(weights)
 
-    return productmap(func=_outer, variables=tuple(arg_names))
+    variables = tuple(arg_names)
+    return productmap(
+        func=_outer, variables=variables, batch_sizes=dict.fromkeys(variables, 0)
+    )
 
 
 def _get_U_and_F(

src/lcm/regime_building/Q_and_F.py

@@ -21,6 +21,7 @@
 def get_max_Q_over_a(
     *,
     Q_and_F: Callable[..., tuple[FloatND, BoolND]],
+    batch_sizes: dict[str, int],
     action_names: tuple[str, ...],
     state_names: tuple[str, ...],
 ) -> MaxQOverAFunction:
@@ -47,6 +48,8 @@ def get_max_Q_over_a(
         Q_and_F: A function that takes a state-action combination and returns the action
             value of that combination and whether the state-action combination is
             feasible.
+        batch_sizes: Mapping of state variable names to batch sizes for the outer
+            productmap over states. A batch size of 0 means no batching.
         action_names: Tuple of action variable names.
         state_names: Tuple of state names.
 
@@ -60,9 +63,12 @@ def get_max_Q_over_a(
         Q_and_F=Q_and_F, action_names=action_names, state_names=state_names
     )
 
+    # Actions are the inner optimization axis — batching applies only to the
+    # outer state loop.
     Q_and_F = productmap(
         func=Q_and_F,
         variables=action_names,
+        batch_sizes=dict.fromkeys(action_names, 0),
     )
 
     @with_signature(
@@ -80,7 +86,7 @@ def max_Q_over_a(
         )
         return Q_arr.max(where=F_arr, initial=-jnp.inf)
 
-    return productmap(func=max_Q_over_a, variables=state_names)
+    return productmap(func=max_Q_over_a, variables=state_names, batch_sizes=batch_sizes)
 
 
 def get_argmax_and_max_Q_over_a(
@@ -130,6 +136,7 @@ def get_argmax_and_max_Q_over_a(
     Q_and_F = productmap(
         func=Q_and_F,
         variables=action_names,
+        batch_sizes=dict.fromkeys(action_names, 0),
     )
 
     @with_signature(

src/lcm/regime_building/max_Q_over_a.py

@@ -102,6 +102,8 @@ def process_regimes(
     )
     all_grids = MappingProxyType({n: get_grids(r) for n, r in regimes.items()})
 
+    _fail_if_action_has_batch_size(regimes)
+
     regime_to_v_interpolation_info = MappingProxyType(
         {n: create_v_interpolation_info(r) for n, r in regimes.items()}
     )
@@ -243,6 +245,7 @@ def _build_solve_functions(
     max_Q_over_a = _build_max_Q_over_a_per_period(
         state_action_space=state_action_space,
         Q_and_F_functions=Q_and_F_functions,
+        grids=all_grids[regime_name],
         enable_jit=enable_jit,
     )
 
@@ -1256,13 +1259,19 @@ def _build_max_Q_over_a_per_period(
     *,
     state_action_space: StateActionSpace,
     Q_and_F_functions: MappingProxyType[int, QAndFFunction],
+    grids: MappingProxyType[str, Grid],
     enable_jit: bool,
 ) -> MappingProxyType[int, MaxQOverAFunction]:
     """Build max-Q-over-a closures for each period."""
     result = {}
     for period, Q_and_F in Q_and_F_functions.items():
         func = get_max_Q_over_a(
             Q_and_F=Q_and_F,
+            batch_sizes={
+                name: grid.batch_size
+                for name, grid in grids.items()
+                if name in state_action_space.state_names
+            },
             action_names=state_action_space.action_names,
             state_names=state_action_space.state_names,
         )
@@ -1323,3 +1332,22 @@ def _build_next_state_vmapped(
     )
 
     return jax.jit(next_state_vmapped) if enable_jit else next_state_vmapped
+
+
+def _fail_if_action_has_batch_size(regimes: Mapping[str, Regime]) -> None:
+    """Raise if any action grid has a non-zero batch_size.
+
+    Batching applies only to the outer state loop during solving, not to the
+    inner action optimization. A non-zero batch_size on an action grid would be
+    silently ignored, so we reject it early.
+
+    """
+    for regime_name, regime in regimes.items():
+        for action_name, grid in regime.actions.items():
+            if grid.batch_size != 0:
+                msg = (
+                    f"batch_size > 0 is not supported on action grids. Only state "
+                    f"grids can be batched. Found batch_size={grid.batch_size} on "
+                    f"action '{action_name}' in regime '{regime_name}'."
+                )
+                raise ValueError(msg)

src/lcm/regime_building/processing.py

@@ -45,7 +45,9 @@ class _ShockGrid(ContinuousGrid):
     @property
     def _param_field_names(self) -> tuple[str, ...]:
         """Names of distribution-specific parameters."""
-        return tuple(f.name for f in fields(self) if f.name != "n_points")
+        return tuple(
+            f.name for f in fields(self) if f.name not in {"n_points", "batch_size"}
+        )
 
     @property
     def params(self) -> MappingProxyType[str, float]: