Writing reusable model components.

[LeeviLindgren]

Hi!

I would like to write separate Python functions for my model components for reusability and modularity. Here is a simplified example on what I tried to do earlier:

def gp(x, length_scale, m=25):
    cov_func = pm.gp.cov.ExpQuad(1, ls=length_scale, active_dims=[0])
    gp = pm.gp.HSGP(m=[m], c=4.0, cov_func=cov_func)
    f = gp.prior("f", X=x)
    return f

with pm.Model() as model:
    t = pm.Data("t", time)
    slow_trend = gp(time, 10.0)
    # Does not work due to overlapping parameter names, "f", in this case
    fast_trend = gp(time, 1.0)
    sigma = pm.Exponential("sigma", lam=1.0)
    pm.Normal("y", slow_trend + fast_trend, sigma)

Numpyro supports the above with the scope handler, which allows users to prepend a prefix to the names of variables created within a specific scope.

How could I address this in PyMC?

[ricardoV94]

You can create a nested named model, see the fourth code example in the docs: https://www.pymc.io/projects/docs/en/stable/api/model/generated/pymc.model.core.Model.html

import pymc as pm

with pm.Model(name="root") as root:
    x = pm.Normal("x")  # Variable wil be named "root::x"

    with pm.Model(name='first') as first:
        # Variable will belong to root and first
        y = pm.Normal("y", mu=x)  # Variable wil be named "root::first::y"

    # Can pass parent model explicitly
    with pm.Model(name='second', model=root) as second:
        # Variable will belong to root and second
        z = pm.Normal("z", mu=y)  # Variable wil be named "root::second::z"

    # Set None for standalone model
    with pm.Model(name="third", model=None) as third:
        # Variable will belong to third only
        w = pm.Normal("w")  # Variable wil be named "third::w"

[LeeviLindgren]

Thanks @ricardoV94!

This indeed works:

def gp(x, length_scale, m=25):
    cov_func = pm.gp.cov.ExpQuad(1, ls=length_scale, active_dims=[0])
    gp = pm.gp.HSGP(m=[m], c=4.0, cov_func=cov_func)
    f = gp.prior("f", X=x)
    return f

with pm.Model() as model:
    t = pm.Data("t", time)
    with pm.Model(name="slow_trend"):
        slow_trend = gp(time, 10.0)
    with pm.Model(name="fast_trend"):
        fast_trend = gp(time, 0.5)
    sigma = pm.Exponential("sigma", lam=1.0)
    pm.Normal("y", slow_trend + fast_trend, sigma)