Custom step method not working for Gamma distributions

[spribitzer]

Hello, working with custom steps method I recently noticed that a RV that was initially created using a Gamma distribution can not be used in a custom step.
As an example, I am fitting the noise level on a linear fit:

import numpy as np
import pymc3 as pm
import matplotlib.pyplot as plt

x = np.linspace(1,100,500)
y = 0.5*x 
sigma = 2
y_obs = y + sigma*np.random.randn(len(x))

plt.plot(x,y,'k')
plt.scatter(x,y_obs)

Normal sampling using NUTS

Setting up a model and analyze it using the built in NUTS sampler works:

with pm.Model() as model:
    sigma = pm.Gamma('sigma', mu=0.7, sigma=2)
    pm.Normal('y',mu = 0.5*x, sigma = sigma, observed = y_obs)
    trace = pm.sample(chains=2, cores=1, draws=500, tune= 100, progressbar = False)
    display(pm.summary(trace))

Incorrect sampling for Gamma prior

Now here is the bug: if I add a custom step method for sigma, pymc3 still samples sigma using NUTS. The custom step sampler is just a mock sampler that returns one fixed value:

from pymc3.step_methods.arraystep import BlockedStep

class SigmaStep(BlockedStep):
    def __init__(self, var):
            self.vars = [var]
            self.var = var

    def step(self, point: dict):
        newpoint = point.copy()
        newpoint[self.var.name] = 2.
        return newpoint

with model:
    trace = pm.sample(step = SigmaStep(sigma), chains=2, cores=1, draws=500, tune= 100, progressbar = False)
    display(pm.summary(trace))

Note how pymc3 states it is going to sample sigma using both, NUTS, and the custom step method.

But the results clearly show that NUTS was used (though the custom step does get executed, I checked that).

Correct behavior for Normal prior

Now if I change the prior on sigma from Gamma to Normal``, sigma` is sampled correctly using the custom step method:

with pm.Model() as model:
    sigma = pm.Normal('sigma', mu=0.7, sigma=2) # change from Gamma to Normal
    pm.Normal('y',mu = 0.5*x, sigma = sigma, observed = y_obs)
    trace = pm.sample(step = SigmaStep(sigma), chains=2, cores=1, draws=500, tune= 100, progressbar = False)
    display(pm.summary(trace))

Sidenote: I am using VS Code for my Jupyternotebooks. I noticed that the progressbar crashes my VS code (or the jupyter extension) pretty much everytime I run more complex, longer models. Therefore I always run sample with progressbar = False.

Versions and main components

• PyMC3 Version: 3.11.2
• Aesara Version: 2.0.7
• Python Version: 3.8.8
• Operating system: Windows 10
• installed via conda-forge

[ricardoV94]

I think you need to assign the gamma_log__, the autotransformed variable, to the sampler. This is not an issue with the normal because that is not auto-transformed.

You can also pass transformed=None when initializing gamma if you don't want the autotransform.

Does it work?

[spribitzer]

I have actually tried the same thing with the transformed variable (I didn't post it as I didn't want the example to be too confusing) it it shows the same behavior. Here is the code for completeness:

from pymc3.step_methods.arraystep import BlockedStep

from pymc3.step_methods.arraystep import BlockedStep

class SigmaStepTransformedSpace(BlockedStep):
    def __init__(self, var):
            self.vars = [var]
            self.var = var

    def step(self, point: dict):
        newpoint = point.copy()
        newpoint[self.var.transformed.name] = np.log(2)
        return newpoint

with pm.Model() as model:
    sigma = pm.Gamma('sigma', mu=0.7, sigma=2)
    pm.Normal('y',mu = 0.5*x, sigma = sigma, observed = y_obs)
    trace = pm.sample(step = SigmaStepTransformedSpace(sigma), chains=2, cores=1, draws=500, tune= 100, progressbar = False)
    display(pm.summary(trace))

Note: I have also tried hardwiring the variable name sigma_log__: newpoint["sigma_log__"] = np.log(2) , without success.

[spribitzer]

Using transform = None does work though:

class SigmaStepUntransformed(BlockedStep):
    def __init__(self, var):
            self.vars = [var]
            self.var = var

    def step(self, point: dict):
        newpoint = point.copy()
        newpoint[self.var.name] = 2
        return newpoint

with pm.Model() as model:
    sigma = pm.Gamma('sigma', mu=0.7, sigma=2, transform = None)
    pm.Normal('y',mu = 0.5*x, sigma = sigma, observed = y_obs)
    trace = pm.sample(step = SigmaStepUntransformedSpace(sigma), chains=2, cores=1, draws=500, tune= 100, progressbar = False)
    display(pm.summary(trace))

[ricardoV94]

Sorry if I was not clear. The problem is that sigma is auto-transformed to a variable contained in model['sigma_log__'] and after initialization, this sigma refers just to a pm.Deterministic which is just the backwards transformation from model['sigma_log__']. You have to assign the later to the step function if you want to use it:

step = SigmaStep(model['sigma_log__']),

If you have more questions, please feel free to open a topic on our discourse forum at https://discourse.pymc.io/

[spribitzer]

Okay, I can see this now. However, I feel like there should be a warning message or similar. If the sampler initializes the compound step and prints out that the same parameter is being sampled using both methods, that is pretty misleading and also makes debugging harder - at least for someone that does not have a complete understanding of the custom step methods.

I know custom steps are probably a pretty exotic thing, but it might be beneficial for the community to add a little more documentation and examples (from simple to more complex) for it. Being new to custom steps, the only illustration I could find was here, which sort of went over my head when I first looked at it and I needed to create a discourse topic to understand how to properly setup a custom sampler. At least a second example that uses a different transformation and shows how to either figure out what kind of transformation is done to the various variables in the model and how to reverse them (or how to do it with transform = None) could really help with getting started on custom sampling methods!

[ricardoV94]

Thanks for the feedback. Certainly the step logic could be more user friendly. I started a PR recently that improves things a bit in that direction. I'll keep your remarks in mind and see if something better can be done. #4701

[spribitzer]

Sweet!