src/loc.rs

/*!
Array containers that support vector space operations on floats.
For working with small vectors in $ℝ^2$ or $ℝ^3$.
*/

use crate::euclidean::*;
use crate::instance::{BasicDecomposition, Instance};
use crate::linops::{Linear, Mapping, AXPY};
use crate::mapping::Space;
use crate::maputil::{map1, map1_mut, map2, map2_mut, FixedLength, FixedLengthMut};
use crate::norms::*;
use crate::types::{Float, Num, SignedNum};
use serde::ser::{Serialize, SerializeSeq, Serializer};
use std::fmt::{Display, Formatter};
use std::ops::{
    Add, AddAssign, Div, DivAssign, Index, IndexMut, Mul, MulAssign, Neg, Sub, SubAssign,
};
use std::slice::{Iter, IterMut};

/// A container type for (short) `N`-dimensional vectors of element type `F`.
///
/// Supports basic operations of an [`Euclidean`] space, several [`Norm`]s, and
/// fused [`AXPY`] operations, among others.
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub struct Loc<const N: usize, F = f64>(
    /// An array of the elements of the vector
    pub [F; N],
);

impl<F: Display, const N: usize> Display for Loc<N, F> {
    // Required method
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "[")?;
        let mut comma = "";
        for e in self.iter() {
            write!(f, "{comma}{e}")?;
            comma = ", ";
        }
        write!(f, "]")
    }
}

// Need to manually implement as [F; N] serialisation is provided only for some N.
impl<F, const N: usize> Serialize for Loc<N, F>
where
    F: Serialize,
{
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let mut seq = serializer.serialize_seq(Some(N))?;
        for e in self.iter() {
            seq.serialize_element(e)?;
        }
        seq.end()
    }
}

impl<F, const N: usize> Loc<N, F> {
    /// Creates a new `Loc` vector from an array.
    #[inline]
    pub fn new(arr: [F; N]) -> Self {
        Loc(arr)
    }

    /// Returns an iterator over the elements of the vector
    #[inline]
    pub fn iter(&self) -> Iter<'_, F> {
        self.0.iter()
    }

    /// Returns an iterator over mutable references to the elements of the vector
    #[inline]
    pub fn iter_mut(&mut self) -> IterMut<'_, F> {
        self.0.iter_mut()
    }
}

impl<F: Copy, const N: usize> Loc<N, F> {
    /// Maps `g` over the elements of the vector, returning a new [`Loc`] vector
    #[inline]
    pub fn map<H>(&self, g: impl Fn(F) -> H) -> Loc<N, H> {
        Loc::new(map1(self, |u| g(*u)))
    }

    /// Maps `g` over pairs of elements of two vectors, retuning a new one.
    #[inline]
    pub fn map2<H>(&self, other: &Self, g: impl Fn(F, F) -> H) -> Loc<N, H> {
        Loc::new(map2(self, other, |u, v| g(*u, *v)))
    }

    /// Maps `g` over mutable references to elements of the vector.
    #[inline]
    pub fn map_mut(&mut self, g: impl Fn(&mut F)) {
        map1_mut(self, g)
    }

    /// Maps `g` over pairs of mutable references to elements of `self, and elements
    /// of `other` vector.
    #[inline]
    pub fn map2_mut(&mut self, other: &Self, g: impl Fn(&mut F, F)) {
        map2_mut(self, other, |u, v| g(u, *v))
    }

    /// Maps `g` over the elements of `self` and returns the product of the results.
    #[inline]
    pub fn product_map<A: Num>(&self, g: impl Fn(F) -> A) -> A {
        match N {
            1 => g(unsafe { *self.0.get_unchecked(0) }),
            2 => g(unsafe { *self.0.get_unchecked(0) }) * g(unsafe { *self.0.get_unchecked(1) }),
            3 => {
                g(unsafe { *self.0.get_unchecked(0) })
                    * g(unsafe { *self.0.get_unchecked(1) })
                    * g(unsafe { *self.0.get_unchecked(2) })
            }
            _ => self.iter().fold(A::ONE, |m, &x| m * g(x)),
        }
    }
}

