crater/csg/

regions.rs

//! [`Region`]s are collections of points inside an [`Isosurface`]

use crate::{csg::prelude::*, primitives::nvector::NVector};
use burn::prelude::*;

use super::{fields::IntoIsosurface, transformations::Translate};

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum Side {
    Negative,
    Positive,
}

/// A region in N-dimensional space defined by constructive solid geometry (CSG).
///
/// [`Region`] represents a collection of points satisfying boolean combinations of implicit
/// surface constraints. It forms the foundation of the CSG system, enabling construction
/// of complex 3D models from simpler primitive shapes through union, intersection, and
/// complement operations.
///
/// # Mathematical Foundation
///
/// A region R ⊂ ℝⁿ is defined recursively:
/// - **Half-space**: R = {x : f(x) ⋈ 0} where ⋈ ∈ {<, >} and f is a scalar field
/// - **Union**: R = R₁ ∪ R₂ using algebraic max operations  
/// - **Intersection**: R = R₁ ∩ R₂ using algebraic min operations
///
/// The choice of [`CSGAlgebra`] determines the specific boolean operators used,
/// affecting surface smoothness and computational properties.
///
/// # Variants
///
/// * [`HalfSpace`](Region::HalfSpace) - Points on one side of an implicit surface
/// * [`Union`](Region::Union) - Points belonging to either of two regions
/// * [`Intersection`](Region::Intersection) - Points belonging to both regions
///
/// # Usage Patterns
///
/// Regions are typically constructed by starting with primitive shapes and combining
/// them using boolean operations:
///
/// ```rust
/// use crater::csg::prelude::*;
/// use crater::primitives::prelude::*;
/// use burn::prelude::*;
/// use burn::backend::ndarray::NdArrayDevice;
///
/// // Create primitive shapes
/// let sphere = Field3D::<burn::backend::ndarray::NdArray>::sphere(1.0, NdArrayDevice::Cpu).into_isosurface(0.0);
/// let cube = Field3D::<burn::backend::ndarray::NdArray>::sphere(1.5, NdArrayDevice::Cpu).into_isosurface(0.0);
///
/// // Build CSG tree: sphere inside cube
/// let sphere_region = Region::HalfSpace(sphere, Side::Negative);
/// let cube_region = Region::HalfSpace(cube, Side::Negative);
/// let rounded_cube = Region::Union(
///     Box::new(sphere_region),
///     Box::new(cube_region)
/// );
/// ```
///
/// # Performance Considerations
///
/// - **Tree depth**: Deep CSG trees increase evaluation cost linearly
/// - **Algebraic choice**: R-functions provide smoothness but higher computational cost
/// - **Surface complexity**: Complex implicit surfaces require careful numerical handling
/// - **Memory layout**: Boxed recursive structure has pointer indirection overhead
///
/// For detailed theory, see the [Constructive Solid Geometry](../../book/theory/csg.md) chapter.
#[derive(Debug, Clone)]
pub enum Region<const N: usize, B: Backend> {
    /// Points on one side of an implicit surface.
    ///
    /// Represents the set {x : side(f(x)) ⋈ 0} where f is the isosurface scalar field.
    /// The [`Side`] determines whether we include points where f < 0 or f > 0.
    HalfSpace(Isosurface<N, B>, Side),

    /// Union of two regions: R₁ ∪ R₂.
    ///
    /// Points belong to the union if they belong to either the left or right region.
    /// The specific implementation depends on the chosen [`CSGAlgebra`].
    Union(Box<Region<N, B>>, Box<Region<N, B>>),

    /// Intersection of two regions: R₁ ∩ R₂.
    ///
    /// Points belong to the intersection only if they belong to both regions.
    /// The specific implementation depends on the chosen [`CSGAlgebra`].
    Intersection(Box<Region<N, B>>, Box<Region<N, B>>),
}

