Circular data: directions, times of day and seasons

import logging
import time
import warnings

import arviz as az
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pymc as pm
import pytensor.tensor as pt
import xarray as xr
from matplotlib.patches import FancyArrowPatch
from scipy import special
from scipy.optimize import brentq

from pymc_challenges import data

RANDOM_SEED = 81
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-variat")
logging.getLogger("pymc").setLevel(logging.WARNING)
warnings.filterwarnings("ignore", category=RuntimeWarning, module="arviz")
warnings.filterwarnings("ignore", message=".*sample_stats.*")
pd.set_option("display.width", 170)
pd.set_option("display.max_columns", 20)
T_START = time.time()

BLUE, ORANGE, AQUA = "#2a78d6", "#eb6834", "#1baf7a"   # categorical slots 1-3, fixed order
INK, MUTED, LIGHT = "#0b0b0b", "#8a8984", "#d9d8d3"
TWO_PI = 2 * np.pi


def wrap(a):
    """Angle(s) to (-pi, pi]."""
    return np.angle(np.exp(1j * np.asarray(a)))


def mean_resultant(th, axis=None):
    """Mean direction and mean resultant length of angles in radians."""
    z = np.exp(1j * np.asarray(th)).mean(axis=axis)
    return np.angle(z), np.abs(z)


def A1(k):
    """Mean resultant length of a von Mises with concentration k: I1(k) / I0(k)."""
    return special.i1e(k) / special.i0e(k)


def A1inv(R):
    return brentq(lambda k: A1(k) - R, 1e-9, 1e4) if R > 1e-9 else 0.0


def vm_logpdf(th, e1, e2):
    """von Mises log density in NATURAL parameters: exp(e1 cos th + e2 sin th) / (2 pi I0(|e|))."""
    k = np.hypot(e1, e2)
    return e1 * np.cos(th) + e2 * np.sin(th) - np.log(TWO_PI) - (np.log(special.i0e(k)) + k)


def pn_logpdf(th, m1, m2):
    """Projected normal (identity covariance) log density of the angle of z ~ N((m1, m2), I)."""
    u = m1 * np.cos(th) + m2 * np.sin(th)
    return -np.log(TWO_PI) - 0.5 * (m1 ** 2 + m2 ** 2) + np.log1p(u * np.sqrt(np.pi / 2) * special.erfcx(-u / np.sqrt(2)))


def vm_lp(th, e1, e2):
    """The same von Mises log density as a PyTensor expression."""
    k = pt.sqrt(e1 ** 2 + e2 ** 2)
    return e1 * np.cos(th) + e2 * np.sin(th) - np.log(TWO_PI) - pt.log(pt.i0(k))


def loo_from_ll(ll):
    """PSIS-LOO from a (draws, observations) log-likelihood matrix computed in NumPy."""
    tree = xr.DataTree.from_dict({
        "posterior": xr.Dataset({"d": (("chain", "draw"), np.zeros((1, ll.shape[0])))}),
        "log_likelihood": xr.Dataset({"y": (("chain", "draw", "obs"), ll[None])}),
    })
    res = az.loo(tree, pointwise=True)
    res.log_weights = None
    return res


def circ_hdi(draws, prob=0.94):
    """Shortest arc containing `prob` of the draws (valid when they cover less than a half circle)."""
    c, _ = mean_resultant(draws)
    d = np.sort(wrap(draws - c))
    m = int(np.ceil(prob * len(d)))
    widths = d[m - 1:] - d[:len(d) - m + 1]
    i = np.argmin(widths)
    return c + d[i], c + d[i + m - 1], c


def compass(ax, labels=True):
    """Polar axes as a compass: north up, clockwise, degrees FROM which the wind blows."""
    ax.set_theta_zero_location("N")
    ax.set_theta_direction(-1)
    ax.set_xticks(np.deg2rad([0, 45, 90, 135, 180, 225, 270, 315]))
    ax.set_xticklabels(["N", "NE", "E", "SE", "S", "SW", "W", "NW"] if labels else [])
    ax.grid(alpha=0.4)


