src/linops.rs

/*!
Abstract linear operators.
*/

use numeric_literals::replace_float_literals;
use std::marker::PhantomData;
use serde::Serialize;
use crate::types::*;
pub use crate::mapping::{Mapping, Space, Composition};
use crate::direct_product::Pair;
use crate::instance::Instance;
use crate::norms::{NormExponent, PairNorm, L1, L2, Linfinity, Norm, HasDual};

/// Trait for linear operators on `X`.
pub trait Linear<X : AXPY> : Mapping<X, Codomain = Self::LinCodomain> {
    type LinCodomain : AXPY; // = Self::Codomain, but with AXPY enforced without trait bound bloat.
}

/// Efficient in-place summation.
#[replace_float_literals(Self::Field::cast_from(literal))]
pub trait AXPY : Space {
    type Field : Num;

    /// Computes  `y = βy + αx`, where `y` is `Self`.
    fn axpy<I : Instance<Self>>(&mut self, α : Self::Field, x : I, β : Self::Field);

    /// Copies `x` to `self`.
    fn copy_from<I : Instance<Self>>(&mut self, x : I) {
        self.axpy(1.0, x, 0.0)
    }

    /// Computes  `y = αx`, where `y` is `Self`.
    fn scale_from<I : Instance<Self>>(&mut self, α : Self::Field, x : I) {
        self.axpy(α, x, 0.0)
    }
}

macro_rules! impl_axpy {
    ($($type:ty)*) => { $(
        impl AXPY for $type {
            type Field = $type;
            fn axpy<I : Instance<Self>>(&mut self, α : $type, x : I, β : $type) {
                *self = *self * α + x.own() * β;
            }
        }
    )* };
}

impl_axpy!(i8 i16 i32 i64 i128 isize f32 f64);


/// Efficient in-place application for [`Linear`] operators.
#[replace_float_literals(X::Field::cast_from(literal))]
pub trait GEMV<X : AXPY> : Linear<X> {
    /// Computes  `y = αAx + βy`, where `A` is `Self`.
    fn gemv<I : Instance<X>>(&self, y : &mut Self::LinCodomain, α : <Self::LinCodomain as AXPY>::Field, x : I, β : <Self::LinCodomain as AXPY>::Field);

    /// Computes `y = Ax`, where `A` is `Self`
    fn apply_mut<I : Instance<X>>(&self, y : &mut Self::LinCodomain, x : I){
        self.gemv(y, 1.0, x, 0.0)
    }

    /// Computes `y += Ax`, where `A` is `Self`
    fn apply_add<I : Instance<X>>(&self, y : &mut Self::LinCodomain, x : I){
        self.gemv(y, 1.0, x, 1.0)
    }
}


/// Bounded linear operators
pub trait BoundedLinear<X, XExp, CodExp> : Linear<X>
where
    F : Num,
    X : AXP + Norm<F, XExp>,
    XExp : NormExponent,
    CodExp : NormExponent
{
    /// A bound on the operator norm $\|A\|$ for the linear operator $A$=`self`.
    /// This is not expected to be the norm, just any bound on it that can be
    /// reasonably implemented. The [`NormExponent`] `xexp`  indicates the norm
    /// in `X`, and `codexp` in the codomain.
    fn opnorm_bound(&self, xexp : XExp, codexp : CodExp) -> X::Field; // TODO: Should pick real subfield
}

// Linear operator application into mutable target. The [`AsRef`] bound
// is used to guarantee compatibility with `Yʹ` and `Self::Codomain`;
// the former is assumed to be e.g. a view into the latter.

/*impl<X,Y,T> Fn(&X) -> Y for T where T : Linear<X,Codomain=Y> {
    fn call(&self, x : &X) -> Y {
        self.apply(x)
    }
}*/

/// Trait for forming the adjoint operator of `Self`.
pub trait Adjointable<X, F : Float = f64> : Linear<X>
where
    X : HasDual<F>,
    Self::Codomain : HasDual<F>,
{
    type AdjointCodomain : Instance<X::DualSpace>;
    type Adjoint<'a> : Linear<
        <Self::Codomain as HasDual<F>>::DualSpace,
        Codomain=Self::AdjointCodomain,
    > where Self : 'a;

    /// Form the adjoint operator of `self`.
    fn adjoint(&self) -> Self::Adjoint<'_>;
}

