Dynamax

https://probml.github.io/dynamax
Admin
Dynamax is a JAX library for probabilistic state space models (SSMs), providing code for inference...
Tokens:221,027
Snippets:1,323
Trust Score:9.7
Update:3 months ago
Show doc for...
Tokens:
Raw
### Install dynamax

Source: https://probml.github.io/dynamax/notebooks/generalized_gaussian_ssm/cmgf_poisson_demo.html

Installs the dynamax library with support for notebooks. This is a prerequisite for running the subsequent code examples.

```python
%%capture
try:
    import dynamax
except ModuleNotFoundError:
    print('installing dynamax')
    %pip install -q dynamax[notebooks]
    import dynamax

```

--------------------------------

### Install and Import Dynamax

Source: https://probml.github.io/dynamax/notebooks/generalized_gaussian_ssm/cmgf_logistic_regression_demo.html

Installs the dynamax library with notebook support if not already present, then imports necessary modules.

```python
%%capture
try:
    import dynamax
except ModuleNotFoundError:
    print('installing dynamax')
    %pip install -q dynamax[notebooks]
    import dynamax

```

```python
import matplotlib.pyplot as plt
import jax
import jax.numpy as jnp
import jax.random as jr

from jax.scipy.optimize import minimize

from dynamax.generalized_gaussian_ssm import ParamsGGSSM, EKFIntegrals, UKFIntegrals, GHKFIntegrals
from dynamax.generalized_gaussian_ssm import conditional_moments_gaussian_filter

```

--------------------------------

### Install DYNAMAX for Development

Source: https://probml.github.io/dynamax/index.html

Clone the DYNAMAX repository and install it in editable mode with development, testing, and documentation dependencies.

```bash
git clone git@github.com:probml/dynamax.git
cd dynamax
pip install -e '.[dev]'
```

--------------------------------

### Install Latest DYNAMAX Release

Source: https://probml.github.io/dynamax/index.html

Install the latest stable version of DYNAMAX and its core dependencies from PyPI. Use the `[notebooks]` extra to include dependencies for demo notebooks.

```bash
pip install dynamax                 # Install dynamax and core dependencies, or
pip install dynamax[notebooks]      # Install with demo notebook dependencies
```

--------------------------------

### Initialize GammaHMMEmissions with KMeans Method

Source: https://probml.github.io/dynamax/_modules/dynamax/hidden_markov_model/models/gamma_hmm.html

Initializes GammaHMMEmissions parameters using the K-Means clustering method on provided emissions. Requires scikit-learn to be installed.

```python
key = jr.PRNGKey(0)
# Assume 'emissions' is a JAX array of shape (num_timesteps,)
emissions_data = emissions.reshape(-1, 1)
params, props = GammaHMMEmissions(num_states).initialize(method="kmeans", emissions=emissions_data, key=key)
```

--------------------------------

### Install Latest DYNAMAX Development Branch

Source: https://probml.github.io/dynamax/index.html

Install the most recent development version of DYNAMAX directly from its GitHub repository.

```bash
pip install git+https://github.com/probml/dynamax.git
```

--------------------------------

### Initialize Logistic Regression HMM Emissions with K-Means

Source: https://probml.github.io/dynamax/_modules/dynamax/hidden_markov_model/models/logreg_hmm.html

Initializes emission parameters using K-Means clustering on provided emissions and inputs. Requires scikit-learn to be installed.

```python
from sklearn.cluster import KMeans

flat_emissions = emissions.reshape(-1,)
flat_inputs = inputs.reshape(-1, self.input_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(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)])
```

--------------------------------

### Import necessary libraries

Source: https://probml.github.io/dynamax/notebooks/linear_gaussian_ssm/lgssm_learning.html

Imports core libraries for LG-SSM modeling, including JAX for numerical computation, Matplotlib for plotting, and Optax for optimization. Ensure these are installed before running.

```python
import jax.numpy as jnp
import jax.random as jr

from jax import vmap
from matplotlib import pyplot as plt
from optax import adam

from dynamax.linear_gaussian_ssm import LinearGaussianSSM
from dynamax.utils.utils import monotonically_increasing, random_rotation

```

--------------------------------

### Get HMM Initial Distribution

Source: https://probml.github.io/dynamax/_modules/dynamax/hidden_markov_model/models/abstractions.html

Returns the initial distribution of the HMM, parameterized by the initial component. This is used to define the starting state probabilities.

```python
    def initial_distribution(self, params: HMMParameterSet, 
                             inputs:Optional[Array] = None) 
                             -> tfd.Distribution:
        """Return the initial distribution."""
        return self.initial_component.distribution(params.initial, inputs=inputs)
```

--------------------------------

### Instantiate MLP with Flattened Parameters

Source: https://probml.github.io/dynamax/notebooks/nonlinear_gaussian_ssm/ekf_mlp.html

Example of how to use the `get_mlp_flattened_params` function to create an MLP model and obtain its flattened parameters and apply function. This setup is typically done before training or inference.

```python
input_dim, hidden_dim, output_dim = 1, 6, 1
model_dims = [input_dim, hidden_dim, output_dim]
_, flat_params, _, apply_fn = get_mlp_flattened_params(model_dims)

```

--------------------------------

### Install Flax

Source: https://probml.github.io/dynamax/notebooks/generalized_gaussian_ssm/cmgf_mlp_classification_demo.html

