location_surface_grids = {}
with location_error_gp_model:
    for multiplier, X_noisy_value in zip(multipliers, noisy_xs):
        pm.set_data({"X_noisy": X_noisy_value, "σ_s": multiplier})
        posterior_mean_point = {
            name: location_error_idatas[multiplier].posterior[name].mean(("chain", "draw")).values
            for name in ["μ", "σ", "ℓ", "σ0", "Δs", "X_true"]
        }
        f_mean, _ = gp_location.predict(Xnew, point=posterior_mean_point, diag=True, pred_noise=False)
        location_surface_grids[multiplier] = f_mean.reshape(n_prediction_grid, n_prediction_grid)

naive_kde_grids = {}
kde_bandwidths = {
    multiplier: location_error_idatas[multiplier].posterior["ℓ"].mean(("chain", "draw")).item()
    for multiplier in multipliers
}
for multiplier, X_noisy_value in zip(multipliers, noisy_xs):
    squared_distance = (
        (Xnew[:, 0, None] - X_noisy_value[:, 0]) ** 2
        + (Xnew[:, 1, None] - X_noisy_value[:, 1]) ** 2
    )
    weights = np.exp(-0.5 * squared_distance / kde_bandwidths[multiplier]**2)
    naive_kde_grids[multiplier] = (weights @ y_walker / weights.sum(axis=1)).reshape(n_prediction_grid, n_prediction_grid)

location_norm = colors.Normalize(
    vmin=min(y_walker.min(), *(grid.min() for grid in location_surface_grids.values())),
    vmax=max(y_walker.max(), *(grid.max() for grid in location_surface_grids.values())),
)
point_colors = [theme["gunmetal"], theme["sepia"], theme["rust"], theme["steel"]]
arrow_head_width = 0.018 * max(x_range, y_range)

fig = plt.figure(figsize=(5, 11.2 * (5 / 7)))
gs = fig.add_gridspec(4, len(multipliers), hspace=0.08, wspace=0.08)
axes = np.array([[fig.add_subplot(gs[row, col]) for col in range(len(multipliers))] for row in range(4)])
panel_idx = 0
surface_meshes = []

for col, (multiplier, X_noisy_value) in enumerate(zip(multipliers, noisy_xs)):
    ax = axes[0, col]
    ax.scatter(X_walker[:, 0], X_walker[:, 1], facecolor=theme["paper"], edgecolor=theme["gunmetal"], s=14, linewidth=0.35, alpha=0.5)
    ax.scatter(X_noisy_value[:, 0], X_noisy_value[:, 1], c=y_walker, cmap=sepia_cmap, norm=location_norm, s=16, edgecolor=theme["paper"], linewidth=0.25, alpha=0.5)
    for point_idx, point_color in zip(selected_location_idx, point_colors):
        circle = Circle(X_noisy_value[point_idx], radius=multiplier, facecolor=colors.to_rgba(theme["steel"], 0.2), edgecolor=theme["gunmetal"], linewidth=0.6)
        ax.add_patch(circle)
        dx, dy = X_noisy_value[point_idx] - X_walker[point_idx]
        ax.arrow(X_walker[point_idx, 0], X_walker[point_idx, 1], dx, dy, color=plot_text_color, linewidth=0.8, length_includes_head=True, head_width=arrow_head_width, head_length=arrow_head_width)
    ax.scatter(X_walker[selected_location_idx, 0], X_walker[selected_location_idx, 1], facecolor=theme["paper"], edgecolor=plot_text_color, s=44, linewidth=0.8)
    ax.scatter(X_noisy_value[selected_location_idx, 0], X_noisy_value[selected_location_idx, 1], c=y_walker[selected_location_idx], cmap=sepia_cmap, norm=location_norm, s=44, edgecolor=plot_text_color, linewidth=0.8)
    ax.xaxis.set_label_position("top"); ax.set_xlabel(rf"$\sigma_s = {multiplier:.0f}$ m")

    ax = axes[2, col]
    mesh = ax.pcolormesh(x_new_mesh, y_new_mesh, location_surface_grids[multiplier], cmap=sepia_cmap, norm=location_norm, shading="auto")
    surface_meshes.append(mesh)

    ax = axes[3, col]
    ax.pcolormesh(x_new_mesh, y_new_mesh, naive_kde_grids[multiplier], cmap=sepia_cmap, norm=location_norm, shading="auto")

    ax = axes[1, col]
    ax.pcolormesh(x_new_mesh, y_new_mesh, location_surface_grids[multiplier], cmap=sepia_cmap, norm=location_norm, shading="auto", alpha=0.25)
    X_true_samples = location_error_idatas[multiplier].posterior["X_true"].stack(sample=("chain", "draw")).transpose("obs", "coord", "sample")
    for point_idx, point_color in zip(selected_location_idx, point_colors):
        x_draws = X_true_samples.isel(obs=point_idx).sel(coord="x").values
        y_draws = X_true_samples.isel(obs=point_idx).sel(coord="y").values
        sns.kdeplot(x=x_draws, y=y_draws, levels=3, color=point_color, linewidths=0.9, fill=False, ax=ax)
        ax.scatter(X_walker[point_idx, 0], X_walker[point_idx, 1], marker="x", color=point_color, s=34, linewidth=1.1)
        ax.scatter(X_noisy_value[point_idx, 0], X_noisy_value[point_idx, 1], marker="o", facecolor=theme["paper"], edgecolor=point_color, s=34, linewidth=1.1)

for row, row_label in enumerate(["Perturbed coordinates", "Posterior location density", r"Posterior mean of $f(s)$", "Naive smoothed KDE"]):
    axes[row, 0].set_ylabel(row_label)

for ax in axes.ravel():
    ax.set_xlim(*x_limits); ax.set_ylim(*y_limits); ax.set_aspect("equal"); ax.grid(False)
    ax.tick_params(axis="both", which="both", bottom=False, left=False, top=False, right=False, labelbottom=False, labelleft=False, labeltop=False, labelright=False)
    panel_idx += 1

cbar = fig.colorbar(surface_meshes[-1], ax=axes.ravel().tolist(), orientation="horizontal", fraction=0.035, pad=0.04)
cbar.set_label(r"Uranium, $\log_{10}(x + 1)$")

figure_dir = Path.home() / "ckrapu.github.io/images/2026-05-24-dont-know-where-your-data-is-from"
figure_dir.mkdir(parents=True, exist_ok=True)
fig.savefig(figure_dir / "error-in-location-grid.png", dpi=220, bbox_inches="tight", facecolor=fig.get_facecolor())