[WIP] Initial implementation of choice-only models

.gitignore

@@ -86,4 +86,8 @@ explorations/
 *.nbconvert.ipynb
 pypi-token
 pre-commit
+
+# AI slops
 .github/copilot-instructions.md
+CLAUDE.md
+.claude/

src/hssm/_types.py

@@ -23,6 +23,8 @@
     "ddm_seq2_no_bias",
     "lba3",
     "lba2",
+    "softmax_inv_temperature_2",
+    "softmax_inv_temperature_3",
 ]
 
 

src/hssm/config.py

@@ -260,6 +260,13 @@ def get_defaults(
         """
         return self.default_priors.get(param), self.bounds.get(param)
 
+    @property
+    def is_choice_only(self) -> bool:
+        """Check if the model is a choice-only model."""
+        if self.response is None:
+            raise ValueError("Response is not defined in the configuration.")
+        return len(self.response) == 1
+
 
 @dataclass
 class RLSSMConfig(BaseModelConfig):

src/hssm/data_validator.py

@@ -3,10 +3,10 @@
 import logging
 import warnings
 
-import numpy as np  # noqa: F401
-import pandas as pd  # noqa: F401
+import numpy as np
+import pandas as pd
 
-from hssm.defaults import MissingDataNetwork  # noqa: F401
+from hssm.defaults import MissingDataNetwork
 
 _logger = logging.getLogger("hssm")
 
@@ -26,31 +26,16 @@ class DataValidatorMixin:
     - missing_data_value: float
     """
 
-    def __init__(
-        self,
-        data: pd.DataFrame,
-        response: list[str] | None = ["rt", "response"],
-        choices: list[int] | None = [0, 1],
-        n_choices: int = 2,
-        extra_fields: list[str] | None = None,
-        deadline: bool = False,
-        deadline_name: str = "deadline",
-        missing_data: bool = False,
-        missing_data_value: float = -999.0,
-    ):
-        """Initialize the DataValidatorMixin.
-
-        Init method kept for testing purposes.
-        """
-        self.data = data
-        self.response = response
-        self.choices = choices
-        self.n_choices = n_choices
-        self.extra_fields = extra_fields
-        self.deadline = deadline
-        self.deadline_name = deadline_name
-        self.missing_data = missing_data
-        self.missing_data_value = missing_data_value
+    data: pd.DataFrame
+    response: list[str]
+    choices: list[int]
+    n_choices: int
+    extra_fields: list[str] | None
+    deadline: bool
+    deadline_name: str
+    missing_data: bool
+    missing_data_value: float
+    is_choice_only: bool
 
     @staticmethod
     def check_fields(a, b):
@@ -131,6 +116,10 @@ def _post_check_data_sanity(self):
     def _handle_missing_data_and_deadline(self):
         """Handle missing data and deadline."""
         if not self.missing_data and not self.deadline:
+            # In the case of choice only model, we don't need to do anything with the
+            # data.
+            if self.is_choice_only:
+                return
             # In the case where missing_data is set to False, we need to drop the
             # cases where rt = na_value
             if pd.isna(self.missing_data_value):

src/hssm/distribution_utils/dist.py

@@ -191,6 +191,7 @@ def make_hssm_rv(
     simulator_fun: Callable | str,
     list_params: list[str],
     lapse: bmb.Prior | None = None,
+    is_choice_only: bool = False,
 ) -> type[RandomVariable]:
     """Build a RandomVariable Op according to the list of parameters.
 
@@ -202,6 +203,8 @@ def make_hssm_rv(
         A list of str of all parameters for this `RandomVariable`.
     lapse : optional
         A bmb.Prior object representing the lapse distribution.
+    is_choice_only : bool
+        Whether the model is a choice-only model. This parameter overrides
 
     Returns
     -------
@@ -225,7 +228,12 @@ class HSSMRV(RandomVariable):
         # parameter is a scalar. The string to the right of the
         # `->` sign describes the output signature, which is `(2)`, which means the
         # random variable is a length-2 array.
-        signature: str = f"{','.join(['()'] * len(list_params))}->({obs_dim_int})"
+
+        # Override the output from ssm_simulator based on whether the model is
+        # choice-only.
+        output = "()" if is_choice_only else f"({obs_dim_int})"
+        signature: str = f"{','.join(['()'] * len(list_params))}->{output}"
+
         dtype: str = "floatX"
         _print_name: tuple[str, str] = ("SSM", "\\operatorname{SSM}")
         _list_params = list_params
@@ -389,6 +397,7 @@ def make_distribution(
     extra_fields: list[np.ndarray] | None = None,
     fixed_vector_params: dict[str, np.ndarray] | None = None,
     params_is_trialwise: list[bool] | None = None,
+    is_choice_only: bool = False,
 ) -> type[pm.Distribution]:
     """Make a `pymc.Distribution`.
 
@@ -437,6 +446,8 @@ def make_distribution(
         that vmapped JAX log-likelihoods receive consistently shaped inputs,
         regardless of whether Bambi produces ``(1,)`` or ``(n_obs,)`` tensors.
         When ``None``, no graph-level broadcasting is applied.
+    is_choice_only : optional
+        Whether the model is a choice-only model.
 
     Returns
     -------
@@ -561,12 +572,15 @@ def logp(data, *dist_params):  # pylint: disable=E0213
                         "lapse_func is not defined. "
                         "Make sure lapse is properly initialized."
                     )
-                lapse_logp = lapse_func(data[:, 0].eval())
+                data_for_lapse = data if is_choice_only else data[:, 0]
+                lapse_logp = lapse_func(data_for_lapse.eval())
+
                 # AF-TODO potentially apply clipping here
                 logp = loglik(data, *dist_params, *extra_fields)
                 # Ensure that non-decision time is always smaller than rt.
                 # Assuming that the non-decision time parameter is always named "t".
-                logp = ensure_positive_ndt(data, logp, list_params, dist_params)
+                if not is_choice_only:
+                    logp = ensure_positive_ndt(data, logp, list_params, dist_params)
                 logp = pt.log(
                     (1.0 - p_outlier) * pt.exp(logp)
                     + p_outlier * pt.exp(lapse_logp)
@@ -575,7 +589,8 @@ def logp(data, *dist_params):  # pylint: disable=E0213
             else:
                 logp = loglik(data, *dist_params, *extra_fields)
                 # Ensure that non-decision time is always smaller than rt.
-                logp = ensure_positive_ndt(data, logp, list_params, dist_params)
+                if not is_choice_only:
+                    logp = ensure_positive_ndt(data, logp, list_params, dist_params)
 
             if bounds is not None:
                 logp = apply_param_bounds_to_loglik(
@@ -593,6 +608,7 @@ def make_distribution_for_supported_model(
     backend: Literal["pytensor", "jax", "other"] = "pytensor",
     reg_params: list[str] | None = None,
     lapse: bmb.Prior | None = None,
+    is_choice_only: bool = False,
 ) -> type[pm.Distribution]:
     """Make a pm.Distribution class for a supported model.
 
@@ -614,6 +630,8 @@ class that can be used for PyMC modeling.
         parameters are assumed.
     lapse : optional
         A bmb.Prior object representing the lapse distribution.
+    is_choice_only : optional
+        Whether the model is a choice-only model.
     """
     supported_models = get_args(SupportedModels)
     if model not in supported_models:
@@ -643,6 +661,7 @@ class that can be used for PyMC modeling.
         list_params=config.list_params,
         bounds=config.bounds,
         lapse=lapse,
+        is_choice_only=is_choice_only,
     )
 
 

src/hssm/hssm.py

@@ -43,6 +43,7 @@
     assemble_callables,
     make_distribution,
     make_family,
+    make_hssm_rv,
     make_likelihood_callable,
     make_missing_data_callable,
 )
@@ -396,50 +397,64 @@ def __init__(
         self.model_config.validate()
 
         # Set up shortcuts so old code will work
-        self.response = self.model_config.response
+        self.response = self.model_config.response[:]
         self.list_params = self.model_config.list_params
         self.choices = self.model_config.choices
         self.model_name = self.model_config.model_name
         self.loglik = self.model_config.loglik
         self.loglik_kind = self.model_config.loglik_kind
         self.extra_fields = self.model_config.extra_fields
 
+        self.response = cast("list[str]", self.response)
+        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.n_choices = len(self.choices)
 
         self._validate_choices()
         self._pre_check_data_sanity()
 
         # Process missing data setting
         # AF-TODO: Could be a function in data validator?
-        if isinstance(missing_data, float):
-            if not ((self.data.rt == missing_data).any()):
-                raise ValueError(
-                    f"missing_data argument is provided as a float {missing_data}, "
-                    f"However, you have no RTs of {missing_data} in your dataset!"
-                )
+        if self.is_choice_only and missing_data:
+            raise ValueError("Choice-only models cannot have missing data.")
+
+        if not self.is_choice_only:
+            if isinstance(missing_data, float):
+                if not ((self.data.rt == missing_data).any()):
+                    raise ValueError(
+                        f"missing_data argument is provided as a float {missing_data}, "
+                        f"However, you have no RTs of {missing_data} in your dataset!"
+                    )
+                else:
+                    self.missing_data = True
+                    self.missing_data_value = missing_data
+            elif isinstance(missing_data, bool):
+                if missing_data and (not (self.data.rt == -999.0).any()):
+                    raise ValueError(
+                        "missing_data argument is provided as True, "
+                        " so RTs of -999.0 are treated as missing. \n"
+                        "However, you have no RTs of -999.0 in your dataset!"
+                    )
+                elif (not missing_data) and (self.data.rt == -999.0).any():
+                    # self.missing_data = True
+                    raise ValueError(
+                        "Missing data provided as False. \n"
+                        "However, you have RTs of -999.0 in your dataset!"
+                    )
+                else:
+                    self.missing_data = missing_data
             else:
-                self.missing_data = True
-                self.missing_data_value = missing_data
-        elif isinstance(missing_data, bool):
-            if missing_data and (not (self.data.rt == -999.0).any()):
                 raise ValueError(
-                    "missing_data argument is provided as True, "
-                    " so RTs of -999.0 are treated as missing. \n"
-                    "However, you have no RTs of -999.0 in your dataset!"
+                    "missing_data argument must be a bool or a float! \n"
+                    f"You provided: {type(missing_data)}"
                 )
-            elif (not missing_data) and (self.data.rt == -999.0).any():
-                # self.missing_data = True
-                raise ValueError(
-                    "Missing data provided as False. \n"
-                    "However, you have RTs of -999.0 in your dataset!"
-                )
-            else:
-                self.missing_data = missing_data
         else:
-            raise ValueError(
-                "missing_data argument must be a bool or a float! \n"
-                f"You provided: {type(missing_data)}"
-            )
+            self.missing_data = False
 
         if isinstance(deadline, str):
             self.deadline = True
@@ -501,7 +516,8 @@ def __init__(
         self.p_outlier = self.params.get("p_outlier")
         self.lapse = lapse if self.has_lapse else None
 
-        self._post_check_data_sanity()
+        if not self.is_choice_only:
+            self._post_check_data_sanity()
 
         self.model_distribution = self._make_model_distribution()
 
@@ -1358,9 +1374,15 @@ def set_alias(self, aliases: dict[str, str | dict]):
 
     @property
     def response_c(self) -> str:
-        """Return the response variable names in c() format."""
+        """Return the response variable names in c() format.
+
+        New in 0.2.12: when model is choice-only and has deadline, the response
+        is not in the form of c(...).
+        """
         if self.response is None:
-            return "c()"
+            raise ValueError("Response is not defined.")
+        if self.is_choice_only and not self.deadline:
+            return self.response[0]
         return f"c({', '.join(self.response)})"
 
     @property
@@ -2096,6 +2118,31 @@ def _make_model_distribution(self) -> type[pm.Distribution]:
                 self.missing_data_value,
             )
 
+        if self.is_choice_only and self.model_config.rv is None:
+            _logger.warning(
+                "You are building a choice-only model without specifying "
+                "a RandomVariable class. Using a dummy simulator function. "
+                "Simulating data from this model will result in an error."
+            )
+
+            def dummy_simulator_func(*args, **kwargs):
+                raise NotImplementedError(
+                    "You are trying to simulate data from a choice-only model "
+                    "without specifying a RandomVariable class. Please specify "
+                    "a RandomVariable class via the `model_config.rv` argument."
+                )
+
+            setattr(dummy_simulator_func, "model_name", self.model_name)
+            setattr(dummy_simulator_func, "choices", self.choices)
+            setattr(dummy_simulator_func, "obs_dim", 1)
+
+            self.model_config.rv = make_hssm_rv(
+                dummy_simulator_func,
+                list_params=self.list_params,
+                lapse=self.lapse,
+                is_choice_only=True,
+            )
+
         self.data = _rearrange_data(self.data)
 
         # Collect fixed-vector params to substitute in the distribution logp
@@ -2118,6 +2165,7 @@ def _make_model_distribution(self) -> type[pm.Distribution]:
             ),
             fixed_vector_params=fixed_vector_params if fixed_vector_params else None,
             params_is_trialwise=params_is_trialwise_base,
+            is_choice_only=self.is_choice_only,
         )
 
     def _get_deterministic_var_names(self, idata) -> list[str]: