rlst 0.9.0

//! Implementation of a general abstract operator.

use num::{One, Zero};

use crate::base_types::{c32, c64};
use crate::{LinearSpace, abstract_operator::OperatorBase};

/// Owning wrapper that gives an [`OperatorBase`] implementation operator algebra.
///
/// Addition, subtraction, products, negation, and scalar multiplication consume
/// their operands. Use [`Self::r`] when an expression should borrow an operator.
pub struct Operator<OpImpl>(OpImpl);

impl<OpImpl> Operator<OpImpl>
where
    OpImpl: OperatorBase,
{
    /// Wrap an operator implementation without changing its domain or range.
    pub fn new(op_impl: OpImpl) -> Self {
        Self(op_impl)
    }

    /// Return a lightweight operator that borrows this operator.
    pub fn r(&self) -> Operator<RlstOperatorReference<'_, OpImpl>> {
        Operator::new(RlstOperatorReference::new(self))
    }
}

/// Operator implementation that forwards all operations to a borrowed [`Operator`].
pub struct RlstOperatorReference<'a, OpImpl>(&'a Operator<OpImpl>);

impl<'a, OpImpl: OperatorBase> RlstOperatorReference<'a, OpImpl> {
    /// Borrow `op` for use in ownership-consuming operator expressions.
    pub fn new(op: &'a Operator<OpImpl>) -> Self {
        Self(op)
    }
}

impl<OpImpl: OperatorBase> OperatorBase for Operator<OpImpl> {
    type Domain = OpImpl::Domain;

    type Range = OpImpl::Range;

    #[inline(always)]
    fn domain(&self) -> &Self::Domain {
        self.0.domain()
    }

    #[inline(always)]
    fn range(&self) -> &Self::Range {
        self.0.range()
    }

    #[inline(always)]
    fn apply(
        &self,
        alpha: <Self::Range as crate::LinearSpace>::F,
        x: &super::element::Element<Self::Domain>,
        beta: <Self::Range as crate::LinearSpace>::F,
        y: &mut super::element::Element<Self::Range>,
    ) {
        self.0.apply(alpha, x, beta, y)
    }
}

impl<OpImpl: OperatorBase> OperatorBase for RlstOperatorReference<'_, OpImpl> {
    type Domain = OpImpl::Domain;

    type Range = OpImpl::Range;

    fn domain(&self) -> &Self::Domain {
        self.0.domain()
    }

    fn range(&self) -> &Self::Range {
        self.0.range()
    }

    #[inline(always)]
    fn apply(
        &self,
        alpha: <Self::Range as crate::LinearSpace>::F,
        x: &super::element::Element<Self::Domain>,
        beta: <Self::Range as crate::LinearSpace>::F,
        y: &mut super::element::Element<Self::Range>,
    ) {
        self.0.apply(alpha, x, beta, y)
    }
}

/// Lazy owning sum of two operators with identical domain and range.
pub struct OperatorSum<OpImpl1, OpImpl2>(Operator<OpImpl1>, Operator<OpImpl2>);

/// Lazy owning scalar multiple of an operator.
pub struct OperatorScalarMul<Scalar, OpImpl>(Scalar, Operator<OpImpl>);

/// Lazy owning difference of operators with identical domain and range.
pub struct OperatorSub<OpImpl1, OpImpl2>(Operator<OpImpl1>, Operator<OpImpl2>);

/// Lazy composition `Op1 * Op2`, applying `Op2` first.
///
/// The range of `Op2` must equal the domain of `Op1`.
pub struct OperatorProduct<OpImpl1, OpImpl2>(Operator<OpImpl1>, Operator<OpImpl2>);

/// Lazy owning negation of an operator.
pub struct OperatorNegation<OpImpl>(Operator<OpImpl>);