/// Trait for forming a preadjoint of an operator.
///
/// For an operator $A$ this is an operator $A\_\*$
/// such that its adjoint $(A\_\*)^\*=A$. The space `X` is the domain of the `Self`
/// operator. The space `Ypre` is the predual of its codomain, and should be the
/// domain of the adjointed operator. `Self::Preadjoint` should be
/// [`Adjointable`]`<'a,Ypre,X>`.
///
/// We do not set further restrictions on the spacds, to allow preadjointing when `X`
/// is on the dual space of `Ypre`, but a subset of it.
pub trait Preadjointable<X : AXPY, Ypre : AXPY> : Linear<X>
//where
//    Ypre : HasDual<F, DualSpace=Self::Codomain>,
{
    type PreadjointCodomain : AXPY;
    type Preadjoint<'a> : Linear<Ypre, Codomain=Self::PreadjointCodomain> where Self : 'a;

    /// Form the adjoint operator of `self`.
    fn preadjoint(&self) -> Self::Preadjoint<'_>;
}

/// The identity operator
#[derive(Clone,Copy,Debug,Serialize,Eq,PartialEq)]
pub struct IdOp<X : AXPY> (PhantomData<X>);

impl<X : AXPY> IdOp<X> {
    pub fn new() -> IdOp<X> { IdOp(PhantomData) }
}

impl<X : Clone + AXPY> Mapping<X> for IdOp<X> {
    type Codomain = X;

    fn apply<I : Instance<X>>(&self, x : I) -> X {
        x.own()
    }
}

impl<X : Clone + AXPY> Linear<X> for IdOp<X> {
    type LinCodomain = X;
}

#[replace_float_literals(F::cast_from(literal))]
impl<X : Clone + AXPY> GEMV<X> for IdOp<X> {
    // Computes  `y = αAx + βy`, where `A` is `Self`.
    fn gemv<I : Instance<X>>(&self, y : &mut Self::LinCodomain, α : X::Field, x : I, β : X::Field) {
        y.axpy(α, x, β)
    }

    fn apply_mut<I : Instance<X>>(&self, y : &mut X, x : I){
        y.copy_from(x);
    }
}

impl<X, E> BoundedLinear<X, E, E> for IdOp<X>
where
    X : AXPY + Clone + Norm<F, E>,
    F : Num + AXPY,
    E : NormExponent
{
    fn opnorm_bound(&self, _xexp : E, _codexp : E) -> X::Field { F::ONE }
}

impl<X : Clone + HasDual<F, DualSpace=X>, F : Float> Adjointable<X, F> for IdOp<X> {
    type AdjointCodomain = X;
    type Adjoint<'a> = IdOp<X::DualSpace> where X : 'a;

    fn adjoint(&self) -> Self::Adjoint<'_> { IdOp::new() }
}

impl<X : Clone + AXPY> Preadjointable<X, X> for IdOp<X> {
    type PreadjointCodomain = X;
    type Preadjoint<'a> = IdOp<X> where X : 'a;

    fn preadjoint(&self) -> Self::Preadjoint<'_> { IdOp::new() }
}


impl<S, T, E, X> Linear<X> for Composition<S, T, E>
where
    X : AXPY,
    T : Linear<X>,
    S : Linear<T::LinCodomain>
{
    type LinCodomain = S::LinCodomain;
}

impl<F, S, T, E, X, Y> GEMV<F, X, Y> for Composition<S, T, E>
where
    F : Num,
    X : Space,
    T : Linear<X>,
    S : GEMV<F, T::Codomain, Y>,
{
    fn gemv<I : Instance<X>>(&self, y : &mut Y, α : F, x : I, β : F) {
        self.outer.gemv(y, α, self.inner.apply(x), β)
    }

    /// Computes `y = Ax`, where `A` is `Self`
    fn apply_mut<I : Instance<X>>(&self, y : &mut Y, x : I){
        self.outer.apply_mut(y, self.inner.apply(x))
    }

    /// Computes `y += Ax`, where `A` is `Self`
    fn apply_add<I : Instance<X>>(&self, y : &mut Y, x : I){
        self.outer.apply_add(y, self.inner.apply(x))
    }
}

