refactor: better format separation

mosure · Jan 3, 2025 · d19d15c · d19d15c
1 parent f0f7bdb
commit d19d15c
Show file tree

Hide file tree

Showing 16 changed files with 262 additions and 263 deletions.
diff --git a/src/gaussian/f16.rs b/src/gaussian/f16.rs
@@ -21,9 +21,9 @@ use crate::gaussian::{
         Rotation,
         ScaleOpacity,
     },
-    packed::{
-        Gaussian3d,
-        Gaussian4d,
+    formats::{
+        planar_3d::Gaussian3d,
+        planar_4d::Gaussian4d,
     },
 };
 

diff --git a/src/gaussian/f32.rs b/src/gaussian/f32.rs
@@ -15,7 +15,10 @@ use serde::{
 
 use crate::gaussian::{
     covariance::compute_covariance_3d,
-    packed::{Gaussian3d, Gaussian4d},
+    formats::{
+        planar_3d::Gaussian3d,
+        planar_4d::Gaussian4d,
+    },
 };
 
 

diff --git a/src/gaussian/formats/planar_3d.rs b/src/gaussian/formats/planar_3d.rs
@@ -1,9 +1,20 @@
+use std::marker::Copy;
 use rand::{
+    prelude::Distribution,
     seq::SliceRandom,
     Rng,
 };
 
+use bevy::prelude::*;
 use bevy_interleave::prelude::*;
+use bytemuck::{
+    Pod,
+    Zeroable,
+};
+use serde::{
+    Deserialize,
+    Serialize,
+};
 
 #[allow(unused_imports)]
 use crate::{
@@ -19,7 +30,6 @@ use crate::{
             TestCloud,
         },
         iter::PositionIter,
-        packed::{Gaussian3d, PlanarGaussian3d},
         settings::CloudSettings,
     },
     material::spherical_harmonics::{
@@ -29,6 +39,35 @@ use crate::{
     },
 };
 
+
+
+#[derive(
+    Clone,
+    Debug,
+    Default,
+    Copy,
+    PartialEq,
+    Planar,
+    ReflectInterleaved,
+    StorageBindings,
+    Reflect,
+    Pod,
+    Zeroable,
+    Serialize,
+    Deserialize,
+)]
+#[repr(C)]
+pub struct Gaussian3d {
+    pub position_visibility: PositionVisibility,
+    pub spherical_harmonic: SphericalHarmonicCoefficients,
+    pub rotation: Rotation,
+    pub scale_opacity: ScaleOpacity,
+}
+
+pub type Gaussian2d = Gaussian3d;  // GaussianMode::Gaussian2d /w Gaussian3d structure
+
+
+
 // #[allow(unused_imports)]
 // #[cfg(feature = "f16")]
 // use crate::gaussian::f16::{
@@ -94,6 +133,66 @@ impl From<Vec<Gaussian3d>> for PlanarGaussian3d {
 }
 
 
+impl Distribution<Gaussian3d> for rand::distributions::Standard {
+    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Gaussian3d {
+        Gaussian3d {
+            rotation: [
+                rng.gen_range(-1.0..1.0),
+                rng.gen_range(-1.0..1.0),
+                rng.gen_range(-1.0..1.0),
+                rng.gen_range(-1.0..1.0),
+            ].into(),
+            position_visibility: [
+                rng.gen_range(-20.0..20.0),
+                rng.gen_range(-20.0..20.0),
+                rng.gen_range(-20.0..20.0),
+                1.0,
+            ].into(),
+            scale_opacity: [
+                rng.gen_range(0.0..1.0),
+                rng.gen_range(0.0..1.0),
+                rng.gen_range(0.0..1.0),
+                rng.gen_range(0.0..0.8),
+            ].into(),
+            spherical_harmonic: SphericalHarmonicCoefficients {
+                coefficients: {
+                    // #[cfg(feature = "f16")]
+                    // {
+                    //     let mut coefficients: [u32; HALF_SH_COEFF_COUNT] = [0; HALF_SH_COEFF_COUNT];
+                    //     for coefficient in coefficients.iter_mut() {
+                    //         let upper = rng.gen_range(-1.0..1.0);
+                    //         let lower = rng.gen_range(-1.0..1.0);
+
+                    //         *coefficient = pack_f32s_to_u32(upper, lower);
+                    //     }
+                    //     coefficients
+                    // }
+
+                    {
+                        let mut coefficients = [0.0; SH_COEFF_COUNT];
+                        for coefficient in coefficients.iter_mut() {
+                            *coefficient = rng.gen_range(-1.0..1.0);
+                        }
+                        coefficients
+                    }
+                },
+            },
+        }
+    }
+}
+
+pub fn random_gaussians_3d(n: usize) -> PlanarGaussian3d {
+    let mut rng = rand::thread_rng();
+    let mut gaussians: Vec<Gaussian3d> = Vec::with_capacity(n);
+
+    for _ in 0..n {
+        gaussians.push(rng.gen());
+    }
+
+    PlanarGaussian3d::from_interleaved(gaussians)
+}
+
+
 impl TestCloud for PlanarGaussian3d {
     fn test_model() -> Self {
         let mut rng = rand::thread_rng();

diff --git a/src/gaussian/formats/planar_4d.rs b/src/gaussian/formats/planar_4d.rs
@@ -1,17 +1,76 @@
+use std::marker::Copy;
+
+use bevy::prelude::*;
 use bevy_interleave::prelude::*;
+use bytemuck::{
+    Pod,
+    Zeroable,
+};
+use rand::{
+    prelude::Distribution,
+    Rng,
+};
+use serde::{
+    Deserialize,
+    Serialize,
+};
 
 use crate::{
     gaussian::{
+        f32::{
+            IsotropicRotations,
+            PositionVisibility,
+            ScaleOpacity,
+            TimestampTimescale,
+        },
         interface::{
             CommonCloud,
             TestCloud,
         },
         iter::PositionIter,
-        packed::{Gaussian4d, PlanarGaussian4d},
     },
-    random_gaussians_4d,
+    material::spherindrical_harmonics::{
+        SH_4D_COEFF_COUNT,
+        SpherindricalHarmonicCoefficients,
+    },
 };
 
+
+
+#[derive(
+    Clone,
+    Debug,
+    Default,
+    Copy,
+    PartialEq,
+    Planar,
+    ReflectInterleaved,
+    StorageBindings,
+    Reflect,
+    Pod,
+    Zeroable,
+    Serialize,
+    Deserialize,
+)]
+#[repr(C)]
+pub struct Gaussian4d {
+    pub position_visibility: PositionVisibility,
+    pub spherindrical_harmonic: SpherindricalHarmonicCoefficients,
+    pub isotropic_rotations: IsotropicRotations,
+    pub scale_opacity: ScaleOpacity,
+    pub timestamp_timescale: TimestampTimescale,
+}
+
+// // TODO: GaussianSpacetime, determine temporal position/rotation structure
+// pub struct GaussianSpacetime {
+//     pub position_visibility: PositionVisibility,
+//     pub color_mlp: ColorMlp,
+//     pub isotropic_rotations: IsotropicRotations,
+//     pub scale_opacity: ScaleOpacity,
+//     pub timestamp_timescale: TimestampTimescale,
+// }
+
+
 // TODO: quantize 4d representation
 // #[derive(
 //     Debug,