Installs the Flax library if it's not already installed. Flax is a neural network library used for building the MLP classifier.

```python
try:
    import flax.linen as nn
except ModuleNotFoundError:
    print('installing flax')
    %pip install -qq flax
import flax.linen as nn

```

--------------------------------

### Initialize HMM Parameters (K-Means Method)

Source: https://probml.github.io/dynamax/api.html

Initializes HMM parameters using k-means clustering on the provided emissions. Requires a PRNGKey if parameters are not fully specified.

```python
initialize(_key =Array([0, 0], dtype=uint32)_, _method ='kmeans'_, _initial_probs =None_, _transition_matrix =None_, _emission_weights =None_, _emission_means =None_, _emission_scale_diags =None_, _emissions =None_)
```

--------------------------------

### Install blackjax Library

Source: https://probml.github.io/dynamax/notebooks/linear_gaussian_ssm/lgssm_hmc.html

Installs the blackjax library if it's not already installed. Blackjax is a library of numerical computation primitives for JAX, often used for MCMC methods.

```python
try:
    import blackjax
except ModuleNotFoundError:
    print('installing blackjax')
    %pip install -qq blackjax
    import blackjax

```

--------------------------------

### Initialize GammaHMMEmissions with Prior Method

Source: https://probml.github.io/dynamax/_modules/dynamax/hidden_markov_model/models/gamma_hmm.html

Initializes GammaHMMEmissions parameters using a prior method. Requires a PRNGKey if parameters are not manually specified.

```python
key = jr.PRNGKey(0)
emission_concentrations = jnp.ones((num_states,))
emission_rates = jr.exponential(key, (num_states,))
params, props = GammaHMMEmissions(num_states).initialize(method="prior", emission_concentrations=emission_concentrations, emission_rates=emission_rates)
```

--------------------------------

### Initialize LGSSM Parameters and Properties

Source: https://probml.github.io/dynamax/_modules/dynamax/linear_gaussian_ssm/models.html

Sets up default parameters and their properties for a Linear Gaussian State Space Model. Use this when creating a new LGSSM instance without providing custom initial parameters.

```python
        # Arbitrary default values, for demo purposes.
        _initial_mean = jnp.zeros(self.state_dim)
        _initial_covariance = jnp.eye(self.state_dim)
        _dynamics_weights = 0.99 * jnp.eye(self.state_dim)
        _dynamics_input_weights = jnp.zeros((self.state_dim, self.input_dim))
        _dynamics_bias = jnp.zeros((self.state_dim,)) if self.has_dynamics_bias else None
        _dynamics_covariance = 0.1 * jnp.eye(self.state_dim)
        _emission_weights = jr.normal(key, (self.emission_dim, self.state_dim))
        _emission_input_weights = jnp.zeros((self.emission_dim, self.input_dim))
        _emission_bias = jnp.zeros((self.emission_dim,)) if self.has_emissions_bias else None
        _emission_covariance = 0.1 * jnp.eye(self.emission_dim)

        # Only use the values above if the user hasn't specified their own
        default = lambda x, x0: x if x is not None else x0

        # Create nested dictionary of params
        params = ParamsLGSSM(
            initial=ParamsLGSSMInitial(
                mean=default(initial_mean, _initial_mean),
                cov=default(initial_covariance, _initial_covariance)),
            dynamics=ParamsLGSSMDynamics(
                weights=default(dynamics_weights, _dynamics_weights),
                bias=default(dynamics_bias, _dynamics_bias),
                input_weights=default(dynamics_input_weights, _dynamics_input_weights),
                cov=default(dynamics_covariance, _dynamics_covariance)),
            emissions=ParamsLGSSMEmissions(
                weights=default(emission_weights, _emission_weights),
                bias=default(emission_bias, _emission_bias),
                input_weights=default(emission_input_weights, _emission_input_weights),
                cov=default(emission_covariance, _emission_covariance))
            )

        # The keys of param_props must match those of params!
        props = ParamsLGSSM(
            initial=ParamsLGSSMInitial(
                mean=ParameterProperties(),
                cov=ParameterProperties(constrainer=RealToPSDBijector())),
            dynamics=ParamsLGSSMDynamics(
                weights=ParameterProperties(),
                bias=ParameterProperties(),
                input_weights=ParameterProperties(),
                cov=ParameterProperties(constrainer=RealToPSDBijector())),
            emissions=ParamsLGSSMEmissions(
                weights=ParameterProperties(),
                bias=ParameterProperties(),
                input_weights=ParameterProperties(),
                cov=ParameterProperties(constrainer=RealToPSDBijector()))
            )
        return params, props
```

--------------------------------

### HMMInitialState.initialize_m_step_state

Source: https://probml.github.io/dynamax/api.html

Initialize any required state for the M step.

```APIDOC
## HMMInitialState.initialize_m_step_state

### Description
Initialize any required state for the M step. For example, this might include the optimizer state for Adam.

### Parameters
- **params** (ParameterSet)
- **props** (PropertySet)
```

--------------------------------

### HMMEmissions.initialize

Source: https://probml.github.io/dynamax/api.html

Initializes the model parameters and their corresponding properties for the emission model.

```APIDOC
## HMMEmissions.initialize

### Description
Initialize the model parameters and their corresponding properties.

### Parameters
* **key** (Array | None) – random number generator
* **method** (str) – specifies the type of initialization
* **kwargs** – additional keyword arguments

### Returns
Tuple of parameters and their corresponding properties

### Return Type
Tuple[ParameterSet, PropertySet]
```

