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