Import from Quazim0t0/neural-raytracing; repoint refs to NeuralVerified

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+neural_raytrace_validation.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,70 @@
+---
+license: mit
+tags:
+  - neural-rendering
+  - ray-tracing
+  - radiance-cache
+  - graphics
+library_name: pytorch
+---
+
+# Neural Ray Tracing — Radiance Cache
+
+The light-transport analog of the
+[neural physics engine](https://huggingface.co/NeuralVerified/neural-physics-engine)
+thesis: keep visibility and direct lighting **analytic**, learn **only the
+indirect transport** with a tiny tied MLP. One analytic ray + next-event
+estimation + a network lookup replaces the many-bounce random walk after
+the first hit.
+
+## How it was done
+
+Render decomposition: `L = emitted + direct(analytic) + indirect(learned)`.
+
+1. **Ground truth**: a small PyTorch path tracer
+   (`engine3d/raytrace.py`) renders high-spp references and splits each
+   pixel into emitted / direct / indirect components.
+2. **Training data**: first-hit surface points + normals paired with the
+   path-traced indirect radiance at those points.
+3. **The cache** (`engine3d/neural_rt.py`): one tiny MLP shared across the
+   whole scene (the tied-embedding structure used everywhere in this
+   project) maps (hit point, normal) → indirect RGB.
+4. **Composition**: at render time, trace one analytic primary ray, add
+   analytic direct lighting (NEE), and look up the cache for the rest.
+5. **Evaluation** (`experiments/w9_neural_radiance_cache.py`): PSNR on a
+   held-out camera view, compared against an *equal-cost* few-spp path
+   trace, plus the spp the baseline needs to match the neural render.
+
+`experiments/w11_world_lighting.py` applies the same recipe to bake a
+world ambient/GI field (`world_light.pt`) over the voxel world — this is
+the "neural GI" used live in the
+[Neural World demo Space](https://huggingface.co/spaces/NeuralVerified/neural-world),
+parsed in-browser by `pt_loader.js`.
+
+## Checkpoints
+
+| file | net | consumed by |
+|---|---|---|
+| `experiments/radiance_cache.pt` | indirect radiance cache MLP | `w9_neural_radiance_cache.py` renders |
+| `experiments/world_light.pt` | world ambient/GI field | Neural World demo (browser) |
+
+## Validation
+
+![neural raytrace validation](neural_raytrace_validation.png)
+![world lighting validation](world_lighting_validation.png)
+
+## Reproduce
+
+```
+python experiments/w9_neural_radiance_cache.py   # trains + renders comparison
+python experiments/w11_world_lighting.py         # bakes world_light.pt
+```
+
+
+---
+
+<!-- neuralverified-relocation-note -->
+> **Now hosted by [NeuralVerified](https://huggingface.co/NeuralVerified).**
+>
+> This repo was moved into the NeuralVerified organization to help organize my profile.
+> Originally published at [`Quazim0t0/neural-raytracing`](https://huggingface.co/Quazim0t0/neural-raytracing).

engine3d/neural_rt.py

@@ -0,0 +1,63 @@
+"""Neural radiance cache (W9 — neural ray tracing).
+
+The same thesis as the physics engine, applied to light transport: keep
+the exact parts analytic (ray/visibility, next-event direct lighting),
+learn only the expensive part (multi-bounce indirect transport).
+
+A path becomes:  L = emitted + direct(analytic NEE) + cache_θ(x, n)
+where cache_θ is one tiny MLP, tied across the whole scene, that maps a
+surface point + normal to its outgoing indirect radiance. This replaces
+the random walk after the first bounce with a single network lookup —
+Müller et al.'s "Real-time Neural Radiance Caching" idea, distilled to
+the project's local-projection skeleton.
+
+Positional (frequency) encoding is essential: a plain MLP on raw xyz
+cannot fit the high-frequency shading near contact shadows and the
+color-bleeding corners. The encoding is the analytic structure we keep;
+the MLP only learns amplitudes.
+"""
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+class FreqEncoding(nn.Module):
+    """NeRF-style sin/cos encoding: x -> [x, sin(2^k x), cos(2^k x)]_k."""
+
+    def __init__(self, n_freq=6):
+        super().__init__()
+        self.register_buffer("bands", 2.0 ** torch.arange(n_freq) * torch.pi)
+        self.out_mult = 1 + 2 * n_freq
+
+    def forward(self, x):
+        proj = x[..., None] * self.bands            # (...,D,F)
+        enc = torch.cat([torch.sin(proj), torch.cos(proj)], -1)
+        return torch.cat([x, enc.flatten(-2)], -1)
+
+
+class RadianceCache(nn.Module):
+    """Tied MLP: (position, normal) -> outgoing indirect radiance (RGB).
+
+    Shared across every surface point in the scene (the tied-embedding
+    thesis, now over spatial location instead of mesh element). Output is
+    non-negative (softplus) — radiance cannot be negative, a guard baked
+    into the architecture rather than clamped after the fact.
+    """
+
+    def __init__(self, n_freq=6, hidden=96):
+        super().__init__()
+        self.enc = FreqEncoding(n_freq)
+        din = self.enc.out_mult * 3 + 3          # encoded position + raw normal
+        self.net = nn.Sequential(
+            nn.Linear(din, hidden), nn.SiLU(),
+            nn.Linear(hidden, hidden), nn.SiLU(),
+            nn.Linear(hidden, hidden), nn.SiLU(),
+            nn.Linear(hidden, 3),
+        )
+
+    def forward(self, pos, nrm):
+        h = torch.cat([self.enc(pos), nrm], -1)
+        return torch.nn.functional.softplus(self.net(h))
+
+    def n_params(self):
+        return sum(p.numel() for p in self.parameters())

engine3d/raytrace.py

@@ -0,0 +1,200 @@
+"""Vectorized torch path tracer (the substrate for neural ray tracing).
+
+Same design rule as the physics engine: keep the exact parts analytic —
+ray intersections and next-event-estimated direct lighting — and reserve
+learning for the expensive part (indirect transport, see W9 experiment).
+
+Scene: a neon Cornell box. Diffuse surfaces, one area light on the
+ceiling, a sphere and a box inside. All ray math is batched over
+(N,3) tensors on the chosen device.
+"""
+import torch
+
+EPS = 1e-4
+
+# ---- scene definition (unit room, front face at z=+1 open) ----
+LIGHT = dict(y=0.999, half=0.42, Le=torch.tensor([11.0, 11.5, 13.0]))
+SPHERE = dict(c=torch.tensor([-0.42, -0.65, -0.22]), r=0.35,
+              alb=torch.tensor([0.85, 0.85, 0.88]))
+BOX = dict(mn=torch.tensor([0.12, -1.0, -0.10]),
+           mx=torch.tensor([0.72, -0.34, 0.52]),
+           alb=torch.tensor([0.30, 0.75, 0.95]))
+WALLS = [  # (axis, value, normal-sign, albedo)
+    (0, -1.0, +1, [0.25, 0.55, 0.95]),   # left  — neon blue
+    (0, +1.0, -1, [0.95, 0.30, 0.40]),   # right — crimson
+    (1, -1.0, +1, [0.70, 0.70, 0.72]),   # floor
+    (1, +1.0, -1, [0.70, 0.70, 0.72]),   # ceiling
+    (2, -1.0, +1, [0.55, 0.60, 0.75]),   # back
+]
+
+
+def _dev(dev):
+    out = {"Le": LIGHT["Le"].to(dev), "sc": SPHERE["c"].to(dev),
+           "sa": SPHERE["alb"].to(dev), "bmn": BOX["mn"].to(dev),
+           "bmx": BOX["mx"].to(dev), "ba": BOX["alb"].to(dev),
+           "wa": torch.tensor([w[3] for w in WALLS], device=dev)}
+    return out
+
+
+def intersect(o, d, S):
+    """Closest hit for rays (N,3),(N,3) -> t, normal, albedo, is_light."""
+    N = len(o)
+    dev = o.device
+    INF = torch.full((N,), 1e9, device=dev)
+    best_t = INF.clone()
+    n = torch.zeros(N, 3, device=dev)
+    alb = torch.zeros(N, 3, device=dev)
+    is_light = torch.zeros(N, dtype=torch.bool, device=dev)
+
+    arange = torch.arange(N, device=dev)
+    for wi, (ax, val, sgn, _) in enumerate(WALLS):
+        denom = d[:, ax]
+        t = (val - o[:, ax]) / torch.where(denom.abs() < 1e-9,
+                                           torch.full_like(denom, 1e-9), denom)
+        p = o + t[:, None] * d
+        oth = [a for a in range(3) if a != ax]
+        ok = (t > EPS) & (t < best_t) \
+            & (p[:, oth[0]].abs() <= 1.0) & (p[:, oth[1]].abs() <= 1.0) \
+            & (p[:, 2] <= 1.0)
+        best_t = torch.where(ok, t, best_t)
+        nw = torch.zeros_like(n); nw[:, ax] = float(sgn)
+        n = torch.where(ok[:, None], nw, n)
+        alb = torch.where(ok[:, None], S["wa"][wi], alb)
+    # sphere
+    oc = o - S["sc"]
+    b = (oc * d).sum(1)
+    c = (oc * oc).sum(1) - SPHERE["r"] ** 2
+    disc = b * b - c
+    sq = torch.sqrt(disc.clamp_min(0))
+    t1 = -b - sq
+    t2 = -b + sq
+    ts = torch.where(t1 > EPS, t1, t2)
+    ok = (disc > 0) & (ts > EPS) & (ts < best_t)
+    best_t = torch.where(ok, ts, best_t)
+    ps = o + ts[:, None] * d
+    n = torch.where(ok[:, None], (ps - S["sc"]) / SPHERE["r"], n)
+    alb = torch.where(ok[:, None], S["sa"], alb)
+    # box (slabs)
+    inv = 1.0 / torch.where(d.abs() < 1e-9, torch.full_like(d, 1e-9), d)
+    t0s = (S["bmn"] - o) * inv
+    t1s = (S["bmx"] - o) * inv
+    tsm = torch.minimum(t0s, t1s).max(1).values
+    tbg = torch.maximum(t0s, t1s).min(1).values
+    ok = (tsm < tbg) & (tsm > EPS) & (tsm < best_t)
+    best_t = torch.where(ok, tsm, best_t)
+    pb = o + tsm[:, None] * d
+    ctr = (S["bmn"] + S["bmx"]) / 2
+    half = (S["bmx"] - S["bmn"]) / 2
+    rel = (pb - ctr) / half
+    axb = rel.abs().argmax(1)
+    nb = torch.zeros_like(pb)
+    nb[arange, axb] = torch.sign(rel[arange, axb])
+    n = torch.where(ok[:, None], nb, n)
+    alb = torch.where(ok[:, None], S["ba"], alb)
+    # light flag: recomputed cleanly from the final hit (ceiling patch)
+    p = o + best_t[:, None] * d
+    is_light = (best_t < 1e8) & (n[:, 1] == -1.0) & (p[:, 1] > 0.99) \
+        & (p[:, 0].abs() <= LIGHT["half"]) & (p[:, 2].abs() <= LIGHT["half"])
+    return best_t, n, alb, is_light
+
+
+def cosine_hemisphere(n, rng):
+    """Cosine-weighted directions about normals n (N,3)."""
+    N = len(n)
+    u1 = torch.rand(N, device=n.device, generator=rng)
+    u2 = torch.rand(N, device=n.device, generator=rng)
+    r = torch.sqrt(u1)
+    phi = 2 * torch.pi * u2
+    x = r * torch.cos(phi)
+    y = r * torch.sin(phi)
+    z = torch.sqrt((1 - u1).clamp_min(0))
+    a = torch.where(n[:, 0:1].abs() > 0.9,
+                    torch.tensor([0.0, 1.0, 0.0], device=n.device).expand_as(n),
+                    torch.tensor([1.0, 0.0, 0.0], device=n.device).expand_as(n))
+    t = torch.linalg.cross(a, n)
+    t = t / t.norm(dim=1, keepdim=True).clamp_min(1e-9)
+    b = torch.linalg.cross(n, t)
+    return x[:, None] * t + y[:, None] * b + z[:, None] * n
+
+
+def nee(p, n, alb, S, rng):
+    """Next-event estimation toward the ceiling light. Returns (N,3)."""
+    N = len(p)
+    dev = p.device
+    u = (torch.rand(N, 2, device=dev, generator=rng) * 2 - 1) * LIGHT["half"]
+    lp = torch.stack([u[:, 0], torch.full((N,), LIGHT["y"], device=dev),
+                      u[:, 1]], 1)
+    dl = lp - p
+    dist = dl.norm(dim=1).clamp_min(1e-6)
+    dl = dl / dist[:, None]
+    cos_s = (n * dl).sum(1).clamp_min(0)
+    # light normal is (0,-1,0); the emission cosine is between it and the
+    # direction light->surface (=-dl), i.e. (0,-1,0)·(-dl) = +dl_y
+    cos_l = dl[:, 1].clamp_min(0)
+    t, _, _, _ = intersect(p + EPS * n, dl, S)
+    vis = t > dist - 3e-3
+    area = (2 * LIGHT["half"]) ** 2
+    g = cos_s * cos_l / (dist ** 2)
+    return (alb / torch.pi) * S["Le"] * (g * vis * area)[:, None]
+
+
+def trace_split(o, d, S, rng, depth=5):
+    """Path trace with NEE; returns (emitted, direct, indirect) per ray.
+
+    emitted:  light seen directly by the given ray
+    direct:   NEE at the FIRST hit (analytic given visibility)
+    indirect: everything after the first bounce — the part the neural
+              cache learns
+    Also returns the first-hit geometry (t, n, albedo) for cache training.
+    """
+    N = len(o)
+    dev = o.device
+    emitted = torch.zeros(N, 3, device=dev)
+    direct = torch.zeros(N, 3, device=dev)
+    indirect = torch.zeros(N, 3, device=dev)
+    tput = torch.ones(N, 3, device=dev)
+    alive = torch.ones(N, dtype=torch.bool, device=dev)
+    first = {}
+    co, cd = o.clone(), d.clone()
+    for depth_i in range(depth):
+        t, n, alb, isl = intersect(co, cd, S)
+        hit = (t < 1e8) & alive
+        p = co + t[:, None] * cd
+        if depth_i == 0:
+            first = dict(t=t.clone(), n=n.clone(), alb=alb.clone(),
+                         p=p.clone(), hit=hit.clone() & ~isl)
+            emitted[hit & isl] = S["Le"]
+        alive = alive & hit & ~isl
+        # no .any() early-out: at these depths the surviving-ray count is
+        # data-dependent and the sync costs more than the wasted work
+        contrib = torch.zeros(N, 3, device=dev)
+        contrib[alive] = nee(p[alive], n[alive], alb[alive], S, rng)
+        contrib = contrib * tput
+        if depth_i == 0:
+            direct = direct + contrib
+        else:
+            indirect = indirect + contrib
+        # bounce
+        nd = torch.zeros_like(cd)
+        nd[alive] = cosine_hemisphere(n[alive], rng)
+        tput = tput * alb
+        co = p + EPS * n
+        cd = nd
+    return emitted, direct, indirect, first
+
+
+def camera_rays(res, jitter, dev, rng, cam=(0.0, 0.0, 3.05), fov=0.62):
+    ys, xs = torch.meshgrid(
+        torch.linspace(1, -1, res, device=dev),
+        torch.linspace(-1, 1, res, device=dev), indexing="ij")
+    if jitter:
+        xs = xs + (torch.rand(res, res, device=dev, generator=rng) - .5) * (2 / res)
+        ys = ys + (torch.rand(res, res, device=dev, generator=rng) - .5) * (2 / res)
+    d = torch.stack([xs * fov, ys * fov, -torch.ones_like(xs)], -1).reshape(-1, 3)
+    d = d / d.norm(dim=1, keepdim=True)
+    o = torch.tensor(cam, device=dev).expand_as(d).contiguous()
+    return o, d
+
+
+def scene_tensors(dev):
+    return _dev(dev)

experiments/radiance_cache.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6c963cdd817d845c9290ced6c1d188887e28b0b41f12707a9f5ea9b7f7e0245d
+size 95959

experiments/w11_world_lighting.py

@@ -0,0 +1,148 @@
+"""W11: neural lighting cache for the voxel world (ray-traced AO -> .pt).
+
+The ray-tracing tie-in for the Neural World demo. Sky visibility (ambient
+occlusion) is TRACED: from each ground point we shoot a hemisphere of
+rays and count how many escape past the tall blocks (house, trees) to the
+sky. That expensive per-point integral is distilled into a tiny MLP that
+the browser evaluates per tile in real time — the same "trace once, learn
+the field, look it up" idea as the W9 radiance cache, specialized to this
+scene's geometry.
+
+Layout MUST match the demo's hand-placed world (index.html): house 3x3 at
+rows/cols 2..4, ten trees at fixed spots — everything else is flat.
+"""
+import sys, os, time
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+torch.manual_seed(0)
+HERE = os.path.dirname(os.path.abspath(__file__))
+ROOT = os.path.dirname(HERE)
+DEV = "cuda" if torch.cuda.is_available() else "cpu"
+N = 32
+NFREQ = 6
+
+# tall occluders (cx, cz, half, top) — must match the demo world layout.
+# house: 6x6 perimeter centred at (14.5, 14.5); trees: the demo's fixed set.
+# x=col, z=row (three.js places tile (row r, col c) at x=c, z=r).
+TREES = [[3, 4], [5, 20], [7, 27], [2, 10], [4, 28], [9, 3], [6, 24], [10, 29],
+         [20, 4], [22, 27], [25, 10], [27, 22], [29, 6], [24, 29], [19, 28],
+         [28, 15], [3, 16], [8, 8], [23, 3], [26, 26], [21, 20], [18, 3],
+         [29, 29], [2, 24]]
+HR0, HC0, HS = 12, 12, 6
+BOXES = [(HC0 + HS / 2, HR0 + HS / 2, HS / 2, 2.2)]       # house block (x,z,half,top)
+for (r, c) in TREES:
+    if not (HR0 - 2 <= r < HR0 + HS + 2 and HC0 - 2 <= c < HC0 + HS + 2):
+        BOXES.append((c + 0.5, r + 0.5, 0.6, 2.4))       # tree (x=col, z=row)
+BOX = torch.tensor(BOXES, device=DEV)                    # (B,4): cx,cz,half,top
+
+
+def sky_visibility(pts, n_rays=64):
+    """pts (P,2) ground xz -> AO in [0,1] via hemisphere ray marching."""
+    P = len(pts)
+    g = torch.Generator(device=DEV); g.manual_seed(1)
+    # cosine-ish hemisphere directions (upper)
+    u1 = torch.rand(n_rays, device=DEV, generator=g)
+    u2 = torch.rand(n_rays, device=DEV, generator=g)
+    r = torch.sqrt(u1); phi = 2 * np.pi * u2
+    dirs = torch.stack([r * torch.cos(phi),
+                        torch.sqrt((1 - u1).clamp_min(0)),   # up = +y
+                        r * torch.sin(phi)], 1)              # (n_rays,3)
+    o = torch.stack([pts[:, 0], torch.full((P,), 0.55, device=DEV),
+                     pts[:, 1]], 1)                          # (P,3)
+    vis = torch.ones(P, n_rays, device=DEV)
+    cx, cz, half, top = BOX[:, 0], BOX[:, 1], BOX[:, 2], BOX[:, 3]
+    # march: does ray (o,d) hit any box slab before escaping upward?
+    for b in range(len(BOX)):
+        # slab intersection in xz, check the entry y is below the box top
+        dx = dirs[:, 0][None, :]; dz = dirs[:, 2][None, :]; dy = dirs[:, 1][None, :]
+        ox = o[:, 0][:, None]; oz = o[:, 2][:, None]; oy = o[:, 1][:, None]
+        inv_x = 1.0 / torch.where(dx.abs() < 1e-6, torch.full_like(dx, 1e-6), dx)
+        inv_z = 1.0 / torch.where(dz.abs() < 1e-6, torch.full_like(dz, 1e-6), dz)
+        tx1 = (cx[b] - half[b] - ox) * inv_x; tx2 = (cx[b] + half[b] - ox) * inv_x
+        tz1 = (cz[b] - half[b] - oz) * inv_z; tz2 = (cz[b] + half[b] - oz) * inv_z
+        tmin = torch.maximum(torch.minimum(tx1, tx2), torch.minimum(tz1, tz2))
+        tmax = torch.minimum(torch.maximum(tx1, tx2), torch.maximum(tz1, tz2))
+        hit_xz = (tmax > tmin) & (tmax > 0)
+        t_enter = tmin.clamp_min(0)
+        y_at = oy + dy * t_enter
+        blocked = hit_xz & (y_at < top[b]) & (t_enter > 1e-3)
+        vis = torch.where(blocked, torch.zeros_like(vis), vis)
+    return vis.mean(1)                                        # (P,) AO
+
+
+def freq_encode(xz):
+    bands = (2.0 ** torch.arange(NFREQ, device=xz.device)) * np.pi
+    proj = xz[..., None] * bands
+    return torch.cat([xz, torch.sin(proj).flatten(-2),
+                      torch.cos(proj).flatten(-2)], -1)
+
+
+class LightNet(nn.Module):
+    def __init__(self, hidden=64):
+        super().__init__()
+        din = 2 + 2 * 2 * NFREQ
+        self.net = nn.Sequential(nn.Linear(din, hidden), nn.SiLU(),
+                                 nn.Linear(hidden, hidden), nn.SiLU(),
+                                 nn.Linear(hidden, 3))
+
+    def forward(self, xz):
+        return torch.nn.functional.softplus(self.net(freq_encode(xz)))
+
+
+# ---- ground-truth AO -> ambient RGB on a dense grid ----
+print(f"=== tracing sky visibility on {DEV} ===")
+t0 = time.time()
+G = 112
+gx, gz = torch.meshgrid(torch.linspace(0, N, G, device=DEV),
+                        torch.linspace(0, N, G, device=DEV), indexing="ij")
+pts = torch.stack([gx.ravel(), gz.ravel()], 1)
+ao = sky_visibility(pts, n_rays=96)
+SKY = torch.tensor([0.55, 0.70, 1.0], device=DEV)         # cool sky
+GROUND = torch.tensor([0.30, 0.32, 0.28], device=DEV)     # warm bounce floor
+amb = (0.30 + 0.70 * ao)[:, None] * SKY + (1 - ao)[:, None] * GROUND * 0.4
+xz_norm = pts / N * 2 - 1
+print(f"  {G*G} points, mean AO {ao.mean():.3f} "
+      f"(min {ao.min():.3f}) ({time.time()-t0:.1f}s)")
+
+# ---- fit the cache ----
+net = LightNet().to(DEV)
+print(f"=== training light cache ({sum(p.numel() for p in net.parameters())} params) ===")
+opt = torch.optim.Adam(net.parameters(), lr=3e-3)
+sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=2500)
+for it in range(2500):
+    idx = torch.randint(0, len(xz_norm), (8192,), device=DEV)
+    pred = net(xz_norm[idx])
+    loss = ((pred - amb[idx]) ** 2).mean()
+    opt.zero_grad(); loss.backward(); opt.step(); sched.step()
+    if it % 500 == 0 or it == 2499:
+        print(f"  it {it:4d}  MSE {loss.item():.6f}")
+with torch.no_grad():
+    err = (net(xz_norm) - amb).abs().mean().item()
+print(f"  final mean abs error {err:.4f}")
+
+# ---- figure: traced vs learned ambient ----
+try:
+    import matplotlib; matplotlib.use("Agg"); import matplotlib.pyplot as plt
+    with torch.no_grad():
+        pr = net(xz_norm).reshape(G, G, 3).clamp(0, 1).cpu().numpy()
+    gt = amb.reshape(G, G, 3).clamp(0, 1).cpu().numpy()
+    fig, ax = plt.subplots(1, 3, figsize=(14, 4.6))
+    ax[0].imshow(ao.reshape(G, G).cpu().numpy(), cmap="bone"); ax[0].set_title("traced sky visibility (AO)")
+    ax[1].imshow(gt); ax[1].set_title("target ambient (AO -> RGB)")
+    ax[2].imshow(pr); ax[2].set_title(f"neural cache (world_light.pt), err {err:.3f}")
+    for a in ax: a.set_xticks([]); a.set_yticks([])
+    fig.suptitle("Neural world lighting — traced ambient occlusion distilled to a .pt")
+    fig.tight_layout()
+    out = os.path.join(ROOT, "world_lighting_validation.png")
+    fig.savefig(out, dpi=110); print(f"wrote {out}")
+except Exception as e:
+    print("plot skipped:", e)
+
+torch.save({"state_dict": net.state_dict(), "arch": {"nfreq": NFREQ, "hidden": 64},
+            "scene": "neural_world_32"},
+           os.path.join(HERE, "world_light.pt"))
+print("saved experiments/world_light.pt")

experiments/w9_neural_radiance_cache.py

@@ -0,0 +1,217 @@
+"""W9: neural radiance cache — neural ray tracing on the project skeleton.
+
+Claim under test (the light-transport analog of the physics thesis): if
+you keep visibility and direct lighting analytic and learn ONLY the
+indirect transport with a tiny tied MLP, you match a many-bounce path
+tracer at a fraction of the cost — one analytic ray + NEE + a network
+lookup replaces the random walk after the first bounce.
+
+Render decomposition:  L = emitted + direct(analytic) + indirect
+  - reference:  indirect from high-spp path tracing (ground truth)
+  - neural:     indirect from cache_theta(first-hit point, normal)
+  - baseline:   an equal-COST path trace (few spp) for a fair comparison
+
+Success criteria:
+  A. cache accuracy on HELD-OUT surface points (never-seen view) — PSNR
+     of predicted vs true indirect radiance
+  B. full-image PSNR: neural-cached vs reference, and it beats the
+     equal-cost path-traced baseline
+  C. speedup at matched quality (spp the baseline needs to reach the
+     neural render's PSNR)
+"""
+import sys, os, time
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+import numpy as np
+import torch
+from engine3d.raytrace import (camera_rays, trace_split, scene_tensors,
+                               intersect)
+from engine3d.neural_rt import RadianceCache
+
+torch.manual_seed(0)
+HERE = os.path.dirname(os.path.abspath(__file__))
+ROOT = os.path.dirname(HERE)
+DEV = "cuda" if torch.cuda.is_available() else "cpu"
+S = scene_tensors(DEV)
+
+
+def rng_of(seed):
+    g = torch.Generator(device=DEV); g.manual_seed(seed); return g
+
+
+def render(res, spp, seed, cam=(0.0, 0.0, 3.05), depth=5):
+    """Full reference render (emitted + direct + indirect), returns (res,res,3)."""
+    g = rng_of(seed)
+    acc = torch.zeros(res * res, 3, device=DEV)
+    for _ in range(spp):
+        o, d = camera_rays(res, True, DEV, g, cam=cam)
+        em, di, ind, _ = trace_split(o, d, S, g, depth=depth)
+        acc += em + di + ind
+    return (acc / spp).reshape(res, res, 3)
+
+
+def render_components(res, spp, seed, cam=(0.0, 0.0, 3.05)):
+    """Per-pixel emitted+direct (low variance) and first-hit geometry, plus
+    a high-spp indirect target. Used to compose the neural image."""
+    g = rng_of(seed)
+    ed = torch.zeros(res * res, 3, device=DEV)      # emitted + direct
+    ind = torch.zeros(res * res, 3, device=DEV)
+    o, d = camera_rays(res, False, DEV, g, cam=cam)  # no jitter: fixed hits
+    for _ in range(spp):
+        em, di, ii, first = trace_split(o, d, S, g, depth=5)
+        ed += em + di
+        ind += ii
+    ed /= spp; ind /= spp
+    o2, d2 = camera_rays(res, False, DEV, rng_of(seed), cam=cam)
+    _, n, alb, isl = intersect(o2, d2, S)
+    t, _, _, _ = intersect(o2, d2, S)
+    p = o2 + t[:, None] * d2
+    hit = (t < 1e8) & ~isl
+    return ed, ind, p, n, hit
+
+
+def psnr(a, b, mask=None):
+    if mask is not None:
+        a, b = a[mask], b[mask]
+    mse = ((a.clamp(0, 4) - b.clamp(0, 4)) ** 2).mean().item()
+    return 10 * np.log10(4.0 ** 2 / max(mse, 1e-12))
+
+
+# ---------- training data: surface points with GT indirect ----------
+print(f"=== generating radiance-cache training data on {DEV} ===")
+t0 = time.time()
+PTS, NRM, TGT = [], [], []
+CAM_TRAIN = [(0.0, 0.0, 3.05), (0.7, 0.2, 2.9), (-0.7, 0.25, 2.9),
+             (0.0, 0.6, 2.8), (0.4, -0.3, 3.0)]
+for ci, cam in enumerate(CAM_TRAIN):
+    g = rng_of(100 + ci)
+    o, d = camera_rays(80, True, DEV, g, cam=cam)
+    t, n, alb, isl = intersect(o, d, S)
+    p = o + t[:, None] * d
+    hit = (t < 1e8) & ~isl
+    p, n = p[hit], n[hit]
+    # GT indirect at these points: average many one-bounce-onward paths
+    ind = torch.zeros(len(p), 3, device=DEV)
+    SPP = 256
+    for _ in range(SPP):
+        from engine3d.raytrace import cosine_hemisphere, nee, EPS
+        nd = cosine_hemisphere(n, g)
+        # radiance arriving from the bounce dir, path-traced (depth 4),
+        # times the cosine-weighted albedo throughput (albedo/pi * pi = albedo)
+        _, ndi, nind, _ = trace_split(p + EPS * n, nd, S, g, depth=4)
+        ind += (ndi + nind)
+    ind /= SPP
+    alb_h = alb[hit]
+    ind = ind * alb_h                     # outgoing = albedo * incident indirect
+    PTS.append(p); NRM.append(n); TGT.append(ind)
+P = torch.cat(PTS); Nrm = torch.cat(NRM); Y = torch.cat(TGT)
+print(f"  {len(P)} surface samples from {len(CAM_TRAIN)} views "
+      f"({time.time()-t0:.1f}s)")
+
+# ---------- train the cache ----------
+cache = RadianceCache().to(DEV)
+print(f"=== training radiance cache ({cache.n_params()} params) ===")
+opt = torch.optim.Adam(cache.parameters(), lr=3e-3)
+sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=1500)
+# tone-mapped (log) loss: indirect radiance spans orders of magnitude and
+# the eye is roughly logarithmic — regress log(1+L), like real NRC
+logY = torch.log1p(Y)
+t0 = time.time()
+for it in range(1500):
+    idx = torch.randint(0, len(P), (16384,), device=DEV)
+    pred = cache(P[idx], Nrm[idx])
+    loss = ((torch.log1p(pred) - logY[idx]) ** 2).mean()
+    opt.zero_grad(); loss.backward(); opt.step(); sched.step()
+    if it % 300 == 0 or it == 1499:
+        print(f"  it {it:4d}  log-MSE {loss.item():.5f}")
+print(f"  ({time.time()-t0:.1f}s)")
+
+# ---------- criterion A: held-out points (novel view) ----------
+g = rng_of(999)
+cam_test = (0.35, -0.15, 2.95)
+o, d = camera_rays(80, True, DEV, g, cam=cam_test)
+t, n, alb, isl = intersect(o, d, S)
+p = o + t[:, None] * d
+hit = (t < 1e8) & ~isl
+p, n, alb = p[hit], n[hit], alb[hit]
+ind = torch.zeros(len(p), 3, device=DEV)
+from engine3d.raytrace import cosine_hemisphere, EPS
+for _ in range(256):
+    nd = cosine_hemisphere(n, g)
+    _, ndi, nind, _ = trace_split(p + EPS * n, nd, S, g, depth=4)
+    ind += (ndi + nind)
+ind = ind / 256 * alb
+with torch.no_grad():
+    pred = cache(p, n)
+a_psnr = psnr(pred, ind)
+print(f"\n[A] held-out indirect radiance PSNR: {a_psnr:.2f} dB "
+      f"(novel view, {len(p)} pts)")
+
+# ---------- criterion B & C: full-image renders ----------
+RES = 112
+REF_SPP = 256
+print(f"\n=== full renders at {RES}x{RES} ===")
+t0 = time.time(); REF = render(RES, REF_SPP, 7); ref_ms = (time.time()-t0)*1000
+print(f"  reference ({REF_SPP} spp): {ref_ms:.0f} ms")
+
+# neural: analytic emitted+direct at low spp + cache indirect
+t0 = time.time()
+ed, _, p, n, hit = render_components(RES, 4, 7)
+with torch.no_grad():
+    ind_pred = torch.zeros(RES * RES, 3, device=DEV)
+    ind_pred[hit] = cache(p[hit], n[hit])
+NEUR = (ed + ind_pred).reshape(RES, RES, 3)
+neur_ms = (time.time()-t0)*1000
+print(f"  neural (4 spp direct + cache): {neur_ms:.0f} ms")
+
+# baseline: equal-cost full path trace
+def spp_for_ms(target_ms):
+    t0 = time.time(); render(RES, 4, 7); one = (time.time()-t0)/4*1000
+    return max(1, round(target_ms / one)), one
+base_spp, per = spp_for_ms(neur_ms)
+BASE = render(RES, base_spp, 7)
+
+pn = psnr(NEUR, REF)
+pb = psnr(BASE, REF)
+print(f"\n[B] full-image PSNR vs reference:")
+print(f"     neural (cache)          : {pn:.2f} dB  ({neur_ms:.0f} ms)")
+print(f"     path trace (equal cost) : {pb:.2f} dB  ({base_spp} spp)")
+
+# criterion C: spp the baseline needs to match the neural PSNR
+match_spp = None
+for spp in [8, 16, 32, 64, 128]:
+    if psnr(render(RES, spp, 7), REF) >= pn:
+        match_spp = spp; break
+base_ms = (match_spp or 128) * per
+print(f"\n[C] path trace needs ~{match_spp or '>256'} spp "
+      f"(~{base_ms:.0f} ms) to reach the neural PSNR "
+      f"-> ~{base_ms/neur_ms:.1f}x speedup at matched quality")
+
+# ---------- figure ----------
+try:
+    import matplotlib; matplotlib.use("Agg"); import matplotlib.pyplot as plt
+    def tm(x): return np.clip((x.clamp(0, 4) ** (1/2.2)).cpu().numpy(), 0, 1)
+    err = (NEUR - REF).abs().mean(-1).cpu().numpy()
+    errb = (BASE - REF).abs().mean(-1).cpu().numpy()
+    fig, ax = plt.subplots(2, 3, figsize=(12, 8))
+    for a in ax.ravel(): a.set_xticks([]); a.set_yticks([])
+    ax[0,0].imshow(tm(REF)); ax[0,0].set_title(f"reference ({REF_SPP} spp)")
+    ax[0,1].imshow(tm(NEUR)); ax[0,1].set_title(f"neural cache — {pn:.1f} dB, {neur_ms:.0f} ms")
+    ax[0,2].imshow(tm(BASE)); ax[0,2].set_title(f"equal-cost path trace ({base_spp} spp) — {pb:.1f} dB")
+    ax[1,0].imshow(tm(ind_pred.reshape(RES,RES,3))); ax[1,0].set_title("learned indirect (cache only)")
+    m=max(err.max(),errb.max())
+    ax[1,1].imshow(err,cmap="inferno",vmax=m); ax[1,1].set_title("neural error")
+    ax[1,2].imshow(errb,cmap="inferno",vmax=m); ax[1,2].set_title("path-trace error")
+    fig.suptitle(f"Neural radiance cache — learn indirect, keep visibility+direct analytic "
+                 f"(held-out cache PSNR {a_psnr:.1f} dB)", fontsize=12)
+    fig.tight_layout()
+    out=os.path.join(ROOT,"neural_raytrace_validation.png")
+    fig.savefig(out,dpi=110); print(f"\nwrote {out}")
+except Exception as e:
+    print("plot skipped:", e)
+
+torch.save({"state_dict": cache.state_dict(),
+            "arch": {"n_freq": 6, "hidden": 96},
+            "scene": "neon_cornell"},
+           os.path.join(HERE, "radiance_cache.pt"))
+print("saved experiments/radiance_cache.pt")

experiments/world_light.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e127cfc99059eae827da9a78cf373e671a1d18253ac95a724c5d5a2000afa194
+size 27184