Excess deaths due to COVID-19#
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
from pymc_extras.prior import Prior
import causalpy as cp
%load_ext autoreload
%autoreload 2
%config InlineBackend.figure_format = 'retina'
seed = 42
The autoreload extension is already loaded. To reload it, use:
%reload_ext autoreload
Load data#
df = (
cp.load_data("covid")
.assign(date=lambda x: pd.to_datetime(x["date"]))
.set_index("date")
)
treatment_time = pd.to_datetime("2020-01-01")
df.head()
| temp | deaths | year | month | t | pre | |
|---|---|---|---|---|---|---|
| date | ||||||
| 2006-01-01 | 3.8 | 49124 | 2006 | 1 | 0 | True |
| 2006-02-01 | 3.4 | 42664 | 2006 | 2 | 1 | True |
| 2006-03-01 | 3.9 | 49207 | 2006 | 3 | 2 | True |
| 2006-04-01 | 7.4 | 40645 | 2006 | 4 | 3 | True |
| 2006-05-01 | 10.7 | 42425 | 2006 | 5 | 4 | True |
The columns are:
date+year: self explanatorymonth: month, numerically encoded. Needs to be treated as a categorical variabletemp: average UK temperature (Celsius)t: timepre: boolean flag indicating pre or post intervention
A linear-trend baseline#
Let’s begin with a simple model that assumes a straight-line trend in t, along with monthly fixed effects, and an adjustment for temperature. Since our outcome variable is standardized, we’ll place unit normal priors on our regression coefficients and a half normal on the standard deviation.
Note
In this example we are going to standardize the data. So we have to be careful in how we interpret the inferred regression coefficients, and the posterior predictions will be in this standardized space.
Note
The random_seed keyword argument for the PyMC sampler is not necessary. We use it here so that the results are reproducible.
model = cp.pymc_models.LinearRegression(
sample_kwargs={"random_seed": seed},
priors={
"beta": Prior("Normal", mu=0, sigma=1, dims=["treated_units", "coeffs"]),
"y_hat": Prior(
"Normal",
sigma=Prior("HalfNormal", sigma=0.5, dims=["treated_units"]),
dims=["obs_ind", "treated_units"],
),
},
)
result_linear = cp.InterruptedTimeSeries(
df,
treatment_time,
formula="standardize(deaths) ~ 0 + standardize(t) + C(month) + standardize(temp)",
model=model,
)
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [beta, y_hat_sigma]
?25l
Step Grad
Progress Draw Diverge… size evals Speed Elapsed Remain…
──────────────────────────────────────────────────────────────────────────────
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 109 0 0.278 15 1041.81 0:00:00 0:00:02
draws/s
━━━━━━━ 124 0 0.299 15 1204.52 0:00:00 0:00:02
draws/s
━━━━━━━ 109 0 0.335 7 1076.59 0:00:00 0:00:02
draws/s
━━━━━━━ 103 0 0.030 63 1039.65 0:00:00 0:00:02
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━╺━━━━━ 340 0 0.285 7 1583.27 0:00:00 0:00:02
draws/s
━╺━━━━━ 345 0 0.442 15 1617.40 0:00:00 0:00:01
draws/s
━╺━━━━━ 289 0 0.355 15 1366.87 0:00:00 0:00:02
draws/s
╸━━━━━━ 267 0 0.262 15 1277.10 0:00:00 0:00:02
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━╺━━━━ 619 0 0.325 15 1908.40 0:00:00 0:00:01
draws/s
━━╺━━━━ 627 0 0.375 7 1940.23 0:00:00 0:00:01
draws/s
━╸━━━━━ 505 0 0.478 7 1573.60 0:00:00 0:00:01
draws/s
━╸━━━━━ 471 0 0.368 7 1477.07 0:00:00 0:00:02
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━╺━━━ 904 0 0.466 7 2085.81 0:00:00 0:00:01
draws/s
━━━╺━━━ 882 0 0.462 7 2043.61 0:00:00 0:00:01
draws/s
━━╸━━━━ 742 0 0.431 7 1727.44 0:00:00 0:00:01
draws/s
━━╺━━━━ 661 0 0.376 7 1545.70 0:00:00 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━╺━━ 1237 0 0.443 7 2288.86 0:00:00 0:00:01
draws/s
━━━━╺━━ 1158 0 0.369 7 2150.06 0:00:00 0:00:01
draws/s
━━━╺━━━ 960 0 0.484 7 1787.68 0:00:00 0:00:01
draws/s
━━╸━━━━ 847 0 0.508 7 1584.64 0:00:00 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━╸━ 1574 0 0.443 7 2423.13 0:00:00 0:00:01
draws/s
━━━━━╺━ 1444 0 0.369 7 2228.48 0:00:00 0:00:01
draws/s
━━━━╺━━ 1197 0 0.432 15 1852.56 0:00:00 0:00:01
draws/s
━━━╸━━━ 1038 0 0.317 7 1612.37 0:00:00 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━╸ 1914 0 0.443 7 2524.88 0:00:00 0:00:01
draws/s
━━━━━━╺ 1746 0 0.369 7 2308.31 0:00:00 0:00:01
draws/s
━━━━━╺━ 1452 0 0.432 7 1924.12 0:00:00 0:00:01
draws/s
━━━━╺━━ 1252 0 0.317 7 1664.31 0:00:00 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.443 7 2535.01 0:00:00 0:00:00
draws/s
━━━━━━╺ 1807 0 0.369 7 2295.64 0:00:00 0:00:01
draws/s
━━━━━╺━ 1508 0 0.432 7 1921.25 0:00:00 0:00:01
draws/s
━━━━╸━━ 1304 0 0.317 7 1666.38 0:00:00 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.443 7 2535.01 0:00:00 0:00:00
draws/s
━━━━━━━ 2000 0 0.369 7 2353.96 0:00:00 0:00:00
draws/s
━━━━━╸━ 1683 0 0.432 7 1986.70 0:00:00 0:00:01
draws/s
━━━━━╺━ 1466 0 0.317 7 1735.11 0:00:00 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.443 7 2535.01 0:00:00 0:00:00
draws/s
━━━━━━━ 2000 0 0.369 7 2353.96 0:00:00 0:00:00
draws/s
━━━━━━╺ 1726 0 0.432 7 1999.38 0:00:00 0:00:01
draws/s
━━━━━╺━ 1503 0 0.317 7 1745.18 0:00:00 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.443 7 2535.01 0:00:00 0:00:00
draws/s
━━━━━━━ 2000 0 0.369 7 2353.96 0:00:00 0:00:00
draws/s
━━━━━━━ 2000 0 0.432 7 2087.88 0:00:00 0:00:00
draws/s
━━━━━━╺ 1767 0 0.317 7 1850.28 0:00:00 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.443 7 2535.01 0:00:00 0:00:00
draws/s
━━━━━━━ 2000 0 0.369 7 2353.96 0:00:00 0:00:00
draws/s
━━━━━━━ 2000 0 0.432 7 2087.88 0:00:00 0:00:00
draws/s
━━━━━━╺ 1802 0 0.317 7 1858.91 0:00:00 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.443 7 2535.01 0:00:00 0:00:00
draws/s
━━━━━━━ 2000 0 0.369 7 2353.96 0:00:00 0:00:00
draws/s
━━━━━━━ 2000 0 0.432 7 2087.88 0:00:00 0:00:00
draws/s
━━━━━━━ 2000 0 0.317 15 1937.76 0:00:01 0:00:00
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.443 7 2535.01 0:00:00 0:00:00
draws/s
━━━━━━━ 2000 0 0.369 7 2353.96 0:00:00 0:00:00
draws/s
━━━━━━━ 2000 0 0.432 7 2087.88 0:00:00 0:00:00
draws/s
━━━━━━━ 2000 0 0.317 15 1937.76 0:00:01 0:00:00
draws/s
?25hSampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 2 seconds.
Sampling: [beta, y_hat, y_hat_sigma]
Sampling: [y_hat]
Sampling: [y_hat]
Sampling: [y_hat]
Sampling: [y_hat]
fig, ax = result_linear.plot()
result_linear.summary()
==================================Pre-Post Fit==================================
Formula: standardize(deaths) ~ 0 + standardize(t) + C(month) + standardize(temp)
Model coefficients:
C(month)[1] 1.5, 94% HDI [1.1, 1.9]
C(month)[2] -0.25, 94% HDI [-0.68, 0.13]
C(month)[3] 0.22, 94% HDI [-0.13, 0.56]
C(month)[4] -0.045, 94% HDI [-0.31, 0.25]
C(month)[5] -0.13, 94% HDI [-0.41, 0.17]
C(month)[6] -0.16, 94% HDI [-0.54, 0.2]
C(month)[7] 0.046, 94% HDI [-0.42, 0.49]
C(month)[8] -0.35, 94% HDI [-0.79, 0.081]
C(month)[9] -0.39, 94% HDI [-0.74, -0.034]
C(month)[10] -0.042, 94% HDI [-0.31, 0.22]
C(month)[11] -0.38, 94% HDI [-0.7, -0.062]
C(month)[12] 0.019, 94% HDI [-0.37, 0.42]
standardize(t) 0.23, 94% HDI [0.15, 0.31]
standardize(temp) -0.49, 94% HDI [-0.74, -0.24]
Adding flexibility#
This model doesn’t look too bad, but with a pre-treatment period Bayesian \(R^2\) of 0.71, we have some room for improvement. One thing we could try is to relax our assumption that the trend is globally linear. In doing so, we might improve pretreatment fit and partially deconfound our other coefficients. To this end, let’s try to swap the linear model for t for a spline. We can use cr() from Patsy to fit a cubic regression spline in the cardinal parameterization. cr() takes an argument for the number of degrees of freedom for the fit. For df=k you get k basis columns, and then coefficient j is simply the value of the fitted curve at knot j.
CausalPy makes fitting this model easy enough, so let’s try it! We’ll swap in cr(t, df=6) and keep our uninformative normal priors on the coefficients.
ridge_spline = cp.pymc_models.LinearRegression(
sample_kwargs={"random_seed": seed},
priors={
"beta": Prior("Normal", mu=0, sigma=50, dims=["treated_units", "coeffs"]),
"y_hat": Prior(
"Normal",
sigma=Prior("HalfNormal", sigma=0.5, dims=["treated_units"]),
dims=["obs_ind", "treated_units"],
),
},
)
result_ridge = cp.InterruptedTimeSeries(
df,
treatment_time,
formula="standardize(deaths) ~ 0 + cr(t, df=6) + C(month) + standardize(temp)",
model=ridge_spline,
)
print(f"pre-intervention R2 = {result_ridge.score['unit_0_r2']:.3f}")
az.summary(result_ridge.model.idata, var_names=["beta"]).sort_values("ess_bulk").head(5)
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [beta, y_hat_sigma]
?25l
Step Grad
Progress Draw Diverge… size evals Speed Elapsed Remain…
──────────────────────────────────────────────────────────────────────────────
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 79 0 0.045 15 789.98 0:00:00 0:00:03
draws/s
━━━━━━━ 55 0 0.013 47 554.48 0:00:00 0:00:04
draws/s
━━━━━━━ 51 0 0.051 15 503.89 0:00:00 0:00:04
draws/s
━━━━━━━ 67 0 0.048 255 690.00 0:00:00 0:00:03
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 131 0 0.008 7 615.51 0:00:00 0:00:03
draws/s
━━━━━━━ 138 0 0.049 95 649.70 0:00:00 0:00:03
draws/s
━━━━━━━ 107 0 0.038 47 518.69 0:00:00 0:00:04
draws/s
━━━━━━━ 131 0 0.047 7 632.00 0:00:00 0:00:03
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
╸━━━━━━ 158 0 0.008 255 494.80 0:00:00 0:00:04
draws/s
╸━━━━━━ 182 0 0.018 7 578.15 0:00:00 0:00:04
draws/s
╸━━━━━━ 157 0 0.027 55 500.70 0:00:00 0:00:04
draws/s
╸━━━━━━ 172 0 0.029 7 550.35 0:00:00 0:00:04
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
╸━━━━━━ 196 0 0.017 7 461.03 0:00:00 0:00:04
draws/s
╸━━━━━━ 226 0 0.022 95 538.35 0:00:00 0:00:04
draws/s
╸━━━━━━ 180 0 0.015 7 432.57 0:00:00 0:00:05
draws/s
╸━━━━━━ 211 0 0.019 127 503.82 0:00:00 0:00:04
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
╸━━━━━━ 228 0 0.011 7 426.59 0:00:00 0:00:05
draws/s
╸━━━━━━ 254 0 0.019 71 478.81 0:00:00 0:00:04
draws/s
╸━━━━━━ 211 0 0.008 79 418.43 0:00:00 0:00:05
draws/s
╸━━━━━━ 225 0 0.012 119 433.29 0:00:00 0:00:05
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
╸━━━━━━ 248 0 0.007 15 397.40 0:00:00 0:00:05
draws/s
╸━━━━━━ 284 0 0.006 7 442.92 0:00:00 0:00:04
draws/s
╸━━━━━━ 231 0 0.007 15 362.98 0:00:00 0:00:05
draws/s
╸━━━━━━ 250 0 0.007 255 406.67 0:00:00 0:00:05
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
╸━━━━━━ 275 0 0.009 15 377.86 0:00:00 0:00:05
draws/s
━╺━━━━━ 303 0 0.008 31 416.62 0:00:00 0:00:05
draws/s
╸━━━━━━ 245 0 0.004 159 327.37 0:00:00 0:00:06
draws/s
╸━━━━━━ 272 0 0.006 7 364.69 0:00:00 0:00:05
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━╺━━━━━ 288 0 0.008 1023 339.16 0:00:00 0:00:06
draws/s
━╺━━━━━ 316 0 0.005 1023 370.40 0:00:00 0:00:05
draws/s
╸━━━━━━ 267 0 0.009 23 315.04 0:00:00 0:00:06
draws/s
━╺━━━━━ 298 0 0.009 307 353.92 0:00:00 0:00:05
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━╺━━━━━ 300 0 0.010 23 317.58 0:00:00 0:00:06
draws/s
━╺━━━━━ 331 0 0.006 79 347.10 0:00:00 0:00:05
draws/s
━╺━━━━━ 291 0 0.007 255 304.99 0:00:00 0:00:06
draws/s
━╺━━━━━ 335 0 0.010 7 349.02 0:00:00 0:00:05
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━╺━━━━━ 315 0 0.006 7 294.52 0:00:01 0:00:06
draws/s
━╺━━━━━ 343 0 0.006 287 326.07 0:00:01 0:00:06
draws/s
━╺━━━━━ 321 0 0.008 655 299.96 0:00:01 0:00:06
draws/s
━╺━━━━━ 347 0 0.009 111 330.56 0:00:01 0:00:05
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━╺━━━━━ 333 0 0.005 143 283.86 0:00:01 0:00:06
draws/s
━╺━━━━━ 364 0 0.006 23 309.27 0:00:01 0:00:06
draws/s
━╺━━━━━ 341 0 0.008 7 289.60 0:00:01 0:00:06
draws/s
━╺━━━━━ 358 0 0.009 7 309.58 0:00:01 0:00:06
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━╺━━━━━ 341 0 0.006 7 269.39 0:00:01 0:00:07
draws/s
━╺━━━━━ 387 0 0.007 15 303.89 0:00:01 0:00:06
draws/s
━╺━━━━━ 365 0 0.010 63 285.06 0:00:01 0:00:06
draws/s
━╺━━━━━ 372 0 0.010 1023 290.49 0:00:01 0:00:06
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━╺━━━━━ 356 0 0.006 311 255.58 0:00:01 0:00:07
draws/s
━╺━━━━━ 396 0 0.007 1023 284.17 0:00:01 0:00:06
draws/s
━╺━━━━━ 381 0 0.011 255 275.99 0:00:01 0:00:06
draws/s
━╺━━━━━ 396 0 0.012 3 285.94 0:00:01 0:00:06
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━╺━━━━━ 367 0 0.008 223 244.30 0:00:01 0:00:07
draws/s
━╺━━━━━ 400 0 0.007 111 270.85 0:00:01 0:00:06
draws/s
━╺━━━━━ 411 0 0.009 7 274.93 0:00:01 0:00:06
draws/s
━╺━━━━━ 416 0 0.011 11 278.84 0:00:01 0:00:06
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━╺━━━━━ 375 0 0.007 303 234.06 0:00:01 0:00:07
draws/s
━╺━━━━━ 420 0 0.012 7 261.44 0:00:01 0:00:07
draws/s
━╸━━━━━ 441 0 0.009 15 276.67 0:00:01 0:00:06
draws/s
━╸━━━━━ 448 0 0.007 3 279.50 0:00:01 0:00:06
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━╺━━━━━ 392 0 0.007 255 227.96 0:00:01 0:00:08
draws/s
━╸━━━━━ 440 0 0.007 935 256.11 0:00:01 0:00:07
draws/s
━╸━━━━━ 458 0 0.010 335 269.64 0:00:01 0:00:06
draws/s
━╸━━━━━ 467 0 0.008 319 274.32 0:00:01 0:00:06
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━╺━━━━━ 410 0 0.013 55 226.97 0:00:01 0:00:07
draws/s
━╸━━━━━ 482 0 0.008 7 264.76 0:00:01 0:00:06
draws/s
━╸━━━━━ 487 0 0.012 239 268.12 0:00:01 0:00:06
draws/s
━╸━━━━━ 487 0 0.008 7 270.27 0:00:01 0:00:06
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━╸━━━━━ 445 0 0.010 31 230.51 0:00:01 0:00:07
draws/s
━╸━━━━━ 500 0 0.007 47 259.96 0:00:01 0:00:06
draws/s
━╸━━━━━ 509 0 0.010 15 264.35 0:00:01 0:00:06
draws/s
━╸━━━━━ 514 0 0.007 7 267.26 0:00:01 0:00:06
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━╸━━━━━ 469 0 0.007 7 230.64 0:00:02 0:00:07
draws/s
━╸━━━━━ 518 0 0.006 7 253.92 0:00:02 0:00:06
draws/s
━╸━━━━━ 537 0 0.008 7 265.06 0:00:02 0:00:06
draws/s
━╸━━━━━ 532 0 0.009 7 261.91 0:00:02 0:00:06
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━╸━━━━━ 481 0 0.014 255 225.21 0:00:02 0:00:07
draws/s
━╸━━━━━ 538 0 0.011 15 253.01 0:00:02 0:00:06
draws/s
━╸━━━━━ 561 0 0.009 7 261.87 0:00:02 0:00:06
draws/s
━╸━━━━━ 540 0 0.009 63 251.93 0:00:02 0:00:06
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━╸━━━━━ 494 0 0.012 7 219.46 0:00:02 0:00:07
draws/s
━╸━━━━━ 553 0 0.008 7 248.26 0:00:02 0:00:06
draws/s
━━╺━━━━ 582 0 0.005 95 258.62 0:00:02 0:00:06
draws/s
━╸━━━━━ 545 0 0.008 567 241.99 0:00:02 0:00:07
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━╸━━━━━ 516 0 0.013 343 218.39 0:00:02 0:00:07
draws/s
━━╺━━━━ 596 0 0.008 7 252.18 0:00:02 0:00:06
draws/s
━━╺━━━━ 605 0 0.008 1023 256.35 0:00:02 0:00:06
draws/s
━╸━━━━━ 561 0 0.008 415 238.66 0:00:02 0:00:07
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━╸━━━━━ 546 0 0.014 135 221.48 0:00:02 0:00:07
draws/s
━━╺━━━━ 617 0 0.011 455 250.80 0:00:02 0:00:06
draws/s
━━╺━━━━ 621 0 0.007 19 252.47 0:00:02 0:00:06
draws/s
━╸━━━━━ 569 0 0.008 511 231.96 0:00:02 0:00:07
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━╺━━━━ 576 0 0.010 15 224.74 0:00:02 0:00:07
draws/s
━━╺━━━━ 639 0 0.008 87 249.85 0:00:02 0:00:06
draws/s
━━╺━━━━ 645 0 0.013 7 251.48 0:00:02 0:00:06
draws/s
━━╺━━━━ 584 0 0.008 7 228.92 0:00:02 0:00:07
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━╺━━━━ 588 0 0.012 111 222.92 0:00:02 0:00:07
draws/s
━━╺━━━━ 652 0 0.011 7 247.78 0:00:02 0:00:06
draws/s
━━╺━━━━ 658 0 0.009 23 250.32 0:00:02 0:00:06
draws/s
━━╺━━━━ 594 0 0.011 223 223.26 0:00:02 0:00:07
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━╺━━━━ 609 0 0.011 15 218.23 0:00:02 0:00:07
draws/s
━━╺━━━━ 660 0 0.009 7 238.38 0:00:02 0:00:06
draws/s
━━╺━━━━ 678 0 0.006 7 242.96 0:00:02 0:00:06
draws/s
━━╺━━━━ 601 0 0.008 151 215.37 0:00:02 0:00:07
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━╺━━━━ 633 0 0.010 55 218.76 0:00:02 0:00:07
draws/s
━━╺━━━━ 681 0 0.011 7 234.88 0:00:02 0:00:06
draws/s
━━╺━━━━ 699 0 0.008 15 241.52 0:00:02 0:00:06
draws/s
━━╺━━━━ 619 0 0.010 183 213.77 0:00:02 0:00:07
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━╺━━━━ 651 0 0.012 7 216.24 0:00:03 0:00:07
draws/s
━━╺━━━━ 705 0 0.011 151 234.17 0:00:03 0:00:06
draws/s
━━╸━━━━ 715 0 0.007 255 237.64 0:00:03 0:00:06
draws/s
━━╺━━━━ 639 0 0.008 143 212.98 0:00:03 0:00:07
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━╺━━━━ 675 0 0.014 7 216.69 0:00:03 0:00:07
draws/s
━━╸━━━━ 726 0 0.012 23 232.98 0:00:03 0:00:06
draws/s
━━╸━━━━ 735 0 0.008 71 236.05 0:00:03 0:00:06
draws/s
━━╺━━━━ 658 0 0.008 15 211.52 0:00:03 0:00:07
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━╺━━━━ 693 0 0.010 79 215.75 0:00:03 0:00:07
draws/s
━━╸━━━━ 749 0 0.011 255 233.19 0:00:03 0:00:06
draws/s
━━╸━━━━ 760 0 0.009 511 236.02 0:00:03 0:00:06
draws/s
━━╺━━━━ 691 0 0.008 103 214.76 0:00:03 0:00:07
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━╺━━━━ 712 0 0.012 207 213.69 0:00:03 0:00:07
draws/s
━━╸━━━━ 767 0 0.009 511 230.17 0:00:03 0:00:06
draws/s
━━╸━━━━ 796 0 0.006 7 238.98 0:00:03 0:00:06
draws/s
━━╸━━━━ 716 0 0.009 39 215.75 0:00:03 0:00:06
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━╸━━━━ 721 0 0.015 551 209.51 0:00:03 0:00:07
draws/s
━━╸━━━━ 798 0 0.012 15 232.12 0:00:03 0:00:06
draws/s
━━╸━━━━ 813 0 0.008 31 236.38 0:00:03 0:00:06
draws/s
━━╸━━━━ 744 0 0.010 7 217.33 0:00:03 0:00:06
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━╸━━━━ 739 0 0.009 7 208.39 0:00:03 0:00:07
draws/s
━━╸━━━━ 809 0 0.010 7 227.90 0:00:03 0:00:06
draws/s
━━╸━━━━ 832 0 0.014 15 235.17 0:00:03 0:00:05
draws/s
━━╸━━━━ 768 0 0.012 7 217.90 0:00:03 0:00:06
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━╸━━━━ 765 0 0.009 31 208.92 0:00:03 0:00:06
draws/s
━━╸━━━━ 841 0 0.009 7 229.74 0:00:03 0:00:06
draws/s
━━━╺━━━ 863 0 0.011 183 235.81 0:00:03 0:00:05
draws/s
━━╸━━━━ 782 0 0.013 15 215.07 0:00:03 0:00:06
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━╸━━━━ 793 0 0.013 7 210.94 0:00:03 0:00:06
draws/s
━━━╺━━━ 865 0 0.010 7 229.33 0:00:03 0:00:05
draws/s
━━━╺━━━ 871 0 0.010 111 230.92 0:00:03 0:00:05
draws/s
━━╸━━━━ 801 0 0.009 751 212.98 0:00:03 0:00:06
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━╸━━━━ 808 0 0.012 511 208.10 0:00:03 0:00:06
draws/s
━━━╺━━━ 897 0 0.010 7 231.19 0:00:03 0:00:05
draws/s
━━━╺━━━ 892 0 0.013 7 230.48 0:00:03 0:00:05
draws/s
━━╸━━━━ 841 0 0.008 15 216.90 0:00:03 0:00:06
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━╸━━━━ 833 0 0.009 7 209.87 0:00:03 0:00:06
draws/s
━━━╺━━━ 919 0 0.008 7 230.21 0:00:03 0:00:05
draws/s
━━━╺━━━ 918 0 0.011 15 230.21 0:00:03 0:00:05
draws/s
━━━╺━━━ 872 0 0.015 519 218.86 0:00:03 0:00:06
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━╺━━━ 858 0 0.010 7 209.66 0:00:04 0:00:06
draws/s
━━━╺━━━ 928 0 0.007 511 226.50 0:00:04 0:00:05
draws/s
━━━╺━━━ 931 0 0.008 15 227.93 0:00:04 0:00:05
draws/s
━━━╺━━━ 893 0 0.014 343 218.05 0:00:04 0:00:06
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━╺━━━ 892 0 0.006 7 211.87 0:00:04 0:00:06
draws/s
━━━╺━━━ 938 0 0.008 7 223.66 0:00:04 0:00:05
draws/s
━━━╺━━━ 947 0 0.006 7 225.16 0:00:04 0:00:05
draws/s
━━━╺━━━ 910 0 0.012 103 217.56 0:00:04 0:00:06
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━╺━━━ 932 0 0.011 327 216.39 0:00:04 0:00:05
draws/s
━━━╺━━━ 985 0 0.008 7 228.81 0:00:04 0:00:05
draws/s
━━━╺━━━ 960 0 0.008 23 223.94 0:00:04 0:00:05
draws/s
━━━╺━━━ 949 0 0.016 95 220.31 0:00:04 0:00:05
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━╺━━━ 943 0 0.009 63 213.85 0:00:04 0:00:05
draws/s
━━━╺━━━ 998 0 0.008 975 225.57 0:00:04 0:00:05
draws/s
━━━╺━━━ 984 0 0.007 7 222.67 0:00:04 0:00:05
draws/s
━━━╺━━━ 961 0 0.013 15 217.65 0:00:04 0:00:05
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━╺━━━ 992 0 0.017 15 218.85 0:00:04 0:00:05
draws/s
━━━╸━━━ 1008 0 0.008 1023 222.54 0:00:04 0:00:05
draws/s
━━━╸━━━ 1009 0 0.008 15 222.80 0:00:04 0:00:05
draws/s
━━━╺━━━ 983 0 0.015 807 217.14 0:00:04 0:00:05
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━╸━━━ 1041 0 0.014 7 223.95 0:00:04 0:00:05
draws/s
━━━╸━━━ 1022 0 0.008 255 220.39 0:00:04 0:00:05
draws/s
━━━╸━━━ 1057 0 0.008 103 228.47 0:00:04 0:00:05
draws/s
━━━╸━━━ 1015 0 0.010 7 219.10 0:00:04 0:00:05
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━╸━━━ 1071 0 0.014 7 225.29 0:00:04 0:00:05
draws/s
━━━╸━━━ 1038 0 0.008 511 218.47 0:00:04 0:00:05
draws/s
━━━╸━━━ 1077 0 0.008 23 226.53 0:00:04 0:00:05
draws/s
━━━╸━━━ 1047 0 0.010 63 220.26 0:00:04 0:00:05
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━╸━━━ 1104 0 0.014 15 227.37 0:00:04 0:00:05
draws/s
━━━╸━━━ 1067 0 0.008 7 219.90 0:00:04 0:00:05
draws/s
━━━╸━━━ 1093 0 0.008 919 225.40 0:00:04 0:00:05
draws/s
━━━╸━━━ 1071 0 0.010 7 220.87 0:00:04 0:00:05
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━╸━━━ 1125 0 0.014 47 226.18 0:00:04 0:00:05
draws/s
━━━╸━━━ 1096 0 0.008 7 220.74 0:00:04 0:00:05
draws/s
━━━╸━━━ 1103 0 0.008 543 222.53 0:00:04 0:00:05
draws/s
━━━╸━━━ 1091 0 0.010 7 220.03 0:00:04 0:00:05
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━╺━━ 1163 0 0.014 7 228.76 0:00:05 0:00:04
draws/s
━━━╸━━━ 1115 0 0.008 103 219.33 0:00:05 0:00:05
draws/s
━━━╸━━━ 1118 0 0.008 7 220.36 0:00:05 0:00:05
draws/s
━━━╸━━━ 1109 0 0.010 7 218.39 0:00:05 0:00:05
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━╺━━ 1187 0 0.014 63 228.59 0:00:05 0:00:04
draws/s
━━━╸━━━ 1134 0 0.008 511 218.60 0:00:05 0:00:05
draws/s
━━━╸━━━ 1127 0 0.008 959 217.10 0:00:05 0:00:05
draws/s
━━━╸━━━ 1134 0 0.010 511 219.14 0:00:05 0:00:05
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━╺━━ 1216 0 0.014 255 229.09 0:00:05 0:00:04
draws/s
━━━━╺━━ 1147 0 0.008 15 217.15 0:00:05 0:00:05
draws/s
━━━━╺━━ 1161 0 0.008 167 219.11 0:00:05 0:00:05
draws/s
━━━━╺━━ 1158 0 0.010 151 218.57 0:00:05 0:00:05
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━╺━━ 1254 0 0.014 15 231.63 0:00:05 0:00:04
draws/s
━━━━╺━━ 1159 0 0.008 71 215.19 0:00:05 0:00:05
draws/s
━━━━╺━━ 1176 0 0.008 7 217.70 0:00:05 0:00:05
draws/s
━━━━╺━━ 1174 0 0.010 495 217.87 0:00:05 0:00:05
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━╸━━ 1315 0 0.014 511 238.39 0:00:05 0:00:04
draws/s
━━━━╺━━ 1169 0 0.008 1023 212.35 0:00:05 0:00:05
draws/s
━━━━╺━━ 1190 0 0.008 7 215.70 0:00:05 0:00:05
draws/s
━━━━╺━━ 1183 0 0.010 71 215.04 0:00:05 0:00:05
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━╸━━ 1343 0 0.014 47 238.48 0:00:05 0:00:03
draws/s
━━━━╺━━ 1197 0 0.008 15 212.53 0:00:05 0:00:05
draws/s
━━━━╺━━ 1211 0 0.008 511 215.15 0:00:05 0:00:05
draws/s
━━━━╺━━ 1199 0 0.010 31 213.26 0:00:05 0:00:05
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━╸━━ 1380 0 0.014 15 240.19 0:00:05 0:00:03
draws/s
━━━━╺━━ 1214 0 0.008 103 211.94 0:00:05 0:00:05
draws/s
━━━━╺━━ 1237 0 0.008 39 215.49 0:00:05 0:00:04
draws/s
━━━━╺━━ 1223 0 0.010 7 213.15 0:00:05 0:00:05
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━╸━━ 1391 0 0.014 119 237.61 0:00:05 0:00:03
draws/s
━━━━╺━━ 1248 0 0.008 15 213.50 0:00:05 0:00:05
draws/s
━━━━╺━━ 1253 0 0.008 159 214.23 0:00:05 0:00:04
draws/s
━━━━╺━━ 1250 0 0.010 135 213.88 0:00:05 0:00:04
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━╸━━ 1405 0 0.014 7 235.53 0:00:05 0:00:03
draws/s
━━━━╺━━ 1265 0 0.008 799 212.46 0:00:05 0:00:04
draws/s
━━━━╺━━ 1272 0 0.008 7 213.55 0:00:05 0:00:04
draws/s
━━━━╺━━ 1266 0 0.010 103 213.28 0:00:05 0:00:04
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━╸━━ 1426 0 0.014 15 234.68 0:00:06 0:00:03
draws/s
━━━━╺━━ 1274 0 0.008 703 210.12 0:00:06 0:00:04
draws/s
━━━━╸━━ 1291 0 0.008 7 212.96 0:00:06 0:00:04
draws/s
━━━━╸━━ 1303 0 0.010 7 214.62 0:00:06 0:00:04
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━╺━ 1454 0 0.014 87 235.24 0:00:06 0:00:03
draws/s
━━━━╸━━ 1289 0 0.008 7 209.21 0:00:06 0:00:04
draws/s
━━━━╸━━ 1305 0 0.008 511 211.34 0:00:06 0:00:04
draws/s
━━━━╸━━ 1322 0 0.010 1023 214.74 0:00:06 0:00:04
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━╺━ 1487 0 0.014 7 236.60 0:00:06 0:00:03
draws/s
━━━━╸━━ 1303 0 0.008 143 207.20 0:00:06 0:00:04
draws/s
━━━━╸━━ 1312 0 0.008 7 209.40 0:00:06 0:00:04
draws/s
━━━━╸━━ 1339 0 0.010 207 213.08 0:00:06 0:00:04
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━╺━ 1502 0 0.014 7 234.95 0:00:06 0:00:03
draws/s
━━━━╸━━ 1332 0 0.008 7 208.28 0:00:06 0:00:04
draws/s
━━━━╸━━ 1328 0 0.008 7 207.96 0:00:06 0:00:04
draws/s
━━━━╸━━ 1374 0 0.010 7 215.00 0:00:06 0:00:04
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━╺━ 1530 0 0.014 143 235.66 0:00:06 0:00:02
draws/s
━━━━╸━━ 1364 0 0.008 191 209.55 0:00:06 0:00:04
draws/s
━━━━╸━━ 1352 0 0.008 847 208.39 0:00:06 0:00:04
draws/s
━━━━╸━━ 1397 0 0.010 7 214.95 0:00:06 0:00:04
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━╸━ 1578 0 0.014 7 238.39 0:00:06 0:00:02
draws/s
━━━━╸━━ 1379 0 0.008 183 208.99 0:00:06 0:00:04
draws/s
━━━━╸━━ 1376 0 0.008 399 208.09 0:00:06 0:00:04
draws/s
━━━━╸━━ 1424 0 0.010 7 215.07 0:00:06 0:00:03
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━╸━ 1610 0 0.014 983 239.32 0:00:06 0:00:02
draws/s
━━━━╸━━ 1408 0 0.008 7 209.02 0:00:06 0:00:04
draws/s
━━━━╸━━ 1394 0 0.008 255 207.54 0:00:06 0:00:04
draws/s
━━━━━╺━ 1448 0 0.010 7 215.23 0:00:06 0:00:03
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━╸━ 1644 0 0.014 47 240.35 0:00:06 0:00:02
draws/s
━━━━━╺━ 1433 0 0.008 7 209.31 0:00:06 0:00:03
draws/s
━━━━╸━━ 1413 0 0.008 7 206.77 0:00:06 0:00:04
draws/s
━━━━━╺━ 1472 0 0.010 7 215.38 0:00:06 0:00:03
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━╸━ 1676 0 0.014 15 240.94 0:00:06 0:00:02
draws/s
━━━━━╺━ 1468 0 0.008 15 211.09 0:00:06 0:00:03
draws/s
━━━━━╺━ 1436 0 0.008 103 206.59 0:00:06 0:00:04
draws/s
━━━━━╺━ 1488 0 0.010 223 214.19 0:00:06 0:00:03
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━╸━ 1710 0 0.014 279 242.08 0:00:07 0:00:02
draws/s
━━━━━╺━ 1479 0 0.008 15 210.08 0:00:07 0:00:03
draws/s
━━━━━╺━ 1452 0 0.008 7 205.70 0:00:07 0:00:03
draws/s
━━━━━╺━ 1509 0 0.010 71 214.15 0:00:07 0:00:03
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━╺ 1735 0 0.014 111 241.94 0:00:07 0:00:01
draws/s
━━━━━╺━ 1489 0 0.008 15 208.21 0:00:07 0:00:03
draws/s
━━━━━╺━ 1474 0 0.008 263 205.51 0:00:07 0:00:03
draws/s
━━━━━╺━ 1529 0 0.010 7 213.50 0:00:07 0:00:03
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━╺ 1764 0 0.014 103 242.79 0:00:07 0:00:01
draws/s
━━━━━╺━ 1511 0 0.008 303 207.51 0:00:07 0:00:03
draws/s
━━━━━╺━ 1486 0 0.008 7 204.39 0:00:07 0:00:03
draws/s
━━━━━╺━ 1554 0 0.010 7 213.45 0:00:07 0:00:03
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━╺ 1786 0 0.014 7 241.56 0:00:07 0:00:01
draws/s
━━━━━╺━ 1538 0 0.008 7 208.28 0:00:07 0:00:03
draws/s
━━━━━╺━ 1515 0 0.008 39 205.02 0:00:07 0:00:03
draws/s
━━━━━╸━ 1581 0 0.010 15 214.22 0:00:07 0:00:03
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━╺ 1793 0 0.014 1023 239.60 0:00:07 0:00:01
draws/s
━━━━━╺━ 1562 0 0.008 7 208.32 0:00:07 0:00:03
draws/s
━━━━━╺━ 1538 0 0.008 7 205.07 0:00:07 0:00:03
draws/s
━━━━━╸━ 1596 0 0.010 679 212.88 0:00:07 0:00:02
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━╺ 1812 0 0.014 47 238.16 0:00:07 0:00:01
draws/s
━━━━━╸━ 1572 0 0.008 167 207.20 0:00:07 0:00:03
draws/s
━━━━━╺━ 1561 0 0.008 111 205.45 0:00:07 0:00:03
draws/s
━━━━━╸━ 1625 0 0.010 15 214.10 0:00:07 0:00:02
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━╺ 1831 0 0.014 975 237.29 0:00:07 0:00:01
draws/s
━━━━━╸━ 1596 0 0.008 15 206.71 0:00:07 0:00:03
draws/s
━━━━━╸━ 1584 0 0.008 415 205.39 0:00:07 0:00:03
draws/s
━━━━━╸━ 1651 0 0.010 103 213.98 0:00:07 0:00:02
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━╸ 1858 0 0.014 7 237.79 0:00:07 0:00:01
draws/s
━━━━━╸━ 1645 0 0.008 7 210.20 0:00:07 0:00:02
draws/s
━━━━━╸━ 1600 0 0.008 7 204.80 0:00:07 0:00:03
draws/s
━━━━━╸━ 1678 0 0.010 23 214.74 0:00:07 0:00:02
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━╸ 1884 0 0.014 719 237.41 0:00:07 0:00:01
draws/s
━━━━━╸━ 1664 0 0.008 7 209.74 0:00:07 0:00:02
draws/s
━━━━━╸━ 1625 0 0.008 191 204.91 0:00:07 0:00:03
draws/s
━━━━━╸━ 1706 0 0.010 7 215.20 0:00:07 0:00:02
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━╸ 1930 0 0.014 31 239.88 0:00:08 0:00:01
draws/s
━━━━━╸━ 1697 0 0.008 487 211.28 0:00:08 0:00:02
draws/s
━━━━━╸━ 1641 0 0.008 207 204.07 0:00:08 0:00:02
draws/s
━━━━━━╺ 1720 0 0.010 151 214.07 0:00:08 0:00:02
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━╸ 1957 0 0.014 7 240.07 0:00:08 0:00:01
draws/s
━━━━━━╺ 1719 0 0.008 7 211.54 0:00:08 0:00:02
draws/s
━━━━━╸━ 1653 0 0.008 111 202.82 0:00:08 0:00:02
draws/s
━━━━━━╺ 1736 0 0.010 15 213.13 0:00:08 0:00:02
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.014 7 241.86 0:00:08 0:00:00
draws/s
━━━━━━╺ 1746 0 0.008 519 211.38 0:00:08 0:00:02
draws/s
━━━━━╸━ 1683 0 0.008 7 203.74 0:00:08 0:00:02
draws/s
━━━━━━╺ 1757 0 0.010 7 212.75 0:00:08 0:00:02
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.014 7 241.86 0:00:08 0:00:00
draws/s
━━━━━━╺ 1747 0 0.008 215 211.32 0:00:08 0:00:02
draws/s
━━━━━╸━ 1684 0 0.008 183 203.74 0:00:08 0:00:02
draws/s
━━━━━━╺ 1759 0 0.010 7 212.88 0:00:08 0:00:02
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.014 7 241.86 0:00:08 0:00:00
draws/s
━━━━━━╺ 1769 0 0.008 15 211.61 0:00:08 0:00:02
draws/s
━━━━━╸━ 1707 0 0.008 79 204.04 0:00:08 0:00:02
draws/s
━━━━━━╺ 1784 0 0.010 15 213.52 0:00:08 0:00:02
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.014 7 241.86 0:00:08 0:00:00
draws/s
━━━━━━╺ 1784 0 0.008 31 210.93 0:00:08 0:00:02
draws/s
━━━━━━╺ 1727 0 0.008 15 203.83 0:00:08 0:00:02
draws/s
━━━━━━╺ 1804 0 0.010 471 213.04 0:00:08 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.014 7 241.86 0:00:08 0:00:00
draws/s
━━━━━━╺ 1798 0 0.008 511 209.85 0:00:08 0:00:02
draws/s
━━━━━━╺ 1745 0 0.008 151 203.42 0:00:08 0:00:02
draws/s
━━━━━━╺ 1840 0 0.010 63 214.47 0:00:08 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.014 7 241.86 0:00:08 0:00:00
draws/s
━━━━━━╺ 1817 0 0.008 23 209.46 0:00:08 0:00:01
draws/s
━━━━━━╺ 1764 0 0.008 95 203.01 0:00:08 0:00:02
draws/s
━━━━━━╸ 1867 0 0.010 7 215.12 0:00:08 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.014 7 241.86 0:00:08 0:00:00
draws/s
━━━━━━╺ 1830 0 0.008 7 208.01 0:00:08 0:00:01
draws/s
━━━━━━╺ 1792 0 0.008 79 204.14 0:00:08 0:00:02
draws/s
━━━━━━╸ 1887 0 0.010 791 214.52 0:00:08 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.014 7 241.86 0:00:08 0:00:00
draws/s
━━━━━━╺ 1842 0 0.008 7 206.84 0:00:08 0:00:01
draws/s
━━━━━━╺ 1819 0 0.008 7 204.56 0:00:08 0:00:01
draws/s
━━━━━━╸ 1906 0 0.010 15 214.22 0:00:08 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.014 7 241.86 0:00:08 0:00:00
draws/s
━━━━━━╸ 1864 0 0.008 1023 206.93 0:00:09 0:00:01
draws/s
━━━━━━╺ 1839 0 0.008 7 204.15 0:00:09 0:00:01
draws/s
━━━━━━╸ 1933 0 0.010 47 214.62 0:00:09 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.014 7 241.86 0:00:08 0:00:00
draws/s
━━━━━━╸ 1897 0 0.008 47 207.99 0:00:09 0:00:01
draws/s
━━━━━━╸ 1859 0 0.008 511 204.10 0:00:09 0:00:01
draws/s
━━━━━━╸ 1956 0 0.010 63 214.63 0:00:09 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.014 7 241.86 0:00:08 0:00:00
draws/s
━━━━━━╸ 1930 0 0.008 15 209.20 0:00:09 0:00:01
draws/s
━━━━━━╸ 1899 0 0.008 23 206.24 0:00:09 0:00:01
draws/s
━━━━━━╸ 1986 0 0.010 15 215.44 0:00:09 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.014 7 241.86 0:00:08 0:00:00
draws/s
━━━━━━╸ 1956 0 0.008 7 210.21 0:00:09 0:00:01
draws/s
━━━━━━╸ 1918 0 0.008 71 206.11 0:00:09 0:00:01
draws/s
━━━━━━━ 2000 0 0.010 39 214.89 0:00:09 0:00:00
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.014 7 241.86 0:00:08 0:00:00
draws/s
━━━━━━╸ 1960 0 0.008 15 210.10 0:00:09 0:00:01
draws/s
━━━━━━╸ 1929 0 0.008 7 206.74 0:00:09 0:00:01
draws/s
━━━━━━━ 2000 0 0.010 39 214.89 0:00:09 0:00:00
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.014 7 241.86 0:00:08 0:00:00
draws/s
━━━━━━╸ 1981 0 0.008 7 209.80 0:00:09 0:00:01
draws/s
━━━━━━╸ 1939 0 0.008 575 205.45 0:00:09 0:00:01
draws/s
━━━━━━━ 2000 0 0.010 39 214.89 0:00:09 0:00:00
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.014 7 241.86 0:00:08 0:00:00
draws/s
━━━━━━━ 2000 0 0.008 199 210.62 0:00:09 0:00:00
draws/s
━━━━━━╸ 1942 0 0.008 863 204.67 0:00:09 0:00:01
draws/s
━━━━━━━ 2000 0 0.010 39 214.89 0:00:09 0:00:00
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.014 7 241.86 0:00:08 0:00:00
draws/s
━━━━━━━ 2000 0 0.008 199 210.62 0:00:09 0:00:00
draws/s
━━━━━━╸ 1958 0 0.008 143 205.42 0:00:09 0:00:01
draws/s
━━━━━━━ 2000 0 0.010 39 214.89 0:00:09 0:00:00
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.014 7 241.86 0:00:08 0:00:00
draws/s
━━━━━━━ 2000 0 0.008 199 210.62 0:00:09 0:00:00
draws/s
━━━━━━╸ 1980 0 0.008 359 205.31 0:00:09 0:00:01
draws/s
━━━━━━━ 2000 0 0.010 39 214.89 0:00:09 0:00:00
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.014 7 241.86 0:00:08 0:00:00
draws/s
━━━━━━━ 2000 0 0.008 199 210.62 0:00:09 0:00:00
draws/s
━━━━━━╸ 1996 0 0.008 863 204.45 0:00:09 0:00:01
draws/s
━━━━━━━ 2000 0 0.010 39 214.89 0:00:09 0:00:00
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.014 7 241.86 0:00:08 0:00:00
draws/s
━━━━━━━ 2000 0 0.008 199 210.62 0:00:09 0:00:00
draws/s
━━━━━━━ 2000 0 0.008 263 204.28 0:00:09 0:00:00
draws/s
━━━━━━━ 2000 0 0.010 39 214.89 0:00:09 0:00:00
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.014 7 241.86 0:00:08 0:00:00
draws/s
━━━━━━━ 2000 0 0.008 199 210.62 0:00:09 0:00:00
draws/s
━━━━━━━ 2000 0 0.008 263 204.28 0:00:09 0:00:00
draws/s
━━━━━━━ 2000 0 0.010 39 214.89 0:00:09 0:00:00
draws/s
?25hSampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 10 seconds.
The rhat statistic is larger than 1.01 for some parameters. This indicates problems during sampling. See https://arxiv.org/abs/1903.08008 for details
The effective sample size per chain is smaller than 100 for some parameters. A higher number is needed for reliable rhat and ess computation. See https://arxiv.org/abs/1903.08008 for details
Sampling: [beta, y_hat, y_hat_sigma]
Sampling: [y_hat]
Sampling: [y_hat]
Sampling: [y_hat]
Sampling: [y_hat]
pre-intervention R2 = 0.746
| mean | sd | hdi_3% | hdi_97% | mcse_mean | mcse_sd | ess_bulk | ess_tail | r_hat | |
|---|---|---|---|---|---|---|---|---|---|
| beta[unit_0, C(month)[10]] | 1.025 | 10.688 | -19.517 | 20.509 | 0.842 | 0.586 | 157.0 | 178.0 | 1.03 |
| beta[unit_0, C(month)[5]] | 0.946 | 10.689 | -19.762 | 20.365 | 0.842 | 0.585 | 157.0 | 177.0 | 1.03 |
| beta[unit_0, C(month)[6]] | 0.989 | 10.691 | -19.363 | 20.621 | 0.844 | 0.585 | 157.0 | 179.0 | 1.03 |
| beta[unit_0, C(month)[7]] | 1.258 | 10.691 | -19.109 | 20.911 | 0.845 | 0.585 | 157.0 | 177.0 | 1.03 |
| beta[unit_0, C(month)[8]] | 0.842 | 10.690 | -19.383 | 20.541 | 0.844 | 0.585 | 157.0 | 173.0 | 1.03 |
The fit improved, but the sampler is unhappy. For at least some of our parameters, the chains have not converged and our effective sample size is weak. Notice also that the estimated standard deviations have almost exactly the same value. Together, these are good hints for a flat ridge in the posterior geometry.
# Zoom in on a few of the spline coefficients to inspect the posterior geometry.
idata = result_ridge.model.idata
spline_coeffs = [
c for c in idata.posterior["beta"].coords["coeffs"].values if c.startswith("cr(t")
]
az.plot_pair(
idata,
group="posterior",
var_names=["beta"],
coords={"treated_units": "unit_0", "coeffs": spline_coeffs[:4]},
kind="scatter",
divergences=True,
marginals=True,
scatter_kwargs={"alpha": 0.1},
figsize=(9, 9),
)
array([[<Axes: ylabel='beta\ncr(t, df=6)[0]'>, <Axes: >, <Axes: >,
<Axes: >],
[<Axes: ylabel='beta\ncr(t, df=6)[1]'>, <Axes: >, <Axes: >,
<Axes: >],
[<Axes: ylabel='beta\ncr(t, df=6)[2]'>, <Axes: >, <Axes: >,
<Axes: >],
[<Axes: xlabel='beta\ncr(t, df=6)[0]', ylabel='beta\ncr(t, df=6)[3]'>,
<Axes: xlabel='beta\ncr(t, df=6)[1]'>,
<Axes: xlabel='beta\ncr(t, df=6)[2]'>,
<Axes: xlabel='beta\ncr(t, df=6)[3]'>]], dtype=object)
In this case, the cause is a rank deficiency in the design matrix. Since cr() is a cardinal basis, its rows sum to one, which means a simple constant function lies in its span. With 0 + in the formula, patsy also gives C(month) all twelve levels rather than eleven, so the month variables also span the constant function. In effect, we’ve accidentally represented the intercept twice, so only the sum of the two levels is identified by the data.
X = result_ridge.pre_design["X"]
spline_cols = [c for c in X.coeffs.values if c.startswith("cr(")]
month_cols = [c for c in X.coeffs.values if c.startswith("C(month)")]
print(f"design matrix rank {np.linalg.matrix_rank(X.values)} of {X.shape[1]} columns")
print(f"cr() row sums: {X.sel(coeffs=spline_cols).sum('coeffs').values.min():.6f}")
design matrix rank 18 of 19 columns
cr() row sums: 1.000000
Bayesian P-spline with a second-order random walk prior#
Okay, so we need to address this redundant intercept. While we’re at it, let’s improve the prior model for our spline too! Rather than treat each spline coefficient as exchangeable, let’s give them more informative, regularizing priors. A P-spline does this by penalising the second differences of neighboring coefficients. One Bayesian implementation of this is the second-order random walk (RW2) prior! $\( \beta_j - 2\beta_{j-1} + \beta_{j-2} \sim \operatorname{Normal}(0,\sigma^2_{\text{smooth}}) \)$
Because cr() coefficients are the curve’s values at the knots, this prior model implies that the trend in t has little curvature, if any, with \(\sigma_{\text{smooth}}\), learned from the data, controlling how much. P-splines let us be really generous with the df and let this penalty do the heavy lifting.
The null space#
Writing the second differences as \(D\beta\), the prior is \(\beta \sim \mathrm{Normal}(0, \sigma_{\text{smooth}}^2 P^{-})\) with \(P = D^\top D\). Here, \(P\) is rank deficient by two because adding a constant or a straight line to \(\beta\) doesn’t change the second difference. We’ll drop the constant, since we know we need to eliminate this redundancy, and let C(month) carry that level. We can keep the linear direction though, as a learned parameter, spline_slope. This lets our model’s counterfactual extrapolate with a reasonable slope once we run past the last knot, adding a bit more flexibility.
Everything else we’ll parameterize non-centered through the whitened penalty \(Z = V_{+} \Lambda_{+}^{-1/2}\), built from the positive part of the eigendecomposition of \(P\). Then \(\operatorname{Cov}(Zz) = P^{-}\) exactly, the columns of \(Z\) are orthogonal to both null directions, and the sampler sees standard normal variables instead of a funnel. Finally, normalizing the eigenvalues by their geometric mean keeps \(\sigma_{\text{smooth}}\) meaning the same thing even if we vary df (which will come in handy later!).
def rw2_basis(n_knots: int) -> np.ndarray:
"""Whitened RW2 penalty, with constant and linear null directions removed.
Returns an ``(n_knots, n_knots - 2)`` matrix ``Z`` such that ``Z @ z`` with
``z ~ Normal(0, 1)`` has covariance equal to the pseudo-inverse of the second
difference penalty, rescaled so that the geometric mean of the marginal
variances is one. This rescaling lets the smoothing standard deviation mean
the same thing regardless of the knot count.
"""
D = np.diff(np.eye(n_knots), n=2, axis=0)
eigenvalues, eigenvectors = np.linalg.eigh(D.T @ D)
keep = eigenvalues > 1e-8 * eigenvalues.max()
lam, V = eigenvalues[keep], eigenvectors[:, keep]
Z = V / np.sqrt(lam)
marginal_variance = (Z**2).sum(axis=1)
return Z / np.sqrt(np.exp(np.mean(np.log(marginal_variance))))
Thankfully, pymc-extras lets us specify this prior, with a little help from the VariableFactory protocol. Any object with a dims attribute and a create_variable method can be swapped in wherever Prior can. CausalPy’s LinearRegression builds coefficients via self.priors["beta"].create_variables("beta"), so we can pass in our own custom prior here!
The variable factory reads the column names from the model’s coeffs coordinate, splits them into spline and non-spline blocks, and reassembles them into a beta container for CausalPy.
class RW2SplinePrior:
"""RW2 (P-spline) prior on the spline columns, Normal on everything else.
Implements the ``VariableFactory`` protocol from ``pymc_extras.prior``, so it can
be passed as the ``beta`` prior of any CausalPy model that builds its coefficients
with ``create_variable``. Columns whose name starts with ``spline_prefix`` get the
smoothness prior; the rest keep an ordinary normal.
"""
def __init__(
self,
spline_prefix="cr(",
smoothing_sigma=0.5,
slope_sigma=2.0,
beta_sigma=1.0,
dims=("treated_units", "coeffs"),
):
self.spline_prefix = spline_prefix
self.smoothing_sigma = smoothing_sigma
self.slope_sigma = slope_sigma
self.beta_sigma = beta_sigma
self.dims = tuple(dims)
def create_variable(self, name, xdist=False):
model = pm.modelcontext(None)
labels = list(model.coords["coeffs"])
spline_idx = np.array(
[
i
for i, label in enumerate(labels)
if label.startswith(self.spline_prefix)
]
)
if spline_idx.size < 4:
raise ValueError(
f"Need at least 4 columns starting with {self.spline_prefix!r}, "
f"found {spline_idx.size} in {labels}."
)
other_idx = np.array(
[i for i in range(len(labels)) if i not in set(spline_idx.tolist())]
)
n_knots = len(spline_idx)
Z = rw2_basis(n_knots)
ramp = np.arange(n_knots) - (n_knots - 1) / 2
ramp = ramp / np.linalg.norm(ramp) # the trend
model.add_coords(
{
"spline_coeffs": [labels[i] for i in spline_idx],
"other_coeffs": [labels[i] for i in other_idx],
"wiggle": np.arange(n_knots - 2),
}
)
# RW2 block gets the smoothness scale, whitened wiggliness, explicit slope
sigma_smooth = pm.HalfNormal(
"sigma_smooth", self.smoothing_sigma, dims="treated_units"
)
z = pm.Normal("z_spline", 0.0, 1.0, dims=["treated_units", "wiggle"])
slope = pm.Normal("spline_slope", 0.0, self.slope_sigma, dims="treated_units")
beta_spline = pm.Deterministic(
"beta_spline",
sigma_smooth[:, None] * pt.dot(z, Z.T) + slope[:, None] * ramp[None, :],
dims=["treated_units", "spline_coeffs"],
)
# everything else keeps an ordinary normal prior.
beta_other = pm.Normal(
"beta_other", 0.0, self.beta_sigma, dims=["treated_units", "other_coeffs"]
)
# packing these together for the `beta` container CausalPy needs
beta = pt.zeros((len(model.coords["treated_units"]), len(labels)))
beta = pt.set_subtensor(beta[:, spline_idx], beta_spline)
beta = pt.set_subtensor(beta[:, other_idx], beta_other)
return pm.Deterministic(name, beta, dims=list(self.dims))
def pspline_model(**sample_kwargs):
return cp.pymc_models.LinearRegression(
sample_kwargs={"random_seed": seed, "target_accept": 0.95, **sample_kwargs},
priors={
"beta": RW2SplinePrior(),
"y_hat": Prior(
"Normal",
sigma=Prior("HalfNormal", sigma=0.5, dims=["treated_units"]),
dims=["obs_ind", "treated_units"],
),
},
)
Because the penalty determines the smoothness, we can afford far more knots than the six we tried earlier. Let’s double it to 12.
formula = "standardize(deaths) ~ 0 + cr(t, df=12) + C(month) + standardize(temp)"
result = cp.InterruptedTimeSeries(
df,
treatment_time,
formula=formula,
model=pspline_model(),
)
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [sigma_smooth, z_spline, spline_slope, beta_other, y_hat_sigma]
?25l
Step Grad
Progress Draw Diverge… size evals Speed Elapsed Remain…
──────────────────────────────────────────────────────────────────────────────
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
━━━━━━━ 0 0 0.000 0 0.00 0:00:00 -:--:--
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 16 0 0.023 31 156.97 0:00:00 0:00:13
draws/s
━━━━━━━ 10 0 0.009 191 105.00 0:00:00 0:00:20
draws/s
━━━━━━━ 14 0 0.015 255 145.73 0:00:00 0:00:14
draws/s
━━━━━━━ 10 0 0.009 255 122.44 0:00:00 0:00:17
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 30 0 0.027 127 144.12 0:00:00 0:00:14
draws/s
━━━━━━━ 24 0 0.038 127 113.65 0:00:00 0:00:18
draws/s
━━━━━━━ 32 0 0.016 255 156.29 0:00:00 0:00:13
draws/s
━━━━━━━ 25 0 0.039 79 121.58 0:00:00 0:00:17
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 56 0 0.022 127 174.30 0:00:00 0:00:12
draws/s
━━━━━━━ 45 0 0.035 95 140.47 0:00:00 0:00:14
draws/s
━━━━━━━ 45 0 0.037 127 141.04 0:00:00 0:00:14
draws/s
━━━━━━━ 50 0 0.029 127 158.12 0:00:00 0:00:13
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 76 0 0.043 63 178.09 0:00:00 0:00:11
draws/s
━━━━━━━ 62 0 0.033 127 146.66 0:00:00 0:00:14
draws/s
━━━━━━━ 68 0 0.023 63 160.96 0:00:00 0:00:12
draws/s
━━━━━━━ 72 0 0.021 127 171.20 0:00:00 0:00:12
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 98 0 0.032 127 182.34 0:00:00 0:00:11
draws/s
━━━━━━━ 83 0 0.035 127 155.77 0:00:00 0:00:13
draws/s
━━━━━━━ 88 0 0.029 127 165.32 0:00:00 0:00:12
draws/s
━━━━━━━ 93 0 0.022 127 175.01 0:00:00 0:00:11
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
╸━━━━━━ 158 0 0.217 15 244.84 0:00:00 0:00:08
draws/s
━━━━━━━ 103 0 0.022 127 160.63 0:00:00 0:00:12
draws/s
━━━━━━━ 132 0 0.170 15 205.80 0:00:00 0:00:10
draws/s
━━━━━━━ 116 0 0.108 31 181.40 0:00:00 0:00:11
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
╸━━━━━━ 233 0 0.177 31 308.47 0:00:00 0:00:06
draws/s
━━━━━━━ 137 0 0.185 15 181.76 0:00:00 0:00:11
draws/s
╸━━━━━━ 212 0 0.164 15 282.47 0:00:00 0:00:07
draws/s
╸━━━━━━ 211 0 0.269 15 281.71 0:00:00 0:00:07
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━╺━━━━━ 324 0 0.110 31 374.81 0:00:00 0:00:05
draws/s
╸━━━━━━ 251 0 0.180 15 290.85 0:00:00 0:00:06
draws/s
━╺━━━━━ 289 0 0.180 15 335.83 0:00:00 0:00:06
draws/s
━╺━━━━━ 309 0 0.171 15 360.06 0:00:00 0:00:05
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━╺━━━━━ 400 0 0.126 31 411.37 0:00:00 0:00:04
draws/s
━╺━━━━━ 342 0 0.156 15 352.14 0:00:00 0:00:05
draws/s
━╺━━━━━ 385 0 0.173 31 397.09 0:00:00 0:00:05
draws/s
━╺━━━━━ 399 0 0.172 31 412.99 0:00:00 0:00:04
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━╸━━━━━ 477 0 0.207 15 440.22 0:00:01 0:00:04
draws/s
━╸━━━━━ 434 0 0.170 31 401.31 0:00:01 0:00:04
draws/s
━╸━━━━━ 473 0 0.179 15 437.92 0:00:01 0:00:04
draws/s
━╸━━━━━ 466 0 0.133 15 432.41 0:00:01 0:00:04
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━╸━━━━━ 557 0 0.175 15 467.99 0:00:01 0:00:04
draws/s
━╸━━━━━ 514 0 0.179 31 432.17 0:00:01 0:00:04
draws/s
━╸━━━━━ 561 0 0.147 31 472.58 0:00:01 0:00:04
draws/s
━╸━━━━━ 536 0 0.156 31 452.25 0:00:01 0:00:04
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━╺━━━━ 646 0 0.217 15 497.90 0:00:01 0:00:03
draws/s
━━╺━━━━ 611 0 0.162 31 472.37 0:00:01 0:00:03
draws/s
━━╺━━━━ 633 0 0.123 31 489.96 0:00:01 0:00:03
draws/s
━━╺━━━━ 622 0 0.153 31 481.57 0:00:01 0:00:03
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━╸━━━━ 740 0 0.197 31 527.96 0:00:01 0:00:03
draws/s
━━╺━━━━ 691 0 0.175 31 493.72 0:00:01 0:00:03
draws/s
━━╺━━━━ 712 0 0.141 31 509.47 0:00:01 0:00:03
draws/s
━━╸━━━━ 718 0 0.126 31 514.29 0:00:01 0:00:03
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━╸━━━━ 818 0 0.160 15 540.93 0:00:01 0:00:03
draws/s
━━╸━━━━ 780 0 0.184 31 516.47 0:00:01 0:00:03
draws/s
━━╸━━━━ 783 0 0.163 31 519.49 0:00:01 0:00:03
draws/s
━━╸━━━━ 795 0 0.104 31 527.97 0:00:01 0:00:03
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━╺━━━ 907 0 0.169 31 558.50 0:00:01 0:00:02
draws/s
━━╸━━━━ 851 0 0.192 15 524.98 0:00:01 0:00:03
draws/s
━━━╺━━━ 870 0 0.153 15 537.17 0:00:01 0:00:03
draws/s
━━╸━━━━ 845 0 0.149 31 522.32 0:00:01 0:00:03
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━╺━━━ 992 0 0.201 15 573.35 0:00:01 0:00:02
draws/s
━━━╺━━━ 932 0 0.209 15 539.62 0:00:01 0:00:02
draws/s
━━━╺━━━ 937 0 0.173 15 542.85 0:00:01 0:00:02
draws/s
━━━╺━━━ 914 0 0.183 31 530.12 0:00:01 0:00:03
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━╸━━━ 1071 0 0.214 15 582.65 0:00:01 0:00:02
draws/s
━━━╸━━━ 1030 0 0.157 15 560.75 0:00:01 0:00:02
draws/s
━━━╸━━━ 1014 0 0.182 31 552.47 0:00:01 0:00:02
draws/s
━━━╺━━━ 986 0 0.157 31 538.13 0:00:01 0:00:02
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━╺━━ 1155 0 0.214 31 592.83 0:00:01 0:00:02
draws/s
━━━╸━━━ 1121 0 0.157 15 575.96 0:00:01 0:00:02
draws/s
━━━╸━━━ 1090 0 0.182 31 560.83 0:00:01 0:00:02
draws/s
━━━╸━━━ 1045 0 0.185 15 538.11 0:00:01 0:00:02
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━╺━━ 1236 0 0.214 31 601.24 0:00:02 0:00:01
draws/s
━━━━╺━━ 1214 0 0.157 15 590.97 0:00:02 0:00:01
draws/s
━━━━╺━━ 1160 0 0.182 15 565.34 0:00:02 0:00:02
draws/s
━━━╸━━━ 1107 0 0.185 47 540.23 0:00:02 0:00:02
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━╸━━ 1329 0 0.214 31 615.14 0:00:02 0:00:01
draws/s
━━━━╸━━ 1312 0 0.157 15 607.92 0:00:02 0:00:01
draws/s
━━━━╺━━ 1238 0 0.182 15 574.08 0:00:02 0:00:02
draws/s
━━━━╺━━ 1184 0 0.185 15 549.36 0:00:02 0:00:02
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━╸━━ 1413 0 0.214 31 622.28 0:00:02 0:00:01
draws/s
━━━━╸━━ 1413 0 0.157 15 622.60 0:00:02 0:00:01
draws/s
━━━━╸━━ 1319 0 0.182 31 582.46 0:00:02 0:00:01
draws/s
━━━━╺━━ 1246 0 0.185 31 550.30 0:00:02 0:00:02
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━╺━ 1499 0 0.214 15 629.72 0:00:02 0:00:01
draws/s
━━━━━╺━ 1514 0 0.157 31 636.55 0:00:02 0:00:01
draws/s
━━━━╸━━ 1393 0 0.182 15 586.03 0:00:02 0:00:01
draws/s
━━━━╸━━ 1313 0 0.185 31 553.03 0:00:02 0:00:02
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━╸━ 1587 0 0.214 15 637.84 0:00:02 0:00:01
draws/s
━━━━━╸━ 1615 0 0.157 31 649.59 0:00:02 0:00:01
draws/s
━━━━━╺━ 1469 0 0.182 31 591.35 0:00:02 0:00:01
draws/s
━━━━╸━━ 1382 0 0.185 31 556.67 0:00:02 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━╸━ 1689 0 0.214 15 650.26 0:00:02 0:00:01
draws/s
━━━━━━╺ 1720 0 0.157 15 662.73 0:00:02 0:00:01
draws/s
━━━━━╺━ 1553 0 0.182 15 598.77 0:00:02 0:00:01
draws/s
━━━━━╺━ 1451 0 0.185 15 559.99 0:00:02 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━╺ 1785 0 0.214 15 659.56 0:00:02 0:00:01
draws/s
━━━━━━╺ 1832 0 0.157 15 677.48 0:00:02 0:00:01
draws/s
━━━━━╸━ 1634 0 0.182 31 604.64 0:00:02 0:00:01
draws/s
━━━━━╺━ 1526 0 0.185 15 565.21 0:00:02 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━╸ 1880 0 0.214 15 668.68 0:00:02 0:00:01
draws/s
━━━━━━╸ 1937 0 0.157 15 689.34 0:00:02 0:00:01
draws/s
━━━━━╸━ 1714 0 0.182 31 610.61 0:00:02 0:00:01
draws/s
━━━━━╸━ 1599 0 0.185 15 569.82 0:00:02 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━╸ 1943 0 0.214 31 672.76 0:00:02 0:00:01
draws/s
━━━━━━━ 2000 0 0.157 15 692.61 0:00:02 0:00:00
draws/s
━━━━━━╺ 1768 0 0.182 31 613.09 0:00:02 0:00:01
draws/s
━━━━━╸━ 1650 0 0.185 15 572.38 0:00:02 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━╸ 1972 0 0.214 15 674.58 0:00:02 0:00:01
draws/s
━━━━━━━ 2000 0 0.157 15 692.61 0:00:02 0:00:00
draws/s
━━━━━━╺ 1794 0 0.182 31 614.63 0:00:02 0:00:01
draws/s
━━━━━╸━ 1672 0 0.185 31 573.11 0:00:02 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.214 31 676.39 0:00:02 0:00:00
draws/s
━━━━━━━ 2000 0 0.157 15 692.61 0:00:02 0:00:00
draws/s
━━━━━━╺ 1821 0 0.182 31 616.62 0:00:02 0:00:01
draws/s
━━━━━╸━ 1696 0 0.185 31 574.81 0:00:02 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.214 31 676.39 0:00:02 0:00:00
draws/s
━━━━━━━ 2000 0 0.157 15 692.61 0:00:02 0:00:00
draws/s
━━━━━━╸ 1879 0 0.182 15 620.72 0:00:03 0:00:01
draws/s
━━━━━━╺ 1748 0 0.185 31 577.73 0:00:03 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.214 31 676.39 0:00:02 0:00:00
draws/s
━━━━━━━ 2000 0 0.157 15 692.61 0:00:02 0:00:00
draws/s
━━━━━━╸ 1960 0 0.182 15 625.30 0:00:03 0:00:01
draws/s
━━━━━━╺ 1822 0 0.185 31 581.68 0:00:03 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.214 31 676.39 0:00:02 0:00:00
draws/s
━━━━━━━ 2000 0 0.157 15 692.61 0:00:02 0:00:00
draws/s
━━━━━━━ 2000 0 0.182 31 628.24 0:00:03 0:00:00
draws/s
━━━━━━╺ 1855 0 0.185 31 583.29 0:00:03 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.214 31 676.39 0:00:02 0:00:00
draws/s
━━━━━━━ 2000 0 0.157 15 692.61 0:00:02 0:00:00
draws/s
━━━━━━━ 2000 0 0.182 31 628.24 0:00:03 0:00:00
draws/s
━━━━━━╸ 1906 0 0.185 31 588.15 0:00:03 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.214 31 676.39 0:00:02 0:00:00
draws/s
━━━━━━━ 2000 0 0.157 15 692.61 0:00:02 0:00:00
draws/s
━━━━━━━ 2000 0 0.182 31 628.24 0:00:03 0:00:00
draws/s
━━━━━━╸ 1989 0 0.185 31 594.10 0:00:03 0:00:01
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.214 31 676.39 0:00:02 0:00:00
draws/s
━━━━━━━ 2000 0 0.157 15 692.61 0:00:02 0:00:00
draws/s
━━━━━━━ 2000 0 0.182 31 628.24 0:00:03 0:00:00
draws/s
━━━━━━━ 2000 0 0.185 15 594.52 0:00:03 0:00:00
draws/s
Step Grad
Progre… Draw Diverg… size evals Speed Elapsed Remaini…
──────────────────────────────────────────────────────────────────────────────
━━━━━━━ 2000 0 0.214 31 676.39 0:00:02 0:00:00
draws/s
━━━━━━━ 2000 0 0.157 15 692.61 0:00:02 0:00:00
draws/s
━━━━━━━ 2000 0 0.182 31 628.24 0:00:03 0:00:00
draws/s
━━━━━━━ 2000 0 0.185 15 594.52 0:00:03 0:00:00
draws/s
?25hSampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 3 seconds.
Sampling: [beta_other, sigma_smooth, spline_slope, y_hat, y_hat_sigma, z_spline]
Sampling: [y_hat]
Sampling: [y_hat]
Sampling: [y_hat]
Sampling: [y_hat]
def diagnostics(name, experiment, var_names):
idata = experiment.model.idata
summary = az.summary(idata, var_names=var_names)
return {
"model": name,
"max r_hat": summary["r_hat"].max(),
"min ess_bulk": summary["ess_bulk"].min(),
"divergences": int(idata.sample_stats["diverging"].sum()),
"pre-period R²": experiment.score["unit_0_r2"],
}
pd.DataFrame(
[
diagnostics("linear trend", result_linear, ["beta", "y_hat_sigma"]),
diagnostics("cr(df=6), ridge prior", result_ridge, ["beta", "y_hat_sigma"]),
diagnostics(
"cr(df=12), RW2 prior",
result,
["sigma_smooth", "z_spline", "spline_slope", "beta_other", "y_hat_sigma"],
),
]
).set_index("model").round(3)
| max r_hat | min ess_bulk | divergences | pre-period R² | |
|---|---|---|---|---|
| model | ||||
| linear trend | 1.01 | 708.0 | 0 | 0.709 |
| cr(df=6), ridge prior | 1.03 | 157.0 | 0 | 0.746 |
| cr(df=12), RW2 prior | 1.01 | 906.0 | 0 | 0.742 |
This looks great! \(\hat{R}\) is back near 1 with a healthy bulk ESS, while the pre-treatment fit maintains the improvement the spline provides. In this notebook, we’ve worked through a simple example of how CausalPy’s integration with the larger PyMC ecosystem allows for rapid development of counterfactual models.
fig, ax = result.plot()
It’s instructive to look at the trend component on its own since the post-intervention period of an ITS regression is our counterfactual model. If we multiply the cr() columns of the design matrices by the beta_spline posterior, we get back the smooth trend over both periods.
def trend_draws(experiment, design):
"""Posterior draws of the smooth trend component for one design matrix."""
X = design["X"]
cols = [c for c in X.coeffs.values if c.startswith("cr(")]
beta_spline = (
experiment.model.idata.posterior["beta_spline"]
.isel(treated_units=0)
.rename({"spline_coeffs": "coeffs"})
)
return (X.sel(coeffs=cols) * beta_spline).sum("coeffs")
trend_pre = trend_draws(result, result.pre_design)
trend_post = trend_draws(result, result.post_design)
fig, ax = plt.subplots(figsize=(10, 4))
for trend, colour, label in [
(trend_pre, "C0", "fitted trend"),
(trend_post, "C1", "extrapolated counterfactual trend"),
]:
stacked = trend.stack(sample=("chain", "draw"))
lower, upper = np.quantile(stacked.values, [0.03, 0.97], axis=-1)
ax.fill_between(trend.obs_ind.values, lower, upper, color=colour, alpha=0.3)
ax.plot(
trend.obs_ind.values, stacked.mean("sample").values, color=colour, label=label
)
ax.axvline(treatment_time, color="k", ls=":", label="treatment time")
ax.set(
title="P-spline trend component, standardized scale",
xlabel="date",
ylabel="trend contribution to standardize(deaths)",
)
ax.legend()
plt.show()
Note that the counterfactual trend is linear throughout the post-treatment period. While cr() is a natural spline (so it extrapolates linearly past the last knot), the RW2 prior’s linear direction (spline_slope) learns that slope. The spline columns earn their place not by bending the counterfactual (they can’t!) but by (1) letting the local slope be governed by the trend near the intervention rather than a global average, and (2) removing pre-treatment misspecification that would otherwise leak into the seasonal terms and the residual scale.
Does the knot count matter?#
This is really the payoff of the penalty. With the ridge prior, the fit depends heavily on df, because df is the only knob controlling how wiggly the curve can get. With the RW2 prior, \(\sigma_{\text{smooth}}\) takes over that job, so df just has to be big enough to resolve the features in the data. Past that point, adding knots buys us nothing but sampling time.
Let’s sweep from 6 knots up to 30. Thirty knots is roughly one per five months of pre-treatment data, which the ridge prior would probably overfit. Under the penalty, the fitted trend is hard to tell apart from the 6-knot one.
fig, ax = plt.subplots(figsize=(10, 4))
sensitivity = []
for degrees_of_freedom, colour in zip([6, 12, 30], ["C0", "C1", "C2"], strict=True):
fit = cp.InterruptedTimeSeries(
df,
treatment_time,
formula=(
f"standardize(deaths) ~ 0 + cr(t, df={degrees_of_freedom}) "
"+ C(month) + standardize(temp)"
),
model=pspline_model(progressbar=False),
)
trend = trend_draws(fit, fit.pre_design).stack(sample=("chain", "draw"))
ax.plot(
trend.obs_ind.values,
trend.mean("sample").values,
color=colour,
label=f"df={degrees_of_freedom}",
)
sensitivity.append(
{
"df": degrees_of_freedom,
"pre-period R²": fit.score["unit_0_r2"],
"sigma_smooth": float(fit.model.idata.posterior["sigma_smooth"].mean()),
}
)
ax.set(title="Posterior mean trend by knot count", xlabel="date", ylabel="trend")
ax.legend()
plt.show()
pd.DataFrame(sensitivity).set_index("df").round(3)
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [sigma_smooth, z_spline, spline_slope, beta_other, y_hat_sigma]
Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 3 seconds.
Sampling: [beta_other, sigma_smooth, spline_slope, y_hat, y_hat_sigma, z_spline]
Sampling: [y_hat]
Sampling: [y_hat]
Sampling: [y_hat]
Sampling: [y_hat]
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [sigma_smooth, z_spline, spline_slope, beta_other, y_hat_sigma]
Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 3 seconds.
Sampling: [beta_other, sigma_smooth, spline_slope, y_hat, y_hat_sigma, z_spline]
Sampling: [y_hat]
Sampling: [y_hat]
Sampling: [y_hat]
Sampling: [y_hat]
Initializing NUTS using jitter+adapt_diag...
Multiprocess sampling (4 chains in 4 jobs)
NUTS: [sigma_smooth, z_spline, spline_slope, beta_other, y_hat_sigma]
Sampling 4 chains for 1_000 tune and 1_000 draw iterations (4_000 + 4_000 draws total) took 4 seconds.
There was 1 divergence after tuning. Increase `target_accept` or reparameterize.
Sampling: [beta_other, sigma_smooth, spline_slope, y_hat, y_hat_sigma, z_spline]
Sampling: [y_hat]
Sampling: [y_hat]
Sampling: [y_hat]
Sampling: [y_hat]
| pre-period R² | sigma_smooth | |
|---|---|---|
| df | ||
| 6 | 0.741 | 0.313 |
| 12 | 0.742 | 0.273 |
| 30 | 0.742 | 0.271 |
The three trends sit almost on top of each other, and sigma_smooth shrinks as we add knots so that the curve stays about as smooth either way. That’s the behavior we want … df is no longer a modeling choice we have to tune. We just set the resolution high enough and forget it.
print(result.effect_summary())
print(result_linear.effect_summary())
EffectSummary(table= mean median hdi_lower hdi_upper p_gt_0 relative_mean \
average 0.987617 0.964089 0.39757 1.652030 0.9995 372.882744
cumulative 28.640904 27.958592 11.52952 47.908873 0.9995 372.886016
relative_hdi_lower relative_hdi_upper
average -1861.079475 2537.880799
cumulative -1861.079257 2537.881246 , text='During the Post-period (2020-01-01 00:00:00 to 2022-05-01 00:00:00), the response variable had an average value of approx. 1.45. By contrast, in the absence of an intervention, we would have expected an average response of 0.46. The 95% interval of this counterfactual prediction is [-0.21, 1.05]. Subtracting this prediction from the observed response yields an estimate of the causal effect the intervention had on the response variable. This effect is 0.99 with a 95% interval of [0.40, 1.65].\n\nSumming up the individual data points during the Post-period, the response variable had an overall value of 41.93. By contrast, had the intervention not taken place, we would have expected a sum of 13.29. The 95% interval of this prediction is [-5.98, 30.40].\n\nThe 95% HDI of the effect [0.40, 1.65] does not include zero. The posterior probability of an increase is 1.000. Relative to the counterfactual, the effect represents a 372.88% change (95% HDI [-1861.08%, 2537.88%]).\n\nThis analysis assumes that the relationship between the time-based predictors and the response observed during the pre-intervention period remains stable throughout the post-intervention period. If the formula includes external covariates, it further assumes they were not themselves affected by the intervention. We recommend inspecting model fit, examining pre-intervention trends, and conducting sensitivity analyses (e.g., placebo tests) to support any causal conclusions drawn from this analysis.')
EffectSummary(table= mean median hdi_lower hdi_upper p_gt_0 relative_mean \
average 0.917606 0.918220 0.744012 1.112998 1.0 183.455552
cumulative 26.610585 26.628381 21.576358 32.276935 1.0 183.455555
relative_hdi_lower relative_hdi_upper
average 90.441709 294.182543
cumulative 90.441710 294.182550 , text='During the Post-period (2020-01-01 00:00:00 to 2022-05-01 00:00:00), the response variable had an average value of approx. 1.45. By contrast, in the absence of an intervention, we would have expected an average response of 0.53. The 95% interval of this counterfactual prediction is [0.33, 0.70]. Subtracting this prediction from the observed response yields an estimate of the causal effect the intervention had on the response variable. This effect is 0.92 with a 95% interval of [0.74, 1.11].\n\nSumming up the individual data points during the Post-period, the response variable had an overall value of 41.93. By contrast, had the intervention not taken place, we would have expected a sum of 15.32. The 95% interval of this prediction is [9.65, 20.36].\n\nThe 95% HDI of the effect [0.74, 1.11] does not include zero. The posterior probability of an increase is 1.000. Relative to the counterfactual, the effect represents a 183.46% change (95% HDI [90.44%, 294.18%]).\n\nThis analysis assumes that the relationship between the time-based predictors and the response observed during the pre-intervention period remains stable throughout the post-intervention period. If the formula includes external covariates, it further assumes they were not themselves affected by the intervention. We recommend inspecting model fit, examining pre-intervention trends, and conducting sensitivity analyses (e.g., placebo tests) to support any causal conclusions drawn from this analysis.')
Similarly, we can unpack each component of the model to see what its cumulative behavior contributes over the pre-treatment period.
X = result.pre_design["X"]
beta = (
result.model.idata.posterior["beta"].isel(treated_units=0).mean(("chain", "draw"))
)
dates = X.obs_ind.values
blocks = {
"spline": [c for c in X.coeffs.values if c.startswith("cr(")],
"month": [c for c in X.coeffs.values if c.startswith("C(month)")],
"temp": [c for c in X.coeffs.values if "temp" in c],
}
fig, axes = plt.subplots(4, 1, figsize=(10, 9), sharex=True)
# weighted basis functions from each cr() column * its coefficient.
for col in blocks["spline"]:
axes[0].plot(dates, X.sel(coeffs=col) * beta.sel(coeffs=col), color="C0", lw=1)
axes[0].set_title("weighted spline basis functions")
# add each block to the fitted mean.
running = 0
for ax, (name, cols) in zip(axes[1:], blocks.items(), strict=True):
running = running + (X.sel(coeffs=cols) * beta.sel(coeffs=cols)).sum("coeffs")
ax.plot(
dates,
result.pre_design["y"].isel(treated_units=0),
color="k",
lw=0.8,
alpha=0.4,
)
ax.plot(dates, running, color="C1")
ax.set_title(f"+ {name}")
axes[-1].set_xlabel("date")
fig.tight_layout()
plt.show()
Effect Summary Reporting#
For decision-making, you often need a concise summary of the causal effect with key statistics. The effect_summary() method provides a decision-ready report with average and cumulative effects, HDI intervals, tail probabilities, and relative effects. This provides a comprehensive summary without manual post-processing.
Note
Note that in this example, the data has been standardized, so the effect estimates are in standardized units. When interpreting the results, keep in mind that the effects are relative to the standardized scale of the outcome variable.
# Generate effect summary for the full post-period
stats = result.effect_summary()
stats.table
| mean | median | hdi_lower | hdi_upper | p_gt_0 | relative_mean | relative_hdi_lower | relative_hdi_upper | |
|---|---|---|---|---|---|---|---|---|
| average | 0.987617 | 0.964089 | 0.39757 | 1.652030 | 0.9995 | 372.882744 | -1861.079475 | 2537.880799 |
| cumulative | 28.640904 | 27.958592 | 11.52952 | 47.908873 | 0.9995 | 372.886016 | -1861.079257 | 2537.881246 |
# View the prose summary
print(stats.text)
During the Post-period (2020-01-01 00:00:00 to 2022-05-01 00:00:00), the response variable had an average value of approx. 1.45. By contrast, in the absence of an intervention, we would have expected an average response of 0.46. The 95% interval of this counterfactual prediction is [-0.21, 1.05]. Subtracting this prediction from the observed response yields an estimate of the causal effect the intervention had on the response variable. This effect is 0.99 with a 95% interval of [0.40, 1.65].
Summing up the individual data points during the Post-period, the response variable had an overall value of 41.93. By contrast, had the intervention not taken place, we would have expected a sum of 13.29. The 95% interval of this prediction is [-5.98, 30.40].
The 95% HDI of the effect [0.40, 1.65] does not include zero. The posterior probability of an increase is 1.000. Relative to the counterfactual, the effect represents a 372.88% change (95% HDI [-1861.08%, 2537.88%]).
This analysis assumes that the relationship between the time-based predictors and the response observed during the pre-intervention period remains stable throughout the post-intervention period. If the formula includes external covariates, it further assumes they were not themselves affected by the intervention. We recommend inspecting model fit, examining pre-intervention trends, and conducting sensitivity analyses (e.g., placebo tests) to support any causal conclusions drawn from this analysis.
# You can also analyze a specific time window, e.g., the first 6 months of 2020
stats_window = result.effect_summary(
window=(pd.to_datetime("2020-01-01"), pd.to_datetime("2020-06-30"))
)
stats_window.table
| mean | median | hdi_lower | hdi_upper | p_gt_0 | relative_mean | relative_hdi_lower | relative_hdi_upper | |
|---|---|---|---|---|---|---|---|---|
| average | 2.021437 | 2.010946 | 1.620051 | 2.426556 | 1.0 | 369.993909 | 104.723970 | 750.71222 |
| cumulative | 12.128620 | 12.065678 | 9.720309 | 14.559337 | 1.0 | 369.993942 | 104.723971 | 750.71224 |
We can get nicely formatted tables from our integration with the maketables package.
from maketables import ETable
ETable(result, coef_fmt="b:.3f")
| standardize(deaths) | |
|---|---|
| (1) | |
| coef | |
| month=1 | 1.358 |
| month=2 | -0.369 |
| month=3 | 0.135 |
| month=4 | -0.061 |
| month=5 | -0.081 |
| month=6 | -0.042 |
| month=7 | 0.214 |
| month=8 | -0.192 |
| month=9 | -0.276 |
| month=10 | -0.005 |
| month=11 | -0.430 |
| month=12 | -0.070 |
| cr(t, df=12)=0 | -0.047 |
| cr(t, df=12)=1 | -0.123 |
| cr(t, df=12)=2 | -0.226 |
| cr(t, df=12)=3 | -0.354 |
| cr(t, df=12)=4 | -0.378 |
| cr(t, df=12)=5 | -0.284 |
| cr(t, df=12)=6 | -0.144 |
| cr(t, df=12)=7 | 0.079 |
| cr(t, df=12)=8 | 0.260 |
| cr(t, df=12)=9 | 0.399 |
| cr(t, df=12)=10 | 0.414 |
| cr(t, df=12)=11 | 0.403 |
| standardize(temp) | -0.601 |
| stats | |
| N | 197 |
| Bayesian R2 | 0.742 |
| Format of coefficient cell: Coefficient | |