rlst 0.10.0

//! Growable frames and rank-adaptive pivoted QR.
//!
//! [`ElementFrame`] forwards operations to a [`FrameImpl`]. [`GenericFrameImpl`]
//! stores owned elements of any linear space; [`ArrayFrameImpl`] stores columns
//! in one column-major matrix. Frames are built outside any space by wrapping
//! the chosen implementation with [`ElementFrame::new`]. QR never modifies the
//! input frame.
//!
//! Element-backed implementations additionally support
//! [`FrameRefAccess`], which borrows the stored columns
//! and allows in-place mutation. Packed storages have no stored element to
//! borrow and expose array views instead.

mod array;
mod generic;
mod qr;

pub use array::ArrayFrameImpl;
pub use generic::GenericFrameImpl;
pub use qr::{FrameQr, FrameQrDecomposition, FrameQrOptions};

use std::{marker::PhantomData, ops::Neg};

use num::One;

use crate::operator::element::Element;
use crate::{
    Array, FrameApply, FrameImpl, FrameRefAccess, InnerProductSpace, LinearSpace, RlstError,
    RlstNum, RlstResult, RlstScalar, Shape, SliceArray, SliceArrayMut, UnsafeRandom1DAccessByValue,
};

#[cfg(test)]
mod tests;

#[cfg(test)]
mod forwarding_tests;

/// A frame that owns an implementation and forwards all operations to it.
///
/// The implementation determines storage and supported capabilities. Application
/// requires [`FrameApply`], and QR requires [`FrameQr`]. Elements and QR results
/// borrow the underlying space for `'a`, independently of the wrapper's lifetime.
pub struct ElementFrame<'a, ElementFrameImpl: FrameImpl<'a>> {
    imp: ElementFrameImpl,
    _space: PhantomData<&'a ElementFrameImpl::Space>,
}

impl<'a, I: FrameImpl<'a>> ElementFrame<'a, I> {
    /// Wrap an existing implementation without copying its storage.
    pub fn new(imp: I) -> Self {
        Self {
            imp,
            _space: PhantomData,
        }
    }

    /// Borrow the implementation.
    pub fn imp(&self) -> &I {
        &self.imp
    }

    /// Mutably borrow the implementation.
    pub fn imp_mut(&mut self) -> &mut I {
        &mut self.imp
    }

    /// Consume the wrapper and return its implementation without copying.
    pub fn into_imp(self) -> I {
        self.imp
    }

    /// Return the space associated with the columns.
    pub fn space(&self) -> &'a I::Space {
        self.imp.space()
    }

    /// Copy a column into an owned element, or return `None` out of bounds.
    pub fn element(&self, index: usize) -> Option<Element<'a, I::Space>> {
        self.imp.element(index)
    }

    /// Replace a column, delegating bounds and compatibility checks.
    pub fn set_element(&mut self, index: usize, elem: &Element<I::Space>) -> RlstResult<()> {
        self.imp.set_element(index, elem)
    }

    /// Number of active columns.
    pub fn len(&self) -> usize {
        self.imp.len()
    }

    /// Whether there are no active columns.
    pub fn is_empty(&self) -> bool {
        self.imp.is_empty()
    }

    /// Number of columns that fit without reallocating storage.
    pub fn capacity(&self) -> usize {
        self.imp.capacity()
    }

    /// Reserve room for at least `additional` further columns.
    pub fn reserve(&mut self, additional: usize) {
        self.imp.reserve(additional);
    }

    /// Copy an element into a new column, delegating compatibility checks.
    pub fn push(&mut self, elem: &Element<I::Space>) -> RlstResult<()> {
        self.imp.push(elem)
    }

    /// Remove the last column, or return `None` if empty.
    pub fn pop(&mut self) -> Option<Element<'a, I::Space>> {
        self.imp.pop()
    }

    /// Remove all columns, retaining capacity.
    pub fn clear(&mut self) {
        self.imp.clear();
    }
}

impl<'a, I: FrameImpl<'a>> FrameImpl<'a> for ElementFrame<'a, I> {
    type Space = I::Space;

    fn space(&self) -> &'a Self::Space {
        self.imp.space()
    }
    fn element(&self, index: usize) -> Option<Element<'a, Self::Space>> {
        self.imp.element(index)
    }
    fn set_element(&mut self, index: usize, elem: &Element<Self::Space>) -> RlstResult<()> {
        self.imp.set_element(index, elem)
    }
    fn len(&self) -> usize {
        self.imp.len()
    }
    fn is_empty(&self) -> bool {
        self.imp.is_empty()
    }
    fn capacity(&self) -> usize {
        self.imp.capacity()
    }
    fn reserve(&mut self, additional: usize) {
        self.imp.reserve(additional);
    }
    fn push(&mut self, elem: &Element<Self::Space>) -> RlstResult<()> {
        self.imp.push(elem)
    }
    fn pop(&mut self) -> Option<Element<'a, Self::Space>> {
        self.imp.pop()
    }
    fn clear(&mut self) {
        self.imp.clear();
    }
}

