Imperfect detection: occupancy, N-mixture and capture-recapture models

keep, Y, D, S = season("99")
site = cb[keep].reset_index(drop=True)
ELE = (site.ele.to_numpy() - 1200) / 600           # standardised with round constants, so the
FOREST = (site.forest.to_numpy() - 35) / 28        # same scale is used in every season
DATE = np.where(S, (D - 55) / 20, 0.0)             # day of season: 0 = 25 April-ish, 1 unit = 20 days
n_site = len(site)
hist = pd.Series(["".join(map(str, r)) if s.all() else "".join(map(str, r[:2])) + "-" for r, s in zip(Y, S)])
print(f"1999: {n_site} quadrats, {S.sum()} visits, {Y.sum()} detections on "
      f"{(Y.sum(1) > 0).sum()} quadrats")
hist.value_counts().rename("quadrats").to_frame().T

fig, axes = plt.subplots(1, 2, figsize=(11, 4))
det = Y.sum(1) > 0
axes[0].scatter(site.ele[~det], site.forest[~det], s=14, color=GREY, alpha=0.6, label="never detected")
axes[0].scatter(site.ele[det], site.forest[det], s=14 + 14 * Y.sum(1)[det], color=BLUE,
                label="detected (size = number of visits)")
axes[0].set(xlabel="elevation (m)", ylabel="forest cover (%)", title="1999: where the crossbill was seen")
axes[0].legend(loc="upper left", fontsize=8)
for k, c in zip(range(3), [BLUE, AQUA, ORANGE]):
    ok = S[:, k]
    axes[1].hist(D[ok, k], bins=np.arange(0, 125, 5), color=c, alpha=0.6, label=f"visit {k + 1}")
axes[1].set(xlabel="day of the season", ylabel="visits", title="When the visits happened")
axes[1].legend();

with occ:
    prior = pm.sample_prior_predictive(1000, random_seed=RANDOM_SEED)
n_det_prior = (prior.prior_predictive["y"].sum("visit") > 0).sum("site").to_numpy().ravel()
fig, ax = plt.subplots(figsize=(7, 3.2))
ax.hist(n_det_prior, bins=40, color=GREY)
ax.axvline(det.sum(), color=ORANGE, lw=2, label=f"observed ({det.sum()})")
ax.set(xlabel="quadrats with at least one detection", ylabel="prior draws",
       title=f"Prior predictive, {n_site} quadrats")
ax.legend();

idata_occ = fit(occ, "occupancy, 1999")
print("warm-up steps:", idata_occ.posterior.attrs.get("tuning_steps"))
az.summary(idata_occ, var_names=["a0", "a", "b0", "b"], ci_kind="hdi", ci_prob=0.94, round_to=2)

with pm.Model(coords={"psi_cov": PSI_COVS}) as naive:
    a0 = pm.Normal("a0", 0.0, 1.5)
    a = pm.Normal("a", 0.0, 1.0, dims="psi_cov")
    psi_naive = pm.Deterministic("psi", pm.math.invlogit(a0 + np.c_[ELE, ELE**2, FOREST] @ a))
    pm.Bernoulli("detected", psi_naive, observed=det.astype(int))
idata_naive = fit(naive, "naive logistic, 1999")

psi_bar = draws(idata_occ, "psi").mean(1)            # average occupancy over the 245 quadrats
psi_bar_naive = draws(idata_naive, "psi").mean(1)
for lab, x in [("naive (detected = present)", psi_bar_naive), ("occupancy model", psi_bar)]:
    print(f"{lab:<28s} mean occupancy {x.mean():.3f}  94% HDI {az.hdi(x, prob=0.94).round(3)}")
print(f"raw share detected: {det.mean():.3f};  occupancy / naive = {np.mean(psi_bar / psi_bar_naive):.2f}")

ele_grid = np.linspace(250, 2750, 101)
eg = (ele_grid - 1200) / 600
date_grid = np.linspace(10, 115, 100)
dg = (date_grid - 55) / 20


