GPU ensembles: ChEES-HMC and MEADS on a forest of 20,000 cells

import os

os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")  # before NumPy: one BLAS thread per nutpie chain

import contextlib
import importlib.util
import logging
import tempfile
import time
import warnings

import arviz as az
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import nutpie
import pandas as pd
import pymc as pm
import pytensor
import pytensor.tensor as pt
import xarray as xr
from jax.flatten_util import ravel_pytree
from pymc.sampling.jax import get_jaxified_graph, get_jaxified_logp

from pymc_challenges import data

RANDOM_SEED = 86
rng = np.random.default_rng(RANDOM_SEED)
az.style.use("arviz-variat")
warnings.filterwarnings("ignore", message=".*Explicitly requested dtype.*")
warnings.filterwarnings("ignore", message=".*truncated to dtype float32.*")
logging.getLogger("pymc").setLevel(logging.WARNING)

# Asking JAX for its backend would start CUDA, and the CPU baseline comes first. So does importing blackjax, which
# touches a JAX array when it loads: its import waits until after the baseline, against the usual first-cell habit.
HAS_GPU = (importlib.util.find_spec("jax_cuda13_plugin") is not None
           and any(os.path.exists(d) for d in ("/dev/nvidia0", "/dev/nvidiactl", "/dev/dxg"))
           and os.environ.get("JAX_PLATFORMS") != "cpu")