impl<'a, I: FrameApply<'a>> ElementFrame<'a, I> {
    /// Form a linear combination using the implementation's application routine.
    ///
    /// Accepts strided arrays and lazy expressions. The coefficient length must
    /// equal [`Self::len`]; an empty combination is the space's zero element.
    pub fn apply<A>(&self, coefficients: &Array<A, 1>) -> RlstResult<Element<'a, I::Space>>
    where
        A: Shape<1> + UnsafeRandom1DAccessByValue<Item = <I::Space as LinearSpace>::F>,
    {
        self.imp.apply(coefficients)
    }

    /// Compute the inner products of `x` with every column, `V^H x`.
    ///
    /// Delegates to the implementation's [`FrameApply::adjoint_apply`].
    pub fn adjoint_apply(
        &self,
        x: &Element<I::Space>,
    ) -> RlstResult<Vec<<I::Space as LinearSpace>::F>>
    where
        I::Space: InnerProductSpace,
    {
        self.imp.adjoint_apply(x)
    }

    /// Compute `output = alpha * self * coefficients + beta * output`.
    ///
    /// Delegates to the implementation's [`FrameApply::apply_into`], preserving
    /// its validation and optimized application. A zero `beta` discards the old
    /// output. Shape errors leave the output unchanged.
    pub fn apply_into<A>(
        &self,
        alpha: <I::Space as LinearSpace>::F,
        coefficients: &Array<A, 1>,
        beta: <I::Space as LinearSpace>::F,
        output: &mut Element<I::Space>,
    ) -> RlstResult<()>
    where
        A: Shape<1> + UnsafeRandom1DAccessByValue<Item = <I::Space as LinearSpace>::F>,
        <I::Space as LinearSpace>::F: num::Zero + num::One + PartialEq + Copy,
    {
        self.imp.apply_into(alpha, coefficients, beta, output)
    }
}

impl<'a, I: FrameApply<'a>> FrameApply<'a> for ElementFrame<'a, I> {
    fn adjoint_apply(
        &self,
        x: &Element<Self::Space>,
    ) -> RlstResult<Vec<<Self::Space as LinearSpace>::F>>
    where
        Self::Space: InnerProductSpace,
    {
        self.imp.adjoint_apply(x)
    }

    fn apply<A>(&self, coefficients: &Array<A, 1>) -> RlstResult<Element<'a, Self::Space>>
    where
        A: Shape<1> + UnsafeRandom1DAccessByValue<Item = <Self::Space as LinearSpace>::F>,
    {
        self.imp.apply(coefficients)
    }

    fn apply_into<A>(
        &self,
        alpha: <Self::Space as LinearSpace>::F,
        coefficients: &Array<A, 1>,
        beta: <Self::Space as LinearSpace>::F,
        output: &mut Element<Self::Space>,
    ) -> RlstResult<()>
    where
        A: Shape<1> + UnsafeRandom1DAccessByValue<Item = <Self::Space as LinearSpace>::F>,
        <Self::Space as LinearSpace>::F: num::Zero + num::One + PartialEq + Copy,
    {
        self.imp.apply_into(alpha, coefficients, beta, output)
    }
}

impl<'a, I: FrameQr<'a>> ElementFrame<'a, I>
where
    I::Space: InnerProductSpace,
    <I::Space as LinearSpace>::F: RlstScalar,
{
    /// Compute pivoted QR with the implementation's algorithm and wrap Q.
    ///
    /// The result satisfies `self[:, permutation] ≈ q * r` with the same
    /// implementation type for Q. See [`FrameQrOptions`] for rank selection.
    pub fn qr_pivoted(
        &self,
        options: FrameQrOptions<<<I::Space as LinearSpace>::F as RlstScalar>::Real>,
    ) -> RlstResult<FrameQrDecomposition<Self, <I::Space as LinearSpace>::F>> {
        let FrameQrDecomposition { q, r, permutation } = self.imp.qr_pivoted(options)?;
        Ok(FrameQrDecomposition {
            q: Self::new(q),
            r,
            permutation,
        })
    }
}

impl<'a, I: FrameQr<'a>> FrameQr<'a> for ElementFrame<'a, I>
where
    I::Space: InnerProductSpace,
    <I::Space as LinearSpace>::F: RlstScalar,
{
    fn qr_pivoted(
        &self,
        options: FrameQrOptions<<<Self::Space as LinearSpace>::F as RlstScalar>::Real>,
    ) -> RlstResult<FrameQrDecomposition<Self, <Self::Space as LinearSpace>::F>> {
        Self::qr_pivoted(self, options)
    }
}