--------------------------------

### Get Initial Distribution

Source: https://probml.github.io/dynamax/api.html

Returns the initial distribution of the state for the GSSM. Requires model parameters and optionally inputs.

```python
initial_distribution(_params_ , _inputs =None_)
```

--------------------------------

### HMMInitialState.initialize

Source: https://probml.github.io/dynamax/api.html

Abstract method to initialize the model parameters and their corresponding properties.

```APIDOC
## HMMInitialState.initialize

### Description
Initialize the model parameters and their corresponding properties.

### Parameters
- **key** (Array | None) – random number generator
- **method** (str) – specifies the type of initialization
- **kwargs** - additional keyword arguments

### Returns
tuple of parameters and their corresponding properties.
```

--------------------------------

### Get Emission Distribution

Source: https://probml.github.io/dynamax/api.html

Returns the emission distribution of the state for the GSSM. Requires model parameters, current state, and optionally inputs.

```python
emission_distribution(_params_ , _state_ , _inputs =None_)
```

--------------------------------

### Initialize and sample from LGSSM

Source: https://probml.github.io/dynamax/notebooks/linear_gaussian_ssm/lgssm_hmc.html

Sets up a 2D latent state and 10D emission LGSSM with a stable, decaying random walk dynamics. It then samples states and emissions from this true model.

```python
state_dim = 2
emission_dim = 10
num_timesteps = 100
k1, k2, k3 = jr.split(jr.PRNGKey(0), 3)

# Construct the true model with randomly initialized parameters
true_A = 0.99 * random_rotation(seed=k1, n=state_dim, theta=jnp.pi / 10)
true_Sigma = 0.01 * jnp.eye(state_dim)
true_model = LinearGaussianConjugateSSM(state_dim, emission_dim)
true_params, param_props = true_model.initialize(
    key=k1, dynamics_weights=true_A, dynamics_covariance=true_Sigma)

# Sample states and emissions from the true model
true_states, emissions = true_model.sample(true_params, k3, num_timesteps)

```

--------------------------------

### Get Transition Distribution

Source: https://probml.github.io/dynamax/api.html

Returns the transition distribution of the state for the GSSM. Requires model parameters, current state, and optionally inputs.

```python
transition_distribution(_params_ , _state_ , _inputs =None_)
```

--------------------------------

### Abstract Emission Distribution Method

Source: https://probml.github.io/dynamax/_modules/dynamax/hidden_markov_model/models/abstractions.html

Abstract method to define how to get the emission distribution for a given state. Must be implemented by subclasses.

```python
    @abstractmethod
    def distribution(self,
                     params: ParameterSet,
                     state: IntScalar,
                     inputs: Optional[Float[Array, " input_dim"]]=None
    ) -> tfd.Distribution:
        """Return a distribution over the emission

        Args:
            params: emission parameters
            state: current latent state
            inputs: current inputs

        Returns:
            conditional distribution of the emission

        """
        raise NotImplementedError
```

--------------------------------

### initialize

Source: https://probml.github.io/dynamax/_modules/dynamax/hidden_markov_model/models/gaussian_hmm.html

Initializes the emission parameters for a Gaussian HMM. Parameters can be specified manually or sampled from a prior distribution. K-Means initialization is also supported if emissions data is provided.

```APIDOC
## initialize

### Description
Initializes the emission parameters for a Gaussian HMM. Parameters can be specified manually or sampled from a prior distribution. K-Means initialization is also supported if emissions data is provided.

### Arguments
* `key` (PRNGKey, optional): Random number generator for unspecified parameters. Must not be None if there are any unspecified parameters.
* `method` (str, optional): Method for initializing unspecified parameters. Currently, only "prior" is allowed. Defaults to "prior".
* `emission_means` (array, optional): Manually specified emission means.
* `emission_scales` (array, optional): Manually specified emission scales (sqrt of diagonal of spherical covariance matrix).
* `emissions` (array, optional): Emissions for initializing the parameters with kmeans.

### Returns
* `params`: Nested dataclasses of arrays containing model parameters.
* `props`: A nested dictionary of ParameterProperties to specify parameter constraints and whether or not they should be trained.
```

--------------------------------

### Initialize and Run EKF for Pendulum

Source: https://probml.github.io/dynamax/notebooks/nonlinear_gaussian_ssm/ekf_ukf_pendulum.html

Sets up the parameters for the Extended Kalman Filter using `ParamsNLGSSM` and then runs the EKF on the observation data. Ensure `PendulumParams` and `jnp` are properly imported and `obs` is defined.

```python
pendulum_params = PendulumParams()

# Define parameters for EKF
ekf_params = ParamsNLGSSM(
    initial_mean=pendulum_params.initial_state,
    initial_covariance=jnp.eye(states.shape[-1]) * 0.1,
    dynamics_function=pendulum_params.dynamics_function,
    dynamics_covariance=pendulum_params.dynamics_covariance,
    emission_function=pendulum_params.emission_function,
    emission_covariance=pendulum_params.emission_covariance,
)

ekf_posterior = ekf(ekf_params, obs)

```

--------------------------------

### Simulate Pendulum Data

