Tackle Typing and Linting Errors

dynamax/generalized_gaussian_ssm/inference.py

@@ -3,7 +3,6 @@
 from jax import jacfwd, vmap, lax
 import jax.numpy as jnp
 from jax import lax
-from tensorflow_probability.substrates.jax.distributions import MultivariateNormalFullCovariance as MVN
 from jaxtyping import Array, Float
 from typing import NamedTuple, Optional, Union, Callable
 
@@ -83,7 +82,7 @@ def compute_weights_and_sigmas(self, m, P):
 
 
 def _predict(m, P, f, Q, u, g_ev, g_cov):
-    """Predict next mean and covariance under an additive-noise Gaussian filter
+    r"""Predict next mean and covariance under an additive-noise Gaussian filter
 
         p(x_{t+1}) = N(x_{t+1} | mu_pred, Sigma_pred)
         where
@@ -117,7 +116,7 @@ def _predict(m, P, f, Q, u, g_ev, g_cov):
 
 
 def _condition_on(m, P, y_cond_mean, y_cond_cov, u, y, g_ev, g_cov, num_iter, emission_dist):
-    """Condition a Gaussian potential on a new observation with arbitrary
+    r"""Condition a Gaussian potential on a new observation with arbitrary
        likelihood with given functions for conditional moments and make a
        Gaussian approximation.
        p(x_t | y_t, u_t, y_{1:t-1}, u_{1:t-1})
@@ -172,7 +171,7 @@ def _step(carry, _):
 
 
 def _statistical_linear_regression(mu, Sigma, m, S, C):
-    """Return moment-matching affine coefficients and approximation noise variance
+    r"""Return moment-matching affine coefficients and approximation noise variance
     given joint moments.
 
         g(x) \approx Ax + b + e where e ~ N(0, Omega)

dynamax/generalized_gaussian_ssm/models.py

@@ -11,11 +11,11 @@
 from dynamax.nonlinear_gaussian_ssm.models import FnStateToState, FnStateAndInputToState
 from dynamax.nonlinear_gaussian_ssm.models import FnStateToEmission, FnStateAndInputToEmission
 
-FnStateToEmission2 = Callable[[Float[Array, "state_dim"]], Float[Array, "emission_dim emission_dim"]]
-FnStateAndInputToEmission2 = Callable[[Float[Array, "state_dim"], Float[Array, "input_dim"]], Float[Array, "emission_dim emission_dim"]]
+FnStateToEmission2 = Callable[[Float[Array, " state_dim"]], Float[Array, "emission_dim emission_dim"]]
+FnStateAndInputToEmission2 = Callable[[Float[Array, " state_dim"], Float[Array, " input_dim"]], Float[Array, "emission_dim emission_dim"]]
 
 # emission distribution takes a mean vector and covariance matrix and returns a distribution
-EmissionDistFn = Callable[ [Float[Array, "state_dim"], Float[Array, "state_dim state_dim"]], tfd.Distribution]
+EmissionDistFn = Callable[ [Float[Array, " state_dim"], Float[Array, "state_dim state_dim"]], tfd.Distribution]
 
 
 class ParamsGGSSM(NamedTuple):