CORES = len(os.sched_getaffinity(0)) if hasattr(os, "sched_getaffinity") else os.cpu_count()
NUTS_CHAINS = min(8, max(4, CORES // 2 - 1))  # one on each physical core, capped to keep the build under 5 GB
NUTS_DRAWS = 4000 // NUTS_CHAINS               # 4,000 draws in total
L_SMOOTH = 800.0  # half-width of the box for the smooth fields: far enough from the plot for their lengthscales
# The full problem on a GPU; a coarser lattice on a CPU, which keeps the same code within budget.
if HAS_GPU:
    CELL, M_FINE, M_SMOOTH = 5.0, (100, 200), (40, 40)
    CHAINS, WARMUP, SNAPS, KEEP = 128, 1000, 30, 10
    MEADS_WARMUP, MEADS_KEEP, MEADS_THIN = 20_000, 13, 100
else:
    CELL, M_FINE, M_SMOOTH = 20.0, (25, 50), (20, 20)
    CHAINS, WARMUP, SNAPS, KEEP = 64, 500, 20, 5  # MEADS needs at least 16 chains per fold here: see section 6
    MEADS_WARMUP, MEADS_KEEP, MEADS_THIN = 10_000, 10, 100
print(f"GPU build of JAX: {HAS_GPU}; lattice {CELL:.0f} m, {CHAINS} chains")

@contextlib.contextmanager
def precision(floatX):
    """Build and run a model in one precision: PyTensor's floatX and JAX's x64 must agree."""
    old = pytensor.config.floatX
    pytensor.config.floatX = floatX
    jax.config.update("jax_enable_x64", floatX == "float64")
    try:
        yield
    finally:
        pytensor.config.floatX = old
        jax.config.update("jax_enable_x64", old == "float64")


pytensor.config.floatX = "float32"
jax.config.update("jax_enable_x64", False)  # importing PyMC switches x64 on

data.describe("bei_trees")
raw = np.load(data.path("bei_trees"))
n_trees = raw["x"].size
print(f"{n_trees:,} trees; elevation grid {raw['elev'].shape}")


def lattice(cell):
    """Counts and standardised terrain on a lattice of `cell` metres."""
    nx, ny = int(1000 / cell), int(500 / cell)
    y = np.histogram2d(raw["y"], raw["x"], bins=[ny, nx], range=[[0, 500], [0, 1000]])[0]
    k = int(cell / 5)

    def to_cells(v):
        c = 0.25 * (v[:-1, :-1] + v[1:, :-1] + v[:-1, 1:] + v[1:, 1:])  # corners to 5 m centres
        return c.reshape(ny, k, nx, k).mean((1, 3))

    standard = lambda a: (a - a.mean()) / a.std()
    return y, standard(to_cells(raw["elev"])), standard(to_cells(np.log(raw["grad"]))), ny, nx


y5, elev5, grad5, _, _ = lattice(5.0)
fig, axes = plt.subplots(2, 1, figsize=(9, 8.2), sharex=True, layout="constrained")
extent = (0, 1000, 0, 500)
grid_y, grid_x = (np.linspace(0, size, n) for size, n in zip((500, 1000), raw["elev"].shape))
levels = np.arange(np.floor(raw["elev"].min() / 2) * 2, raw["elev"].max() + 2, 2)  # every 2 m

im = axes[0].imshow(raw["elev"], origin="lower", extent=extent, cmap="YlOrBr", alpha=0.75)
lines = axes[0].contour(grid_x, grid_y, raw["elev"], levels=levels, colors="#5b3a1a")
lines.set_linewidths([1.3 if level % 10 == 0 else 0.5 for level in lines.levels])  # index contours
axes[0].clabel(lines, levels=[level for level in levels if level % 10 == 0], fmt="%d m", fontsize=8)
axes[0].scatter(raw["x"], raw["y"], s=4, c="#0b3d2e", linewidths=0)
axes[0].set(title=f"{n_trees:,} Beilschmiedia pendula trees on the terrain, contours every 2 m", ylabel="y (m)",
            aspect="equal", xlim=(0, 1000), ylim=(0, 500))
fig.colorbar(im, ax=axes[0], shrink=0.8, label="elevation (m)")

im = axes[1].imshow(np.log(raw["grad"]), origin="lower", extent=extent, cmap="viridis")
axes[1].contour(grid_x, grid_y, raw["elev"], levels=levels, colors="w", linewidths=0.4, alpha=0.6)
axes[1].set(title="log terrain gradient, with the same contours", xlabel="x (m)", ylabel="y (m)")
fig.colorbar(im, ax=axes[1], shrink=0.8)
plt.show()


def trees_per_ha(terrain):
    """Mean tree density in the four quartiles of a terrain variable, on the 5 m cells."""
    quartile = pd.qcut(terrain.ravel(), 4, labels=["lowest quarter", "second", "third", "highest quarter"])
    return (pd.Series(y5.ravel()).groupby(quartile, observed=True).mean() * 400).round(1)  # 400 cells a hectare


print("Trees per hectare, by quarter of the plot:")
print(pd.DataFrame({"elevation": trees_per_ha(elev5), "log gradient": trees_per_ha(grad5)}).to_string());

def basis(n, cell, half_width, m):
    """Laplacian eigenfunctions on [-L, L] at `n` cell centres, and their frequencies."""
    centres = (np.arange(n) + 0.5) * cell
    x = centres - centres.mean()
    w = np.pi * np.arange(1, m + 1) / (2 * half_width)
    return np.sqrt(1 / half_width) * np.sin(w[None, :] * (x[:, None] + half_width)), w


def matern_sqrt_density(w2, ell, nu):
    """Square root of the 2-D Matern spectral density, nu = 1/2 or 3/2."""
    constant = {0.5: 2 * np.pi / ell, 1.5: 6 * np.pi * 3**1.5 / ell**3}[nu]
    return pt.sqrt(constant) * ((2 * nu) / ell**2 + w2) ** (-(nu + 1) / 2)


def make_model(cell=CELL, m_fine=M_FINE, m_smooth=M_SMOOTH):
    y, elev, grad, ny, nx = lattice(cell)
    fx = pytensor.config.floatX
    Ly, Lx = 1.25 * 250.0, 1.25 * 500.0
    Py, wy = basis(ny, cell, Ly, m_fine[0])
    Px, wx = basis(nx, cell, Lx, m_fine[1])
    Qy, vy = basis(ny, cell, L_SMOOTH, m_smooth[0])
    Qx, vx = basis(nx, cell, L_SMOOTH, m_smooth[1])
    w2 = (wy[:, None] ** 2 + wx[None, :] ** 2).astype(fx)
    v2 = (vy[:, None] ** 2 + vx[None, :] ** 2).astype(fx)
    Py, Px, Qy, Qx = Py.astype(fx), Px.astype(fx), Qy.astype(fx), Qx.astype(fx)
    elev, grad = elev.astype(fx), grad.astype(fx)
    log_area = np.log(cell**2 / 1e4)  # intensity is per hectare

    with pm.Model() as model:
        alpha = pm.Normal("alpha", np.log(30.0), 1.0)
        b_elev = pm.Normal("b_elev", 0.0, 1.0)
        b_grad = pm.Normal("b_grad", 0.0, 1.0)
        s_broad = pm.LogNormal("s_broad", np.log(0.7), 0.5)
        ell_broad = pm.LogNormal("ell_broad", np.log(120.0), 0.4)
        s_fine = pm.LogNormal("s_fine", np.log(0.7), 0.5)
        ell_fine = pm.LogNormal("ell_fine", np.log(20.0), 0.6)
        s_slope_elev = pm.LogNormal("s_slope_elev", np.log(0.5), 0.5)
        s_slope_grad = pm.LogNormal("s_slope_grad", np.log(0.5), 0.5)
        ell_slope = pm.LogNormal("ell_slope", np.log(200.0), 0.4)
        z_broad = pm.Normal("z_broad", 0.0, 1.0, shape=m_smooth)
        z_fine = pm.Normal("z_fine", 0.0, 1.0, shape=m_fine)
        z_elev = pm.Normal("z_elev", 0.0, 1.0, shape=m_smooth)
        z_grad = pm.Normal("z_grad", 0.0, 1.0, shape=m_smooth)

        smooth = lambda z, sd: Qy @ (sd * z) @ Qx.T
        broad = smooth(z_broad, s_broad * matern_sqrt_density(v2, ell_broad, 1.5))
        fine = Py @ (s_fine * matern_sqrt_density(w2, ell_fine, 0.5) * z_fine) @ Px.T
        sd_slope = matern_sqrt_density(v2, ell_slope, 1.5)
        slope_elev = b_elev + smooth(z_elev, s_slope_elev * sd_slope)
        slope_grad = b_grad + smooth(z_grad, s_slope_grad * sd_slope)
        habitat = alpha + slope_elev * elev + slope_grad * grad + broad
        eta = habitat + fine + log_area
        pm.Poisson("y", pt.exp(eta), observed=y)
    outputs = {"eta": eta, "habitat": habitat, "slope_elev": slope_elev, "slope_grad": slope_grad}
    return model, y, outputs


model, y, outputs = make_model()
print(f"{y.size:,} cells, {sum(v.size for v in model.initial_point().values()):,} parameters")

import blackjax  # CUDA starts here, not before
import optax
from blackjax.adaptation.base import get_filter_adapt_info_fn

key = jax.random.key(RANDOM_SEED)
if (jax.default_backend() == "gpu") != HAS_GPU:
    raise RuntimeError(f"expected GPU={HAS_GPU} but JAX chose {jax.default_backend()}: if your GPU is not usable, "
                       "run with JAX_PLATFORMS=cpu")
print("JAX backend:", jax.default_backend(), jax.devices())


def jaxify(m):
    """One flat parameter vector -> log density, plus the unravel function."""
    logp = get_jaxified_logp(m)  # NB: the default negative_logp=True returns +logp
    start = m.initial_point()
    names = [v.name for v in m.value_vars]
    q0, unravel = ravel_pytree([jnp.asarray(start[n]) for n in names])
    return (lambda q: logp(unravel(q))), q0, unravel, names


logdensity, q0, unravel, names = jaxify(model)
print("value variables:", names[:10], "...")
small, _, _ = make_model(25.0, (20, 40), (16, 16))  # the check on a small lattice: compile_logp is heavy
f_small, q_small, _, _ = jaxify(small)
print(f"25 m lattice: PyMC logp {small.compile_logp()(small.initial_point()):.2f}, JAX logp {float(f_small(q_small)):.2f}")
del small, f_small, q_small
NH = len(HYPER)
value_is_log = np.array([n.endswith("_log__") for n in names[:NH]])


def timed(fn, *args, repeat=5):
    """(compile seconds, best-of-`repeat` run seconds). block_until_ready defeats async dispatch."""
    t0 = time.perf_counter()
    jax.block_until_ready(fn(*args))
    compile_s = time.perf_counter() - t0
    runs = []
    for _ in range(repeat):
        t0 = time.perf_counter()
        jax.block_until_ready(fn(*args))
        runs.append(time.perf_counter() - t0)
    return compile_s, min(runs)


def batch_gradient(f, q, chains, device):
    Q = jax.device_put(q + 0.1 * jnp.asarray(rng.normal(size=(chains, q.size)), q.dtype), device)
    result = timed(jax.jit(jax.vmap(jax.value_and_grad(f))), Q)
    jax.clear_caches()  # every chain count is a new compiled program, kept alive otherwise
    return result


cpu = jax.devices("cpu")[0]
gpu = jax.devices("gpu")[0] if HAS_GPU else None
chain_grid = sorted({1, 8, 32, 128, 512, CHAINS}) if HAS_GPU else sorted({1, 8, CHAINS})
rows = []
for chains in chain_grid:
    compile_cpu, cpu_s = batch_gradient(logdensity, q0, chains, cpu) if chains <= 128 else (np.nan, np.nan)
    row = {"chains": chains, "CPU float32": cpu_s * 1e3, "compile CPU s": compile_cpu}
    if HAS_GPU:
        compile_gpu, gpu_s = batch_gradient(logdensity, q0, chains, gpu)
        row |= {"GPU float32": gpu_s * 1e3, "compile GPU s": compile_gpu}
    rows.append(row)
timing = pd.DataFrame(rows).set_index("chains")

if HAS_GPU:  # the same model in double precision
    with precision("float64"):
        model64, _, _ = make_model()
        f64, q64, _, _ = jaxify(model64)
        timing["GPU float64"] = [batch_gradient(f64, q64, c, gpu)[1] * 1e3 for c in chain_grid]
    del model64, f64, q64
    timing["GPU speed-up"] = timing["CPU float32"] / timing["GPU float32"]
    timing["float32 speed-up on GPU"] = timing["GPU float64"] / timing["GPU float32"]
print(timing.round(2).to_string())

key, init_key = jax.random.split(key)
spread = jnp.concatenate([jnp.full(NH, 0.3), jnp.ones(q0.size - NH)]).astype(q0.dtype)
cold_start = q0 + spread * jax.random.normal(init_key, (CHAINS, q0.size), q0.dtype)


def summarise(draws):
    """Posterior summary of hyperparameters, draws of shape (chain, draw, NH) on the value scale."""
    d = np.asarray(draws, float).copy()
    d[..., value_is_log] = np.exp(d[..., value_is_log])
    ds = xr.Dataset({n: (("chain", "draw"), d[..., i]) for i, n in enumerate(HYPER)})
    rhat, ess = az.rhat(ds), az.ess(ds)
    return pd.DataFrame({"mean": d.mean((0, 1)), "sd": d.std((0, 1)),
                         "r_hat": [float(rhat[n]) for n in HYPER],
                         "ess": [float(ess[n]) for n in HYPER]}, index=HYPER)


def sample(step, state, k, snaps, keep, thin):
    """Run the vectorised kernel; every `thin` steps keep hyperparameters, every `keep * thin` all of q."""
    def one(s, ki):
        s, info = jax.vmap(step)(jax.random.split(ki, CHAINS), s)
        steps = jnp.mean(getattr(info, "num_integration_steps", jnp.ones(())))
        return s, (info.acceptance_rate, info.is_divergent, steps)

    def keep_block(s, ki):
        s, (acc, div, steps) = jax.lax.scan(one, s, jax.random.split(ki, thin))
        return s, (s.position[:, :NH], acc.mean(), div.sum(), steps.sum())

    def snap_block(s, ki):
        s, out = jax.lax.scan(keep_block, s, jax.random.split(ki, keep))
        return s, (s.position, out)

    state, (positions, (kept, acc, div, steps)) = jax.lax.scan(snap_block, state, jax.random.split(k, snaps))
    kept = np.asarray(kept).reshape(snaps * keep, CHAINS, NH).transpose(1, 0, 2)
    return state, positions, kept, {"acceptance": float(acc.mean()), "divergences": int(div.sum()),
                                    "gradients per chain": float(steps.sum())}


def warm_chees(k, start, steps, **options):
    warmup = blackjax.chees_adaptation(logdensity, CHAINS, adaptation_info_fn=get_filter_adapt_info_fn(), **options)
    (state, params), _ = warmup.run(k, start, step_size=0.01, optim=optax.adam(0.05), num_steps=steps)
    return jax.block_until_ready(state), params


def run_chees(k, start, **options):
    """Plain ChEES: one warmup, then sampling."""
    k1, k2 = jax.random.split(k)
    t0 = time.perf_counter()
    state, params = warm_chees(k1, start, WARMUP, **options)
    t1 = time.perf_counter()
    state, positions, kept, stats = sample(blackjax.dynamic_hmc(logdensity, **params).step, state, k2, SNAPS, KEEP, 1)
    jax.block_until_ready(positions)
    t2 = time.perf_counter()
    stats |= {"step size": float(params["step_size"]), "L": float(params["integration_steps_params"][0]),
              "warmup s": t1 - t0, "sampling s": t2 - t1}
    return state, params, positions, kept, stats


key, k = jax.random.split(key)
_, _, positions, kept_id, stats_id = run_chees(k, cold_start)
del positions
summary_id = summarise(kept_id)
print({a: round(b, 4) for a, b in stats_id.items()})
print(summary_id.round(3).to_string())

log_density_batch = jax.jit(jax.vmap(logdensity))


def stragglers(state, threshold=6.0):
    """Chains more than `threshold` robust sds from the ensemble in log density or any hyperparameter."""
    h = np.c_[np.asarray(log_density_batch(state.position)), np.asarray(state.position[:, :NH])]
    median = np.median(h, 0)
    z = np.abs(h - median) / (1.4826 * np.median(np.abs(h - median), 0))
    return (z > threshold).any(1), z


def cull(state, bad):
    """Warmup only: restart flagged chains from copies of healthy ones (their paths separate at once)."""
    source = np.arange(CHAINS)
    source[bad] = rng.choice(np.flatnonzero(~bad), bad.sum())
    return jax.tree.map(lambda leaf: leaf[source], state)


key, k1, k2, k3 = jax.random.split(key, 4)
options = {"mass_matrix_estimation": "diagonal"}
t0 = time.perf_counter()
state, params = warm_chees(k1, cold_start, int(0.3 * WARMUP), **options)
bad, z = stragglers(state)
print(f"stage 1 ({int(0.3 * WARMUP)} steps): step size {float(params['step_size']):.4f}, L = {float(params['integration_steps_params'][0]):.0f}, {bad.sum()} straggler(s); "
      f"worst log density z = {z[:, 0].max():.1f}; stragglers' log density {np.asarray(log_density_batch(state.position))[bad].round(0)}, "
      f"median {np.median(np.asarray(log_density_batch(state.position))):.0f}")
state = cull(state, bad)
state, params_chees = warm_chees(k2, state.position, WARMUP - int(0.3 * WARMUP), **options)
bad, z = stragglers(state)
print(f"stage 2 ({WARMUP - int(0.3 * WARMUP)} steps): step size {float(params_chees['step_size']):.4f}, L = {float(params_chees['integration_steps_params'][0]):.0f}, {bad.sum()} straggler(s) culled")
state = cull(state, bad)
t1 = time.perf_counter()
state_chees, pos_chees, kept_chees, stats_chees = sample(blackjax.dynamic_hmc(logdensity, **params_chees).step, state, k3, SNAPS, KEEP, 1)
jax.block_until_ready(pos_chees)
stats_chees |= {"step size": float(params_chees["step_size"]), "L": float(params_chees["integration_steps_params"][0]),
                "warmup s": t1 - t0, "sampling s": time.perf_counter() - t1}
summary_chees = summarise(kept_chees)
print({a: round(b, 4) for a, b in stats_chees.items()})
print(summary_chees.round(3).to_string())

fig, axes = plt.subplots(1, 2, figsize=(10, 3.6), sharex=True, sharey=True, layout="constrained")
i = HYPER.index("s_fine")
means = [np.exp(kept[..., i]).mean(1) for kept in (kept_id, kept_chees)]
bins = np.linspace(min(m.min() for m in means), max(m.max() for m in means), 40)
for ax, m, title in zip(axes, means, ("ChEES, default settings", "ChEES, diagonal metric + floor")):
    ax.hist(m, bins=bins, color="C0")
    ax.set(title=title, xlabel="per-chain mean of s_fine")
axes[0].set_ylabel("chains");

steps_nuts = 30 if HAS_GPU else 10
nuts = blackjax.nuts(logdensity, step_size=params_chees["step_size"], inverse_mass_matrix=params_chees["inverse_mass_matrix"])
nuts_state = jax.vmap(nuts.init)(state_chees.position)


def nuts_iteration(s, ki):
    s, info = jax.vmap(nuts.step)(jax.random.split(ki, CHAINS), s)
    return s, info.num_integration_steps


nuts_run = jax.jit(lambda k: jax.lax.scan(nuts_iteration, nuts_state, jax.random.split(k, steps_nuts))[1])
key, k = jax.random.split(key)
tree_sizes = np.asarray(jax.block_until_ready(nuts_run(k)))   # compiles
key, k = jax.random.split(key)
t0 = time.perf_counter()
tree_sizes = np.asarray(jax.block_until_ready(nuts_run(k)))
nuts_s = time.perf_counter() - t0

chees_step = blackjax.dynamic_hmc(logdensity, **params_chees).step


def chees_iteration(s, ki):
    s, info = jax.vmap(chees_step)(jax.random.split(ki, CHAINS), s)
    return s, (info.num_integration_steps, s.position.mean())  # the sizes alone are a fixed sequence: without
    #                                                             the position XLA deletes the whole sampler


chees_run = jax.jit(lambda k: jax.lax.scan(chees_iteration, state_chees, jax.random.split(k, steps_nuts))[1])
key, k = jax.random.split(key)
jax.block_until_ready(chees_run(k))
key, k = jax.random.split(key)
t0 = time.perf_counter()
chees_sizes = np.asarray(jax.block_until_ready(chees_run(k))[0])
chees_s = time.perf_counter() - t0
bare = timing.loc[CHAINS, "GPU float32" if HAS_GPU else "CPU float32"]
print(f"vectorised NUTS: {tree_sizes.mean():.0f} leapfrog steps per iteration, largest tree per iteration {tree_sizes.max(1).mean():.0f}; "
      f"tree sizes seen {np.unique(tree_sizes)}; {tree_sizes.sum() / (tree_sizes.max(1).sum() * CHAINS):.0%} of the work needed")
print(pd.DataFrame({"gradients per iteration": [tree_sizes.max(1).mean(), chees_sizes.max(1).mean()],
                    "ms per gradient": [1e3 * nuts_s / tree_sizes.max(1).sum(), 1e3 * chees_s / chees_sizes.max(1).sum()]},
                   index=["NUTS", "ChEES"]).round(2).to_string())
print(f"a bare batched gradient costs {bare:.2f} ms")
del nuts_run, chees_run
jax.clear_caches()

graph = get_jaxified_graph(inputs=model.value_vars, outputs=model.replace_rvs_by_values(list(outputs.values())))
fields = jax.jit(jax.vmap(lambda q: graph(*unravel(q))))
k_ha = int(100 / CELL)
ny_ha, nx_ha = y.shape[0] // k_ha, y.shape[1] // k_ha


def hectares(cellwise):
    n = cellwise.shape[0]
    return cellwise.reshape(n, ny_ha, k_ha, nx_ha, k_ha).sum((2, 4))


t0 = time.perf_counter()
sum_eta = sum_hab = 0.0
pos_elev, pos_grad = 0, 0
habitat_ha, contrast = [], []
sim_rng = np.random.default_rng(RANDOM_SEED)
cell_of = lambda x, y_m: (int(y_m / CELL), int(x / CELL))
WEST, PLATEAU = cell_of(200, 250), cell_of(680, 150)
for snapshot in pos_chees:
    eta, habitat, slope_elev, slope_grad = [np.asarray(a) for a in fields(snapshot)]
    sum_eta = sum_eta + np.exp(eta).sum(0)
    sum_hab = sum_hab + habitat.sum(0)
    pos_elev = pos_elev + (slope_elev > 0).sum(0)
    pos_grad = pos_grad + (slope_grad > 0).sum(0)
    habitat_ha.append(hectares(np.exp(habitat) * (CELL**2 / 1e4)))
    contrast.append(slope_elev[(slice(None), *WEST)] - slope_elev[(slice(None), *PLATEAU)])
n_draws = SNAPS * CHAINS
mean_lambda = sum_eta / n_draws
mean_habitat = sum_hab / n_draws
p_elev, p_grad = pos_elev / n_draws, pos_grad / n_draws
habitat_ha = np.concatenate(habitat_ha).reshape(n_draws, -1)
contrast = np.concatenate(contrast)
in_sample = sim_rng.poisson(np.exp(eta)).astype(float)   # the last snapshot: CHAINS independent draws
# In-sample replicates inherit the fitted clumps. A fair test draws NEW clumps: fresh z_fine, same hyperparameters.
z0 = NH + int(np.prod(M_SMOOTH))
z1 = z0 + int(np.prod(M_FINE))
key, k = jax.random.split(key)
fresh_clumps = pos_chees[-1].at[:, z0:z1].set(jax.random.normal(k, (CHAINS, z1 - z0), q0.dtype))
fresh = sim_rng.poisson(np.exp(np.asarray(fields(fresh_clumps)[0]))).astype(float)
postprocess_s = time.perf_counter() - t0
print(f"{n_draws:,} draws of four {y.shape} fields in {postprocess_s:.1f} s")
print(f"expected total trees: {mean_lambda.sum():,.0f} (observed {n_trees:,})")
print(f"elevation effect, west (200, 250) minus plateau (680, 150): mean {contrast.mean():.2f}, "
      f"90% interval {np.percentile(contrast, 5):.2f} to {np.percentile(contrast, 95):.2f}, P(> 0) = {(contrast > 0).mean():.3f}")

blocks = (1, 2, 4, 10, 20) if HAS_GPU else (1, 5)
ppc = pd.DataFrame({"block m": [b * CELL for b in blocks],
                    "observed": [vmr(y, b) for b in blocks],
                    "fitted clumps 5%": [np.percentile([vmr(r, b) for r in in_sample], 5) for b in blocks],
                    "fitted clumps 95%": [np.percentile([vmr(r, b) for r in in_sample], 95) for b in blocks],
                    "new clumps 5%": [np.percentile([vmr(r, b) for r in fresh], 5) for b in blocks],
                    "new clumps 95%": [np.percentile([vmr(r, b) for r in fresh], 95) for b in blocks]})
print(ppc.round(2).to_string(index=False))

fig, axes = plt.subplots(3, 1, figsize=(9, 11), sharex=True, layout="constrained")
panels = ((mean_habitat, "habitat: posterior mean log trees per hectare, without fine clusters", "viridis", {}),
          (p_elev, "P(local elevation effect > 0)", "RdBu_r", {"vmin": 0, "vmax": 1}),
          (p_grad, "P(local gradient effect > 0)", "RdBu_r", {"vmin": 0, "vmax": 1}))
for ax, (v, title, cmap, opts) in zip(axes, panels):
    im = ax.imshow(v, origin="lower", extent=extent, cmap=cmap, **opts)
    ax.set(title=title, ylabel="y (m)")
    fig.colorbar(im, ax=ax, shrink=0.8)
axes[2].set_xlabel("x (m)");