def clockface(ax):
    """Polar axes as a 24-hour clock: midnight up, clockwise."""
    ax.set_theta_zero_location("N")
    ax.set_theta_direction(-1)
    ax.set_xticks(np.deg2rad(np.arange(0, 360, 45)))
    ax.set_xticklabels([f"{h:02d}:00" for h in range(0, 24, 3)], fontsize=8)
    ax.grid(alpha=0.4)

R_target = 0.6
grid = np.linspace(-np.pi, np.pi, 721)
k_vm = A1inv(R_target)
sig_wn = np.sqrt(-2 * np.log(R_target))                     # wrapped normal: R = exp(-sigma^2 / 2)
dens = {
    f"von Mises (kappa = {k_vm:.2f})": np.exp(vm_logpdf(grid, k_vm, 0.0)),
    f"wrapped normal (sigma = {sig_wn:.2f})": sum(np.exp(-(grid + TWO_PI * j) ** 2 / (2 * sig_wn ** 2))
                                                 for j in range(-4, 5)) / np.sqrt(TWO_PI * sig_wn ** 2),
    "wrapped Cauchy (rho = 0.60)": (1 - R_target ** 2) / (TWO_PI * (1 + R_target ** 2 - 2 * R_target * np.cos(grid))),
}
# projected normal: find |m| giving R = 0.6 (R = E cos(theta) under the density)
dth = grid[1] - grid[0]
pn_R = lambda m: np.sum(np.cos(grid[:-1]) * np.exp(pn_logpdf(grid[:-1], m, 0.0))) * dth
m_pn = brentq(lambda m: pn_R(m) - R_target, 0.01, 10)
dens[f"projected normal (|m| = {m_pn:.2f})"] = np.exp(pn_logpdf(grid, m_pn, 0.0))

fig = plt.figure(figsize=(11, 4.2))
ax = fig.add_subplot(1, 2, 1)
axp = fig.add_subplot(1, 2, 2, projection="polar")
cols = [BLUE, ORANGE, AQUA, INK]
for (lab, d), c in zip(dens.items(), cols):
    print(f"{lab:36s} integrates to {np.sum(d[:-1]) * dth:.4f}; peak {d.max():.3f}; "
          f"density opposite the mean {d[0]:.3f}")
    ax.plot(np.rad2deg(grid), d, color=c, lw=2, label=lab)
    axp.plot(grid, d, color=c, lw=2)
ax.set(xlabel="angle from the mean direction (degrees)", ylabel="density (per radian)",
       xticks=[-180, -90, 0, 90, 180], title="Same mean direction, same mean resultant length 0.6")
ax.legend(fontsize=8)
axp.set_theta_zero_location("N")
axp.set_theta_direction(-1)
axp.set_title("The same densities on the circle", fontsize=10);

quakes = data.load("oklahoma_quakes")
t_local = pd.to_datetime(quakes.time, utc=True).dt.tz_convert("America/Chicago")
q_hour = (t_local.dt.hour + t_local.dt.minute / 60).to_numpy()
q_th = TWO_PI * q_hour / 24
n_q = len(q_th)
q_mean, q_R = mean_resultant(q_th)
Z = n_q * q_R ** 2
print(f"n = {n_q}; mean resultant length R = {q_R:.4f}; mean hour {np.rad2deg(q_mean) % 360 / 15:.1f}")
print(f"expected R under uniformity, sqrt(pi / 4n) = {np.sqrt(np.pi / (4 * n_q)):.4f}; "
      f"Rayleigh statistic nR^2 = {Z:.2f}, p ~ exp(-nR^2) = {np.exp(-Z):.2f}")

with pm.Model() as m_quake:
    eta = pm.Normal("eta", 0, 1, shape=2)                  # natural parameters of a von Mises
    kappa = pm.Deterministic("kappa", pt.sqrt(pt.sum(eta ** 2)))
    pm.Potential("ll", pt.sum(vm_lp(q_th, eta[0], eta[1])))
    idata_q = pm.sample(random_seed=RANDOM_SEED, progressbar=False)