impl<Domain, Range, OpImpl1, OpImpl2> OperatorBase for OperatorSum<OpImpl1, OpImpl2>
where
    Domain: LinearSpace,
    Range: LinearSpace,
    Range::F: One + Copy,
    OpImpl1: OperatorBase<Domain = Domain, Range = Range>,
    OpImpl2: OperatorBase<Domain = Domain, Range = Range>,
{
    type Domain = Domain;

    type Range = Range;

    #[inline(always)]
    fn domain(&self) -> &Self::Domain {
        self.0.domain()
    }

    #[inline(always)]
    fn range(&self) -> &Self::Range {
        self.0.range()
    }

    #[inline(always)]
    fn apply(
        &self,
        alpha: <Self::Range as crate::LinearSpace>::F,
        x: &super::element::Element<Self::Domain>,
        beta: <Self::Range as crate::LinearSpace>::F,
        y: &mut super::element::Element<Self::Range>,
    ) {
        self.0.apply(alpha, x, beta, y);
        self.1.apply(alpha, x, One::one(), y);
    }
}

impl<Domain, Range, OpImpl1, OpImpl2> OperatorBase for OperatorSub<OpImpl1, OpImpl2>
where
    Domain: LinearSpace,
    Range: LinearSpace,
    Range::F: One + Copy + std::ops::Neg<Output = Range::F>,
    OpImpl1: OperatorBase<Domain = Domain, Range = Range>,
    OpImpl2: OperatorBase<Domain = Domain, Range = Range>,
{
    type Domain = Domain;

    type Range = Range;

    #[inline(always)]
    fn domain(&self) -> &Self::Domain {
        self.0.domain()
    }

    #[inline(always)]
    fn range(&self) -> &Self::Range {
        self.0.range()
    }

    #[inline(always)]
    fn apply(
        &self,
        alpha: <Self::Range as crate::LinearSpace>::F,
        x: &super::element::Element<Self::Domain>,
        beta: <Self::Range as crate::LinearSpace>::F,
        y: &mut super::element::Element<Self::Range>,
    ) {
        self.0.apply(alpha, x, beta, y);
        self.1.apply(std::ops::Neg::neg(alpha), x, One::one(), y);
    }
}

impl<Range, Scalar, OpImpl> OperatorBase for OperatorScalarMul<Scalar, OpImpl>
where
    Range: LinearSpace<F = Scalar>,
    Scalar: std::ops::Mul<Output = Scalar> + Copy,
    OpImpl: OperatorBase<Range = Range>,
{
    type Domain = OpImpl::Domain;

    type Range = OpImpl::Range;

    #[inline(always)]
    fn domain(&self) -> &Self::Domain {
        self.1.domain()
    }

    #[inline(always)]
    fn range(&self) -> &Self::Range {
        self.1.range()
    }

    #[inline(always)]
    fn apply(
        &self,
        alpha: <Self::Range as LinearSpace>::F,
        x: &super::element::Element<Self::Domain>,
        beta: <Self::Range as LinearSpace>::F,
        y: &mut super::element::Element<Self::Range>,
    ) {
        self.1.apply(self.0 * alpha, x, beta, y)
    }
}

impl<Range, OpImpl1, OpImpl2> OperatorBase for OperatorProduct<Operator<OpImpl1>, Operator<OpImpl2>>
where
    Range: LinearSpace,
    Range::F: One + Zero,
    OpImpl1: OperatorBase<Domain = Range>,
    OpImpl2: OperatorBase<Range = Range>,
{
    type Domain = OpImpl2::Domain;

    type Range = OpImpl1::Range;

    #[inline(always)]
    fn domain(&self) -> &Self::Domain {
        self.1.domain()
    }

    #[inline(always)]
    fn range(&self) -> &Self::Range {
        self.0.range()
    }

    #[inline(always)]
    fn apply(
        &self,
        alpha: <Self::Range as LinearSpace>::F,
        x: &super::element::Element<Self::Domain>,
        beta: <Self::Range as LinearSpace>::F,
        y: &mut super::element::Element<Self::Range>,
    ) {
        let mut tmp = self.1.range().zero();

        self.1.apply(One::one(), x, Zero::zero(), &mut tmp);
        self.0.apply(alpha, &tmp, beta, y);
    }
}

