Incorrect step assignments with custom step function

Questions v5
1

In order to port a complex PyMC3 model with custom conjugate steppers involving transformed variables to v5, I am trying to update this excellent old PyMC3 example from the docs to v5.

The first issue I am running into (I am sure there will be more) is that pm.sample doesn’t appear to set up the steppers correctly. Here is the code, a reduced version of the example:

import numpy as np
import pymc as pm

from pymc.distributions.transforms import sum_to_1
from pymc.step_methods.arraystep import BlockedStep

def sample_dirichlet(c):
    gamma = np.random.gamma(c)
    p = gamma / gamma.sum(axis=-1, keepdims=True)
    return p
    
class ConjugateStep(BlockedStep):
    def __init__(self, var, counts: np.ndarray, concentration):
        self.vars = [var]
        self.name = var.name
        self.counts = counts
        self.conc_prior = concentration

    def step(self, point: dict):
        conc_posterior = np.exp(point[self.conc_prior.transformed.name]) + self.counts
        draw = sample_dirichlet(conc_posterior)
        point[self.name] = sum_to_1.forward(draw)
        return point, []
        
# Generate data
J = 10
N = 500
ncounts = 20
tau_true = 0.5
alpha = tau_true * np.ones([N, J])
p_true = sample_dirichlet(alpha)
counts = np.zeros([N, J])
for i in range(N):
    counts[i] = np.random.multinomial(ncounts, p_true[i])
    
# Set up model and sample
with pm.Model() as model:
    tau = pm.Exponential("tau", lam=1, initval=1.0)
    alpha = pm.Deterministic("alpha", tau * np.ones([N, J]))
    p = pm.Dirichlet("p", a=alpha)
    step = [ConjugateStep(p, counts, tau)]
    trace = pm.sample(step=step, chains=2, cores=1, return_inferencedata=True)

Here is the initial output before it crashes

Sequential sampling (2 chains in 1 job)
CompoundStep
>ConjugateStep: [p]
>NUTS: [tau, p]

The parameter p shows up in both steps, which is not correct.

If I explicitly assign steppers for all parameters using step=[ConjugateStep(p, counts, tau), pm.NUTS(tau)], I get

Sequential sampling (2 chains in 1 job)
CompoundStep
>ConjugateStep: [p]
>NUTS: [tau]
>NUTS: [p]

with the same basic issue.

The problem seems to be somewhere in the function assign_step_methods, but my understanding of the PyMC codebase is insufficient to figure out more.

How can this problem be fixed? Any help would be appreciated!

2

Does seems like a bug, @ricardoV94 any idea?

3

You need to assign the value variable to self.vars and not the RV. So self.vars = [model.rvs_to_values[var]]`. During sampling it’s the value variable that matters not the RV.

Also to get the name of self.conc.prior, you can do model.rvs_to_values[concentration].name during the initialization.

Finally sum_to_1.forward won’t work. The signature is different and it would return a symbolic TensorVariable. You should implement the constraint with a NumPy function.

5

I removed the Dirichlet transform. Here is a working snippet:

import numpy as np
import pymc as pm
from pymc.step_methods.arraystep import BlockedStep

def sample_dirichlet(c):
    gamma = np.random.gamma(c)
    p = gamma / gamma.sum(axis=-1, keepdims=True)
    return p

class ConjugateStep(BlockedStep):
    def __init__(self, var, counts: np.ndarray, concentration):
        model = pm.modelcontext(None)
        value_var = model.rvs_to_values[var]
        self.vars = [value_var]
        self.name = value_var.name
        self.counts = counts
        self.conc_prior_name = model.rvs_to_values[concentration].name

    def step(self, point: dict):
        conc_posterior = np.exp(point[self.conc_prior_name]) + self.counts
        draw = sample_dirichlet(conc_posterior)
        point[self.name] = draw
        return point, []

# Generate data
J = 10
N = 500
ncounts = 20
tau_true = 0.5
alpha = tau_true * np.ones([N, J])
p_true = sample_dirichlet(alpha)
counts = np.zeros([N, J])
for i in range(N):
    counts[i] = np.random.multinomial(ncounts, p_true[i])

# Set up model and sample
with pm.Model() as model:
    tau = pm.Exponential("tau", lam=1, initval=1.0)
    alpha = pm.Deterministic("alpha", tau * np.ones([N, J]))
    p = pm.Dirichlet("p", a=alpha, transform=None)
    step = [ConjugateStep(p, counts, tau)]
    trace = pm.sample(step=step, chains=2, cores=1, return_inferencedata=True)
3 Likes
6

This is great - thanks a lot!

7

The originial PyMC3 example also runs this model with auto-assigned NUTS for both parameters, for comparison.

In your script, when I modfiy step = [ConjugateStep(p, counts, tau)] to step = [], the sampler crashes during initialization with SamplingError: Initial evaluation of model at starting point failed!. Just before that, it issues multiple RuntimeWarning: invalid value encountered in log return np.log(x).

What needs to be adjusted in the script to make this run just as it did on the original PyMC3 example?

8

I imagine the sampler fails because of the untransformed Dirichlet? If you’re sampling with NUTS you should definitely use the transform. I removed it to simplify the case with the transformed Dirichlet