1#define_import_path bevy_gaussian_splatting::spherindrical_harmonics
2
3#import bevy_gaussian_splatting::bindings::gaussian_uniforms
4#import bevy_gaussian_splatting::spherical_harmonics::{
5    shc,
6    spherical_harmonics_lookup,
7}
8
9
10const PI = radians(180.0);
11
12
13fn spherindrical_harmonics_lookup(
14    ray_direction: vec3<f32>,
15    dir_t: f32,
16    sh: array<f32, #{SH_4D_COEFF_COUNT}>,
17) -> vec3<f32> {
18    let rds = ray_direction * ray_direction;
19
20    var color = vec3<f32>(0.5);
21
22    // TODO: reinterpret sh as vec3<f32>
23    color += shc[ 0] * vec3<f32>(sh[0], sh[1], sh[2]);
24
25#if SH_DEGREE > 0
26    let x = ray_direction.x;
27    let y = ray_direction.y;
28    let z = ray_direction.z;
29
30    let l1m1 = shc[1] * y;
31    let l1m0 = shc[2] * z;
32    let l1p1 = shc[3] * x;
33
34    color += l1m1 * vec3<f32>(sh[ 3], sh[ 4], sh[ 5]);
35    color += l1m0 * vec3<f32>(sh[ 6], sh[ 7], sh[ 8]);
36    color += l1p1 * vec3<f32>(sh[ 9], sh[10], sh[11]);
37#endif
38
39#if SH_DEGREE > 1
40    let xx = x * x;
41    let yy = y * y;
42    let zz = z * z;
43    let xy = x * y;
44    let xz = x * z;
45    let yz = y * z;
46
47    let l2m2 = shc[4] * xy;
48    let l2m1 = shc[5] * yz;
49    let l2m0 = shc[6] * (2.0 * zz - xx - yy);
50    let l2p1 = shc[7] * xz;
51    let l2p2 = shc[8] * (xx - yy);
52
53    color += l2m2 * vec3<f32>(sh[12], sh[13], sh[14]);
54    color += l2m1 * vec3<f32>(sh[15], sh[16], sh[17]);
55    color += l2m0 * vec3<f32>(sh[18], sh[19], sh[20]);
56    color += l2p1 * vec3<f32>(sh[21], sh[22], sh[23]);
57    color += l2p2 * vec3<f32>(sh[24], sh[25], sh[26]);
58#endif
59
60#if SH_DEGREE > 2
61    let l3m3 = shc[9] * y * (3.0 * xx - yy);
62    let l3m2 = shc[10] * z * xy;
63    let l3m1 = shc[11] * y * (4.0 * zz - xx - yy);
64    let l3m0 = shc[12] * z * (2.0 * zz - 3.0 * xx - 3.0 * yy);
65    let l3p1 = shc[13] * x * (4.0 * zz - xx - yy);
66    let l3p2 = shc[14] * z * (xx - yy);
67    let l3p3 = shc[15] * x * (xx - 3.0 * yy);
68
69    color += l3m3 * vec3<f32>(sh[27], sh[28], sh[29]);
70    color += l3m2 * vec3<f32>(sh[30], sh[31], sh[32]);
71    color += l3m1 * vec3<f32>(sh[33], sh[34], sh[35]);
72    color += l3m0 * vec3<f32>(sh[36], sh[37], sh[38]);
73    color += l3p1 * vec3<f32>(sh[39], sh[40], sh[41]);
74    color += l3p2 * vec3<f32>(sh[42], sh[43], sh[44]);
75    color += l3p3 * vec3<f32>(sh[45], sh[46], sh[47]);
76#endif
77
78#if SH_DEGREE_TIME > 0
79    let duration = gaussian_uniforms.time_stop - gaussian_uniforms.time_start;
80    let theta = dir_t / duration;
81
82    let t1 = cos(2.0 * PI * theta);
83
84    let l0m0 = shc[0];
85
86    color += t1 * (
87        l0m0 * vec3<f32>(sh[48], sh[49], sh[50]) +
88        l1m1 * vec3<f32>(sh[51], sh[52], sh[53]) +
89        l1m0 * vec3<f32>(sh[54], sh[55], sh[56]) +
90        l1p1 * vec3<f32>(sh[57], sh[58], sh[59]) +
91        l2m2 * vec3<f32>(sh[60], sh[61], sh[62]) +
92        l2m1 * vec3<f32>(sh[63], sh[64], sh[65]) +
93        l2m0 * vec3<f32>(sh[66], sh[67], sh[68]) +
94        l2p1 * vec3<f32>(sh[69], sh[70], sh[71]) +
95        l2p2 * vec3<f32>(sh[72], sh[73], sh[74]) +
96        l3m3 * vec3<f32>(sh[75], sh[76], sh[77]) +
97        l3m2 * vec3<f32>(sh[78], sh[79], sh[80]) +
98        l3m1 * vec3<f32>(sh[81], sh[82], sh[83]) +
99        l3m0 * vec3<f32>(sh[84], sh[85], sh[86]) +
100        l3p1 * vec3<f32>(sh[87], sh[88], sh[89]) +
101        l3p2 * vec3<f32>(sh[90], sh[91], sh[92]) +
102        l3p3 * vec3<f32>(sh[93], sh[94], sh[95])
103    );
104
105    #if SH_DEGREE_TIME > 1
106        let t2 = cos(4.0 * PI * theta);
107
108        color += t1 * (
109            l0m0 * vec3<f32>(sh[ 96], sh[ 97], sh[ 98]) +
110            l1m1 * vec3<f32>(sh[ 99], sh[100], sh[101]) +
111            l1m0 * vec3<f32>(sh[102], sh[103], sh[104]) +
112            l1p1 * vec3<f32>(sh[105], sh[106], sh[107]) +
113            l2m2 * vec3<f32>(sh[108], sh[109], sh[110]) +
114            l2m1 * vec3<f32>(sh[111], sh[112], sh[113]) +
115            l2m0 * vec3<f32>(sh[114], sh[115], sh[116]) +
116            l2p1 * vec3<f32>(sh[117], sh[118], sh[119]) +
117            l2p2 * vec3<f32>(sh[120], sh[121], sh[122]) +
118            l3m3 * vec3<f32>(sh[123], sh[124], sh[125]) +
119            l3m2 * vec3<f32>(sh[126], sh[127], sh[128]) +
120            l3m1 * vec3<f32>(sh[129], sh[130], sh[131]) +
121            l3m0 * vec3<f32>(sh[132], sh[133], sh[134]) +
122            l3p1 * vec3<f32>(sh[135], sh[136], sh[137]) +
123            l3p2 * vec3<f32>(sh[138], sh[139], sh[140]) +
124            l3p3 * vec3<f32>(sh[141], sh[142], sh[143])
125        );
126    #endif
127#endif
128
129    return color;
130}