Prove and document mixed response domains

docs/api/model_config.md

@@ -1 +1,94 @@
+# Model configuration and response domains
+
+`hssm.ModelConfig.response` is the ordered sequence of **physical DataFrame
+columns** that make up one observation. The names are never semantic aliases:
+each name must be present as its own scalar column in `data`. Built-in models
+with one non-RT response keep the physical column name `response`; declaring a
+different domain does not rename that column.
+
+`response_domains` annotates every non-RT response column by that exact physical
+name. Do not include `rt` in this mapping: response-time positivity, deadline,
+and missing-data handling continue to use the existing RT contract.
+
+## Domain specifications
+
+| Kind | Required metadata | Accepted values |
+|------|-------------------|-----------------|
+| `categorical` | `values` | A non-empty list or tuple of distinct integer labels. Observations must match one of them exactly. |
+| `continuous` | Optional `bounds` | Without bounds, any finite value. With finite, increasing bounds, both endpoints are included: `[lower, upper]`. |
+| `circular` | `bounds` | Finite, increasing bounds defining a half-open interval: `[lower, upper)`. |
+
+The mapping keys must match the non-RT entries in `response` exactly. HSSM
+normalizes mapping insertion order to the order declared by `response`.
+
+## Custom RT and two-coordinate example
+
+The domain mapping below deliberately lists `azimuth` before `polar`. The
+resolved order is still `polar`, then `azimuth`, because `response` is the sole
+source of truth for observation order.
+
+```python
+import numpy as np
+import pandas as pd
+
+from hssm import ModelConfig
+
+model_config = ModelConfig(
+    response=("rt", "polar", "azimuth"),
+    list_params=["v"],
+    response_domains={
+        "azimuth": {
+            "kind": "circular",
+            "bounds": (-np.pi, np.pi),
+        },
+        "polar": {
+            "kind": "continuous",
+            "bounds": (0.0, np.pi),
+        },
+    },
+)
+
+data = pd.DataFrame(
+    {
+        "rt": [0.42, 0.73],
+        "polar": [0.80, 1.20],
+        "azimuth": [-2.40, 0.90],
+    }
+)
+```
+
+HSSM uses the normalized `("rt", "polar", "azimuth")` order for likelihood
+observations and predictive output labels. Custom predictive generators must
+return exactly `len(response)` scalars in that declared order. HSSM checks a
+callable generator's declared width against `len(response)`, but it cannot infer
+the semantic meaning of each position. Authors who pass a prebuilt random
+variable or distribution own both its width and order contract.
+
+## Migrating from `choices`
+
+`choices` remains a legacy shorthand only for a configuration with exactly one
+categorical non-RT response. Canonical `ModelConfig` callers declare
+`response_domains` and omit `choices`; never provide both. The current
+`register_model` signature retains the compatibility argument, so canonical
+registrations pass `choices=None` alongside `response_domains`.
+
+After resolution, a derived `choices` view exists only when there is exactly one
+response domain and it is categorical. Continuous, circular, mixed, and
+multiple-domain configurations have no derived `choices` view.
+
+## Current limits
+
+- In a multicolumn RT layout, `rt` must appear exactly once and first.
+- A choice-only model supports exactly one scalar non-RT response column.
+- RT-less multicolumn responses are not supported.
+- Each coordinate must occupy its own scalar DataFrame column; packed array- or
+  object-valued response cells are not supported.
+- No built-in model currently combines different response-domain kinds.
+- `p_outlier` must be `None` or `0` outside the established one-categorical
+  response layouts.
+- Plotting, KDE utilities, and the legacy `hssm.simulate_data` interface are not
+  generalized to wider mixed-response layouts.
+
+## API reference
+
 ::: hssm.ModelConfig

docs/api/model_registry.md

@@ -1,5 +1,9 @@
 # Model discovery and registration
 
+Custom registry entries use the same configuration contract as direct model
+construction. Read [Model configuration and response domains](model_config.md)
+before registering a model; this page documents the registry functions.
+
 ## `hssm.list_models`
 
 ::: hssm.list_models

docs/how_to/external_trainers.md

@@ -37,6 +37,7 @@ and transfer the pattern:
 
 ## See also
 