def band(ax, x, samples, color, label):
    lo, hi = np.quantile(samples, [0.055, 0.945], axis=0)
    ax.fill_between(x, lo, hi, color=color, alpha=0.25, lw=0)
    ax.plot(x, samples.mean(0), color=color, lw=2, label=label)


fig, axes = plt.subplots(1, 2, figsize=(12, 4.2))
for idata_, c, lab in [(idata_naive, ORANGE, "naive: detected = present"), (idata_occ, BLUE, "occupancy model")]:
    a0_, a_ = draws(idata_, "a0"), draws(idata_, "a")
    band(axes[0], ele_grid, expit(a0_[:, None] + a_[:, [0]] * eg + a_[:, [1]] * eg**2), c, lab)
bins = np.arange(250, 2800, 250)
idx = np.digitize(site.ele, bins)
axes[0].scatter([site.ele[idx == k].mean() for k in np.unique(idx)], [det[idx == k].mean() for k in np.unique(idx)],
                color="k", s=18, zorder=3, label="share detected (250 m bins)")
axes[0].set(xlabel="elevation (m)", ylabel="occupancy ψ", ylim=(0, 1),
            title="Occupancy vs elevation (forest at 35%), 89% bands")
axes[0].legend(fontsize=8)
b0_, b_ = draws(idata_occ, "b0"), draws(idata_occ, "b")
p_date = expit(b0_[:, None] + b_[:, [0]] * dg + b_[:, [1]] * dg**2)
band(axes[1], date_grid, p_date, BLUE, "p per visit")
band(axes[1], date_grid, 1 - (1 - p_date) ** 3, AQUA, "P(at least one detection in 3 such visits)")
axes[1].set(xlabel="day of the season", ylabel="probability", ylim=(0, 1),
            title="Detection given presence")
axes[1].legend(fontsize=8);

with pm.Model(coords={"site": np.arange(n_site), "psi_cov": PSI_COVS}) as occ_z:
    a0 = pm.Normal("a0", 0.0, 1.5)
    a = pm.Normal("a", 0.0, 1.0, dims="psi_cov")
    b0 = pm.Normal("b0", 0.0, 1.5)
    b = pm.Normal("b", 0.0, 1.0, shape=2)
    psi = pm.math.invlogit(a0 + np.c_[ELE, ELE**2, FOREST] @ a)
    p = pm.math.invlogit(b0 + b[0] * DATE + b[1] * DATE**2) * S          # p = 0 on visits that did not happen
    z = pm.Bernoulli("z", psi, dims="site")
    pm.Bernoulli("y", z[:, None] * p, observed=Y)
occ_marg = marginalize(occ_z, ["z"])
idata_marg = fit(occ_marg, "same model via marginalize")
with occ_marg:
    idata_rec = recover(idata_marg.isel(draw=slice(None, None, 4)), var_names=["z"])
cmp = pd.DataFrame({v: [f"{float(idata_occ.posterior[v].mean()):.3f}", f"{float(idata_marg.posterior[v].mean()):.3f}"]
                    for v in ["a0", "b0"]}, index=["CustomDist", "marginalize"])
cmp["alpha (ele, ele², forest)"] = [draws(idata_occ, "a").mean(0).round(2), draws(idata_marg, "a").mean(0).round(2)]
cmp

psi_d = draws(idata_occ, "psi", 1000)                                   # (draw, site)
b0_d, b_d = draws(idata_occ, "b0", 1000), draws(idata_occ, "b", 1000)
p_d = expit(b0_d[:, None, None] + b_d[:, None, None, 0] * DATE + b_d[:, None, None, 1] * DATE**2)
miss_all = np.prod(np.where(S, 1 - p_d, 1.0), axis=2)                   # P(all visits missed | z = 1)
q = np.where(det, 1.0, psi_d * miss_all / (psi_d * miss_all + 1 - psi_d))   # P(z_i = 1 | y)
z_draws = rng.binomial(1, q)
n_occ = z_draws.sum(1)
q_rec = idata_rec.posterior["z"].mean(("chain", "draw")).to_numpy()
print(f"P(z=1|y) for never-detected quadrats: NumPy vs recover() max abs difference "
      f"{np.abs(q.mean(0) - q_rec)[~det].max():.3f} (recover uses 1000 draws of z)")