/// Construct a [`Loc`].
///
/// Use as
/// ```
/// # use alg_tools::loc::Loc;
/// # use alg_tools::loc;
/// let x = loc![1.0, 2.0];
/// ```
#[macro_export]
macro_rules! loc {
    ($($x:expr),+ $(,)?) => { Loc::new([$($x),+]) }
}

impl<F, const N: usize> From<[F; N]> for Loc<N, F> {
    #[inline]
    fn from(other: [F; N]) -> Loc<N, F> {
        Loc(other)
    }
}

/*impl<F : Copy, const N : usize> From<&[F; N]> for Loc<N, F> {
    #[inline]
    fn from(other: &[F; N]) -> Loc<N, F> {
        Loc(*other)
    }
}*/

impl<F> From<F> for Loc<1, F> {
    #[inline]
    fn from(other: F) -> Loc<1, F> {
        Loc([other])
    }
}

impl<F> Loc<1, F> {
    #[inline]
    pub fn flatten1d(self) -> F {
        let Loc([v]) = self;
        v
    }
}

impl<F, const N: usize> From<Loc<N, F>> for [F; N] {
    #[inline]
    fn from(other: Loc<N, F>) -> [F; N] {
        other.0
    }
}

/*impl<F : Copy, const N : usize> From<&Loc<N, F>> for [F; N] {
    #[inline]
    fn from(other : &Loc<N, F>) -> [F; N] {
        other.0
    }
}*/

impl<F, const N: usize> IntoIterator for Loc<N, F> {
    type Item = <[F; N] as IntoIterator>::Item;
    type IntoIter = <[F; N] as IntoIterator>::IntoIter;

    #[inline]
    fn into_iter(self) -> Self::IntoIter {
        self.0.into_iter()
    }
}

// Indexing

impl<F, Ix, const N: usize> Index<Ix> for Loc<N, F>
where
    [F; N]: Index<Ix>,
{
    type Output = <[F; N] as Index<Ix>>::Output;

    #[inline]
    fn index(&self, ix: Ix) -> &Self::Output {
        self.0.index(ix)
    }
}

impl<F, Ix, const N: usize> IndexMut<Ix> for Loc<N, F>
where
    [F; N]: IndexMut<Ix>,
{
    #[inline]
    fn index_mut(&mut self, ix: Ix) -> &mut Self::Output {
        self.0.index_mut(ix)
    }
}

// Arithmetic

macro_rules! make_binop {
    ($trait:ident, $fn:ident, $trait_assign:ident, $fn_assign:ident) => {
        impl<F: Num, const N: usize> $trait<Loc<N, F>> for Loc<N, F> {
            type Output = Loc<N, F>;
            #[inline]
            fn $fn(mut self, other: Loc<N, F>) -> Self::Output {
                self.$fn_assign(other);
                self
            }
        }

        impl<'a, F: Num, const N: usize> $trait<&'a Loc<N, F>> for Loc<N, F> {
            type Output = Loc<N, F>;
            #[inline]
            fn $fn(mut self, other: &'a Loc<N, F>) -> Self::Output {
                self.$fn_assign(other);
                self
            }
        }

        impl<'b, F: Num, const N: usize> $trait<Loc<N, F>> for &'b Loc<N, F> {
            type Output = Loc<N, F>;
            #[inline]
            fn $fn(self, other: Loc<N, F>) -> Self::Output {
                self.map2(&other, |a, b| a.$fn(b))
            }
        }

        impl<'a, 'b, F: Num, const N: usize> $trait<&'a Loc<N, F>> for &'b Loc<N, F> {
            type Output = Loc<N, F>;
            #[inline]
            fn $fn(self, other: &'a Loc<N, F>) -> Self::Output {
                self.map2(other, |a, b| a.$fn(b))
            }
        }

        impl<F: Num, const N: usize> $trait_assign<Loc<N, F>> for Loc<N, F> {
            #[inline]
            fn $fn_assign(&mut self, other: Loc<N, F>) {
                self.map2_mut(&other, |a, b| a.$fn_assign(b))
            }
        }

        impl<'a, F: Num, const N: usize> $trait_assign<&'a Loc<N, F>> for Loc<N, F> {
            #[inline]
            fn $fn_assign(&mut self, other: &'a Loc<N, F>) {
                self.map2_mut(other, |a, b| a.$fn_assign(b))
            }
        }
    };
}