Source: https://probml.github.io/dynamax/notebooks/nonlinear_gaussian_ssm/ekf_ukf_pendulum.html

Simulates pendulum states and observations using the defined parameters. This function is based on Särkkä's Example 3.7.

```python
# Pendulum simulation (Särkkä Example 3.7)
def simulate_pendulum(params=PendulumParams(), key=0, num_steps=400):
    if isinstance(key, int):
        key = jr.PRNGKey(key)
    # Unpack parameters
    M, N = params.initial_state.shape[0], params.emission_covariance.shape[0]
    f, h = params.dynamics_function, params.emission_function
    Q, R = params.dynamics_covariance, params.emission_covariance

    def _step(carry, rng):
        state = carry
        rng1, rng2 = jr.split(rng, 2)

        next_state = f(state) + jr.multivariate_normal(rng1, jnp.zeros(M), Q)
        obs = h(next_state) + jr.multivariate_normal(rng2, jnp.zeros(N), R)
        return next_state, (next_state, obs)

    rngs = jr.split(key, num_steps)
    _, (states, observations) = lax.scan(_step, params.initial_state, rngs)
    return states, observations


states, obs = simulate_pendulum()

```

--------------------------------

### initialize

Source: https://probml.github.io/dynamax/_modules/dynamax/hidden_markov_model/models/abstractions.html

Abstract method to initialize the model parameters and their corresponding properties.

```APIDOC
## initialize

### Description
Initialize the model parameters and their corresponding properties.

### Method
`initialize(key: Optional[Array]=None, method: str="prior", **kwargs) -> Tuple[ParameterSet, PropertySet]`

### Parameters
* **key** (Optional[Array]) - random number generator
* **method** (str) - specifies the type of initialization

### Returns
* **Tuple[ParameterSet, PropertySet]** - tuple of parameters and their corresponding properties
```

--------------------------------

### initialize_m_step_state (Transition Distribution)

Source: https://probml.github.io/dynamax/api.html

Initializes any required state for the M-step of the transition distribution. This might include setting up optimizer states.

```APIDOC
## initialize_m_step_state (Transition Distribution)

### Description
Initialize any required state for the M step.
For example, this might include the optimizer state for Adam.

### Parameters
* **params** (ParameterSet) –
* **props** (PropertySet) –

### Returns
Any required state for the M-step.

### Return Type
Any
```

--------------------------------

### Get HMM Transition Distribution

Source: https://probml.github.io/dynamax/_modules/dynamax/hidden_markov_model/models/abstractions.html

Returns the transition distribution for a given state, parameterized by the transition component. This defines the probabilities of moving between states.

```python
    def transition_distribution(self, params: HMMParameterSet, 
                                state: IntScalar, 
                                inputs:Optional[Array] = None) 
                                -> tfd.Distribution:
        """Return the transition distribution."""
        return self.transition_component.distribution(params.transitions, state, inputs=inputs)
```

--------------------------------

### Initialize HMM Parameters (Prior Method)

Source: https://probml.github.io/dynamax/api.html

Initializes HMM parameters, sampling unspecified ones from their prior distributions. Requires a PRNGKey if parameters are not fully specified.

```python
initialize(_key =Array([0, 0], dtype=uint32)_, _method ='prior'_, _initial_probs =None_, _transition_matrix =None_, _emission_weights =None_, _emission_means =None_, _emission_covariances =None_, _emissions =None_)
```

--------------------------------

### Get PyTree Leaf Count

Source: https://probml.github.io/dynamax/_modules/dynamax/utils/utils.html

Calculates the total number of leaves within a JAX PyTree. Handles None PyTrees by returning 0.

```python
def pytree_len(pytree):
    """Return the number of leaves in a PyTree."""
    if pytree is None:
        return 0
    else:
        return len(tree_leaves(pytree)[0])
```

--------------------------------

### Get HMM Emission Distribution

Source: https://probml.github.io/dynamax/_modules/dynamax/hidden_markov_model/models/abstractions.html

Returns the emission distribution for a given state, parameterized by the emission component. This defines the probabilities of observing emissions from each state.

```python
    def emission_distribution(self, params: HMMParameterSet, 
                              state: IntScalar, 
                              inputs:Optional[Array] = None) 
                              -> tfd.Distribution:
        """Return the emission distribution."""
        return self.emission_component.distribution(params.emissions, state, inputs=inputs)
```

--------------------------------

### initialize

Source: https://probml.github.io/dynamax/api.html

Initialize the model parameters and their corresponding properties. You can either specify parameters manually via the keyword arguments, or you can have them set automatically. If any parameters are not specified, you must supply a PRNGKey. Parameters will then be sampled from the prior (if method==prior).

```APIDOC
## initialize

### Description
Initialize the model parameters and their corresponding properties. You can either specify parameters manually via the keyword arguments, or you can have them set automatically. If any parameters are not specified, you must supply a PRNGKey. Parameters will then be sampled from the prior (if method==prior). 

### Parameters
* **key** (Array) – random number generator for unspecified parameters. Must not be None if there are any unspecified parameters.
* **method** (str) – method for initializing unspecified parameters. Both “prior” and “kmeans” are supported.
* **initial_probs** (Float [jaxlib._jax.Array, 'num_states'] | None) – manually specified initial state probabilities.
* **transition_matrix** (Float [jaxlib._jax.Array, 'num_states num_states'] | None) – manually specified transition matrix.
* **emission_means** (Float [jaxlib._jax.Array, 'num_states emission_dim'] | None) – manually specified emission means.
* **emission_scales** (Float [jaxlib._jax.Array, 'num_states'] | None) – manually specified emission scales (sqrt of diagonal of covariance matrix).
* **emissions** (Float [jaxlib._jax.Array, 'num_timesteps emission_dim'] | None) – emissions for initializing the parameters with kmeans.

### Returns
Model parameters and their properties.

### Return type
Tuple[HMMParameterSet, HMMPropertySet]
```

