diffable/

flat.rs

//! Flat manifolds obtained from Euclidean spaces by discrete identifications.
//!
//! [`S1`] is the quotient `R/Z`; [`Torus`] combines two circles with straight
//! gluing, while [`KleinBottle`] twists one identification. Their cover types
//! implement the [`Bounded`] machinery used by [`NerveComplex`] to recover
//! global topology.
//!
//! [`NerveComplex`]: crate::traits::simplicial::NerveComplex
//! [`Bounded`]: crate::traits::simplicial::Bounded

use crate::{
    discrete::Z,
    impl_lie_group_via_quotient,
    traits::{Interval, Tensor},
};
use core::marker::PhantomData;

use crate::traits::{Chart, Euclidean, Group, LieGroup, Quotient, Real, Smooth};
use num_traits::{Euclid, NumCast, One, Zero, real::Real as _};

/// The circle `S¹`, as the quotient of the line `V` by the integer lattice
/// [`Z`]. One-dimensional (`From<[F; 1]>`).
#[derive(Copy, Clone, Debug, PartialEq)]
pub struct S1<V: Euclidean + From<[<V as Tensor>::F; 1]>>(V);

impl<V: Euclidean<F: Real> + From<[V::F; 1]>> Interval for S1<V> {
    type R = V::F;

    fn interval_squared(&self, other: &Self) -> V::F {
        self.to_local(other).unwrap().norm_squared()
    }
}

impl<V: Euclidean<F: Real> + From<[<V as Tensor>::F; 1]>> Quotient<V, Z<V>, V> for S1<V> {
    fn new(g: V) -> Self {
        let one = V::F::one();
        let mut d = g[0].rem_euclid(&one);
        if d == one {
            // floating point garbage check gotta do it :(
            d = V::F::zero();
        }
        Self([d].into())
    }

    fn lift(&self) -> V {
        // nearest representative to the identity (0), not just the
        // canonical [0,1) one — reduce into (-1/2, 1/2] instead.
        let half = V::F::one() / (V::F::one() + V::F::one());
        let mut d = self.0[0];
        if d > half {
            d = d - V::F::one();
        }
        [d].into()
    }

    fn embed(h: Z<V>) -> V {
        [<V::F as NumCast>::from(h.0).unwrap()].into()
    }
}

impl_lie_group_via_quotient!(S1<V>, V, Z<V>, V, V: Euclidean + From<[<V as Tensor>::F; 1]>);

/// The 2-torus `T² = S¹ × S¹`.
#[derive(Debug, Copy, Clone, PartialEq)]
pub struct Torus<I: Euclidean + From<[I::F; 1]> + From<[V::F; 1]>, V: Euclidean + From<[I::F; 2]>>(
    S1<I>,
    S1<I>,
    PhantomData<V>,
);

impl<I: ICompatible<V>, V: VCompatible<I>> Torus<I, V> {
    /// Constructs a torus point from its two circle coordinates.
    pub fn new(a: S1<I>, b: S1<I>) -> Self {
        Self(a, b, PhantomData)
    }
}

impl<I: ICompatible<V>, V: VCompatible<I>> Interval for Torus<I, V>
where
    I::F: From<V::F>,
{
    type R = I::F;

    fn interval_squared(&self, other: &Self) -> I::F {
        self.to_local(other).unwrap().norm_squared().into()
    }
}

impl<I: ICompatible<V>, V: VCompatible<I>> Group for Torus<I, V> {
    fn identity() -> Self {
        Self::new(S1::identity(), S1::identity())
    }

    fn compose(&self, other: &Self) -> Self {
        Self::new(self.0.compose(&other.0), self.1.compose(&other.1))
    }

    fn inverse(&self) -> Self {
        Self::new(self.0.inverse(), self.1.inverse())
    }
}

