One family for every response curve: a Weibull transform in JAX for media models

import logging
import time
import warnings

import arviz as az
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import optax
import pandas as pd
import pymc as pm
import pytensor
import pytensor.tensor as pt
from IPython.display import display
from matplotlib.colors import ListedColormap
from matplotlib.ticker import FixedLocator, NullLocator

from pymc_challenges import data

jax.config.update("jax_enable_x64", True)
RANDOM_SEED = 45
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-variat")
logging.getLogger("pymc").setLevel(logging.WARNING)  # several fits: no sampler banner per fit
# Expected and harmless: Numba runs the JAX node in object mode (section 9 measures the cost)
warnings.filterwarnings("ignore", message="Numba will use object mode")
warnings.filterwarnings("ignore", message="Skipping Check")  # JAX drops PyMC's parameter checks
BLUE, ORANGE, AQUA, GREY, PURPLE, RED = "#2a78d6", "#eb6834", "#1baf7a", "#8a8a86", "#8c5ac8", "#c8384e"
print(f"PyMC {pm.__version__}, ArviZ {az.__version__}, PyTensor {pytensor.__version__}, "
      f"JAX {jax.__version__}, optax {optax.__version__}")

X_PLOT = np.linspace(0, 3, 301)


def hill_np(x, K, n):
    return x**n / (K**n + x**n)


def logistic_np(x, lam):
    return (1 - np.exp(-lam * x)) / (1 + np.exp(-lam * x))


standard = {
    "Hill, S-shaped (K=1.2, n=3)": lambda x: hill_np(x, 1.2, 3.0),
    "Hill, concave (K=1, n=1)": lambda x: hill_np(x, 1.0, 1.0),
    "logistic saturation (λ=1)": lambda x: logistic_np(x, 1.0),
    "power, concave (x^0.5)": lambda x: np.sqrt(x),
    "power, convex (x^1.5)": lambda x: x**1.5,
    "log(1 + 2x)": lambda x: np.log1p(2 * x),
}
fig, axes = plt.subplots(1, 2, figsize=(12, 4.2))
colors = [BLUE, ORANGE, AQUA, PURPLE, RED, GREY]
slopes = {}
for (name, f), col in zip(standard.items(), colors):
    y_ = f(X_PLOT) / f(1.0)
    axes[0].plot(X_PLOT, y_, color=col, lw=2, label=name)
    slopes[name] = (f(1.001) - f(0.999)) / 0.002 / f(1.0)
axes[0].axvline(1, color=GREY, lw=0.8, ls=":")
axes[0].set(xlabel="spend (1 = reference week)", ylabel="response (1 at the reference)",
            title="Same response at the reference spend...", ylim=(0, 3.2))
axes[0].legend(fontsize=8)
names = list(slopes)
axes[1].barh(names, [slopes[n] for n in names], color=colors)
axes[1].axvline(1, color=GREY, lw=0.8, ls=":")
axes[1].invert_yaxis()
axes[1].set(xlabel="marginal return at the reference (1 = the average return)",
            title="...but marginal returns that differ by a factor of 4")

X_RANGE = 2.5  # the spend range we care about, in units of an average week


def weibull_np(x, k, c):
    """Reference NumPy version, for plotting only (section 3 writes the careful one)."""
    c = np.maximum(c, 1e-12)
    return -np.expm1(-c * x**k) / -np.expm1(-c)


def inflection(k, c):
    """x* = ((k - 1) / (c k))^(1/k) for k > 1; inf where there is none."""
    with np.errstate(invalid="ignore", divide="ignore", over="ignore"):
        return np.where(k > 1, ((k - 1) / np.maximum(c * k, 1e-300)) ** (1 / k), np.inf)


def shape_class(k, c, x_max=X_RANGE, tol=0.03):
    """0 concave, 1 ~linear, 2 convex, 3 S-shaped (inflection inside (0, x_max))."""
    k, c = np.broadcast_arrays(np.asarray(k, float), np.asarray(c, float))
    xs = np.linspace(0, x_max, 101)
    curves = weibull_np(xs[:, None], k.ravel()[None], c.ravel()[None])  # (x, points)
    near_linear = (np.abs(curves - xs[:, None]).max(0) < tol * x_max).reshape(k.shape)
    x_star = inflection(k, c)
    out = np.where(k <= 1, 0, np.where(x_star < x_max, 3, 2))
    return np.where(near_linear, 1, out)