--------------------------------

### Get LinearAutoregressiveHMM Inputs Shape

Source: https://probml.github.io/dynamax/_modules/dynamax/hidden_markov_model/models/arhmm.html

Returns the expected shape of inputs for a single time step in the LinearAutoregressiveHMM, which is determined by the number of lags and emission dimensions.

```python
@property
def inputs_shape(self):
        """Return a pytree matching the pytree of tuples specifying the shape(s)
        of a single time step's inputs.
        """
        return (self.num_lags * self.emission_dim,)

```

--------------------------------

### LGSSM Initialization with Priors

Source: https://probml.github.io/dynamax/_modules/dynamax/linear_gaussian_ssm/models.html

Initializes a Linear Gaussian State Space Model, setting up prior distributions for initial state, dynamics, and emissions. Supports optional bias terms and custom prior distributions via keyword arguments.

```python
def __init__(self, state_dim, emission_dim, input_dim=0, has_dynamics_bias=True, has_emissions_bias=True, **kw_priors):
    super().__init__(state_dim=state_dim, emission_dim=emission_dim, input_dim=input_dim,
         has_dynamics_bias=has_dynamics_bias, has_emissions_bias=has_emissions_bias)

    # Initialize prior distributions
    def default_prior(arg, default):
        return kw_priors[arg] if arg in kw_priors else default

    self.initial_prior = default_prior(
        'initial_prior',
        NIW(loc=jnp.zeros(self.state_dim),
            mean_concentration=1.,
            df=self.state_dim + 0.1,
            scale=jnp.eye(self.state_dim)))

    self.dynamics_prior = default_prior(
        'dynamics_prior',
        MNIW(loc=jnp.zeros((self.state_dim, self.state_dim + self.input_dim + self.has_dynamics_bias)),
             col_precision=jnp.eye(self.state_dim + self.input_dim + self.has_dynamics_bias),
             df=self.state_dim + 0.1,
             scale=jnp.eye(self.state_dim)))

    self.emission_prior = default_prior(
        'emission_prior',
        MNIW(loc=jnp.zeros((self.emission_dim, self.state_dim + self.input_dim + self.has_emissions_bias)),
             col_precision=jnp.eye(self.state_dim + self.input_dim + self.has_emissions_bias),
             df=self.emission_dim + 0.1,
             scale=jnp.eye(self.emission_dim)))
```

--------------------------------

### Get Transition Matrix at Time t

Source: https://probml.github.io/dynamax/_modules/dynamax/hidden_markov_model/inference.html

Retrieves the transition matrix for a given time step `t`. It supports both fixed transition matrices and time-varying transition functions.

```python
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"]:
    """
    Helper function to get the transition matrix at time `t`.
    """
    if transition_fn is not None:
        return transition_fn(t)
    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.")

```

--------------------------------

### initialize_m_step_state (Emission Distribution)

Source: https://probml.github.io/dynamax/api.html

Initializes any required state for the M-step of the emission distribution. This might include setting up optimizer states.

```APIDOC
## initialize_m_step_state (Emission Distribution)

### Description
Initialize any required state for the M step.
For example, this might include the optimizer state for Adam.

### Parameters
* **params** (ParameterSet) –
* **props** (PropertySet) –

### Returns
Any required state for the M-step.

### Return Type
Any
```

--------------------------------

### Helper to get one parameter at time t

Source: https://probml.github.io/dynamax/_modules/dynamax/linear_gaussian_ssm/inference.html

Retrieves a specific parameter at a given timestep. Handles both callable parameters (functions of time) and array-based parameters.

```python
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
```

--------------------------------

### Initialize Gaussian HMM Parameters

Source: https://probml.github.io/dynamax/_modules/dynamax/hidden_markov_model/models/gaussian_hmm.html

Initializes the parameters for a Low-Rank Gaussian HMM. Supports initialization via prior sampling or K-Means clustering on emissions. Requires a PRNGKey if parameters are not manually specified.