make_binop!(Add, add, AddAssign, add_assign);
make_binop!(Sub, sub, SubAssign, sub_assign);

impl<F: Float, const N: usize> std::iter::Sum for Loc<N, F> {
    fn sum<I: Iterator<Item = Loc<N, F>>>(mut iter: I) -> Self {
        match iter.next() {
            None => Self::ORIGIN,
            Some(mut v) => {
                for w in iter {
                    v += w
                }
                v
            }
        }
    }
}

macro_rules! make_scalarop_rhs {
    ($trait:ident, $fn:ident, $trait_assign:ident, $fn_assign:ident) => {
        impl<F: Num, const N: usize> $trait<F> for Loc<N, F> {
            type Output = Loc<N, F>;
            #[inline]
            fn $fn(self, b: F) -> Self::Output {
                self.map(|a| a.$fn(b))
            }
        }

        impl<'a, F: Num, const N: usize> $trait<&'a F> for Loc<N, F> {
            type Output = Loc<N, F>;
            #[inline]
            fn $fn(self, b: &'a F) -> Self::Output {
                self.map(|a| a.$fn(*b))
            }
        }

        impl<'b, F: Num, const N: usize> $trait<F> for &'b Loc<N, F> {
            type Output = Loc<N, F>;
            #[inline]
            fn $fn(self, b: F) -> Self::Output {
                self.map(|a| a.$fn(b))
            }
        }

        impl<'a, 'b, F: Float, const N: usize> $trait<&'a F> for &'b Loc<N, F> {
            type Output = Loc<N, F>;
            #[inline]
            fn $fn(self, b: &'a F) -> Self::Output {
                self.map(|a| a.$fn(*b))
            }
        }

        impl<F: Num, const N: usize> $trait_assign<F> for Loc<N, F> {
            #[inline]
            fn $fn_assign(&mut self, b: F) {
                self.map_mut(|a| a.$fn_assign(b));
            }
        }

        impl<'a, F: Num, const N: usize> $trait_assign<&'a F> for Loc<N, F> {
            #[inline]
            fn $fn_assign(&mut self, b: &'a F) {
                self.map_mut(|a| a.$fn_assign(*b));
            }
        }
    };
}

make_scalarop_rhs!(Mul, mul, MulAssign, mul_assign);
make_scalarop_rhs!(Div, div, DivAssign, div_assign);

macro_rules! make_unaryop {
    ($trait:ident, $fn:ident) => {
        impl<F: SignedNum, const N: usize> $trait for Loc<N, F> {
            type Output = Loc<N, F>;
            #[inline]
            fn $fn(mut self) -> Self::Output {
                self.map_mut(|a| *a = (*a).$fn());
                self
            }
        }

        impl<'a, F: SignedNum, const N: usize> $trait for &'a Loc<N, F> {
            type Output = Loc<N, F>;
            #[inline]
            fn $fn(self) -> Self::Output {
                self.map(|a| a.$fn())
            }
        }
    };
}

make_unaryop!(Neg, neg);

macro_rules! make_scalarop_lhs {
    ($trait:ident, $fn:ident; $($f:ident)+) => { $(
        impl<const N : usize> $trait<Loc<N, $f>> for $f {
            type Output = Loc<N, $f>;
            #[inline]
            fn $fn(self, v : Loc<N, $f>) -> Self::Output {
                v.map(|b| self.$fn(b))
            }
        }

        impl<'a, const N : usize> $trait<&'a Loc<N, $f>> for $f {
            type Output = Loc<N, $f>;
            #[inline]
            fn $fn(self, v : &'a Loc<N, $f>) -> Self::Output {
                v.map(|b| self.$fn(b))
            }
        }

        impl<'b, const N : usize> $trait<Loc<N, $f>> for &'b $f {
            type Output = Loc<N, $f>;
            #[inline]
            fn $fn(self, v : Loc<N, $f>) -> Self::Output {
                v.map(|b| self.$fn(b))
            }
        }

        impl<'a, 'b, const N : usize> $trait<&'a Loc<N, $f>> for &'b $f {
            type Output = Loc<N, $f>;
            #[inline]
            fn $fn(self, v : &'a Loc<N, $f>) -> Self::Output {
                v.map(|b| self.$fn(b))
            }
        }
    )+ }
}