print(f"occupied quadrats in 1999: {n_occ.mean():.1f}, 94% HDI {az.hdi(n_occ, prob=0.94)} "
      f"(detected on {det.sum()}; landscape expectation sum(psi) = {psi_d.sum(1).mean():.1f})")

fig, axes = plt.subplots(1, 2, figsize=(12, 4.2), gridspec_kw={"width_ratios": [1.6, 1]})
sc = axes[0].scatter(site.ele[~det], q.mean(0)[~det], c=site.forest[~det], cmap="viridis", s=22)
axes[0].scatter(site.ele[det], np.ones(det.sum()), color=ORANGE, marker="|", s=80, label="detected: z = 1")
fig.colorbar(sc, ax=axes[0], label="forest cover (%)")
axes[0].set(xlabel="elevation (m)", ylabel="P(occupied | never detected)", ylim=(-0.03, 1.05),
            title="Never-detected quadrats: how likely is the bird there anyway?")
axes[0].legend(loc="center right", fontsize=8)
axes[1].hist(n_occ, bins=np.arange(n_occ.min() - 0.5, n_occ.max() + 1.5), color=BLUE)
axes[1].axvline(det.sum(), color=ORANGE, lw=2, label=f"detected ({det.sum()})")
axes[1].set(xlabel="occupied quadrats (of 245)", ylabel="posterior draws", title="Finite-sample occupancy, 1999")
axes[1].legend();

with occ:
    ppc = pm.sample_posterior_predictive(idata_occ.isel(draw=slice(None, None, 4)), random_seed=RANDOM_SEED,
                                         progressbar=False)
yrep = ppc.posterior_predictive["y"].stack(sample=("chain", "draw")).transpose("sample", ...).to_numpy().astype(int)
three = S.all(1)
codes = ["".join(c) for c in np.array(np.meshgrid(*[["0", "1"]] * 3, indexing="ij")).reshape(3, -1).T]
w = np.array([4, 2, 1])
obs_f = np.bincount(Y[three] @ w, minlength=8)
rep_f = np.stack([np.bincount(r @ w, minlength=8) for r in yrep[:, three]])
fig, axes = plt.subplots(1, 2, figsize=(12, 4), gridspec_kw={"width_ratios": [2.2, 1]})
lo, hi = np.quantile(rep_f, [0.03, 0.97], axis=0)
axes[0].bar(range(1, 8), obs_f[1:], color=GREY, label="observed")
axes[0].errorbar(range(1, 8), rep_f.mean(0)[1:], yerr=[rep_f.mean(0)[1:] - lo[1:], hi[1:] - rep_f.mean(0)[1:]],
                 fmt="o", color=BLUE, capsize=4, label="posterior predictive, 94%")
axes[0].set_xticks(range(1, 8), codes[1:])
axes[0].set(xlabel="detection history (visits 1-2-3)", ylabel="quadrats",
            title=f"History frequencies, {three.sum()} quadrats with 3 visits (000 in the text)")
axes[0].legend()
n_rep = (yrep.sum(2) > 0).sum(1)
axes[1].hist(n_rep, bins=25, color=BLUE, alpha=0.7)
axes[1].axvline(det.sum(), color=ORANGE, lw=2, label="observed")
axes[1].set(xlabel="quadrats with a detection", title="All 245 quadrats")
axes[1].legend()
print(f"000 histories: observed {obs_f[0]}, predicted {rep_f[:, 0].mean():.1f} "
      f"[{lo[0]:.0f}, {hi[0]:.0f}]");