kq = idata_q.posterior["kappa"].values.ravel()
eq = az.extract(idata_q, var_names=["eta"]).values.T
ll_q = vm_logpdf(q_th[None, :], eq[::4, 0:1], eq[::4, 1:2])
loo_q = loo_from_ll(ll_q)
print(f"posterior kappa: median {np.median(kq):.3f}, 94% upper bound {np.quantile(kq, 0.94):.3f}; "
      f"implied R at the upper bound {A1(np.quantile(kq, 0.94)):.3f}")
print(f"elpd: uniform {-n_q * np.log(TWO_PI):.1f}, von Mises {loo_q.elpd:.1f} +- {loo_q.se:.1f} "
      f"(difference {loo_q.elpd + n_q * np.log(TWO_PI):.1f})")

fig = plt.figure(figsize=(10, 4))
ax = fig.add_subplot(1, 2, 1, projection="polar")
clockface(ax)
cnt, edges = np.histogram(q_th, bins=24, range=(0, TWO_PI))
ax.bar(edges[:-1] + np.pi / 24, cnt / n_q, width=TWO_PI / 24 * 0.9, color=BLUE, alpha=0.8)
ax.plot(np.linspace(0, TWO_PI, 200), np.full(200, 1 / 24), color=INK, lw=1, ls="--")
ax.set_yticklabels([])
ax.set_title("Oklahoma M3+ earthquakes by hour\n(dashed: uniform, 1/24)", fontsize=10)
ax = fig.add_subplot(1, 2, 2)
ax.hist(kq, bins=60, color=BLUE, alpha=0.8)
ax.set(xlabel="von Mises concentration kappa", ylabel="posterior draws",
       title="Posterior of the concentration kappa");

x_sp = np.log(train.sknt.to_numpy())
X_MEAN, X_SD = x_sp.mean(), x_sp.std()
x_sp = (x_sp - X_MEAN) / X_SD


def tanhalf_profile(beta, th=th, x=x_sp):
    mu0, R = mean_resultant(th - 2 * np.arctan(beta * x))
    k = A1inv(R)
    return len(th) * (k * R - np.log(TWO_PI) - np.log(special.i0e(k)) - k), mu0, k


betas = np.r_[-np.logspace(2, -2, 90), 0.0, np.logspace(-2, 2, 90)]
prof = np.array([tanhalf_profile(b) for b in betas])
i_best = np.argmax(prof[:, 0])
neg = betas < -0.2
i_step = np.flatnonzero(neg)[np.argmax(prof[neg, 0])]
print(f"best: beta = {betas[i_best]:.2f}, mu0 = {np.rad2deg(prof[i_best, 1]) % 360:.0f} deg, "
      f"kappa = {prof[i_best, 2]:.2f}, log-lik {prof[i_best, 0]:.0f}")
print(f"second mode: beta = {betas[i_step]:.0f}, mu0 = {np.rad2deg(prof[i_step, 1]) % 360:.0f} deg, "
      f"kappa = {prof[i_step, 2]:.2f}, log-lik {prof[i_step, 0]:.0f}")
i_valley = np.flatnonzero((betas < 0) & (betas > betas[i_step]))[np.argmin(prof[(betas < 0) & (betas > betas[i_step]), 0])]
print(f"the valley between them: beta = {betas[i_valley]:.2f}, log-lik {prof[i_valley, 0]:.0f}")

fig, axes = plt.subplots(1, 2, figsize=(12, 4))
axes[0].plot(betas, prof[:, 0] - prof[i_best, 0], color=BLUE, lw=2)
axes[0].set_xscale("symlog", linthresh=0.05)
axes[0].set_xticks([-100, -10, -1, -0.1, 0, 0.1, 1, 10, 100])
axes[0].set_xticklabels(["-100", "-10", "-1", "-0.1", "0", "0.1", "1", "10", "100"])
axes[0].set(xlabel="beta (symmetric log scale)", ylabel="profile log-likelihood (relative to the best)",
            title="Profile likelihood of the tan-half slope", ylim=(-2000, 150))