impl<const N: usize, B: Backend> Region<N, B> {
    /// Evaluates the region's isosurface function at multiple origins.
    ///
    /// # Parameters
    ///
    /// * `points` - Tensor of shape `[M, N]` where `M` is the number of origins and `N` is the spatial dimension
    /// * `algebra` - CSG algebra defining how boolean operations are computed numerically
    ///
    /// # Returns
    ///
    /// Tensor of shape `[M]` containing the isosurface function values for each input origin.
    ///
    /// # Examples
    ///
    /// ```rust
    /// use crater::csg::prelude::*;
    /// use burn::backend::ndarray::NdArrayDevice;
    /// use burn::prelude::*;
    ///
    /// // Create a unit sphere region
    /// let sphere = Field3D::<burn::backend::ndarray::NdArray>::sphere(1.0, NdArrayDevice::Cpu).into_isosurface(0.0);
    /// let region = Region::HalfSpace(sphere, Side::Negative);
    /// let algebra = Algebra::default();
    ///
    /// // Test points: center, surface, and exterior
    /// let points = Tensor::from_data([
    ///     [0.0, 0.0, 0.0],  // Center (inside)
    ///     [1.0, 0.0, 0.0],  // Surface
    ///     [2.0, 0.0, 0.0],  // Outside
    /// ], &NdArrayDevice::Cpu);
    ///
    /// let values = region.evaluate(points, &algebra);
    /// // values ≈ [-1.0, 0.0, 1.0] (inside, on surface, outside)
    /// ```
    pub fn evaluate(
        &self,
        points: Tensor<B, 2, Float>,
        algebra: &impl CSGAlgebra<B>,
    ) -> Tensor<B, 1, Float> {
        match self {
            Region::HalfSpace(surf, side) => {
                // When we're at a raw halfspace, we just let the
                // surface do the classification, and change the sign based
                // on whether we're inside or outside.
                // NOTE: Negative means we want the f < 0 points to be considered "inside"
                //       Positive means we want the f > 0 points to be considered "inside"
                //       So we negate positive halfspaces.
                match side {
                    Side::Negative => surf.evaluate(points),
                    Side::Positive => -surf.evaluate(points),
                }
            }
            // For Unions and Intersections, we evaluate the two regions, and
            // use the algebra to define the combination
            Region::Union(l, r) => {
                let a = l.evaluate(points.clone(), algebra);
                let b = r.evaluate(points.clone(), algebra);
                algebra.union(a, b)
            }
            Region::Intersection(l, r) => {
                let a = l.evaluate(points.clone(), algebra);
                let b = r.evaluate(points.clone(), algebra);
                algebra.intersection(a, b)
            }
        }
    }

    /// Classify a point as being inside, outside, or on the region.
    pub fn classify_point(
        &self,
        point: &NVector<N>,
        algebra: &impl CSGAlgebra<B>,
    ) -> Classification {
        let device = self.device();
        self.evaluate(Tensor::from_data([*point], device), algebra)
            .classification_of_index(0)
    }

    /// Get the device on which the [`Region`] is allocated.
    pub fn device(&self) -> &B::Device {
        // Always the same device for all sub-regions
        match self {
            Region::HalfSpace(surf, _) => surf.device(),
            Region::Union(l, _) => l.device(),
            Region::Intersection(l, _) => l.device(),
        }
    }

    pub fn halfspaces(&self) -> HalfSpaceIter<'_, N, B> {
        HalfSpaceIter::new(self)
    }
}

pub struct HalfSpaceIter<'a, const N: usize, B: Backend> {
    stack: Vec<&'a Region<N, B>>,
}

impl<'a, const N: usize, B: Backend> HalfSpaceIter<'a, N, B> {
    fn new(region: &'a Region<N, B>) -> Self {
        Self {
            stack: vec![region],
        }
    }
}

impl<'a, const N: usize, B: Backend> Iterator for HalfSpaceIter<'a, N, B> {
    type Item = &'a Region<N, B>;

    fn next(&mut self) -> Option<Self::Item> {
        while let Some(node) = self.stack.pop() {
            match node {
                Region::HalfSpace(_, _) => return Some(node),
                Region::Union(l, r) | Region::Intersection(l, r) => {
                    self.stack.push(r);
                    self.stack.push(l);
                }
            }
        }
        None
    }
}

impl<const N: usize, B: Backend> IntoRegion<N, B> for Isosurface<N, B> {
    fn into_region(self, _device: B::Device) -> Region<N, B> {
        Region::HalfSpace(self, Side::Negative)
    }
}

impl<const N: usize, B: Backend> Isosurface<N, B> {
    pub fn region(self) -> Region<N, B> {
        let device = self.device().clone();
        self.into_region(device)
    }
}

impl std::ops::Neg for Side {
    type Output = Side;
    fn neg(self) -> Self::Output {
        match self {
            Side::Negative => Side::Positive,
            Side::Positive => Side::Negative,
        }
    }
}