impl<F, S, T, X, Z, Xexp, Yexp, Zexp> BoundedLinear<X, Xexp, Yexp, F> for Composition<S, T, Zexp>
where
    F : Num,
    X : Space + Norm<F, Xexp>,
    Z : Space + Norm<F, Zexp>,
    Xexp : NormExponent,
    Yexp : NormExponent,
    Zexp : NormExponent,
    T : BoundedLinear<X, Xexp, Zexp, F, Codomain=Z>,
    S : BoundedLinear<Z, Zexp, Yexp, F>,
{
    fn opnorm_bound(&self, xexp : Xexp, yexp : Yexp) -> F {
        let zexp = self.intermediate_norm_exponent;
        self.outer.opnorm_bound(zexp, yexp) * self.inner.opnorm_bound(xexp, zexp)
    }
}

/// “Row operator” $(S, T)$; $(S, T)(x, y)=Sx + Ty$.
pub struct RowOp<S, T>(pub S, pub T);

use std::ops::Add;

impl<A, B, S, T, Y> Mapping<Pair<A, B>> for RowOp<S, T>
where
    A : Space,
    B : Space,
    Y : Space,
    S : Mapping<A>,
    T : Mapping<B>,
    S::Codomain : Add<T::Codomain, Output = Y>,
{
    type Codomain = Y;

    fn apply<I : Instance<Pair<A, B>>>(&self, x : I) -> Self::Codomain {
        let Pair(a, b) = x.decompose();
        self.0.apply(a) + self.1.apply(b)
    }
}

impl<F, A, B, S, T, Y> Linear<Pair<A, B>> for RowOp<S, T>
where
    F : Num + AXPY,
    A : AXPY<Field=F>,
    B : AXPY<Field=F>,
    Y : AXPY,
    S : Linear<A>,
    T : Linear<B>,
    S::LinCodomain : Add<T::LinCodomain, Output=Y>,
{
    type LinCodomain = Y;
}


impl<'b, F, S, T, Y, G, U, V> GEMV<Pair<U, V>> for RowOp<S, T>
where
    U : AXPY<Field=G>,
    V : AXPY<Field=G>,
    Y : AXPY<Field=F>,
    S : GEMV<U>,
    T : GEMV<V>,
    F : Num + AXPY,
    G : Num + AXPY,
    S::LinCodomain : Add<T::LinCodomain, Output=Y>,
{
    fn gemv<I : Instance<Pair<U, V>>>(&self, y : &mut Y, α : F, x : I, β : F) {
        let Pair(u, v) = x.decompose();
        self.0.gemv(y, α, u, β);
        self.1.gemv(y, α, v, F::ONE);
    }

    fn apply_mut<I : Instance<Pair<U, V>>>(&self, y : &mut Y, x : I) {
        let Pair(u, v) = x.decompose();
        self.0.apply_mut(y, u);
        self.1.apply_mut(y, v);
    }

    /// Computes `y += Ax`, where `A` is `Self`
    fn apply_add<I : Instance<Pair<U, V>>>(&self, y : &mut Y, x : I) {
        let Pair(u, v) = x.decompose();
        self.0.apply_add(y, u);
        self.1.apply_add(y, v);
    }
}

/// “Column operator” $(S; T)$; $(S; T)x=(Sx, Tx)$.
pub struct ColOp<S, T>(pub S, pub T);

impl<X, S, T> Mapping<X> for ColOp<S, T>
where
    X : AXPY,
    S : Mapping<X>,
    T : Mapping<X>,
{
    type Codomain = Pair<S::Codomain, T::Codomain>;

    fn apply<I : Instance<X>>(&self, a : I) -> Self::Codomain {
        Pair(self.0.apply(a.ref_instance()), self.1.apply(a))
    }
}

impl<X, S, T, A, B, F> Linear<X> for ColOp<S, T>
where
    F : Num + AXPY,
    X : AXPY,
    S : Linear<X, LinCodomain=A>,
    T : Linear<X, LinCodomain=B>,
    A : AXPY<Field=F>,
    B : AXPY<Field=F>,
{
    type LinCodomain = Pair<S::LinCodomain, T::LinCodomain>;
}