axes[0].annotate("main mode", (betas[i_best], 0), xytext=(2, -600), arrowprops=dict(arrowstyle="->"))
axes[0].annotate("step-function mode", (betas[i_step], prof[i_step, 0] - prof[i_best, 0]),
                 xytext=(-60, -300), arrowprops=dict(arrowstyle="->"))
xs = np.linspace(-2.5, 2.5, 200)
sub = rng.choice(n_tr, 3000, replace=False)
axes[1].scatter(x_sp[sub] + rng.uniform(-0.05, 0.05, 3000), np.mod(train.drct.to_numpy()[sub] + rng.uniform(-5, 5, 3000), 360),
                s=2, color=MUTED, alpha=0.4)
for i, c, lab in ((i_best, BLUE, "main mode"), (i_step, ORANGE, "step-function mode")):
    mu = np.rad2deg(prof[i, 1] + 2 * np.arctan(betas[i] * xs)) % 360
    mu[np.abs(np.diff(mu, append=mu[-1])) > 180] = np.nan
    axes[1].plot(xs, mu, color=c, lw=2.5, label=f"{lab}: beta = {betas[i]:.2g}")
axes[1].set(xlabel="standardised log wind speed", ylabel="direction FROM (degrees)", yticks=range(0, 361, 90),
            title="The two modes (grey: 3,000 reports, jittered)")
axes[1].legend(fontsize=8, loc="lower right");

edges_kn = np.array([0, 3, 4, 5, 6, 8, 10, 13, 17, 60])     # the lightest reported wind is 3 knots
kn = np.zeros(len(edges_kn) - 1)
speeds = np.linspace(3, 25, 120)
xs_kn = (np.log(speeds) - X_MEAN) / X_SD
e_grid = canon_eta(cd[:300], xs_kn)
mu_c = np.arctan2(e_grid[..., 1], e_grid[..., 0])
R_c = A1(np.hypot(e_grid[..., 0], e_grid[..., 1]))
mu_t = mu0_d[:300, None] + 2 * np.arctan(beta_d[:300, None] * xs_kn[None])
R_t = np.broadcast_to(A1(kap_d[:300, None]), mu_t.shape)
bin_id = np.digitize(train.sknt.to_numpy(), edges_kn[1:-1], right=True)
ob_mu, ob_R = np.array([mean_resultant(th[bin_id == j]) for j in range(len(kn))]).T
kn = np.array([np.median(train.sknt.to_numpy()[bin_id == j]) for j in range(len(kn))])
print(pd.DataFrame({"median speed (kn)": kn, "reports": np.bincount(bin_id),
                    "mean direction": np.rad2deg(ob_mu) % 360, "R": ob_R}).round(2).T.to_string())

fig = plt.figure(figsize=(12, 5))
ax = fig.add_subplot(1, 2, 1)
for mu_s, c, lab in ((mu_c, BLUE, "canonical, quadratic"), (mu_t, ORANGE, "tan-half, main mode")):
    cm = mean_resultant(mu_s, axis=0)[0]
    dev = wrap(mu_s - cm)
    lo, hi = np.quantile(dev, [0.03, 0.97], axis=0)
    centre = np.rad2deg(np.unwrap(cm)) % 360
    ax.fill_between(speeds, centre + np.rad2deg(lo), centre + np.rad2deg(hi), color=c, alpha=0.3)
    ax.plot(speeds, centre, color=c, lw=2, label=lab)
ax.errorbar(kn, np.rad2deg(ob_mu) % 360, fmt="o", color=INK, ms=5, label="observed circular mean per speed bin")
ax.set(xscale="log", xlabel="wind speed (knots, log scale)", ylabel="mean direction FROM (degrees)",
       yticks=range(0, 361, 45), title="Mean direction against speed (94% bands)")