SHAPE_NAMES = ["concave", "≈ linear", "convex", "S-shaped"]
SHAPE_CMAP = ListedColormap(["#cfe0f5", "#e6e6e3", "#f7d5c6", "#d5ecdf"])
kg = np.exp(np.linspace(np.log(0.3), np.log(5), 240))
cg = np.exp(np.linspace(np.log(1e-3), np.log(20), 240))
KG, CG = np.meshgrid(kg, cg, indexing="xy")
SHAPE_GRID = shape_class(KG, CG)


def draw_shape_map(ax, labels=True):
    ax.pcolormesh(kg, cg, SHAPE_GRID, cmap=SHAPE_CMAP, vmin=-0.5, vmax=3.5, shading="auto",
                  rasterized=True)
    ax.set(xscale="log", yscale="log", xlabel="shape k", ylabel="saturation c")
    ax.xaxis.set_major_locator(FixedLocator([0.3, 0.5, 1, 2, 3, 5]))
    ax.xaxis.set_minor_locator(NullLocator())
    ax.set_xticklabels(["0.3", "0.5", "1", "2", "3", "5"])
    if labels:
        for (kx, cy), name in zip([(0.45, 0.3), (1.0, 0.0022), (3.2, 0.0025), (3.0, 2.0)], SHAPE_NAMES):
            ax.text(kx, cy, name, ha="center", fontsize=10, color="#333333", weight="bold")


examples = {"A": (0.5, 0.002), "B": (0.6, 1.5), "C": (1.0, 0.8), "D": (1.8, 0.005),
            "E": (2.5, 0.4), "F": (3.5, 3.0)}
fig, axes = plt.subplots(1, 2, figsize=(12, 4.8), gridspec_kw={"width_ratios": [1.1, 1]})
draw_shape_map(axes[0])
for lab, (k_, c_) in examples.items():
    axes[0].plot(k_, c_, "o", color="k", ms=5)
    axes[0].annotate(lab, (k_, c_), xytext=(4, 4), textcoords="offset points", weight="bold")
axes[0].set_title(f"The shape map: what the curve looks like on 0 ≤ x ≤ {X_RANGE}")
for (lab, (k_, c_)), col in zip(examples.items(), [PURPLE, BLUE, AQUA, RED, ORANGE, GREY]):
    cls = SHAPE_NAMES[int(shape_class(k_, c_))]
    axes[1].plot(X_PLOT, weibull_np(X_PLOT, k_, c_), color=col, lw=2,
                 label=f"{lab}: k={k_}, c={c_} ({cls})")
axes[1].axvspan(X_RANGE, 3, color=GREY, alpha=0.12)
axes[1].set(xlabel="spend x", ylabel="f(x)", ylim=(0, 3.2), title="Curves at the marked points")
axes[1].legend(fontsize=8)

SERIES_BELOW = 1e-4


def phi(u):
    """(1 - exp(-u)) / u for u >= 0, with phi(0) = 1; stable at all u, including under grad."""
    small = u < SERIES_BELOW
    u_safe = jnp.where(small, 1.0, u)  # double-where: the untaken branch never sees u = 0
    return jnp.where(small, 1 - u / 2 + u**2 / 6 - u**3 / 24, -jnp.expm1(-u_safe) / u_safe)


def safe_pow(x, k):
    """x**k for x >= 0 with finite gradients in x and k at x = 0 (value and gradients 0 there)."""
    pos = x > 0
    x_safe = jnp.where(pos, x, 1.0)
    return jnp.where(pos, jnp.exp(k * jnp.log(x_safe)), 0.0)


def weibull(x, k, c):
    """Weibull response f(x) = (1 - exp(-c x^k)) / (1 - exp(-c)), exact at c = 0 (f = x^k)."""
    z = safe_pow(x, k)
    return z * phi(c * z) / phi(c)


def geometric_adstock(x, theta):
    """x (week, channel), theta (channel,) -> a_t = theta a_{t-1} + (1 - theta) x_t.

    The (1 - theta) weight keeps the long-run level of a constant spend unchanged, so x = 1 still
    means "an average week". The recursion starts from the mean of the first 8 weeks."""
    def step(a, x_t):
        a = theta * a + (1 - theta) * x_t
        return a, a

    _, out = jax.lax.scan(step, x[:8].mean(axis=0), x)
    return out


def media_transform(x, theta, k, c):
    """Spend (week, channel) -> adstocked, saturated response curves (week, channel)."""
    return weibull(geometric_adstock(x, theta), k, c)