```python
def initialize(self, 
                   key: Array = jr.PRNGKey(0),
                   method: str = "prior",
                   emission_means: Optional[Float[Array, "num_states emission_dim"]] = None,
                   emission_cov_diag_factors: Optional[Float[Array, "num_states emission_dim"]] = None,
                   emission_cov_low_rank_factors: Optional[Float[Array, "num_states emission_dim emission_rank"]] = None,
                   emissions: Optional[Float[Array, "num_timesteps emission_dim"]] = None
        ) -> Tuple[ParamsLowRankGaussianHMMEmissions, ParamsLowRankGaussianHMMEmissions]:
        """Initialize the model parameters and their corresponding properties.

        You can either specify parameters manually via the keyword arguments, or you can have
        them set automatically. If any parameters are not specified, you must supply a PRNGKey.
        Parameters will then be sampled from the prior (if `method==prior`).

        Note: in the future we may support more initialization schemes, like K-Means.

        Args:
            key (PRNGKey, optional): random number generator for unspecified parameters. Must not be None if there are any unspecified parameters.
            method (str, optional): method for initializing unspecified parameters. Currently, only "prior" is allowed. Defaults to "prior".
            emission_means (array, optional): manually specified emission means.
            emission_cov_diag_factors (array, optional): manually specified diagonals of the emission covariances.
            emission_cov_low_rank_factors (array, optional): manually specified low rank factors of the emission covariances.
            emissions (array, optional): emissions for initializing the parameters with kmeans.

        Returns:
            params: nested dataclasses of arrays containing model parameters.
            props: a nested dictionary of ParameterProperties to specify parameter constraints and whether or not they should be trained.
        """
        if method.lower() == "kmeans":
            assert emissions is not None, "Need emissions to initialize the model with K-Means!"
            from sklearn.cluster import KMeans
            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))

        elif method.lower() == "prior":
            # We don't have a real prior
            key1, key2, key3 = jr.split(key, 3)
            _emission_means = jr.normal(key1, (self.num_states, self.emission_dim))
            _emission_cov_diag_factors = \
                tfd.Gamma(self.emission_diag_factor_conc, self.emission_diag_factor_rate)
                    .sample(seed=key2, sample_shape=((self.num_states, self.emission_dim)))
            _emission_cov_low_rank_factors = jr.normal(key3, (self.num_states, self.emission_dim, self.emission_rank))

        else:
            raise Exception("Invalid initialization method: {}".format(method))

        # Only use the values above if the user hasn't specified their own
        default = lambda x, x0: x if x is not None else x0
        params = ParamsLowRankGaussianHMMEmissions(
            means=default(emission_means, _emission_means),
            cov_diag_factors=default(emission_cov_diag_factors, _emission_cov_diag_factors),
            cov_low_rank_factors=default(emission_cov_low_rank_factors, _emission_cov_low_rank_factors))
        props = ParamsLowRankGaussianHMMEmissions(
            means=ParameterProperties(),
            cov_diag_factors=ParameterProperties(constrainer=tfb.Softplus()),
            cov_low_rank_factors=ParameterProperties())
        return params, props
```

--------------------------------

### Get HMM Emission Shape

Source: https://probml.github.io/dynamax/_modules/dynamax/hidden_markov_model/models/abstractions.html

Retrieves the emission shape from the HMM's emission component. This property is essential for understanding the expected output format of the emissions.

```python
    # Implement the SSM abstract methods by passing on to the components
    @property
    def emission_shape(self):
        """Return the shape of the emission distribution."""
        return self.emission_component.emission_shape
```

--------------------------------

### Initialize Bernoulli HMM Parameters

Source: https://probml.github.io/dynamax/_modules/dynamax/hidden_markov_model/models/bernoulli_hmm.html

Initializes parameters for a Bernoulli HMM, including initial probabilities, transition matrices, and emission probabilities. Requires JAX random keys for stochastic initialization.

```python
def initialize(self, key, method=None, initial_probs=None, transition_matrix=None, emission_probs=None):
    """
    Initialize parameters for the Bernoulli HMM.

    Args:
        key: JAX random key.
        method: Initialization method. Defaults to None.
        initial_probs: Initial state probabilities. Defaults to None.
        transition_matrix: Transition matrix between states. Defaults to None.
        emission_probs: Emission probabilities for Bernoulli emissions. Defaults to None.

    Returns:
        Model parameters and their properties.
    """
    key1, key2, key3 = jr.split(key, 3)
    params, props = dict(), dict()
    params["initial"], props["initial"] = self.initial_component.initialize(
        key1, method=method, initial_probs=initial_probs
    )
    params["transitions"], props["transitions"] = self.transition_component.initialize(
        key2, method=method, transition_matrix=transition_matrix
    )
    params["emissions"], props["emissions"] = self.emission_component.initialize(
        key3, method=method, emission_probs=emission_probs
    )
    return ParamsBernoulliHMM(**params), ParamsBernoulliHMM(**props)
```

--------------------------------

### Setup Plotting Styles and Colormaps

Source: https://probml.github.io/dynamax/notebooks/hmm/autoregressive_hmm.html

Configures plotting aesthetics using Seaborn and creates a custom colormap for visualizations. This is useful for distinguishing between different states in plots.

```python
sns.set_style("white")

color_names = [
    "windows blue",
    "red",
    "amber",
    "faded green",
    "dusty purple",
    "orange",
    "brown",
    "pink"
]

colors = sns.xkcd_palette(color_names)
cmap = gradient_cmap(colors)

```

--------------------------------

### HMMTransitions.initialize

Source: https://probml.github.io/dynamax/api.html

Abstract method to initialize the model parameters and their corresponding properties.

```APIDOC
## HMMTransitions.initialize

### Description
Initialize the model parameters and their corresponding properties.

### Parameters
- **key** (Array | None) – random number generator
- **method** (str) – specifies the type of initialization
- **kwargs** - additional keyword arguments

### Returns
tuple of parameters and their corresponding properties.
```

--------------------------------

### Check JAX Devices

Source: https://probml.github.io/dynamax/notebooks/linear_gaussian_ssm/lgssm_parallel_inference.html

Prints the available JAX devices and their platforms. Sets a flag to indicate if running on CPU.

