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