@@ -42,7 +42,7 @@ class ParamsGGSSM(NamedTuple):
 
     """
 
-    initial_mean: Float[Array, "state_dim"]
+    initial_mean: Float[Array, " state_dim"]
     initial_covariance: Float[Array, "state_dim state_dim"]
     dynamics_function: Union[FnStateToState, FnStateAndInputToState]
     dynamics_covariance: Float[Array, "state_dim state_dim"]
@@ -97,15 +97,15 @@ def covariates_shape(self):
     def initial_distribution(
         self,
         params: ParamsGGSSM,
-        inputs: Optional[Float[Array, "input_dim"]]=None
+        inputs: Optional[Float[Array, " input_dim"]]=None
     ) -> tfd.Distribution:
         return MVN(params.initial_mean, params.initial_covariance)
 
     def transition_distribution(
         self,
         params: ParamsGGSSM,
-        state: Float[Array, "state_dim"],
-        inputs: Optional[Float[Array, "input_dim"]]=None
+        state: Float[Array, " state_dim"],
+        inputs: Optional[Float[Array, " input_dim"]]=None
     ) -> tfd.Distribution:
         f = params.dynamics_function
         if inputs is None:
@@ -117,8 +117,8 @@ def transition_distribution(
     def emission_distribution(
         self,
         params: ParamsGGSSM,
-        state: Float[Array, "state_dim"],
-        inputs: Optional[Float[Array, "input_dim"]]=None
+        state: Float[Array, " state_dim"],
+        inputs: Optional[Float[Array, " input_dim"]]=None
     ) -> tfd.Distribution:
         h = params.emission_mean_function
         R = params.emission_cov_function
@@ -128,4 +128,4 @@ def emission_distribution(
         else:
             mean = h(state, inputs)
             cov = R(state, inputs)
-        return params.emission_dist(mean, cov)
+        return params.emission_dist(mean, cov)

dynamax/hidden_markov_model/demos/categorical_glm_hmm_demo.py

@@ -58,12 +58,12 @@
     plt.figure()
     plt.imshow(jnp.vstack((states[None, :], most_likely_states[None, :])),
                 aspect="auto", interpolation='none', cmap="Greys")
-    plt.yticks([0.0, 1.0], ["$z$", "$\hat{z}$"])
+    plt.yticks([0.0, 1.0], ["$z$", r"$\hat{z}$"])
     plt.xlabel("time")
     plt.xlim(0, 500)
 
 
     print("true log prob: ", hmm.marginal_log_prob(true_params, emissions, inputs=inputs))
     print("test log prob: ", test_hmm.marginal_log_prob(params, emissions, inputs=inputs))
 
-    plt.show()
+    plt.show()

dynamax/hidden_markov_model/inference.py

@@ -1,25 +1,29 @@
+from functools import partial
+from typing import Callable, NamedTuple, Optional, Tuple, Union
 import jax.numpy as jnp
 import jax.random as jr
-from jax import lax
-from jax import vmap
-from jax import jit
-from functools import partial
-
-from typing import Callable, Optional, Tuple, Union, NamedTuple
+from jax import jit, lax, vmap
 from jaxtyping import Int, Float, Array
 
-from dynamax.types import Scalar, PRNGKey
+from dynamax.types import IntScalar, Scalar
 
 _get_params = lambda x, dim, t: x[t] if x.ndim == dim + 1 else x
 
-def get_trans_mat(transition_matrix, transition_fn, t):
+def get_trans_mat(
+        transition_matrix: Optional[Union[Float[Array, "num_states num_states"],
+                                          Float[Array, "num_timesteps_minus_1 num_states num_states"]]],
+        transition_fn: Optional[Callable[[IntScalar], Float[Array, "num_states num_states"]]],
+        t: IntScalar
+) -> Float[Array, "num_states num_states"]:
     if transition_fn is not None:
         return transition_fn(t)
-    else:
-        if transition_matrix.ndim == 3: # (T,K,K)
+    elif transition_matrix is not None:
+        if transition_matrix.ndim == 3: # (T-1,K,K)
             return transition_matrix[t]
         else:
             return transition_matrix
+    else:
+        raise ValueError("Either `transition_matrix` or `transition_fn` must be specified.")
 
 class HMMPosteriorFiltered(NamedTuple):
     r"""Simple wrapper for properties of an HMM filtering posterior.
@@ -50,12 +54,12 @@ class HMMPosterior(NamedTuple):
     filtered_probs: Float[Array, "num_timesteps num_states"]
     predicted_probs: Float[Array, "num_timesteps num_states"]
     smoothed_probs: Float[Array, "num_timesteps num_states"]