@@ -197,8 +256,62 @@ impl From<Vec<Gaussian4d>> for PlanarGaussian4d {
 }
 
 
+impl Distribution<Gaussian4d> for rand::distributions::Standard {
+    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Gaussian4d {
+        let mut coefficients = [0.0; SH_4D_COEFF_COUNT];
+        for coefficient in coefficients.iter_mut() {
+            *coefficient = rng.gen_range(-1.0..1.0);
+        }
+
+        Gaussian4d {
+            isotropic_rotations: [
+                rng.gen_range(-1.0..1.0),
+                rng.gen_range(-1.0..1.0),
+                rng.gen_range(-1.0..1.0),
+                rng.gen_range(-1.0..1.0),
+                rng.gen_range(-1.0..1.0),
+                rng.gen_range(-1.0..1.0),
+                rng.gen_range(-1.0..1.0),
+                rng.gen_range(-1.0..1.0),
+            ].into(),
+            position_visibility: [
+                rng.gen_range(-20.0..20.0),
+                rng.gen_range(-20.0..20.0),
+                rng.gen_range(-20.0..20.0),
+                1.0,
+            ].into(),
+            scale_opacity: [
+                rng.gen_range(0.0..1.0),
+                rng.gen_range(0.0..1.0),
+                rng.gen_range(0.0..1.0),
+                rng.gen_range(0.0..0.8),
+            ].into(),
+            spherindrical_harmonic: coefficients.into(),
+            timestamp_timescale: [
+                rng.gen_range(0.0..1.0),
+                rng.gen_range(-1.0..1.0),
+                0.0,
+                0.0,
+            ].into(),
+        }
+    }
+}
+
+pub fn random_gaussians_4d(n: usize) -> PlanarGaussian4d {
+    let mut rng = rand::thread_rng();
+    let mut gaussians: Vec<Gaussian4d> = Vec::with_capacity(n);
+
+    for _ in 0..n {
+        gaussians.push(rng.gen());
+    }
+
+    PlanarGaussian4d::from_interleaved(gaussians)
+}
+
+
 impl TestCloud for PlanarGaussian4d {
     fn test_model() -> Self {
         random_gaussians_4d(512)
     }
 }
+
diff --git a/src/gaussian/formats/planar_4d_hierarchy.rs b/src/gaussian/formats/planar_4d_hierarchy.rs
@@ -1,5 +1,5 @@
 // TODO: gaussian cloud 4d with temporal hierarchy
-use crate::gaussian::packed::PlanarGaussian4dHandle;
+use crate::gaussian::formats::planar_4d::PlanarGaussian4dHandle;
 
 
 pub struct TemporalGaussianLevel {

diff --git a/src/gaussian/mod.rs b/src/gaussian/mod.rs
@@ -7,8 +7,6 @@ pub mod f32;
 pub mod formats;
 pub mod interface;
 pub mod iter;
-pub mod packed;
-pub mod rand;
 pub mod settings;