JointCategorical#

class pymc_extras.distributions.JointCategorical(name: str, *args, rng=None, dims: str | Sequence[str | None] | None = None, initval=None, observed=None, total_size=None, transform=UNSET, default_transform=UNSET, **kwargs)[source]#

Joint distribution over n_lags categorical states.

A draw is a vector \((x_0, \dots, x_{n\_lags - 1})\), each entry taking one of n_states values. Unlike a product of independent categoricals, the entries can be arbitrarily correlated: the distribution is parameterized by the probability of every one of the n_states ** n_lags configurations, one array axis per entry.

\[\Pr(x_0, \dots, x_{n\_lags - 1}) = p[x_0, \dots, x_{n\_lags - 1}]\]
Parameters:

  • p (tensor_like of float) – Probability of each state configuration, with shape (n_states,) * n_lags: one axis per entry of a draw, so p[i, j, ...] is the probability of drawing [i, j, ...]. Any leading axes are batch dimensions. Note the parameter grows exponentially in n_lags, as the joint is stored explicitly.

  • logit_p (tensor_like of float) – Alternative parametrization, as unnormalized log-probabilities of the same shape. Only one of p or logit_p may be given.

  • n_lags (int) – Number of states in the joint. This is the support dimensionality: a single draw is a vector of length n_lags. The number of values each entry can take is taken from the trailing axes of p / logit_p.

Examples

A joint over two binary states. Of the four possible paths, only the two that stay put, [0, 0] and [1, 1], carry appreciable mass, so the pair is strongly correlated even though each state is marginally 50/50:

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

# One axis per lag, so p[i, j] is the probability of drawing [i, j].
p = np.array(
    [
        [0.45, 0.05],  # [0, 0], [0, 1]
        [0.05, 0.45],  # [1, 0], [1, 1]
    ]
)
dist = pmx.JointCategorical.dist(p=p, n_lags=2)

pm.draw(dist, draws=5)
# array([[1, 1],
#        [0, 0],
#        [0, 0],
#        [1, 1],
#        [0, 1]])

Each row is one draw: a vector of n_lags=2 states, the first entry being \(x_0\) and the second \(x_1\). Switching paths like [0, 1] show up in roughly one draw in ten. Setting those entries to exactly 0.0 instead would make them impossible, leaving a joint supported on two configurations only.

Notes

This is the distribution of the smoothed initial states of a marginalized DiscreteMarkovChain with n_lags > 1, where the first n_lags states are correlated a posteriori.

__init__()#

Methods

__init__()

dist([p, logit_p])

Create a tensor variable corresponding to the cls distribution.

rv_op(logit_p, *, n_lags[, size, rng])