ax.set_xticks([3, 5, 10, 20])
ax.set_xticklabels(["3", "5", "10", "20"])
ax.xaxis.set_minor_formatter(plt.NullFormatter())
ax.legend(fontsize=8, loc="upper left")
axp = fig.add_subplot(1, 2, 2, projection="polar")
compass(axp)
for i in range(60):
    axp.plot(mu_c[i], R_c[i], color=BLUE, lw=0.6, alpha=0.25)
    axp.plot(mu_t[i], R_t[i], color=ORANGE, lw=0.6, alpha=0.25)
axp.plot(ob_mu, ob_R, "o", color=INK, ms=5)
for m_, r_, k_ in zip(ob_mu, ob_R, kn):
    axp.annotate(f"{k_:g} kn", (m_, r_), fontsize=7, xytext=(4, 4), textcoords="offset points")
axp.set_ylim(0, 0.85)
axp.set_yticks([0.2, 0.4, 0.6])
axp.tick_params(axis="y", labelsize=7)
axp.set_title("Mean resultant vector as speed grows\n(blue: canonical, orange: tan-half; 60 draws each)", fontsize=10);

def features(df):
    h = TWO_PI * df.hour.to_numpy() / 24
    d = TWO_PI * df.doy.to_numpy()
    cols, names = [np.ones_like(h)], ["1"]
    for k in (1, 2, 3):
        cols += [np.cos(k * h), np.sin(k * h)]
        names += [f"cos{k}h", f"sin{k}h"]
    for k in (1, 2):
        cols += [np.cos(k * d), np.sin(k * d)]
        names += [f"cos{k}d", f"sin{k}d"]
    for k in (1, 2):
        for a, an in ((np.cos(k * h), f"cos{k}h"), (np.sin(k * h), f"sin{k}h")):
            for b, bn in ((np.cos(d), "cos1d"), (np.sin(d), "sin1d")):
                cols.append(a * b)
                names.append(f"{an}*{bn}")
    return np.column_stack(cols), names


F_tr, FEAT = features(train)
F_te, _ = features(test)
SMALL = [0, 1, 2, 3, 4, 7, 8]                              # 1, cos1h, sin1h, cos2h, sin2h, cos1d, sin1d
prior_sd = np.r_[2.0, np.ones(len(FEAT) - 1)]
coords = {"feat": FEAT, "comp": ["eta1", "eta2"], "sfeat": [FEAT[i] for i in SMALL],
          "regime": ["easterly", "westerly"]}

# prior predictive: implied concentration at the training rows
B_prior = rng.normal(0, prior_sd[None, :, None], size=(200, len(FEAT), 2))
kap_prior = np.hypot(*np.moveaxis(np.einsum("nf,sfc->snc", F_tr[:2000], B_prior), -1, 0))
print("prior concentration at training rows: quantiles 5/50/95% =", np.round(np.quantile(kap_prior, [0.05, 0.5, 0.95]), 2))

t0 = time.time()
with pm.Model(coords=coords) as m_vm:
    B = pm.Normal("B", 0, prior_sd[:, None], dims=("feat", "comp"))
    eta = pt.dot(F_tr, B)
    pm.Potential("ll", pt.sum(vm_lp(th, eta[:, 0], eta[:, 1])))
    idata_vm = pm.sample(random_seed=RANDOM_SEED, progressbar=False)
print(f"von Mises regression: {time.time() - t0:.0f} s")

t0 = time.time()
with pm.Model(coords=coords) as m_pn:
    B = pm.Normal("B", 0, prior_sd[:, None], dims=("feat", "comp"))
    m = pt.dot(F_tr, B)
    u = m[:, 0] * np.cos(th) + m[:, 1] * np.sin(th)
    pm.Potential("ll", pt.sum(-np.log(TWO_PI) - 0.5 * pt.sum(m ** 2, axis=1)
                              + pt.log1p(u * np.sqrt(np.pi / 2) * pt.erfcx(-u / np.sqrt(2)))))
    idata_pn = pm.sample(random_seed=RANDOM_SEED, progressbar=False)