+- [Model configuration and response domains](../api/model_config.md) — declare the physical observation columns and their domains
 - [The ONNX likelihood contract](custom_onnx_likelihoods.md) — the rules an ONNX artifact must satisfy
 - [Custom models from JAX callables](../tutorials/jax_callable_contribution_onnx_example.ipynb) — the same callable gesture without an external trainer
 - [Likelihood kinds in HSSM](../explanations/likelihoods.md) — where `approx_differentiable` fits among the likelihood kinds

docs/reference/index.md

@@ -11,8 +11,8 @@ This section records the public HSSM interfaces and project metadata. Use the
   attentional drift-diffusion interface.
 - [`hssm.Param`](../api/param.md), [`hssm.Prior`](../api/prior.md), and
   [`hssm.Link`](../api/link.md) describe parameter formulas, priors, and links.
-- [`hssm.ModelConfig`](../api/model_config.md) describes registered model
-  configurations.
+- [Model configuration and response domains](../api/model_config.md) defines
+  physical response-column ordering and canonical domain metadata.
 - [`hssm.rl`](../api/rl.md) contains the reinforcement-learning interfaces.
 - [Built-in models and likelihoods](models-and-likelihoods.md) records every
   `hssm.HSSM(model=...)` name, configured likelihood kind, parameter, and choice.

mkdocs.yml

@@ -92,7 +92,7 @@ nav:
           - hssm.HSSM: api/hssm.md
           - hssm.aDDM and hssm.aDDMConfig: api/addm.md
           - hssm.Param: api/param.md
-          - hssm.ModelConfig: api/model_config.md
+          - Model configuration and response domains: api/model_config.md
           - hssm.Prior: api/prior.md
           - hssm.Link: api/link.md
           - hssm.rl: api/rl.md

src/hssm/hssm.py

@@ -57,21 +57,20 @@ class HSSM(HSSMBase):
     Parameters
     ----------
     data
-        A pandas DataFrame with the minimum requirements of containing the data with the
-        columns "rt" and "response".
+        A pandas DataFrame containing the physical response columns declared in
+        `hssm.ModelConfig.response`, plus any configured extra fields. Built-in RT
+        models typically use the columns "rt" and "response".
     model
         The name of the model to use. Currently supported models are "ddm", "ddm_sdv",
         "full_ddm", "angle", "levy", "ornstein", "weibull", "race_no_bias_angle_4",
         "ddm_seq2_no_bias", "gamma_drift". If any other string is passed, the model
         will be considered custom, in which case all `model_config`, `loglik`, and
         `loglik_kind` have to be provided by the user.
     choices : optional
-        When an `int`, the number of choices that the participants can make. If `2`, the
-        choices are [-1, 1] by default. If anything greater than `2`, the choices are
-        [0, 1, ..., n_choices - 1] by default. If a `list` is provided, it should be the
-        list of choices that the participants can make. Defaults to `2`. If any value
-        other than the choices provided is found in the "response" column of the data,
-        an error will be raised.
+        Legacy shorthand for the allowed integer labels when a model has exactly one
+        categorical non-RT response. Canonical custom configurations declare domains
+        through `hssm.ModelConfig.response_domains` and leave `choices` unset; do not
+        provide both.
     include : optional
         A list of dictionaries specifying parameter specifications to include in the
         model. If left unspecified, defaults will be used for all parameter
@@ -223,8 +222,8 @@ class HSSM(HSSMBase):
     Attributes
     ----------
     data
-        A pandas DataFrame with at least two columns of "rt" and "response" indicating
-        the response time and responses.
+        The pandas DataFrame containing the physical response columns declared by the
+        model configuration.
     list_params
         The list of strs of parameter names.
     model_name

tests/test_observation_dimensions.py

@@ -14,15 +14,28 @@
 
 
 def _synthetic_simulator(obs_dim: int, *, scalar_output: bool = False) -> Callable:
-    """Return a deterministic-contract simulator with seeded random values."""
+    """Return a seeded simulator with domain-valid, axis-coded values."""
 
     def simulator(theta, random_state, n_replicas, **kwargs):
         del kwargs
         theta_array = np.asarray(theta)
         n_rows = (
             n_replicas if theta_array.ndim == 1 else theta_array.shape[0] * n_replicas
         )
