Multilevel Elasticities for a Single SKU - Part III.

Prepare Notebook

import arviz as az
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import numpyro
import numpyro.distributions as dist
import polars as pl
import seaborn as sns
from jax import random
from numpyro.infer import MCMC, NUTS, Predictive
from sklearn.preprocessing import LabelEncoder

az.style.use("arviz-darkgrid")
plt.rcParams["figure.figsize"] = [12, 7]
plt.rcParams["figure.dpi"] = 100
plt.rcParams["figure.facecolor"] = "white"

numpyro.set_host_device_count(n=4)

rng_key = random.PRNGKey(seed=42)

%load_ext autoreload
%autoreload 2
%config InlineBackend.figure_format = "retina"

Read Data

We read the data from cour case study. YOu can find an explicit data generating process in the first notebook Multilevel Elasticities for a Single SKU - Part I.

market_df = pl.read_csv(
    "https://raw.githubusercontent.com/juanitorduz/website_projects/master/data/multilevel_elasticities_single_sku_data.csv",
    schema_overrides={"region_id": pl.Categorical},
)

market_df.head()

Recall that our objective is to estimate the elasticity of a single SKU across different regions.

g = sns.lmplot(
    data=market_df,
    x="log_price",
    y="log_quantities",
    hue="region_id",
    height=8,
    aspect=1.2,
    scatter=False,
)
g.set_axis_labels(x_var="Log Prices", y_var="Log Quantities")
legend = g.legend
legend.set_title(title="Region ID", prop={"size": 10})
g.fig.suptitle(
    "Log Prices vs Log Quantities by Region", y=1.05, fontsize=18, fontweight="bold"
);

We expect this elasticity to be influenced by the region’s median income.

fig, ax = plt.subplots(figsize=(9, 6))
(
    market_df.group_by("region_id")
    .agg(pl.col("median_income").mean())
    .to_pandas()
    .pipe(
        (sns.barplot, "data"), x="region_id", y="median_income", hue="region_id", ax=ax
    )
)
ax.set(xlabel="Region ID", ylabel="Median Income")
ax.set_title(label="Median Income by Region", fontsize=18, fontweight="bold");

Model Specification

As described in the previous notebooks, we can use a multilevel covariance model to estimate the elasticity by constraining how the intercept and slops vary. In the first notebook we did this just for the median_income variable. Here we extend it so that we have two covariance components for the intercept and two for the slope. We do thin in the non-centered parametrization (see Statistical Rethinking in NumPyro - Chapter 14. Adventures in Covariance). Here is a mathematical description of the model:

\[\begin{align*} \log(q) & \sim \text{Normal}(\mu, \sigma ^2) \\ \mu & = \alpha + \beta \log(p) \\ \alpha & = \alpha_{\text{intercept}, \text{region}} + \alpha_{\text{slope}, \text{region}} m \\ \beta & = \beta_{\text{intercept}, \text{region}} + \beta_{\text{slope}, \text{region}} m \\ \alpha_{\text{intercept}, \text{region}} & = z'_{\text{intercept}}[0] \\ \alpha_{\text{slope}, \text{region}} & = z'_{\text{slope}}[0] \\ \beta_{\text{intercept}, \text{region}} & = z'_{\text{intercept}}[1] \\ \beta_{\text{slope}, \text{region}} & = z'_{\text{slope}}[1] \\ z'_{\text{intercept}} &= \text{diag}(\sigma_{\alpha_{\text{intercept}}}^2, \sigma_{\beta_{\text{intercept}}}^2) L_{\text{intercept}} \: z_{\text{intercept}} \\ z'_{\text{slope}} &= \text{diag}(\sigma_{\alpha_{\text{slope}}}^2, \sigma_{\beta_{\text{slope}}}^2) L_{\text{slope}} \: z_{\text{slope}} \\ L_{\text{slope}} & \sim \text{LKJCholesky}(1) \\ z_{\text{intercept}} & \sim \text{Normal}(0, 1) \\ z_{\text{slope}} & \sim \text{Normal}(0, 1) \end{align*}\]

We prepare the data (very similar to the previous notebooks) and then implement the model in NumPyro.

obs_idx = np.arange(market_df.shape[0])

# For the region_id we use a label encoder to transform the categorical
# variable into integers (in more applications these tags won't be always integers so
# it is good to get used to these encoders.
regions_encoder = LabelEncoder()
regions_encoder.fit(market_df["region_id"].to_numpy())
regions_idx = regions_encoder.transform(market_df["region_id"].to_numpy())
regions = regions_encoder.classes_

log_price = market_df["log_price"].to_jax()
log_quantities = market_df["log_quantities"].to_jax()

# This is just an auxiliary variable to indicate we are modeling
# two covariance components.
effect = ["intercept", "median_income"]

Besides the usual data preparation, we need a map from the encoded category regions to the corresponding median income values:

region_median_income_mapping_df = market_df.select(
    ["region_id", "median_income"]
).unique()

region_median_income_mapping_df = region_median_income_mapping_df.with_columns(
    region_idx=regions_encoder.transform(
        region_median_income_mapping_df["region_id"].to_numpy()
    )
).sort("region_idx")

median_income = region_median_income_mapping_df["median_income"].to_jax()

Now we implement the model:

def model(log_price, regions_idx, median_income, log_quantities=None):
    dim = 2
    effects = 2
    n_obs = log_price.size
    n_regions = np.unique(regions_idx).size

    # Prior for the degrees of freedom of the Student-T distribution.
    nu = numpyro.sample("nu", dist.Gamma(concentration=20.0, rate=3))
    # Prior for the scale parameter of the Student-T distribution.
    sigma = numpyro.sample("sigma", dist.Exponential(rate=1.0 / 0.5))

    # concentration prior for the LKJ prior (this value assumes uniform correlation),
    concentration = jnp.ones(1)

    with numpyro.plate("effect", effects, dim=-1):
        # standard deviation of the effect.
        sigma_effect = numpyro.sample(
            "sigma_effect", dist.HalfNormal(scale=0.2 * jnp.ones(dim))
        )
        # correlation matrix
        chol_corr = numpyro.sample(
            "chol_corr",
            dist.LKJCholesky(dimension=dim, concentration=concentration),
        )

        with numpyro.plate("regions", n_regions, dim=-2):
            # z is the latent variable that we will use to compute the covariance matrix
            # in the non-centered parameterization.
            z = numpyro.sample("z", dist.Normal(loc=0, scale=1))

    # We compute the covariance matrix using the cholesky decomposition
    chol_cov = z @ (sigma_effect[None, ...] * chol_corr)
    # We extract the region effects
    alpha_region_effect = chol_cov[:, :, 0]
    beta_region_effect = chol_cov[:, :, 1]
    # Global intercept component (intercept and median income term)
    alpha_region_intercept = alpha_region_effect[0, :]
    alpha_region_median_income = alpha_region_effect[1, :]
    # Global slope factor of `log_price` (intercept and median income term)
    beta_region_intercept = beta_region_effect[0, :]
    beta_region_median_income = beta_region_effect[1, :]

    # Region specific intercept and slope
    alpha_region = numpyro.deterministic(
        "alpha_region",
        alpha_region_intercept + alpha_region_median_income * median_income,
    )
    beta_region = numpyro.deterministic(
        "beta_region",
        beta_region_intercept + beta_region_median_income * median_income,
    )

    # Likelihood mean
    mu = alpha_region[regions_idx] + beta_region[regions_idx] * log_price

    with numpyro.plate("data", n_obs):
        # We use a Student-T likelihood to model the residuals.
        numpyro.sample(
            "log_quantities",
            dist.StudentT(df=nu, loc=mu, scale=sigma),
            obs=log_quantities,
        )

Let’s visualize the model:

numpyro.render_model(
    model=model,
    model_kwargs={
        "log_price": log_price,
        "regions_idx": regions_idx,
        "median_income": median_income,
        "log_quantities": log_quantities,
    },
    render_distributions=True,
    render_params=True,
)

Inference

We now use NUTS to sample from the posterior distribution of the model.

sampler = NUTS(model)
mcmc = MCMC(
    sampler=sampler,
    num_warmup=1_000,
    num_samples=4_000,
    num_chains=4,
)

rng_key, rng_subkey = random.split(rng_key)
mcmc.run(rng_subkey, log_price, regions_idx, median_income, log_quantities)

We collect the results in an az.InferenceData object:

idata = az.from_numpyro(
    mcmc,
    coords={"effect": effect, "regions": regions},
    dims={
        "sigma_effect": ["effect"],
        "z": ["regions", "effect"],
        "alpha_region": ["regions"],
        "beta_region": ["regions"],
    },
)

Model Diagnostics

# Count divergences
idata.sample_stats.diverging.sum().item()

0

az.summary(idata, var_names=["~z", "~chol_corr"])

We do not have divergences and the r-hat values look good!

axes = az.plot_trace(
    data=idata,
    var_names=["~z", "~chol_corr"],
    compact=True,
    kind="rank_bars",
    backend_kwargs={"figsize": (15, 9), "layout": "constrained"},
)
plt.gcf().suptitle("Covariance Model", fontsize=20, fontweight="bold", y=1.05);

Intercepts and Elasticities

Let’s look into the estimates for the intercepts and elasticities:

ax, *_ = az.plot_forest(
    data=idata,
    var_names=["alpha_region"],
    combined=True,
    figsize=(10, 6),
)
ax.set_title(
    label="Posterior Distribution of the Intercepts",
    fontsize=18,
    fontweight="bold",
);

ax, *_ = az.plot_forest(
    data=idata,
    var_names=["beta_region"],
    combined=True,
    figsize=(10, 6),
)
ax.set_title(
    label="Posterior Distribution of the Elasticities",
    fontsize=18,
    fontweight="bold",
);

These results are very similar to the ones we obtained in the previous notebooks!

Posterior Predictive Distribution

Finally, we compute the posterior predictive distribution and compare it to the observed data.

predictive = Predictive(
    model=model,
    posterior_samples=mcmc.get_samples(),
    return_sites=["log_quantities"],
)

rng_key, rng_subkey = random.split(rng_key)
predictions = predictive(rng_subkey, log_price, regions_idx, median_income)

idata.extend(
    az.from_numpyro(
        posterior_predictive=predictions,
        coords={"obs_idx": obs_idx},
        dims={"log_quantities": ["obs_idx"]},
    )
)

fig, ax = plt.subplots()
az.plot_ppc(idata, kind="kde", num_pp_samples=1_000, ax=ax)
ax.set_title(label="Posterior Predictive", fontsize=18, fontweight="bold");

The model seems to be capturing most of the variation in the data.