make_scalarop_lhs!(Mul, mul; f32 f64 i8 i16 i32 i64 isize u8 u16 u32 u64 usize);
make_scalarop_lhs!(Div, div; f32 f64 i8 i16 i32 i64 isize u8 u16 u32 u64 usize);

// Norms

macro_rules! domination {
    ($norm:ident, $dominates:ident) => {
        impl<F: Float, const N: usize> Dominated<F, $dominates, Loc<N, F>> for $norm {
            #[inline]
            fn norm_factor(&self, _p: $dominates) -> F {
                F::ONE
            }
            #[inline]
            fn from_norm(&self, p_norm: F, _p: $dominates) -> F {
                p_norm
            }
        }
    };
    ($norm:ident, $dominates:ident, $fn:path) => {
        impl<F: Float, const N: usize> Dominated<F, $dominates, Loc<N, F>> for $norm {
            #[inline]
            fn norm_factor(&self, _p: $dominates) -> F {
                $fn(F::cast_from(N))
            }
        }
    };
}

domination!(L1, L1);
domination!(L2, L2);
domination!(Linfinity, Linfinity);

domination!(L1, L2, F::sqrt);
domination!(L2, Linfinity, F::sqrt);
domination!(L1, Linfinity, std::convert::identity);

domination!(Linfinity, L1);
domination!(Linfinity, L2);
domination!(L2, L1);

impl<F: Float, const N: usize> Euclidean<F> for Loc<N, F> {
    /// This implementation is not stabilised as it's meant to be used for very small vectors.
    /// Use [`nalgebra`] for larger vectors.
    #[inline]
    fn dot<I: Instance<Self>>(&self, other: I) -> F {
        other.eval_ref_decompose(|r| {
            self.0
                .iter()
                .zip(r.0.iter())
                .fold(F::ZERO, |m, (&v, &w)| m + v * w)
        })
    }

    /// This implementation is not stabilised as it's meant to be used for very small vectors.
    /// Use [`nalgebra`] for larger vectors.
    #[inline]
    fn norm2_squared(&self) -> F {
        self.iter().fold(F::ZERO, |m, &v| m + v * v)
    }

    fn dist2_squared<I: Instance<Self>>(&self, other: I) -> F {
        other.eval_ref_decompose(|r| {
            self.iter().zip(r.iter()).fold(F::ZERO, |m, (&v, &w)| {
                let d = v - w;
                m + d * d
            })
        })
    }

    #[inline]
    fn norm2(&self) -> F {
        // Optimisation for N==1 that avoids squaring and square rooting.
        if N == 1 {
            unsafe { self.0.get_unchecked(0) }.abs()
        } else {
            self.norm2_squared().sqrt()
        }
    }

    #[inline]
    fn dist2<I: Instance<Self>>(&self, other: I) -> F {
        // Optimisation for N==1 that avoids squaring and square rooting.
        if N == 1 {
            other.eval_ref_decompose(|r| {
                unsafe { *self.0.get_unchecked(0) - *r.0.get_unchecked(0) }.abs()
            })
        } else {
            self.dist2_squared(other).sqrt()
        }
    }
}

impl<F: Num, const N: usize> Loc<N, F> {
    /// Origin point
    pub const ORIGIN: Self = Loc([F::ZERO; N]);
}

impl<F: Num + std::ops::Neg<Output = F>, const N: usize> Loc<N, F> {
    /// Reflects along the given coordinate
    pub fn reflect(mut self, i: usize) -> Self {
        self[i] = -self[i];
        self
    }

    /// Reflects along the given coordinate sequences
    pub fn reflect_many<I: IntoIterator<Item = usize>>(mut self, idxs: I) -> Self {
        for i in idxs {
            self[i] = -self[i]
        }
        self
    }
}

