Rlssm class make model dist

.pre-commit-config.yaml

@@ -9,9 +9,11 @@ repos:
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]
+        exclude: ^tests/
       - id: ruff-format
   - repo: https://github.com/pre-commit/mirrors-mypy
     rev: v1.19.1 # Use the sha / tag you want to point at
     hooks:
       - id: mypy
         args: [--no-strict-optional, --ignore-missing-imports]
+        exclude: ^tests/

src/hssm/__init__.py

@@ -11,14 +11,15 @@
 import logging
 import sys
 
-from .config import ModelConfig
+from .config import ModelConfig, RLSSMConfig
 from .datasets import load_data
 from .defaults import show_defaults
 from .hssm import HSSM
 from .link import Link
 from .param import UserParam as Param
 from .prior import Prior
 from .register import register_model
+from .rl import RLSSM
 from .simulator import simulate_data
 from .utils import check_data_for_rl, set_floatX
 
@@ -31,6 +32,8 @@
 
 __all__ = [
     "HSSM",
+    "RLSSM",
+    "RLSSMConfig",
     "Link",
     "load_data",
     "ModelConfig",

src/hssm/base.py

@@ -264,8 +264,6 @@ class HSSMBase(ABC, DataValidatorMixin, MissingDataMixin):
         The jitter value for the initial values.
     """
 
-    config_class = Config
-
     def __init__(
         self,
         data: pd.DataFrame,
@@ -529,7 +527,8 @@ def _build_model_config(
             A complete Config object with choices and other settings applied.
         """
         # Start with defaults
-        config = cls.config_class.from_defaults(model, loglik_kind)
+        # get_config_class is provided by Config/RLSSMConfig mixin through MRO
+        config = cls.get_config_class().from_defaults(model, loglik_kind)  # type: ignore[attr-defined]
 
         # Handle user-provided model_config
         if model_config is not None:

src/hssm/config.py

@@ -38,6 +38,7 @@
     "decision_process_loglik_kind",
     "learning_process_loglik_kind",
     "extra_fields",
+    "ssm_logp_func",
 )
 
 ParamSpec = Union[float, dict[str, Any], Prior, None]
@@ -67,6 +68,16 @@ class BaseModelConfig(ABC):
     # Additional data requirements
     extra_fields: list[str] | None = None
 
+    @classmethod
+    @abstractmethod
+    def get_config_class(cls) -> type["BaseModelConfig"]:
+        """Return the config class for this model type.
+
+        This enables polymorphic config resolution without circular imports.
+        Each subclass returns itself as the config class.
+        """
+        ...
+
     @abstractmethod
     def validate(self) -> None:
         """Validate configuration. Must be implemented by subclasses."""
@@ -91,6 +102,11 @@ def __post_init__(self):
         if self.loglik_kind is None:
             raise ValueError("loglik_kind is required for Config")
 
+    @classmethod
+    def get_config_class(cls) -> type["Config"]:
+        """Return Config as the config class for HSSM models."""
+        return Config
+
     @classmethod
     def from_defaults(
         cls, model_name: SupportedModels | str, loglik_kind: LoglikKind | None
@@ -266,19 +282,41 @@ class RLSSMConfig(BaseModelConfig):
 
     This configuration class is designed for models that combine reinforcement
     learning processes with sequential sampling decision models (RLSSM).
+
+    The ``ssm_logp_func`` field holds the fully annotated JAX SSM log-likelihood
+    function (an :class:`AnnotatedFunction`) that is passed directly to
+    ``make_rl_logp_op``.  It supersedes the ``loglik`` / ``loglik_kind`` workflow
+    used by :class:`HSSM`: the Op is built from ``ssm_logp_func`` and therefore
+    no ``loglik`` callable needs to be provided.
     """
 
     decision_process_loglik_kind: str = field(kw_only=True)
     learning_process_loglik_kind: str = field(kw_only=True)
     params_default: list[float] = field(kw_only=True)
     decision_process: str | ModelConfig = field(kw_only=True)
     learning_process: dict[str, Any] = field(kw_only=True)
+    # The fully annotated SSM log-likelihood function used by make_rl_logp_op.
+    # Type is Any to avoid a hard dependency on the AnnotatedFunction Protocol at
+    # import time; validated at runtime in validate().
+    ssm_logp_func: Any = field(default=None, kw_only=True)
 
     def __post_init__(self):
         """Set default loglik_kind for RLSSM models if not provided."""
         if self.loglik_kind is None:
             self.loglik_kind = "approx_differentiable"
 
+    @classmethod
+    def get_config_class(cls) -> type["RLSSMConfig"]:
+        """Return RLSSMConfig as the config class for RLSSM models."""
+        return RLSSMConfig
+
+    @classmethod
+    def from_defaults(
+        cls, model_name: SupportedModels | str, loglik_kind: LoglikKind | None
+    ) -> Config:
+        """Return the shared Config defaults (delegated to :class:`Config`)."""
+        return Config.from_defaults(model_name, loglik_kind)
+
     @property
     def n_params(self) -> int | None:
         """Return the number of parameters."""
@@ -290,17 +328,16 @@ def n_extra_fields(self) -> int | None:
         return len(self.extra_fields) if self.extra_fields else None
 
     @classmethod
-    def from_rlssm_dict(cls, model_name: str, config_dict: dict[str, Any]):
+    def from_rlssm_dict(cls, config_dict: dict[str, Any]) -> "RLSSMConfig":
         """
         Create RLSSMConfig from a configuration dictionary.
 
         Parameters
         ----------
-        model_name : str
-            The name of the RLSSM model.
         config_dict : dict[str, Any]
             Dictionary containing model configuration. Expected keys:
-                - description: Model description (optional)
+                - model_name: Model identifier (required)
+                - description: Model description (required)
                 - list_params: List of parameter names (required)
                 - extra_fields: List of extra field names from data (required)
                 - params_default: Default parameter values (required)
@@ -332,6 +369,7 @@ def from_rlssm_dict(cls, model_name: str, config_dict: dict[str, Any]):
             params_default=config_dict["params_default"],
             decision_process=config_dict["decision_process"],
             learning_process=config_dict["learning_process"],
+            ssm_logp_func=config_dict["ssm_logp_func"],
             bounds=config_dict.get("bounds", {}),
             response=config_dict["response"],
             choices=config_dict["choices"],
@@ -355,6 +393,32 @@ def validate(self) -> None:
             raise ValueError("Please provide `choices` in the configuration.")
         if self.decision_process is None:
             raise ValueError("Please specify a `decision_process`.")
+        if self.ssm_logp_func is None:
+            raise ValueError(
+                "Please provide `ssm_logp_func`: the fully annotated JAX SSM "
+                "log-likelihood function required by `make_rl_logp_op`."
+            )
+        if not callable(self.ssm_logp_func):
+            raise ValueError(
+                "`ssm_logp_func` must be a callable, "
+                f"but got {type(self.ssm_logp_func)!r}."
+            )
+        missing_attrs = [
+            attr
+            for attr in ("inputs", "outputs", "computed")
+            if not hasattr(self.ssm_logp_func, attr)
+        ]
+        if missing_attrs:
+            raise ValueError(
+                "`ssm_logp_func` must be decorated with `@annotate_function` "
+                "so that it carries the attributes required by `make_rl_logp_op`. "
+                f"Missing attribute(s): {missing_attrs}. "
+                "Decorate the function like:\n\n"
+                "    @annotate_function(\n"
+                "        inputs=[...], outputs=[...], computed={...}\n"
+                "    )\n"
+                "    def my_ssm_logp(lan_matrix): ..."
+            )
 
         # Validate parameter defaults consistency
         if self.params_default and self.list_params:
@@ -445,6 +509,27 @@ def to_config(self) -> Config:
             loglik=self.loglik,
         )
 
