MaskedPrior#

class pymc_marketing.special_priors.MaskedPrior(prior, mask, active_dim=None)[source]#

Create variables from a prior over only the active entries of a boolean mask.

Warning

This class is experimental and its API may change in future versions.

Parameters:
priorPrior

Base prior whose variable is defined over prior.dims. Internally, the variable is created only for the active entries given by mask and then expanded back to the full shape with zeros at inactive positions.

maskxarray.DataArray

Boolean array with the same dims and shape as prior.dims marking active (True) and inactive (False) entries.

active_dimstr, optional

Name of the coordinate indexing the active subset. If not provided, a name is generated as "non_null_dims:<dim1>_<dim2>_...". If an existing coordinate with the same name has a different length, a suffix with the active length is appended.

Examples

Simple 1D masking.

import numpy as np
import xarray as xr
import pymc as pm
from pymc_extras.prior import Prior
from pymc_marketing.special_priors import MaskedPrior

coords = {"country": ["Venezuela", "Colombia"]}
mask = xr.DataArray(
    [True, False],
    dims=["country"],
    coords={"country": coords["country"]},
)
intercept = Prior("Normal", mu=0, sigma=10, dims=("country",))
with pm.Model(coords=coords):
    masked = MaskedPrior(intercept, mask)
    intercept_full = masked.create_variable("intercept")

Nested parameter priors with dims remapped to the active subset.

import numpy as np
import xarray as xr
import pymc as pm
from pymc_extras.prior import Prior
from pymc_marketing.special_priors import MaskedPrior

coords = {"country": ["Venezuela", "Colombia"]}
mask = xr.DataArray(
    [True, False],
    dims=["country"],
    coords={"country": coords["country"]},
)
intercept = Prior(
    "Normal",
    mu=Prior("HalfNormal", sigma=1, dims=("country",)),
    sigma=10,
    dims=("country",),
)
with pm.Model(coords=coords):
    masked = MaskedPrior(intercept, mask)
    intercept_full = masked.create_variable("intercept")

All entries masked (returns deterministic zeros with original dims).

import numpy as np
import xarray as xr
import pymc as pm
from pymc_extras.prior import Prior
from pymc_marketing.special_priors import MaskedPrior

coords = {"country": ["Venezuela", "Colombia"]}
mask = xr.DataArray(
    [False, False],
    dims=["country"],
    coords={"country": coords["country"]},
)
prior = Prior("Normal", mu=0, sigma=10, dims=("country",))
with pm.Model(coords=coords):
    masked = MaskedPrior(prior, mask)
    zeros = masked.create_variable("intercept")

Apply over a saturation function priors:

from pymc_marketing.mmm import LogisticSaturation
from pymc_marketing.special_priors import MaskedPrior

coords = {
    "country": ["Colombia", "Venezuela"],
    "channel": ["x1", "x2", "x3", "x4"],
}

mask_excluded_x4_colombia = xr.DataArray(
    [[True, False, True, False], [True, True, True, True]],
    dims=["country", "channel"],
    coords=coords,
)

saturation = LogisticSaturation(
    priors={
        "lam": MaskedPrior(
            Prior(
                "Gamma",
                mu=2,
                sigma=0.5,
                dims=("country", "channel"),
            ),
            mask=mask_excluded_x4_colombia,
        ),
        "beta": Prior(
            "Gamma",
            mu=3,
            sigma=0.5,
            dims=("country", "channel"),
        ),
    }
)

prior = saturation.sample_prior(coords=coords, random_seed=10)
curve = saturation.sample_curve(prior)
saturation.plot_curve(
    curve,
    subplot_kwargs={
        "ncols": 4,
        "figsize": (12, 18),
    },
)

Masked likelihood over an arbitrary subset of entries (2D example over (date, country)):

import numpy as np
import xarray as xr
import pymc as pm
from pymc_extras.prior import Prior
from pymc_marketing.special_priors import MaskedPrior

coords = {
    "date": np.array(["2021-01-01", "2021-01-02"], dtype="datetime64[D]"),
    "country": ["Venezuela", "Colombia"],
}

mask = xr.DataArray(
    [[True, False], [True, False]],
    dims=["date", "country"],
    coords={"date": coords["date"], "country": coords["country"]},
)

intercept = Prior("Normal", mu=0, sigma=10, dims=("country",))
likelihood = Prior(
    "Normal", sigma=Prior("HalfNormal", sigma=1), dims=("date", "country")
)
observed = np.random.normal(0, 1, size=(2, 2))

with pm.Model(coords=coords):
    mu = intercept.create_variable("intercept")
    masked = MaskedPrior(likelihood, mask)
    y = masked.create_likelihood_variable("y", mu=mu, observed=observed)

Methods

MaskedPrior.__init__(prior, mask[, active_dim])

MaskedPrior.create_likelihood_variable(name, ...)

Create an observed variable over the active subset and expand to full dims.

MaskedPrior.create_variable(name)

Create a deterministic variable with full dims using the active subset.

MaskedPrior.from_dict(data)

Deserialize MaskedPrior from dictionary created by to_dict.

MaskedPrior.to_dict()

Serialize MaskedPrior to a JSON-serializable dictionary.