print(f"projected normal regression: {time.time() - t0:.0f} s")

MU_PRIOR = wrap(np.deg2rad([90.0, 250.0]))                 # easterly from E, westerly from WSW
t0 = time.time()
with pm.Model(coords=coords) as m_mix:
    B0 = pm.Normal("B0", 2.0 * np.column_stack([np.cos(MU_PRIOR), np.sin(MU_PRIOR)]), 1.0, dims=("regime", "comp"))
    Bh = pm.Normal("Bh", 0, 1, shape=(2, len(SMALL) - 1, 2))
    Bc = pt.concatenate([B0[:, None, :], Bh], axis=1)      # regime x small feature x component
    e_e = pt.dot(F_tr[:, SMALL], Bc[0])
    e_w = pt.dot(F_tr[:, SMALL], Bc[1])
    g = pm.Normal("g", 0, np.r_[1.5, np.ones(len(FEAT) - 1)], dims="feat")
    z = pt.dot(F_tr, g)
    ll = pt.logaddexp(-pm.math.log1pexp(z) + vm_lp(th, e_e[:, 0], e_e[:, 1]),
                      -pm.math.log1pexp(-z) + vm_lp(th, e_w[:, 0], e_w[:, 1]))
    pm.Potential("ll", pt.sum(ll))
    idata_mix = pm.sample(random_seed=RANDOM_SEED, progressbar=False)
print(f"mixture of regimes: {time.time() - t0:.0f} s")

for name, idd, vs in (("von Mises", idata_vm, ["B"]), ("projected normal", idata_pn, ["B"]),
                      ("mixture", idata_mix, ["B0", "Bh", "g"])):
    s = pd.concat([az.summary(idd, var_names=[v], round_to=4) for v in vs])
    print(f"{name:17s} max r_hat {s.r_hat.max():.3f}, min bulk ESS {s.ess_bulk.min():.0f}, "
          f"divergences {int(idd.sample_stats['diverging'].sum())}, mean tree depth "
          f"{float(idd.sample_stats['depth'].mean()):.1f}, warm-up {idd.posterior.attrs.get('tuning_steps')}")
print(per_chain(idata_mix, []).round(1).to_string())
print("regime intercept directions per chain (deg):",
      np.round(np.rad2deg(np.arctan2(idata_mix.posterior["B0"].values[..., 1],
                                     idata_mix.posterior["B0"].values[..., 0]).mean(1)) % 360).tolist())

S_LL = 1000
draws = {
    "vm": az.extract(idata_vm, var_names=["B"], num_samples=S_LL, random_seed=1).transpose("sample", ...).values,
    "pn": az.extract(idata_pn, var_names=["B"], num_samples=S_LL, random_seed=1).transpose("sample", ...).values,
}
mx = az.extract(idata_mix, var_names=["B0", "Bh", "g"], num_samples=S_LL, random_seed=1)
draws["mix"] = (np.concatenate([mx["B0"].transpose("sample", ...).values[:, :, None, :],
                                mx["Bh"].transpose("sample", ...).values], axis=2),
                mx["g"].transpose("sample", ...).values)


def loglik(model, F, th_obs):
    if model == "vm":
        e = np.einsum("nf,sfc->snc", F, draws["vm"])
        return vm_logpdf(th_obs[None], e[..., 0], e[..., 1])
    if model == "pn":
        m = np.einsum("nf,sfc->snc", F, draws["pn"])
        return pn_logpdf(th_obs[None], m[..., 0], m[..., 1])
    Bc, g = draws["mix"]
    ee = np.einsum("nf,srfc->srnc", F[:, SMALL], Bc)
    z = np.einsum("nf,sf->sn", F, g)
    return np.logaddexp(-np.logaddexp(0, z) + vm_logpdf(th_obs[None], ee[:, 0, :, 0], ee[:, 0, :, 1]),
                        -np.logaddexp(0, -z) + vm_logpdf(th_obs[None], ee[:, 1, :, 0], ee[:, 1, :, 1]))


