rustytorch/demos/rtx-digital-twin-demo/src/presets.rs

//! Preset geometries for Digital Twin demo.

use rtx_digital_twin::{OrganGeometry, TissueLabel, TissueType};
use rtx_digital_twin_shared::{DigitalTwinDemoError, DigitalTwinDemoResult};

/// Preset geometry generator.
pub struct GeometryPresets;

impl GeometryPresets {
    /// Create geometry from preset name.
    pub fn create(
        name: &str,
        resolution: [usize; 3],
        spacing: [f32; 3],
    ) -> DigitalTwinDemoResult<OrganGeometry> {
        match name.to_lowercase().as_str() {
            "liver_tumor" => Self::liver_with_tumor(resolution, spacing),
            "kidney" => Self::kidney(resolution, spacing),
            "brain" => Self::brain(resolution, spacing),
            "simple_sphere" => Self::simple_sphere(resolution, spacing),
            _ => Err(DigitalTwinDemoError::InvalidConfig(format!(
                "Unknown preset: {name}"
            ))),
        }
    }

    /// List available preset names.
    #[must_use]
    pub fn available() -> Vec<&'static str> {
        vec!["liver_tumor", "kidney", "brain", "simple_sphere"]
    }

    /// Create liver with embedded tumor.
    fn liver_with_tumor(
        resolution: [usize; 3],
        spacing: [f32; 3],
    ) -> DigitalTwinDemoResult<OrganGeometry> {
        let [nx, ny, nz] = resolution;
        let mut geometry = OrganGeometry::new(resolution, spacing);

        // Create ellipsoidal liver
        let center = [nx as f32 / 2.0, ny as f32 / 2.0, nz as f32 / 2.0];
        let radii = [nx as f32 * 0.4, ny as f32 * 0.3, nz as f32 * 0.35];

        for z in 0..nz {
            for y in 0..ny {
                for x in 0..nx {
                    let dx = (x as f32 - center[0]) / radii[0];
                    let dy = (y as f32 - center[1]) / radii[1];
                    let dz = (z as f32 - center[2]) / radii[2];

                    if dx * dx + dy * dy + dz * dz <= 1.0 {
                        geometry.set_label(x, y, z, TissueLabel::from(TissueType::Liver));
                    }
                }
            }
        }

        // Add tumor
        let tumor_center = [center[0] + radii[0] * 0.2, center[1], center[2]];
        let tumor_radius = nx as f32 * 0.1;

        geometry.create_sphere(
            [
                tumor_center[0] * spacing[0],
                tumor_center[1] * spacing[1],
                tumor_center[2] * spacing[2],
            ],
            tumor_radius * spacing[0],
            TissueLabel::from(TissueType::Tumor),
        );

        Ok(geometry)
    }

    /// Create kidney geometry.
    fn kidney(resolution: [usize; 3], spacing: [f32; 3]) -> DigitalTwinDemoResult<OrganGeometry> {
        let [nx, ny, nz] = resolution;
        let mut geometry = OrganGeometry::new(resolution, spacing);

        // Bean-shaped kidney approximation
        let center = [nx as f32 / 2.0, ny as f32 / 2.0, nz as f32 / 2.0];

        for z in 0..nz {
            for y in 0..ny {
                for x in 0..nx {
                    let dx = (x as f32 - center[0]) / (nx as f32 * 0.25);
                    let dy = (y as f32 - center[1]) / (ny as f32 * 0.35);
                    let dz = (z as f32 - center[2]) / (nz as f32 * 0.2);

                    // Outer ellipsoid
                    let dist = dx * dx + dy * dy + dz * dz;

                    // Create indentation on one side
                    let indent = if dx > 0.0 {
                        0.3 * (-dy * dy * 4.0).exp()
                    } else {
                        0.0
                    };

                    if dist <= 1.0 + indent {
                        geometry.set_label(x, y, z, TissueLabel::from(TissueType::Kidney));
                    }
                }
            }
        }

        Ok(geometry)
    }

    /// Create brain geometry.
    fn brain(resolution: [usize; 3], spacing: [f32; 3]) -> DigitalTwinDemoResult<OrganGeometry> {
        let [nx, ny, nz] = resolution;
        let mut geometry = OrganGeometry::new(resolution, spacing);

        // Simplified brain as two hemispheres
        let center = [nx as f32 / 2.0, ny as f32 / 2.0, nz as f32 / 2.0];

        for z in 0..nz {
            for y in 0..ny {
                for x in 0..nx {
                    let dx = (x as f32 - center[0]) / (nx as f32 * 0.35);
                    let dy = (y as f32 - center[1]) / (ny as f32 * 0.4);
                    let dz = (z as f32 - center[2]) / (nz as f32 * 0.3);

                    // Main brain shape
                    let dist = dx * dx + dy * dy + dz * dz;

                    // Create central fissure
                    let x_f = x as f32;
                    let fissure = if x_f > center[0] - 2.0 && x_f < center[0] + 2.0 {
                        0.1
                    } else {
                        0.0
                    };

                    if dist <= 1.0 - fissure {
                        geometry.set_label(x, y, z, TissueLabel::from(TissueType::BrainGrayMatter));
                    }
                }
            }
        }

        Ok(geometry)
    }

    /// Create simple sphere for testing.
    fn simple_sphere(
        resolution: [usize; 3],
        spacing: [f32; 3],
    ) -> DigitalTwinDemoResult<OrganGeometry> {
        let [nx, ny, nz] = resolution;
        let mut geometry = OrganGeometry::new(resolution, spacing);

        let center = [
            nx as f32 / 2.0 * spacing[0],
            ny as f32 / 2.0 * spacing[1],
            nz as f32 / 2.0 * spacing[2],
        ];
        let radius = nx.min(ny).min(nz) as f32 / 3.0 * spacing[0];

        geometry.create_sphere(center, radius, TissueLabel::from(TissueType::Muscle));

        Ok(geometry)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_available_presets() {
        let presets = GeometryPresets::available();
        assert!(presets.contains(&"liver_tumor"));
        assert!(presets.contains(&"kidney"));
    }

    #[test]
    fn test_create_liver_tumor() {
        let geometry =
            GeometryPresets::create("liver_tumor", [32, 32, 32], [1.0, 1.0, 1.0]).unwrap();

        assert_eq!(geometry.shape(), [32, 32, 32]);

        // Should have liver and tumor tissue
        let histogram = geometry.tissue_histogram();
        assert!(histogram.get(&TissueType::Liver).unwrap_or(&0) > &0);
        assert!(histogram.get(&TissueType::Tumor).unwrap_or(&0) > &0);
    }

    #[test]
    fn test_create_kidney() {
        let geometry = GeometryPresets::create("kidney", [32, 32, 32], [1.0, 1.0, 1.0]).unwrap();

        let histogram = geometry.tissue_histogram();
        assert!(histogram.get(&TissueType::Kidney).unwrap_or(&0) > &0);
    }

    #[test]
    fn test_create_brain() {
        let geometry = GeometryPresets::create("brain", [32, 32, 32], [1.0, 1.0, 1.0]).unwrap();

        let histogram = geometry.tissue_histogram();
        assert!(histogram.get(&TissueType::BrainGrayMatter).unwrap_or(&0) > &0);
    }

    #[test]
    fn test_unknown_preset() {
        let result = GeometryPresets::create("unknown", [32, 32, 32], [1.0, 1.0, 1.0]);
        assert!(result.is_err());
    }
}