HSSM integration API

ssms.hssm_support adapts classic ssms simulator functions to HSSM's random variable interface. It decorates simulator callables with model_name, choices, and obs_dim, validates those attributes, broadcasts parameter and trial covariate shapes, and returns (rt, response) observations in HSSM's expected layout.

This module is the classic SSM bridge. RLSSM inference uses the separate assembled-model contract.

ssms.hssm_support

Functions:

• decorate_atomic_simulator –

  Decorator to add metadata attributes to simulator functions.

• get_simulator_fun_internal –

  Get the internal simulator function for a given model.

• hssm_sim_wrapper –

  Wrap a ssms simulator function to match HSSM's expected interface.

• rng_fn –

  Generate random variables from this distribution using the provided simulator function.

• validate_simulator_fun –

  Validate that the simulator function has required attributes.

ssms.hssm_support.decorate_atomic_simulator

decorate_atomic_simulator(model_name: str, choices: list | ndarray | None = None, obs_dim: int = 2)

Decorator to add metadata attributes to simulator functions.

This decorator attaches the following attributes to the decorated function as expected of simulators in HSSM: - model_name: Name of the model. - choices: List or array of possible choices/responses. - obs_dim: Number of observation dimensions.

Parameters:

• model_name (str) –

  Name of the model.

• choices (list or ndarray, default: None ) –

  List or array of possible choices/responses (default: [-1, 1]).

• obs_dim (int, default: 2 ) –

  Number of observation dimensions (default: 2).

Returns:

• Callable –

  Decorator that adds attributes to the simulator function.

ssms.hssm_support.get_simulator_fun_internal

get_simulator_fun_internal(simulator_fun: Callable | str)

Get the internal simulator function for a given model.

Parameters:

• simulator_fun (Callable or str) –

  The simulator function or the name of the model as a string.

Returns:

• Callable –

  The decorated simulator function.

Raises:

• ValueError –

  If the simulator argument is not a string or a callable.

ssms.hssm_support.hssm_sim_wrapper

hssm_sim_wrapper(simulator_fun, theta, model, n_replicas, random_state, **kwargs)

Wrap a ssms simulator function to match HSSM's expected interface.

Parameters:

• simulator_fun (callable) –

  The simulator function to wrap, which should have the following interface: - theta: array-like, shape (n_trials, n_parameters) - model: str, name of the model to simulate - n_samples: int, number of replica datasets to generate - random_state: int, to be used as the random seed internally - **kwargs: additional keyword arguments

• theta (array - like) –

  Model parameters, shape (n_trials, n_parameters)

• model (str) –

  Name of the model to simulate

• n_replicas (int) –

  Number of replica datasets to generate

• random_state (int or Generator) –

  Random seed or random number generator

• **kwargs –

  Additional keyword arguments passed to simulator_fun

Returns:

• array - like –

  Array of shape (n_trials, 2) containing reaction times and choices stacked column-wise

ssms.hssm_support.rng_fn

rng_fn(arg_arrays: list[ndarray], size: int | tuple | None, rng: Generator, simulator_fun: Callable, obs_dim_int: int, *args, **kwargs) -> np.ndarray

Generate random variables from this distribution using the provided simulator function.

Parameters:

• arg_arrays (list of np.ndarray) –

  List of argument arrays corresponding to model parameters.

• size (int, tuple, or None) –

  The total number of samples to be drawn. If None or 1, only one replica

• rng (Generator) –

  Random number generator for reproducibility.

• simulator_fun (Callable) –

  The simulator function to generate samples.

• obs_dim_int (int) –

  Number of observation dimensions.

• *args (tuple, default: () ) –

  Model parameters, in the order of _list_params, with the last argument as size.

• **kwargs (dict, default: {} ) –

  Additional keyword arguments passed to the simulator function.

Returns:

• tuple[ndarray, ndarray] –

  An array of shape (..., obs_dim_int) containing generated (rt, response) pairs and the p_outlier values if applicable.

ssms.hssm_support.validate_simulator_fun

validate_simulator_fun(simulator_fun: Any) -> tuple[str, list, int]

Validate that the simulator function has required attributes.

Parameters:

• simulator_fun (Any) –

  The simulator function or object to validate.

Returns:

• tuple –

  A tuple containing model_name, choices, and obs_dim_int.

Raises:

• ValueError –

  If any required attribute is missing or invalid.