impl<const N: usize, B: Backend> std::ops::Neg for &Region<N, B> {
    type Output = Region<N, B>;
    fn neg(self) -> Self::Output {
        match self {
            Region::HalfSpace(surf, side) => Region::HalfSpace(surf.clone(), -*side),
            Region::Union(a, b) => {
                Region::Intersection(Box::new(-*a.clone()), Box::new(-*b.clone()))
            }
            Region::Intersection(a, b) => {
                Region::Union(Box::new(-*a.clone()), Box::new(-*b.clone()))
            }
        }
    }
}
impl<const N: usize, B: Backend> std::ops::Neg for Region<N, B> {
    type Output = Region<N, B>;
    fn neg(self) -> Self::Output {
        match self {
            Region::HalfSpace(surf, side) => Region::HalfSpace(surf, -side),
            Region::Union(a, b) => Region::Intersection(Box::new(-*a), Box::new(-*b)),
            Region::Intersection(a, b) => Region::Union(Box::new(-*a), Box::new(-*b)),
        }
    }
}
impl<const N: usize, B: Backend> std::ops::BitAnd for &Region<N, B> {
    type Output = Region<N, B>;
    fn bitand(self, rhs: Self) -> Self::Output {
        Region::Intersection(Box::new(self.clone()), Box::new(rhs.clone()))
    }
}
impl<const N: usize, B: Backend> std::ops::BitAnd for Region<N, B> {
    type Output = Region<N, B>;
    fn bitand(self, rhs: Self) -> Self::Output {
        Region::Intersection(Box::new(self), Box::new(rhs))
    }
}
impl<const N: usize, B: Backend> std::ops::BitOr for &Region<N, B> {
    type Output = Region<N, B>;
    fn bitor(self, rhs: Self) -> Self::Output {
        Region::Union(Box::new(self.clone()), Box::new(rhs.clone()))
    }
}
impl<const N: usize, B: Backend> std::ops::BitOr for Region<N, B> {
    type Output = Region<N, B>;
    fn bitor(self, rhs: Self) -> Self::Output {
        Region::Union(Box::new(self), Box::new(rhs))
    }
}

// impl<'de, const N: usize, B: Backend> Deserialize<'de> for Region<N, B> {
//     fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
//     where
//         D: serde::Deserializer<'de>,
//     {
//         // Deserialize and parse into region
//         let r_string = String::deserialize(deserializer)?.to_owned();
//         str_to_region(&r_string).map_err(serde::de::Error::custom)
//     }
// }

/// Convert a type into a [`Region`]
pub trait IntoRegion<const N: usize, B: Backend> {
    fn into_region(self, device: B::Device) -> Region<N, B>;
}

