Source code for dynamax.hidden_markov_model.models.gamma_hmm
if method.lower() == "kmeans":
assert emissions is not None, "Need emissions to initialize the model with K-Means!"
from sklearn.cluster import KMeans
- km = KMeans(self.num_states).fit(emissions.reshape(-1, 1))
+ key, subkey = jr.split(key) # Create a random seed for SKLearn.
+ sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647) # Max int32 value.
+ km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, 1))
_emission_concentrations = jnp.ones((self.num_states,))
_emission_rates = jnp.ravel(1.0 / km.cluster_centers_)
diff --git a/_modules/dynamax/hidden_markov_model/models/gaussian_hmm.html b/_modules/dynamax/hidden_markov_model/models/gaussian_hmm.html
index 44687ce8..3b2950da 100644
--- a/_modules/dynamax/hidden_markov_model/models/gaussian_hmm.html
+++ b/_modules/dynamax/hidden_markov_model/models/gaussian_hmm.html
@@ -353,7 +353,9 @@
diff --git a/notebooks/hmm/autoregressive_hmm.html b/notebooks/hmm/autoregressive_hmm.html
index 8cf315c7..bc582b6a 100644
--- a/notebooks/hmm/autoregressive_hmm.html
+++ b/notebooks/hmm/autoregressive_hmm.html
@@ -464,10 +464,10 @@ Source code for dynamax.hidden_markov_model.models.gaussian_hmm
if method.lower() == "kmeans":
assert emissions is not None, "Need emissions to initialize the model with K-Means!"
from sklearn.cluster import KMeans
- km = KMeans(self.num_states).fit(emissions.reshape(-1, self.emission_dim))
+ key, subkey = jr.split(key) # Create a random seed for SKLearn.
+ sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647) # Max int32 value.
+ km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, self.emission_dim))
_emission_means = jnp.array(km.cluster_centers_)
_emission_covs = jnp.tile(jnp.eye(self.emission_dim)[None, :, :], (self.num_states, 1, 1))
@@ -449,7 +451,9 @@
Source code for dynamax.hidden_markov_model.models.gaussian_hmm
if method.lower() == "kmeans":
assert emissions is not None, "Need emissions to initialize the model with K-Means!"
from sklearn.cluster import KMeans
- km = KMeans(self.num_states).fit(emissions.reshape(-1, self.emission_dim))
+ key, subkey = jr.split(key) # Create a random seed for SKLearn.
+ sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647) # Max int32 value.
+ km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, self.emission_dim))
_emission_means = jnp.array(km.cluster_centers_)
_emission_scale_diags = jnp.ones((self.num_states, self.emission_dim))
@@ -568,7 +572,9 @@ CMGF-EKF Inference
-
Source code for dynamax.hidden_markov_model.models.gaussian_hmm
if method.lower() == "kmeans":
assert emissions is not None, "Need emissions to initialize the model with K-Means!"
from sklearn.cluster import KMeans
- km = KMeans(self.num_states).fit(emissions.reshape(-1, self.emission_dim))
+ key, subkey = jr.split(key) # Create a random seed for SKLearn.
+ sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647) # Max int32 value.
+ km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, self.emission_dim))
_emission_means = jnp.array(km.cluster_centers_)
_emission_scales = jnp.ones((self.num_states,))
@@ -668,7 +674,9 @@
@@ -653,7 +653,7 @@ Source code for dynamax.hidden_markov_model.models.gaussian_hmm
if method.lower() == "kmeans":
assert emissions is not None, "Need emissions to initialize the model with K-Means!"
from sklearn.cluster import KMeans
- km = KMeans(self.num_states).fit(emissions.reshape(-1, self.emission_dim))
+ key, subkey = jr.split(key) # Create a random seed for SKLearn.
+ sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647) # Max int32 value.
+ km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, self.emission_dim))
_emission_means = jnp.array(km.cluster_centers_)
_emission_cov = jnp.eye(self.emission_dim)
@@ -788,7 +796,9 @@
Source code for dynamax.hidden_markov_model.models.gaussian_hmm
if method.lower() == "kmeans":
assert emissions is not None, "Need emissions to initialize the model with K-Means!"
from sklearn.cluster import KMeans
- km = KMeans(self.num_states).fit(emissions.reshape(-1, self.emission_dim))
+ key, subkey = jr.split(key) # Create a random seed for SKLearn.
+ sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647) # Max int32 value.
+ km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, self.emission_dim))
_emission_means = jnp.array(km.cluster_centers_)
_emission_cov_diag_factors = jnp.ones((self.num_states, self.emission_dim))
_emission_cov_low_rank_factors = jnp.zeros((self.num_states, self.emission_dim, self.emission_rank))
diff --git a/_modules/dynamax/hidden_markov_model/models/gmm_hmm.html b/_modules/dynamax/hidden_markov_model/models/gmm_hmm.html
index 5ddd6bfb..0f65c9cf 100644
--- a/_modules/dynamax/hidden_markov_model/models/gmm_hmm.html
+++ b/_modules/dynamax/hidden_markov_model/models/gmm_hmm.html
@@ -360,7 +360,9 @@
Source code for dynamax.hidden_markov_model.models.gmm_hmm
if method.lower() == "kmeans":
assert emissions is not None, "Need emissions to initialize the model with K-Means!"
from sklearn.cluster import KMeans
- km = KMeans(self.num_states).fit(emissions.reshape(-1, self.emission_dim))
+ key, subkey = jr.split(key) # Create a random seed for SKLearn.
+ sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647) # Max int32 value.
+ km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, self.emission_dim))
_emission_weights = jnp.ones((self.num_states, self.num_components)) / self.num_components
_emission_means = jnp.tile(jnp.array(km.cluster_centers_)[:, None, :], (1, self.num_components, 1))
_emission_covs = jnp.tile(jnp.eye(self.emission_dim), (self.num_states, self.num_components, 1, 1))
@@ -580,7 +582,9 @@
Source code for dynamax.hidden_markov_model.models.gmm_hmm
if method.lower() == "kmeans":
assert emissions is not None, "Need emissions to initialize the model with K-Means!"
from sklearn.cluster import KMeans
- km = KMeans(self.num_states).fit(emissions.reshape(-1, self.emission_dim))
+ key, subkey = jr.split(key) # Create a random seed for SKLearn.
+ sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647) # Max int32 value.
+ km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, self.emission_dim))
_emission_weights = jnp.ones((self.num_states, self.num_components)) / self.num_components
_emission_means = jnp.tile(jnp.array(km.cluster_centers_)[:, None, :], (1, self.num_components, 1))
_emission_scale_diags = jnp.ones((self.num_states, self.num_components, self.emission_dim))
diff --git a/_modules/dynamax/hidden_markov_model/models/linreg_hmm.html b/_modules/dynamax/hidden_markov_model/models/linreg_hmm.html
index 450471c4..cd519942 100644
--- a/_modules/dynamax/hidden_markov_model/models/linreg_hmm.html
+++ b/_modules/dynamax/hidden_markov_model/models/linreg_hmm.html
@@ -341,7 +341,9 @@ Online Training Using CMGF-EKF
-
Source code for dynamax.hidden_markov_model.models.linreg_hmm
if method.lower() == "kmeans":
assert emissions is not None, "Need emissions to initialize the model with K-Means!"
from sklearn.cluster import KMeans
- km = KMeans(self.num_states).fit(emissions.reshape(-1, self.emission_dim))
+ key, subkey = jr.split(key) # Create a random seed for SKLearn.
+ sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647) # Max int32 value.
+ km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(emissions.reshape(-1, self.emission_dim))
_emission_weights = jnp.zeros((self.num_states, self.emission_dim, self.input_dim))
_emission_biases = jnp.array(km.cluster_centers_)
_emission_covs = jnp.tile(jnp.eye(self.emission_dim)[None, :, :], (self.num_states, 1, 1))
diff --git a/_modules/dynamax/hidden_markov_model/models/logreg_hmm.html b/_modules/dynamax/hidden_markov_model/models/logreg_hmm.html
index bc582d26..e8bf51f4 100644
--- a/_modules/dynamax/hidden_markov_model/models/logreg_hmm.html
+++ b/_modules/dynamax/hidden_markov_model/models/logreg_hmm.html
@@ -334,7 +334,9 @@
Source code for dynamax.hidden_markov_model.models.logreg_hmm
flat_emissions = emissions.reshape(-1,)
flat_inputs = inputs.reshape(-1, self.input_dim)
- km = KMeans(self.num_states).fit(flat_inputs)
+ key, subkey = jr.split(key) # Create a random seed for SKLearn.
+ sklearn_key = jr.randint(subkey, shape=(), minval=0, maxval=2147483647) # Max int32 value.
+ km = KMeans(self.num_states, random_state=int(sklearn_key)).fit(flat_inputs)
_emission_weights = jnp.zeros((self.num_states, self.input_dim))
_emission_biases = jnp.array([tfb.Sigmoid().inverse(flat_emissions[km.labels_ == k].mean())
for k in range(self.num_states)])
diff --git a/_modules/dynamax/linear_gaussian_ssm/inference.html b/_modules/dynamax/linear_gaussian_ssm/inference.html
index 4c0c7167..fda844fc 100644
--- a/_modules/dynamax/linear_gaussian_ssm/inference.html
+++ b/_modules/dynamax/linear_gaussian_ssm/inference.html
@@ -283,9 +283,13 @@ Online Training Using CMGF-EKF
-
Source code for dynamax.linear_gaussian_ssm.inference
import jax.numpy as jnp
import jax.random as jr
from jax import lax
-from tensorflow_probability.substrates.jax.distributions import MultivariateNormalFullCovariance as MVN
from functools import wraps
import inspect
+import warnings
+
+from tensorflow_probability.substrates.jax.distributions import (
+ MultivariateNormalDiagPlusLowRankCovariance as MVNLowRank,
+ MultivariateNormalFullCovariance as MVN)
from jax.tree_util import tree_map
from jaxtyping import Array, Float
@@ -323,10 +327,22 @@
Source code for dynamax.linear_gaussian_ssm.inference
:param cov: dynamics covariance $Q$
"""
- weights: Union[Float[Array, "state_dim state_dim"], Float[Array, "ntime state_dim state_dim"], ParameterProperties]
- bias: Union[Float[Array, "state_dim"], Float[Array, "ntime state_dim"], ParameterProperties]
- input_weights: Union[Float[Array, "state_dim input_dim"], Float[Array, "ntime state_dim input_dim"], ParameterProperties]
- cov: Union[Float[Array, "state_dim state_dim"], Float[Array, "ntime state_dim state_dim"], Float[Array, "state_dim_triu"], ParameterProperties]
+ weights: Union[ParameterProperties,
+ Float[Array, "state_dim state_dim"],
+ Float[Array, "ntime state_dim state_dim"]]
+
+ bias: Union[ParameterProperties,
+ Float[Array, "state_dim"],
+ Float[Array, "ntime state_dim"]]
+
+ input_weights: Union[ParameterProperties,
+ Float[Array, "state_dim input_dim"],
+ Float[Array, "ntime state_dim input_dim"]]
+
+ cov: Union[ParameterProperties,
+ Float[Array, "state_dim state_dim"],
+ Float[Array, "ntime state_dim state_dim"],
+ Float[Array, "state_dim_triu"]][docs]class ParamsLGSSMEmissions(NamedTuple):
@@ -342,11 +358,24 @@
Source code for dynamax.linear_gaussian_ssm.inference
:param cov: emission covariance $R$
"""
- weights: Union[Float[Array, "emission_dim state_dim"], Float[Array, "ntime emission_dim state_dim"], ParameterProperties]
- bias: Union[Float[Array, "emission_dim"], Float[Array, "ntime emission_dim"], ParameterProperties]
- input_weights: Union[Float[Array, "emission_dim input_dim"], Float[Array, "ntime emission_dim input_dim"], ParameterProperties]
- cov: Union[Float[Array, "emission_dim emission_dim"], Float[Array, "ntime emission_dim emission_dim"], Float[Array, "emission_dim_triu"], ParameterProperties]
-
+ weights: Union[ParameterProperties,
+ Float[Array, "emission_dim state_dim"],
+ Float[Array, "ntime emission_dim state_dim"]]
+
+ bias: Union[ParameterProperties,
+ Float[Array, "emission_dim"],
+ Float[Array, "ntime emission_dim"]]
+
+ input_weights: Union[ParameterProperties,
+ Float[Array, "emission_dim input_dim"],
+ Float[Array, "ntime emission_dim input_dim"]]
+
+ cov: Union[ParameterProperties,
+ Float[Array, "emission_dim emission_dim"],
+ Float[Array, "ntime emission_dim emission_dim"],
+ Float[Array, "emission_dim"],
+ Float[Array, "ntime emission_dim"],
+ Float[Array, "emission_dim_triu"]][docs]class ParamsLGSSM(NamedTuple):
@@ -397,14 +426,46 @@
@@ -622,7 +622,7 @@ Source code for dynamax.linear_gaussian_ssm.inference
# Helper functions
-# _get_params = lambda x, dim, t: x[t] if x.ndim == dim + 1 else x
-def _get_params(x, dim, t):
+
+def _get_one_param(x, dim, t):
+ """Helper function to get one parameter at time t."""
if callable(x):
return x(t)
elif x.ndim == dim + 1:
return x[t]
else:
return x
+
+def _get_params(params, num_timesteps, t):
+ """Helper function to get parameters at time t."""
+ assert not callable(params.emissions.cov), "Emission covariance cannot be a callable."
+
+ F = _get_one_param(params.dynamics.weights, 2, t)
+ B = _get_one_param(params.dynamics.input_weights, 2, t)
+ b = _get_one_param(params.dynamics.bias, 1, t)
+ Q = _get_one_param(params.dynamics.cov, 2, t)
+ H = _get_one_param(params.emissions.weights, 2, t)
+ D = _get_one_param(params.emissions.input_weights, 2, t)
+ d = _get_one_param(params.emissions.bias, 1, t)
+
+ if len(params.emissions.cov.shape) == 1:
+ R = _get_one_param(params.emissions.cov, 1, t)
+ elif len(params.emissions.cov.shape) > 2:
+ R = _get_one_param(params.emissions.cov, 2, t)
+ elif params.emissions.cov.shape[0] != num_timesteps:
+ R = _get_one_param(params.emissions.cov, 2, t)
+ elif params.emissions.cov.shape[1] != num_timesteps:
+ R = _get_one_param(params.emissions.cov, 1, t)
+ else:
+ R = _get_one_param(params.emissions.cov, 2, t)
+ warnings.warn(
+ "Emission covariance has shape (N,N) where N is the number of timesteps. "
+ "The covariance will be interpreted as static and non-diagonal. To "
+ "specify a dynamic and diagonal covariance, pass it as a 3D array.")
+
+ return F, B, b, Q, H, D, d, R
+
+
_zeros_if_none = lambda x, shape: x if x is not None else jnp.zeros(shape)
@@ -481,7 +542,6 @@ Create data
-
Source code for dynamax.linear_gaussian_ssm.inference
S = (R + H * P * H')
K = P * H' * S^{-1}
PP = P - K S K' = Sigma_cond
- **Note! This can be done more efficiently when R is diagonal.**
Args:
m (D_hid,): prior mean.
@@ -497,9 +557,25 @@
diff --git a/notebooks/generalized_gaussian_ssm/cmgf_mlp_classification_demo.html b/notebooks/generalized_gaussian_ssm/cmgf_mlp_classification_demo.html
index 8f7011b5..65c4bba9 100644
--- a/notebooks/generalized_gaussian_ssm/cmgf_mlp_classification_demo.html
+++ b/notebooks/generalized_gaussian_ssm/cmgf_mlp_classification_demo.html
@@ -492,7 +492,7 @@ Source code for dynamax.linear_gaussian_ssm.inference
mu_pred (D_hid,): predicted mean.
Sigma_pred (D_hid,D_hid): predicted covariance.
"""
- # Compute the Kalman gain
- S = R + H @ P @ H.T
- K = psd_solve(S, H @ P).T
+ if R.ndim == 2:
+ S = R + H @ P @ H.T
+ K = psd_solve(S, H @ P).T
+ else:
+ # Optimization using Woodbury identity with A=R, U=H@chol(P), V=U.T, C=I
+ # (see https://en.wikipedia.org/wiki/Woodbury_matrix_identity)
+ I = jnp.eye(P.shape[0])
+ U = H @ jnp.linalg.cholesky(P)
+ X = U / R[:, None]
+ S_inv = jnp.diag(1.0 / R) - X @ psd_solve(I + U.T @ X, X.T)
+ """
+ # Could alternatively use U=H and C=P
+ R_inv = jnp.diag(1.0 / R)
+ P_inv = psd_solve(P, jnp.eye(P.shape[0]))
+ S_inv = R_inv - R_inv @ H @ psd_solve(P_inv + H.T @ R_inv @ H, H.T @ R_inv)
+ """
+ K = P @ H.T @ S_inv
+ S = jnp.diag(R) + H @ P @ H.T
+
Sigma_cond = P - K @ S @ K.T
mu_cond = m + K @ (y - D @ u - d - H @ m)
return mu_cond, symmetrize(Sigma_cond)
@@ -567,6 +643,8 @@
Source code for dynamax.linear_gaussian_ssm.inference
return wrapper
+
+
def lgssm_joint_sample(
params: ParamsLGSSM,
key: PRNGKey,
@@ -584,7 +662,6 @@
@@ -756,10 +756,10 @@ Source code for dynamax.linear_gaussian_ssm.inference
latent states and emissions
"""
-
params, inputs = preprocess_params_and_inputs(params, num_timesteps, inputs)
def _sample_transition(key, F, B, b, Q, x_tm1, u):
@@ -593,6 +670,7 @@
Source code for dynamax.linear_gaussian_ssm.inference
def _sample_emission(key, H, D, d, R, x, u):
mean = H @ x + D @ u + d
+ R = jnp.diag(R) if R.ndim==1 else R
return MVN(mean, R).sample(seed=key)
def _sample_initial(key, params, inputs):
@@ -600,10 +678,7 @@
Source code for dynamax.linear_gaussian_ssm.inference
initial_state = MVN(params.initial.mean, params.initial.cov).sample(seed=key1)
- H0 = _get_params(params.emissions.weights, 2, 0)
- D0 = _get_params(params.emissions.input_weights, 2, 0)
- d0 = _get_params(params.emissions.bias, 1, 0)
- R0 = _get_params(params.emissions.cov, 2, 0)
+ H0, D0, d0, R0 = _get_params(params, num_timesteps, 0)[4:]
u0 = tree_map(lambda x: x[0], inputs)
initial_emission = _sample_emission(key2, H0, D0, d0, R0, initial_state, u0)
@@ -613,15 +688,8 @@
@@ -727,10 +727,10 @@ Source code for dynamax.linear_gaussian_ssm.inference
key, t, inpt = args
key1, key2 = jr.split(key, 2)
- # Shorthand: get parameters and inputs for time index t
- F = _get_params(params.dynamics.weights, 2, t)
- B = _get_params(params.dynamics.input_weights, 2, t)
- b = _get_params(params.dynamics.bias, 1, t)
- Q = _get_params(params.dynamics.cov, 2, t)
- H = _get_params(params.emissions.weights, 2, t)
- D = _get_params(params.emissions.input_weights, 2, t)
- d = _get_params(params.emissions.bias, 1, t)
- R = _get_params(params.emissions.cov, 2, t)
+ # Get parameters and inputs for time index t
+ F, B, b, Q, H, D, d, R = _get_params(params, num_timesteps, t)
# Sample from transition and emission distributions
state = _sample_transition(key1, F, B, b, Q, prev_state, inpt)
@@ -668,23 +736,26 @@
Source code for dynamax.linear_gaussian_ssm.inference
num_timesteps = len(emissions)
inputs = jnp.zeros((num_timesteps, 0)) if inputs is None else inputs
+ def _log_likelihood(pred_mean, pred_cov, H, D, d, R, u, y):
+ m = H @ pred_mean + D @ u + d
+ if R.ndim==2:
+ S = R + H @ pred_cov @ H.T
+ return MVN(m, S).log_prob(y)
+ else:
+ L = H @ jnp.linalg.cholesky(pred_cov)
+ return MVNLowRank(m, R, L).log_prob(y)
+
+
def _step(carry, t):
ll, pred_mean, pred_cov = carry
# Shorthand: get parameters and inputs for time index t
- F = _get_params(params.dynamics.weights, 2, t)
- B = _get_params(params.dynamics.input_weights, 2, t)
- b = _get_params(params.dynamics.bias, 1, t)
- Q = _get_params(params.dynamics.cov, 2, t)
- H = _get_params(params.emissions.weights, 2, t)
- D = _get_params(params.emissions.input_weights, 2, t)
- d = _get_params(params.emissions.bias, 1, t)
- R = _get_params(params.emissions.cov, 2, t)
+ F, B, b, Q, H, D, d, R = _get_params(params, num_timesteps, t)
u = inputs[t]
y = emissions[t]
# Update the log likelihood
- ll += MVN(H @ pred_mean + D @ u + d, H @ pred_cov @ H.T + R).log_prob(y)
+ ll += _log_likelihood(pred_mean, pred_cov, H, D, d, R, u, y)
# Condition on this emission
filtered_mean, filtered_cov = _condition_on(pred_mean, pred_cov, H, D, d, R, u, y)
@@ -732,11 +803,8 @@
Source code for dynamax.linear_gaussian_ssm.inference
smoothed_mean_next, smoothed_cov_next = carry
t, filtered_mean, filtered_cov = args
- # Shorthand: get parameters and inputs for time index t
- F = _get_params(params.dynamics.weights, 2, t)
- B = _get_params(params.dynamics.input_weights, 2, t)
- b = _get_params(params.dynamics.bias, 1, t)
- Q = _get_params(params.dynamics.cov, 2, t)
+ # Get parameters and inputs for time index t
+ F, B, b, Q = _get_params(params, num_timesteps, t)[:4]
u = inputs[t]
# This is like the Kalman gain but in reverse
@@ -804,10 +872,7 @@
@@ -699,10 +699,10 @@ Source code for dynamax.linear_gaussian_ssm.inference
key, filtered_mean, filtered_cov, t = args
# Shorthand: get parameters and inputs for time index t
- F = _get_params(params.dynamics.weights, 2, t)
- B = _get_params(params.dynamics.input_weights, 2, t)
- b = _get_params(params.dynamics.bias, 1, t)
- Q = _get_params(params.dynamics.cov, 2, t)
+ F, B, b, Q = _get_params(params, num_timesteps, t)[:4]
u = inputs[t]
# Condition on next state
diff --git a/_modules/dynamax/utils/utils.html b/_modules/dynamax/utils/utils.html
index f96ea7ef..ff3a69ec 100644
--- a/_modules/dynamax/utils/utils.html
+++ b/_modules/dynamax/utils/utils.html
@@ -291,6 +291,7 @@
Source code for dynamax.utils.utils
from jaxtyping import Array, Int
from scipy.optimize import linear_sum_assignment
from typing import Optional
+from jax.scipy.linalg import cho_factor, cho_solve
def has_tpu():
try:
@@ -480,10 +481,12 @@ Source code for dynamax.utils.utils
return perm
-def psd_solve(A,b):
+def psd_solve(A, b, diagonal_boost=1e-9):
"""A wrapper for coordinating the linalg solvers used in the library for psd matrices."""
- A = A + 1e-6
- return jnp.linalg.solve(A,b)
+ A = symmetrize(A) + diagonal_boost * jnp.eye(A.shape[-1])
+ L, lower = cho_factor(A, lower=True)
+ x = cho_solve((L, lower), b)
+ return x
def symmetrize(A):
"""Symmetrize one or more matrices."""
diff --git a/_sources/notebooks/linear_gaussian_ssm/kf_linreg.ipynb b/_sources/notebooks/linear_gaussian_ssm/kf_linreg.ipynb
index ea352631..ea705f19 100644
--- a/_sources/notebooks/linear_gaussian_ssm/kf_linreg.ipynb
+++ b/_sources/notebooks/linear_gaussian_ssm/kf_linreg.ipynb
@@ -10,7 +10,7 @@
"\n",
"We perform sequential (recursive) Bayesian inference for the parameters of a linear regression model\n",
"using the Kalman filter. (This algorithm is also known as recursive least squares.)\n",
- "To do this, we treat the parameers of the model as the unknown hidden states.\n",
+ "To do this, we treat the parameters of the model as the unknown hidden states.\n",
"We assume that these are constant over time.\n",
"The graphical model is shown below.\n",
"\n",
diff --git a/api.html b/api.html
index bd4cf4af..cea3bb0e 100644
--- a/api.html
+++ b/api.html
@@ -3281,10 +3281,10 @@ Types#
Parameters:
-weights (Union[Float[Array, 'state_dim state_dim'], Float[Array, 'ntime state_dim state_dim'], ParameterProperties]) – dynamics weights \(F\)
bias (Union[Float[Array, 'state_dim'], Float[Array, 'ntime state_dim'], ParameterProperties]) – dynamics bias \(b\)
input_weights (Union[Float[Array, 'state_dim input_dim'], Float[Array, 'ntime state_dim input_dim'], ParameterProperties]) – dynamics input weights \(B\)
cov (Union[Float[Array, 'state_dim state_dim'], Float[Array, 'ntime state_dim state_dim'], Float[Array, 'state_dim_triu'], ParameterProperties]) – dynamics covariance \(Q\)
weights (Union[ParameterProperties, Float[Array, 'state_dim state_dim'], Float[Array, 'ntime state_dim state_dim']]) – dynamics weights \(F\)
bias (Union[ParameterProperties, Float[Array, 'state_dim'], Float[Array, 'ntime state_dim']]) – dynamics bias \(b\)
input_weights (Union[ParameterProperties, Float[Array, 'state_dim input_dim'], Float[Array, 'ntime state_dim input_dim']]) – dynamics input weights \(B\)
cov (Union[ParameterProperties, Float[Array, 'state_dim state_dim'], Float[Array, 'ntime state_dim state_dim'], Float[Array, 'state_dim_triu']]) – dynamics covariance \(Q\)
Types#
Parameters:
-weights (Union[Float[Array, 'emission_dim state_dim'], Float[Array, 'ntime emission_dim state_dim'], ParameterProperties]) – emission weights \(H\)
bias (Union[Float[Array, 'emission_dim'], Float[Array, 'ntime emission_dim'], ParameterProperties]) – emission bias \(d\)
input_weights (Union[Float[Array, 'emission_dim input_dim'], Float[Array, 'ntime emission_dim input_dim'], ParameterProperties]) – emission input weights \(D\)
cov (Union[Float[Array, 'emission_dim emission_dim'], Float[Array, 'ntime emission_dim emission_dim'], Float[Array, 'emission_dim_triu'], ParameterProperties]) – emission covariance \(R\)
weights (Union[ParameterProperties, Float[Array, 'emission_dim state_dim'], Float[Array, 'ntime emission_dim state_dim']]) – emission weights \(H\)
bias (Union[ParameterProperties, Float[Array, 'emission_dim'], Float[Array, 'ntime emission_dim']]) – emission bias \(d\)
input_weights (Union[ParameterProperties, Float[Array, 'emission_dim input_dim'], Float[Array, 'ntime emission_dim input_dim']]) – emission input weights \(D\)
cov (Union[ParameterProperties, Float[Array, 'emission_dim emission_dim'], Float[Array, 'ntime emission_dim emission_dim'], Float[Array, 'emission_dim'], Float[Array, 'ntime emission_dim'], Float[Array, 'emission_dim_triu']]) – emission covariance \(R\)
Simulation and Plotting
-
Laplace Estimate
-
/opt/hostedtoolcache/Python/3.9.17/x64/lib/python3.9/site-packages/jax/_src/random.py:458: FutureWarning: jax.random.shuffle is deprecated and will be removed in a future release. Use jax.random.permutation with independent=True.
+/opt/hostedtoolcache/Python/3.9.17/x64/lib/python3.9/site-packages/jax/_src/random.py:480: FutureWarning: jax.random.shuffle is deprecated and will be removed in a future release. Use jax.random.permutation with independent=True.
warnings.warn(msg, FutureWarning)
@@ -502,7 +502,7 @@ Simulation and Plotting
-
+
+
Let us define a grid on which we compute the predictive distribution.
@@ -622,7 +622,7 @@Laplace Estimate
-
+
@@ -637,7 +637,7 @@ Laplace Estimate
-
Laplace Estimate
-
+
EKF#
-
-
-
-
+
+
+
+
-
-
-
+
+
+
+
UKF#
-
-
-
-
+
+
+
+
-
-
-
+
+
+
+
GHKF#
-
-
-
-
+
+
+
+
-
-
-
+
+
+
+
Create data
-
+
Online Training Using CMGF-EKF
-
+
Next, we visualize the training procedure by evaluating the intermediate steps.
@@ -639,7 +639,7 @@Online Training Using CMGF-EKF
-
+
Finally, we generate a video of the MLP-Classifier being trained.
diff --git a/notebooks/generalized_gaussian_ssm/cmgf_poisson_demo.html b/notebooks/generalized_gaussian_ssm/cmgf_poisson_demo.html index 3c923c5a..f6a98b17 100644 --- a/notebooks/generalized_gaussian_ssm/cmgf_poisson_demo.html +++ b/notebooks/generalized_gaussian_ssm/cmgf_poisson_demo.html @@ -631,7 +631,7 @@Model#
-
+
+
CMGF-EKF Inference
-
+
@@ -672,12 +672,12 @@ CMGF-EKF Inference
-
CMGF-EKF Inference
-
-
-
-
-
-
+
+
+
+
+
+
Manually construct an ARHMM
-<matplotlib.colorbar.Colorbar at 0x7f837ca09880>
+<matplotlib.colorbar.Colorbar at 0x7f5bf015efd0>
-
+
-
+
<matplotlib.colorbar.Colorbar at 0x7f837ca09880>
+<matplotlib.colorbar.Colorbar at 0x7f5bf015efd0>
-
+
+
+
+
@@ -620,8 +620,8 @@ 
+
+
+
@@ -767,7 +767,7 @@ 
+
diff --git a/notebooks/hmm/casino_hmm_inference.html b/notebooks/hmm/casino_hmm_inference.html
index b04f5061..9761f94e 100644
--- a/notebooks/hmm/casino_hmm_inference.html
+++ b/notebooks/hmm/casino_hmm_inference.html
@@ -413,8 +413,8 @@ 
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
-
-
+
+
+
+
+
+
+
+
+
+
@@ -476,7 +476,7 @@ 
+
@@ -514,7 +514,7 @@ 
-
+
+
@@ -554,8 +554,8 @@ 
-
+
+
diff --git a/notebooks/linear_gaussian_ssm/lgssm_parallel_inference.html b/notebooks/linear_gaussian_ssm/lgssm_parallel_inference.html
index 77f46855..bbb5f32f 100644
--- a/notebooks/linear_gaussian_ssm/lgssm_parallel_inference.html
+++ b/notebooks/linear_gaussian_ssm/lgssm_parallel_inference.html
@@ -496,7 +496,7 @@ 
+
+
+
-
-
-
+
+
+
+
+
+
+
+
+
+
+
+
+
+