weibull_jit = jax.jit(weibull)
print("c = 0 now gives x^k:          ", weibull_jit(xs, 1.5, 0.0), "vs", xs**1.5)
print("d/dk at x = 0:                ", jax.grad(weibull, argnums=1)(0.0, 1.5, 0.5))
print("d/dx at x = 0, k = 0.5:       ", jax.grad(weibull, argnums=0)(0.0, 0.5, 0.5))

# 1. limit and continuity
for c_ in (1e-2, 1e-4, 1e-8, 0.0):
    err = float(jnp.max(jnp.abs(weibull(xs, 1.5, c_) - xs**1.5)))
    print(f"c = {c_:.0e}: max |f - x^k| = {err:.1e}")
u_edge = SERIES_BELOW
jump = abs(float(phi(u_edge * (1 - 1e-12)) - phi(u_edge * (1 + 1e-12))))
print(f"jump in phi across the series switch: {jump:.1e}")

# 2. gradients vs finite differences, both sides of the switch
test_points = [(0.7, 1.3, 0.0), (0.7, 1.3, 0.99e-4), (0.7, 1.3, 1.01e-4), (1.8, 0.6, 2.0),
               (2.4, 3.5, 5.0), (0.0, 0.8, 1.0)]
worst = 0.0
for x_, k_, c_ in test_points:
    g = jax.grad(weibull, argnums=(1, 2))(x_, k_, c_)
    h = 1e-6
    fd_k = (weibull(x_, k_ + h, c_) - weibull(x_, k_ - h, c_)) / (2 * h)
    fd_c = (weibull(x_, k_, c_ + h) - weibull(x_, k_, max(c_ - h, 0.0))) / (h + min(h, c_))
    worst = max(worst, abs(float(g[0] - fd_k)), abs(float(g[1] - fd_c)))
print(f"worst |autodiff - finite difference| over {len(test_points)} points: {worst:.1e}")

# 3. float32 vs float64
xg64 = jnp.linspace(0, 3, 50)
with jax.enable_x64(False):
    f32 = np.asarray(weibull(jnp.asarray(xg64, jnp.float32), jnp.float32(2.2), jnp.float32(3e-5)))
f64 = np.asarray(weibull(xg64, 2.2, 3e-5))
print(f"float32 vs float64, c = 3e-5: max relative difference {np.max(np.abs(f32[1:] / f64[1:] - 1)):.1e}")

# 4. inflection point
k_, c_ = 2.5, 0.4
d2 = jax.vmap(jax.grad(jax.grad(weibull)), in_axes=(0, None, None))
xfine = jnp.linspace(0.05, 3, 20001)
x_sign_change = float(xfine[jnp.argmax(d2(xfine, k_, c_) < 0)])
print(f"inflection: second derivative changes sign at x = {x_sign_change:.4f}, "
      f"closed form x* = {((k_ - 1) / (c_ * k_)) ** (1 / k_):.4f}")

X_FIT = jnp.linspace(0, 3, 121)
STARTS = jnp.array([[0.0, 0.0, -1.0], [0.0, 0.7, 0.0], [0.0, -0.5, 1.0]])
adam = optax.adam(0.05)


def weibull_curve(p, x=X_FIT):
    return jnp.exp(p[0]) * weibull(x, jnp.exp(p[1]), jnp.exp(p[2]))


def fit_from(target, p0, steps=3000):
    loss = lambda p: jnp.mean((weibull_curve(p) - target) ** 2)  # noqa: E731

    def step(carry, _):
        p, state = carry
        updates, state = adam.update(jax.grad(loss)(p), state, p)
        return (optax.apply_updates(p, updates), state), None

    (p, _), _ = jax.lax.scan(step, (p0, adam.init(p0)), None, length=steps)
    return p, jnp.max(jnp.abs(weibull_curve(p) - target)) / jnp.max(target)


def fit_best(target):
    ps, errs = jax.vmap(fit_from, in_axes=(None, 0))(target, STARTS)
    best = jnp.argmin(errs)
    return ps[best], errs[best]


fit_many = jax.jit(jax.vmap(fit_best))  # (n_curves, 121) -> params (n_curves, 3), errors (n_curves,)