impl<F, S, T, A, B, X> GEMV<X> for ColOp<S, T>
where
    F : Num + AXPY,
    X : AXPY,
    S : GEMV<X, LinCodomain=A>,
    T : GEMV<X, LinCodomain=B>,
    A : AXPY<Field=F>,
    B : AXPY<Field=F>,
{
    fn gemv<I : Instance<X>>(&self, y : &mut Pair<A, B>, α : F, x : I, β : F) {
        self.0.gemv(&mut y.0, α, x.ref_instance(), β);
        self.1.gemv(&mut y.1, α, x, β);
    }

    fn apply_mut<I : Instance<X>>(&self, y : &mut Pair<A, B>, x : I){
        self.0.apply_mut(&mut y.0, x.ref_instance());
        self.1.apply_mut(&mut y.1, x);
    }

    /// Computes `y += Ax`, where `A` is `Self`
    fn apply_add<I : Instance<X>>(&self, y : &mut Pair<A, B>, x : I){
        self.0.apply_add(&mut y.0, x.ref_instance());
        self.1.apply_add(&mut y.1, x);
    }
}


impl<F, A, B, Yʹ, S, T> Adjointable<Pair<A, B>, F> for RowOp<S, T>
where
    F : Float,
    A : HasDual<F>,
    B : HasDual<F>,
    S : Adjointable<A, F>,
    T : Adjointable<B, F>,
    Yʹ : AXPY,
    S :: Codomain : HasDual<F, DualSpace=Yʹ>,
    T :: Codomain : HasDual<F, DualSpace=Yʹ>,
    S::Codomain : Add<T::Codomain>,
    <S::Codomain as Add<T::Codomain>>::Output : HasDual<F, DualSpace=Yʹ>,
    Self::Codomain : HasDual<F, DualSpace=Yʹ>,
    //Self : Linear<Pair<A, B>>,
    // for<'a> ColOp<S::Adjoint<'a>, T::Adjoint<'a>> : Linear<
    //     Yʹ,
    //     //Codomain=Pair<S::AdjointCodomain, T::AdjointCodomain>
    // >,
{
    type AdjointCodomain = Pair<S::AdjointCodomain, T::AdjointCodomain>;
    type Adjoint<'a> = ColOp<S::Adjoint<'a>, T::Adjoint<'a>> where Self : 'a;

    fn adjoint(&self) -> Self::Adjoint<'_> {
        ColOp(self.0.adjoint(), self.1.adjoint())
    }
}

impl<F, A, B, Y, Yʹ, S, T> Preadjointable<Pair<A,B>, Yʹ> for RowOp<S, T>
wherea
    A : AXPY,
    B : AXPY,
    Yʹ : AXPY,
    S : Preadjointable<A, Yʹ>,
    T : Preadjointable<B, Yʹ>,
    S::Codomain : Add<T::Codomain>,
    <S::Codomain as Add<T::Codomain>>::Output : AXPY,
    //Self : Linear<Pair<A, B>, Codomain=Y>,
    //for<'a> ColOp<S::Preadjoint<'a>, T::Preadjoint<'a>> : Adjointable<Yʹ, F>,
{
    type PreadjointCodomain = Pair<S::PreadjointCodomain, T::PreadjointCodomain>;
    type Preadjoint<'a> = ColOp<S::Preadjoint<'a>, T::Preadjoint<'a>> where Self : 'a;

    fn preadjoint(&self) -> Self::Preadjoint<'_> {
        ColOp(self.0.preadjoint(), self.1.preadjoint())
    }
}


