//! Linear operator algebra

use crate::generate::hstack;

use crate::types::*;

use ndarray::*;

/// Abstracted linear operator as an action to vector (`ArrayBase<S, Ix1>`) and matrix

/// (`ArrayBase<S, Ix2`)

pub trait LinearOperator {

type Elem: Scalar;

/// Apply operator out-place

fn apply<S>(&self, a: &ArrayBase<S, Ix1>) -> Array1<S::Elem>

where

S: Data<Elem = Self::Elem>,

{

let mut a = a.to_owned();

self.apply_mut(&mut a);

a

}

/// Apply operator in-place

fn apply_mut<S>(&self, a: &mut ArrayBase<S, Ix1>)

where

S: DataMut<Elem = Self::Elem>,

{

let b = self.apply(a);

azip!((a in a, &b in &b) *a = b);

}

/// Apply operator with move

fn apply_into<S>(&self, mut a: ArrayBase<S, Ix1>) -> ArrayBase<S, Ix1>

where

S: DataOwned<Elem = Self::Elem> + DataMut,

{

self.apply_mut(&mut a);

a

}

/// Apply operator to matrix out-place

fn apply2<S>(&self, a: &ArrayBase<S, Ix2>) -> Array2<S::Elem>

where

S: Data<Elem = Self::Elem>,

{

let cols: Vec<_> = a.axis_iter(Axis(1)).map(|col| self.apply(&col)).collect();

hstack(&cols).unwrap()

}

/// Apply operator to matrix in-place

fn apply2_mut<S>(&self, a: &mut ArrayBase<S, Ix2>)

where

S: DataMut<Elem = Self::Elem>,

{

for mut col in a.axis_iter_mut(Axis(1)) {

self.apply_mut(&mut col)

}

}

/// Apply operator to matrix with move

fn apply2_into<S>(&self, mut a: ArrayBase<S, Ix2>) -> ArrayBase<S, Ix2>

where

S: DataOwned<Elem = Self::Elem> + DataMut,

{

self.apply2_mut(&mut a);

a

}

}

impl<A, Sa> LinearOperator for ArrayBase<Sa, Ix2>

where

A: Scalar,

Sa: Data<Elem = A>,

{

type Elem = A;

fn apply<S>(&self, a: &ArrayBase<S, Ix1>) -> Array1<A>

where

S: Data<Elem = A>,

{

self.dot(a)

}

}