loos, fits_p = {}, {}
for label, covs in [("p constant", ()), ("p ~ date + date²", ("date", "date²")),
                    ("p ~ date + date² + ele", ("date", "date²", "ele"))]:
    if covs == ("date", "date²"):
        m_, idata_ = occ, idata_occ
    else:
        m_ = occupancy_model(Y, S, ELE, FOREST, DATE, p_covs=covs)
        idata_ = fit(m_, label)
    pm.compute_log_likelihood(idata_, model=m_, progressbar=False)
    loos[label] = az.loo(idata_)
    fits_p[label] = idata_
b_ele = draws(fits_p["p ~ date + date² + ele"], "b")[:, 2]
print(f"detection slope on elevation: {b_ele.mean():.2f}, 94% HDI {az.hdi(b_ele, prob=0.94).round(2)}")
print(f"mean occupancy (245 quadrats): " + ", ".join(
    f"{k}: {draws(v, 'psi').mean():.3f}" for k, v in fits_p.items()))
az.compare(loos, round_to=1)

m_pe = fits_p["p ~ date + date² + ele"]
fig, axes = plt.subplots(1, 2, figsize=(12, 4.2))
for lab, c in [("p ~ date + date²", BLUE), ("p ~ date + date² + ele", PURPLE)]:
    a0_, a_ = draws(fits_p[lab], "a0"), draws(fits_p[lab], "a")
    band(axes[0], ele_grid, expit(a0_[:, None] + a_[:, [0]] * eg + a_[:, [1]] * eg**2), c, lab)
axes[0].set(xlabel="elevation (m)", ylabel="occupancy ψ", ylim=(0, 1), title="Occupancy vs elevation (forest 35%)")
axes[0].legend(fontsize=8)
b0_, b_ = draws(m_pe, "b0"), draws(m_pe, "b")
band(axes[1], ele_grid, expit(b0_[:, None] + b_[:, [2]] * eg), PURPLE, "p ~ ... + ele (at day 55)")
b0_, b_ = draws(idata_occ, "b0"), draws(idata_occ, "b")
band(axes[1], ele_grid, expit(b0_[:, None] + 0 * eg), BLUE, "p ~ date + date² (at day 55)")
axes[1].set(xlabel="elevation (m)", ylabel="detection per visit", ylim=(0, 1),
            title="Detection vs elevation")
axes[1].legend(fontsize=8);

rows = []
for yy in YEARS:
    keep_, Y_, D_, S_ = season(yy)
    s_ = cb[keep_]
    ele_, forest_ = (s_.ele.to_numpy() - 1200) / 600, (s_.forest.to_numpy() - 35) / 28
    date_ = np.where(S_, (D_ - 55) / 20, 0.0)
    year = 1999 if yy == "99" else 2000 + int(yy)
    for lab, covs in [("date", ("date", "date²")), ("date+ele", ("date", "date²", "ele"))]:
        m_ = occupancy_model(Y_, S_, ele_, forest_, date_, p_covs=covs)
        id_ = fit(m_, f"{year}, p ~ {lab}")
        pb = draws(id_, "psi").mean(1)
        b0_, b_ = draws(id_, "b0", 1000), draws(id_, "b", 1000)
        X_ = np.stack([{"date": date_, "date²": date_**2, "ele": ele_[:, None] + 0 * date_}[c] for c in covs], -1)
        p_vis = expit(b0_[:, None, None] + np.einsum("svk,dk->dsv", X_, b_))[:, S_]   # p on the visits made
        rows.append({"year": year, "p model": lab, "naive": (Y_.sum(1) > 0).mean(), "psi_mean": pb.mean(),
                     "psi_lo": np.quantile(pb, 0.055), "psi_hi": np.quantile(pb, 0.945),
                     "p_mean": p_vis.mean(), "p_lo": np.quantile(p_vis.mean(1), 0.055),
                     "p_hi": np.quantile(p_vis.mean(1), 0.945)})
        del id_
