1#define_import_path bevy_gaussian_splatting::planar
2
3#ifdef GAUSSIAN_3D_STRUCTURE
4    #ifdef PRECOMPUTE_COVARIANCE_3D
5        #import bevy_gaussian_splatting::bindings::{
6            position_visibility,
7            spherical_harmonics,
8            covariance_3d_opacity,
9        }
10    #else
11        #import bevy_gaussian_splatting::bindings::{
12            position_visibility,
13            spherical_harmonics,
14            rotation,
15            rotation_scale_opacity,
16            scale_opacity,
17        }
18    #endif
19
20    #import bevy_gaussian_splatting::spherical_harmonics::{
21        spherical_harmonics_lookup,
22        srgb_to_linear,
23    }
24#else ifdef GAUSSIAN_4D
25    #import bevy_gaussian_splatting::bindings::{
26        position_visibility,
27        spherindrical_harmonics,
28        isotropic_rotations,
29        scale_opacity,
30        timestamp_timescale,
31    }
32
33    #import bevy_gaussian_splatting::spherical_harmonics::srgb_to_linear
34    #import bevy_gaussian_splatting::spherindrical_harmonics::spherindrical_harmonics_lookup
35#endif
36
37
38#ifdef GAUSSIAN_3D_STRUCTURE
39    #ifdef PLANAR_F16
40        fn get_color(
41            index: u32,
42            ray_direction: vec3<f32>,
43        ) -> vec3<f32> {
44            let sh = get_spherical_harmonics(index);
45            let color = spherical_harmonics_lookup(ray_direction, sh);
46            return srgb_to_linear(color);
47        }
48
49        fn get_position(index: u32) -> vec3<f32> {
50            return position_visibility[index].xyz;
51        }
52
53        fn get_spherical_harmonics(index: u32) -> array<f32, #{SH_COEFF_COUNT}> {
54            var coefficients: array<f32, #{SH_COEFF_COUNT}>;
55
56            for (var i = 0u; i < #{HALF_SH_COEFF_COUNT}u; i = i + 1u) {
57                let values = unpack2x16float(spherical_harmonics[index][i]);
58
59                coefficients[i * 2u] = values[0];
60                coefficients[i * 2u + 1u] = values[1];
61            }
62
63            return coefficients;
64        }
65
66        #ifdef PRECOMPUTE_COVARIANCE_3D
67            fn get_cov3d(index: u32) -> array<f32, 6> {
68                let c0 = unpack2x16float(covariance_3d_opacity[index].x);
69                let c1 = unpack2x16float(covariance_3d_opacity[index].y);
70                let c2 = unpack2x16float(covariance_3d_opacity[index].z);
71
72                var cov3d: array<f32, 6>;
73
74                cov3d[0] = c0.y;
75                cov3d[1] = c0.x;
76                cov3d[2] = c1.y;
77                cov3d[3] = c1.x;
78                cov3d[4] = c2.y;
79                cov3d[5] = c2.x;
80
81                return cov3d;
82            }
83        #else
84            fn get_rotation(index: u32) -> vec4<f32> {
85                let q0 = unpack2x16float(rotation_scale_opacity[index].x);
86                let q1 = unpack2x16float(rotation_scale_opacity[index].y);
87
88                return vec4<f32>(
89                    q0.yx,
90                    q1.yx,
91                );
92            }
93
94            fn get_scale(index: u32) -> vec3<f32> {
95                let s0 = unpack2x16float(rotation_scale_opacity[index].z);
96                let s1 = unpack2x16float(rotation_scale_opacity[index].w);
97
98                return vec3<f32>(
99                    s0.yx,
100                    s1.y,
101                );
102            }
103        #endif
104
105        fn get_opacity(index: u32) -> f32 {
106            #ifdef PRECOMPUTE_COVARIANCE_3D
107                return unpack2x16float(covariance_3d_opacity[index].w).y;
108            #else
109                return unpack2x16float(rotation_scale_opacity[index].w).x;
110            #endif
111        }
112
113        fn get_visibility(index: u32) -> f32 {
114            return position_visibility[index].w;
115        }
116    #else ifdef PLANAR_F32
117        fn get_color(
118            index: u32,
119            ray_direction: vec3<f32>,
120        ) -> vec3<f32> {
121            let sh = get_spherical_harmonics(index);
122            let color = spherical_harmonics_lookup(ray_direction, sh);
123            return srgb_to_linear(color);
124        }
125
126        fn get_position(index: u32) -> vec3<f32> {
127            return position_visibility[index].xyz;
128        }
129
130        fn get_spherical_harmonics(index: u32) -> array<f32, #{SH_COEFF_COUNT}> {
131            return spherical_harmonics[index];
132        }
133
134        fn get_rotation(index: u32) -> vec4<f32> {
135            return rotation[index];
136        }
137
138        fn get_scale(index: u32) -> vec3<f32> {
139            return scale_opacity[index].xyz;
140        }
141
142        fn get_opacity(index: u32) -> f32 {
143            return scale_opacity[index].w;
144        }
145
146        fn get_visibility(index: u32) -> f32 {
147            return position_visibility[index].w;
148        }
149    #endif
150#else ifdef GAUSSIAN_4D
151    #ifdef PLANAR_F32
152        fn get_color(
153            index: u32,
154            dir_t: f32,
155            ray_direction: vec3<f32>,
156        ) -> vec3<f32> {
157            let sh = get_spherindrical_harmonics(index);
158            let color = spherindrical_harmonics_lookup(ray_direction, dir_t, sh);
159            return srgb_to_linear(color);
160        }
161
162        fn get_isotropic_rotations(index: u32) -> mat2x4<f32> {
163            let r1x = isotropic_rotations[index][0];
164            let r1y = isotropic_rotations[index][1];
165            let r1z = isotropic_rotations[index][2];
166            let r1w = isotropic_rotations[index][3];
167
168            let r2x = isotropic_rotations[index][4];
169            let r2y = isotropic_rotations[index][5];
170            let r2z = isotropic_rotations[index][6];
171            let r2w = isotropic_rotations[index][7];
172
173            return mat2x4<f32>(
174                r1x, r1y, r1z, r1w,
175                r2x, r2y, r2z, r2w,
176            );
177        }
178
179        fn get_scale(index: u32) -> vec3<f32> {
180            return scale_opacity[index].xyz;
181        }
182
183        fn get_opacity(index: u32) -> f32 {
184            return scale_opacity[index].w;
185        }
186
187        fn get_position(index: u32) -> vec3<f32> {
188            return position_visibility[index].xyz;
189        }
190
191        fn get_visibility(index: u32) -> f32 {
192            return position_visibility[index].w;
193        }
194
195        fn get_spherindrical_harmonics(index: u32) -> array<f32, #{SH_4D_COEFF_COUNT}> {
196            return spherindrical_harmonics[index];
197        }
198
199        fn get_timestamp(index: u32) -> f32 {
200            return timestamp_timescale[index].x;
201        }
202
203        fn get_time_scale(index: u32) -> f32 {
204            return timestamp_timescale[index].y;
205        }
206    #endif
207
208    // TODO: PLANAR_F16 for GAUSSIAN_4D
209#endif