1#define_import_path bevy_gaussian_splatting::gaussian_4d
2
3#import bevy_gaussian_splatting::bindings::{
4    view,
5    globals,
6    gaussian_uniforms,
7}
8
9#ifdef BUFFER_STORAGE
10    #import bevy_gaussian_splatting::planar::{
11        get_isotropic_rotations,
12        get_scale,
13        get_timestamp,
14        get_time_scale,
15    }
16#endif
17
18
19struct DecomposedGaussian4d {
20    cov3d: array<f32, 6>,
21    delta_mean: vec3<f32>,
22    opacity_modifier: f32,
23    dir_t: f32,
24    mask: bool,
25}
26
27
28fn outer_product(
29    a: vec3<f32>,
30    b: vec3<f32>,
31) -> mat3x3<f32> {
32    return mat3x3<f32>(
33        a.x * b.x, a.x * b.y, a.x * b.z,
34        a.y * b.x, a.y * b.y, a.y * b.z,
35        a.z * b.x, a.z * b.y, a.z * b.z,
36    );
37}
38
39fn conditional_cov3d(
40    position: vec3<f32>,
41    index: u32,
42    time: f32,
43) -> DecomposedGaussian4d {
44    let isotropic_rotations = get_isotropic_rotations(index);
45    let rotation = isotropic_rotations[0];
46    let rotation_r = isotropic_rotations[1];
47    let scale = get_scale(index);
48
49    let dt = time - get_timestamp(index);
50
51    let S = mat4x4<f32>(
52        gaussian_uniforms.global_scale * scale.x, 0.0, 0.0, 0.0,
53        0.0, gaussian_uniforms.global_scale * scale.y, 0.0, 0.0,
54        0.0, 0.0, gaussian_uniforms.global_scale * scale.z, 0.0,
55        0.0, 0.0, 0.0, get_time_scale(index),
56    );
57
58    let w = rotation.x;
59    let x = rotation.y;
60    let y = rotation.z;
61    let z = rotation.w;
62
63    let wr = rotation_r.x;
64    let xr = rotation_r.y;
65    let yr = rotation_r.z;
66    let zr = rotation_r.w;
67
68    let M_l = mat4x4<f32>(
69        w, -x, -y, -z,
70        x,  w, -z,  y,
71        y,  z,  w, -x,
72        z, -y,  x,  w,
73    );
74
75    let M_r = mat4x4<f32>(
76        wr, -xr, -yr, -zr,
77        xr,  wr,  zr, -yr,
78        yr, -zr,  wr,  xr,
79        zr,  yr, -xr,  wr,
80    );
81
82    let R = M_r * M_l;
83    let M = R * S;
84    let Sigma = transpose(M) * M;
85
86    let cov_t = Sigma[3][3];
87    let marginal_t = exp(-0.5 * dt * dt / cov_t);
88
89    let mask = marginal_t > 0.05;
90    if !mask {
91        return DecomposedGaussian4d(
92            array<f32, 6>(0.0, 0.0, 0.0, 0.0, 0.0, 0.0),
93            vec3<f32>(0.0, 0.0, 0.0),
94            0.0,
95            dt,
96            mask,
97        );
98    }
99
100    let opacity_modifier = marginal_t;
101
102    let cov11 = mat3x3<f32>(
103        Sigma[0][0], Sigma[0][1], Sigma[0][2],
104        Sigma[1][0], Sigma[1][1], Sigma[1][2],
105        Sigma[2][0], Sigma[2][1], Sigma[2][2],
106    );
107    let cov12 = vec3<f32>(Sigma[0][3], Sigma[1][3], Sigma[2][3]);
108    let cov_op = outer_product(cov12, cov12);
109    let cov_op_t = mat3x3<f32>(
110        cov_op[0].x / cov_t, cov_op[0].y / cov_t, cov_op[0].z / cov_t,
111        cov_op[1].x / cov_t, cov_op[1].y / cov_t, cov_op[1].z / cov_t,
112        cov_op[2].x / cov_t, cov_op[2].y / cov_t, cov_op[2].z / cov_t,
113    );
114    let cov3d_condition = cov11 - cov_op_t;
115
116    let delta_mean = (cov12 / cov_t) * dt;
117
118    return DecomposedGaussian4d(
119        array<f32, 6>(
120            cov3d_condition[0][0],
121            cov3d_condition[0][1],
122            cov3d_condition[0][2],
123            cov3d_condition[1][1],
124            cov3d_condition[1][2],
125            cov3d_condition[2][2]
126        ),
127        delta_mean,
128        opacity_modifier,
129        dt,
130        mask,
131    );
132}