Support arbitrary observation dimensions

src/hssm/addm/addm.py

@@ -152,6 +152,7 @@ def _make_model_distribution(self) -> type[pm.Distribution]:
             lapse=self.lapse,
             extra_fields=extra_fields_data,
             params_is_trialwise=params_is_trialwise,
+            expected_obs_dim=self._obs_dim,
         )
         # Expose the observed fixations to the RV's generative path so
         # posterior-predictive draws condition on them (see _push_rv_extra_fields).

src/hssm/base.py

@@ -311,14 +311,13 @@ def __init__(
         # region ===== Store the pre-built config =====
         self.model_config: BaseModelConfig = copy(model_config)
         self.model_config.response_domains = deepcopy(model_config.response_domains)
+        configured_response = self.model_config.response
+        assert configured_response is not None
+        self._obs_dim = len(configured_response)
         # endregion
 
         # region ===== Set up shortcuts so old code will work ======
-        self.response: list[str] = (  # type: ignore[assignment]
-            list(self.model_config.response)
-            if self.model_config.response is not None
-            else []
-        )
+        self.response: list[str] = list(configured_response)  # type: ignore[assignment]
         self.list_params = (
             list(self.model_config.list_params)
             if self.model_config.list_params is not None
@@ -391,6 +390,7 @@ def __init__(
             self.list_params,
             self.link,
             self._parent,
+            obs_dim=self._obs_dim,
         )
 
         # Targeted checks against the user's prior dict:
@@ -1277,18 +1277,7 @@ def sample_do(
         )
         do_dt = pm.sample_prior_predictive(model=do_model, draws=draws, **kwargs)
 
-        # clean up `rt,response_mean` to `v`
-        do_dt = self._drop_parent_str_from_datatree(dt=do_dt)
-
-        # rename otherwise inconsistent dims and coords
-        if "rt,response_extra_dim_0" in do_dt["prior_predictive"].dims:
-            do_dt["prior_predictive"] = do_dt["prior_predictive"].ds.rename_dims(
-                {"rt,response_extra_dim_0": "rt,response_dim"}
-            )
-        if "rt,response_extra_dim_0" in do_dt["prior_predictive"].coords:
-            do_dt["prior_predictive"] = do_dt["prior_predictive"].ds.rename_vars(
-                {"rt,response_extra_dim_0": "rt,response_dim"}
-            )
+        do_dt = self._clean_predictive_datatree(dt=do_dt)
 
         if return_model:
             return do_dt, do_model
@@ -1346,18 +1335,7 @@ def sample_prior_predictive(
                     continue
                 self._inference_obj[group_name] = prior_predictive[group_name]
 
-        # clean up `rt,response_mean` to `v`
-        dt = self._drop_parent_str_from_datatree(dt=self._inference_obj)
-
-        # rename otherwise inconsistent dims and coords
-        if "rt,response_extra_dim_0" in dt["prior_predictive"].dims:
-            dt["prior_predictive"] = dt["prior_predictive"].ds.rename_dims(
-                {"rt,response_extra_dim_0": "rt,response_dim"}
-            )
-        if "rt,response_extra_dim_0" in dt["prior_predictive"].coords:
-            dt["prior_predictive"] = dt["prior_predictive"].ds.rename_vars(
-                name_dict={"rt,response_extra_dim_0": "rt,response_dim"}
-            )
+        dt = self._clean_predictive_datatree(dt=self._inference_obj)
 
         # Update self._inference_obj to match the cleaned datatree
         self._inference_obj = dt
@@ -2022,8 +2000,8 @@ def _get_deterministic_var_names(self, dt) -> list[str]:
 
         return var_names
 
-    def _drop_parent_str_from_datatree(self, dt: DataTree | None) -> DataTree:
-        """Drop the parent_str variable from a DataTree object.
+    def _clean_predictive_datatree(self, dt: DataTree | None) -> DataTree:
+        """Normalize generated response names and dimensions.
 
         Parameters
         ----------
@@ -2032,20 +2010,32 @@ def _drop_parent_str_from_datatree(self, dt: DataTree | None) -> DataTree:
 
         Returns
         -------
-        xr.Dataset
+        DataTree
             The modified DataTree object.
         """
         if dt is None:
             raise ValueError("Please provide a DataTree (traces) object.")
-        else:
-            for group in dt.groups:
-                if group == "/":
-                    continue
-                if ("rt,response_mean" in dt[group].data_vars) and (
-                    self._parent not in dt[group].data_vars
-                ):
-                    dt[group] = dt[group].ds.rename({"rt,response_mean": self._parent})
-            return dt
+
+        response_mean = f"{self.response_str}_mean"
+        raw_response_dim = f"{self.response_str}_extra_dim_0"
+        response_dim = f"{self.response_str}_dim"
+        for group in dt.groups:
+            if group == "/":
+                continue
+            dataset = dt[group].ds
+            rename = {}
+            if response_mean in dataset.data_vars and self._parent not in dataset:
+                rename[response_mean] = self._parent
+            if group.rsplit("/", maxsplit=1)[-1] in {
+                "posterior_predictive",
+                "prior_predictive",
+            } and (
+                raw_response_dim in dataset.dims or raw_response_dim in dataset.coords
+            ):
+                rename[raw_response_dim] = response_dim
+            if rename:
+                dt[group] = dataset.rename(rename)
+        return dt
 
     def _postprocess_initvals_deterministic(
         self, initval_settings: dict = INITVAL_SETTINGS

src/hssm/distribution_utils/dist.py

@@ -7,6 +7,7 @@
 
 import logging
 from collections.abc import Callable
+from functools import partial
 from os import PathLike
 from typing import Any, Literal, Protocol, cast, get_args
 
@@ -195,6 +196,7 @@ def make_hssm_rv(
     list_params: list[str],
     lapse: bmb.Prior | float | None = None,
     is_choice_only: bool = False,
+    expected_obs_dim: int | None = None,
 ) -> type[RandomVariable]:
     """Build a RandomVariable Op according to the list of parameters.
 
@@ -208,14 +210,28 @@ def make_hssm_rv(
         A bmb.Prior object representing the lapse distribution.
     is_choice_only : bool
         Whether the model is a choice-only model.
+    expected_obs_dim : int, optional
+        Configured number of physical response columns. When supplied, it must
+        agree with the simulator's effective observation width.
 
     Returns
     -------
     type[RandomVariable]
         A class of RandomVariable that are to be used in a `pm.Distribution`.
     """
     simulator_fun_internal = get_simulator_fun_internal(simulator_fun)
-    model_name, choices, obs_dim_int = validate_simulator_fun(simulator_fun_internal)
+    model_name, choices, declared_obs_dim = validate_simulator_fun(
+        simulator_fun_internal
+    )
+    effective_obs_dim = 1 if is_choice_only else declared_obs_dim
+    comparison_obs_dim = (
+        declared_obs_dim if callable(simulator_fun) else effective_obs_dim
+    )
+    if expected_obs_dim is not None and comparison_obs_dim != expected_obs_dim:
+        raise ValueError(
+            f"The simulator observation width {comparison_obs_dim} does not match "
+            f"the configured response width {expected_obs_dim}."
+        )
 
     if lapse is not None and list_params[-1] != "p_outlier":
         list_params.append("p_outlier")
@@ -249,7 +265,7 @@ class HSSMRV(RandomVariable):
 
         # Override the output from ssm_simulator based on whether the model is
         # choice-only.
-        output = "()" if is_choice_only else f"({obs_dim_int})"
+        output = "()" if is_choice_only else f"({effective_obs_dim})"
         signature = f"{','.join(['()'] * len(list_params))}->{output}"
 
         dtype = "floatX"
@@ -319,7 +335,7 @@ def rng_fn(
                 size,
                 rng,
                 simulator_fun_internal,
-                obs_dim_int,
+                effective_obs_dim,
                 *args,
                 **kwargs,
             )
@@ -440,6 +456,7 @@ def make_distribution(
     fixed_vector_params: dict[str, np.ndarray] | None = None,
     params_is_trialwise: list[bool] | None = None,
     is_choice_only: bool = False,
+    expected_obs_dim: int | None = None,
 ) -> type[pm.Distribution]:
     """Make a `pymc.Distribution`.
 
@@ -486,6 +503,9 @@ def make_distribution(
         When ``None``, no graph-level broadcasting is applied.
     is_choice_only : optional
         Whether the model is a choice-only model.
+    expected_obs_dim : int, optional
+        Configured number of physical response columns. This is checked against
+        generated RandomVariable metadata.
 
     Returns
     -------
@@ -502,6 +522,7 @@ def make_distribution(
             list_params=list_params,
             lapse=lapse,
             is_choice_only=is_choice_only,
+            expected_obs_dim=expected_obs_dim,
         )
         rv_instance = random_variable()
     elif isinstance(rv, str):
@@ -510,6 +531,7 @@ def make_distribution(
             list_params=list_params,
             lapse=lapse,
             is_choice_only=is_choice_only,
+            expected_obs_dim=expected_obs_dim,
         )
         rv_instance = random_variable()
     else:
@@ -718,6 +740,7 @@ def make_family(
     parent: str = "v",
     likelihood_name: str = "SSM Likelihood",
     family_name="SSM Family",
+    obs_dim: int = 2,
 ) -> bmb.Family:
     """Build a family in bambi.
 
@@ -735,6 +758,8 @@ def make_family(
         the name of the likelihood function. Defaults to "SSM Likelihood".
     family_name
         the name of the family. Defaults to "SSM Family".
+    obs_dim
+        Number of physical response columns in posterior-predictive draws.
 
     Returns
     -------
@@ -745,17 +770,23 @@ def make_family(
         likelihood_name, parent=parent, params=list_params, dist=dist
     )
 
-    family = SSMFamily(family_name, likelihood=likelihood, link=link)
+    family = SSMFamily(
+        family_name,
+        likelihood=likelihood,
+        link=link,
+        obs_dim=obs_dim,
+    )
 
     return family
 
 
 class SSMFamily(bmb.Family):
     """Extends bmb.Family to get around the dimensionality mismatch."""
 
-    def create_extra_pps_coord(self):
-        """Create an extra dimension."""
-        return np.arange(2)
+    def __init__(self, name, likelihood, link, obs_dim: int = 2):
+        super().__init__(name, likelihood, link)
+        if obs_dim > 1:
+            self.create_extra_pps_coord = partial(np.arange, obs_dim)
 
 
 def make_likelihood_callable(

src/hssm/hssm.py

@@ -435,6 +435,7 @@ def dummy_simulator_func(*args, **kwargs):
                 list_params=self.list_params or [],
                 lapse=self.lapse,
                 is_choice_only=True,
+                expected_obs_dim=self._obs_dim,
             )
 
         self.data = typing_cast("pd.DataFrame", _rearrange_data(self.data))
@@ -469,4 +470,5 @@ def dummy_simulator_func(*args, **kwargs):
             params_is_trialwise=params_is_trialwise_base,
             # TODO: add to HSSMBase
             is_choice_only=self.is_choice_only,
+            expected_obs_dim=self._obs_dim,
         )

src/hssm/rl/rlssm.py

@@ -307,6 +307,7 @@ def _make_model_distribution(self) -> type[pm.Distribution]:
             extra_fields=extra_fields_data,
             params_is_trialwise=params_is_trialwise,
             is_choice_only=self.model_config.is_choice_only,
+            expected_obs_dim=self._obs_dim,
         )
 
 

src/hssm/utils.py

@@ -576,19 +576,35 @@ def predictive_dt_to_dataframe(
         The DataTree object to convert.
     predictive_group : Literal["posterior_predictive", "prior_predictive"]
         The predictive group to convert.
+    response_str
+        Comma-separated physical response column names in configured order.
+    response_dim
+        Name of the vector response coordinate. Scalar responses have no such
+        coordinate.
 
     Returns
     -------
     pd.DataFrame:
         A dataframe with the predictive samples.
     """
+    predictive = dt[predictive_group].ds[response_str]
     df = dt[predictive_group].ds.to_dataframe().reset_index(drop=False)
+    response_names = response_str.split(",")
+    if response_dim not in df.columns:
+        if any(dim not in {"chain", "draw", "__obs__"} for dim in predictive.dims):
+            raise ValueError(
+                f"Predictive response coordinate {response_dim!r} is missing."
+            )
+        return df.loc[:, ["chain", "draw", "__obs__", response_str]].rename(
+            columns={response_str: response_names[0]}
+        )
+
     df_wide = df.pivot_table(
         index=["chain", "draw", "__obs__"], columns=response_dim, values=response_str
     ).reset_index()
 
     df_wide.columns.name = None
-    df_wide = df_wide.rename(columns={0: "rt", 1: "response"})
+    df_wide = df_wide.rename(columns=dict(enumerate(response_names)))
     return df_wide
 
 

tests/distribution_utils/test_distribution_utils.py

@@ -113,8 +113,15 @@ def test_make_distribution_forwards_choice_only_to_generated_rv(monkeypatch, rv)
     """Generated RVs must keep the choice-only support-shape contract."""
     captured = {}
 
-    def fake_make_hssm_rv(simulator_fun, list_params, lapse=None, is_choice_only=False):
+    def fake_make_hssm_rv(
+        simulator_fun,
+        list_params,
+        lapse=None,
+        is_choice_only=False,
+        expected_obs_dim=None,
+    ):
         captured["is_choice_only"] = is_choice_only
+        captured["expected_obs_dim"] = expected_obs_dim
 
         class FakeRV:
             def __call__(self):
@@ -132,6 +139,7 @@ def __call__(self):
     )
 
     assert captured["is_choice_only"] is True
+    assert captured["expected_obs_dim"] is None
 
 
 @pytest.mark.slow

tests/rl/test_choice_only_rl.py

@@ -266,6 +266,7 @@ def _make_shell_rlssm(config, lapse):
     model = _RLSSM.__new__(_RLSSM)
     model.list_params = [*config.list_params, "p_outlier"]
     model.model_config = config
+    model._obs_dim = len(config.response)
     model.bounds = dict(config.bounds)
     model.lapse = lapse
     model.data = pd.DataFrame({"response": [0, 1], "rt": [0.5, 0.6]})