public/math/trace.mjs
与本次数学页面一起发布的源码快照。SHA-256: c1369ede9a2303a9d1f39f73bfd50b1bb9b7e14ad689308f8388033f77f18a2c
1// Executable single-pixel lesson: STG-Lite fields -> time -> ellipse -> pixel ->
2// ordered over -> MSE -> analytic gradient -> SGD -> the same 32-float PLY.
3// Independent CPU arithmetic, not a WebGPU renderer or the official STG trainer.
4export const properties = (
5 'x y z trbf_center trbf_scale nx ny nz ' +
6 'motion_0 motion_1 motion_2 motion_3 motion_4 motion_5 motion_6 motion_7 motion_8 ' +
7 'f_dc_0 f_dc_1 f_dc_2 opacity scale_0 scale_1 scale_2 ' +
8 'rot_0 rot_1 rot_2 rot_3 omega_0 omega_1 omega_2 omega_3'
9).split(' ');
10export const camera = { focal: 100, width: 128, height: 96, cx: 64, cy: 48 };
11export const background = [0.04, 0.05, 0.08];
12export const target = [0.45, 0.18, 0.22];
13export const defaults = { time: 0.5, pixel: [74.5, 68.5], rate: 0.5 };
14export const dot = (a, b) => a.reduce((sum, v, i) => sum + v * b[i], 0);
15export const transpose = a => a[0].map((_, j) => a.map(row => row[j]));
16export const multiply = (a, b) => a.map(row => transpose(b).map(col => dot(row, col)));
17const outer = a => a.map(v => a.map(w => v * w));
18const matrixVector = (a, v) => a.map(row => dot(row, v));
19const sigmoid = v => 1 / (1 + Math.exp(-v));
20const plus = (a, b) => a.map((v, i) => v + b[i]);
21const scaled = (a, s) => a.map(v => v * s);
22
23export function initialRecords() {
24 const a = Array(32).fill(0), b = Array(32).fill(0);
25 for (const row of [a, b]) {
26 row[3] = 0.5; row[4] = Math.log(0.55);
27 row.splice(21, 3, ...[0.20, 0.12, 0.10].map(Math.log));
28 row[24] = 1;
29 }
30 a.splice(0, 3, 1, 2, 10);
31 a.splice(8, 9, 2, 0, 0, 0, 4, 0, 0, 0, 8);
32 a.splice(17, 3, 0.9, 0.2, 0.1); a[20] = Math.log(4);
33 // Thirty degrees around Z, with a small four-dimensional coefficient velocity.
34 a.splice(24, 4, Math.cos(Math.PI / 12), 0, 0, Math.sin(Math.PI / 12));
35 a[31] = 0.2;
36 b.splice(0, 3, 1.12, 2.04, 11);
37 b.splice(8, 3, -0.2, 0.1, 0);
38 b.splice(17, 3, 0.1, 0.3, 0.9); b[20] = 0;
39 return [a, b];
40}
41
42export function evaluateRecord(row, time, pixel) {
43 const dt = time - row[3], rho = Math.max(Math.exp(row[4]), 1e-6);
44 const temporal = Math.exp(-((dt / rho) ** 2));
45 const mean = row.slice(0, 3).map((v, i) =>
46 v + row[8 + i] * dt + row[11 + i] * dt ** 2 + row[14 + i] * dt ** 3);
47 const u = row.slice(24, 28).map((v, i) => v + dt * row[28 + i]);
48 const norm = Math.hypot(...u);
49 if (!(norm > 0)) throw new Error('Undefined quaternion');
50 const q = u.map(v => v / norm), [w, x, y, z] = q;
51 const scales = row.slice(21, 24).map(Math.exp);
52 // Columns of the active, right-handed rotation, exactly as splat.wgsl axes.
53 const axes = [
54 [1 - 2 * (y*y + z*z), 2 * (x*y + w*z), 2 * (x*z - w*y)],
55 [2 * (x*y - w*z), 1 - 2 * (x*x + z*z), 2 * (y*z + w*x)],
56 [2 * (x*z + w*y), 2 * (y*z - w*x), 1 - 2 * (x*x + y*y)],
57 ].map((axis, i) => scaled(axis, scales[i]));
58 const sigma = axes.map(outer).reduce((a, b) => a.map((row, i) => plus(row, b[i])));
59 const [mx, my, mz] = mean, f = camera.focal;
60 if (!(mz > 0)) throw new Error('Teaching camera requires positive depth');
61 const uv = [f * mx / mz + camera.cx, f * my / mz + camera.cy];
62 const J = [[f / mz, 0, -f * mx / mz ** 2], [0, f / mz, -f * my / mz ** 2]];
63 const projectedAxes = axes.map(axis => matrixVector(J, axis));
64 const covariance = multiply(multiply(J, sigma), transpose(J))
65 .map((row, i) => row.map((v, j) => v + (i === j ? 0.3 : 0)));
66 const [[a, b], [, c]] = covariance, determinant = a*c - b*b;
67 const inverse = [[c / determinant, -b / determinant], [-b / determinant, a / determinant]];
68 const delta = pixel.map((v, i) => v - uv[i]);
69 const qd = matrixVector(inverse, delta);
70 const radiusSquared = dot(delta, qd), spatial = Math.exp(-0.5 * radiusSquared);
71 const baseOpacity = sigmoid(row[20]), rawAlpha = baseOpacity * temporal * spatial;
72 const discriminant = Math.sqrt((a-c)**2 + 4*b*b);
73 const eigenvalues = [(a+c+discriminant)/2, (a+c-discriminant)/2];
74 let major = [1, 0];
75 if (Math.abs(b) > 1e-5) {
76 major = [b, Math.max(0.1, eigenvalues[0])-a];
77 major = scaled(major, 1 / Math.hypot(...major));
78 } else if (c > a) major = [0, 1];
79 const minor = [-major[1], major[0]];
80 const deviations = eigenvalues.map(v => Math.min(Math.sqrt(Math.max(0.1, v)), 192));
81 // Teaching coordinates are downwards; simultaneous Y reflection leaves r² unchanged.
82 const local = [dot(delta, major) / deviations[0], dot(delta, minor) / deviations[1]];
83 const covered = local.every(v => Math.abs(v) <= 3);
84 const capped = Math.min(0.99, rawAlpha);
85 const alpha = covered && capped >= 1/255 ? capped : 0;
86 return { dt, rho, temporal, mean, q, scales, axes, sigma, J, projectedAxes, uv,
87 covariance, determinant, inverse, delta, qd, radiusSquared, spatial,
88 baseOpacity, rawAlpha, eigenvalues, local, covered, alpha,
89 branch: !covered ? 'quad 外：不产生片元' : capped < 1/255 ? 'alpha 低于阈值：丢弃' :
90 rawAlpha >= 0.99 ? 'alpha 上限生效' : '光滑区间',
91 color: row.slice(17, 20).map(v => Math.max(0, Math.min(1, v))) };
92}
93
94export function forward(records, time = defaults.time, pixel = defaults.pixel, options = {}) {
95 const expected = options.target ?? target;
96 const splats = records.map((row, id) => ({ id, ...evaluateRecord(row, time, pixel) }))
97 .sort((a, b) => a.mean[2] - b.mean[2]);
98 let remaining = 1, color = [0, 0, 0];
99 for (const splat of splats) {
100 splat.transmittance = remaining;
101 splat.contribution = scaled(splat.color, remaining * splat.alpha);
102 color = plus(color, splat.contribution);
103 remaining *= 1 - splat.alpha;
104 }
105 color = plus(color, scaled(background, remaining));
106 const error = color.map((v, i) => v - expected[i]);
107 return { splats, remaining, color, target: [...expected], loss: dot(error, error) / 3 };
108}
109
110export function gradient(records, parameter, time = defaults.time, pixel = defaults.pixel, options = {}) {
111 // Differentiate record A, while holding the visibility/sort branches fixed.
112 // Supported fields: base mean x/y/z or opacity logit. Scales/rotation remain fixed.
113 if (![0, 1, 2, 20].includes(parameter)) throw new Error('Unsupported teaching parameter');
114 const result = forward(records, time, pixel, options);
115 const index = result.splats.findIndex(s => s.id === 0), s = result.splats[index];
116 let tail = [...background];
117 for (let j = result.splats.length - 1; j > index; j--) {
118 const layer = result.splats[j];
119 tail = plus(scaled(layer.color, layer.alpha), scaled(tail, 1-layer.alpha));
120 }
121 const colorSlope = scaled(s.color.map((v, i) => v - tail[i]), s.transmittance);
122 const lossColorSlope = result.color.map((v, i) => 2 * (v-result.target[i]) / 3);
123 const lossAlphaSlope = dot(lossColorSlope, colorSlope);
124 let meanPath = 0, covariancePath = 0, opacityPath = 0;
125 if (s.branch === '光滑区间') {
126 if (parameter === 20) opacityPath = s.alpha * (1-s.baseOpacity);
127 else {
128 const [x, y, z] = s.mean, f = camera.focal;
129 const dJ = parameter === 0 ? [[0, 0, -f/z**2], [0, 0, 0]] :
130 parameter === 1 ? [[0, 0, 0], [0, 0, -f/z**2]] :
131 [[-f/z**2, 0, 2*f*x/z**3], [0, -f/z**2, 2*f*y/z**3]];
132 const left = multiply(multiply(dJ, s.sigma), transpose(s.J));
133 const right = multiply(multiply(s.J, s.sigma), transpose(dJ));
134 const dC = left.map((row, i) => plus(row, right[i]));
135 meanPath = s.alpha * dot(s.qd, s.J.map(row => row[parameter]));
136 covariancePath = s.alpha / 2 * dot(s.qd, matrixVector(dC, s.qd));
137 }
138 }
139 const alphaSlope = meanPath + covariancePath + opacityPath;
140 const analytic = lossAlphaSlope * alphaSlope;
141 const finite = (options.finite === false ? [] : [1e-4, 1e-5, 1e-6]).map(h => {
142 const plusRecords = records.map(row => [...row]), minusRecords = records.map(row => [...row]);
143 plusRecords[0][parameter] += h; minusRecords[0][parameter] -= h;
144 const value = (forward(plusRecords, time, pixel, options).loss - forward(minusRecords, time, pixel, options).loss) / (2*h);
145 return { h, value, error: Math.abs(value-analytic) };
146 });
147 return { result, tail, colorSlope, lossColorSlope, lossAlphaSlope, meanPath, covariancePath,
148 opacityPath, alphaSlope, analytic, finite };
149}
150
151export function encodePly(records) {
152 const header = new TextEncoder().encode(
153 'ply\nformat binary_little_endian 1.0\nelement vertex ' + records.length + '\n' +
154 properties.map(name => 'property float ' + name + '\n').join('') + 'end_header\n');
155 const output = new Uint8Array(header.length + records.length * 32 * 4);
156 output.set(header);
157 const view = new DataView(output.buffer);
158 records.forEach((row, i) => row.forEach((value, j) =>
159 view.setFloat32(header.length + (i * 32 + j) * 4, value, true)));
160 return output;
161}
162