impl<const N: usize, B: Backend> IntoRegion<N, B> for crate::primitives::bounding::BoundingBox<N> {
    fn into_region(self, device: B::Device) -> Region<N, B> {
        // A box in dimension N has 2N faces. Generate each as a hyperplane,
        // then take the intersection of all of them.
        let faces = (0..N)
            .flat_map(|i| {
                let mut normal = [0.0; N];
                let mut offset_min = [0.0; N];
                let mut offset_max = [0.0; N];

                normal[i] = 1.0;
                offset_min[i] = self.min()[i];
                offset_max[i] = self.max()[i];
                [
                    -FieldND::hyperplane(normal, device.clone())
                        .into_isosurface(0.0)
                        .transform(Translate(offset_min))
                        .region(),
                    FieldND::hyperplane(normal, device.clone())
                        .into_isosurface(0.0)
                        .transform(Translate(offset_max))
                        .region(),
                ]
            })
            .collect::<Vec<_>>();

        faces.into_iter().reduce(|a, b| a & b).unwrap()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::primitives::bounding::BoundingBox;
    use backend_macro::with_backend;

    use rstest::rstest;

    #[with_backend]
    #[rstest]
    fn test_halfspace() {
        // Try out a ZPlane halfspace
        // z > 0 halfspace (using Side::Positive)
        let surface = Field3D::<Backend>::plane([0.0, 0.0, 1.0], device());
        let halfspace = Region::HalfSpace(surface.into_isosurface(0.0), Side::Positive);
        assert_eq!(
            halfspace.classify_point(&[0.0, 0.0, 0.0], &Algebra::default()),
            Classification::On
        );
        assert_eq!(
            halfspace.classify_point(&[0.0, 0.0, 1.0], &Algebra::default()),
            Classification::Inside(-1.0)
        );
        assert_eq!(
            halfspace.classify_point(&[0.0, 0.0, -2.0], &Algebra::default()),
            Classification::Outside(2.0)
        );

        // Try a Sphere halfspace (everything inside the sphere)
        let surface = Field3D::<Backend>::sphere(1.0, device());
        let halfspace = Region::HalfSpace(surface.into_isosurface(0.0), Side::Negative);

        assert_eq!(
            halfspace.classify_point(&[0.0, 0.0, 0.0], &Algebra::default()),
            Classification::Inside(-1.0)
        );
        assert_eq!(
            halfspace.classify_point(&[0.0, 0.0, 1.0], &Algebra::default()),
            Classification::On
        );
        assert_eq!(
            halfspace.classify_point(&[0.0, 0.0, -2.0], &Algebra::default()),
            Classification::Outside(3.0)
        );
    }

    #[with_backend]
    #[rstest]
    fn test_regions() {
        // Try out a ZPlane halfspace
        // z > 0 halfspace
        let surface = Field3D::<Backend>::plane([0.0, 0.0, 1.0], device());
        let halfspace = Region::HalfSpace(surface.into_isosurface(0.0), Side::Positive);

        // Try a Sphere halfspace (everything inside the sphere)
        let surface = Field3D::<Backend>::sphere(1.0, device());
        let halfspace2 = Region::HalfSpace(surface.into_isosurface(0.0), Side::Negative);

        // Everything inside the sphere OR above the z=0 plane
        let union = Region::Union(Box::new(halfspace), Box::new(halfspace2));

        assert_eq!(
            union.classify_point(&[0.0, 0.0, 0.0], &Algebra::default()),
            Classification::Inside(-1.0) // We're "on" the plane, but inside the sphere, so the sphere wins
        );
        assert_eq!(
            union.classify_point(&[0.0, 0.0, 1.0], &Algebra::default()),
            Classification::Inside(-1.0) // We're "on" the sphere, but inside the plane halfspace, so the plane wins
        );
        assert_eq!(
            union.classify_point(&[-1.0, 0.0, 0.0], &Algebra::default()),
            Classification::On
        );
        assert_eq!(
            union.classify_point(&[0.0, 0.0, 0.5], &Algebra::default()),
            Classification::Inside(-0.75)
        );
        assert_eq!(
            union.classify_point(&[0.0, 0.0, -2.0], &Algebra::default()),
            Classification::Outside(2.0)
        );
        assert_eq!(
            union.classify_point(&[0.0, 0.0, 3.0], &Algebra::default()),
            Classification::Inside(-3.0)
        );

        // Try out some Intersections
        let surface = Field3D::<Backend>::plane([0.0, 0.0, 1.0], device());
        let halfspace = Region::HalfSpace(surface.into_isosurface(0.0), Side::Positive);

        // Try a Sphere halfspace (everything inside the sphere)
        let surface = Field3D::<Backend>::sphere(1.0, device());
        let halfspace2 = Region::HalfSpace(surface.into_isosurface(0.0), Side::Negative);

        // Everything inside the sphere AND above the z=0 plane
        let intersection = Region::Intersection(Box::new(halfspace), Box::new(halfspace2));

        assert_eq!(
            intersection.classify_point(&[0.0, 0.0, 0.0], &Algebra::default()),
            Classification::On
        );
        assert_eq!(
            intersection.classify_point(&[0.0, 0.0, 2.0], &Algebra::default()),
            Classification::Outside(3.0)
        );
    }

    #[with_backend]
    #[rstest]
    fn test_bounding_box() {
        // 1D
        let bbox = BoundingBox::new([0.0], [1.0]);
        let region: Region<1, Backend> = bbox.into_region(device());
        assert_eq!(
            region.classify_point(&[0.0], &Algebra::default()),
            Classification::On
        );
        assert_eq!(
            region.classify_point(&[1.0], &Algebra::default()),
            Classification::On
        );
        assert_eq!(
            region.classify_point(&[0.5], &Algebra::default()),
            Classification::Inside(-0.5)
        );

        let bbox = BoundingBox::new([0.0, 0.0, 0.0], [1.0, 1.0, 1.0]);
        let region: Region<3, Backend> = bbox.into_region(device());
        assert_eq!(
            region.classify_point(&[0.0, 0.0, 0.0], &Algebra::default()),
            Classification::On
        );
        assert_eq!(
            region.classify_point(&[1.0, 1.0, 1.0], &Algebra::default()),
            Classification::On
        );
        assert_eq!(
            region.classify_point(&[0.5, 0.5, 0.5], &Algebra::default()),
            Classification::Inside(-0.5)
        );
    }
}