-    initial_probs: Float[Array, "num_states"]
-    trans_probs: Optional[Union[Float[Array, "num_timesteps num_states num_states"],
-                                Float[Array, "num_states num_states"]]] = None
+    initial_probs: Float[Array, " num_states"]
+    trans_probs: Optional[Union[Float[Array, "num_states num_states"],
+                                Float[Array, "num_timesteps_minus_1 num_states num_states"]]] = None
 
 
-def _normalize(u, axis=0, eps=1e-15):
+def _normalize(u: Array, axis=0, eps=1e-15):
     """Normalizes the values within the axis in a way that they sum up to 1.
 
     Args:
@@ -97,11 +101,11 @@ def _predict(probs, A):
 
 @partial(jit, static_argnames=["transition_fn"])
 def hmm_filter(
-    initial_distribution: Float[Array, "num_states"],
-    transition_matrix: Union[Float[Array, "num_timesteps num_states num_states"],
-                             Float[Array, "num_states num_states"]],
+    initial_distribution: Float[Array, " num_states"],
+    transition_matrix: Optional[Union[Float[Array, "num_states num_states"],
+                                      Float[Array, "num_timesteps_minus_1 num_states num_states"]]],
     log_likelihoods: Float[Array, "num_timesteps num_states"],
-    transition_fn: Optional[Callable[[Int], Float[Array, "num_states num_states"]]] = None
+    transition_fn: Optional[Callable[[IntScalar], Float[Array, "num_states num_states"]]] = None
 ) -> HMMPosteriorFiltered:
     r"""Forwards filtering
 
@@ -145,11 +149,11 @@ def _step(carry, t):
 
 @partial(jit, static_argnames=["transition_fn"])
 def hmm_backward_filter(
-    transition_matrix: Union[Float[Array, "num_timesteps num_states num_states"],
-                             Float[Array, "num_states num_states"]],
+    transition_matrix: Optional[Union[Float[Array, "num_states num_states"],
+                                      Float[Array, "num_timesteps_minus_1 num_states num_states"]]],
     log_likelihoods: Float[Array, "num_timesteps num_states"],
-    transition_fn: Optional[Callable[[Int], Float[Array, "num_states num_states"]]]= None
-) -> Tuple[Float, Float[Array, "num_timesteps num_states"]]:
+    transition_fn: Optional[Callable[[int], Float[Array, "num_states num_states"]]]= None
+) -> Tuple[Scalar, Float[Array, "num_timesteps num_states"]]:
     r"""Run the filter backwards in time. This is the second step of the forward-backward algorithm.
 
     Transition matrix may be either 2D (if transition probabilities are fixed) or 3D
@@ -191,11 +195,11 @@ def _step(carry, t):
 
 @partial(jit, static_argnames=["transition_fn"])
 def hmm_two_filter_smoother(
-    initial_distribution: Float[Array, "num_states"],
-    transition_matrix: Union[Float[Array, "num_timesteps num_states num_states"],
-                             Float[Array, "num_states num_states"]],
+    initial_distribution: Float[Array, " num_states"],
+    transition_matrix: Optional[Union[Float[Array, "num_states num_states"],
+                                      Float[Array, "num_timesteps_minus_1 num_states num_states"]]],
     log_likelihoods: Float[Array, "num_timesteps num_states"],
-    transition_fn: Optional[Callable[[Int], Float[Array, "num_states num_states"]]]= None,
+    transition_fn: Optional[Callable[[IntScalar], Float[Array, "num_states num_states"]]]= None,
     compute_trans_probs: bool = True
 ) -> HMMPosterior:
     r"""Computed the smoothed state probabilities using the two-filter
