From ebf9463f1174b2ec5dc3b9e88f1c562977522f64 Mon Sep 17 00:00:00 2001 From: Alexander Fengler Date: Sun, 30 Aug 2026 20:01:44 -0400 Subject: [PATCH 1/2] test: prove mixed-response lifecycle (#1265) --- tests/test_observation_dimensions.py | 150 ++++++++++++++++++++++++--- 1 file changed, 137 insertions(+), 13 deletions(-) diff --git a/tests/test_observation_dimensions.py b/tests/test_observation_dimensions.py index a094e43f2..ba8a4b202 100644 --- a/tests/test_observation_dimensions.py +++ b/tests/test_observation_dimensions.py @@ -14,7 +14,7 @@ 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 @@ -22,7 +22,20 @@ def simulator(theta, random_state, n_replicas, **kwargs): 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,7 +48,11 @@ 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( @@ -43,26 +60,53 @@ def _synthetic_model( *, 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(): From e49290861e041efb4077f146abf9a7f48d3fe4fd Mon Sep 17 00:00:00 2001 From: Alexander Fengler Date: Sun, 30 Aug 2026 20:05:19 -0400 Subject: [PATCH 2/2] docs: document canonical response domains (#1265) --- docs/api/model_config.md | 93 ++++++++++++++++++++++++++++++++ docs/api/model_registry.md | 4 ++ docs/how_to/external_trainers.md | 1 + docs/reference/index.md | 4 +- mkdocs.yml | 2 +- src/hssm/hssm.py | 19 ++++--- 6 files changed, 110 insertions(+), 13 deletions(-) diff --git a/docs/api/model_config.md b/docs/api/model_config.md index 03a1fe1cd..d011190d9 100644 --- a/docs/api/model_config.md +++ b/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 diff --git a/docs/api/model_registry.md b/docs/api/model_registry.md index 49de633fd..c0dc938de 100644 --- a/docs/api/model_registry.md +++ b/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 diff --git a/docs/how_to/external_trainers.md b/docs/how_to/external_trainers.md index c53766517..753cea3b0 100644 --- a/docs/how_to/external_trainers.md +++ b/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 diff --git a/docs/reference/index.md b/docs/reference/index.md index 0a93d474e..1345bcb1c 100644 --- a/docs/reference/index.md +++ b/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. diff --git a/mkdocs.yml b/mkdocs.yml index edd26b609..774c08723 100644 --- a/mkdocs.yml +++ b/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 diff --git a/src/hssm/hssm.py b/src/hssm/hssm.py index 4dae1d7eb..491f7c956 100644 --- a/src/hssm/hssm.py +++ b/src/hssm/hssm.py @@ -57,8 +57,9 @@ 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", @@ -66,12 +67,10 @@ class HSSM(HSSMBase): 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