File size: 7,928 Bytes
ac9d706 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 | """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)
|