const_mu, const_R = mean_resultant(th)
const_k = A1inv(const_R)
loos, lpd_test = {}, {}
th_te = test.theta.to_numpy()
for mname in ("vm", "pn", "mix"):
    loos[mname] = loo_from_ll(loglik(mname, F_tr, th))
    lt = loglik(mname, F_te, th_te)
    lpd_test[mname] = special.logsumexp(lt, axis=0) - np.log(S_LL)
    del lt
lpd_test["one direction"] = vm_logpdf(th_te, const_k * np.cos(const_mu), const_k * np.sin(const_mu))
lpd_test["uniform"] = np.full(len(th_te), -np.log(TWO_PI))
names_long = {"vm": "von Mises (canonical)", "pn": "projected normal", "mix": "mixture of 2 regimes"}
print(az.compare({names_long[k]: v for k, v in loos.items()}, round_to=1).to_string())

day = test.lst.dt.dayofyear.to_numpy()
rows = []
for k, v in lpd_test.items():
    dd = pd.Series(v - lpd_test["mix"]).groupby(day).sum()
    rows.append({"model": names_long.get(k, k), "2023 log score per report": v.mean(),
                 "total vs mixture": v.sum() - lpd_test["mix"].sum(),
                 "se (by day)": dd.std() * np.sqrt(len(dd))})
print(pd.DataFrame(rows).set_index("model").round(3).to_string())

def simulate(mname, F, n_rep, seed=1):
    r = np.random.default_rng(seed)
    if mname == "vm":
        e = np.einsum("nf,sfc->snc", F, draws["vm"][:n_rep])
        return r.vonmises(np.arctan2(e[..., 1], e[..., 0]), np.hypot(e[..., 0], e[..., 1]))
    Bc, g = draws["mix"]
    ee = np.einsum("nf,srfc->srnc", F[:, SMALL], Bc[:n_rep])
    west = r.random((n_rep, len(F))) < special.expit(np.einsum("nf,sf->sn", F, g[:n_rep]))
    e1 = np.where(west, ee[:, 1, :, 0], ee[:, 0, :, 0])
    e2 = np.where(west, ee[:, 1, :, 1], ee[:, 0, :, 1])
    return r.vonmises(np.arctan2(e2, e1), np.hypot(e1, e2))


N_REP = 200
rep = {m_: simulate(m_, F_tr, N_REP) for m_ in ("vm", "mix")}
bins12 = np.linspace(-np.pi / 12, TWO_PI - np.pi / 12, 13)
blk = (train.hour.to_numpy() // 3).astype(int)
fig, axes = plt.subplots(2, 4, figsize=(12, 6.6), subplot_kw={"projection": "polar"})
centres = bins12[:-1] + np.pi / 12
for b, ax in enumerate(axes.ravel()):
    sel = blk == b
    obs = np.histogram(np.mod(th[sel], TWO_PI), bins=bins12)[0] / sel.sum()
    ax.bar(centres, obs, width=TWO_PI / 12 * 0.92, color=LIGHT, edgecolor=MUTED, lw=0.5)
    for m_, c in (("vm", ORANGE), ("mix", BLUE)):
        h = np.array([np.histogram(np.mod(rep[m_][i, sel], TWO_PI), bins=bins12)[0] for i in range(N_REP)]) / sel.sum()
        lo, md, hi = np.quantile(h, [0.05, 0.5, 0.95], axis=0)
        cc = np.r_[centres, centres[0]]
        ax.fill_between(cc, np.r_[lo, lo[0]], np.r_[hi, hi[0]], color=c, alpha=0.35)
        ax.plot(cc, np.r_[md, md[0]], color=c, lw=1.5)
    compass(ax)
    ax.set_yticklabels([])
    ax.set_title(f"{3 * b:02d}:00-{3 * b + 3:02d}:00", fontsize=10)
fig.suptitle("Posterior predictive rose diagrams, 2021-2022 (grey bars: observed; orange: von Mises regression; "
             "blue: mixture of regimes; bands: 90% of 200 replicates)", fontsize=10);