years = pd.DataFrame(rows)
years.pivot(index="year", columns="p model", values=["naive", "psi_mean", "p_mean"]).round(3)

fig, axes = plt.subplots(2, 1, figsize=(10, 6.5), sharex=True)
for (lab, c, dx) in [("date", BLUE, -0.1), ("date+ele", PURPLE, 0.1)]:
    d_ = years[years["p model"] == lab]
    axes[0].errorbar(d_.year + dx, d_.psi_mean, yerr=[d_.psi_mean - d_.psi_lo, d_.psi_hi - d_.psi_mean],
                     fmt="o-", color=c, capsize=3, label=f"occupancy model, p ~ {lab}")
    axes[1].errorbar(d_.year + dx, d_.p_mean, yerr=[d_.p_mean - d_.p_lo, d_.p_hi - d_.p_mean],
                     fmt="o-", color=c, capsize=3, label=f"p ~ {lab}")
d_ = years[years["p model"] == "date"]
axes[0].plot(d_.year, d_.naive, "s--", color=ORANGE, label="naive: share of quadrats detected")
axes[0].set(ylabel="mean occupancy", ylim=(0, 0.85), title="Crossbill occupancy by season (89% intervals)")
axes[0].legend(fontsize=8, ncol=3, loc="upper center")
axes[1].set(ylabel="mean detection per visit", xlabel="season", ylim=(0, 1))
axes[1].legend(fontsize=8);

d_ = years[years["p model"] == "date"]
print("correlation across seasons, naive vs mean detection p:", np.corrcoef(d_.naive, d_.p_mean)[0, 1].round(2))
print("correlation across seasons, occupancy (p ~ date) vs mean p:", np.corrcoef(d_.psi_mean, d_.p_mean)[0, 1].round(2))
print("range across seasons: naive", d_.naive.min().round(3), "-", d_.naive.max().round(3),
      "| occupancy (p ~ date)", d_.psi_mean.min().round(3), "-", d_.psi_mean.max().round(3))

PRIORS = {"default N(0, 1.5)": (0.0, 1.5), "'hard' N(-2, 0.5)": (-2.0, 0.5), "'easy' N(1, 0.5)": (1.0, 0.5)}
rows = []
for n_vis in (1, 2, 3):
    S_k = S.copy()
    S_k[:, n_vis:] = False
    ok = S_k.any(1)
    for plab, pr in PRIORS.items():
        m_ = occupancy_model(Y[ok] * S_k[ok], S_k[ok], ELE[ok], FOREST[ok], np.where(S_k, DATE, 0)[ok], b0_prior=pr)
        id_ = fit(m_, f"{n_vis} visit(s), {plab}")
        pb = draws(id_, "psi").mean(1)
        rows.append({"visits": n_vis, "prior": plab, "mean": pb.mean(), "lo": np.quantile(pb, 0.055),
                     "hi": np.quantile(pb, 0.945), "p_day55": expit(draws(id_, "b0")).mean()})
        del id_
conf = pd.DataFrame(rows)
conf.round(3)

fig, ax = plt.subplots(figsize=(9, 4))
for k, (plab, c) in enumerate(zip(PRIORS, [BLUE, ORANGE, AQUA])):
    d_ = conf[conf.prior == plab]
    ax.errorbar(d_.visits + (k - 1) * 0.1, d_["mean"], yerr=[d_["mean"] - d_.lo, d_.hi - d_["mean"]],
                fmt="o", color=c, capsize=4, ms=7, label=f"prior on detection: {plab}")
naive_by_k = [((Y * (np.arange(3) < k)).sum(1) > 0).mean() for k in (1, 2, 3)]
ax.plot([1, 2, 3], naive_by_k, "s--", color=GREY, label="naive share detected")
ax.set(xticks=[1, 2, 3], xlabel="visits used per quadrat", ylabel="mean occupancy, 1999", ylim=(0, 1),
       title="Few visits: the detection prior decides the occupancy (89% intervals)")
