Dealing with multiple closely related models

Questions version agnostic
1

What is a good way to deal with defining multiple closely related PYMC models while avoiding duplicate code? For example , suppose I want to compare a model that uses Student_T versus Normal likelihood, and also I want to compare including certain features or not, etc.

My current approach is to use a function that creates the model, with a bunch of flags and enums… for example something like this :slight_smile: but worse:

def build_model(data, mode = "condition", baseline = False, skew = False)
    ... 
    # set up the data (centering etc), define coordinates etc.
    with pm.Model(coords = coords): 
             ... 
             if not baseline:
                    # add priors for more features

             if skew:
                     alpha = pm.Normal("alpha", mu = 0, sigma = 4)
                     pm.SkewNormal("y_obs", alpha = alpha, mu = mu, sigma = sigma, observed = y)
              else: 
                     pm.Normal("y_obs", mu = mu, sigma = sigma, observed = y)  
     return model

I must confess I have not fully searched the examples gallery, maybe there are some tips in there.

2

You may want different functions that build distinct parts of the model, but in the end it will look something like that yes

1 Like
3

Oh! Yes this would help refactor quite a bit ! Thanks !

4

I will often split out bits of my models into little functions for organization. A nice helper function is pm.modelcontext for this. You can make a function like this:

def add_skew_likelihood(mu, sigma, y, alpha_mu=0, alpha_sigma=4, model=None):
    model = pm.modelcontext(model)
    with model:
        alpha = pm.Normal("alpha", mu = alpha_mu, sigma = alpha_sigma)
        pm.SkewNormal("y_obs", alpha = alpha, mu = mu, sigma = sigma, observed = y)

If you pass model = None (the default), modelcontext will look for an open model context, and raise an error if it doesn’t find one. So if you call this function inside a with pm.Model(): block, alpha and y_obs will be added. Alternatively, you can pass a model to the function, and it will open it up and and add these variables. So both of these work:

model = pm.Model(coords = coords)
add_skew_likelihood(0, 1, y, model=model)

And:

with pm.Model(coords=coords) as model:
    add_skew_likelihood(0, 1, y)

If you have tons of if/else conditions, it can be nice to claw back one level of indentation using the first method. You can also mix and match; open a context and add all the “shared” stuff, then use control flow to finish off the rest.

3 Likes
5

Cool, I was not aware of this function! Thanks!

6

Is there an advantage to doing it this way rather than defining separate models that just incorporate functions like priors and likelihoods? I would think doing it this way you run into all of the usual problems of using conditionals rather than object-oriented or functional programming.

PyMC is so much more flexible than Stan for this kind of thing, I’m surprised I don’t see it more often on this list. I’d be writing my models with functions or object inheritance all over the place for things like hierarchical priors or compound likelihoods.

7

I guess it’s a question of preference. For more complex models like pymc-marketing they go that route: pymc-marketing/pymc_marketing/mmm/mmm.py at f334980ecf66c28fa796ced52e75bf5005fdceb2 · pymc-labs/pymc-marketing · GitHub

I personally feel it makes it inscrutable. For me PyMC already has the right level of abstraction, just sprinkled sometimes with a helper to make the non-centered parametrization of a hierarchical prior.

It may also have to do with how complex the models one works with are.

8

That looks like a huge model and I don’t read Python or PyMC well enough to make a lot of sense of it.

I was thinking specifically about shared and reusable model patterns. For example, I’d like to be able to write a function hierarchical_normal where I could then write

alpha = hierarchical_normal("alpha", dims=(5,), center=False)

rather than

mu_alpha = pm.Normal("mu_alpha", mu=0, sigma = 4)
sigma_alpha = pm.LogNormal("sigma_alpha", mu=0, sigma=1)
alpha_std = pm.Normal(mu=0, sigma=1, dims=(5,))
alpha = mu_alpha + sigma_alpha * alpha_std

Please forgive any PyMC or Python errors here—I hope the intent is clear. I’m not exactly sure how to write the function, but I’m imagining something like this:

def hierarchical_normal(name, dims=(1,), center=False):
    if not center:
        mu = pm.Normal("mu_alpha", mu=0, sigma = 4)
        sigma = pm.LogNormal("sigma_alpha", mu=0, sigma=1)
        alpha_std = pm.Normal("name" + "_std", mu=0, sigma=1, dims=(5,))
        alpha = mu_alpha + sigma_alpha * alpha_std
    else:
        mu = pm.Normal("mu_" + name, mu=0, sigma = 4)
        sigma = pm.LogNormal("sigma_" + name, mu=0, sigma=1)
        alpha = pm.Normal(name, mu=mu_population sigma="sigma_population")
    return alpha

I’m unsure what this would do to the with pm.Model() scoping.

I’m pretty sure I asked a question very much like this a year or so ago.

1 Like
9

This function would work as-in. Inside a with pm.Model() scope, it would automatically add all the declared RVs in the active branch to the model as expected.

While I don’t recommend it, if you really wanted to lean into all the automagic stuff that PyMC does, you could return nothing from the function and access the created RVs using the model, e.g:

with pm.Model() as m:
    hierarchical_normal('mu', dims=(1,), center=True)
    alpha = m['mu'] # it's saved as mu because that's the name in the pm.Normal

This wouldn’t work if center = False, because alpha = mu_alpha + sigma_alpha * alpha_std isn’t registered (We have pm.Deterministic to let users register arbitrary computation to a model).

Obviously it works fine if it is used as you originally intended: alpha = hierarchical_normal('mu', dims=(1,), center=False)