Added the ability to use safe priors for hierarchical models

.github/workflows/run_tests.yml

@@ -2,7 +2,6 @@ name: Run tests
 
 on:
   pull_request:
-  push:
 
 jobs:
   run_tests:

mkdocs.yml

@@ -126,5 +126,5 @@ markdown_extensions:
   - pymdownx.superfences
   - attr_list
   - pymdownx.emoji:
-      emoji_index: !!python/name:materialx.emoji.twemoji
-      emoji_generator: !!python/name:materialx.emoji.to_svg
+      emoji_index: !!python/name:material.extensions.emoji.twemoji
+      emoji_generator: !!python/name:material.extensions.emoji.to_svg

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "HSSM"
-version = "0.1.5"
+version = "0.2.0"
 description = "Bayesian inference for hierarchical sequential sampling models."
 authors = [
     "Alexander Fengler <[email protected]>",
@@ -23,14 +23,14 @@ numpy = ">=1.23.4,<1.26"
 onnx = "^1.12.0"
 jax = "^0.4.0"
 jaxlib = "^0.4.0"
-ssm-simulators = "0.5.1"
+ssm-simulators = "0.5.3"
 huggingface-hub = "^0.15.1"
 onnxruntime = "^1.15.0"
 bambi = "^0.12.0"
 numpyro = "^0.12.1"
 hddm-wfpt = "^0.1.1"
 seaborn = "^0.13.0"
-pytensor = "<=2.17.3"
+pytensor = "<2.17.4"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^7.3.1"
@@ -69,7 +69,7 @@ profile = "black"
 
 [tool.ruff]
 line-length = 88
-target-version = "py39"
+target-version = "py310"
 unfixable = ["E711"]
 
 select = [
@@ -132,6 +132,8 @@ ignore = [
     "B020",
     # Function definition does not bind loop variable
     "B023",
+    # zip()` without an explicit `strict=
+    "B905",
     # Functions defined inside a loop must not use variables redefined in the loop
     # "B301",  # not yet implemented
     # Too many arguments to function call
@@ -166,14 +168,7 @@ ignore = [
     "TID252",
 ]
 
-exclude = [
-    ".github",
-    "docs",
-    "notebook",
-    "tests",
-    "src/hssm/likelihoods/hddm_wfpt/cdfdif_wrapper.c",
-    "src/hssm/likelihoods/hddm_wfpt/wfpt.cpp",
-]
+exclude = [".github", "docs", "notebook", "tests"]
 
 [tool.ruff.pydocstyle]
 convention = "numpy"

src/hssm/hssm.py

@@ -24,6 +24,7 @@
 import seaborn as sns
 import xarray as xr
 from bambi.model_components import DistributionalComponent
+from bambi.transformations import transformations_namespace
 
 from hssm.defaults import (
     LoglikKind,
@@ -164,6 +165,9 @@ class HSSM:
         recommended when you are using hierarchical models.
         The default value is `None` when `hierarchical` is `False` and `"safe"` when
         `hierarchical` is `True`.
+    extra_namespace : optional
+        Additional user supplied variables with transformations or data to include in
+        the environment where the formula is evaluated. Defaults to `None`.
     **kwargs
         Additional arguments passed to the `bmb.Model` object.
 
@@ -214,6 +218,7 @@ def __init__(
         hierarchical: bool = False,
         link_settings: Literal["log_logit"] | None = None,
         prior_settings: Literal["safe"] | None = None,
+        extra_namespace: dict[str, Any] | None = None,
         **kwargs,
     ):
         self.data = data
@@ -232,6 +237,11 @@ def __init__(
         self.link_settings = link_settings
         self.prior_settings = prior_settings
 
+        additional_namespace = transformations_namespace.copy()
+        if extra_namespace is not None:
+            additional_namespace.update(extra_namespace)
+        self.additional_namespace = additional_namespace
+
         responses = self.data["response"].unique().astype(int)
         self.n_responses = len(responses)
         if self.n_responses == 2:
@@ -312,7 +322,12 @@ def __init__(
         )
 
         self.model = bmb.Model(
-            self.formula, data, family=self.family, priors=self.priors, **other_kwargs
+            self.formula,
+            data=data,
+            family=self.family,
+            priors=self.priors,
+            extra_namespace=extra_namespace,
+            **other_kwargs,
         )
 
         self._aliases = get_alias_dict(self.model, self._parent_param)
@@ -322,6 +337,7 @@ def sample(
         self,
         sampler: Literal["mcmc", "nuts_numpyro", "nuts_blackjax", "laplace", "vi"]
         | None = None,
+        init: str | None = None,
         **kwargs,
     ) -> az.InferenceData | pm.Approximation:
         """Perform sampling using the `fit` method via bambi.Model.
@@ -335,6 +351,9 @@ def sample(
             sampler will automatically be chosen: when the model uses the
             `approx_differentiable` likelihood, and `jax` backend, "nuts_numpyro" will
             be used. Otherwise, "mcmc" (the default PyMC NUTS sampler) will be used.
+        init: optional
+            Initialization method to use for the sampler. If any of the NUTS samplers
+            is used, defaults to `"adapt_diag"`. Otherwise, defaults to `"auto"`.
         kwargs
             Other arguments passed to bmb.Model.fit(). Please see [here]
             (https://bambinos.github.io/bambi/api_reference.html#bambi.models.Model.fit)
@@ -370,7 +389,7 @@ def sample(
                 )
 
             if "step" not in kwargs:
-                kwargs["step"] = pm.Slice(model=self.pymc_model)
+                kwargs |= {"step": pm.Slice(model=self.pymc_model)}
 
         if (
             self.loglik_kind == "approx_differentiable"
@@ -387,7 +406,15 @@ def sample(
         if self._check_extra_fields():
             self._update_extra_fields()
 
-        self._inference_obj = self.model.fit(inference_method=sampler, **kwargs)
+        if init is None:
+            if sampler in ["mcmc", "nuts_numpyro", "nuts_blackjax"]:
+                init = "adapt_diag"
+            else:
+                init = "auto"
+
+        self._inference_obj = self.model.fit(
+            inference_method=sampler, init=init, **kwargs
+        )
 
         return self.traces
 
@@ -629,11 +656,11 @@ def plot_trace(
         data : optional
             An ArviZ InferenceData object. If None, the traces stored in the model will
             be used.
-        include deterministic : optional
+        include_deterministic : optional
             Whether to include deterministic variables in the plot. Defaults to False.
             Note that if include deterministic is set to False and and `var_names` is
             provided, the `var_names` provided will be modified to also exclude the
-         deterministic values. If this is not desirable, set
+            deterministic values. If this is not desirable, set
             `include deterministic` to True.
         tight_layout : optional
             Whether to call plt.tight_layout() after plotting. Defaults to True.
@@ -852,6 +879,8 @@ def _add_kwargs_and_p_outlier_to_include(
         """Process kwargs and p_outlier and add them to include."""
         if include is None:
             include = []
+        else:
+            include = include.copy()
         params_in_include = [param["name"] for param in include]
 
         # Process kwargs
@@ -913,7 +942,7 @@ def _preprocess_rest(self, processed: dict[str, Param]) -> dict[str, Param]:
                     bounds = self.model_config.bounds.get(param_str)
                     param = Param(
                         param_str,
-                        formula="1 + (1|participant_id)",
+                        formula=f"{param_str} ~ 1 + (1|participant_id)",
                         bounds=bounds,
                     )
                 else:
@@ -956,15 +985,27 @@ def _find_parent(self) -> tuple[str, Param]:
 
     def _override_defaults(self):
         """Override the default priors or links."""
+        is_ddm = (
+            self.model_name in ["ddm", "ddm_sdv", "ddm_full"]
+            and self.loglik_kind != "approx_differentiable"
+        )
         for param in self.list_params:
             param_obj = self.params[param]
             if self.prior_settings == "safe":
-                param_obj.override_default_priors(self.data)
+                if is_ddm:
+                    param_obj.override_default_priors_ddm(
+                        self.data, self.additional_namespace
+                    )
+                else:
+                    param_obj.override_default_priors(
+                        self.data, self.additional_namespace
+                    )
             elif self.link_settings == "log_logit":
                 param_obj.override_default_link()
 
     def _process_all(self):
         """Process all params."""
+        assert self.list_params is not None
         for param in self.list_params:
             self.params[param].convert()