diff --git a/design/response_domains.md b/design/response_domains.md index 7a243bb93..c241a500c 100644 --- a/design/response_domains.md +++ b/design/response_domains.md @@ -83,9 +83,15 @@ part of this contract. 3. HSSM has one normalization boundary. Legacy `choices` can create a categorical domain only when there is exactly one non-RT response column. Multiple response columns require canonical metadata. -4. Canonical metadata combined with explicitly supplied `choices` is ambiguous and - fails. Resolved configs have no implicit choice default: built-in factories supply - their existing labels explicitly. +4. Canonical metadata combined with explicitly supplied `choices` at a raw input + boundary is ambiguous and fails. A resolved config may retain an exactly matching + derived `choices` view so ordinary dataclass copying and validation remain + idempotent; conflicting values fail. Resolved configs have no implicit choice + default: built-in factories supply their existing labels explicitly. Resolved + configs are construction snapshots, not mutable domain registries. Directly + mutating nested canonical metadata is unsupported and a later validation fails if + the derived view no longer matches; construct a replacement with canonical metadata + and `choices=None` instead. 5. On `main`, `choices` is derived only for exactly one categorical domain. It is `None` for continuous, circular, or multiple domains, including multiple categorical domains. diff --git a/src/hssm/_types.py b/src/hssm/_types.py index 579f7a348..31f363edf 100644 --- a/src/hssm/_types.py +++ b/src/hssm/_types.py @@ -1,7 +1,7 @@ """Type definitions for the HSSM package.""" from os import PathLike -from typing import Any, Callable, Literal, Optional, TypedDict, Union +from typing import Any, Callable, Literal, NotRequired, Optional, TypedDict, Union import bambi as bmb import numpy as np @@ -48,12 +48,21 @@ class LoglikConfig(TypedDict): LoglikConfigs = dict[LoglikKind, LoglikConfig] +class ResponseDomainSpec(TypedDict): + """Canonical metadata for one physical response column.""" + + kind: Literal["categorical", "continuous", "circular"] + values: NotRequired[tuple[int, ...]] + bounds: NotRequired[tuple[float, float]] + + class DefaultConfig(TypedDict): """Type for the value of DefaultConfig.""" response: list[str] list_params: list[str] - choices: list[int] + choices: NotRequired[list[int]] + response_domains: NotRequired[dict[str, ResponseDomainSpec]] description: Optional[str] likelihoods: LoglikConfigs diff --git a/src/hssm/addm/config.py b/src/hssm/addm/config.py index c7823f61e..fe7ecd4cb 100644 --- a/src/hssm/addm/config.py +++ b/src/hssm/addm/config.py @@ -13,11 +13,12 @@ """ from collections.abc import Callable +from copy import deepcopy from dataclasses import dataclass, field, fields from typing import Any from .._types import LoglikKind -from ..config import BaseModelConfig +from ..config import BaseModelConfig, _resolve_response_domains from .attention_process import resolve_attention_process @@ -60,7 +61,7 @@ class aDDMConfig(BaseModelConfig): model_name: str = "addm" description: str | None = "Attentional Drift Diffusion Model" response: list[str] = field(default_factory=lambda: ["rt", "response"]) - choices: tuple[int, ...] = (-1, 1) + choices: tuple[int, ...] | None = (-1, 1) list_params: list[str] = field( default_factory=lambda: ["eta", "kappa", "a", "b", "x0", "t"] ) @@ -90,10 +91,17 @@ class aDDMConfig(BaseModelConfig): continuation_mode: str = "prolong_last_fixation" continuation_params: dict | None = None # ``loglik`` and ``backend`` are inherited (default ``None``) and injected by - # ``aDDM.__init__`` via ``dataclasses.replace`` in Commit 4 — not redeclared here. + # ``aDDM.__init__`` via ``dataclasses.replace`` — not redeclared here. def validate(self) -> None: """Validate the configuration (mirrors ``RLSSMConfig.validate``).""" + response_domains, choices = _resolve_response_domains( + self.response, self.response_domains, self.choices + ) + if choices is None: + raise ValueError("aDDM requires one categorical response domain.") + self.response_domains = response_domains + self.choices = choices if not self.list_params: raise ValueError("Please provide `list_params` in the configuration.") # Raises ValueError (unknown name) / TypeError (bad type) on failure. @@ -128,5 +136,15 @@ def get_defaults(self, param: str) -> tuple[None, tuple[float, float] | None]: @classmethod def from_addm_dict(cls, config_dict: dict[str, Any]) -> "aDDMConfig": """Build an ``aDDMConfig`` from a dict, ignoring unknown keys.""" + if ( + config_dict.get("response_domains") is not None + and config_dict.get("choices") is not None + ): + raise ValueError( + "Provide either `response_domains` or legacy `choices`, not both." + ) field_names = {f.name for f in fields(cls)} - return cls(**{k: v for k, v in config_dict.items() if k in field_names}) + init_kwargs = {k: v for k, v in config_dict.items() if k in field_names} + if init_kwargs.get("response_domains") is not None: + init_kwargs["response_domains"] = deepcopy(init_kwargs["response_domains"]) + return cls(**init_kwargs) diff --git a/src/hssm/base.py b/src/hssm/base.py index 59572b9d9..1291c56f2 100644 --- a/src/hssm/base.py +++ b/src/hssm/base.py @@ -10,7 +10,7 @@ import logging import warnings from abc import ABC, abstractmethod -from copy import deepcopy +from copy import copy, deepcopy from os import PathLike from pathlib import Path from typing import TYPE_CHECKING, Any, Callable, Literal, Union, cast @@ -309,7 +309,8 @@ def __init__( self.initval_jitter = initval_jitter # region ===== Store the pre-built config ===== - self.model_config: BaseModelConfig = model_config + self.model_config: BaseModelConfig = copy(model_config) + self.model_config.response_domains = deepcopy(model_config.response_domains) # endregion # region ===== Set up shortcuts so old code will work ====== @@ -324,6 +325,7 @@ def __init__( else None ) self.choices = self.model_config.choices # type: ignore[assignment] + self.response_domains = self.model_config.response_domains or {} self.model_name = self.model_config.model_name self.loglik = self.model_config.loglik self.loglik_kind = self.model_config.loglik_kind @@ -333,13 +335,6 @@ def __init__( # TODO: add to HSSMBase self.is_choice_only: bool = self.model_config.is_choice_only - if self.choices is None: - raise ValueError( - "`choices` must be provided either in `model_config` or as an argument." - ) - - self._validate_choices() - # region Avoid mypy error later (None.append). Should list_params be Optional? if self.list_params is None: raise ValueError( @@ -347,7 +342,7 @@ def __init__( ) # endregion - self.n_choices = len(self.choices) # type: ignore[arg-type] + self.n_choices = len(self.choices) if self.choices is not None else None self._pre_check_data_sanity() @@ -565,6 +560,23 @@ def _store_init_args( exclude_keys = {"self", "kwargs", "__class__"} result = {k: v for k, v in local_vars.items() if k not in exclude_keys} result.update(extra_kwargs) + model_config = result.get("model_config") + if ( + isinstance(model_config, dict) + and model_config.get("response_domains") is not None + ): + model_config = model_config.copy() + model_config["response_domains"] = deepcopy( + model_config["response_domains"] + ) + result["model_config"] = model_config + elif ( + model_config is not None + and getattr(model_config, "response_domains", None) is not None + ): + model_config = copy(model_config) + model_config.response_domains = deepcopy(model_config.response_domains) + result["model_config"] = model_config return result def find_MAP(self, **kwargs): @@ -1958,8 +1970,20 @@ def _check_lapse(self, lapse): + "parameter is not None" ) if self.has_lapse: + domain = ( + next(iter(self.response_domains.values())) + if len(self.response_domains) == 1 + else None + ) + supports_lapse = domain is not None and domain["kind"] == "categorical" + if not supports_lapse: + raise ValueError( + "`p_outlier` is supported only for one categorical response " + "column. Set `p_outlier=None` or `p_outlier=0` for this model." + ) if lapse is None: if self.is_choice_only: + assert self.n_choices is not None self.lapse = 1 / self.n_choices else: self.lapse = bmb.Prior("Uniform", lower=0.0, upper=20.0) diff --git a/src/hssm/config.py b/src/hssm/config.py index 11cbc6a96..1714c8cfc 100644 --- a/src/hssm/config.py +++ b/src/hssm/config.py @@ -3,14 +3,17 @@ # This is necessary to enable forward looking from __future__ import annotations +import math from abc import ABC, abstractmethod +from collections.abc import Mapping from copy import deepcopy from dataclasses import dataclass, field +from numbers import Integral, Real from typing import TYPE_CHECKING, Any, Literal, Union, cast, get_args from bambi import Prior -from ._types import LogLik, LoglikKind, SupportedModels +from ._types import LogLik, LoglikKind, ResponseDomainSpec, SupportedModels from .defaults import ( default_model_config, ) @@ -44,7 +47,7 @@ class BaseModelConfig(ABC): # Data specification response: list[str] | None = field(default_factory=DEFAULT_SSM_OBSERVED_DATA.copy) - choices: tuple[int, ...] | None = DEFAULT_SSM_CHOICES + choices: tuple[int, ...] | None = None # Parameter specification list_params: list[str] | None = None @@ -61,6 +64,11 @@ class BaseModelConfig(ABC): # Random variable (simulator) for posterior predictive sampling rv: Any | None = None + # Canonical per-column response metadata. Appended for positional compatibility. + response_domains: dict[str, ResponseDomainSpec] | None = field( + default=None, kw_only=True + ) + @abstractmethod def validate(self) -> None: """Validate configuration. Must be implemented by subclasses.""" @@ -142,7 +150,14 @@ def from_defaults( model_name=model_name, loglik_kind=kind, response=list(default_config["response"]), - choices=tuple(default_config["choices"]), + choices=( + tuple(default_config["choices"]) + if default_config.get("choices") is not None + else None + ), + response_domains=deepcopy( + default_config.get("response_domains") + ), list_params=default_config["list_params"], description=default_config["description"], **loglik_config, @@ -170,7 +185,14 @@ def from_defaults( model_name=model_name, loglik_kind=loglik_kind, response=list(default_config["response"]), - choices=tuple(default_config["choices"]), + choices=( + tuple(default_config["choices"]) + if default_config.get("choices") is not None + else None + ), + response_domains=deepcopy( + default_config.get("response_domains") + ), list_params=default_config["list_params"], description=default_config["description"], **loglik_config, @@ -179,7 +201,12 @@ def from_defaults( model_name=model_name, loglik_kind=loglik_kind, response=list(default_config["response"]), - choices=tuple(default_config["choices"]), + choices=( + tuple(default_config["choices"]) + if default_config.get("choices") is not None + else None + ), + response_domains=deepcopy(default_config.get("response_domains")), list_params=default_config["list_params"], description=default_config["description"], ) @@ -213,6 +240,10 @@ def update_choices(self, choices: tuple[int, ...] | None) -> None: """ if choices is None: return + if self.response_domains is not None: + raise ValueError( + "Provide either `response_domains` or legacy `choices`, not both." + ) self.choices = choices @@ -228,7 +259,18 @@ def update_config(self, user_config: ModelConfig) -> None: self.response = list(user_config.response) # type: ignore[assignment] if user_config.list_params is not None: self.list_params = user_config.list_params - if user_config.choices is not None: + if user_config.response_domains is not None: + if user_config.choices is not None: + raise ValueError( + "Provide either `response_domains` or legacy `choices`, not both." + ) + self.response_domains = deepcopy(user_config.response_domains) + self.choices = None + elif user_config.choices is not None: + if self.response_domains is not None: + raise ValueError( + "Provide either `response_domains` or legacy `choices`, not both." + ) self.choices = user_config.choices if user_config.rv is not None: self.rv = user_config.rv @@ -247,10 +289,11 @@ def validate(self) -> None: """Ensure that mandatory fields are not None.""" if self.response is None: raise ValueError("Please provide `response` columns in the configuration.") + self.response_domains, self.choices = _resolve_response_domains( + self.response, self.response_domains, self.choices + ) if self.list_params is None: raise ValueError("Please provide `list_params`.") - if self.choices is None: - raise ValueError("Please provide `choices`.") if self.loglik is None: raise ValueError("Please provide a log-likelihood function via `loglik`.") if self.loglik_kind == "approx_differentiable" and self.backend is None: @@ -304,7 +347,7 @@ def _build_model_config( if model not in get_args(SupportedModels): if choices is not None: config.update_choices(choices) - elif model in ssms_model_config: + elif config.response_domains is None and model in ssms_model_config: config.update_choices(ssms_model_config[model]["choices"]) _logger.info( "choices argument passed as None, " @@ -331,6 +374,7 @@ class ModelConfig: backend: Literal["jax", "pytensor"] | None = None rv: RandomVariable | None = None extra_fields: list[str] | None = None + response_domains: dict[str, ResponseDomainSpec] | None = None def _normalize_model_config_with_choices( @@ -353,6 +397,16 @@ def _normalize_model_config_with_choices( else: mc = model_config.copy() + if mc.get("response_domains") is not None: + mc["response_domains"] = deepcopy(mc["response_domains"]) + + if mc.get("response_domains") is not None and ( + mc.get("choices") is not None or choices is not None + ): + raise ValueError( + "Provide either `response_domains` or legacy `choices`, not both." + ) + # Coerce any existing choices on the input to a tuple for immutability if mc.get("choices") is not None: mc["choices"] = tuple(mc["choices"]) @@ -374,3 +428,130 @@ def _normalize_model_config_with_choices( mc["choices"] = tuple(choices) return ModelConfig(**{k: v for k, v in mc.items() if v is not None}) + + +def _resolve_response_domains( + response: list[str] | tuple[str, ...] | None, + response_domains: Mapping[str, Mapping[str, object]] | None, + choices: list[int] | tuple[int, ...] | None, +) -> tuple[dict[str, ResponseDomainSpec], tuple[int, ...] | None]: + """Return detached, response-ordered domain metadata and legacy choices.""" + if not response: + raise ValueError("Please provide at least one `response` column.") + if any(not isinstance(name, str) or not name for name in response): + raise ValueError("Every `response` column name must be a non-empty string.") + if len(set(response)) != len(response): + raise ValueError("`response` column names must be unique.") + + rt_count = response.count("rt") + if rt_count: + if rt_count != 1 or response[0] != "rt": + raise ValueError("RT-based models require `rt` exactly once at index zero.") + elif len(response) != 1: + raise ValueError( + "Models without `rt` currently support exactly one response column." + ) + + response_columns = [name for name in response if name != "rt"] + if not response_columns: + raise ValueError("At least one non-RT response column is required.") + + if response_domains is None: + if len(response_columns) != 1 or choices is None: + raise ValueError( + "Provide `response_domains`; legacy `choices` can describe only one " + "non-RT response column." + ) + raw_domains: Mapping[str, Mapping[str, object]] = { + response_columns[0]: {"kind": "categorical", "values": choices} + } + else: + if not isinstance(response_domains, Mapping): + raise ValueError("`response_domains` must be a mapping.") + missing = set(response_columns) - set(response_domains) + extra = set(response_domains) - set(response_columns) + if missing or extra: + details = [] + if missing: + details.append(f"missing {sorted(missing)}") + if extra: + details.append(f"unexpected {sorted(extra)}") + raise ValueError( + "`response_domains` keys must match non-RT response columns: " + + ", ".join(details) + + "." + ) + raw_domains = response_domains + + resolved: dict[str, ResponseDomainSpec] = {} + for column in response_columns: + raw_spec = raw_domains[column] + if not isinstance(raw_spec, Mapping): + raise ValueError(f"Response domain for {column!r} must be a mapping.") + kind = raw_spec.get("kind") + if kind not in {"categorical", "continuous", "circular"}: + raise ValueError( + f"Response domain for {column!r} has invalid kind {kind!r}." + ) + + allowed = {"kind", "values"} if kind == "categorical" else {"kind", "bounds"} + unknown = set(raw_spec) - allowed + if unknown: + raise ValueError( + f"Response domain for {column!r} has unknown fields {sorted(unknown)}." + ) + + if kind == "categorical": + values = raw_spec.get("values") + if not isinstance(values, (list, tuple)) or not values: + raise ValueError( + f"Categorical response domain for {column!r} requires values." + ) + if any( + isinstance(value, bool) or not isinstance(value, Integral) + for value in values + ): + raise ValueError(f"Categorical values for {column!r} must be integers.") + normalized_values = tuple(int(value) for value in values) + if len(set(normalized_values)) != len(normalized_values): + raise ValueError(f"Categorical values for {column!r} must be distinct.") + resolved[column] = { + "kind": "categorical", + "values": normalized_values, + } + continue + + bounds = raw_spec.get("bounds") + if "bounds" not in raw_spec and kind == "continuous": + resolved[column] = {"kind": "continuous"} + continue + if not isinstance(bounds, (list, tuple)) or len(bounds) != 2: + raise ValueError( + f"{kind.capitalize()} response domain for {column!r} requires " + "two bounds." + ) + if any( + isinstance(bound, bool) or not isinstance(bound, Real) for bound in bounds + ): + raise ValueError(f"Bounds for {column!r} must be real numbers.") + lower, upper = (float(bounds[0]), float(bounds[1])) + if not math.isfinite(lower) or not math.isfinite(upper) or lower >= upper: + raise ValueError( + f"Bounds for {column!r} must be finite and strictly increasing." + ) + if kind == "continuous": + resolved[column] = {"kind": "continuous", "bounds": (lower, upper)} + else: + resolved[column] = {"kind": "circular", "bounds": (lower, upper)} + + only_domain = next(iter(resolved.values())) if len(resolved) == 1 else None + resolved_choices = ( + tuple(only_domain["values"]) + if only_domain is not None and only_domain["kind"] == "categorical" + else None + ) + if choices is not None and tuple(choices) != resolved_choices: + raise ValueError( + "Provide either `response_domains` or legacy `choices`, not both." + ) + return resolved, resolved_choices diff --git a/src/hssm/data_validator.py b/src/hssm/data_validator.py index 4f6871c2e..c54a3bdd5 100644 --- a/src/hssm/data_validator.py +++ b/src/hssm/data_validator.py @@ -2,10 +2,13 @@ import logging import warnings +from numbers import Integral, Real import numpy as np import pandas as pd +from ._types import ResponseDomainSpec + _logger = logging.getLogger("hssm") @@ -14,8 +17,9 @@ class DataValidatorMixin: data: pd.DataFrame response: list[str] - choices: list[int] - n_choices: int + response_domains: dict[str, ResponseDomainSpec] + choices: tuple[int, ...] | None + n_choices: int | None extra_fields: list[str] | None deadline: bool deadline_name: str @@ -49,53 +53,93 @@ def _pre_check_data_sanity(self): def _post_check_data_sanity(self): """Check if the data is clean enough for the model.""" if self.is_choice_only: - return - if self.deadline or self.missing_data: - if -999.0 not in self.data["rt"].unique(): + valid_rows = np.ones(len(self.data), dtype=bool) + else: + if self.deadline or self.missing_data: + if -999.0 not in self.data["rt"].unique(): + raise ValueError( + "You have no missing data in your dataset, " + + "which is not allowed when `missing_data` or `deadline` " + "is set to True." + ) + rt_filtered = self.data.rt[self.data.rt != -999.0] + else: + rt_filtered = self.data.rt + + if np.any(rt_filtered.isna(), axis=None): raise ValueError( - "You have no missing data in your dataset, " - + "which is not allowed when `missing_data` or `deadline` is set to" - + " True." + "You have NaN response times in your dataset, " + + "which is not allowed." ) - rt_filtered = self.data.rt[self.data.rt != -999.0] - else: - rt_filtered = self.data.rt - - if np.any(rt_filtered.isna(), axis=None): - raise ValueError( - "You have NaN response times in your dataset, " - + "which is not allowed." - ) - - if not np.all(rt_filtered >= 0): - raise ValueError( - "You have negative response times in your dataset, " - + "which is not allowed." - ) - - valid_responses = self.data.loc[self.data["rt"] != -999.0, "response"] - unique_responses = valid_responses.unique().astype(int) - - if np.any(~np.isin(unique_responses, self.choices)): - invalid_responses = sorted( - unique_responses[~np.isin(unique_responses, self.choices)].tolist() - ) - raise ValueError( - f"Invalid responses found in your dataset: {invalid_responses}" - ) - if len(unique_responses) != self.n_choices: - missing_responses = sorted( - np.setdiff1d(self.choices, unique_responses).tolist() - ) - warnings.warn( - ( - f"You set choices to be {self.choices}, but {missing_responses} " - "are missing from your dataset." - ), - UserWarning, - stacklevel=2, + if not np.all(rt_filtered >= 0): + raise ValueError( + "You have negative response times in your dataset, " + + "which is not allowed." + ) + valid_rows = self.data["rt"].to_numpy() != -999.0 + + for column, domain in self.response_domains.items(): + observed = self.data.loc[valid_rows, column].to_numpy() + if any( + isinstance(value, (bool, np.bool_)) + or not isinstance(value, Real) + or (not isinstance(value, Integral) and not np.isfinite(value)) + for value in observed + ): + raise ValueError( + f"Response column {column!r} must contain finite numeric values." + ) + if domain["kind"] == "categorical": + allowed = domain["values"] + observed_values = set(observed.tolist()) + invalid = sorted(observed_values - set(allowed)) + if invalid: + invalid_responses = [ + int(value) + if isinstance(value, Integral) or float(value).is_integer() + else float(value) + for value in invalid + ] + if column == "response" and len(self.response_domains) == 1: + raise ValueError( + "Invalid responses found in your dataset: " + f"{invalid_responses}" + ) + raise ValueError( + f"Invalid responses found in column {column!r}: " + f"{invalid_responses}" + ) + + missing = sorted(set(allowed) - observed_values) + if missing: + if column == "response" and len(self.response_domains) == 1: + message = ( + f"You set choices to be {allowed}, but {missing} " + "are missing from your dataset." + ) + else: + message = ( + f"Categorical response domain for {column!r} declares " + f"{allowed}, but {missing} are missing from your dataset." + ) + warnings.warn(message, UserWarning, stacklevel=2) + continue + + numeric = observed.astype(float, copy=False) + bounds = domain.get("bounds") + if bounds is None: + continue + lower, upper = bounds + outside = (numeric < lower) | ( + numeric >= upper if domain["kind"] == "circular" else numeric > upper ) + if np.any(outside): + interval = "half-open" if domain["kind"] == "circular" else "closed" + raise ValueError( + f"Response column {column!r} has values outside its {interval} " + f"bounds {bounds}." + ) # AF-TODO: We probably want to incorporate some of the # remaining check on missing data @@ -122,14 +166,3 @@ def _update_extra_fields(self, new_data: pd.DataFrame | None = None): self.model_distribution.extra_fields = [ # type: ignore[attr-defined] new_data[field].values for field in self.extra_fields ] - - def _validate_choices(self): - """ - Ensure that `choices` is provided (not None). - - Raises ValueError if choices is None. - """ - if self.choices is None: - raise ValueError( - "`choices` must be provided either in `model_config` or as an argument." - ) diff --git a/src/hssm/register.py b/src/hssm/register.py index 479624fc4..07ce22e45 100644 --- a/src/hssm/register.py +++ b/src/hssm/register.py @@ -6,6 +6,7 @@ from ._types import ( DefaultConfig, LoglikConfigs, + ResponseDomainSpec, SupportedModels, ) from .defaults import ( @@ -17,9 +18,10 @@ def register_model( name: SupportedModels, response: list[str], list_params: list[str], - choices: list[int], + choices: list[int] | None, likelihoods: LoglikConfigs, description: str | None, + response_domains: dict[str, ResponseDomainSpec] | None = None, ) -> None: """Register a new model in HSSM. @@ -31,8 +33,10 @@ def register_model( List of response variables list_params : list[str] List of parameters - choices : list[int] - List of possible choices + choices : list[int] or None + Legacy list of possible choices for one categorical response. + response_domains : dict or None + Canonical metadata keyed by each physical non-RT response column. description : str Description of the model likelihoods : LoglikConfigs @@ -50,9 +54,27 @@ def register_model( # Ensure no collisions with existing models if name in registered_models: raise ValueError(f"Model '{name}' already exists") + if response_domains is not None and choices is not None: + raise ValueError( + "Provide either `response_domains` or legacy `choices`, not both." + ) - _config = {k: v for k, v in locals().items() if k != "name"} - config = cast("DefaultConfig", _config) + from .config import _resolve_response_domains # noqa: PLC0415 + + resolved_domains, resolved_choices = _resolve_response_domains( + response, response_domains, choices + ) + config: DefaultConfig = { + "response": list(response), + "list_params": list(list_params), + "likelihoods": dict(likelihoods), + "description": description, + } + if response_domains is None: + assert resolved_choices is not None + config["choices"] = list(resolved_choices) + else: + config["response_domains"] = resolved_domains # TODO: validate provided configs? diff --git a/src/hssm/rl/config.py b/src/hssm/rl/config.py index c6dffe72c..73285ca20 100644 --- a/src/hssm/rl/config.py +++ b/src/hssm/rl/config.py @@ -10,6 +10,7 @@ import importlib import logging +from copy import deepcopy from dataclasses import MISSING, dataclass, field, fields from typing import TYPE_CHECKING, Any, Literal @@ -17,7 +18,12 @@ from .._types import LoglikKind, SupportedModels from ..config import ModelConfig -from ..config import DEFAULT_SSM_CHOICES, DEFAULT_SSM_OBSERVED_DATA, BaseModelConfig +from ..config import ( + DEFAULT_SSM_CHOICES, + DEFAULT_SSM_OBSERVED_DATA, + BaseModelConfig, + _resolve_response_domains, +) from ..utils import annotate_function _logger = logging.getLogger("hssm") @@ -164,7 +170,16 @@ def _get_or_warn(key: str, default: Any) -> None: init_kwargs[key] = config_dict.get(key, default) _get_or_warn("response", DEFAULT_SSM_OBSERVED_DATA) - _get_or_warn("choices", DEFAULT_SSM_CHOICES) + response_domains = config_dict.get("response_domains") + if response_domains is not None: + if config_dict.get("choices") is not None: + raise ValueError( + "Provide either `response_domains` or legacy `choices`, not both." + ) + init_kwargs["response_domains"] = deepcopy(response_domains) + init_kwargs["choices"] = None + else: + _get_or_warn("choices", DEFAULT_SSM_CHOICES) return cls(**init_kwargs) @@ -266,6 +281,13 @@ def from_ssms_model( def validate(self) -> None: # noqa: D102 if self.response is None: raise ValueError("Please provide `response` columns in the configuration.") + if self.response_domains is None and self.choices is None: + raise ValueError( + "Please provide `choices` or `response_domains` in the configuration." + ) + self.response_domains, self.choices = _resolve_response_domains( + self.response, self.response_domains, self.choices + ) if self.list_params is None: raise ValueError("Please provide `list_params` in the configuration.") if self.choices is None: diff --git a/src/hssm/rl/rlssm.py b/src/hssm/rl/rlssm.py index 65411aec9..34fdb17e6 100644 --- a/src/hssm/rl/rlssm.py +++ b/src/hssm/rl/rlssm.py @@ -22,7 +22,6 @@ """ import logging -import warnings from dataclasses import replace from typing import TYPE_CHECKING, Any, Callable, Literal, cast @@ -181,9 +180,6 @@ def __init__( # Infer panel structure and validate balance BEFORE calling super so any # error surfaces before the expensive model-build steps. n_participants, n_trials = validate_balanced_panel(data, participant_col) - if model_config.is_choice_only: - self._validate_choice_only_responses(data, model_config) - # Store RL-specific state on self BEFORE super().__init__() so that # _make_model_distribution() (called from super) can access them. self.n_participants = n_participants @@ -236,59 +232,6 @@ def __init__( **kwargs, ) - @staticmethod - def _validate_choice_only_responses( - data: pd.DataFrame, - model_config: RLSSMConfig, - ) -> None: - """Validate response labels for RLSSM choice-only observed data.""" - if not model_config.response: - return - response_col = model_config.response[0] - if response_col not in data.columns: - return - if model_config.choices is None: - return - - responses = data[response_col] - numeric_responses = pd.to_numeric(responses, errors="coerce") - non_numeric = numeric_responses.isna() & ~responses.isna() - if non_numeric.any(): - raise ValueError( - "Choice-only RLSSM response labels must be numeric. " - f"Invalid values: {responses[non_numeric].unique().tolist()}" - ) - - response_values = numeric_responses.to_numpy(dtype=float) - if not np.all(np.isfinite(response_values)): - raise ValueError("Choice-only RLSSM response labels must be finite.") - - if not np.all(np.equal(response_values, np.round(response_values))): - raise ValueError("Choice-only RLSSM response labels must be integral.") - - response_ints = response_values.astype(int) - unique_responses = np.unique(response_ints) - choices = np.asarray(model_config.choices, dtype=int) - - if np.any(~np.isin(unique_responses, choices)): - invalid_responses = sorted( - unique_responses[~np.isin(unique_responses, choices)].tolist() - ) - raise ValueError( - f"Invalid responses found in your dataset: {invalid_responses}" - ) - - missing_responses = sorted(np.setdiff1d(choices, unique_responses).tolist()) - if missing_responses: - warnings.warn( - ( - f"You set choices to be {model_config.choices}, but " - f"{missing_responses} are missing from your dataset." - ), - UserWarning, - stacklevel=2, - ) - def _make_model_distribution(self) -> type[pm.Distribution]: """Build a pm.Distribution using the pre-built RL log-likelihood Op. diff --git a/tests/addm/test_addm_config.py b/tests/addm/test_addm_config.py index a7e3f81c6..005db3feb 100644 --- a/tests/addm/test_addm_config.py +++ b/tests/addm/test_addm_config.py @@ -10,6 +10,8 @@ Patterned after tests/test_rlssm_config.py. """ +from dataclasses import replace + import pytest from hssm.addm.config import aDDMConfig @@ -50,6 +52,28 @@ def test_validate_ok(): aDDMConfig().validate() # must not raise +def test_response_domain_compatibility_view_must_match(): + domains = {"response": {"kind": "categorical", "values": (-1, 1)}} + with pytest.raises(ValueError, match="either `response_domains` or legacy"): + aDDMConfig.from_addm_dict({"response_domains": domains, "choices": (-1, 1)}) + + config = aDDMConfig(response_domains=domains) + config.validate() + assert replace(config).choices == (-1, 1) + + mismatched = aDDMConfig( + response_domains={"response": {"kind": "categorical", "values": (0, 1)}} + ) + with pytest.raises(ValueError, match="either `response_domains` or legacy"): + mismatched.validate() + + noncategorical = aDDMConfig( + response_domains={"response": {"kind": "continuous"}}, choices=None + ) + with pytest.raises(ValueError, match="requires one categorical"): + noncategorical.validate() + + def test_validate_rejects_unknown_attention_process(): with pytest.raises(ValueError): aDDMConfig(attention_process="bogus").validate() @@ -94,6 +118,18 @@ def test_from_addm_dict_roundtrip(): assert rebuilt.model_name == src.model_name +def test_from_addm_dict_detaches_nested_response_domains(): + """The dict adapter owns its canonical metadata before validation.""" + domains = {"response": {"kind": "categorical", "values": [-1, 1]}} + config = aDDMConfig.from_addm_dict({"response_domains": domains, "choices": None}) + + domains["response"]["values"][0] = 0 + + assert config.response_domains == { + "response": {"kind": "categorical", "values": [-1, 1]} + } + + if __name__ == "__main__": for fn in ( test_defaults, diff --git a/tests/rl/test_choice_only_rl.py b/tests/rl/test_choice_only_rl.py index d9c09ad41..190973cb9 100644 --- a/tests/rl/test_choice_only_rl.py +++ b/tests/rl/test_choice_only_rl.py @@ -206,14 +206,14 @@ def test_inv_temp_softmax_preserves_lan_matrix_float_dtype(): @pytest.mark.parametrize( ("responses", "match"), [ - ([0, 1.5, 0, 1], "integral"), - ([0, "left", 0, 1], "numeric"), - ([0, np.inf, 0, 1], "finite"), + ([0, 1.5, 0, 1], "Invalid responses"), + ([0, "left", 0, 1], "finite numeric"), + ([0, np.inf, 0, 1], "finite numeric"), ([0, 2, 0, 1], "Invalid responses"), ], ) def test_choice_only_rlssm_validates_response_labels(responses, match): - """Choice-only RLSSM response labels are checked before logp evaluation.""" + """Choice-only labels use the same exact domain checks as other models.""" with pytest.raises(ValueError, match=match): hssm.RLSSM( data=_fake_choice_only_data(responses), diff --git a/tests/rl/test_rlssm.py b/tests/rl/test_rlssm.py index 5e9fe0533..ab34a2e55 100644 --- a/tests/rl/test_rlssm.py +++ b/tests/rl/test_rlssm.py @@ -379,31 +379,6 @@ def capturing_make_distribution(*args, **kwargs): assert captured.get("extra_fields") is None - def test_choice_only_response_validation_exits_when_metadata_is_incomplete( - self, - ) -> None: - """Choice-only validation is a no-op until response metadata is complete.""" - data = pd.DataFrame({"response": [0, 1], "feedback": [1.0, 0.0]}) - - assert ( - _RLSSM._validate_choice_only_responses( - data, SimpleNamespace(response=[], choices=[0, 1]) - ) - is None - ) - assert ( - _RLSSM._validate_choice_only_responses( - data, SimpleNamespace(response=["missing"], choices=[0, 1]) - ) - is None - ) - assert ( - _RLSSM._validate_choice_only_responses( - data, SimpleNamespace(response=["response"], choices=None) - ) - is None - ) - def test_choice_only_scalar_lapse_must_be_probability(self) -> None: """Choice-only scalar lapse values must satisfy the documented bounds.""" model = object.__new__(_RLSSM) diff --git a/tests/rl/test_rlssm_config.py b/tests/rl/test_rlssm_config.py index 6082e3ea6..08faaef2b 100644 --- a/tests/rl/test_rlssm_config.py +++ b/tests/rl/test_rlssm_config.py @@ -2,6 +2,7 @@ import sys import types +from dataclasses import replace import jax.numpy as jnp import pytest @@ -69,6 +70,60 @@ def create_config_dict( ) +def test_from_rlssm_dict_preserves_canonical_response_domains(): + """Canonical RL metadata is copied without injecting legacy choices.""" + domains = {"response": {"kind": "categorical", "values": [0, 1]}} + config_dict = create_config_dict( + "canonical_rlssm", + ["alpha"], + [0.5], + bounds={"alpha": (0.0, 1.0)}, + ) + config_dict.pop("choices") + config_dict["response_domains"] = domains + + config = RLSSMConfig.from_rlssm_dict(config_dict) + domains["response"]["values"][0] = -1 + config.validate() + + assert config.response_domains == { + "response": {"kind": "categorical", "values": (0, 1)} + } + assert config.choices == (0, 1) + + +def test_from_rlssm_dict_rejects_canonical_and_legacy_metadata(): + """RL dictionary construction rejects two response-domain sources.""" + config_dict = create_config_dict( + "ambiguous_rlssm", + ["alpha"], + [0.5], + bounds={"alpha": (0.0, 1.0)}, + ) + config_dict["response_domains"] = { + "response": {"kind": "categorical", "values": [0, 1]} + } + + with pytest.raises(ValueError, match="either `response_domains` or legacy"): + RLSSMConfig.from_rlssm_dict(config_dict) + + +def test_resolved_config_round_trip_accepts_matching_derived_choices( + valid_rlssmconfig_kwargs, +): + """Resolved canonical RL configs remain replace-compatible.""" + valid_rlssmconfig_kwargs["response_domains"] = { + "response": {"kind": "categorical", "values": (0, 1)} + } + + config = RLSSMConfig(**valid_rlssmconfig_kwargs) + config.validate() + + copied = replace(config) + copied.validate() + assert copied.choices == (0, 1) + + # region fixtures and helpers @pytest.fixture def valid_rlssmconfig_kwargs(): diff --git a/tests/test_data_validator.py b/tests/test_data_validator.py index e37151296..1e2bc2b56 100644 --- a/tests/test_data_validator.py +++ b/tests/test_data_validator.py @@ -1,10 +1,13 @@ -from typing import Callable +"""Tests for response-data validation.""" + +from collections.abc import Callable -import pytest -import pandas as pd import numpy as np +import pandas as pd +import pytest + +from hssm._types import ResponseDomainSpec from hssm.data_validator import DataValidatorMixin -from hssm.defaults import MissingDataNetwork class DataValidatorTester(DataValidatorMixin): @@ -18,17 +21,45 @@ def __init__( missing_data: bool = False, choices: list[int] | None = None, n_choices: int | None = None, + response: list[str] | None = None, + response_domains: dict[str, ResponseDomainSpec] | None = None, + is_choice_only: bool = False, ): self.data = data - self.response = ["rt", "response"] - self.choices = choices if choices is not None else [0, 1] - self.n_choices = n_choices if n_choices is not None else len(self.choices) + self.response = response or ["rt", "response"] + if response_domains is None: + self.choices = tuple(choices) if choices is not None else (0, 1) + self.response_domains = { + self.response[-1]: { + "kind": "categorical", + "values": self.choices, + } + } + else: + self.response_domains = response_domains + only_domain = ( + next(iter(response_domains.values())) + if len(response_domains) == 1 + else None + ) + self.choices = ( + tuple(only_domain["values"]) + if only_domain is not None and only_domain["kind"] == "categorical" + else None + ) + self.n_choices = ( + n_choices + if n_choices is not None + else len(self.choices) + if self.choices is not None + else None + ) self.extra_fields = extra_fields self.deadline = deadline self.deadline_name = "deadline" self.missing_data = missing_data self.missing_data_value = -999.0 - self.is_choice_only = False + self.is_choice_only = is_choice_only def _base_data(): @@ -88,7 +119,7 @@ def test_constructor(base_data): assert dv.data.equals(_base_data()) assert dv.response == ["rt", "response"] - assert dv.choices == [0, 1] + assert dv.choices == (0, 1) assert dv.n_choices == 2 assert dv.extra_fields == ["extra"] assert dv.deadline is True @@ -184,16 +215,207 @@ class DummyModelDist: assert (dv.model_distribution.extra_fields[i] == data[field].values).all() -def test_validate_choices(): - # ====== Valid choices ===== - dv = DataValidatorTester( - data=_base_data(), - choices=[0, 1], - n_choices=2, +def test_mixed_domains_validate_in_physical_column_order_without_mutation(): + """Mixed scalar columns validate against their own domains without coercion.""" + data = pd.DataFrame( + { + "rt": [0.2, 0.3, 0.4], + "polar": [0.0, 1.0, 2.0], + "azimuth": [-np.pi, 0.0, np.nextafter(np.pi, -np.inf)], + } + ) + original = data.copy(deep=True) + validator = DataValidatorTester( + data, + response=["rt", "polar", "azimuth"], + response_domains={ + "polar": {"kind": "continuous", "bounds": (0.0, 2.0)}, + "azimuth": {"kind": "circular", "bounds": (-np.pi, np.pi)}, + }, + ) + + validator._post_check_data_sanity() + + pd.testing.assert_frame_equal(data, original) + + +def test_four_observation_columns_validate_independently(): + """Each scalar coordinate in a wider observation has its own domain.""" + validator = DataValidatorTester( + pd.DataFrame( + { + "rt": [0.2, 0.3], + "confidence": [0.0, 1.0], + "angle": [-np.pi, 0.0], + "choice": [0.0, 1.0], + } + ), + response=["rt", "confidence", "angle", "choice"], + response_domains={ + "confidence": {"kind": "continuous", "bounds": (0.0, 1.0)}, + "angle": {"kind": "circular", "bounds": (-np.pi, np.pi)}, + "choice": {"kind": "categorical", "values": (0, 1)}, + }, + ) + + validator._post_check_data_sanity() + + +def test_first_invalid_physical_column_follows_declared_order(): + """Simultaneous failures report the first configured response coordinate.""" + validator = DataValidatorTester( + pd.DataFrame({"rt": [0.2], "first": [2.0], "second": [2.0]}), + response=["rt", "first", "second"], + response_domains={ + "first": {"kind": "continuous", "bounds": (0.0, 1.0)}, + "second": {"kind": "continuous", "bounds": (0.0, 1.0)}, + }, + ) + + with pytest.raises(ValueError, match="column 'first'.*bounds"): + validator._post_check_data_sanity() + + +@pytest.mark.parametrize( + ("kind", "bounds", "value", "passes"), + [ + ("continuous", (0.0, 1.0), 0.0, True), + ("continuous", (0.0, 1.0), 1.0, True), + ("continuous", (0.0, 1.0), np.nextafter(0.0, -np.inf), False), + ("continuous", (0.0, 1.0), np.nextafter(1.0, np.inf), False), + ("circular", (-np.pi, np.pi), -np.pi, True), + ("circular", (-np.pi, np.pi), np.nextafter(np.pi, -np.inf), True), + ("circular", (-np.pi, np.pi), np.pi, False), + ("circular", (-np.pi, np.pi), np.nextafter(-np.pi, -np.inf), False), + ], +) +def test_continuous_and_circular_endpoint_semantics( + kind: str, bounds: tuple[float, float], value: float, passes: bool +): + """Continuous bounds are closed while circular upper bounds are excluded.""" + validator = DataValidatorTester( + pd.DataFrame({"rt": [0.2], "coordinate": [value]}), + response=["rt", "coordinate"], + response_domains={ + "coordinate": {"kind": kind, "bounds": bounds} # type: ignore[typeddict-item] + }, + ) + + if passes: + validator._post_check_data_sanity() + else: + with pytest.raises(ValueError, match="column 'coordinate'.*bounds"): + validator._post_check_data_sanity() + + +@pytest.mark.parametrize("value", [np.nan, np.inf, -np.inf, "1", True]) +def test_domains_reject_nonfinite_or_nonnumeric_values(value): + """Every domain rejects nonfinite, nonnumeric, and Boolean observations.""" + validator = DataValidatorTester( + pd.DataFrame({"rt": [0.2], "coordinate": [value]}), + response=["rt", "coordinate"], + response_domains={"coordinate": {"kind": "continuous"}}, + ) + + with pytest.raises(ValueError, match="column 'coordinate'.*finite numeric"): + validator._post_check_data_sanity() + + +def test_categorical_membership_does_not_integer_cast_fractional_values(): + """A fractional category cannot pass by truncation to an integer label.""" + validator = DataValidatorTester( + pd.DataFrame({"rt": [0.2, 0.3], "response": [0.0, 0.5]}), + response_domains={"response": {"kind": "categorical", "values": (0, 1)}}, + ) + + with pytest.raises(ValueError, match=r"Invalid responses.*\[0\.5\]"): + validator._post_check_data_sanity() + + +def test_categorical_membership_preserves_large_integer_precision(): + """Adjacent integer labels above float precision remain distinguishable.""" + allowed = 2**53 + validator = DataValidatorTester( + pd.DataFrame({"rt": [0.2], "response": [allowed + 1]}), + response_domains={"response": {"kind": "categorical", "values": (allowed,)}}, + ) + + with pytest.raises(ValueError, match=str(allowed + 1)): + validator._post_check_data_sanity() + + +def test_categorical_membership_supports_arbitrary_python_integers(): + """Categorical labels are not constrained to NumPy fixed-width integers.""" + allowed = 10**100 + validator = DataValidatorTester( + pd.DataFrame( + { + "rt": [0.2], + "response": pd.Series([allowed], dtype=object), + } + ), + response_domains={"response": {"kind": "categorical", "values": (allowed,)}}, + ) + + validator._post_check_data_sanity() + + validator.data.loc[0, "response"] = allowed + 1 + with pytest.raises(ValueError, match=str(allowed + 1)): + validator._post_check_data_sanity() + + +def test_multidomain_categorical_failure_names_physical_response_column(): + """A legacy-named column is still identified in a wider response.""" + validator = DataValidatorTester( + pd.DataFrame({"rt": [0.2], "response": [2], "confidence": [0.5]}), + response=["rt", "response", "confidence"], + response_domains={ + "response": {"kind": "categorical", "values": (0, 1)}, + "confidence": {"kind": "continuous", "bounds": (0, 1)}, + }, + ) + + with pytest.raises(ValueError, match="column 'response'"): + validator._post_check_data_sanity() + + +def test_missing_rt_rows_are_omitted_but_observed_rows_remain_validated(): + """Missing-data sentinels exempt only their own response row.""" + data = pd.DataFrame({"rt": [-999.0, 0.3], "coordinate": [99.0, 0.5]}) + validator = DataValidatorTester( + data, + missing_data=True, + response=["rt", "coordinate"], + response_domains={"coordinate": {"kind": "continuous", "bounds": (0.0, 1.0)}}, + ) + validator._post_check_data_sanity() + + validator.data.loc[1, "coordinate"] = 99.0 + with pytest.raises(ValueError, match="column 'coordinate'.*bounds"): + validator._post_check_data_sanity() + + +def test_custom_missing_marker_uses_processed_internal_sentinel(): + """A custom missing marker remains omitted after preprocessing to -999.""" + validator = DataValidatorTester( + pd.DataFrame({"rt": [-999.0, 0.3], "coordinate": [99.0, 0.5]}), + missing_data=True, + response=["rt", "coordinate"], + response_domains={"coordinate": {"kind": "continuous", "bounds": (0, 1)}}, + ) + validator.missing_data_value = -123.0 + + validator._post_check_data_sanity() + + +def test_choice_only_domain_uses_the_shared_validation_loop(): + """A one-column choice-only response is checked by the canonical path.""" + validator = DataValidatorTester( + pd.DataFrame({"choice": [0.0, 1.0, 0.5]}), + response=["choice"], + response_domains={"choice": {"kind": "categorical", "values": (0, 1)}}, + is_choice_only=True, ) - dv._validate_choices() # Should not raise an exception - # ===== Invalid choices ===== - dv.choices = None # type: ignore[assignment] - with pytest.raises(ValueError, match="`choices` must be provided*."): - dv._validate_choices() + with pytest.raises(ValueError, match=r"column 'choice'.*\[0\.5\]"): + validator._post_check_data_sanity() diff --git a/tests/test_response_domains.py b/tests/test_response_domains.py new file mode 100644 index 000000000..4bf6c2428 --- /dev/null +++ b/tests/test_response_domains.py @@ -0,0 +1,467 @@ +"""Tests for canonical response-domain configuration.""" + +from collections.abc import Mapping +from dataclasses import fields +from typing import Any + +import bambi as bmb +import pytest + +import hssm +import hssm.config as config_module +from hssm.config import Config, ModelConfig +from hssm.defaults import default_model_config +from hssm.register import register_model + + +def _config( + response: list[str], + *, + response_domains: Mapping[str, Mapping[str, object]] | None = None, + choices: tuple[int, ...] | None = None, +) -> Config: + return Config( + model_name="custom_response_domains", + loglik_kind="analytical", + response=response, + response_domains=response_domains, # type: ignore[arg-type] + choices=choices, + list_params=["v"], + loglik=lambda *args: 0.0, + ) + + +def test_legacy_choices_resolve_to_one_categorical_domain(): + """Legacy choices normalize to canonical metadata and remain projected.""" + config = _config(["rt", "response"], choices=(-1, 1)) + + config.validate() + + assert config.response_domains == { + "response": {"kind": "categorical", "values": (-1, 1)} + } + assert config.choices == (-1, 1) + + +def test_resolved_config_accepts_matching_derived_choices(): + """Resolved configs remain idempotent with an exact compatibility view.""" + config = _config( + ["rt", "response"], + response_domains={"response": {"kind": "categorical", "values": (0, 1)}}, + choices=(0, 1), + ) + + config.validate() + assert config.choices == (0, 1) + + +def test_mutating_a_resolved_domain_snapshot_fails_closed(): + """Nested config mutation cannot silently stale a compatibility view.""" + config = _config( + ["rt", "response"], + response_domains={"response": {"kind": "categorical", "values": (0, 1)}}, + ) + config.validate() + assert config.response_domains is not None + + config.response_domains["response"]["values"] = (2, 3) + + with pytest.raises(ValueError, match="either `response_domains` or legacy"): + config.validate() + + +def test_domains_follow_physical_response_order_and_derive_no_global_choices(): + """Canonical domains follow physical response order without global choices.""" + config = _config( + ["rt", "polar", "azimuth"], + response_domains={ + "azimuth": {"kind": "circular", "bounds": (-3.14, 3.14)}, + "polar": {"kind": "continuous", "bounds": (0, 3.14)}, + }, + ) + + config.validate() + + assert list(config.response_domains) == ["polar", "azimuth"] + assert config.response_domains == { + "polar": {"kind": "continuous", "bounds": (0.0, 3.14)}, + "azimuth": {"kind": "circular", "bounds": (-3.14, 3.14)}, + } + assert config.choices is None + + +def test_single_categorical_domain_derives_legacy_choices(): + """One categorical domain retains the established choices projection.""" + config = _config( + ["response"], + response_domains={"response": {"kind": "categorical", "values": [0, 2, 4]}}, + ) + + config.validate() + + assert config.response_domains == { + "response": {"kind": "categorical", "values": (0, 2, 4)} + } + assert config.choices == (0, 2, 4) + assert config.is_choice_only + + config.validate() + assert config.choices == (0, 2, 4) + + +@pytest.mark.parametrize( + ("response", "domains", "choices", "message"), + [ + (["rt", "response"], None, None, "Provide `response_domains`"), + (["rt"], {}, None, "At least one non-RT"), + (["response", "other"], {}, None, "without `rt`"), + (["response", "rt"], {}, None, "index zero"), + (["rt", "rt", "response"], {}, None, "unique"), + ( + ["rt", "response"], + {"response": {"kind": "continuous"}}, + (0, 1), + "either `response_domains` or legacy `choices`", + ), + ( + ["rt", "response"], + {"other": {"kind": "continuous"}}, + None, + "keys must match", + ), + ( + ["rt", "response"], + {"response": {"kind": "ordinal"}}, + None, + "invalid kind", + ), + ( + ["rt", "response"], + {"response": {"kind": "continuous", "values": [0, 1]}}, + None, + "unknown fields", + ), + ( + ["rt", "response"], + {"response": {"kind": "categorical", "values": []}}, + None, + "requires values", + ), + ( + ["rt", "response"], + {"response": {"kind": "categorical", "values": [0, 0]}}, + None, + "must be distinct", + ), + ( + ["rt", "response"], + {"response": {"kind": "categorical", "values": [0, 0.5]}}, + None, + "must be integers", + ), + ( + ["rt", "response"], + {"response": {"kind": "circular"}}, + None, + "requires two bounds", + ), + ( + ["rt", "response"], + {"response": {"kind": "continuous", "bounds": None}}, + None, + "requires two bounds", + ), + ( + ["rt", "response"], + {"response": {"kind": "circular", "bounds": (0, float("inf"))}}, + None, + "finite and strictly increasing", + ), + ( + ["rt", "response"], + {"response": {"kind": "continuous", "bounds": (1, 1)}}, + None, + "finite and strictly increasing", + ), + ], +) +def test_invalid_response_domain_contracts_fail( + response: list[str], + domains: dict[str, dict[str, Any]] | None, + choices: tuple[int, ...] | None, + message: str, +): + """Malformed, incomplete, and ambiguous domain declarations fail closed.""" + config = _config(response, response_domains=domains, choices=choices) + + with pytest.raises(ValueError, match=message): + config.validate() + + +def test_model_config_rejects_canonical_and_legacy_inputs_together(): + """Canonical metadata cannot be combined with explicit legacy choices.""" + model_config = ModelConfig( + response=("rt", "response"), + response_domains={"response": {"kind": "continuous"}}, + ) + + with pytest.raises(ValueError, match="either `response_domains` or legacy"): + Config._build_model_config( + "ddm", None, model_config, choices=(-1, 1), loglik=None + ) + + +def test_model_config_detaches_nested_domain_input(): + """ModelConfig normalization owns a detached canonical mapping.""" + domains: dict[str, Any] = { + "response": {"kind": "circular", "bounds": [-3.14, 3.14]} + } + model_config = ModelConfig(response=("rt", "response"), response_domains=domains) + + config = Config._build_model_config("ddm", None, model_config, None) + domains["response"]["bounds"][0] = 0.0 + + assert config.response_domains == { + "response": {"kind": "circular", "bounds": (-3.14, 3.14)} + } + assert config.choices is None + + +def test_registered_domains_are_detached_and_resolve_canonically(): + """Registration detaches caller input and reconstructs canonical config.""" + name = "custom_registered_response_domains" + domains: dict[str, Any] = {"response": {"kind": "continuous", "bounds": [0.0, 1.0]}} + likelihoods = { + "analytical": { + "loglik": lambda *args: 0.0, + "backend": None, + "default_priors": {}, + "bounds": {}, + "extra_fields": None, + } + } + register_model( + name=name, # type: ignore[arg-type] + response=["rt", "response"], + list_params=["v"], + choices=None, + response_domains=domains, + likelihoods=likelihoods, # type: ignore[arg-type] + description=None, + ) + try: + domains["response"]["bounds"][0] = -1.0 + config = Config.from_defaults(name, "analytical") + config.validate() + + assert config.response_domains == { + "response": {"kind": "continuous", "bounds": (0.0, 1.0)} + } + assert config.choices is None + finally: + default_model_config.pop(name, None) # type: ignore[arg-type] + + +def test_registered_canonical_model_rejects_top_level_choices_override(): + """A registered canonical source cannot be combined with legacy choices.""" + name = "custom_registered_choice_conflict" + register_model( + name=name, # type: ignore[arg-type] + response=["rt", "response"], + list_params=["v"], + choices=None, + response_domains={"response": {"kind": "categorical", "values": (0, 1)}}, + likelihoods={ + "analytical": { + "loglik": lambda *args: 0.0, + "backend": None, + "default_priors": {}, + "bounds": {}, + "extra_fields": None, + } + }, # type: ignore[arg-type] + description=None, + ) + try: + with pytest.raises(ValueError, match="either `response_domains` or legacy"): + Config._build_model_config( + name, "analytical", None, choices=(0, 1), loglik=None + ) + finally: + default_model_config.pop(name, None) # type: ignore[arg-type] + + +def test_registered_canonical_model_ignores_ssms_legacy_fallback(monkeypatch): + """Canonical registration wins over an ssms registry name collision.""" + name = "custom_registered_ssms_collision" + register_model( + name=name, # type: ignore[arg-type] + response=["rt", "response"], + list_params=["v"], + choices=None, + response_domains={"response": {"kind": "categorical", "values": (0, 1)}}, + likelihoods={ + "analytical": { + "loglik": lambda *args: 0.0, + "backend": None, + "default_priors": {}, + "bounds": {}, + "extra_fields": None, + } + }, # type: ignore[arg-type] + description=None, + ) + monkeypatch.setitem(config_module.ssms_model_config, name, {"choices": (8, 9)}) + try: + config = Config._build_model_config(name, "analytical", None, None) + + assert config.choices == (0, 1) + finally: + default_model_config.pop(name, None) # type: ignore[arg-type] + + +@pytest.mark.parametrize( + ("model", "loglik_kind", "choices"), + [ + ("ddm", "analytical", (-1, 1)), + ("ddm_seq2_no_bias", "approx_differentiable", (0, 1, 2, 3)), + ("lba4", "analytical", (0, 1, 2, 3)), + ("softmax_inv_temperature_3", "analytical", (0, 1, 2)), + ], +) +def test_legacy_builtins_resolve_without_changing_response_or_choices( + model, loglik_kind, choices +): + """Existing categorical factories acquire canonical internal metadata only.""" + config = Config.from_defaults(model, loglik_kind) + + config.validate() + + response_column = config.response[-1] + assert config.response_domains == { + response_column: {"kind": "categorical", "values": choices} + } + assert config.choices == choices + + +@pytest.mark.parametrize("as_dict", [False, True]) +def test_constructor_snapshot_detaches_nested_response_domains(as_dict): + """Save/load constructor arguments own their nested domain metadata.""" + domains: dict[str, Any] = {"response": {"kind": "continuous", "bounds": [0.0, 1.0]}} + model_config: ModelConfig | dict[str, Any] + if as_dict: + model_config = {"response_domains": domains} + else: + model_config = ModelConfig(response_domains=domains) + + snapshot = hssm.HSSM._store_init_args( + {"self": object(), "model_config": model_config}, {} + ) + domains["response"]["bounds"][0] = -1.0 + stored = snapshot["model_config"] + stored_domains = ( + stored["response_domains"] + if isinstance(stored, dict) + else stored.response_domains + ) + + assert stored_domains["response"]["bounds"] == [0.0, 1.0] + + stored_domains["response"]["bounds"][1] = 2.0 + assert domains["response"]["bounds"] == [-1.0, 1.0] + + +def test_model_config_positional_arguments_keep_their_legacy_meaning(): + """Appending response domains does not shift existing positional fields.""" + config = ModelConfig(("rt", "response"), ["v"], (-1, 1)) + + assert config.response == ("rt", "response") + assert config.list_params == ["v"] + assert config.choices == (-1, 1) + assert config.response_domains is None + + domain_field = next( + field for field in fields(Config) if field.name == "response_domains" + ) + assert domain_field.kw_only + + +def test_live_model_owns_one_detached_response_domain_mapping(): + """Live config and validation share one mapping detached from the caller.""" + domains: dict[str, Any] = {"response": {"kind": "categorical", "values": [-1, 1]}} + model = hssm.HSSM( + data=hssm.load_data("cavanagh_theta").head(8), + model_config=ModelConfig(response_domains=domains), + p_outlier=None, + process_initvals=False, + ) + + domains["response"]["values"][0] = -2 + assert model.response_domains == { + "response": {"kind": "categorical", "values": (-1, 1)} + } + assert model.response_domains is model.model_config.response_domains + + +def _lapse_shell( + domains: dict[str, dict[str, Any]], *, choice_only: bool = False +) -> hssm.HSSM: + model = object.__new__(hssm.HSSM) + model.list_params = ["v"] + model.response_domains = domains # type: ignore[assignment] + model.has_lapse = True + model.is_choice_only = choice_only + model.n_choices = 2 if choice_only else None + return model + + +def test_lapse_remains_available_for_established_categorical_layouts(): + """RT+choice and choice-only models retain their established lapse behavior.""" + domains = {"response": {"kind": "categorical", "values": (0, 1)}} + + rt_model = _lapse_shell(domains) + rt_model._check_lapse(None) + assert isinstance(rt_model.lapse, bmb.Prior) + + choice_only_model = _lapse_shell(domains, choice_only=True) + choice_only_model._check_lapse(None) + assert choice_only_model.lapse == 0.5 + + +@pytest.mark.parametrize( + "domains", + [ + {"response": {"kind": "continuous"}}, + {"response": {"kind": "circular", "bounds": (-3.14, 3.14)}}, + { + "first": {"kind": "categorical", "values": (0, 1)}, + "second": {"kind": "categorical", "values": (0, 1)}, + }, + { + "first": {"kind": "continuous"}, + "second": {"kind": "categorical", "values": (0, 1)}, + }, + ], +) +def test_lapse_rejects_noncategorical_or_multiresponse_layouts(domains): + """Active outlier mixtures fail before unsupported domain layouts build.""" + model = _lapse_shell(domains) + + with pytest.raises(ValueError, match="only for one categorical response"): + model._check_lapse(None) + + +def test_inactive_lapse_accepts_mixed_domains(): + """None or zero p_outlier maps to an inactive lapse for mixed responses.""" + model = _lapse_shell( + { + "first": {"kind": "continuous"}, + "second": {"kind": "categorical", "values": (0, 1)}, + } + ) + model.has_lapse = False + + model._check_lapse(None) + + assert model.lapse is None + assert model.list_params == ["v"]