""" Plotting helpers for the geo-test calibration post. Same conventions as the previous two posts, so the figures of the series read as one: - channels keep fixed colours: PLA blue, Meta green, TV yellow (the Okabe-Ito palette, readable for colour-blind readers); - three fixed roles: truth green, estimate yellow, fit blue; - in this post's experiment figures, the control region borrows the fit blue and the treatment region the estimate yellow; - series are labelled directly instead of legend boxes; where direct labels would clutter, the encoding is stated in the title; - axes carry only the ink they need: no top/right spines, at most a light horizontal grid, no dual axes. The helpers own no model quantities: parameters, traces, and the adstock and saturation functions are passed in from the notebook, which remains the single source of truth. Sections are labelled and ordered by the step of the post in which each figure appears. Steps 3-6 (the differencing equation, the model, data preparation, estimation) produce no figures of their own, which is why the numbering runs 1, 2, 7, 8. """ import numpy as np import pandas as pd import matplotlib.pyplot as plt import arviz as az # Okabe-Ito palette, plus the roles of this series COLORS = { 'pla': '#0072B2', # blue 'meta': '#009E73', # green 'tv': '#D55E00', # yellow 'sales': '#000000', 'truth': '#009E73', 'estimate': '#D55E00', 'fit': '#0072B2', 'control': '#0072B2', 'treatment': '#D55E00', 'muted': '#808080', } # --------------------------------------------------------------------------- # Shared styling utilities (identical to the earlier helpers) # --------------------------------------------------------------------------- def _tufte(ax, grid='y'): """Strip an axis down to the ink it needs.""" for side in ('top', 'right'): ax.spines[side].set_visible(False) if grid == 'y': ax.grid(axis='y', color='0.92', linewidth=0.6) ax.set_axisbelow(True) else: ax.grid(False) ax.tick_params(length=3, labelsize=9) def _spread_labels(ax, x, entries, dx=2.0, fontsize=9): """Direct end labels with a minimum vertical gap, computed in display space so labels never collide regardless of panel height. entries: list of (y, text, color), any order.""" ax.figure.canvas.draw() height_px = fontsize * 1.45 * ax.figure.dpi / 72.0 inv = ax.transData.inverted() min_gap = abs(inv.transform((0, height_px))[1] - inv.transform((0, 0))[1]) entries = sorted(entries, key=lambda e: e[0]) ys = [e[0] for e in entries] for i in range(1, len(ys)): if ys[i] - ys[i - 1] < min_gap: ys[i] = ys[i - 1] + min_gap shift = (ys[-1] - entries[-1][0]) / 2 if len(ys) > 1 else 0.0 ys = [y - shift for y in ys] for (y_orig, text, color), y in zip(entries, ys): ax.annotate(text, xy=(x, y_orig), xytext=(x + dx, y), va='center', ha='left', color=color, fontsize=fontsize, annotation_clip=False) def _title(ax, text): """Left-aligned, regular-weight sentence titles.""" ax.set_title(text, loc='left', fontsize=10.5, fontweight='normal', pad=8) def _hdi90(draws): """90% HDI, tolerant of the arviz API change (hdi_prob -> prob).""" try: lo, hi = az.hdi(draws, hdi_prob=0.90) except TypeError: lo, hi = np.asarray(az.hdi(draws, prob=0.90)).ravel()[:2] return float(lo), float(hi) # --------------------------------------------------------------------------- # Step 1: the test calendar over the full history # --------------------------------------------------------------------------- def visualize_test_calendar(control_universe, treatment_universe, test_calendar, test_channel='pla'): """The tested channel's spend over the full history in both universes, with the go-dark windows shaded. One picture of the design: staggered windows on a drifting spending path probe the saturation curve at different levels.""" weeks = control_universe['week'].values spend_col = f'{test_channel}_spend' fig, ax = plt.subplots(figsize=(9.5, 3.4)) ax.plot(weeks, control_universe[spend_col], color=COLORS['control'], lw=1.1) ax.plot(weeks, treatment_universe[spend_col], color=COLORS['treatment'], lw=1.1, alpha=0.9) _spread_labels(ax, weeks[-1], [(control_universe[spend_col].iloc[-1], 'control', COLORS['control']), (treatment_universe[spend_col].iloc[-1], 'treatment (go-dark)', COLORS['treatment'])], dx=(weeks[-1] - weeks[0]) * 0.01, fontsize=8.5) for test in test_calendar.itertuples(): ax.axvspan(test.test_start_week - 0.5, test.test_end_week + 0.5, alpha=0.10, color='black', zorder=0) ax.annotate(f'test {test.test_id}', xy=((test.test_start_week + test.test_end_week) / 2, ax.get_ylim()[1]), ha='center', va='top', fontsize=8.5, color='0.45') ax.set_xlabel('Week', fontsize=9.5) ax.set_ylabel(f'{test_channel.upper()} spend (EUR K)', fontsize=9.5) ax.set_xlim(weeks[0], weeks[-1] * 1.14) _tufte(ax) _title(ax, 'The test calendar: four go-dark windows, staggered across ' 'a drifting spending path') plt.show() # --------------------------------------------------------------------------- # Step 2: the anatomy of one experiment # --------------------------------------------------------------------------- def visualize_geo_test(geo_test): """Three panels sharing a time axis: sales per region, the tested channel's spend per region, and the weekly sales gap between them. The gap panel is the object the model will be fit to.""" control = geo_test[geo_test['region'] == 'control'].sort_values('week') treatment = geo_test[geo_test['region'] == 'treatment'].sort_values('week') test_weeks = geo_test[geo_test['period'] == 'test']['week'] weeks = control['week'].values label_dx = (weeks[-1] - weeks[0]) * 0.02 label_room = (weeks[-1] - weeks[0]) * 0.18 fig, axes = plt.subplots(3, 1, figsize=(9.5, 7.5), sharex=True, gridspec_kw={'hspace': 0.4}) # Panel 1: sales in both regions ax = axes[0] ax.plot(weeks, control['sales'], 'o-', color=COLORS['control'], ms=3.5, lw=1.2) ax.plot(weeks, treatment['sales'], 's-', color=COLORS['treatment'], ms=3.5, lw=1.2) _spread_labels(ax, weeks[-1], [(control['sales'].iloc[-1], 'control', COLORS['control']), (treatment['sales'].iloc[-1], 'treatment (go-dark)', COLORS['treatment'])], dx=label_dx) ax.set_ylabel('Sales (EUR K)', fontsize=9.5) _title(ax, 'During the test, sales in the go-dark group fall behind...') # Panel 2: the tested channel's spend in both regions ax = axes[1] ax.plot(weeks, control['pla_spend'], 'o-', color=COLORS['control'], ms=3.5, lw=1.2) ax.plot(weeks, treatment['pla_spend'], 's-', color=COLORS['treatment'], ms=3.5, lw=1.2) _spread_labels(ax, weeks[-1], [(control['pla_spend'].iloc[-1], 'control', COLORS['control']), (treatment['pla_spend'].iloc[-1], 'treatment', COLORS['treatment'])], dx=label_dx) ax.set_ylabel('PLA spend (EUR K)', fontsize=9.5) _title(ax, '...because its PLA spend is switched off...') # Panel 3: the weekly sales gap, observed and true ax = axes[2] gap = control['sales'].values - treatment['sales'].values true_gap = (control['pla_effect'].values - treatment['pla_effect'].values) ax.plot(weeks, true_gap, color=COLORS['truth'], lw=1.6) ax.plot(weeks, gap, 'd-', color=COLORS['sales'], ms=3.5, lw=1.0, alpha=0.85) ax.axhline(0, color='0.7', lw=0.8) _spread_labels(ax, weeks[-1], [(gap[-1], 'observed gap', COLORS['sales']), (true_gap[-1], 'true effect gap', COLORS['truth'])], dx=label_dx, fontsize=8.5) ax.set_ylabel('Sales gap (EUR K)', fontsize=9.5) ax.set_xlabel('Week', fontsize=9.5) _title(ax, '...leaving a weekly sales gap between the groups — ' 'the object we will model') for ax in axes: ax.axvspan(test_weeks.min() - 0.5, test_weeks.max() + 0.5, alpha=0.07, color='black', zorder=0) _tufte(ax) ax.set_xlim(weeks[0], weeks[-1] + label_room) axes[0].annotate('test window', xy=(test_weeks.min() - 0.2, axes[0].get_ylim()[1]), fontsize=8.5, color='0.45', va='top', ha='left') plt.show() # --------------------------------------------------------------------------- # Step 7: the response curve the experiments identify # --------------------------------------------------------------------------- def visualize_identified_response_curve(trace, spend, scale_factors, true_config, adstock_fn, saturation_fn, experiment_weeks, n_lags=6, n_draws=200, seed=1): """Steady-state response curve for the tested channel: posterior draws (spaghetti), the posterior mean, and the true curve. The rug marks the adstocked spends during the experiment weeks — the levels at which the tests actually probed the curve. `adstock_fn` and `saturation_fn` are the notebook's own functions, passed in so the definitions live in exactly one place.""" import matplotlib.transforms as mtransforms y_max = scale_factors['total_sales'] x_max = scale_factors['pla_spend'] spend = np.asarray(spend) alpha_draws = trace.posterior['alpha'].values.flatten() lambda_draws = trace.posterior['lam'].values.flatten() beta_draws = trace.posterior['beta'].values.flatten() idx = np.random.default_rng(seed).choice( len(alpha_draws), min(n_draws, len(alpha_draws)), replace=False) # adstocked spend at the posterior-mean carryover rate, in EUR K adstocked = adstock_fn(spend, alpha_draws.mean(), n_lags) probed = adstocked[np.isin(np.arange(1, len(spend) + 1), experiment_weeks)] x_grid = np.linspace(0, adstocked.max() * 1.1, 200) # EUR K x_scaled = x_grid / x_max fig, ax = plt.subplots(figsize=(7.0, 4.2)) draws = np.array([beta_draws[i] * saturation_fn(x_scaled, lambda_draws[i]) * y_max for i in idx]) for response in draws: ax.plot(x_grid, response, color=COLORS['fit'], alpha=0.03, lw=0.8) response_mean = draws.mean(axis=0) ax.plot(x_grid, response_mean, color=COLORS['fit'], lw=1.8) true_response = (true_config['coefficient'] * saturation_fn(x_grid, true_config['saturation_lambda'])) ax.plot(x_grid, true_response, color=COLORS['truth'], lw=1.6, linestyle='--') rug_tf = mtransforms.blended_transform_factory(ax.transData, ax.transAxes) ax.plot(probed, np.zeros_like(probed), marker='|', linestyle='none', color='0.55', markersize=7, markeredgewidth=0.9, alpha=0.6, transform=rug_tf, clip_on=False) ax.annotate('adstocked spend during the experiments', xy=(probed.mean(), 0), xycoords=rug_tf, xytext=(4, 4), textcoords='offset points', fontsize=7.5, color='0.5', va='bottom') _spread_labels(ax, x_grid[-1], [(response_mean[-1], 'estimated', COLORS['fit']), (true_response[-1], 'true', COLORS['truth'])], dx=x_grid[-1] * 0.02, fontsize=8.5) ax.set_xlim(0, x_grid[-1] * 1.22) ax.set_xlabel('Adstocked weekly spend (EUR K)', fontsize=9.5) ax.set_ylabel('Weekly incremental sales (EUR K)', fontsize=9.5) _tufte(ax) _title(ax, 'The experiments recover the response curve where they probed it') plt.show() # --------------------------------------------------------------------------- # Step 7: parameter posteriors against the truth # --------------------------------------------------------------------------- def visualize_parameter_posteriors(traces_by_label, parameters): """One panel per parameter: the posterior under each set of experiments, against the true value (dashed line). `traces_by_label` is an ordered dict mapping labels ('2 tests', '4 tests') to InferenceData objects; `parameters` is an ordered dict mapping display titles to {'var': posterior variable name, 'truth': scaled true value}. Earlier runs take the fit blue, the final run the estimate yellow — the two ends of the Okabe-Ito pair, chosen to stay apart when the histograms overlap.""" labels = list(traces_by_label.keys()) shades = [COLORS['fit']] * (len(labels) - 1) + [COLORS['estimate']] fig, axes = plt.subplots(1, len(parameters), figsize=(9.5, 3.3), gridspec_kw={'wspace': 0.25}) for (title, spec), ax in zip(parameters.items(), np.atleast_1d(axes)): peaks = [] for label, shade in zip(labels, shades): draws = traces_by_label[label].posterior[spec['var']].values.flatten() counts, edges, _ = ax.hist(draws, bins=35, density=True, alpha=0.6, color=shade, linewidth=0) peak = np.argmax(counts) peaks.append(((edges[peak] + edges[peak + 1]) / 2, counts[peak], label, shade)) # label the runs at their modes; stack vertically if the modes collide x_span = np.ptp(ax.get_xlim()) collide = len(peaks) > 1 and abs(peaks[0][0] - peaks[1][0]) < 0.15 * x_span for i, (x, y, label, color) in enumerate(peaks): ax.annotate(label, xy=(x, y), xytext=(0, 4 + (14 * i if collide else 0)), textcoords='offset points', ha='center', va='bottom', fontsize=8.5, color=color) ax.axvline(spec['truth'], color=COLORS['truth'], linestyle='--', lw=1.4) ax.margins(y=0.20) _tufte(ax) _title(ax, title) ax.set_yticks([]) np.atleast_1d(axes)[0].set_ylabel('Density', fontsize=9.5) fig.suptitle('Posteriors under two and four tests; ' 'dashed lines mark the truths', x=0.01, ha='left', fontsize=10.5, fontweight='normal') plt.subplots_adjust(top=0.82) plt.show() # --------------------------------------------------------------------------- # Step 7: pairwise joint posteriors — the trade-offs, drawn # --------------------------------------------------------------------------- def visualize_joint_posteriors(trace, parameters, n_draws=2000, seed=2): """One panel per parameter pair: the joint posterior as a point cloud of draws, the truth as a dashed crosshair. `parameters` follows the same convention as visualize_parameter_posteriors: an ordered dict mapping display titles to {'var': posterior variable name, 'truth': scaled true value}. The point of the figure is the lambda-beta panel, where weak identification appears as a ridge rather than a blob.""" from itertools import combinations items = list(parameters.items()) draws = {title: trace.posterior[spec['var']].values.flatten() for title, spec in items} n_total = len(next(iter(draws.values()))) idx = np.random.default_rng(seed).choice( n_total, min(n_draws, n_total), replace=False) pairs = list(combinations(items, 2)) fig, axes = plt.subplots(1, len(pairs), figsize=(9.5, 3.3), gridspec_kw={'wspace': 0.38}) for ((title_x, spec_x), (title_y, spec_y)), ax in zip(pairs, np.atleast_1d(axes)): ax.scatter(draws[title_x][idx], draws[title_y][idx], s=4, alpha=0.12, color=COLORS['estimate'], linewidths=0, rasterized=True) ax.axvline(spec_x['truth'], color=COLORS['truth'], linestyle='--', lw=1.1, alpha=0.9) ax.axhline(spec_y['truth'], color=COLORS['truth'], linestyle='--', lw=1.1, alpha=0.9) ax.set_xlabel(title_x, fontsize=9.5) ax.set_ylabel(title_y, fontsize=9.5) _tufte(ax, grid='n') fig.suptitle('Pairwise joint posteriors under four tests; ' 'dashed crosshairs mark the truths', x=0.01, ha='left', fontsize=10.5, fontweight='normal') plt.subplots_adjust(top=0.84) plt.show() # --------------------------------------------------------------------------- # Step 8 (the verdict): every method of the series, side by side # --------------------------------------------------------------------------- def visualize_method_comparison(results_by_method, true_roas): """Estimated ROAS (dot, 90% interval) under each method of the series, against the true value (dashed line). Each estimate is labelled directly with its mean and interval, so the figure carries the exact numbers and no separate table is needed. `results_by_method` is an ordered dict mapping labels either to an array of posterior draws or to a (mean, lower, upper) tuple for results quoted from an earlier post.""" labels = list(results_by_method.keys()) x_pos = np.arange(len(labels)) fig, ax = plt.subplots(figsize=(7.5, 4.2)) for x, label in zip(x_pos, labels): result = results_by_method[label] if isinstance(result, (tuple, list)): # quoted from an earlier post mean, lo, hi = result color = '0.55' else: # posterior draws from this one draws = np.asarray(result).flatten() mean = draws.mean() lo, hi = _hdi90(draws) color = COLORS['estimate'] ax.errorbar(x, mean, yerr=[[mean - lo], [hi - mean]], fmt='o', capsize=3, markersize=6, lw=1.2, color=color) # direct labels above the interval: the mean, then the 90% bounds ax.annotate(f"{mean:.1f}x", xy=(x, hi), xytext=(0, 14), textcoords='offset points', ha='center', va='bottom', fontsize=8.5, color=color) ax.annotate(f"[{lo:.1f}, {hi:.1f}]", xy=(x, hi), xytext=(0, 3), textcoords='offset points', ha='center', va='bottom', fontsize=7.5, color='0.55') ax.margins(y=0.18) ax.axhline(true_roas, color=COLORS['truth'], lw=1.2, linestyle='--') ax.annotate(f"true {true_roas:.1f}x", xy=(x_pos[0] - 0.45, true_roas), xytext=(0, 4), textcoords='offset points', ha='left', va='bottom', fontsize=8.5, color=COLORS['truth']) ax.axhline(0, color='0.85', lw=0.7) ax.set_xticks(x_pos) ax.set_xticklabels(labels, fontsize=8.5) ax.set_xlim(-0.5, len(labels) - 0.5) ax.set_ylabel('PLA ROAS (90% interval)', fontsize=9.5) _tufte(ax) _title(ax, 'Estimated PLA ROAS by method: observational (grey)\n' 'and experimental (yellow), against the truth') plt.show()