ax.legend(fontsize=8);

data.describe("mallard")
mal = read_rdata(data.path("mallard"))
C_raw = mal["mallard.y"].reshape(3, -1).T                         # R stores matrices column by column
ivel_raw = np.asarray(mal["mallard.obs"]["ivel"]).reshape(3, -1).T
date_raw = np.asarray(mal["mallard.obs"]["date"]).reshape(3, -1).T
S_m = ~np.isnan(C_raw) & ~np.isnan(ivel_raw) & ~np.isnan(date_raw)
keep_m = S_m.any(1)
print(f"{keep_m.sum()} of {len(C_raw)} routes with at least one count; {(~S_m[keep_m]).sum()} missing counts")
C = np.where(S_m, C_raw, 0).astype(int)[keep_m]
S_m = S_m[keep_m]
IVEL, DATE_M = np.where(S_m, ivel_raw[keep_m], 0), np.where(S_m, date_raw[keep_m], 0)
X_lam = pd.DataFrame(mal["mallard.site"])[keep_m].to_numpy()
LAM_COVS = ["elev", "length", "forest"]
n_route = len(C)
print("count distribution:", dict(zip(*np.unique(C[S_m], return_counts=True))))
print(f"routes with any mallard: {(C.max(1) > 0).sum()};  sum over routes of the max count: {C.max(1).sum()}")

az.summary(nm_fits["NegBin"], var_names=["a0", "a", "b0", "b", "inv_r"], ci_kind="hdi", ci_prob=0.94, round_to=2)

def n_posterior(idata, kind, n=200, K_=K):
    """P(N_i = N | counts, parameters) for n posterior draws: array (draw, route, N)."""
    lam_d, p_d = draws(idata, "lam", n), draws(idata, "p", n)
    Ng = np.arange(K_ + 1)
    Nb, cb_ = Ng[None, None, :, None], C[None, :, None, :]
    lb = (gammaln(Nb + 1) - gammaln(cb_ + 1) - gammaln(np.maximum(Nb - cb_, 0) + 1)
          + cb_ * np.log(p_d[:, :, None, :]) + (Nb - cb_) * np.log1p(-p_d[:, :, None, :]))
    lb = np.where(Nb >= cb_, lb, -np.inf)
    ll = np.sum(np.where(S_m[None, :, None, :], lb, 0.0), axis=-1)
    lam_ = lam_d[..., None]
    if kind == "Poisson":
        lprior = Ng * np.log(lam_) - lam_ - gammaln(Ng + 1)
    else:
        r = 1 / draws(idata, "inv_r", n)[:, None, None]
        lprior = gammaln(Ng + r) - gammaln(r) - gammaln(Ng + 1) + r * np.log(r / (r + lam_)) + Ng * np.log(lam_ / (r + lam_))
    lj = lprior + ll
    return np.exp(lj - logsumexp(lj, axis=-1, keepdims=True))


N_post, N_tot = {}, {}
for kind, id_ in nm_fits.items():
    P = n_posterior(id_, kind)
    N_post[kind] = P
    cdf = P.cumsum(-1)
    N_draw = (rng.random(P.shape[:2] + (1,)) > cdf).sum(-1)              # one N_i per draw and route
    N_tot[kind] = N_draw.sum(1)
    tail = P[..., -10:].sum(-1)                                          # mass in the top 10 values of N
    print(f"{kind:8s} max over routes/draws of P(N in [91, 100] | data): {tail.max():.1e};  "
          f"mallards on the {n_route} routes: {N_tot[kind].mean():.0f}, 94% HDI {az.hdi(N_tot[kind], prob=0.94)}"
          f"  (sum of max counts {C.max(1).sum()})")