src/loc.rs

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

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


/// 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<F, const N : usize>(
    /// An array of the elements of the vector
    pub [F; N]
);

impl<F : Display, const N : usize> Display for Loc<F, N>{
    // 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<F, N>
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<F, N> {
    /// 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<F, N> {
    /// 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<H, N> {
        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<H, N> {
        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<F, N> {
    #[inline]
    fn from(other: [F; N]) -> Loc<F, N> {
        Loc(other)
    }
}

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

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

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

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

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


impl<F, const N : usize> IntoIterator for Loc<F, N> {
    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<F,N>
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<F,N>
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<F,N>> for Loc<F, N> {
            type Output = Loc<F, N>;
            #[inline]
            fn $fn(mut self, other : Loc<F, N>) -> Self::Output {
                self.$fn_assign(other);
                self
            }
        }

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

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

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

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

        impl<'a, F : Num, const N : usize> $trait_assign<&'a Loc<F,N>> for Loc<F, N> {
            #[inline]
            fn $fn_assign(&mut self, other : &'a Loc<F, N>) {
                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<F, N> {
    fn sum<I: Iterator<Item = Loc<F, N>>>(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<F, N> {
            type Output = Loc<F, N>;
            #[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<F, N> {
            type Output = Loc<F, N>;
            #[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<F, N> {
            type Output = Loc<F, N>;
            #[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<F, N> {
            type Output = Loc<F, N>;
            #[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<F, N> {
            #[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<F, N> {
            #[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<F, N> {
            type Output = Loc<F, N>;
            #[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<F, N> {
            type Output = Loc<F, N>;
            #[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<$f,N>> for $f {
            type Output = Loc<$f, N>;
            #[inline]
            fn $fn(self, v : Loc<$f,N>) -> Self::Output {
                v.map(|b| self.$fn(b))
            }
        }

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

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

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

    /// 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 {
        self.0.iter()
              .zip(other.ref_instance().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 {
        self.iter()
            .zip(other.ref_instance().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.
        let otherr = other.ref_instance();
        if N==1 {
            unsafe { *self.0.get_unchecked(0) - *otherr.0.get_unchecked(0) }.abs()
        } else {
            self.dist2_squared(other).sqrt()
        }
    }
}

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

impl<F : Num + std::ops::Neg<Output=F>, const N : usize> Loc<F, N> {
    /// 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<F, 2> {
    /// 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<F, 2> {
    /// 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<F, 3> {
    /// 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<F, 3> {
    /// 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<F, N> {
    #[inline]
    fn origin() -> Self {
        Self::ORIGIN
    }
}


/// The default norm for `Loc` is [`L2`].
impl<F : Float, const N : usize> Normed<F> for Loc<F, N> {
    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<F, N> {
    type DualSpace = Self;
}

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

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

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


// Misc.

impl<A, const N : usize> FixedLength<N> for Loc<A,N> {
    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<A,N> {
    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<A,N> {
    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<F, N> {
    type Decomp = BasicDecomposition;
}

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

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

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

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

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

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

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