impl<F, A, Aʹ, S, T> Adjointable<A, F> for ColOp<S, T>
where
    F : Float,
    A : HasDual<F>,
    S : Adjointable<A, F>,
    T : Adjointable<A, F>,
    T::Codomain : HasDual<F>,
    S::Codomain : HasDual<F>,
    Aʹ : Space + Instance<A::DualSpace>,
    <S as Adjointable<A, F>>::AdjointCodomain : Add<<T as Adjointable<A, F>>::AdjointCodomain, Output=Aʹ>,
    // for<'a> RowOp<S::Adjoint<'a>, T::Adjoint<'a>> : Linear<
    //     Pair<<T::Codomain as HasDual<F>>::DualSpace, <S::Codomain as HasDual<F>>::DualSpace>,
    //     Codomain=Aʹ
    // >,
{
    type AdjointCodomain = Aʹ;
    type Adjoint<'a> = RowOp<S::Adjoint<'a>, T::Adjoint<'a>> where Self : 'a;

    fn adjoint(&self) -> Self::Adjoint<'_> {
        RowOp(self.0.adjoint(), self.1.adjoint())
    }
}

impl<A, Aʹ, Xʹ, Yʹ, S, T> Preadjointable<A,Pair<Xʹ,Yʹ>> for ColOp<S, T>
where
    A : AXPY,
    Aʹ : AXPY,
    Xʹ : AXPY,
    Yʹ : AXPY,
    S : Preadjointable<A, Xʹ, PreadjointCodomain=Aʹ>,
    T : Preadjointable<A, Yʹ, PreadjointCodomain=Aʹ>,
    Aʹ : ClosedAdd,
    //for<'a> RowOp<S::Preadjoint<'a>, T::Preadjoint<'a>> : Adjointable<Pair<Xʹ,Yʹ>, F>,
{
    type PreadjointCodomain = Aʹ;
    type Preadjoint<'a> = RowOp<S::Preadjoint<'a>, T::Preadjoint<'a>> where Self : 'a;

    fn preadjoint(&self) -> Self::Preadjoint<'_> {
        RowOp(self.0.preadjoint(), self.1.preadjoint())
    }
}

/// Diagonal operator
pub struct DiagOp<S, T>(pub S, pub T);

impl<A, B, S, T> Mapping<Pair<A, B>> for DiagOp<S, T>
where
    A : AXPY,
    B : AXPY,
    S : Mapping<A>,
    T : Mapping<B>,
{
    type Codomain = Pair<S::Codomain, T::Codomain>;

    fn apply<I : Instance<Pair<A, B>>>(&self, x : I) -> Self::Codomain {
        let Pair(a, b) = x.decompose();
        Pair(self.0.apply(a), self.1.apply(b))
    }
}

impl<F, A, B, S, T, G, U, V> Linear<Pair<A, B>> for DiagOp<S, T>
where
    F : Num + AXPY,
    G : Num + AXPY,
    A : AXPY<Field=F>,
    B : AXPY<Field=F>,
    U : AXPY<Field=G>,
    V : AXPY<Field=G>,
    S : Linear<A, LinCodomain=U>,
    T : Linear<B, LinCodomain=V>,
{
    type LinCodomain = Pair<U, V>;
}

impl<F, A, B, S, T, G, U, V> GEMV<Pair<U, V>> for DiagOp<S, T>
where
    F : Num + AXPY,
    G : Num + AXPY,
    A : AXPY<Field=F>,
    B : AXPY<Field=F>,
    U : AXPY<Field=G>,
    V : AXPY<Field=G>,
    S : GEMV<A, LinCodomain=U>,
    T : GEMV<B, LinCodomain=V>,
{
    fn gemv<I : Instance<Pair<U, V>>>(&self, y : &mut Pair<A, B>, α : G, x : I, β : G) {
        let Pair(u, v) = x.decompose();
        self.0.gemv(&mut y.0, α, u, β);
        self.1.gemv(&mut y.1, α, v, β);
    }

    fn apply_mut<I : Instance<Pair<U, V>>>(&self, y : &mut Pair<A, B>, x : I){
        let Pair(u, v) = x.decompose();
        self.0.apply_mut(&mut y.0, u);
        self.1.apply_mut(&mut y.1, v);
    }

    /// Computes `y += Ax`, where `A` is `Self`
    fn apply_add<I : Instance<Pair<U, V>>>(&self, y : &mut Pair<A, B>, x : I){
        let Pair(u, v) = x.decompose();
        self.0.apply_add(&mut y.0, u);
        self.1.apply_add(&mut y.1, v);
    }
}

