Problem with NormalMixture

Questions v5
4

Here’s a reproducible example of what I’m seeing. This version only uses two predictors, but has the same number of observations.

import pymc as pm
import numpy as np
import pandas as pd

# Create fake data
mus = np.array([4, 14])
sigmas = np.array([2, 3])

predictors = pd.DataFrame({
    "male": pm.draw(pm.Categorical.dist(p=[0.5, 0.5]), draws=129),
    "age": pm.draw(pm.Normal.dist(mu=35, sigma=3), draws=129) 
})

p = pm.draw(pm.Normal.dist(mu=1 + predictors["age"] * 0.25 + predictors["male"] * 1.0))
p = p / (np.max(p) + 0.1)

index = np.array(pm.draw(pm.Categorical.dist(p=np.column_stack((p, 1-p))), draws=1))
data = pm.draw(pm.Normal.dist(mu=mus[index], sigma=2), draws=1)

# Model
classes = 2
with pm.Model(coords={"classes": range(classes), "predictors": predictors.columns}) as model:
    # Priors
    sigma = pm.HalfNormal("sigma", sigma=1)
    centers = pm.Normal(
        "centers",
        mu=[data.mean(), data.mean()],
        sigma=10,
        transform=pm.distributions.transforms.univariate_ordered,
        initval=[1, 2],
        dims="classes"
    )

    intercept = pm.Normal('intercept', mu=0, sigma=20)
    beta = pm.Normal("beta", mu=0, sigma=1, dims="predictors")
    theta = intercept + pm.math.dot(predictors, beta)
    p = pm.Deterministic("p", pm.math.invlogit(theta))
    w = [p, 1-p]

    # Likelihood
    y = pm.NormalMixture("obs", w=w, mu=centers, sigma=sigma, observed=data)

# Train
with model:
    idata = pm.sample()