Merge pull request #428 from flatironinstitute/development

Development

Author: BalzaniEdoardo

.gitignore

@@ -153,3 +153,7 @@ docs/data/
 
 # rst generated files
 docs/stubs
+
+# NPZ/NPY files
+*.npz
+*.npy

pyproject.toml

@@ -61,6 +61,7 @@ dev = [
     "dandi",                        # Required by doctest for fetch module
     "seaborn",                      # Required by doctest for _documentation_utils module
     "myst-nb",                      # Test myst_nb utils for glue
+    "hmmlearn",                     # Test algorithmic implementations of HMM
 ]
 docs = [
     "numpydoc",

scripts/check_parameter_naming.py

@@ -77,12 +77,6 @@
     {"flat_dict", "flat_map_dict"},
     {"fit_params", "init_params"},
     {"args", "arg"},
-    {"tol", "atol"},
-    {"tol", "rtol"},
-    {"fit_params", "flat_params"},
-    {"solver_kwargs", "solver_init_kwargs"},
-    {"unaccepted_name", "accepted_name"},
-    {"fn", "fun"},
     {"initialize_init_proba", "initialize_transition_proba"},
     *(
         {a, b}
@@ -95,6 +89,14 @@
             r=2,
         )
     ),
+    {"likelihood_func", "log_likelihood_func"},
+    {"negative_log_likelihood_func", "log_likelihood_func"},
+    {"tol", "atol"},
+    {"tol", "rtol"},
+    {"fit_params", "flat_params"},
+    {"solver_kwargs", "solver_init_kwargs"},
+    {"unaccepted_name", "accepted_name"},
+    {"fn", "fun"},
 ]
 
 

src/nemos/basis/_transformer_basis.py

@@ -335,7 +335,7 @@ def __getattr__(self, name: str):
         if name in self._wrapped_methods:
             return self._wrapped_methods[name]
 
-        if not hasattr(self._basis, name) or name == "to_transformer":
+        if not hasattr(self._basis, name):
             raise AttributeError(f"'TransformerBasis' object has no attribute '{name}'")
 
         # Get the original attribute from the basis
@@ -463,7 +463,6 @@ def __dir__(self) -> list[str]:
         """Extend the list of properties of methods with the ones from the underlying Basis."""
         unique_attrs = set(list(super().__dir__()) + list(self.basis.__dir__()))
         # discard without raising errors if not present
-        unique_attrs.discard("to_transformer")
         return list(unique_attrs)
 
     def __add__(self, other: TransformerBasis | Basis) -> TransformerBasis:

src/nemos/glm_hmm/expectation_maximization.py

@@ -338,6 +338,38 @@ def body_fn(carry, xs):
     return betas
 
 