impl<F: std::ops::Neg<Output = F>> Loc<2, F> {
    /// Reflect x coordinate
    pub fn reflect_x(self) -> Self {
        let Loc([x, y]) = self;
        [-x, y].into()
    }

    /// Reflect y coordinate
    pub fn reflect_y(self) -> Self {
        let Loc([x, y]) = self;
        [x, -y].into()
    }
}

impl<F: Float> Loc<2, F> {
    /// Rotate by angle φ
    pub fn rotate(self, φ: F) -> Self {
        let Loc([x, y]) = self;
        let sin_φ = φ.sin();
        let cos_φ = φ.cos();
        [cos_φ * x - sin_φ * y, sin_φ * x + cos_φ * y].into()
    }
}

impl<F: std::ops::Neg<Output = F>> Loc<3, F> {
    /// Reflect x coordinate
    pub fn reflect_x(self) -> Self {
        let Loc([x, y, z]) = self;
        [-x, y, z].into()
    }

    /// Reflect y coordinate
    pub fn reflect_y(self) -> Self {
        let Loc([x, y, z]) = self;
        [x, -y, z].into()
    }

    /// Reflect y coordinate
    pub fn reflect_z(self) -> Self {
        let Loc([x, y, z]) = self;
        [x, y, -z].into()
    }
}

impl<F: Float> Loc<3, F> {
    /// Rotate by angles (yaw, pitch, roll)
    pub fn rotate(self, Loc([φ, ψ, θ]): Self) -> Self {
        let Loc([mut x, mut y, mut z]) = self;
        let sin_φ = φ.sin();
        let cos_φ = φ.cos();
        [x, y, z] = [cos_φ * x - sin_φ * y, sin_φ * x + cos_φ * y, z];
        let sin_ψ = ψ.sin();
        let cos_ψ = ψ.cos();
        [x, y, z] = [cos_ψ * x + sin_ψ * z, y, -sin_ψ * x + cos_ψ * z];
        let sin_θ = θ.sin();
        let cos_θ = θ.cos();
        [x, y, z] = [x, cos_θ * y - sin_θ * z, sin_θ * y + cos_θ * z];
        [x, y, z].into()
    }
}

impl<F: Float, const N: usize> StaticEuclidean<F> for Loc<N, F> {
    #[inline]
    fn origin() -> Self {
        Self::ORIGIN
    }
}

/// The default norm for `Loc` is [`L2`].
impl<F: Float, const N: usize> Normed<F> for Loc<N, F> {
    type NormExp = L2;

    #[inline]
    fn norm_exponent(&self) -> Self::NormExp {
        L2
    }

    // #[inline]
    // fn similar_origin(&self) -> Self {
    //     [F::ZERO; N].into()
    // }

    #[inline]
    fn is_zero(&self) -> bool {
        self.norm2_squared() == F::ZERO
    }
}

impl<F: Float, const N: usize> HasDual<F> for Loc<N, F> {
    type DualSpace = Self;

    fn dual_origin(&self) -> Self::DualSpace {
        self.similar_origin()
    }
}

impl<F: Float, const N: usize> Norm<L2, F> for Loc<N, F> {
    #[inline]
    fn norm(&self, _: L2) -> F {
        self.norm2()
    }
}

impl<F: Float, const N: usize> Dist<L2, F> for Loc<N, F> {
    #[inline]
    fn dist<I: Instance<Self>>(&self, other: I, _: L2) -> F {
        self.dist2(other)
    }
}

/* Implemented automatically as Euclidean.
impl<F : Float, const N : usize> Projection<F, L2> for Loc<N, F> {
    #[inline]
    fn proj_ball_mut(&mut self, ρ : F, nrm : L2) {
        let n = self.norm(nrm);
        if n > ρ {
            *self *= ρ/n;
        }
    }
}*/

impl<F: Float, const N: usize> Norm<L1, F> for Loc<N, F> {
    /// This implementation is not stabilised as it's meant to be used for very small vectors.
    /// Use [`nalgebra`] for larger vectors.
    #[inline]
    fn norm(&self, _: L1) -> F {
        self.iter().fold(F::ZERO, |m, v| m + v.abs())
    }
}