```python
print(jax.devices())
print(jax.devices()[0].platform)
if jax.devices()[0].platform == 'cpu':
    cpu_mode = True
else:
    cpu_mode = False
    

```

--------------------------------

### Setup Plotting Style

Source: https://probml.github.io/dynamax/notebooks/hmm/custom_hmm.html

Configures the plotting style using Seaborn and defines a custom colormap for visualizations. This helps in creating aesthetically pleasing plots for HMM analysis.

```python
sns.set_style("white")

color_names = [
    "windows blue",
    "red",
    "amber"
]

colors = sns.xkcd_palette(color_names)
cmap = gradient_cmap(colors)

```

--------------------------------

### Initialize Gaussian HMM Parameters

Source: https://probml.github.io/dynamax/_modules/dynamax/hidden_markov_model/models/gaussian_hmm.html

Initializes the parameters and properties for a Gaussian HMM. Supports prior or k-means initialization for unspecified parameters. Requires a PRNGKey if parameters are not manually provided.

```python
def initialize(self, key: Array=jr.PRNGKey(0),
                   method: str="prior",
                   initial_probs: Optional[Float[Array, " num_states"]]=None,
                   transition_matrix: Optional[Float[Array, "num_states num_states"]]=None,
                   emission_means: Optional[Float[Array, "num_states emission_dim"]]=None,
                   emission_scales: Optional[Float[Array, " num_states"]]=None,
                   emissions:  Optional[Float[Array, "num_timesteps emission_dim"]]=None
        ) -> Tuple[HMMParameterSet, HMMPropertySet]:
        """Initialize the model parameters and their corresponding properties.

        You can either specify parameters manually via the keyword arguments, or you can have
        them set automatically. If any parameters are not specified, you must supply a PRNGKey.
        Parameters will then be sampled from the prior (if `method==prior`).

        Args:
            key: random number generator for unspecified parameters. Must not be None if there are any unspecified parameters.
            method: method for initializing unspecified parameters. Both "prior" and "kmeans" are supported.
            initial_probs: manually specified initial state probabilities.
            transition_matrix: manually specified transition matrix.
            emission_means: manually specified emission means.
            emission_scales: manually specified emission scales (sqrt of diagonal of covariance matrix).
            emissions: emissions for initializing the parameters with kmeans.

        Returns:
            Model parameters and their properties.

        """
        key1, key2, key3 = jr.split(key , 3)
        params, props = dict(), dict()
        params["initial"], props["initial"] = self.initial_component.initialize(
            key1, method=method, initial_probs=initial_probs
        )
        params["transitions"], props["transitions"] = self.transition_component.initialize(
            key2, method=method, transition_matrix=transition_matrix
        )
        params["emissions"], props["emissions"] = self.emission_component.initialize(
            key3, method=method, emission_means=emission_means, emission_scales=emission_scales, emissions=emissions
        )
        return ParamsSphericalGaussianHMM(**params), ParamsSphericalGaussianHMM(**props)
```

--------------------------------

### Example Usage: Smoothing, Forecasting, and Plotting

Source: https://probml.github.io/dynamax/notebooks/linear_gaussian_ssm/lgssm_hmc.html

Demonstrates the usage of the `smooth_and_forecast` function and then visualizes the results using `plot_emissions_and_forecast`. This snippet shows a complete workflow from model prediction to visualization.

```python
(true_smooth_emissions, true_smooth_emissions_std, 
 true_forecast_emissions, true_forecast_emissions_std) = \
    smooth_and_forecast(true_model, true_params, emissions)

plot_emissions_and_forecast(true_smooth_emissions,
                            true_smooth_emissions_std,
                            true_forecast_emissions,
                            true_forecast_emissions_std)

```

--------------------------------

### HMMInitialState.m_step

Source: https://probml.github.io/dynamax/api.html

Perform an M-step on the initial distribution parameters.

```APIDOC
## HMMInitialState.m_step

### Description
Perform an M-step on the initial distribution parameters.

### Parameters
- **params** (ParameterSet) – current initial distribution parameters
- **props** (PropertySet) – parameter properties
- **batch_stats** (PyTree) – PyTree of sufficient statistics from each sequence, as output by `collect_suff_stats()`.
- **m_step_state** (Any) – any state required for the M-step
- **scale** (float) – how to scale the objective

### Returns
Parameters that maximize the expected log joint probability.
```

--------------------------------

### GammaHMMEmissions Initialization and Methods

Source: https://probml.github.io/dynamax/_modules/dynamax/hidden_markov_model/models/gamma_hmm.html

This snippet details the initialization of the GammaHMMEmissions class and its core methods for setting up and using Gamma distributions within an HMM.

```APIDOC
## GammaHMMEmissions

### Description
Models Gamma emissions for a Hidden Markov Model (HMM).

### Class Definition
```python
class GammaHMMEmissions(HMMEmissions):
    """Gamma emissions for an HMM.
    
    :param num_states: number of discrete states $K$
    :param m_step_optimizer: ``optax`` optimizer, like Adam.
    :param m_step_num_iters: number of optimizer steps per M-step.
    """
    def __init__(
        self,
        num_states: int,
        m_step_optimizer: optax.GradientTransformation = optax.adam(1e-2),
        m_step_num_iters: int = 50,
    ):
        ...
