DiscreteMarkovChain#

class pymc_extras.distributions.DiscreteMarkovChain(*args, steps=None, n_lags=1, **kwargs)[source]#

A Discrete Markov Chain is a sequence of random variables

\[\{x_t\}_{t=0}^T\]

Where transition probability \(P(x_t | x_{t-1})\) depends only on the state of the system at \(x_{t-1}\). With n_lags > 1 the chain is of higher order, and the transition probability \(P(x_t | x_{t-1}, \dots, x_{t-n\_lags})\) depends on the last n_lags states.

Parameters:

P (tensor) –
Matrix of transition probabilities between states. Rows must sum to 1. One of P or P_logits must be provided.

When time_varying_P=False (default), P is a k x k matrix shared across all transitions, with optional leading batch dimensions: (*batch, k, k).

When time_varying_P=True, P carries an extra time axis just before the two state axes: (*batch, steps, k, k). P[..., t, :, :] is the transition matrix used to go from state t to state t + 1, so the time axis must have length steps (one transition matrix per step). When steps is not given explicitly, it is inferred from this axis.

With n_lags > 1 the two state axes become n_lags + 1 axes, all of length k: P[..., x_{t-n_lags}, ..., x_{t-1}, x_t], so the shape is (*batch, k, ..., k) (or (*batch, steps, k, ..., k) when time-varying).
P_logit (tensor, optional) – Matrix of transition logits. Converted to probabilities via Softmax activation. One of P or P_logits must be provided.
steps (tensor, optional) – Length of the markov chain. Only needed if state is not provided.
init_dist (unnamed distribution, optional) –
Distribution over the n_lags initial states. Unnamed refers to distributions created with the .dist() API. A scalar-support distribution (e.g. Categorical) is broadcast IID across the initial states; a vector-support distribution (e.g. JointCategorical) provides their joint distribution and must have support length n_lags.

Warning

init_dist will be cloned, rendering it independent of the one passed as input.
n_lags (int, default 1) – Order of the chain: how many previous states the transition probability conditions on. P gains one state axis per extra lag (see P above) and the chain starts from n_lags initial states drawn from init_dist, so it has n_lags + steps states in total.
time_varying_P (bool, default False) – If True, P is interpreted as a sequence of transition matrices, one per step, with shape (*batch, steps, k, k) (see P above). This disambiguates the time axis from a leading batch dimension.

Notes

The initial distribution will be cloned, rendering it distinct from the one passed as input.

Examples

Create a Markov Chain of length 100 with 3 states. The number of states is given by the shape of P, 3 in this case.

import numpy as np
import pymc as pm
import pymc_extras as pmx

with pm.Model() as markov_chain:
    P = pm.Dirichlet("P", a=[1, 1, 1], size=(3,))
    init_dist = pm.Categorical.dist(p=np.full(3, 1 / 3))
    markov_chain = pmx.DiscreteMarkovChain(
        "markov_chain", P=P, init_dist=init_dist, shape=(100,)
    )

Use a time-varying transition matrix. P carries a leading time axis of length steps (here 99 transitions for a chain of 100 states), so its shape is (steps, k, k). Below the chain switches regime partway through: a “sticky” kernel governs the first 40 transitions and a “mixing” kernel the remaining 59, assembled by repeating each kernel along the time axis.

import numpy as np
import pymc as pm
import pymc_extras as pmx
import pytensor.tensor as pt

with pm.Model() as regime_chain:
    # Two 2 x 2 kernels: states persist under "sticky", shuffle under "mixing".
    P_sticky = pm.Dirichlet("P_sticky", a=np.eye(2) * 8 + 1, size=2)
    P_mixing = pm.Dirichlet("P_mixing", a=np.ones((2, 2)), size=2)

    # Stack one kernel per step: 40 sticky then 59 mixing -> shape (99, 2, 2)
    P = pt.concatenate(
        [pt.repeat(P_sticky[None], 40, axis=0), pt.repeat(P_mixing[None], 59, axis=0)],
        axis=0,
    )

    init_dist = pm.Categorical.dist(p=np.full(2, 0.5))
    markov_chain = pmx.DiscreteMarkovChain(
        "markov_chain",
        P=P,
        init_dist=init_dist,
        time_varying_P=True,
        shape=(100,),
    )

Use a second-order chain, where each state depends on the previous two. P gains one state axis per lag, so it is (2, 2, 2) here: P[x_{t-2}, x_{t-1}, x_t]. The chain starts from n_lags=2 initial states, so shape=(100,) means 98 transitions.

import numpy as np
import pymc as pm
import pymc_extras as pmx

with pm.Model() as second_order_chain:
    P = pm.Dirichlet("P", a=np.ones((2, 2, 2)))
    init_dist = pm.Categorical.dist(p=np.full(2, 0.5))
    markov_chain = pmx.DiscreteMarkovChain(
        "markov_chain",
        P=P,
        init_dist=init_dist,
        n_lags=2,
        shape=(100,),
    )

__init__()#

Methods

`__init__`()
`dist`([P, logit_P, steps, init_dist, n_lags, ...])	Create a tensor variable corresponding to the cls distribution.
`rv_op`(P, steps, init_dist, n_lags[, size, ...])	The type of the None singleton.