Evaluating parts of a model after sampling

In Dan-Foreman Mackey’s recent modeling package built on top of PyMC3 there’s a helpful function called
eval_in_model for evaluating variables in a PyMC3 model given some values of the parameters:
https://exoplanet.dfm.io/en/latest/tutorials/pymc3-extras/#evaluating-model-components-for-specific-samples

I would like to have the model in a separate function or class and be able to evaluate it after sampling but I’ m not sure how to do that. I modified Dan’s example to

def model():
    with pm.Model() as model:
        logs = pm.Normal("logs", mu=-3.0, sd=1.0)
        a0 = pm.Normal("a0")
        a1 = pm.Normal("a1")
        a2 = pm.Normal("a2")
        mod = a0 + a1 * x + a2 * x**2

        # Sample from the prior
        prior_sample = pm.sample_prior_predictive(samples=1)
        y = xo.eval_in_model(mod, prior_sample)
        y += np.exp(prior_sample["logs"]) * np.random.randn(len(y))

        # Add the likelihood
        pm.Normal("obs", mu=mod, sd=pm.math.exp(logs), observed=y)

    return model

with model():
    # Fit the data
    map_soln = pm.find_MAP()
    trace = pm.sample()

x_grid = np.linspace(-1.1, 1.1, 5000)
with model():
    pred = xo.eval_in_model(a0 + a1 * x_grid + a2 * x_grid**2, map_soln)

but I’m getting an error:

NameError: name 'a0' is not defined

What am I doing wrong?

The reason you are getting the NameError is that you are defining the variable a0 inside the scope of the function you called model. Outside said function, you have not defined any variable named a0, and that causes the error.

I have to point out that your approach has a really crucial problem, each time you call your function model(), you get a different Model instance. This means that the tensor variables named a0 in the first model call will be different from the tensor named a0 on the second model call. This means that if you somehow got the a0, a1 and a2 tensors defined from your first call, and ran the second part as

with model():
    pred = xo.eval_in_model(a0 + a1 * x_grid + a2 * x_grid**2, map_soln)

you would not be able to resolve the inputs of the output tensor.

I strongly recommend that you first change the name of the function model to something that clashes less, like model_factory_function, and to replace the second part of your code with this:

with model() as model_instance:
    # Fit the data
    map_soln = pm.find_MAP()
    trace = pm.sample()

x_grid = np.linspace(-1.1, 1.1, 5000)
with model_instance:
    a0 = model_instance['a0']
    a1 = model_instance['a1']
    a2 = model_instance['a2']
    pred = xo.eval_in_model(a0 + a1 * x_grid + a2 * x_grid**2, map_soln)

As a final comment you can replace this:

prior_sample = pm.sample_prior_predictive(samples=1)
y = xo.eval_in_model(mod, prior_sample)

with this:

y = pm.distributions.distribution.draw_values([mod], size=1)[0]

draw_values is the function that does all of the sample_prior_predictive logic, it compiles the theano function if it needs to, it calls the random method of random variables that the tensor conditionally depends on if needed in the middle. If you also want to know the other variable’s value’s do this:

a0_val, a1_val, a2_val, y = pm.distributions.distribution.draw_values([a0, a1, a2, mod], size=1)
4 Likes

Thanks a lot for the the detailed answer! I didn’t really understand the with keyword properly.

Do you mean that you don’t really understand context managers in general or what the with statement did in the code snippet?

Context managers in general, it’s clear to me now though!