docs/experiments/EXP-3D-CAJU3D-JEPAGATE-V1/analysis/gate_dropped_profile.py
"""EXP-3D-CAJU3D-JEPAGATE-V1 descriptive diagnostic (after the DEV decision; arrived frames only, no ground truth, no
TTT): which measurements each gate drops. Recomputes the DEV episodes' scores and gates exactly as gate_ttt.py does
(same seeds, same generator order) and profiles the dropped pixels against all valid pixels:
no cross-view hit no point of the other six views lands on the pixel (DEPTH's z-buffer; DEPTH scores it 0)
no JEPA cell hit no source latent cell is warped onto the pixel's latent cell
cosine, edge incidence cosine and depth jump >= 0.1 (engine_core.incidence_weight)
relative depth depth / the view's median valid depth
stream view one of the 4 stream views (else one of the 3 context views)
log P, log err the JEPA's predicted log variance and log mean (z - z_hat)^2 of the pixel's cell
depth / photometric disagreement | hit DEPTH's score and the channel-mean |RGB - RGB of the nearest point of the
other views|, over the pixels with a cross-view hit
dropped neighbour share the share of a dropped pixel's valid 4-neighbours that are dropped too (spatial clumping;
a scattered 10 % drop gives ~0.10)
Run with the DepthSplat environment from the DepthSplat repository root (PYTHONPATH=.):
python gate_dropped_profile.py ROOT OUT.json
"""
import hashlib
import json
import sys
from pathlib import Path
import numpy as np
import torch
import torch.nn.functional as F
HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE.parent / "code"))
import gate_ttt as T # noqa: E402 (path set-up and imports of the replay)
G, J, v1, engine_core, ds_predict = T.G, T.J, T.v1, T.engine_core, T.ds_predict
ARMS = ("JEPA", "DEPTH", "INC", "RANDOM")
COHORTS = ("DEV_SYSTEM", "DEV_SYSTEM2")
def cross_view_hit(planar, rgb, k, c2w, min_depth=0.05):
"""The z-buffer of gates.depth_disagreement: the hit mask [V, H, W] and the photometric disagreement [V, H, W]
(channel-mean |RGB_v - RGB of the nearest point of the other six views|, 0 where no hit)."""
v_n, h, w = planar.shape
rows = torch.arange(h, dtype=torch.float32, device=planar.device).view(h, 1)
cols = torch.arange(w, dtype=torch.float32, device=planar.device).view(1, w)
pts, ok, col = [], [], []
for s in range(v_n):
d = planar[s]
cam = torch.stack([(cols - k[s, 0, 2]) / k[s, 0, 0] * d, (rows - k[s, 1, 2]) / k[s, 1, 1] * d, d], dim=-1)
pts.append(cam.reshape(-1, 3) @ c2w[s, :3, :3].T + c2w[s, :3, 3])
ok.append((d > 0).reshape(-1))
col.append(rgb[s].reshape(3, -1).T)
hits, photo = [], []
for v in range(v_n):
others = [s for s in range(v_n) if s != v]
p = torch.cat([pts[s] for s in others])
valid = torch.cat([ok[s] for s in others])
c = torch.cat([col[s] for s in others])
cam = (p - c2w[v, :3, 3]) @ c2w[v, :3, :3]
z = cam[:, 2]
zs = z.clamp_min(1e-6)
u = torch.round(k[v, 0, 0] * cam[:, 0] / zs + k[v, 0, 2]).long()
r = torch.round(k[v, 1, 1] * cam[:, 1] / zs + k[v, 1, 2]).long()
good = valid & (z > min_depth) & (u >= 0) & (u < w) & (r >= 0) & (r < h)
flat = r.clamp(0, h - 1) * w + u.clamp(0, w - 1)
key = torch.where(good, z, torch.full_like(z, float("inf")))
nearest = torch.full((h * w,), float("inf"), device=planar.device).scatter_reduce(0, flat, key, reduce="amin")
win = (good & (key <= nearest[flat] * (1 + 1e-6))).float()
count = torch.zeros(h * w, device=planar.device).scatter_add(0, flat, win)
total = torch.zeros(h * w, 3, device=planar.device).scatter_add(0, flat[:, None].expand(-1, 3), c * win[:, None])
hit = (count > 0).view(h, w) & (planar[v] > 0)
rep = (total / count.clamp_min(1.0)[:, None]).T.reshape(3, h, w)
hits.append(hit)
photo.append(torch.where(hit, (rgb[v] - rep).abs().mean(dim=0), torch.zeros_like(planar[v])))
return torch.stack(hits), torch.stack(photo)
def neighbour_share(mask, valid):
"""Per pixel of `mask`: the share of its valid 4-neighbours that are also in `mask`."""
m, ok = mask.float(), valid.float()
pad = lambda x: F.pad(x[:, None], (1, 1, 1, 1))[:, 0] # noqa: E731
mp, op = pad(m), pad(ok)
same = mp[:, :-2, 1:-1] + mp[:, 2:, 1:-1] + mp[:, 1:-1, :-2] + mp[:, 1:-1, 2:]
total = op[:, :-2, 1:-1] + op[:, 2:, 1:-1] + op[:, 1:-1, :-2] + op[:, 1:-1, 2:]
return same / total.clamp_min(1.0)
@torch.no_grad()
def jepa_terms(model, rgb, ray, k, c2w):
"""Per view: log P, log mean latent error and the cell hit of the leave-one-out prediction, at the pixels."""
out = {"log_p": [], "log_err": [], "cell_hit": []}
for v in range(rgb.shape[0]):
others = [s for s in range(rgb.shape[0]) if s != v]
res = model(rgb[others][None], ray[others][None], k[others][None], c2w[others][None],
rgb[v][None], ray[v][None], k[v][None], c2w[v][None])
err = ((res["z"] - res["z_hat"]) ** 2).mean(dim=1, keepdim=True)
for name, x in (("log_p", res["log_p"]), ("log_err", torch.log(err.clamp_min(1e-12))),
("cell_hit", res["hit"])):
out[name].append(F.interpolate(x, size=(J.H, J.W), mode="nearest")[0, 0])
return {n: torch.stack(x) for n, x in out.items()}
def main():
root, out_path = Path(sys.argv[1]).resolve(), Path(sys.argv[2])
device = "cuda"
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
locks = json.loads((root / "frame_role_lock.json").read_text())["cohorts"]
jepa = J.GeometricReverseJEPA()
jepa.load_state_dict(torch.load(root / "model/hjepa_final.pth", map_location="cpu", weights_only=False)["state_dict"])
jepa = jepa.to(device).eval()
sums = {name: {} for name in ("ALL",) + ARMS}
pooled = {"log_nu": [], "log_p": [], "log_err": []}
agree = {arm: [] for arm in ARMS}
episodes = 0
for cohort in COHORTS:
lock = locks[cohort]
if T.sha(lock["manifest"]) != lock["manifest_sha256"]:
raise PermissionError(f"{cohort}: manifest changed")
records = {r["scene_id"]: r for r in json.loads(Path(lock["manifest"]).read_text())["scenes"]}
sealed = {}
for part in sorted((root / "gate/DEV").glob("gate_seal_*of*.json")):
sealed.update(json.loads(part.read_text())["episodes_log"])
for e in lock["episodes"]:
frames = records[e["scene_id"]]["frames"]
rows = {f["frame_id"]: f for f in frames}
inputs = e["context"] + e["stream"]
key = f"{e['scene_id']}/{e['role']}"
generator = torch.Generator(device=device).manual_seed(int(hashlib.sha256(key.encode()).hexdigest()[:8], 16))
scores, valid = T.episode_scores(jepa, frames, inputs, device, generator)
gates = {arm: G.matched_gate(scores[arm], valid, T.SHARE, generator) for arm in ARMS}
dropped = {arm: valid & (g == 0) for arm, g in gates.items()}
overlap = {f"{a}|{b}": float((dropped[a] & dropped[b]).sum() / max(int(dropped[a].sum()), 1))
for i, a in enumerate(ARMS) for b in ARMS[i + 1:]}
for name, v in sealed[key]["arms"]["dropped_overlap"].items():
agree[name.split("|")[0]].append(abs(v - overlap[name]))
planar = T.measured_z(frames, inputs).to(device)
k = torch.tensor(np.array([rows[i]["intrinsics"] for i in inputs], dtype=np.float32), device=device)
c2w = torch.tensor(np.array([ds_predict.orthonormal(np.array(rows[i]["c2w"], dtype=np.float64))
for i in inputs], dtype=np.float32), device=device)
rgb = T.views(frames, inputs, None, v1.NEAR, v1.FAR, device)["image"][0]
terms = jepa_terms(jepa, rgb, planar * J.ray_factor(k, device=device), k, c2w)
_, cosine, jump = engine_core.incidence_weight(planar, k, 0.1, 0.4, 0.1)
med = torch.stack([planar[v][planar[v] > 0].median() for v in range(planar.shape[0])]).view(-1, 1, 1)
stream = torch.zeros_like(valid)
stream[len(e["context"]):] = True
hit, photo = cross_view_hit(planar, rgb, k, c2w)
# name -> (per-pixel value, condition): conditional features average over the set's pixels with a hit
feats = {"no_cross_view_hit": ((~hit).float(), None),
"no_jepa_cell_hit": ((terms["cell_hit"] == 0).float(), None), "cosine": (cosine, None),
"edge": ((jump >= 0.1).float(), None), "relative_depth": (planar / med, None),
"stream_view": (stream.float(), None), "log_p": (terms["log_p"], None),
"log_err": (terms["log_err"], None), "depth_disagreement|hit": (scores["DEPTH"], hit),
"photometric_disagreement|hit": (photo, hit)}
for name, mask in (("ALL", valid),) + tuple(dropped.items()):
s = sums[name]
s["pixels"] = s.get("pixels", 0) + int(mask.sum())
for f, (x, cond) in feats.items():
m = mask if cond is None else mask & cond
s[f] = s.get(f, 0.0) + float(x[m].double().sum())
if cond is not None:
s[f"#{f}"] = s.get(f"#{f}", 0) + int(m.sum())
share = neighbour_share(mask, valid)
s["dropped_neighbour_share"] = s.get("dropped_neighbour_share", 0.0) + float(share[mask].double().sum())
pick = valid & (torch.rand(valid.shape, generator=generator, device=device) < 0.02)
pooled["log_nu"].append(torch.log(scores["JEPA"][pick].clamp_min(1e-12)).cpu())
pooled["log_p"].append(terms["log_p"][pick].cpu())
pooled["log_err"].append(terms["log_err"][pick].cpu())
episodes += 1
profile = {}
for name, s in sums.items():
profile[name] = {"pixels": s["pixels"]}
for f, v in s.items():
if f == "pixels" or f.startswith("#"):
continue
profile[name][f] = v / (s[f"#{f}"] if f"#{f}" in s else s["pixels"])
x = {n: torch.cat(v).double().numpy() for n, v in pooled.items()}
corr = {f"log_nu~{n}": float(np.corrcoef(x["log_nu"], x[n])[0, 1]) for n in ("log_p", "log_err")}
result = {"episodes": episodes, "profile": profile, "pooled_correlations": corr,
"overlap_reproduction_max_abs_diff": {a: float(max(v)) for a, v in agree.items() if v},
"ground_truth_read": False, "query_frames_read": False}
out_path.write_text(json.dumps(result, indent=1))
names = ("ALL",) + ARMS
print("feature".ljust(30), "".join(n.rjust(10) for n in names))
for f in ("pixels",) + tuple(feats) + ("dropped_neighbour_share",):
print(f.ljust(30), "".join(f"{profile[n][f]:10.4f}" if f != "pixels" else f"{profile[n][f]:10d}" for n in names))
print("pooled correlations", json.dumps({k: round(v, 3) for k, v in corr.items()}))
print("overlap reproduction max |diff|", json.dumps(result["overlap_reproduction_max_abs_diff"]))
if __name__ == "__main__":
main()