How to add correlated noise to a mathematical model of observed data?

[mjedrz]

Hi I would like to do NUTS sampling with correlated noise added to my model of the observed data (the model is represented by “mu”, the noise is represented by “c” variable). When I try to do that, the noise is sampled only once and its value is repeated in all samples. How to make it so that c is sampled for every NUTS draw? Here is the current code:

rng = np.random.default_rng()
data = [151.2, 181.9, 219.7]
basic_model = Model()
with basic_model:
    # Priors for unknown model parameters
    a = pm.Normal("a", mu=50, sigma=5)
    b = pm.Normal("b", mu=40, sigma=7)
    #c is supposed to simulate noise
    c = rng.normal(15, 3)
    mu = [a + 2*b + 3*c,
          2*a + 2*b + 2*c,
          4*a + b + c]
    # Likelihood (sampling distribution) of observations
    Y_obs = Normal("Y_obs", mu=mu, sigma=0.1, observed=data)
    trace = sample(draws=1000)

This works well on the other hand for SMC_ABC, for which the model of observed data comes in the form of simulator:

data = [151.2, 181.9, 219.7]
def normal_sim(rng, a, b, size=1):
    c=rng.normal(15, 3, 1)
    x1=a+2*b+3*c
    x2=2*a+2*b+2*c
    x3=4*a+b+c
    return [x1, x2, x3]
with pm.Model() as example:
    a = pm.Normal("a", mu=50, sigma=5)
    b = pm.Normal("b", mu=40, sigma=7)
    s = pm.Simulator("s", normal_sim, params=(a, b), sum_stat="sort", epsilon=0.1, observed=data)
    idata = pm.sample_smc(5000)
    idata.extend(pm.sample_posterior_predictive(idata))

[iavicenna]

This is something I am a bit curious about too. I initially had asked my self if supplying

c = pm.Normal("c",15, 3)

would achieve a similar affect. Probably not to full effect and highly dependant on how the model evolves. Differences from a “regular noise” are: this affects your likelihood by multiplying the prior by the probability of sampled c given this distribution. Moreover, given your model, it is possible that the posterior for c could be something quite different which means that eventually then your model will be sampling from that posterior rather than something that looks like this normal.

I think one way to achieve this would be the suggestion here by @ricardoV94 where you define the step function of your pm.Normal manually (in the link below this is shown for Dirichlet):

Prevent prior from updating?

Currently the only way to achieve this is to assign said variable to a custom sampler that provides proposals from the prior. Here is one recent example that should be similar (in that case it provided proposals from the conjugate posterior, but it’s otherwise similar): Incorrect step assignments with custom step function - #5 by ricardoV94

What I am interested in trying our but never did is whether one can use random number generator of pytensor for this purpose:

from pytensor.tensor.random.utils import RandomStream

rng = RandomStream()
sample = rng.normal(0, 1, size=(2, 2))

It seems it used to work when pymc was with Theano via

srng = tt.shared_randomstreams.RandomStreams(seed=234)
mu = pm.Deterministic('mu', srng.normal(avg=4.9, std=.1))

but never had the chance to try this. See also, again with Theano,

Frozen or non-adaptive distribution for random variable

Hello, I recently purchased Allen Downey’s Think Bayes - I’m working through Chapter 20 - Counting Cells (the chapter can be opened in Google Colab here). There is a question I’m pondering that affects some of my models. It seems that we do not want all random variables to be adaptive (i.e., have changeable parameters). For example, in the snippet below, we want yeast_conc to be adaptive because that is what we are inferring. But we do not want shaker1_vol to be adaptive because it is measu…

Also one has got to ask then if you use such a non-conventional addition to your model, how to evaluate such stuff as convergence statistics, r values etc etc. They might come out with unexpected values and then would that mean your model is not converging correctly or is it because you are stepping out side of the known territory of Bayesian modelling? Anyone?

[iavicenna]

I just tried the suggestion above it gives an error:

ValueError: Random variables detected in the logp graph: {normal_rv{0, (0, 0), floatX, False}.out}.
This can happen when DensityDist logp or Interval transform functions reference nonlocal variables,
or when not all rvs have a corresponding value variable.

But would be better if @ricardoV94 confirms that his suggestions is still the only way and random streams cant be used for this purpose.

[mjedrz]

I got the same error, and after commenting it out deep in the Pymc source code, there started appearing divergences when running NUTS. It could be the case that NUTS, unlike SMC-ABC is fundamentally unable to include such random correlated noise. Could someone confirm?
It would be a shame because adding correlated noise is important in two cases: when observed data has correlated uncertainties and when some of the mathematical model’s input parameters are not designated to have their priors updated. SMC-ABC works but is much more computationally heavy than NUTS.

[iavicenna]

Defining your own sampler is still on the table though. When I had to do a very similar thing but adding multinomial noise to my model it looked like this:

from pymc.step_methods.arraystep import BlockedStep

#this is required specific to my case, you can skip it
def simplex_backward(value):
    value = np.concatenate([value, -np.sum(value, -1, keepdims=True)], axis=-1)
    exp_value_max = np.exp(value - np.max(value, -1, keepdims=True))
    return exp_value_max / np.sum(exp_value_max, -1, keepdims=True)

class MultinomialNoiseStep(BlockedStep):
    def __init__(self, var, N, props, size):
        model = pm.modelcontext(None)
        value_var = model.rvs_to_values[var]
        self.vars = [value_var]
        self.name = value_var.name
        self.N = N
        self.size = size
        self.p_prior_name = model.rvs_to_values[props].name

    def step(self, point: dict):
        p_posterior = simplex_backward(point[self.p_prior_name])
        draw = np.random.multinomial(self.N, p_posterior, size=self.size)
        point[self.name] = draw
        return point, []