impl<I: ICompatible<V>, V: VCompatible<I>> LieGroup<V> for Torus<I, V> {
    fn identity_exp(v: V) -> Self {
        let v0 = [v[0]].into();
        let v1 = [v[1]].into();
        Self::new(S1::identity_exp(v0), S1::identity_exp(v1))
    }

    fn identity_log(p: &Self) -> Option<V> {
        let a = S1::identity_log(&p.0)?;
        let b = S1::identity_log(&p.1)?;

        Some([a[0], b[0]].into())
    }
}

/// Dimensional-compatibility bound for the *inner circle* type of a
/// [`KleinBottle`]/[`Torus`], relating it to the ambient type `V`. See
/// [`VCompatible`] for the dual constraint.
pub trait ICompatible<V: Euclidean<F: Real + Send + Sync> + From<[Self::F; 2]>>:
    Euclidean<F: Real + From<V::F>> + From<[Self::F; 1]> + From<[V::F; 1]> + 'static + Send + Sync
{
}

/// Dimensional-compatibility bound for the *ambient* type of a
/// [`KleinBottle`]/[`Torus`], relating it to the inner circle type `I`.
pub trait VCompatible<I: Euclidean<F: Real + From<Self::F>> + From<[I::F; 1]> + From<[Self::F; 1]>>:
    Euclidean<F: Real + Send + Sync> + From<[I::F; 2]> + 'static + Send + Sync
{
}

impl<
    I: Euclidean<F: Real + From<V::F>> + From<[I::F; 1]> + From<[V::F; 1]> + 'static + Send + Sync,
    V: Euclidean<F: Real + Send + Sync> + From<[I::F; 2]> + 'static + Send + Sync,
> ICompatible<V> for I
{
}

impl<
    I: Euclidean<F: Real + From<V::F>> + From<[I::F; 1]> + From<[V::F; 1]> + 'static + Send + Sync,
    V: Euclidean<F: Real + Send + Sync> + From<[I::F; 2]> + 'static + Send + Sync,
> VCompatible<I> for V
{
}

/// The Klein bottle — the non-orientable quotient of the plane, built as two
/// circles with a twist. `I` is the "inner" circle coordinate type and `V` the
/// ambient embedding type; [`ICompatible`]/[`VCompatible`] pin the dimensional
/// relationship between them that the twist requires.
#[derive(Debug, Copy, Clone, PartialEq)]
pub struct KleinBottle<I: ICompatible<V>, V: VCompatible<I>>(S1<I>, S1<I>, PhantomData<V>);

impl<I: ICompatible<V>, V: VCompatible<I>> KleinBottle<I, V> {
    /// Constructs a Klein-bottle point from its two circle coordinates.
    pub fn new(a: S1<I>, b: S1<I>) -> Self {
        Self(a, b, PhantomData)
    }
}

impl<I: ICompatible<V>, V: VCompatible<I>> Smooth<V> for KleinBottle<I, V> {
    type Global = Self;

    fn exp(&self, v: V) -> Self {
        let (x, y) = self.coords();
        let vx: I::F = v[0].into();
        let vy: I::F = v[1].into();
        Self::from_cover(x + vx, y + vy)
    }

    fn log(&self, other: &Self) -> Option<V> {
        let one = I::F::one();
        let two = one + one;
        let (sx, sy) = self.coords();
        let (ox, oy) = other.coords();
        let mut best: Option<(I::F, I::F)> = None;
        let mut best_sq = I::F::zero();

        for n in [-one, I::F::zero(), one] {
            let n_odd = n.rem_euclid(&two) != I::F::zero();
            // Reflection formula in the (-1/2,1/2]-centered
            // convention is `-ox`, not `1 - ox` — reflecting about
            // 0 (the domain's center), not about 1/2 (which was
            // only the reflection point under the old [0,1)
            // convention).
            let base_ox = if n_odd { -ox } else { ox };
            for m in [-one, I::F::zero(), one] {
                let cx = base_ox + m;
                let cy = oy + n;
                let dx = cx - sx;
                let dy = cy - sy;
                let sq = dx * dx + dy * dy;
                if best.is_none() || sq < best_sq {
                    best = Some((dx, dy));
                    best_sq = sq;
                }
            }
        }
        best.map(|(dx, dy)| [dx, dy].into())
    }
}

