pymc.Mixture

class pymc.Mixture(name, *args, rng=None, dims=None, initval=None, observed=None, total_size=None, transform=UNSET, default_transform=UNSET, **kwargs)

Mixture log-likelihood.

Often used to model subpopulation heterogeneity

\[f(x \mid w, \theta) = \sum_{i = 1}^n w_i f_i(x \mid \theta_i)\]

Support	\(\cup_{i = 1}^n \textrm{support}(f_i)\)
Mean	\(\sum_{i = 1}^n w_i \mu_i\)

Parameters:

wtensor_like of float

w >= 0 and w <= 1 the mixture weights

comp_distsiterable of unnamed distributions or single batched distribution

Distributions should be created via the .dist() API. If a single distribution is passed, the last size dimension (not shape) determines the number of mixture components (e.g. pm.Poisson.dist(…, size=components)) \(f_1, \ldots, f_n\)

Warning

comp_dists will be cloned, rendering them independent of the ones passed as input.

Examples

# Mixture of 2 Poisson variables
with pm.Model() as model:
    w = pm.Dirichlet("w", a=np.array([1, 1]))  # 2 mixture weights

    lam1 = pm.Exponential("lam1", lam=1)
    lam2 = pm.Exponential("lam2", lam=1)

    # As we just need the logp, rather than add a RV to the model, we need to call `.dist()`
    # These two forms are equivalent, but the second benefits from vectorization
    components = [
        pm.Poisson.dist(mu=lam1),
        pm.Poisson.dist(mu=lam2),
    ]
    # `shape=(2,)` indicates 2 mixture components
    components = pm.Poisson.dist(mu=pm.math.stack([lam1, lam2]), shape=(2,))

    like = pm.Mixture("like", w=w, comp_dists=components, observed=data)

# Mixture of Normal and StudentT variables
with pm.Model() as model:
    w = pm.Dirichlet("w", a=np.array([1, 1]))  # 2 mixture weights

    mu = pm.Normal("mu", 0, 1)

    components = [
        pm.Normal.dist(mu=mu, sigma=1),
        pm.StudentT.dist(nu=4, mu=mu, sigma=1),
    ]

    like = pm.Mixture("like", w=w, comp_dists=components, observed=data)

# Mixture of (5 x 3) Normal variables
with pm.Model() as model:
    # w is a stack of 5 independent size 3 weight vectors
    # If shape was `(3,)`, the weights would be shared across the 5 replication dimensions
    w = pm.Dirichlet("w", a=np.ones(3), shape=(5, 3))

    # Each of the 3 mixture components has an independent mean
    mu = pm.Normal("mu", mu=np.arange(3), sigma=1, shape=3)

    # These two forms are equivalent, but the second benefits from vectorization
    components = [
        pm.Normal.dist(mu=mu[0], sigma=1, shape=(5,)),
        pm.Normal.dist(mu=mu[1], sigma=1, shape=(5,)),
        pm.Normal.dist(mu=mu[2], sigma=1, shape=(5,)),
    ]
    components = pm.Normal.dist(mu=mu, sigma=1, shape=(5, 3))

    # The mixture is an array of 5 elements
    # Each element can be thought of as an independent scalar mixture of 3
    # components with different means
    like = pm.Mixture("like", w=w, comp_dists=components, observed=data)

# Mixture of 2 Dirichlet variables
with pm.Model() as model:
    w = pm.Dirichlet("w", a=np.ones(2))  # 2 mixture weights

    # These two forms are equivalent, but the second benefits from vectorization
    components = [
        pm.Dirichlet.dist(a=[1, 10, 100], shape=(3,)),
        pm.Dirichlet.dist(a=[100, 10, 1], shape=(3,)),
    ]
    components = pm.Dirichlet.dist(a=[[1, 10, 100], [100, 10, 1]], shape=(2, 3))

    # The mixture is an array of 3 elements
    # Each element comes from only one of the two core Dirichlet components
    like = pm.Mixture("like", w=w, comp_dists=components, observed=data)

Methods

Mixture.dist(w, comp_dists, **kwargs)

Create a tensor variable corresponding to the cls distribution.