impl<OpImpl> OperatorBase for OperatorNegation<OpImpl>
where
    OpImpl: OperatorBase,
    <OpImpl::Range as LinearSpace>::F: std::ops::Neg<Output = <OpImpl::Range as LinearSpace>::F>,
{
    type Domain = OpImpl::Domain;

    type Range = OpImpl::Range;

    fn domain(&self) -> &Self::Domain {
        self.0.domain()
    }

    fn range(&self) -> &Self::Range {
        self.0.range()
    }

    fn apply(
        &self,
        alpha: <Self::Range as LinearSpace>::F,
        x: &super::element::Element<Self::Domain>,
        beta: <Self::Range as LinearSpace>::F,
        y: &mut super::element::Element<Self::Range>,
    ) {
        self.0.apply(std::ops::Neg::neg(alpha), x, beta, y);
    }
}

impl<Domain, Range, OpImpl1, OpImpl2> std::ops::Add<Operator<OpImpl2>> for Operator<OpImpl1>
where
    Domain: LinearSpace,
    Range: LinearSpace,
    Range::F: One + Copy,
    OpImpl1: OperatorBase<Domain = Domain, Range = Range>,
    OpImpl2: OperatorBase<Domain = Domain, Range = Range>,
{
    type Output = Operator<OperatorSum<OpImpl1, OpImpl2>>;

    fn add(self, rhs: Operator<OpImpl2>) -> Self::Output {
        Operator::new(OperatorSum(self, rhs))
    }
}

impl<Domain, Range, OpImpl1, OpImpl2> std::ops::Sub<Operator<OpImpl2>> for Operator<OpImpl1>
where
    Domain: LinearSpace,
    Range: LinearSpace,
    Range::F: One + Copy + std::ops::Neg<Output = Range::F>,
    OpImpl1: OperatorBase<Domain = Domain, Range = Range>,
    OpImpl2: OperatorBase<Domain = Domain, Range = Range>,
{
    type Output = Operator<OperatorSub<OpImpl1, OpImpl2>>;

    fn sub(self, rhs: Operator<OpImpl2>) -> Self::Output {
        Operator::new(OperatorSub(self, rhs))
    }
}

impl<Scalar, OpImpl> std::ops::Mul<Scalar> for Operator<OpImpl>
where
    OpImpl: OperatorBase,
    OpImpl::Range: LinearSpace<F = Scalar>,
    Scalar: std::ops::Mul<Output = Scalar> + Copy,
{
    type Output = Operator<OperatorScalarMul<Scalar, OpImpl>>;

    fn mul(self, rhs: Scalar) -> Self::Output {
        Operator::new(OperatorScalarMul(rhs, self))
    }
}

macro_rules! impl_scalar_mult {
    ($scalar:ty) => {
        impl<OpImpl> std::ops::Mul<Operator<OpImpl>> for $scalar
        where
            OpImpl: OperatorBase,
            OpImpl::Range: LinearSpace<F = $scalar>,
        {
            type Output = Operator<OperatorScalarMul<$scalar, OpImpl>>;

            fn mul(self, rhs: Operator<OpImpl>) -> Self::Output {
                Operator::new(OperatorScalarMul(self, rhs))
            }
        }
    };
}

impl_scalar_mult!(f64);
impl_scalar_mult!(f32);
impl_scalar_mult!(c64);
impl_scalar_mult!(c32);
impl_scalar_mult!(usize);
impl_scalar_mult!(i8);
impl_scalar_mult!(i16);
impl_scalar_mult!(i32);
impl_scalar_mult!(i64);
impl_scalar_mult!(u8);
impl_scalar_mult!(u16);
impl_scalar_mult!(u32);
impl_scalar_mult!(u64);