impl<F: Float, const N: usize> Dist<L1, F> for Loc<N, F> {
    #[inline]
    fn dist<I: Instance<Self>>(&self, other: I, _: L1) -> F {
        other.eval_ref_decompose(|r| {
            self.iter()
                .zip(r.iter())
                .fold(F::ZERO, |m, (&v, &w)| m + (v - w).abs())
        })
    }
}

impl<F: Float, const N: usize> Projection<F, Linfinity> for Loc<N, F> {
    #[inline]
    fn proj_ball_mut(&mut self, ρ: F, _: Linfinity) {
        self.iter_mut()
            .for_each(|v| *v = num_traits::clamp(*v, -ρ, ρ))
    }
}

impl<F: Float, const N: usize> Norm<Linfinity, F> for Loc<N, F> {
    /// This implementation is not stabilised as it's meant to be used for very small vectors.
    /// Use [`nalgebra`] for larger vectors.
    #[inline]
    fn norm(&self, _: Linfinity) -> F {
        self.iter().fold(F::ZERO, |m, v| m.max(v.abs()))
    }
}

impl<F: Float, const N: usize> Dist<Linfinity, F> for Loc<N, F> {
    #[inline]
    fn dist<I: Instance<Self>>(&self, other: I, _: Linfinity) -> F {
        other.eval_ref_decompose(|r| {
            self.iter()
                .zip(r.iter())
                .fold(F::ZERO, |m, (&v, &w)| m.max((v - w).abs()))
        })
    }
}

// Misc.

impl<A, const N: usize> FixedLength<N> for Loc<N, A> {
    type Iter = std::array::IntoIter<A, N>;
    type Elem = A;
    #[inline]
    fn fl_iter(self) -> Self::Iter {
        self.into_iter()
    }
}

impl<A, const N: usize> FixedLengthMut<N> for Loc<N, A> {
    type IterMut<'a>
        = std::slice::IterMut<'a, A>
    where
        A: 'a;
    #[inline]
    fn fl_iter_mut(&mut self) -> Self::IterMut<'_> {
        self.iter_mut()
    }
}

impl<'a, A, const N: usize> FixedLength<N> for &'a Loc<N, A> {
    type Iter = std::slice::Iter<'a, A>;
    type Elem = &'a A;
    #[inline]
    fn fl_iter(self) -> Self::Iter {
        self.iter()
    }
}

impl<F: Num, const N: usize> Space for Loc<N, F> {
    type Decomp = BasicDecomposition;
}

impl<F: Float, const N: usize> Mapping<Loc<N, F>> for Loc<N, F> {
    type Codomain = F;

    fn apply<I: Instance<Loc<N, F>>>(&self, x: I) -> Self::Codomain {
        x.eval_decompose(|x̃| self.dot(x̃))
    }
}

impl<F: Float, const N: usize> Linear<Loc<N, F>> for Loc<N, F> {}

impl<F: Float, const N: usize> AXPY<Loc<N, F>> for Loc<N, F> {
    type Field = F;
    type Owned = Self;

    #[inline]
    fn axpy<I: Instance<Loc<N, F>>>(&mut self, α: F, x: I, β: F) {
        x.eval(|x̃| {
            if β == F::ZERO {
                map2_mut(self, x̃, |yi, xi| *yi = α * (*xi))
            } else {
                map2_mut(self, x̃, |yi, xi| *yi = β * (*yi) + α * (*xi))
            }
        })
    }

    #[inline]
    fn copy_from<I: Instance<Loc<N, F>>>(&mut self, x: I) {
        x.eval(|x̃| map2_mut(self, x̃, |yi, xi| *yi = *xi))
    }

    // #[inline]
    // fn make_origin_generator(&self) -> StaticEuclideanOriginGenerator {
    //     StaticEuclideanOriginGenerator
    // }

    #[inline]
    fn similar_origin(&self) -> Self::Owned {
        Self::ORIGIN
    }

    #[inline]
    fn similar_origin_inst<I: Instance<Self>>(_: I) -> Self::Owned {
        Self::ORIGIN
    }

    #[inline]
    fn set_zero(&mut self) {
        *self = Self::ORIGIN;
    }
}