K_GRID = np.linspace(0.3, 3.0, 28)
N_GRID = np.linspace(0.5, 4.0, 29)
KK, NN = np.meshgrid(K_GRID, N_GRID, indexing="ij")
hill_targets = jnp.asarray(hill_np(np.asarray(X_FIT)[None], KK.ravel()[:, None], NN.ravel()[:, None]))
t0 = time.perf_counter()
hill_params, hill_err = jax.block_until_ready(fit_many(hill_targets))
print(f"{len(hill_targets)} Hill curves x 3 starts x 3000 Adam steps: {time.perf_counter() - t0:.1f} s")
HILL_ERR = np.asarray(hill_err).reshape(KK.shape)
print(f"max error over the Hill grid: median {100 * np.median(HILL_ERR):.1f}%, "
      f"90th percentile {100 * np.quantile(HILL_ERR, 0.9):.1f}%, worst {100 * HILL_ERR.max():.1f}%")

N_PRIOR = 1000
K_prior = np.exp(rng.normal(np.log(1.5), 0.5, N_PRIOR))
n_prior = np.exp(rng.normal(np.log(1.5), 0.4, N_PRIOR))
X_DATA_RANGE = jnp.linspace(0, X_RANGE, 101)
prior_targets = hill_np(np.asarray(X_DATA_RANGE)[None], K_prior[:, None], n_prior[:, None])
prior_targets /= hill_np(1.0, K_prior, n_prior)[:, None]


def fit_best_range(target):
    """As fit_best, on the data range [0, 2.5]."""
    def fit_from_r(p0):
        loss = lambda p: jnp.mean((weibull_curve(p, X_DATA_RANGE) - target) ** 2)  # noqa: E731

        def step(carry, _):
            p, state = carry
            updates, state = adam.update(jax.grad(loss)(p), state, p)
            return (optax.apply_updates(p, updates), state), None

        (p, _), _ = jax.lax.scan(step, (p0, adam.init(p0)), None, length=3000)
        return p, jnp.max(jnp.abs(weibull_curve(p, X_DATA_RANGE) - target)) / jnp.max(target)

    ps, errs = jax.vmap(fit_from_r)(STARTS)
    return ps[jnp.argmin(errs)], jnp.min(errs)


prior_params, prior_err = jax.jit(jax.vmap(fit_best_range))(jnp.asarray(prior_targets))
log_kc = np.asarray(prior_params)[:, 1:]  # (log k, log c)
print(f"translation error: median {100 * np.median(prior_err):.1f}%, worst {100 * np.max(prior_err):.1f}%")
mean_kc, cov_kc = log_kc.mean(0), np.cov(log_kc.T)
print("translated prior: (log k, log c) ~ MvNormal(mean =", mean_kc.round(2), ", sd =",
      np.sqrt(np.diag(cov_kc)).round(2), ", corr =", round(cov_kc[0, 1] / np.sqrt(np.prod(np.diag(cov_kc))), 2), ")")

def build_mmm(y, x_spend, curve="weibull"):
    coords = {"channel": CHANNELS, "week": np.arange(len(y)), "fourier": np.arange(FOURIER.shape[1]),
              "holiday": ["black_friday", "christmas"]}
    x_jax = jnp.asarray(x_spend)
    n_ch = len(CHANNELS)
    with pm.Model(coords=coords) as model:
        theta = pm.Beta("theta", 2.0, 2.0, dims="channel")  # adstock retention per week
        cac = pm.LogNormal("cac", np.log(60.0), 0.8, dims="channel")
        beta = pm.Deterministic("beta", SPEND_MEAN / cac, dims="channel")
        theta_ = pt.specify_shape(theta, (n_ch,))
        if curve == "weibull":
            k = pm.LogNormal("k", 0.0, 0.5, dims="channel")
            c = pm.HalfNormal("c", 1.5, dims="channel")
            transform = pytensor.wrap_jax(lambda th, k_, c_: media_transform(x_jax, th, k_, c_))
            f = transform(theta_, pt.specify_shape(k, (n_ch,)), pt.specify_shape(c, (n_ch,)))
        else:
            adstock = pytensor.wrap_jax(lambda th: geometric_adstock(x_jax, th))(theta_)
            if curve == "hill":
                K = pm.LogNormal("K", 0.0, 0.7, dims="channel")
                n = pm.LogNormal("n", 0.3, 0.5, dims="channel")
                f = adstock**n / (K**n + adstock**n) * (K**n + 1)
            elif curve == "logistic":
                lam = pm.LogNormal("lam", 0.0, 1.0, dims="channel")
                sat = lambda z: (1 - pt.exp(-z)) / (1 + pt.exp(-z))  # noqa: E731
                f = sat(lam * adstock) / sat(lam)
            else:
                raise ValueError(curve)
        a0 = pm.Normal("a0", np.log(np.median(y) / 2), 0.5)
        b_trend = pm.Normal("b_trend", 0.0, 0.3)
        b_fourier = pm.Normal("b_fourier", 0.0, 0.2, dims="fourier")
        b_holiday = pm.Normal("b_holiday", 0.0, 0.5, dims="holiday")
        baseline = pm.Deterministic(
            "baseline", pt.exp(a0 + b_trend * TREND + FOURIER @ b_fourier + HOLIDAYS @ b_holiday),
            dims="week")
        mu = baseline + f @ beta
        alpha = pm.Gamma("alpha", 2.0, 0.05)
        pm.NegativeBinomial("y", mu=mu, alpha=alpha, observed=y, dims="week")
    return model