```

### Methods

#### `initialize`

Initialize the model parameters and their corresponding properties.

**Parameters:**

*   `key` (Array, optional): Random number generator for unspecified parameters. Defaults to `jr.PRNGKey(0)`.
*   `method` (str, optional): Method for initializing unspecified parameters. Supports "prior" and "kmeans". Defaults to "prior".
*   `emission_concentrations` (Float[Array, " num_states"], optional): Manually specified emission concentrations.
*   `emission_rates` (Float[Array, " num_states"], optional): Manually specified emission rates.
*   `emissions` (Float[Array, " num_timesteps"], optional): Emissions for initializing the parameters with kmeans.

**Returns:**

*   `Tuple[ParamsGammaHMMEmissions, ParamsGammaHMMEmissions]`: Model parameters and their properties.

**Example Usage:**

```python
# Assuming 'hmm' is an instance of GammaHMMEmissions
params, props = hmm.initialize(key=jr.PRNGKey(1), method="kmeans", emissions=sample_emissions)
```

#### `log_prior`

Compute the log-prior probability of the parameters.

**Parameters:**

*   `params` (ParamsGammaHMMEmissions): The parameters of the emission model.

**Returns:**

*   `float`: The log-prior probability.

#### `distribution`

Return the emission distribution for a given state.

**Parameters:**

*   `params` (ParamsGammaHMMEmissions): The parameters of the emission model.
*   `state` (int): The current discrete state.
*   `inputs` (Optional, optional): Additional inputs to the distribution. Defaults to None.

**Returns:**

*   `tfd.Distribution`: The emission distribution (Gamma) for the given state.

### Parameter Structure (`ParamsGammaHMMEmissions`)

```python
class ParamsGammaHMMEmissions(NamedTuple):
    """Parameters for the Gamma emissions of an HMM."""
    concentration: Union[Float[Array, " state_dim"], ParameterProperties]
    rate: Union[Float[Array, " state_dim"], ParameterProperties]
```

### Emission Shape

*   `emission_shape` (Tuple): Shape of the emission distribution, which is `()` for scalar Gamma distributions.
```

--------------------------------

### MLP Parameter and Apply Function Setup

Source: https://probml.github.io/dynamax/notebooks/generalized_gaussian_ssm/cmgf_mlp_classification_demo.html

Initializes MLP parameters, flattens them, and creates an apply function for use with flattened parameters. Requires model dimensions and a PRNG key.

```python
def get_mlp_flattened_params(model_dims, key=0):
    if isinstance(key, int):
        key = jr.PRNGKey(key)

    # Define MLP model
    input_dim, features = model_dims[0], model_dims[1:]
    model = MLP(features)
    dummy_input = jnp.ones((input_dim,))

    # Initialize parameters using dummy input
    params = model.init(key, dummy_input)
    flat_params, unflatten_fn = ravel_pytree(params)

    # Define apply function
    def apply(flat_params, x, model, unflatten_fn):
        return model.apply(unflatten_fn(flat_params), jnp.atleast_1d(x))

    apply_fn = partial(apply, model=model, unflatten_fn=unflatten_fn)

    return model, flat_params, unflatten_fn, apply_fn
```

--------------------------------

### Example Output of Marginal Log-Likelihood Estimates

Source: https://probml.github.io/dynamax/notebooks/nonlinear_gaussian_ssm/ekf_ukf_spiral.html

Shows the typical output format for the marginal log-likelihood estimates from EKF and UKF. This serves as a reference for expected results when running the estimation code.

```text
EKF Marginal LL Estimate: -7.378
UKF Marginal LL Estimate: -7.052
```

--------------------------------

### Initialize Poisson HMM Parameters

Source: https://probml.github.io/dynamax/_modules/dynamax/hidden_markov_model/models/poisson_hmm.html

Initializes the parameters and their properties for a Poisson HMM. It uses provided keys and optional manual specifications for initial probabilities, transition matrix, and emission rates. The initialization method defaults to 'prior'.

```python
key1, key2, key3 = jr.split(key , 3)
params, props = dict(), dict()
params["initial"], props["initial"] = self.initial_component.initialize(key1, method=method, initial_probs=initial_probs)
params["transitions"], props["transitions"] = self.transition_component.initialize(key2, method=method, transition_matrix=transition_matrix)
params["emissions"], props["emissions"] = self.emission_component.initialize(key3, method=method, emission_rates=emission_rates)
return ParamsPoissonHMM(**params), ParamsPoissonHMM(**props)
```

--------------------------------

### Initialize Gaussian HMM Parameters

Source: https://probml.github.io/dynamax/_modules/dynamax/hidden_markov_model/models/gaussian_hmm.html

Initializes parameters for a Gaussian HMM using a specified key and initialization method. Supports prior or k-means initialization for initial probabilities, transition matrix, and emissions.

```python
key1, key2, key3 = jr.split(key , 3)
        params, props = dict(), dict()
        params["initial"], props["initial"] = self.initial_component.initialize(
            key1, method=method, initial_probs=initial_probs
        )
        params["transitions"], props["transitions"] = self.transition_component.initialize(
            key2, method=method, transition_matrix=transition_matrix
        )
        params["emissions"], props["emissions"] = self.emission_component.initialize(
            key3, method=method, emission_means=emission_means, emission_covariances=emission_covariances, emissions=emissions
        )
        return ParamsGaussianHMM(**params), ParamsGaussianHMM(**props)
```
Tokens:
Raw