impl<F, A, B, S, T> Adjointable<Pair<A,B>, F> for DiagOp<S, T>
where
    F: Float,
    A : HasDual<F>,
    B : HasDual<F>,
    S : Adjointable<A, F>,
    T : Adjointable<B, F>,
    T::Codomain : HasDual<F>,
    S::Codomain : HasDual<F>,
{
    type AdjointCodomain = Pair<S::AdjointCodomain, T::AdjointCodomain>;
    type Adjoint<'a> = DiagOp<S::Adjoint<'a>, T::Adjoint<'a>> where Self : 'a;

    fn adjoint(&self) -> Self::Adjoint<'_> {
        DiagOp(self.0.adjoint(), self.1.adjoint())
    }
}

impl<A, B, Xʹ, Yʹ, S, T> Preadjointable<Pair<A,B>, Pair<Xʹ,Yʹ>> for DiagOp<S, T>
where
    A : AXPY,
    B : AXPY,
    Xʹ : AXPY,
    Yʹ : AXPY,
    S : Preadjointable<A, Xʹ>,
    T : Preadjointable<B, Yʹ>,
{
    type PreadjointCodomain = Pair<S::PreadjointCodomain, T::PreadjointCodomain>;
    type Preadjoint<'a> = DiagOp<S::Preadjoint<'a>, T::Preadjoint<'a>> where Self : 'a;

    fn preadjoint(&self) -> Self::Preadjoint<'_> {
        DiagOp(self.0.preadjoint(), self.1.preadjoint())
    }
}

/// Block operator
pub type BlockOp<S11, S12, S21, S22> = ColOp<RowOp<S11, S12>, RowOp<S21, S22>>;


macro_rules! pairnorm {
    ($expj:ty) => {
        impl<F, A, B, S, T, ExpA, ExpB, ExpR>
        BoundedLinear<Pair<A, B>, PairNorm<ExpA, ExpB, $expj>, ExpR>
        for RowOp<S, T>
        where
            F : Float,
            A : AXPY + Norm<F, ExpA>,
            B : AXPY + Norm<F, ExpB>,
            S : BoundedLinear<A, ExpA, ExpR, F>,
            T : BoundedLinear<B, ExpB, ExpR, F>,
            S::Codomain : Add<T::Codomain>,
            <S::Codomain as Add<T::Codomain>>::Output : AXPY,
            ExpA : NormExponent,
            ExpB : NormExponent,
            ExpR : NormExponent,
        {
            fn opnorm_bound(
                &self,
                PairNorm(expa, expb, _) : PairNorm<ExpA, ExpB, $expj>,
                expr : ExpR
            ) -> <Pair<A, B> as AXPY>::Field {
                // An application of the triangle inequality bounds the norm by the maximum
                // of the individual norms. A simple observation shows this to be exact.
                let na = self.0.opnorm_bound(expa, expr);
                let nb = self.1.opnorm_bound(expb, expr);
                na.max(nb)
            }
        }
        
        impl<A, S, T, ExpA, ExpS, ExpT>
        BoundedLinear<A, ExpA, PairNorm<ExpS, ExpT, $expj>>
        for ColOp<S, T>
        where
            F : Float,
            A : AXPY + Norm<F, ExpA>,
            S : BoundedLinear<A, ExpA, ExpS, F>,
            T : BoundedLinear<A, ExpA, ExpT, F>,
            ExpA : NormExponent,
            ExpS : NormExponent,
            ExpT : NormExponent,
        {
            fn opnorm_bound(
                &self,
                expa : ExpA,
                PairNorm(exps, expt, _) : PairNorm<ExpS, ExpT, $expj>
            ) -> A::Field {
                // This is based on the rule for RowOp and ‖A^*‖ = ‖A‖, hence,
                // for A=[S; T], ‖A‖=‖[S^*, T^*]‖ ≤ max{‖S^*‖, ‖T^*‖} = max{‖S‖, ‖T‖}
                let ns = self.0.opnorm_bound(expa, exps);
                let nt = self.1.opnorm_bound(expa, expt);
                ns.max(nt)
            }
        }
    }
}

pairnorm!(L1);
pairnorm!(L2);
pairnorm!(Linfinity);