+def initialize_new_session(n_samples, is_new_session):
+    """Initialize new session indicator."""
+    # Revise if the data is one single session or multiple sessions.
+    # If new_sess is not provided, assume one session
+    if is_new_session is None:
+        # default: all False, but first time bin must be True
+        is_new_session = jax.lax.dynamic_update_index_in_dim(
+            jnp.zeros(n_samples, dtype=bool), True, 0, axis=0
+        )
+    else:
+        # use the user-provided tree, but force the first time bin to be True
+        is_new_session = jax.lax.dynamic_update_index_in_dim(
+            jnp.asarray(is_new_session, dtype=bool), True, 0, axis=0
+        )
+
+    return is_new_session
+
+
+def compute_rate_per_state(X, glm_params, inverse_link_function):
+    """Compute the GLM mean per state."""
+    coef, intercept = glm_params
+
+    # Predicted y
+    if coef.ndim > 2:
+        predicted_rate_given_state = inverse_link_function(
+            jnp.einsum("ik, kjw->ijw", X, coef) + intercept
+        )
+    else:
+        predicted_rate_given_state = inverse_link_function(X @ coef + intercept)
+    return predicted_rate_given_state
+
+
 @partial(jax.jit, static_argnames=["inverse_link_function", "likelihood_func"])
 def forward_backward(
     X: Array,
@@ -412,30 +444,12 @@ def forward_backward(
     ----------
     .. [1] Bishop, C. M. (2006). *Pattern recognition and machine learning*. Springer.
     """
-    coef, intercept = glm_params
     # Initialize variables
     n_time_bins = X.shape[0]
-
-    # Revise if the data is one single session or multiple sessions.
-    # If new_sess is not provided, assume one session
-    if is_new_session is None:
-        # default: all False, but first time bin must be True
-        is_new_session = jax.lax.dynamic_update_index_in_dim(
-            jnp.zeros(y.shape[0], dtype=bool), True, 0, axis=0
-        )
-    else:
-        # use the user-provided tree, but force the first time bin to be True
-        is_new_session = jax.lax.dynamic_update_index_in_dim(
-            jnp.asarray(is_new_session, dtype=bool), True, 0, axis=0
-        )
-
-    # Predicted y
-    if coef.ndim > 2:
-        predicted_rate_given_state = inverse_link_function(
-            jnp.einsum("ik, kjw->ijw", X, coef) + intercept
-        )
-    else:
-        predicted_rate_given_state = inverse_link_function(X @ coef + intercept)
+    is_new_session = initialize_new_session(y.shape[0], is_new_session)
+    predicted_rate_given_state = compute_rate_per_state(
+        X, glm_params, inverse_link_function
+    )
 
     # Compute likelihood given the fixed weights
     # Data likelihood p(y|z) from emissions model
@@ -636,3 +650,138 @@ def run_m_step(
     optimized_projection_weights, state = solver_run(glm_params, X, y, posteriors)
 
     return optimized_projection_weights, new_initial_prob, new_transition_prob, state
+
+
+def max_sum(
+    X: Array,
+    y: Array,
+    initial_prob: Array,
+    transition_prob: Array,
+    glm_params: Tuple[Array, Array],
+    inverse_link_function: Callable,
+    log_likelihood_func: Callable[[Array, Array], Array],
+    is_new_session: Array | None = None,
+    return_index: bool = False,
+):
+    """
+    Find maximum a posteriori (MAP) state path via the max-sum algorithm.
+
+    This function implements the max-sum algorithm for a GLM-HMM, also known as Viterbi algorithm.
+
+    Parameters
+    ----------
+    X :
+        Design matrix, pytree with leaves of shape ``(n_time_bins, n_features)``.
+
+    y :
+        Observations, pytree with leaves of shape ``(n_time_bins,)``.
+
+    initial_prob :
+        Initial latent state probability, pytree with leaves of shape ``(n_states, 1)``.
+
+    transition_prob :
+        Latent state transition matrix, pytree with leaves of shape ``(n_states, n_states)``.
+        ``transition_prob[i, j]`` is the probability of transitioning from state ``i`` to state ``j``.
+
+    glm_params :
+        Length two tuple with the GLM coefficients of shape ``(n_features, n_states)``
+        and intercept of shape ``(n_states,)``.
+
+    inverse_link_function :
+        Function mapping linear predictors to the mean of the observation distribution
+        (e.g., exp for Poisson, sigmoid for Bernoulli).
+
+    is_new_session :
+        Boolean array marking the start of a new session.
+        If unspecified or empty, treats the full set of trials as a single session.
+
+    return_index:
+        If False, return 1-hot encoded map states, if True, return map state indices.
+
+    Returns
+    -------
+    map_path:
+        The MAP state path.
+
+    """
+    is_new_session = initialize_new_session(y.shape[0], is_new_session)
+    predicted_rate_given_state = compute_rate_per_state(
+        X, glm_params, inverse_link_function
+    )
+    log_emission = log_likelihood_func(y, predicted_rate_given_state)
+
+    log_transition = jnp.log(transition_prob)
+    log_init = jnp.log(initial_prob)
+    n_states = initial_prob.shape[0]
+
+    def forward_max_sum(omega_prev, xs):
+        log_em, is_new_sess = xs
+
+        def reset_chain(omega_prev, log_em):
+            # New session: reset to initial distribution
+            omega = log_init + log_em
+            max_prob_state = jnp.full(n_states, -1)  # Boundary marker
+            return omega, max_prob_state
+
+        def continue_chain(omega_prev, log_em):
+            # Continue existing session: Viterbi step
+            step = log_em[None, :] + log_transition + omega_prev[:, None]
+            max_prob_state = jnp.argmax(step, axis=0)
+            omega = step[max_prob_state, jnp.arange(n_states)]
+            return omega, max_prob_state
+
+        omega, max_prob_state = jax.lax.cond(
+            is_new_sess,
+            reset_chain,
+            continue_chain,
+            omega_prev,
+            log_em,
+        )
+
+        return omega, (omega, max_prob_state)
+
+    init_omega = log_init + log_emission[0]
+    _, (omegas, max_prob_states) = jax.lax.scan(
+        forward_max_sum, init_omega, (log_emission[1:], is_new_session[1:])
+    )
+
+    # Backward pass
+    best_final_state = jnp.argmax(omegas[-1])
+    # Prepend initial omega and exclude last one, which is already considered.
+    omegas = jnp.concatenate([init_omega[None, :], omegas[:-1]], axis=0)
+
+    def backward_max_sum(current_state_idx, xs):
+        max_prob_st, omega_t = xs
+
+        def session_boundary(state_idx, max_prob, omega):
+            # Hit a session start, pick best state at this boundary
+            return jnp.argmax(omega)
+
+        def continue_backward(state_idx, max_prob, omega):
+            # Normal backtracking
+            return max_prob[state_idx]
+
+        is_boundary = max_prob_st[current_state_idx] == -1
+
+        prev_state_idx = jax.lax.cond(
+            is_boundary,
+            session_boundary,
+            continue_backward,
+            current_state_idx,
+            max_prob_st,
+            omega_t,
+        )
+
+        return prev_state_idx, prev_state_idx
+
+    _, map_path = jax.lax.scan(
+        backward_max_sum, best_final_state, (max_prob_states, omegas), reverse=True
+    )
+
+    # Append the final state
+    map_path = jnp.concatenate([map_path, jnp.array([best_final_state])])
+
+    if not return_index:
+        map_path = jax.nn.one_hot(map_path, n_states, dtype=jnp.int32)
+
+    return map_path