use lax::*;
use ndarray::*;
use crate::error::*;
use crate::layout::*;
use crate::types::*;
#[derive(Debug, Clone)]
pub struct LeastSquaresResult<E: Scalar, I: Dimension> {
pub singular_values: Array1<E::Real>,
pub solution: Array<E, I>,
pub rank: i32,
pub residual_sum_of_squares: Option<Array<E::Real, I::Smaller>>,
}
pub trait LeastSquaresSvd<D, E, I>
where
D: Data<Elem = E>,
E: Scalar + Lapack,
I: Dimension,
{
fn least_squares(&self, rhs: &ArrayBase<D, I>) -> Result<LeastSquaresResult<E, I>>;
}
pub trait LeastSquaresSvdInto<D, E, I>
where
D: Data<Elem = E>,
E: Scalar + Lapack,
I: Dimension,
{
fn least_squares_into(self, rhs: ArrayBase<D, I>) -> Result<LeastSquaresResult<E, I>>;
}
pub trait LeastSquaresSvdInPlace<D, E, I>
where
D: Data<Elem = E>,
E: Scalar + Lapack,
I: Dimension,
{
fn least_squares_in_place(
&mut self,
rhs: &mut ArrayBase<D, I>,
) -> Result<LeastSquaresResult<E, I>>;
}
impl<E, D1, D2> LeastSquaresSvd<D2, E, Ix1> for ArrayBase<D1, Ix2>
where
E: Scalar + Lapack,
D1: Data<Elem = E>,
D2: Data<Elem = E>,
{
fn least_squares(&self, rhs: &ArrayBase<D2, Ix1>) -> Result<LeastSquaresResult<E, Ix1>> {
let a = self.to_owned();
let b = rhs.to_owned();
a.least_squares_into(b)
}
}
impl<E, D1, D2> LeastSquaresSvd<D2, E, Ix2> for ArrayBase<D1, Ix2>
where
E: Scalar + Lapack,
D1: Data<Elem = E>,
D2: Data<Elem = E>,
{
fn least_squares(&self, rhs: &ArrayBase<D2, Ix2>) -> Result<LeastSquaresResult<E, Ix2>> {
let a = self.to_owned();
let b = rhs.to_owned();
a.least_squares_into(b)
}
}
impl<E, D1, D2> LeastSquaresSvdInto<D2, E, Ix1> for ArrayBase<D1, Ix2>
where
E: Scalar + Lapack,
D1: DataMut<Elem = E>,
D2: DataMut<Elem = E>,
{
fn least_squares_into(
mut self,
mut rhs: ArrayBase<D2, Ix1>,
) -> Result<LeastSquaresResult<E, Ix1>> {
self.least_squares_in_place(&mut rhs)
}
}
impl<E, D1, D2> LeastSquaresSvdInto<D2, E, Ix2> for ArrayBase<D1, Ix2>
where
E: Scalar + Lapack,
D1: DataMut<Elem = E>,
D2: DataMut<Elem = E>,
{
fn least_squares_into(
mut self,
mut rhs: ArrayBase<D2, Ix2>,
) -> Result<LeastSquaresResult<E, Ix2>> {
self.least_squares_in_place(&mut rhs)
}
}
impl<E, D1, D2> LeastSquaresSvdInPlace<D2, E, Ix1> for ArrayBase<D1, Ix2>
where
E: Scalar + Lapack,
D1: DataMut<Elem = E>,
D2: DataMut<Elem = E>,
{
fn least_squares_in_place(
&mut self,
rhs: &mut ArrayBase<D2, Ix1>,
) -> Result<LeastSquaresResult<E, Ix1>> {
if self.shape()[0] != rhs.shape()[0] {
return Err(ShapeError::from_kind(ErrorKind::IncompatibleShape).into());
}
let (m, n) = (self.shape()[0], self.shape()[1]);
if n > m {
let mut new_rhs = Array1::<E>::zeros((n,));
new_rhs.slice_mut(s![0..m]).assign(rhs);
compute_least_squares_srhs(self, &mut new_rhs)
} else {
compute_least_squares_srhs(self, rhs)
}
}
}
fn compute_least_squares_srhs<E, D1, D2>(
a: &mut ArrayBase<D1, Ix2>,
rhs: &mut ArrayBase<D2, Ix1>,
) -> Result<LeastSquaresResult<E, Ix1>>
where
E: Scalar + Lapack,
D1: DataMut<Elem = E>,
D2: DataMut<Elem = E>,
{
let LeastSquaresOwned::<E> {
singular_values,
rank,
} = E::least_squares(
a.layout()?,
a.as_allocated_mut()?,
rhs.as_slice_memory_order_mut()
.ok_or(LinalgError::MemoryNotCont)?,
)?;
let (m, n) = (a.shape()[0], a.shape()[1]);
let solution = rhs.slice(s![0..n]).to_owned();
let residual_sum_of_squares = compute_residual_scalar(m, n, rank, rhs);
Ok(LeastSquaresResult {
solution,
singular_values: Array::from_shape_vec((singular_values.len(),), singular_values)?,
rank,
residual_sum_of_squares,
})
}
fn compute_residual_scalar<E: Scalar, D: Data<Elem = E>>(
m: usize,
n: usize,
rank: i32,
b: &ArrayBase<D, Ix1>,
) -> Option<Array<E::Real, Ix0>> {
if m < n || n != rank as usize {
return None;
}
let mut arr: Array<E::Real, Ix0> = Array::zeros(());
arr[()] = b.slice(s![n..]).mapv(|x| x.powi(2).abs()).sum();
Some(arr)
}
impl<E, D1, D2> LeastSquaresSvdInPlace<D2, E, Ix2> for ArrayBase<D1, Ix2>
where
E: Scalar + Lapack,
D1: DataMut<Elem = E>,
D2: DataMut<Elem = E>,
{
fn least_squares_in_place(
&mut self,
rhs: &mut ArrayBase<D2, Ix2>,
) -> Result<LeastSquaresResult<E, Ix2>> {
if self.shape()[0] != rhs.shape()[0] {
return Err(ShapeError::from_kind(ErrorKind::IncompatibleShape).into());
}
let (m, n) = (self.shape()[0], self.shape()[1]);
if n > m {
let k = rhs.shape()[1];
let mut new_rhs = match self.layout()? {
MatrixLayout::C { .. } => Array2::<E>::zeros((n, k)),
MatrixLayout::F { .. } => Array2::<E>::zeros((n, k).f()),
};
new_rhs.slice_mut(s![0..m, ..]).assign(rhs);
compute_least_squares_nrhs(self, &mut new_rhs)
} else {
compute_least_squares_nrhs(self, rhs)
}
}
}
fn compute_least_squares_nrhs<E, D1, D2>(
a: &mut ArrayBase<D1, Ix2>,
rhs: &mut ArrayBase<D2, Ix2>,
) -> Result<LeastSquaresResult<E, Ix2>>
where
E: Scalar + Lapack,
D1: DataMut<Elem = E>,
D2: DataMut<Elem = E>,
{
let a_layout = a.layout()?;
let rhs_layout = rhs.layout()?;
let LeastSquaresOwned::<E> {
singular_values,
rank,
} = E::least_squares_nrhs(
a_layout,
a.as_allocated_mut()?,
rhs_layout,
rhs.as_allocated_mut()?,
)?;
let solution: Array2<E> = rhs.slice(s![..a.shape()[1], ..]).to_owned();
let singular_values = Array::from_shape_vec((singular_values.len(),), singular_values)?;
let (m, n) = (a.shape()[0], a.shape()[1]);
let residual_sum_of_squares = compute_residual_array1(m, n, rank, rhs);
Ok(LeastSquaresResult {
solution,
singular_values,
rank,
residual_sum_of_squares,
})
}
fn compute_residual_array1<E: Scalar, D: Data<Elem = E>>(
m: usize,
n: usize,
rank: i32,
b: &ArrayBase<D, Ix2>,
) -> Option<Array1<E::Real>> {
if m < n || n != rank as usize {
return None;
}
Some(
b.slice(s![n.., ..])
.mapv(|x| x.powi(2).abs())
.sum_axis(Axis(0)),
)
}
#[cfg(test)]
mod tests {
use crate::{error::LinalgError, *};
use approx::AbsDiffEq;
use ndarray::*;
fn assert_result<D1: Data<Elem = f64>, D2: Data<Elem = f64>>(
a: &ArrayBase<D1, Ix2>,
b: &ArrayBase<D2, Ix1>,
res: &LeastSquaresResult<f64, Ix1>,
) {
assert_eq!(res.rank, 2);
let b_hat = a.dot(&res.solution);
let rssq = (b - &b_hat).mapv(|x| x.powi(2)).sum();
assert!(res.residual_sum_of_squares.as_ref().unwrap()[()].abs_diff_eq(&rssq, 1e-12));
assert!(res
.solution
.abs_diff_eq(&array![-0.428571428571429, 0.85714285714285], 1e-12));
}
#[test]
fn on_arc() {
let a: ArcArray2<f64> = array![[1., 2.], [4., 5.], [3., 4.]].into_shared();
let b: ArcArray1<f64> = array![1., 2., 3.].into_shared();
let res = a.least_squares(&b).unwrap();
assert_result(&a, &b, &res);
}
#[test]
fn on_cow() {
let a = CowArray::from(array![[1., 2.], [4., 5.], [3., 4.]]);
let b = CowArray::from(array![1., 2., 3.]);
let res = a.least_squares(&b).unwrap();
assert_result(&a, &b, &res);
}
#[test]
fn on_view() {
let a: Array2<f64> = array![[1., 2.], [4., 5.], [3., 4.]];
let b: Array1<f64> = array![1., 2., 3.];
let av = a.view();
let bv = b.view();
let res = av.least_squares(&bv).unwrap();
assert_result(&av, &bv, &res);
}
#[test]
fn on_view_mut() {
let mut a: Array2<f64> = array![[1., 2.], [4., 5.], [3., 4.]];
let mut b: Array1<f64> = array![1., 2., 3.];
let av = a.view_mut();
let bv = b.view_mut();
let res = av.least_squares(&bv).unwrap();
assert_result(&av, &bv, &res);
}
#[test]
fn on_cow_view() {
let a = CowArray::from(array![[1., 2.], [4., 5.], [3., 4.]]);
let b: Array1<f64> = array![1., 2., 3.];
let bv = b.view();
let res = a.least_squares(&bv).unwrap();
assert_result(&a, &bv, &res);
}
#[test]
fn into_on_owned() {
let a: Array2<f64> = array![[1., 2.], [4., 5.], [3., 4.]];
let b: Array1<f64> = array![1., 2., 3.];
let ac = a.clone();
let bc = b.clone();
let res = ac.least_squares_into(bc).unwrap();
assert_result(&a, &b, &res);
}
#[test]
fn into_on_arc() {
let a: ArcArray2<f64> = array![[1., 2.], [4., 5.], [3., 4.]].into_shared();
let b: ArcArray1<f64> = array![1., 2., 3.].into_shared();
let a2 = a.clone();
let b2 = b.clone();
let res = a2.least_squares_into(b2).unwrap();
assert_result(&a, &b, &res);
}
#[test]
fn into_on_cow() {
let a = CowArray::from(array![[1., 2.], [4., 5.], [3., 4.]]);
let b = CowArray::from(array![1., 2., 3.]);
let a2 = a.clone();
let b2 = b.clone();
let res = a2.least_squares_into(b2).unwrap();
assert_result(&a, &b, &res);
}
#[test]
fn into_on_owned_cow() {
let a: Array2<f64> = array![[1., 2.], [4., 5.], [3., 4.]];
let b = CowArray::from(array![1., 2., 3.]);
let ac = a.clone();
let b2 = b.clone();
let res = ac.least_squares_into(b2).unwrap();
assert_result(&a, &b, &res);
}
#[test]
fn in_place_on_owned() {
let a = array![[1., 2.], [4., 5.], [3., 4.]];
let b = array![1., 2., 3.];
let mut a2 = a.clone();
let mut b2 = b.clone();
let res = a2.least_squares_in_place(&mut b2).unwrap();
assert_result(&a, &b, &res);
}
#[test]
fn in_place_on_cow() {
let a = CowArray::from(array![[1., 2.], [4., 5.], [3., 4.]]);
let b = CowArray::from(array![1., 2., 3.]);
let mut a2 = a.clone();
let mut b2 = b.clone();
let res = a2.least_squares_in_place(&mut b2).unwrap();
assert_result(&a, &b, &res);
}
#[test]
fn in_place_on_mut_view() {
let a = array![[1., 2.], [4., 5.], [3., 4.]];
let b = array![1., 2., 3.];
let mut a2 = a.clone();
let mut b2 = b.clone();
let av = &mut a2.view_mut();
let bv = &mut b2.view_mut();
let res = av.least_squares_in_place(bv).unwrap();
assert_result(&a, &b, &res);
}
#[test]
fn in_place_on_owned_cow() {
let a = array![[1., 2.], [4., 5.], [3., 4.]];
let b = CowArray::from(array![1., 2., 3.]);
let mut a2 = a.clone();
let mut b2 = b.clone();
let res = a2.least_squares_in_place(&mut b2).unwrap();
assert_result(&a, &b, &res);
}
#[test]
fn incompatible_shape_error_on_mismatching_num_rows() {
let a: Array2<f64> = array![[1., 2.], [4., 5.], [3., 4.]];
let b: Array1<f64> = array![1., 2.];
match a.least_squares(&b) {
Err(LinalgError::Shape(e)) if e.kind() == ErrorKind::IncompatibleShape => {}
_ => panic!("Should be raise IncompatibleShape"),
}
}
}