@@ -245,11 +249,11 @@ def hmm_two_filter_smoother(
 
 @partial(jit, static_argnames=["transition_fn"])
 def hmm_smoother(
-    initial_distribution: Float[Array, "num_states"],
-    transition_matrix: Union[Float[Array, "num_timesteps num_states num_states"],
-                             Float[Array, "num_states num_states"]],
+    initial_distribution: Float[Array, " num_states"],
+    transition_matrix: Optional[Union[Float[Array, "num_states num_states"],
+                                      Float[Array, "num_timesteps_minus_1 num_states num_states"]]],
     log_likelihoods: Float[Array, "num_timesteps num_states"],
-    transition_fn: Optional[Callable[[Int], Float[Array, "num_states num_states"]]]= None,
+    transition_fn: Optional[Callable[[IntScalar], Float[Array, "num_states num_states"]]]= None,
     compute_trans_probs: bool = True
 ) -> HMMPosterior:
     r"""Computed the smoothed state probabilities using a general
@@ -325,12 +329,12 @@ def _step(carry, args):
 
 @partial(jit, static_argnames=["transition_fn", "window_size"])
 def hmm_fixed_lag_smoother(
-    initial_distribution: Float[Array, "num_states"],
-    transition_matrix: Union[Float[Array, "num_timesteps num_states num_states"],
-                             Float[Array, "num_states num_states"]],
+    initial_distribution: Float[Array, " num_states"],
+    transition_matrix: Optional[Union[Float[Array, "num_states num_states"],
+                                      Float[Array, "num_timesteps_minus_1 num_states num_states"]]],
     log_likelihoods: Float[Array, "num_timesteps num_states"],
-    window_size: Int,
-    transition_fn: Optional[Callable[[Int], Float[Array, "num_states num_states"]]]= None
+    window_size: int,
+    transition_fn: Optional[Callable[[IntScalar], Float[Array, "num_states num_states"]]]= None
 ) -> HMMPosterior:
     r"""Compute the smoothed state probabilities using the fixed-lag smoother.
 
@@ -439,12 +443,12 @@ def compute_posterior(filtered_probs, beta):
 
 @partial(jit, static_argnames=["transition_fn"])
 def hmm_posterior_mode(
-    initial_distribution: Float[Array, "num_states"],
-    transition_matrix: Union[Float[Array, "num_timesteps num_states num_states"],
-                             Float[Array, "num_states num_states"]],
+    initial_distribution: Float[Array, " num_states"],
+    transition_matrix: Optional[Union[Float[Array, "num_states num_states"],
+                                     Float[Array, "num_timesteps_minus_1 num_states num_states"]]],
     log_likelihoods: Float[Array, "num_timesteps num_states"],
-    transition_fn: Optional[Callable[[Int], Float[Array, "num_states num_states"]]]= None
-) -> Int[Array, "num_timesteps"]:
+    transition_fn: Optional[Callable[[IntScalar], Float[Array, "num_states num_states"]]]= None
+) -> Int[Array, " num_timesteps"]:
     r"""Compute the most likely state sequence. This is called the Viterbi algorithm.
 
     Args:
@@ -486,13 +490,13 @@ def _forward_pass(state, best_next_state):
 
 @partial(jit, static_argnames=["transition_fn"])
 def hmm_posterior_sample(
-    rng: jr.PRNGKey,
-    initial_distribution: Float[Array, "num_states"],
-    transition_matrix: Union[Float[Array, "num_timesteps num_states num_states"],
-                             Float[Array, "num_states num_states"]],
+    key: Array,
+    initial_distribution: Float[Array, " num_states"],
+    transition_matrix: Optional[Union[Float[Array, "num_states num_states"],
+                                      Float[Array, "num_timesteps_minus_1 num_states num_states"]]],
     log_likelihoods: Float[Array, "num_timesteps num_states"],
-    transition_fn: Optional[Callable[[Int], Float[Array, "num_states num_states"]]] = None
-) -> Int[Array, "num_timesteps"]:
+    transition_fn: Optional[Callable[[IntScalar], Float[Array, "num_states num_states"]]] = None
+) -> Tuple[Scalar, Int[Array, " num_timesteps"]]:
     r"""Sample a latent sequence from the posterior.
 
     Args:
@@ -515,7 +519,7 @@ def hmm_posterior_sample(
     # Run the sampler backward in time
     def _step(carry, args):
         next_state = carry
-        t, rng, filtered_probs = args
+        t, subkey, filtered_probs = args
 
         A = get_trans_mat(transition_matrix, transition_fn, t)
 
@@ -524,15 +528,15 @@ def _step(carry, args):
         smoothed_probs /= smoothed_probs.sum()
 
         # Sample current state
-        state = jr.choice(rng, a=num_states, p=smoothed_probs)
+        state = jr.choice(subkey, a=num_states, p=smoothed_probs)
 
         return state, state
 
     # Run the HMM smoother
-    rngs = jr.split(rng, num_timesteps)
-    last_state = jr.choice(rngs[-1], a=num_states, p=filtered_probs[-1])
+    keys = jr.split(key, num_timesteps)
+    last_state = jr.choice(keys[-1], a=num_states, p=filtered_probs[-1])
     _, states = lax.scan(
-        _step, last_state, (jnp.arange(1, num_timesteps), rngs[:-1], filtered_probs[:-1]),
+        _step, last_state, (jnp.arange(1, num_timesteps), keys[:-1], filtered_probs[:-1]),
         reverse=True
     )
 
@@ -544,12 +548,13 @@ def _compute_sum_transition_probs(
     transition_matrix: Float[Array, "num_states num_states"],
     hmm_posterior: HMMPosterior) -> Float[Array, "num_states num_states"]:
     """Compute the transition probabilities from the HMM posterior messages.
+
     Args:
         transition_matrix (_type_): _description_
         hmm_posterior (_type_): _description_
     """
 
-    def _step(carry, args):
+    def _step(carry, args: Tuple[Array, Array, Array, Int[Array, ""]]):
         filtered_probs, smoothed_probs_next, predicted_probs_next, t = args
 
         # Get parameters for time t
@@ -580,11 +585,13 @@ def _step(carry, args):
 
 
 def _compute_all_transition_probs(
-    transition_matrix: Float[Array, "num_timesteps num_states num_states"],
+    transition_matrix: Optional[Union[Float[Array, "num_states num_states"],
+                                      Float[Array, "num_timesteps_minus_1 num_states num_states"]]],
     hmm_posterior: HMMPosterior,
-    transition_fn: Optional[Callable[[Int], Float[Array, "num_states num_states"]]] = None
+    transition_fn: Optional[Callable[[IntScalar], Float[Array, "num_states num_states"]]] = None
     ) -> Float[Array, "num_timesteps num_states num_states"]:
     """Compute the transition probabilities from the HMM posterior messages.
+
     Args:
         transition_matrix (_type_): _description_
         hmm_posterior (_type_): _description_
@@ -598,20 +605,21 @@ def _compute_probs(t):
         A = get_trans_mat(transition_matrix, transition_fn, t)
         return jnp.einsum('i,ij,j->ij', filtered_probs[t], A, relative_probs_next[t])
 
-    transition_probs = vmap(_compute_probs)(jnp.arange(len(filtered_probs)-1))
+    transition_probs = vmap(_compute_probs)(jnp.arange(len(filtered_probs)))
     return transition_probs
 
 
-# TODO: Consider alternative annotation for return type:
-#  Float[Array, "*num_timesteps num_states num_states"] I think this would allow multiple prepended dims.
-#  Float[Array, "#num_timesteps num_states num_states"] this might accept (1, sd, sd) but not (sd, sd).
+# TODO: This is a candidate for @overload however at present I think we would need to use 
+#       `@beartype.typing.overload` and beartype is currently not a core dependency.
+#       Support for `typing.overload` might change in the future: 
+#         https://github.com/beartype/beartype/issues/54
 def compute_transition_probs(
-    transition_matrix: Union[Float[Array, "num_timesteps num_states num_states"],
-                             Float[Array, "num_states num_states"]],
+    transition_matrix: Optional[Union[Float[Array, "num_states num_states"],
+                                      Float[Array, "num_timesteps_minus_1 num_states num_states"]]],
     hmm_posterior: HMMPosterior,
-    transition_fn: Optional[Callable[[Int], Float[Array, "num_states num_states"]]] = None
-) -> Union[Float[Array, "num_timesteps num_states num_states"],
-            Float[Array, "num_states num_states"]]:
+    transition_fn: Optional[Callable[[IntScalar], Float[Array, "num_states num_states"]]] = None
+) -> Union[Float[Array, "num_states num_states"],
+           Float[Array, "num_timesteps_minus_1 num_states num_states"]]:
     r"""Compute the posterior marginal distributions $p(z_{t+1}, z_t \mid y_{1:T}, u_{1:T}, \theta)$.
 
     Args:
@@ -622,8 +630,10 @@ def compute_transition_probs(
     Returns:
         array of smoothed transition probabilities.
     """
-    reduce_sum = transition_matrix is not None and transition_matrix.ndim == 2
-    if reduce_sum:
+    if transition_matrix is None and transition_fn is None:
+        raise ValueError("Either `transition_matrix` or `transition_fn` must be specified.")
+
+    if transition_matrix is not None and transition_matrix.ndim == 2:
         return _compute_sum_transition_probs(transition_matrix, hmm_posterior)
     else:
         return _compute_all_transition_probs(transition_matrix, hmm_posterior, transition_fn=transition_fn)

dynamax/hidden_markov_model/inference_test.py

@@ -1,4 +1,3 @@
-import pytest
 import itertools as it
 import jax.numpy as jnp
 import jax.random as jr