model_weibull = build_mmm(y_obs, X_SPEND, "weibull")
print(model_weibull.point_logps())

def curve_draws(idata, curve, j, x, n=1000):
    """Posterior response curves (draw, x) in customers/week for channel j, and slope at x = 1."""
    post = az.extract(idata, num_samples=n, random_seed=RANDOM_SEED)
    beta = post["beta"].values[j]
    if curve == "weibull":
        k, c = post["k"].values[j], post["c"].values[j]
        f = jax.vmap(weibull, in_axes=(None, 0, 0))(jnp.asarray(x), k, c)
        slope = jax.vmap(jax.grad(weibull), in_axes=(None, 0, 0))(1.0, k, c)
    elif curve == "hill":
        K, n_ = post["K"].values[j][:, None], post["n"].values[j][:, None]
        hill = lambda z: z**n_ / (K**n_ + z**n_) * (K**n_ + 1)  # noqa: E731
        f, slope = hill(x[None]), ((hill(1.0 + 1e-5) - hill(1.0 - 1e-5)) / 2e-5)[:, 0]
    else:
        lam = post["lam"].values[j][:, None]
        sat = lambda z: (1 - np.exp(-lam * z)) / (1 + np.exp(-lam * z)) / logistic_np(1.0, lam)  # noqa: E731
        f, slope = sat(x[None]), ((sat(1.0 + 1e-5) - sat(1.0 - 1e-5)) / 2e-5)[:, 0]
    return beta[:, None] * np.asarray(f), beta * np.asarray(slope)


CURVE_COLORS = {"weibull": BLUE, "hill": AQUA, "logistic": ORANGE}
xg = np.linspace(0, 3, 200)
fig, axes = plt.subplots(1, 2, figsize=(13, 4.5))
rows = []
for j, (ax, ch) in enumerate(zip(axes, CHANNELS)):
    for curve, idata in fake_fits.items():
        f, slope = curve_draws(idata, curve, j, xg)
        lo, md, hi = np.quantile(f, [0.05, 0.5, 0.95], axis=0)
        ax.fill_between(xg, lo, hi, color=CURVE_COLORS[curve], alpha=0.15, lw=0)
        ax.plot(xg, md, color=CURVE_COLORS[curve], lw=2, label=curve)
        rows.append({"channel": ch, "model": curve, "customers at x=1": np.median(f[:, np.searchsorted(xg, 1.0)]),
                     "marginal at x=1 (median)": np.median(slope),
                     "marginal 90% interval": f"{np.quantile(slope, 0.05):.0f} - {np.quantile(slope, 0.95):.0f}"})
    true_f = truth_curves[ch](xg)
    ax.plot(xg, true_f, "k--", lw=1.5, label="truth")
    rows.append({"channel": ch, "model": "TRUTH", "customers at x=1": truth_curves[ch](1.0),
                 "marginal at x=1 (median)": (truth_curves[ch](1.001) - truth_curves[ch](0.999)) / 0.002,
                 "marginal 90% interval": ""})
    ax2 = ax.twinx()
    ax2.hist(ad_true[:, j], bins=30, range=(0, 3), color=GREY, alpha=0.25)
    ax2.set_yticks([])
    ax.set(xlabel="adstocked spend (1 = average week); grey: weeks observed", ylabel="customers / week",
           title=f"{ch.title()}: truth is {'S-shaped (Hill)' if ch == 'google' else 'a power curve'}")
    ax.set_zorder(ax2.get_zorder() + 1)
    ax.patch.set_visible(False)
axes[0].legend(fontsize=8, loc="upper left")
display(pd.DataFrame(rows).set_index(["channel", "model"]).round(0));