impl<I: ICompatible<V>, V: VCompatible<I>> KleinBottle<I, V> {
    /// Reduce a cover point (x, y) ∈ ℝ² to the fundamental domain via
    /// Γ = ⟨A, B⟩, A: (x,y) ↦ (x+1, y), B: (x,y) ↦ (−x, y+1).
    ///
    /// Uses the SAME (-1/2, 1/2]-centered convention as `S1::lift`
    /// throughout: seam-crossing count is `y.round()` (nearest
    /// integer), not `y.floor()`, since the fundamental domain is
    /// centered at 0 rather than starting at 0. Parity of that
    /// rounded count decides the flip, exactly as before — only the
    /// rounding function and the reflection formula's center point
    /// (0, not 1/2) changed.
    fn from_cover(x: I::F, y: I::F) -> Self {
        let one = I::F::one();
        let two = one + one;
        let ky = y.round(); // nearest-centered seam count
        let y_red = y - ky; // in (-1/2, 1/2]

        let ky_parity_odd = ky.rem_euclid(&two) != I::F::zero();
        let x_oriented = if ky_parity_odd { -x } else { x };

        // S1::new performs the (-1/2,1/2] reduction itself now, so
        // x_oriented can be handed to it directly, unreduced.
        Self(
            S1::new([x_oriented].into()),
            S1::new([y_red].into()),
            PhantomData,
        )
    }

    fn coords(&self) -> (I::F, I::F) {
        (self.0.lift()[0], self.1.lift()[0])
    }
}

impl<I: ICompatible<V>, V: VCompatible<I>> Interval for KleinBottle<I, V> {
    type R = I::F;

    fn interval_squared(&self, other: &Self) -> I::F {
        self.to_local(other).unwrap().norm_squared().into()
    }
}

#[cfg(feature = "simplicial")]
mod simplicial {
    use super::*;
    use crate::{
        impl_tangent_bundle_via_bounded,
        traits::simplicial::{Bounded, BuildNodes, NerveComplexParameters},
    };
    use std::vec::Vec;
    /// A bounded chart domain in the regular finite cover of [`Torus`].
    #[derive(Debug, Copy, Clone)]
    pub struct TorusCover<I: ICompatible<V>, V: VCompatible<I>>(Torus<I, V>);

    impl<I: ICompatible<V>, V: VCompatible<I>> From<Torus<I, V>> for TorusCover<I, V> {
        fn from(value: Torus<I, V>) -> Self {
            Self(value)
        }
    }

    impl<I: ICompatible<V>, V: VCompatible<I>> AsRef<Torus<I, V>> for TorusCover<I, V> {
        fn as_ref(&self) -> &Torus<I, V> {
            &self.0
        }
    }

    const S: usize = 4;

    impl<I: ICompatible<V>, V: VCompatible<I>> Bounded<Torus<I, V>, Torus<I, V>, V>
        for TorusCover<I, V>
    {
        fn sdf(&self, v: &V) -> <V as Tensor>::F {
            let to = |x| <V::F as NumCast>::from(x).unwrap();
            v.norm() - (to(2).sqrt() + to(2)) / to(4 * S)
        }
    }

    impl_tangent_bundle_via_bounded!(
        TorusCover<I, V>,
        Torus<I, V>,
        Torus<I, V>,
        V,
        I: ICompatible<V>,
    V: VCompatible<I>
    );

    impl<I: ICompatible<V>, V: VCompatible<I>> BuildNodes<TorusCover<I, V>> for TorusCover<I, V> {
        fn build_nodes() -> Vec<Self> {
            let to = |x| <I::F as NumCast>::from(x).unwrap();
            let s = to(S);
            let offset = to(1) / (to(2) * s);

            (0..S)
                .flat_map(|y| (0..S).map(move |x| (x, y)))
                .map(|(x, y)| {
                    Torus::new(
                        S1([offset + to(x) / s].into()),
                        S1([offset + to(y) / s].into()),
                    )
                    .into()
                })
                .collect()
        }
    }

