File size: 2,493 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 | """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())
|