and inside the model I would use it as such

with pm.Model() as model:
   #some priors and constants including props, frac0, N0 etc
   
   #here is the part where you first define your noise as a pymc distribution
   multinomial_noise =  pm.Multinomial("multinomial_noise", N, props,
                                                               shape=(nrepeats,ncategories))

   # but change its step so that it is just sampled
   steps = [MultinomialNoiseStep(multinomial_noise,  N, props, size)]

   # you can then add multinomial_noise to other variables etc
   # add the steps to the sampler via
   trace = pm.sample(**sample_parameters, step=steps)

Note that simplex_backward is simply the inverse of a transformation that pymc uses internally. I suspect you would not need that in your case since mu and sd are fixed. If you use sd and mu not fixed, you would need to retrieve it as such in init

mu_prior_name = model.rvs_to_values[mu].name

and if these variables are internally transformed (I think mu is not and sd might be logged)
inside step you need to do

mu_posterior = inverse_trans(point[self.mu_prior_name])

if not transformed you can just use

mu_posterior = point[self.mu_prior_name]

But as I said, if you are supplying these as constants, you don’t need to deal with these. I think something like

class NormalNoiseStep(BlockedStep):
    def __init__(self, var, mu, sd, size):
        model = pm.modelcontext(None)
        value_var = model.rvs_to_values[var]
        self.vars = [value_var]
        self.name = value_var.name
        self.mu = mu
        self.size = size
        self.sd = sd

    def step(self, point: dict):

        draw = np.random.normal(self.mu, self.sd, size=self.size)
        point[self.name] = draw
        return point, []

might be enough. Also perhaps better to use rng instead of np.random.

[mjedrz]

Thank you! I will test your suggested code and let you know if it works in a day or two.

[mjedrz]

Unfortunately, it didn’t work for me. The console printed information that something called “CompoundStep” was taking place, and from this point, Python appeared to freeze.
I was able to run DEMetropolis after disabling the “ValueError” you mentioned in the source code, but it doesn’t work as well as SMC-ABC. It produces way larger posterior uncertainties for the parameters “a” and “b”. I suppose that I will have to create or find another sampler because SMC-ABC feels like an overkill for cases when I know the posterior distributions are not multimodal.

[iavicenna]

Which version of pymc are you on? There were some fixes in the latest versions that fixed some sampling issues with Multivariate variables. I don’t know if that applies to you. The code where I used with Multinomial noise sampled fine, did not try it yet with a normal one though.

[mjedrz]

The problem turned out to be that my computer was having trouble when the number of cores was not specified to be 1. I managed to obtain the results I needed with the help of the code you provided. Apart from the issue with the core, I was wrong to use Y_obs = Normal(“Y_obs”, mu=mu, sigma=0.1, observed=data). What I should have been doing was using Y_obs = MvNormal(“Y_obs”, mu=mu, cov=cov, observed=data), where cov is created by propagating uncertainties from “c” to observed data (and including the uncertainty of observed data equal to 0.1 for each item in data). In the end, this allowed for obtaining posterior results that were exactly the same as for SMC-ABC (It should also be noted here, that true values of a, b and c are 43, 36 and 12 for this example). Here is the final code for NUTS (but it also works for Metropolis):

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

rng = np.random.default_rng()
class NormalNoiseStep(BlockedStep):
    def __init__(self, var, mu, sd, size):
        self.model = pm.modelcontext(None)
        self.value_var = self.model.rvs_to_values[var]
        self.vars = [self.value_var]
        self.name = self.value_var.name
        self.mu = mu
        self.sd = sd
        self.size = size

    def step(self, point: dict):
        draw = rng.normal(self.mu, self.sd, size=self.size)
        point[self.name] = draw
        return point, []

data = [151.2, 181.9, 219.7]
cov = np.array([[92.18, 61.41, 30.72],
               [61.41, 40.92, 20.46],
               [30.72, 20.46, 10.25]])

with pm.Model() as basic_model:
    # Priors for unknown model parameters
    a = pm.Normal("a", mu=50, sigma=5)
    b = pm.Normal("b", mu=40, sigma=7)
    # Initialize 'c' with a placeholder value. It will be updated by our custom step method.
    c = pm.Normal("c", mu=0, sigma=1, shape=1)  # Placeholder distribution

    mu = [(a + 2*b + 3*c[0]), (2*a + 2*b + 2*c[0]), (4*a + b + c[0])]

    Y_obs = pm.MvNormal("Y_obs", mu=mu, cov=cov, observed=data)

    steps = [NormalNoiseStep(c, 15, 3, size=1), pm.NUTS(vars=[a, b])]
    trace = pm.sample(tune=1000, draws=15000, step=steps, cores=1, chains=4)

print(az.summary(trace, kind="stats"))
az.plot_trace(trace);

For SMC-ABC the following code should produce the same posterior results (I didn’t check if posterior covariances between a and b are also the same but I’m guessing they should be):

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pymc as pm

az.style.use("arviz-darkgrid")

data = [151.2, 181.9, 219.7]
rng = np.random.default_rng()

def normal_sim(rng, a, b, size=1):
    c=rng.normal(15, 3, size=1)
    x1=a+2*b+3*c
    x2=2*a+2*b+2*c
    x3=4*a+b+c
    return [x1, x2, x3]

with pm.Model() as example:
    a = pm.Normal("a", mu=50, sigma=5)
    b = pm.Normal("b", mu=40, sigma=7)
    s = pm.Simulator("s", normal_sim, params=(a, b), sum_stat="sort", epsilon=0.1, observed=data)

    idata = pm.sample_smc(40000, cores=12, chains=1)
    idata.extend(pm.sample_posterior_predictive(idata))

az.plot_trace(idata, kind="rank_vlines");
print(az.summary(idata, kind="stats"))