    impl<I: ICompatible<V>, V: VCompatible<I>>
        NerveComplexParameters<Torus<I, V>, V, Torus<I, V>, TorusCover<I, V>> for TorusCover<I, V>
    {
    }

    /// A bounded chart domain in the regular finite cover of [`KleinBottle`].
    #[derive(Debug, Copy, Clone)]
    pub struct KleinBottleCover<I: ICompatible<V>, V: VCompatible<I>>(KleinBottle<I, V>);

    impl<I: ICompatible<V>, V: VCompatible<I>> From<KleinBottle<I, V>> for KleinBottleCover<I, V> {
        fn from(value: KleinBottle<I, V>) -> Self {
            Self(value)
        }
    }

    impl<I: ICompatible<V>, V: VCompatible<I>> AsRef<KleinBottle<I, V>> for KleinBottleCover<I, V> {
        fn as_ref(&self) -> &KleinBottle<I, V> {
            &self.0
        }
    }

    impl<I: ICompatible<V>, V: VCompatible<I>> Bounded<KleinBottle<I, V>, KleinBottle<I, V>, V>
        for KleinBottleCover<I, V>
    {
        fn sdf(&self, v: &V) -> <V as Tensor>::F {
            let to = |x| <V::F as NumCast>::from(x).unwrap();
            v.norm() - (to(2).sqrt() + to(2)) / to(4 * S)
        }
    }

    impl_tangent_bundle_via_bounded!(
        KleinBottleCover<I, V>,
        KleinBottle<I, V>,
        KleinBottle<I, V>,
        V,
        I: ICompatible<V>, V: VCompatible<I>
    );

    impl<I: ICompatible<V>, V: VCompatible<I>> BuildNodes<KleinBottleCover<I, V>>
        for KleinBottleCover<I, V>
    {
        fn build_nodes() -> Vec<Self> {
            let to = |x| <I::F as NumCast>::from(x).unwrap();
            let s = to(S);
            let offset = to(1) / (to(2) * s);

            (0..S)
                .flat_map(|y| (0..S).map(move |x| (x, y)))
                .map(|(x, y)| {
                    KleinBottle::new(
                        S1([offset + to(x) / s].into()),
                        S1([offset + to(y) / s].into()),
                    )
                    .into()
                })
                .collect()
        }
    }

    impl<I: ICompatible<V>, V: VCompatible<I>>
        NerveComplexParameters<KleinBottle<I, V>, V, KleinBottle<I, V>, KleinBottleCover<I, V>>
        for KleinBottleCover<I, V>
    {
    }

    /// A deliberately overlapping bounded domain on [`Torus`].
    ///
    /// This exercises [`NerveComplex`](crate::traits::simplicial::NerveComplex)
    /// when many cover nodes see the same region rather than forming the regular
    /// cover represented by [`TorusCover`].
    #[derive(Debug, Clone)]
    pub struct MyopicTorus<I: ICompatible<V>, V: VCompatible<I>>(pub Torus<I, V>);

    impl<I: ICompatible<V>, V: VCompatible<I>> MyopicTorus<I, V> {
        /// Returns the number of cover samples along each torus coordinate.
        pub fn s() -> usize {
            8
        }

        fn radius() -> V::F {
            // 2/s, quite a lot larger than the lattice spacing.
            (V::F::one() + V::F::one()) / <V::F as NumCast>::from(Self::s()).unwrap()
        }
    }

    impl<I: ICompatible<V>, V: VCompatible<I>> AsRef<Torus<I, V>> for MyopicTorus<I, V> {
        fn as_ref(&self) -> &Torus<I, V> {
            &self.0
        }
    }