+    def to_model_config(self) -> "ModelConfig":
+        """Build a :class:`ModelConfig` from this :class:`RLSSMConfig`.
+
+        All fields are sourced from ``self``; the backend is fixed to ``"jax"``
+        because RLSSM exclusively uses the JAX backend.
+
+        ``default_priors`` is intentionally left empty so the
+        ``prior_settings="safe"`` mechanism in :class:`~hssm.base.HSSMBase`
+        assigns sensible priors from bounds rather than fixing every parameter
+        to a constant scalar.
+        """
+        return ModelConfig(
+            response=tuple(self.response),  # type: ignore[arg-type]
+            list_params=list(self.list_params),  # type: ignore[arg-type]
+            choices=tuple(self.choices),  # type: ignore[arg-type]
+            default_priors={},
+            bounds=self.bounds,
+            extra_fields=self.extra_fields,
+            backend="jax",
+        )
+
 
 @dataclass
 class ModelConfig:

src/hssm/data_validator.py

@@ -17,7 +17,7 @@ class DataValidatorMixin:
     This class expects subclasses to define the following attributes:
     - data: pd.DataFrame
     - response: list[str]
-    - choices: list[int]
+    - choices: list[int] | tuple[int, ...]
     - n_choices: int
     - extra_fields: list[str] | None
     - deadline: bool
@@ -30,7 +30,7 @@ def __init__(
         self,
         data: pd.DataFrame,
         response: list[str] | None = ["rt", "response"],
-        choices: list[int] | None = [0, 1],
+        choices: tuple[int, ...] | None = (0, 1),
         n_choices: int = 2,
         extra_fields: list[str] | None = None,
         deadline: bool = False,

src/hssm/hssm.py

@@ -30,6 +30,7 @@
 )
 
 from .base import HSSMBase
+from .config import Config
 
 if TYPE_CHECKING:
     from os import PathLike
@@ -74,7 +75,7 @@ def __get__(self, instance, owner):  # noqa: D105
         return self.fget(owner)
 
 
-class HSSM(HSSMBase):
+class HSSM(HSSMBase, Config):
     """The basic Hierarchical Sequential Sampling Model (HSSM) class.
 
     Parameters

src/hssm/rl/__init__.py

@@ -0,0 +1,18 @@
+"""Reinforcement learning extensions for HSSM.
+
+This sub-package provides:
+
+- :class:`~hssm.rl.rlssm.RLSSM` — the RL + SSM model class.
+- :func:`~hssm.rl.utils.validate_balanced_panel` — panel-balance utility.
+- :mod:`hssm.rl.likelihoods` — log-likelihood builders
+  (:func:`~hssm.rl.likelihoods.builder.make_rl_logp_func`,
+  :func:`~hssm.rl.likelihoods.builder.make_rl_logp_op`).
+"""
+
+from .rlssm import RLSSM
+from .utils import validate_balanced_panel
+
+__all__ = [
+    "RLSSM",
+    "validate_balanced_panel",
+]