-        values = np.random.default_rng(random_state).normal(size=(n_rows, obs_dim))
+        rng = np.random.default_rng(random_state)
+        if obs_dim == 1:
+            values = rng.integers(0, 2, size=(n_rows, 1)).astype(float)
+        else:
+            locations = {
+                2: (0.30, -0.50),
+                3: (0.30, 0.80, -2.40),
+                4: (0.30, 0.20, 0.80, -2.40),
+            }[obs_dim]
+            values = np.asarray(locations) + rng.uniform(
+                0.0,
+                0.05,
+                size=(n_rows, obs_dim),
+            )
         return values[:, 0] if scalar_output else values
 
     simulator.model_name = f"synthetic_{obs_dim}d"  # type: ignore[attr-defined]
@@ -35,34 +48,65 @@ def _synthetic_logp(data, v):
     """Return one finite log-likelihood value per observation."""
     if data.ndim == 1:
         return -pt.square(data - v)
-    return -pt.sum(pt.square(data - pt.shape_padright(v)), axis=-1)
+    weights = pt.arange(1, data.shape[-1] + 1, dtype=data.dtype)
+    return -pt.sum(
+        pt.square(data - pt.shape_padright(v)) * weights,
+        axis=-1,
+    )
 
 
 def _synthetic_model(
     configured_obs_dim: int,
     *,
     simulator_obs_dim: int | None = None,
 ) -> hssm.HSSM:
-    """Build a small custom model with scalar physical response columns."""
+    """Build a small custom model with ordered physical response columns."""
     simulator_obs_dim = simulator_obs_dim or configured_obs_dim
     if configured_obs_dim == 1:
         response = ("response",)
         data = pd.DataFrame({"response": [0, 1, 0, 1, 0]})
         response_domains = {"response": {"kind": "categorical", "values": (0, 1)}}
