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/generalized_gaussian_ssm/models_test.py

@@ -21,7 +21,7 @@ def test_poisson_emission(key, kwargs):
     keys = jr.split(key, 3)
     state_dim = kwargs['state_dim']
     emission_dim = 1 # Univariate Poisson
-    poisson_weights = jr.normal(keys[0], shape=(emission_dim, state_dim))
+    poisson_weights = jr.normal(keys[0], shape=(emission_dim, state_dim)) /  jnp.sqrt(state_dim)
     model = GeneralizedGaussianSSM(state_dim, emission_dim)
 
     # Define model parameters with Poisson emission
@@ -51,7 +51,7 @@ def test_poisson_emission(key, kwargs):
 
     # Fit model with Gaussian emission
     gaussian_marginal_lls = conditional_moments_gaussian_filter(gaussian_params, EKFIntegrals(), emissions).marginal_loglik
-
+    
     # Check that the marginal log-likelihoods under Poisson emission are higher
     assert pois_marginal_lls > gaussian_marginal_lls
 

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()