    impl<I: ICompatible<V>, V: VCompatible<I>> From<Torus<I, V>> for MyopicTorus<I, V> {
        fn from(value: Torus<I, V>) -> Self {
            Self(value)
        }
    }

    impl<I: ICompatible<V>, V: VCompatible<I>> Bounded<Torus<I, V>, Torus<I, V>, V>
        for MyopicTorus<I, V>
    {
        fn sdf(&self, v: &V) -> <V as Tensor>::F {
            v.norm() - Self::radius()
        }
    }

    impl_tangent_bundle_via_bounded!(
        MyopicTorus<I, V>,
        Torus<I, V>,
        Torus<I, V>,
        V,
        I: ICompatible<V>, V: VCompatible<I>
    );

    /// The tangent-bundle chart wrapper associated with [`MyopicTorus`].
    #[derive(Debug, Clone)]
    pub struct MyopicTorusCover<I: ICompatible<V>, V: VCompatible<I>>(MyopicTorus<I, V>);

    impl<I: ICompatible<V>, V: VCompatible<I>> AsRef<MyopicTorus<I, V>> for MyopicTorusCover<I, V> {
        fn as_ref(&self) -> &MyopicTorus<I, V> {
            &self.0
        }
    }

    impl<I: ICompatible<V>, V: VCompatible<I>> From<MyopicTorus<I, V>> for MyopicTorusCover<I, V> {
        fn from(value: MyopicTorus<I, V>) -> Self {
            Self(value)
        }
    }

    impl<I: ICompatible<V>, V: VCompatible<I>> Bounded<MyopicTorus<I, V>, Torus<I, V>, V>
        for MyopicTorusCover<I, V>
    {
        fn sdf(&self, v: &V) -> <V as Tensor>::F {
            let to = |x| <V::F as NumCast>::from(x).unwrap();
            v.norm() - (to(2).sqrt() + to(2)) / to(4 * MyopicTorus::<I, V>::s())
        }
    }

    impl_tangent_bundle_via_bounded!(
        MyopicTorusCover<I, V>,
        MyopicTorus<I, V>,
        Torus<I, V>,
        V,
        I: ICompatible<V>, V: VCompatible<I>
    );

    impl<I: ICompatible<V>, V: VCompatible<I>> BuildNodes<MyopicTorusCover<I, V>>
        for MyopicTorusCover<I, V>
    {
        fn build_nodes() -> Vec<Self> {
            let to = |x| <I::F as NumCast>::from(x).unwrap();
            let s_usize = MyopicTorus::<I, V>::s();
            let s = to(s_usize);
            let offset = to(1) / (to(2) * s);

            (0..s_usize)
                .flat_map(|y| (0..s_usize).map(move |x| (x, y)))
                .map(|(x, y)| {
                    MyopicTorus(Torus::new(
                        S1([offset + to(x) / s].into()),
                        S1([offset + to(y) / s].into()),
                    ))
                    .into()
                })
                .collect()
        }
    }

    impl<I: ICompatible<V>, V: VCompatible<I>>
        NerveComplexParameters<Torus<I, V>, V, MyopicTorus<I, V>, MyopicTorusCover<I, V>>
        for MyopicTorusCover<I, V>
    {
        fn overestimation_bound() -> Option<(V::F, V::F)> {
            let to = |x| <V::F as NumCast>::from(x).unwrap();
            let s = to(MyopicTorus::<I, V>::s());
            // κ: king-graph worst case at 22.5°, √(4 − 2√2). Scale-free.
            let kappa = (to(4) - to(2) * to(2).sqrt()).sqrt();
            // C = (1+κ)·2δ_s, with δ_s = √2/(2S) the lattice half-diagonal.
            let delta_s = to(2).sqrt() / (to(2) * s);
            Some((kappa, (V::F::one() + kappa) * to(2) * delta_s))
        }
    }
}

#[cfg(feature = "simplicial")]
pub use simplicial::*;