impl<'a, I: FrameRefAccess<'a>> ElementFrame<'a, I> {
    /// Borrow a column without copying it, or return `None` out of bounds.
    pub fn element_ref(&self, index: usize) -> Option<&Element<'a, I::Space>> {
        self.imp.element_ref(index)
    }

    /// Mutably borrow a column, or return `None` out of bounds.
    ///
    /// Mutating a column through this reference must preserve the space's
    /// invariants; see [`FrameRefAccess`].
    pub fn element_mut<'s>(&'s mut self, index: usize) -> Option<&'s mut Element<'a, I::Space>> {
        self.imp.element_mut(index)
    }

    /// Iterate over stored elements without copying them.
    pub fn iter<'s>(&'s self) -> impl ExactSizeIterator<Item = &'s Element<'a, I::Space>>
    where
        'a: 's,
    {
        self.imp.iter()
    }

    /// Mutably iterate over stored elements.
    ///
    /// The same invariant contract applies as for [`Self::element_mut`].
    pub fn iter_mut<'s>(
        &'s mut self,
    ) -> impl ExactSizeIterator<Item = &'s mut Element<'a, I::Space>>
    where
        'a: 's,
    {
        self.imp.iter_mut()
    }
}

impl<'a, I: FrameRefAccess<'a>> FrameRefAccess<'a> for ElementFrame<'a, I> {
    fn element_ref(&self, index: usize) -> Option<&Element<'a, Self::Space>> {
        self.imp.element_ref(index)
    }
    fn element_mut<'s>(&'s mut self, index: usize) -> Option<&'s mut Element<'a, Self::Space>> {
        self.imp.element_mut(index)
    }
    fn iter<'s>(&'s self) -> impl ExactSizeIterator<Item = &'s Element<'a, Self::Space>>
    where
        'a: 's,
    {
        self.imp.iter()
    }
    fn iter_mut<'s>(&'s mut self) -> impl ExactSizeIterator<Item = &'s mut Element<'a, Self::Space>>
    where
        'a: 's,
    {
        self.imp.iter_mut()
    }
}

impl<'a, Item: RlstNum + Neg<Output = Item>> ElementFrame<'a, ArrayFrameImpl<'a, Item>> {
    /// Borrow the active matrix without copying.
    pub fn matrix(&self) -> SliceArray<'_, Item, 2> {
        self.imp.matrix()
    }

    /// Mutably borrow active entries without permitting storage resizing.
    pub fn matrix_mut(&mut self) -> SliceArrayMut<'_, Item, 2> {
        self.imp.matrix_mut()
    }

    /// Borrow an active column, or return `None` out of bounds.
    pub fn column(&self, index: usize) -> Option<SliceArray<'_, Item, 1>> {
        self.imp.column(index)
    }

    /// Mutably borrow an active column without permitting storage resizing.
    pub fn column_mut(&mut self, index: usize) -> Option<SliceArrayMut<'_, Item, 1>> {
        self.imp.column_mut(index)
    }
}

/// Orthogonalize `candidate` against the columns of `frame` and return the
/// projection coefficients, one per column.
///
/// Two passes of classical Gram-Schmidt (CGS2) are used. Each pass computes all
/// inner products through [`FrameApply::adjoint_apply`] before subtracting the
/// combination through [`FrameApply::apply_into`], so array-backed frames run
/// two dense matrix-vector products per pass and no column is copied. The
/// columns are assumed to be orthonormal, as for a Krylov basis or a Q factor.
/// Shared by GMRES and the generic pivoted QR.
pub(crate) fn orthogonalize_twice<'a, F>(
    frame: &F,
    candidate: &mut Element<F::Space>,
) -> RlstResult<Vec<<F::Space as LinearSpace>::F>>
where
    F: FrameApply<'a>,
    F::Space: InnerProductSpace,
    <F::Space as LinearSpace>::F: RlstScalar,
{
    let mut coefficients = Vec::new();
    for pass in 0..2 {
        let projection = frame.adjoint_apply(candidate)?;
        frame.apply_into(
            -<<F::Space as LinearSpace>::F as One>::one(),
            &SliceArray::from_shape(&projection, [projection.len()]),
            One::one(),
            candidate,
        )?;
        if pass == 0 {
            coefficients = projection;
        } else {
            for (total, correction) in coefficients.iter_mut().zip(projection) {
                *total += correction;
            }
        }
    }
    Ok(coefficients)
}

fn index_error() -> RlstError {
    RlstError::IndexLayoutError("Frame column index out of bounds".into())
}