CustomDist does not work with pm.mixture

Questions v5
1

Hi team,

I’m trying to implement a mixture using a customDist, but sampling fails with the error

AttributeError: 'float' object has no attribute 'name'

The following code produces the error on pymc 5.17.0

from scipy.stats import dirichlet, poisson, multinomial
from scipy.stats._distn_infrastructure import rv_discrete_frozen

import numpy as np

import matplotlib.pyplot as plt

import pymc as pm
from pymc.distributions.dist_math import factln, logpow


n_sims = 10_000

sim_data = [None] * n_sims

m = dirichlet([2.5,1.4,.1])
p_1 = poisson(5)

p_2 = poisson(.2, loc = 5000)

for i in range(n_sims): 

    m_samp = m.rvs(1).flatten()

    cat = multinomial(p = m_samp, n = 1).rvs(1).flatten()

    if np.argmax(cat) == 0:
        sim_data[i] = 0
    else:
        sim_data[i] = [p_1, p_2][np.argmax(cat) - 1].rvs(1)[0]

    

#functions for customdist
def shiftedpoisson_logp(y,dist_params:list[float]):
    lam = dist_params[0]
    shift = dist_params[1]
    if (lam < 0) | (shift < 0):
        return -np.inf
    else:
        logp = -lam + logpow(lam, y - shift) - factln(y - shift)

def random(
    lam: float,
    shift: int,
    size:int = 1,
) -> np.ndarray | float:
    rng = poisson(lam, loc=shift)
    return rng.rvs(size=size)

with pm.Model() as mod:
    lam_1 = pm.Exponential('lam_1', 1)
    lam_2 = pm.Exponential('lam_2', 1)
    shift = pm.DiscreteUniform('shift', 3000, 5000)
    m = pm.Dirichlet('m', [1,1,1])
    zero = pm.DiracDelta.dist(0)
    p_1 = pm.Poisson.dist(lam_1)
    p_2 =  pm.CustomDist.dist(
        lam_2,
        shift,
        logp=shiftedpoisson_logp,
        random=random,
        signature="()->()",
        size=1,
        dtype = 'int'
    )

    y = pm.Mixture('y', w = m, comp_dists = [zero, p_1, p_2], observed = sim_data)
    samples = pm.sample()

I’ve attached the full error message. It seems to be related to logp looking for a name, something that I guess doesn’t happen when .dist() is used. How can I avoid this issue?

error.txt (5.8 KB)

2

Your logp is mixing numpy and pytensor variables, and more improtantly using python if statements, which won’t work.

The good news is that you can define a Shifted Poisson much more simply these days, with the dist argument of CustomDist which only requires you to define the generative graph using PyTensor / PyMC variables:

import pymc as pm

def shifted_poisson(lam, shift, size):
    return pm.Poisson.dist(lam, size=size) + shift

with pm.Model() as mod:
    lam_1 = pm.Exponential('lam_1', 1)
    lam_2 = pm.Exponential('lam_2', 1)
    shift = pm.DiscreteUniform('shift', 3000, 5000)
    m = pm.Dirichlet('m', [1,1,1])
    zero = pm.DiracDelta.dist(0)
    p_1 = pm.Poisson.dist(lam_1)
    p_2 =  pm.CustomDist.dist(
        lam_2,
        shift,
        dist = shifted_poisson,
        dtype="int",
    )

    y = pm.Mixture('y', w = m, comp_dists = [zero, p_1, p_2], observed = [1, 2, 5, 1, 9000, 3])
    samples = pm.sample()

PyMC will figure out the logp automatically for you.

There are other examples in the docs: pymc.CustomDist — PyMC v5.14.0 documentation

1 Like
3

Cheers @ricardoV94.