File size: 5,768 Bytes
9d6c005
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""Small physical shadow-ray oracle with an exact finite-light reference.

Lambertian point receivers on z=0, three opaque sphere occluders, and a 6x6
array of point lights at z=2.2. This is direct illumination only. The scene
does not contain indirect transport, specular receivers, or camera antialiasing.
"""
from dataclasses import dataclass
import numpy as np


@dataclass
class Scene:
    seed: int
    spheres: np.ndarray  # [3,4] = x,y,z,radius, renderer-owned geometry

    def __post_init__(self):
        self.spheres=np.asarray(self.spheres,dtype=float)
        if self.spheres.shape!=(3,4) or not np.isfinite(self.spheres).all() or (self.spheres[:,3]<=0).any():
            raise ValueError("Reference scene requires three finite spheres with positive radii")

    @classmethod
    def create(cls, seed):
        rng = np.random.default_rng(seed)
        xyz = rng.uniform([-0.72,-0.72,0.45], [0.72,0.72,0.95], size=(3,3))
        radius = rng.uniform(0.17,0.36, size=(3,1))
        spheres = np.concatenate([xyz, radius], axis=1)
        return cls(int(seed), spheres[np.argsort(spheres[:, 0])])

    def changed(self):
        moved = self.spheres.copy()
        moved[0, 0] += 0.65
        moved[1, 1] -= 0.6
        moved[2, 3] *= 1.25
        return Scene(self.seed, moved)

    def visibility(self, receivers, lights):
        """Exact segment/sphere visibility for matching broadcastable arrays."""
        p, l = np.broadcast_arrays(np.asarray(receivers, float), np.asarray(lights, float))
        if p.shape[-1]!=3 or not np.isfinite(p).all() or not np.isfinite(l).all():
            raise ValueError("Physical segment endpoints must be finite 3D points")
        d = l-p
        a = (d*d).sum(-1)
        if (a <= 0).any():
            raise ValueError("A physical shadow segment must have positive length")
        visible = np.ones(a.shape, bool)
        for sphere in self.spheres:
            oc = p-sphere[:3]
            b = (oc*d).sum(-1)
            c = (oc*oc).sum(-1)-sphere[3]**2
            discriminant = b*b-a*c
            root = np.sqrt(np.maximum(discriminant,0))
            near, far = (-b-root)/a, (-b+root)/a
            intersects = (discriminant >= 0) & (far > 1e-6) & (near < 1-1e-6)
            visible &= ~intersects
        return visible.astype(np.float64)

    def features(self, receivers, lights):
        p, l = np.broadcast_arrays(np.asarray(receivers, float), np.asarray(lights, float))
        geom = np.broadcast_to(self.spheres.ravel(), p.shape[:-1]+(12,))
        return np.concatenate([p[..., :2], l[..., :2], geom], axis=-1).astype(np.float32)


def receiver_grid(height=32, width=64):
    x = np.linspace(-1,1,width)
    y = np.linspace(-1,1,height)
    xx,yy = np.meshgrid(x,y)
    return np.stack([xx,yy,np.zeros_like(xx)],axis=-1).reshape(-1,3)


def light_grid(side=6):
    x = np.linspace(-0.85,0.85,side)
    xx,yy = np.meshgrid(x,x)
    return np.stack([xx,yy,np.full_like(xx,2.2)],axis=-1).reshape(-1,3)


def albedo(points, phase=0):
    x,y = points[:,0],points[:,1]
    checker = ((np.floor((x+1)*10)+np.floor((y+1)*10))%2)
    r = 0.22+0.42*checker
    g = 0.25+0.25*(0.5+0.5*np.sin(11*x+phase))
    b = 0.2+0.38*(0.5+0.5*np.sin(13*y-0.3+phase))
    return np.stack([r,g,b],axis=-1)


def lighting(lights, changed=False):
    x,y=lights[:,0],lights[:,1]
    if changed:
        rgb=np.stack([0.3+1.7*(x>0),0.4+0.3*np.cos(3*y)**2,0.5+1.2*(x<0)],axis=-1)
    else:
        rgb=np.stack([0.8+0.4*np.cos(2*x)**2,0.9+0.2*np.sin(y)**2,0.7+0.2*np.cos(3*y)**2],axis=-1)
    return 11.0*rgb/len(lights)  # each emitter's radiant intensity


def unoccluded(points, lights, changed_lighting=False, material_phase=0):
    d=lights[None,:,:]-points[:,None,:]
    distance=np.linalg.norm(d,axis=-1)
    cosine=np.maximum(d[...,2]/distance,0)
    geometry=cosine/(np.pi*distance**2)
    return geometry[...,None]*albedo(points,material_phase)[:,None,:]*lighting(lights,changed_lighting)[None,:,:]


def physical_table(scene, points, lights, bound):
    """Privileged offline reference; do not pass this table to online policy."""
    vis=scene.visibility(points[:,None,:],lights[None,:,:])
    return bound*vis[...,None]


class VisibilityPrior:
    """Portable NumPy evaluation of the trained 16-48-48-1 MLP."""
    def __init__(self,path):
        with np.load(path,allow_pickle=False) as data:
            self.arrays={k:np.array(data[k]) for k in data.files}
        required={"mean","scale","w0","b0","w1","b1","w2","b2"}
        if set(self.arrays)!=required or not all(np.isfinite(v).all() for v in self.arrays.values()):
            raise ValueError("Malformed prior weights")
        shapes={"mean":(16,),"scale":(16,),"w0":(48,16),"b0":(48,),"w1":(48,48),"b1":(48,),"w2":(1,48),"b2":(1,)}
        if any(self.arrays[k].shape!=s for k,s in shapes.items()) or (self.arrays["scale"]<=0).any():
            raise ValueError("Invalid prior dimensions/scaling")

    def __call__(self,features):
        features=np.asarray(features,np.float32)
        if features.ndim<1 or features.shape[-1]!=16 or not np.isfinite(features).all():
            raise ValueError("Visibility prior requires finite 16-dimensional features")
        original_shape=features.shape[:-1]
        x=np.asarray(features,np.float32).reshape(-1,16)
        a=self.arrays
        outputs=[]
        for start in range(0,len(x),8192):
            h=(x[start:start+8192]-a["mean"])/a["scale"]
            h=np.maximum(h@a["w0"].T+a["b0"],0)
            h=np.maximum(h@a["w1"].T+a["b1"],0)
            h=h@a["w2"].T+a["b2"]
            outputs.append(1/(1+np.exp(-np.clip(h,-40,40))))
        if not outputs:
            return np.empty(original_shape,dtype=np.float32)
        return np.concatenate(outputs).reshape(original_shape)