Update HSSM class for choice-only models

.github/setup-env-notebooks/action.yml

@@ -28,5 +28,5 @@ runs:
 
     - name: Install hssm
       if: steps.cache.outputs.cache-hit != 'true'
-      run: uv sync --group test --group notebook
+      run: uv sync --group notebook
       shell: bash

.github/setup-env/action.yml

@@ -28,5 +28,5 @@ runs:
 
     - name: Install hssm
       if: steps.cache.outputs.cache-hit != 'true'
-      run: uv sync --group test
+      run: uv sync
       shell: bash

.github/workflows/build_and_publish.yml

@@ -51,4 +51,4 @@ jobs:
           uv publish --token ${{ secrets.PYPI_TOKEN }}
 
       - name: Build and publish docs
-        run: uv run mkdocs gh-deploy --force
+        run: uv run --group notebook --group docs mkdocs gh-deploy --force

.github/workflows/build_docs.yml

@@ -22,4 +22,4 @@ jobs:
           python-version: ${{ matrix.python-version }}
 
       - name: Run mkdocs
-        run: uv run mkdocs gh-deploy --force
+        run: uv run --group notebook --group docs mkdocs gh-deploy --force

.github/workflows/check_notebooks.yml

@@ -54,7 +54,7 @@ jobs:
 
             echo "Cleaning $notebook"
             uv run nb-clean clean -o "$notebook"
-            if ! uv run jupyter nbconvert --ExecutePreprocessor.timeout=10000 --to notebook --execute "$notebook"; then
+            if ! uv run --group notebook jupyter nbconvert --ExecutePreprocessor.timeout=10000 --to notebook --execute "$notebook"; then
               echo "::error::Failed to execute notebook: $notebook"
               EXIT_CODE=1
             fi

pyproject.toml

@@ -28,7 +28,6 @@ dependencies = [
     "arviz>=0.22.0",
     "bambi>=0.17.2",
     "cloudpickle>=3.0.0",
-    "formulae>=0.6.0",
     "hddm-wfpt>=0.1.6",
     "huggingface-hub>=0.34.0",
     "jaxonnxruntime>=0.3.0",
@@ -52,39 +51,40 @@ bayesflow = ["bayesflow", "keras>=3"]
 [dependency-groups]
 dev = [
     "coverage>=7.6.4",
+    "mypy>=1.11.1",
+    "pytest-cov>=6.0.0",
+    "pytest-xdist>=3.6.1",
+    "pytest>=8.3.1",
+    "pytest-random-order>=1.1.1",
+    "pytest-rerunfailures>=15.0",
+    "ruff>=0.15.0",
+    "pre-commit>=4.1.0",
+    "onnxruntime>=1.17.1",
+]
+
+notebook = [
     "graphviz>=0.20.3",
     "ipykernel>=6.29.5",
     "ipython>=8.31.0",
     "ipywidgets>=8.1.2",
     "jupyterlab>=4.2.4",
     "mistune>=3.0.2",
-    "mkdocs-material>=9.5.21",
-    "mkdocs>=1.6.0",
-    "mkdocs-jupyter>=0.25.1",
-    "mkdocstrings-python>=1.10.0",
     "nbconvert>=7.16.5",
-    "onnxruntime>=1.17.1",
     "nb-clean>=4.0.1",
     "nbval>=0.11.0",
-    "pre-commit>=2.20.0",
     "ptpython>=3.0.29",
     "pyarrow>=20.0.0",
     "lanfactory>=0.5.3",
-    "HSSM[test]",
+    "zeus-mcmc>=2.5.4",
 ]
 
-test = [
-    "mypy>=1.11.1",
-    "pytest-cov>=6.0.0",
-    "pytest-xdist>=3.6.1",
-    "pytest>=8.3.1",
-    "ruff>=0.15.0",
-    "pytest-random-order>=1.1.1",
-    "pytest-rerunfailures>=15.0",
+docs = [
+    "mkdocs-material>=9.5.21",
+    "mkdocs>=1.6.0",
+    "mkdocs-jupyter>=0.25.1",
+    "mkdocstrings-python>=1.10.0",
 ]
 
-notebook = ["ipykernel>=6.29.5", "zeus-mcmc>=2.5.4"]
-
 [tool.ruff]
 line-length = 88
 target-version = "py310"

src/hssm/data_validator.py

@@ -50,6 +50,8 @@ 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():
                 raise ValueError(
@@ -106,9 +108,8 @@ def _handle_missing_data_and_deadline(self):
         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.
-            # TODO: commented out for now for tests to pass
-            # if self.is_choice_only:
-            #     return
+            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.
 
     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
@@ -385,10 +393,11 @@ def make_distribution(
     loglik: LogLikeFunc | pytensor.graph.Op,
     list_params: list[str],
     bounds: dict | None = None,
-    lapse: bmb.Prior | None = None,
+    lapse: float | bmb.Prior | None = None,
     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`.
 
@@ -418,7 +427,7 @@ def make_distribution(
         A dictionary with parameters as keys (a string) and its boundaries as values.
         Example: {"parameter": (lower_boundary, upper_boundary)}.
     lapse : optional
-        A bmb.Prior object representing the lapse distribution.
+        A float or bmb.Prior object representing the lapse distribution.
     extra_fields : optional
         An optional list of arrays that are stored in the class created and will be
         used in likelihood calculation. Defaults to None.
@@ -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
     -------
@@ -481,15 +492,18 @@ def make_distribution(
         if list_params[-1] != "p_outlier":
             list_params.append("p_outlier")
 
-        data_vector = pt.dvector()
-        lapse_logp = pm.logp(
-            get_distribution_from_prior(lapse).dist(**lapse.args),
-            data_vector,
-        )
-        lapse_func = pytensor.function(
-            [data_vector],
-            lapse_logp,
-        )
+        if isinstance(lapse, float):
+            lapse_func = lambda data: np.full_like(data, lapse)
+        else:
+            data_vector = pt.dvector()
+            lapse_logp = pm.logp(
+                get_distribution_from_prior(lapse).dist(**lapse.args),
+                data_vector,
+            )
+            lapse_func = pytensor.function(
+                [data_vector],
+                lapse_logp,
+            )
     else:
         lapse_func = None
 
@@ -561,12 +575,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 +592,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 +611,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 +633,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 +664,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,
     )