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

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.

1 Like