-    else:
-        response = ("rt",) + tuple(
-            f"response_{index}" for index in range(1, configured_obs_dim)
+    elif configured_obs_dim == 2:
+        response = ("rt", "response_1")
+        data = pd.DataFrame(
+            {
+                "rt": np.linspace(0.3, 0.7, 5),
+                "response_1": np.linspace(-0.5, 0.5, 5),
+            }
+        )
+        response_domains = {"response_1": {"kind": "continuous", "bounds": (-1.0, 1.0)}}
+    elif configured_obs_dim in (3, 4):
+        has_confidence = configured_obs_dim == 4
+        response = (
+            ("rt", "confidence", "polar", "azimuth")
+            if has_confidence
+            else ("rt", "polar", "azimuth")
         )
         data = pd.DataFrame(
             {
                 "rt": np.linspace(0.3, 0.7, 5),
-                **{column: np.linspace(-0.5, 0.5, 5) for column in response[1:]},
+                **({"confidence": np.linspace(0.0, 1.0, 5)} if has_confidence else {}),
+                "polar": np.linspace(0.0, np.pi, 5),
+                "azimuth": np.linspace(
+                    -np.pi,
+                    np.nextafter(np.pi, -np.inf),
+                    5,
+                ),
             }
         )
         response_domains = {
-            column: {"kind": "continuous", "bounds": (-1.0, 1.0)}
-            for column in response[1:]
+            **(
+                {"confidence": {"kind": "continuous", "bounds": (0.0, 1.0)}}
+                if has_confidence
+                else {}
+            ),
+            "polar": {"kind": "continuous", "bounds": (0.0, np.pi)},
+            "azimuth": {"kind": "circular", "bounds": (-np.pi, np.pi)},
         }
+    else:
+        raise ValueError(
+            f"Unsupported synthetic observation width: {configured_obs_dim}"
+        )
 
     return hssm.HSSM(
         data=data,
@@ -144,6 +188,44 @@ def test_model_rejects_callable_width_mismatch_before_building_distribution():
         _synthetic_model(3, simulator_obs_dim=2)
 
 
+@pytest.mark.parametrize(
+    ("expected_response", "expected_domain_kinds"),
+    [
+        (("rt", "polar", "azimuth"), ("continuous", "circular")),
+        (
+            ("rt", "confidence", "polar", "azimuth"),
+            ("continuous", "continuous", "circular"),
+        ),
+    ],
+)
+def test_mixed_response_model_compiles_ordered_logp(
+    expected_response,
+    expected_domain_kinds,
+):
+    """Mixed physical domains retain order through likelihood construction."""
+    obs_dim = len(expected_response)
+    model = _synthetic_model(obs_dim)
+
+    assert model._obs_dim == obs_dim
+    assert tuple(model.response) == expected_response
+    assert tuple(model.response_domains) == expected_response[1:]
+    assert tuple(domain["kind"] for domain in model.response_domains.values()) == (
+        expected_domain_kinds
+    )
+    point = model.initial_point(transformed=True)
+    logp = model.compile_logp(
+        keep_transformed=True,
+        vars=model.pymc_model.observed_RVs,
+        sum=False,
+    )
+    actual = logp(point)[0]
+    observed = model.data.loc[:, expected_response].to_numpy()
+    weights = np.arange(1, obs_dim + 1)
+    expected = -np.sum(np.square(observed - point["v"]) * weights, axis=1)
+    np.testing.assert_allclose(actual, expected)
+    assert np.isfinite(actual).all()
+
+
 @pytest.mark.parametrize("obs_dim", [1, 2, 3, 4])
 def test_predictive_shapes_follow_configured_response_width(obs_dim):
     """Prior and posterior predictions follow the physical response width."""
@@ -169,15 +251,26 @@ def test_predictive_shapes_follow_configured_response_width(obs_dim):
         assert not hasattr(model.family, "create_extra_pps_coord")
         assert prior_values.dims == ("chain", "draw", "__obs__")
         assert posterior_values.dims == ("chain", "draw", "__obs__")
+        assert prior_values.shape == (1, 3, len(model.data))
+        assert posterior_values.shape == (1, 2, len(model.data))
     else:
         np.testing.assert_array_equal(
             model.family.create_extra_pps_coord(), np.arange(obs_dim)
         )
-        assert prior_values.dims[-1] == response_dim
-        assert posterior_values.dims[-1] == response_dim
+        assert prior_values.dims == ("chain", "draw", "__obs__", response_dim)
+        assert posterior_values.dims == (
+            "chain",
+            "draw",
+            "__obs__",
+            response_dim,
+        )
+        assert prior_values.shape == (1, 3, len(model.data), obs_dim)
+        assert posterior_values.shape == (1, 2, len(model.data), obs_dim)
         np.testing.assert_array_equal(
             posterior_values.coords[response_dim], np.arange(obs_dim)
         )
+    assert np.isfinite(prior_values).all()
+    assert np.isfinite(posterior_values).all()
 
 
 def test_safe_mode_preserves_width_four_draw_and_response_coordinates():
@@ -235,6 +328,10 @@ def test_choice_only_string_keeps_scalar_legacy_width():
 def test_sample_do_dataframe_uses_physical_response_order(obs_dim):
     """Intervention samples expose each configured physical response column."""
     model = _synthetic_model(obs_dim)
+    expected_response = {
+        1: ("response",),
+        4: ("rt", "confidence", "polar", "azimuth"),
+    }[obs_dim]
     predictive = model.sample_do(params={"v": 0.5}, draws=3)
     frame = hssm.utils.predictive_dt_to_dataframe(
         predictive,
@@ -243,8 +340,35 @@ def test_sample_do_dataframe_uses_physical_response_order(obs_dim):
         response_dim=f"{model.response_str}_dim",
     )
 
-    assert list(frame.columns) == ["chain", "draw", "__obs__", *model.response]
+    assert tuple(model.response) == expected_response
+    assert list(frame.columns) == ["chain", "draw", "__obs__", *expected_response]
     assert len(frame) == 3 * len(model.data)
+    values = predictive["prior_predictive"][model.response_str]
+    expected_dims = ("chain", "draw", "__obs__")
+    if obs_dim > 1:
+        expected_dims += (f"{model.response_str}_dim",)
+    assert values.dims == expected_dims
+    assert values.shape == (
+        1,
+        3,
+        len(model.data),
+        *(() if obs_dim == 1 else (obs_dim,)),
+    )
+    np.testing.assert_array_equal(
+        frame[list(expected_response)].to_numpy(),
+        values.to_numpy().reshape(-1, obs_dim),
+    )
+    if obs_dim == 1:
+        assert set(frame["response"]) <= {0, 1}
+    else:
+        expected_intervals = {
+            "rt": (0.30, 0.35),
+            "confidence": (0.20, 0.25),
+            "polar": (0.80, 0.85),
+            "azimuth": (-2.40, -2.35),
+        }
+        for column, (lower, upper) in expected_intervals.items():
+            assert frame[column].between(lower, upper, inclusive="left").all()
 
 
 def test_scalar_dataframe_ignores_deadline_technical_suffix():