Custom step function does not store value in trace

Questions v5
1

Hello,

I am trying to create a custom step function for my model so that I can include a variable without making said variable a free random variable. I could not adapt the solution provided in here for my model; running the notebook with PyMC 5.6.1 throws this error:

RemoteTraceback                           Traceback (most recent call last)
RemoteTraceback: 
"""
Traceback (most recent call last):
  File "/usr/local/lib/python3.8/dist-packages/pymc/sampling/parallel.py", line 122, in run
    self._start_loop()
  File "/usr/local/lib/python3.8/dist-packages/pymc/sampling/parallel.py", line 174, in _start_loop
    point, stats = self._step_method.step(self._point)
  File "/usr/local/lib/python3.8/dist-packages/pymc/step_methods/compound.py", line 236, in step
    sts.pop("model_logp", None)
AttributeError: 'tuple' object has no attribute 'pop'
"""

The above exception was the direct cause of the following exception:

AttributeError                            Traceback (most recent call last)
File /usr/local/lib/python3.8/dist-packages/pymc/sampling/parallel.py:122, in run()
    121     self._point = self._make_numpy_refs()
--> 122     self._start_loop()
    123 except KeyboardInterrupt:

File /usr/local/lib/python3.8/dist-packages/pymc/sampling/parallel.py:174, in _start_loop()
    173 try:
--> 174     point, stats = self._step_method.step(self._point)
    175 except SamplingError as e:

File /usr/local/lib/python3.8/dist-packages/pymc/step_methods/compound.py:236, in step()
    235 for sts in stats[:-1]:
--> 236     sts.pop("model_logp", None)
    237 return point, stats

AttributeError: 'tuple' object has no attribute 'pop'

I managed to write something that works based on this question. But there are two issue that I cannot figure out how to solve:

  1. The custom step values are not recorded in the inference data
  2. I would like to control the random seed in my custom step when using multiple cores

I am using PyMC version 5.6.1.

This is what I have so far:

import numpy as np
import pymc as pm
import arviz as az
import corner

########################
## generating test data
## not really important 

def gauss_1d(x, mu, sig):
    return np.exp(-np.square(x-mu)/sig)

def func(mass, Phi_s, Ms, a):
    mvec = np.power(10, mass-Ms)
    return Phi_s * np.power(mvec, a+1) * np.exp(-np.power(10,mvec))*1e4

nbin = 10
ndraw = 2 + nbin
mat = np.array([gauss_1d(np.arange(ndraw), i-1.1/(i+1), (i+1)/(3+i))*2/(ndraw-i) for i in range(ndraw)])
mat = mat/np.sum(mat, axis = 1, keepdims = True)*np.sin(np.linspace(np.pi/ndraw, np.pi/2, ndraw))[:,None]
mat = mat[1:ndraw-1,1:ndraw-1]

ml, mu = 8, 11
mbins = np.linspace(ml, mu, nbin+1)
mcens = 0.5*(mbins[1:]+mbins[:-1])

sel_par = [.01, 10.75, -1.1]
sel_mat = np.random.normal(mat, mat/10, mat.shape)
obs_val = np.sum(func(mcens, *sel_par)[:,None]*sel_mat, axis = 0, dtype = int)

###################
## custom step function

from pymc.step_methods.arraystep import BlockedStep

class FixedDistStep(BlockedStep):
    def __init__(self, var, md, mu, seed=4):
        model = pm.modelcontext(None)
        value_var = model.rvs_to_values[var]
        self.vars = [value_var]
        self.name = value_var.name

        self.rng = np.random.default_rng(seed=seed)
        self.md = md
        self.mu = mu
        
    def step(self, point: dict):
        mat = self.rng.normal(self.md, self.mu, self.md.shape)
        mat[mat<0] = 0
        mat[mat>1] = 1
        point[self.name] = mat
                
        return point, []

###################
## the model

with pm.Model() as rocky:

    a = pm.Uniform('a', lower=-2, upper=-.5)
    Ms = pm.Uniform('Ms', lower=8, upper=14)
    Phi_s = pm.Uniform('Phi_s', lower=0.0, upper=.02)
    
    mat_prc = pm.TruncatedNormal('mat_prc', mu=mat, sigma=(mat+1e-14)/10,
                                 lower=0, upper=1, shape=mat.shape, transform=None)

    phi_true = pm.Deterministic('phi_true', func(mcens, Phi_s, Ms, a))
    
    obs_ns = pm.Deterministic('obs_ns', pm.math.sum(phi_true[:,None]*mat_prc, axis = 0))
        
    logl_R = pm.Normal('rec_pars', mu=obs_ns, observed=obs_val,
                       sigma=np.power(obs_ns, .888))
     
    fixed_step = [FixedDistStep(mat_prc, mat, mat/10)]   

with rocky:
    idata = pm.sample(draws = 4096, tune = 1024, chains = 4, step=fixed_step,
                      cores = 4, random_seed=np.random.default_rng(seed = 4))

In the inference data mat_prc (the custom step variable) is a matrix of 0.5s for all the steps. The step is actually working and doing what it should, however, the values are not being recorded. How can I store the values correctly? And how can I assign a different numpy random seed to each chain?

I would appreciate any help. Thanks!

-Amir

2

I have had a go at this as I have also implemented Custom step methods before but never tried looking at the statistics of fixed variables. Like you, I am also seeing constant sampled values for fixed step variables. On the other hand if you do

pm.draw(custom_step_var, 1000)

you do get different values for each draw. So, I tried to go into the bowels of pymc starting with a single core. In the single core setup it seems that your drawn samples are recorded inside _iter_sample in mcmc.py right at this point

point, stats = step.step(point)
trace.record(point, stats)

breaking here actually did suggest that the drawn and recorded points are different each time (for the fixed step parameter too). Indeed if you run sample with chains=4, cores=1, then the values you get for your fixed step are not constant.

So this made me think is there something we are doing wrong in initialization of the random number generator that maybe doesn’t work as intended when it is shared across cores or something. The analogous operation in the case of parallelized sampler seems to be also in mcmc.py, _mp_sample

for draw in sampler:
    strace = traces[draw.chain]
    strace.record(draw.point, draw.stats)
    log_warning_stats(draw.stats)
    if draw.is_last:
        strace.close()

    if callback is not None:
        callback(trace=strace, draw=draw)

which does indeed always draw constant values for fixed step parameters (if you look at draw.point). Tried to understand this by going into the details of the ParallelSampler which draws the samples. Within here the drawing seems to occur at

draw = ProcessAdapter.recv_draw(self._active)

where draw is tuple whose first element is something called _proc and if you check proc._point it does contain again our constant values for the fixed step variable.
After here things became harder for me to understand but it seems that at least writing of sampled points happens inside _write_point of instance of _Process and when you check here the sampled values for fixed step variables are different! On the other hand when everything is said and done for some reason, draw has only constant values for these parameters. I am completely puzzled.

So for the time being I can only offer single core run as a solution. Maybe some of the devs will have a better idea on this.