lnccbrown · cpaniaguam · Oct 9, 2025 · Oct 9, 2025 · Oct 9, 2025 · Oct 9, 2025
diff --git a/src/hssm/likelihoods/rldm_optimized_abstraction.py b/src/hssm/likelihoods/rldm_optimized_abstraction.py
@@ -0,0 +1,291 @@
+"""The log-likelihood function for the RLDM model."""
+
+import functools
+from typing import Any, Callable
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from jax.lax import scan
+from pytensor.graph import Op
+
+from hssm.distribution_utils.func_utils import make_vjp_func
+
+from ..distribution_utils.jax import make_jax_logp_ops
+from ..distribution_utils.onnx import make_jax_matrix_logp_funcs_from_onnx
+
+# Obtain the angle log-likelihood function from an ONNX model.
+angle_logp_jax_func = make_jax_matrix_logp_funcs_from_onnx(
+    model="angle.onnx",
+)
+
+
+def annotate_function(**kwargs):
+    """Attach arbitrary metadata as attributes to a function.
+
+    Parameters
+    ----------
+    **kwargs
+        Arbitrary keyword arguments to attach as attributes.
+
+    Returns
+    -------
+    Callable
+        Decorator that adds metadata attributes to the wrapped function.
+    """
+
+    def decorator(func: Callable) -> Callable:
+        @functools.wraps(func)
+        def wrapper(*args, **inner_kwargs):
+            return func(*args, **inner_kwargs)
+
+        for key, value in kwargs.items():
+            setattr(wrapper, key, value)
+        return wrapper
+
+    return decorator
+
+
+# Inner function to compute the drift rate and update q-values for each trial.
+# This function is used with `jax.lax.scan` to process each trial in the RLDM model.
+def compute_v_trial_wise(
+    q_val: jnp.ndarray, inputs: jnp.ndarray
+) -> tuple[jnp.ndarray, jnp.ndarray]:
+    """Compute the drift rate and updates the q-values for each trial.
+
+    This function is used with `jax.lax.scan` to process each trial. It takes the
+    current q-values and the RL parameters (rl_alpha, scaler), action (response),
+    and reward (feedback) for the current trial, computes the drift rate, and
+    updates the q-values. The q_values are updated in each iteration and carried
+    forward to the next one.
+
+    Parameters
+    ----------
+    q_val
+        A length-2 jnp array containing the current q-values for the two alternatives.
+        These values are updated in each iteration and carried forward to the next
+        trial.
+    inputs
+        A 2D jnp array containing the RL parameters (rl_alpha, scaler),
+        action (response), and reward (feedback) for the current trial.
+
+    Returns
+    -------
+    tuple
+        A tuple containing the updated q-values and the computed drift rate (v).
+    """
+    rl_alpha, scaler, action, reward = inputs
+    action = jnp.astype(action, jnp.int32)
+
+    # drift rate on each trial depends on difference in expected rewards for
+    # the two alternatives:
+    # drift rate = (q_up - q_low) * scaler where
+    # the scaler parameter describes the weight to put on the difference in
+    # q-values.
+    computed_v = (q_val[1] - q_val[0]) * scaler
+
+    # compute the reward prediction error
+    delta_RL = reward - q_val[action]
+
+    # update the q-values using the RL learning rule (here, simple TD rule)
+    q_val = q_val.at[action].set(q_val[action] + rl_alpha * delta_RL)
+
+    return q_val, computed_v
+
+
+# This function computes the drift rates (v) for each subject by processing
+# their trials one by one. It uses `jax.lax.scan` to efficiently iterate over
+# the trials and compute the drift rates based on the RL parameters, actions,
+# and rewards for each trial.
+def compute_v_subject_wise(
+    subj_trials: jnp.ndarray,
+) -> jnp.ndarray:
+    """Compute the drift rates (v) for a given subject.
+
+    Parameters
+    ----------
+    subj_trials:
+        A jnp array of dimension (n_trials, 4) containing rl_alpha, scaler,
+        action (response), and reward (feedback) for each trial of the subject.
+
+    Returns
+    -------
+    jnp.ndarray
+        The computed drift rates (v) for the RLDM model for the given subject.
+    """
+    _, v = scan(
+        compute_v_trial_wise,
+        jnp.ones(2) * 0.5,  # initial q-values for the two alternatives
+        subj_trials,
+    )
+
+    return v
+
+
+def _get_column_indices(
+    cols_to_look_up: list[str],
+    data_cols: list[str],
+    list_params: list[str] | None,
+    extra_fields: list[str] | None,
+) -> dict[str, tuple[str, int]]:
+    """Return indices for required columns.
+
+    Parameters
+    ----------
+    cols_to_look_up : list[str]
+        Columns to find indices for
+    data_cols : list[str]
+        Available data columns
+    list_params : list[str] | None
+        Available list parameters
+    extra_fields : list[str] | None
+        Available extra fields
+
+    Returns
+    -------
+    dict[str, tuple[str, int]]
+        Mapping of column names to (source, index) tuples
+    """
+    list_params = list_params or []
+    extra_fields = extra_fields or []
+    list_params_extra_fields = list_params + extra_fields
+    colidxs = {}
+    for col in cols_to_look_up:
+        if col in data_cols:
+            colidxs[col] = ("data", data_cols.index(col))
+        elif col in list_params_extra_fields:
+            colidxs[col] = ("args", list_params_extra_fields.index(col))
+        else:
+            raise ValueError(
+                f"Column '{col}' not found in any of `data`, `list_params`, "
+                f"or `extra_fields`."
+            )
+    return colidxs
+
+
+def _collect_cols_arrays(data, _args, colidxs):
+    collected = []
+    for col in colidxs:
+        source, idx = colidxs[col]
+        if source == "data":
+            collected.append(data[:, idx])
+        else:
+            collected.append(_args[idx])
+    return collected
+
+
+def make_rl_logp_func(
+    subject_wise_func: Callable[..., Any],
+    n_participants: int,
+    n_trials: int,
+    data_cols: list[str] = ["rt", "response"],
+    list_params: list[str] | None = None,
+    extra_fields: list[str] | None = None,
+) -> Callable:
+    """Create a function to compute the drift rates (v) for the RLDM model.
+
+    Parameters
+    ----------
+    subject_wise_func : Callable
+        Function that computes drift rates for a subject's trials.
+    n_participants : int
+        Number of participants in the dataset.
+    n_trials : int
+        Number of trials per participant.
+    data_cols : list[str] | None
+        List of column names in the data array.
+    dist_params : list[str] | None
+        List of distribution parameter names required by the RL model.
+    extra_fields : list[str] | None
+        List of extra field names required by the RL model.
+
+    Returns
+    -------
+    Callable
+        A function that computes drift rates (v) for all subjects given their trial data
+        and RLDM parameters.
+    """
+    inputs = subject_wise_func.inputs  # type: ignore[attr-defined]
+    # _validate_columns(data_cols, inputs)
+    colidxs = _get_column_indices(
+        inputs,
+        data_cols,
+        list_params,
+        extra_fields,
+    )
+
+    # Vectorized version of  subject_wise_func to handle multiple subjects.
+    subject_wise_vmapped = jax.vmap(subject_wise_func, in_axes=0)
+
+    def logp(data, *args) -> np.ndarray:
+        """Compute the drift rates (v) for each trial in a reinforcement learning model.
+
+        data : np.ndarray
+            A 2D array containing trial data.
+
+        args: Model parameters included in list_params and extra_fields.
+
+        Notes
+        -----
+        - The function internally reshapes the input data to group trials by
+          participant and applies a vectorized mapping function to compute drift
+          rates.
+        - The function assumes that `n_participants`, `n_trials`, `idxs`, and
+          `subject_wise_vmapped` are defined in the surrounding scope.
+
+        Returns
+        -------
+        np.ndarray
+            The computed drift rates for each trial, reshaped as a 2D array.
+        """
+        # Reshape subj_trials into a 3D array of shape
+        # (n_participants, n_trials, len(args))
+        # so we can act on this object with the vmapped version of the mapping function
+        _data = _collect_cols_arrays(data, args, colidxs)
+
+        subj_trials = jnp.stack(_data, axis=1).reshape(n_participants, n_trials, -1)
+
+        drift_rates = subject_wise_vmapped(subj_trials).reshape((-1, 1))
+        return drift_rates
+
+        # TODO: reintroduce workflow using a jax function and handling the selection
+        # dist_params to stack
+        # create parameter arrays to be passed to the likelihood function
+        # ddm_params_matrix = jnp.stack(dist_params[2:6], axis=1)
+        # lan_matrix = jnp.concatenate((v, ddm_params_matrix, data), axis=1)
+        # return _logp_jax_func(lan_matrix)
+
+    return logp
+
+
+# TODO[CP]: Adapt this function given the changes to make_rl_logp_func
+# pragma: no cover
+def make_rldm_logp_op(
+    subject_wise_func: Callable[..., Any],
+    n_participants: int,
+    n_trials: int,
+    n_params: int,
+) -> Op:
+    """Create a pytensor Op for the likelihood function of RLDM model.
+
+    Parameters
+    ----------
+    n_participants : int
+        The number of participants in the dataset.
+    n_trials : int
+        The number of trials per participant.
+
+    Returns
+    -------
+    Op
+        A pytensor Op that computes the log likelihood for the RLDM model.
+    """
+    logp = make_rl_logp_func(subject_wise_func, n_participants, n_trials)
+    vjp_logp = make_vjp_func(logp, params_only=False, n_params=n_params)
+
+    return make_jax_logp_ops(
+        logp=jax.jit(logp),
+        logp_vjp=jax.jit(vjp_logp),
+        logp_nojit=logp,
+        n_params=n_params,  # rl_alpha, scaler, a, z, t, theta
+    )