src/forward_model/sensor_grid.rs

/*!
Sensor grid forward model
*/

use numeric_literals::replace_float_literals;
use nalgebra::base::{
    DMatrix,
    DVector
};
use std::iter::Zip;
use std::ops::RangeFrom;

pub use alg_tools::linops::*;
use alg_tools::norms::{
    L1, Linfinity, L2, Norm
};
use alg_tools::bisection_tree::*;
use alg_tools::mapping::{
    RealMapping,
    DifferentiableMapping
};
use alg_tools::lingrid::*;
use alg_tools::iter::{MapX, Mappable};
use alg_tools::nalgebra_support::ToNalgebraRealField;
use alg_tools::tabledump::write_csv;
use alg_tools::error::DynError;
use alg_tools::maputil::map2;
use alg_tools::instance::Instance;

use crate::types::*;
use crate::measures::{DiscreteMeasure, Radon};
use crate::seminorms::{
    ConvolutionOp,
    SimpleConvolutionKernel,
};
use crate::kernels::{
    Convolution,
    AutoConvolution,
    BoundedBy,
};
use crate::types::L2Squared;
use crate::transport::TransportLipschitz;
use crate::preadjoint_helper::PreadjointHelper;
use super::{
    ForwardModel,
    LipschitzValues,
    AdjointProductBoundedBy
};
use crate::frank_wolfe::FindimQuadraticModel;

type RNDM<F, const N : usize> = DiscreteMeasure<Loc<F,N>, F>;

pub type ShiftedSensor<F, S, P, const N : usize> = Shift<Convolution<S, P>, F, N>;

/// Trait for physical convolution models. Has blanket implementation for all cases.
pub trait Spread<F : Float, const N : usize>
: 'static + Clone + Support<F, N> + RealMapping<F, N> + Bounded<F> {}

impl<F, T, const N : usize> Spread<F, N> for T
where F : Float,
      T : 'static + Clone + Support<F, N> + Bounded<F> + RealMapping<F, N> {}

/// Trait for compactly supported sensors. Has blanket implementation for all cases.
pub trait Sensor<F : Float, const N : usize> : Spread<F, N> + Norm<F, L1> + Norm<F, Linfinity> {}

impl<F, T, const N : usize> Sensor<F, N> for T
where F : Float,
      T : Spread<F, N> + Norm<F, L1> + Norm<F, Linfinity> {}


pub trait SensorGridBT<F, S, P, const N : usize> :
Clone + BTImpl<F, N, Data=usize, Agg=Bounds<F>>
where F : Float,
      S : Sensor<F, N>,
      P : Spread<F, N> {}

impl<F, S, P, T, const N : usize>
SensorGridBT<F, S, P, N>
for T
where T : Clone + BTImpl<F, N, Data=usize, Agg=Bounds<F>>,
      F : Float,
      S : Sensor<F, N>,
      P : Spread<F, N> {}

// We need type alias bounds to access associated types
#[allow(type_alias_bounds)]
pub type SensorGridBTFN<F, S, P, BT : SensorGridBT<F, S, P, N>, const N : usize>
= BTFN<F, SensorGridSupportGenerator<F, S, P, N>, BT, N>;

/// Sensor grid forward model
#[derive(Clone)]
pub struct SensorGrid<F, S, P, BT, const N : usize>
where F : Float,
      S : Sensor<F, N>,
      P : Spread<F, N>,
      Convolution<S, P> : Spread<F, N>,
      BT : SensorGridBT<F, S, P, N>,
{
    domain : Cube<F, N>,
    sensor_count : [usize; N],
    sensor : S,
    spread : P,
    base_sensor : Convolution<S, P>,
    bt : BT,
}

impl<F, S, P, BT, const N : usize> SensorGrid<F, S, P, BT, N>
where F : Float,
      BT : SensorGridBT<F, S, P, N>,
      S : Sensor<F, N>,
      P : Spread<F, N>,
      Convolution<S, P> : Spread<F, N> + LocalAnalysis<F, BT::Agg, N>,
{

    /// Create a new sensor grid.
    ///
    /// The parameter `depth` indicates the search depth of the created [`BT`]s
    /// for the adjoint values.
    pub fn new(
        domain : Cube<F, N>,
        sensor_count : [usize; N],
        sensor : S,
        spread : P,
        depth : BT::Depth
    ) -> Self {
        let base_sensor = Convolution(sensor.clone(), spread.clone());
        let bt = BT::new(domain, depth);
        let mut sensorgrid = SensorGrid {
            domain,
            sensor_count,
            sensor,
            spread,
            base_sensor,
            bt,
        };

        for (x, id) in sensorgrid.grid().into_iter().zip(0usize..) {
            let s = sensorgrid.shifted_sensor(x);
            sensorgrid.bt.insert(id, &s);
        }

        sensorgrid
    }
}


impl<F, S, P, BT, const N : usize> SensorGrid<F, S, P, BT, N>
where F : Float,
      BT : SensorGridBT<F, S, P, N>,
      S : Sensor<F, N>,
      P : Spread<F, N>,
      Convolution<S, P> : Spread<F, N>
{
      
    /// Return the grid of sensor locations.
    pub fn grid(&self) -> LinGrid<F, N> {
        lingrid_centered(&self.domain, &self.sensor_count)
    }

    /// Returns the number of sensors (number of grid points)
    pub fn n_sensors(&self) -> usize {
        self.sensor_count.iter().product()
    }

    /// Constructs a sensor shifted by `x`.
    #[inline]
    fn shifted_sensor(&self, x : Loc<F, N>) -> ShiftedSensor<F, S, P, N> {
        self.base_sensor.clone().shift(x)
    }

    #[inline]
    fn _zero_observable(&self) -> DVector<F> {
        DVector::zeros(self.n_sensors())
    }

    /// Returns the maximum number of overlapping sensors $N_\psi$.
    pub fn max_overlapping(&self) -> F {
        let w = self.base_sensor.support_hint().width();
        let d = map2(self.domain.width(), &self.sensor_count, |wi, &i| wi/F::cast_from(i));
        w.iter()
         .zip(d.iter())
         .map(|(&wi, &di)| (wi/di).ceil())
         .reduce(F::mul)
         .unwrap()
    }
}

impl<F, S, P, BT, const N : usize> Mapping<RNDM<F, N>> for SensorGrid<F, S, P, BT, N>
where
    F : Float,
    BT : SensorGridBT<F, S, P, N>,
    S : Sensor<F, N>,
    P : Spread<F, N>,
    Convolution<S, P> : Spread<F, N>,
{

    type Codomain =  DVector<F>;

    #[inline]
    fn apply<I : Instance<RNDM<F, N>>>(&self, μ : I) -> DVector<F> {
        let mut y = self._zero_observable();
        self.apply_add(&mut y, μ);
        y
    }
}


impl<F, S, P, BT, const N : usize> Linear<RNDM<F, N>> for SensorGrid<F, S, P, BT, N>
where
    F : Float,
    BT : SensorGridBT<F, S, P, N>,
    S : Sensor<F, N>,
    P : Spread<F, N>,
    Convolution<S, P> : Spread<F, N>,
{ }


#[replace_float_literals(F::cast_from(literal))]
impl<F, S, P, BT, const N : usize> GEMV<F, RNDM<F, N>, DVector<F>> for SensorGrid<F, S, P, BT, N>
where F : Float,
      BT : SensorGridBT<F, S, P, N>,
      S : Sensor<F, N>,
      P : Spread<F, N>,
      Convolution<S, P> : Spread<F, N>,
{

    fn gemv<I : Instance<RNDM<F, N>>>(
        &self, y : &mut DVector<F>, α : F, μ : I, β : F
    ) {
        let grid = self.grid();
        if β == 0.0 {
            y.fill(0.0)
        } else if β != 1.0 {
            *y *= β; // Need to multiply first, as we have to be able to add to y.
        }
        if α == 1.0 {
            self.apply_add(y, μ)
        } else {
            for δ in μ.ref_instance() {
                for &d in self.bt.iter_at(&δ.x) {
                    let sensor = self.shifted_sensor(grid.entry_linear_unchecked(d));
                    y[d] += sensor.apply(&δ.x) * (α * δ.α);
                }
            }
        }
    }

    fn apply_add<I : Instance<RNDM<F, N>>>(
        &self, y : &mut DVector<F>, μ : I
    ) {
        let grid = self.grid();
        for δ in μ.ref_instance() {
            for &d in self.bt.iter_at(&δ.x) {
                let sensor = self.shifted_sensor(grid.entry_linear_unchecked(d));
                y[d] += sensor.apply(&δ.x) * δ.α;
            }
        }
    }

}


impl<F, S, P, BT, const N : usize>
BoundedLinear<RNDM<F, N>, Radon, L2, F>
for SensorGrid<F, S, P, BT, N>
where F : Float,
      BT : SensorGridBT<F, S, P, N, Agg=Bounds<F>>,
      S : Sensor<F, N>,
      P : Spread<F, N>,
      Convolution<S, P> : Spread<F, N> + LocalAnalysis<F, BT::Agg, N>
{

    /// An estimate on the operator norm in $𝕃(ℳ(Ω); ℝ^n)$ with $ℳ(Ω)$ equipped
    /// with the Radon norm, and $ℝ^n$ with the Euclidean norm.
    fn opnorm_bound(&self, _ : Radon, _ : L2) -> F {
        // With {x_i}_{i=1}^n the grid centres and φ the kernel, we have
        // |Aμ|_2 = sup_{|z|_2 ≤ 1} ⟨z,Αμ⟩ = sup_{|z|_2 ≤ 1} ⟨A^*z|μ⟩
        // ≤ sup_{|z|_2 ≤ 1} |A^*z|_∞ |μ|_ℳ
        // = sup_{|z|_2 ≤ 1} |∑ φ(· - x_i)z_i|_∞ |μ|_ℳ
        // ≤ sup_{|z|_2 ≤ 1} |φ(y)| ∑_{i:th sensor active at y}|z_i| |μ|_ℳ
        //      where the supremum of |∑ φ(· - x_i)z_i|_∞  is reached at y
        // ≤ sup_{|z|_2 ≤ 1} |φ|_∞ √N_ψ |z|_2 |μ|_ℳ
        //      where N_ψ is the maximum number of sensors that overlap, and
        //      |z|_2 is restricted to the active sensors.
        // = |φ|_∞ √N_ψ |μ|_ℳ.
        // Hence
        let n = self.max_overlapping();
        self.base_sensor.bounds().uniform() * n.sqrt()
    }
}

type SensorGridPreadjoint<'a, A, F, const N : usize> = PreadjointHelper<'a, A, RNDM<F,N>>;


impl<F, S, P, BT, const N : usize>
Preadjointable<RNDM<F, N>, DVector<F>>
for SensorGrid<F, S, P, BT, N>
where F : Float,
      BT : SensorGridBT<F, S, P, N>,
      S : Sensor<F, N>,
      P : Spread<F, N>,
      Convolution<S, P> : Spread<F, N> + LocalAnalysis<F, BT::Agg, N>
{
    type PreadjointCodomain = BTFN<F, SensorGridSupportGenerator<F, S, P, N>, BT, N>;
    type Preadjoint<'a> = SensorGridPreadjoint<'a, Self, F, N> where Self : 'a;

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

#[replace_float_literals(F::cast_from(literal))]
impl<'a, F, S, P, BT, const N : usize> LipschitzValues
for SensorGridPreadjoint<'a, SensorGrid<F, S, P, BT, N>, F, N>
where F : Float,
      BT : SensorGridBT<F, S, P, N>,
      S : Sensor<F, N>,
      P : Spread<F, N>,
      Convolution<S, P> : Spread<F, N> + Lipschitz<L2, FloatType=F> + DifferentiableMapping<Loc<F,N>> + LocalAnalysis<F, BT::Agg, N>,
      for<'b> <Convolution<S, P> as DifferentiableMapping<Loc<F,N>>>::Differential<'b> : Lipschitz<L2, FloatType=F>,
{
    
    type FloatType = F;

    fn value_unit_lipschitz_factor(&self) -> Option<F> {
        // The Lipschitz factor of the sensors has to be scaled by the square root of twice
        // the number of overlapping sensors at a single ponit, as Lipschitz estimates involve
        // two points.
        let fw = self.forward_op;
        let n = fw.max_overlapping();
        fw.base_sensor.lipschitz_factor(L2).map(|l| (2.0 * n).sqrt() * l)
    }

    fn value_diff_unit_lipschitz_factor(&self) -> Option<F> {
        // The Lipschitz factor of the sensors has to be scaled by the square root of twice
        // the number of overlapping sensors at a single ponit, as Lipschitz estimates involve
        // two points.
        let fw = self.forward_op;
        let n = fw.max_overlapping();
        fw.base_sensor.diff_ref().lipschitz_factor(L2).map(|l| (2.0 * n).sqrt() * l)
    }
}

#[derive(Clone,Debug)]
pub struct SensorGridSupportGenerator<F, S, P, const N : usize>
where F : Float,
      S : Sensor<F, N>,
      P : Spread<F, N>
{
    base_sensor : Convolution<S, P>,
    grid : LinGrid<F, N>,
    weights : DVector<F>
}

impl<F, S, P, const N : usize> SensorGridSupportGenerator<F, S, P, N>
where F : Float,
      S : Sensor<F, N>,
      P : Spread<F, N>,
      Convolution<S, P> : Spread<F, N>
{

    #[inline]
    fn construct_sensor(&self, id : usize, w : F) -> Weighted<ShiftedSensor<F, S, P, N>, F> {
        let x = self.grid.entry_linear_unchecked(id);
        self.base_sensor.clone().shift(x).weigh(w)
    }

    #[inline]
    fn construct_sensor_and_id<'a>(&'a self, (id, w) : (usize, &'a F))
    -> (usize, Weighted<ShiftedSensor<F, S, P, N>, F>) {
        (id.into(), self.construct_sensor(id, *w))
    }
}

impl<F, S, P, const N : usize> SupportGenerator<F, N>
for SensorGridSupportGenerator<F, S, P, N>
where F : Float,
      S : Sensor<F, N>,
      P : Spread<F, N>,
      Convolution<S, P> : Spread<F, N>
{
    type Id = usize;
    type SupportType = Weighted<ShiftedSensor<F, S, P, N>, F>;
    type AllDataIter<'a> = MapX<'a, Zip<RangeFrom<usize>,
                                        std::slice::Iter<'a, F>>,
                                Self,
                                (Self::Id, Self::SupportType)>
                           where Self : 'a;

    #[inline]
    fn support_for(&self, d : Self::Id) -> Self::SupportType {
        self.construct_sensor(d, self.weights[d])
    }

    #[inline]
    fn support_count(&self) -> usize {
        self.weights.len()
    }

    #[inline]
    fn all_data(&self) -> Self::AllDataIter<'_> {
        (0..).zip(self.weights.as_slice().iter()).mapX(self, Self::construct_sensor_and_id)
    }
}

impl<F, S, P, BT, const N : usize> ForwardModel<DiscreteMeasure<Loc<F, N>, F>, F>
for SensorGrid<F, S, P, BT, N>
where F : Float + ToNalgebraRealField<MixedType=F> + nalgebra::RealField,
      BT : SensorGridBT<F, S, P, N>,
      S : Sensor<F, N>,
      P : Spread<F, N>,
      Convolution<S, P> : Spread<F, N> + LocalAnalysis<F, BT::Agg, N>,
{
    type Observable = DVector<F>;

    fn write_observable(&self, b : &Self::Observable, prefix : String) -> DynError {
        let it = self.grid().into_iter().zip(b.iter()).map(|(x, &v)| (x, v));
        write_csv(it, prefix + ".txt")
    }

    #[inline]
    fn zero_observable(&self) -> Self::Observable {
        self._zero_observable()
    }
}

impl<F, S, P, BT, const N : usize> FindimQuadraticModel<Loc<F, N>, F>
for SensorGrid<F, S, P, BT, N>
where F : Float + ToNalgebraRealField<MixedType=F> + nalgebra::RealField,
      BT : SensorGridBT<F, S, P, N>,
      S : Sensor<F, N>,
      P : Spread<F, N>,
      Convolution<S, P> : Spread<F, N> + LocalAnalysis<F, BT::Agg, N>
{

    fn findim_quadratic_model(
        &self,
        μ : &DiscreteMeasure<Loc<F, N>, F>,
        b : &Self::Observable
    ) -> (DMatrix<F::MixedType>, DVector<F::MixedType>) {
        assert_eq!(b.len(), self.n_sensors());
        let mut mA = DMatrix::zeros(self.n_sensors(), μ.len());
        let grid = self.grid();
        for (mut mAcol, δ) in mA.column_iter_mut().zip(μ.iter_spikes()) {
            for &d in self.bt.iter_at(&δ.x) {
                let sensor = self.shifted_sensor(grid.entry_linear_unchecked(d));
                mAcol[d] += sensor.apply(&δ.x);
            }
        }
        let mAt = mA.transpose();
        (&mAt * mA, &mAt * b)
    }
}

/// Implements the calculation a factor $L$ such that $A_*A ≤ L 𝒟$ for $A$ the forward model
/// and $𝒟$ a seminorm of suitable form.
///
/// **This assumes (but does not check) that the sensors are not overlapping.**
#[replace_float_literals(F::cast_from(literal))]
impl<F, BT, S, P, K, const N : usize>
AdjointProductBoundedBy<RNDM<F, N>, ConvolutionOp<F, K, BT, N>>
for SensorGrid<F, S, P, BT, N>
where F : Float + nalgebra::RealField + ToNalgebraRealField,
      BT : SensorGridBT<F, S, P, N>,
      S : Sensor<F, N>,
      P : Spread<F, N>,
      Convolution<S, P> : Spread<F, N>,
      K : SimpleConvolutionKernel<F, N>,
      AutoConvolution<P> : BoundedBy<F, K>
{

    type FloatType = F;

    fn adjoint_product_bound(&self, seminorm : &ConvolutionOp<F, K, BT, N>) -> Option<F> {
        // Sensors should not take on negative values to allow
        // A_*A to be upper bounded by a simple convolution of `spread`.
        if self.sensor.bounds().lower() < 0.0 {
            return None
        }

        // Calculate the factor $L_1$ for betwee $ℱ[ψ * ψ] ≤ L_1 ℱ[ρ]$ for $ψ$ the base spread
        // and $ρ$ the kernel of the seminorm.
        let l1 = AutoConvolution(self.spread.clone()).bounding_factor(seminorm.kernel())?;

        // Calculate the factor for transitioning from $A_*A$ to `AutoConvolution<P>`, where A
        // consists of several `Convolution<S, P>` for the physical model `P` and the sensor `S`.
        let l0 = self.sensor.norm(Linfinity) * self.sensor.norm(L1);

        // The final transition factor is:
        Some(l0 * l1)
    }
}

#[replace_float_literals(F::cast_from(literal))]
impl<F, BT, S, P, const N : usize> TransportLipschitz<L2Squared>
for SensorGrid<F, S, P, BT, N>
where F : Float + ToNalgebraRealField,
      BT : SensorGridBT<F, S, P, N>,
      S : Sensor<F, N>,
      P : Spread<F, N>,
      Convolution<S, P> : Spread<F, N> + Lipschitz<L2, FloatType = F>
{
    type FloatType = F;

    fn transport_lipschitz_factor(&self, L2Squared : L2Squared) -> Self::FloatType {
        // We estimate the factor by N_ψL^2, where L is the 2-norm Lipschitz factor of
        // the base sensor (sensor * base_spread), and N_ψ the maximum overlap.
        // The factors two comes from Lipschitz estimates having two possible
        // points of overlap.
        let l = self.base_sensor.lipschitz_factor(L2).unwrap();
        2.0 * self.max_overlapping() * l.powi(2)
    }
}


macro_rules! make_sensorgridsupportgenerator_scalarop_rhs {
    ($trait:ident, $fn:ident, $trait_assign:ident, $fn_assign:ident) => {
        impl<F, S, P, const N : usize>
        std::ops::$trait_assign<F>
        for SensorGridSupportGenerator<F, S, P, N>
        where F : Float,
              S : Sensor<F, N>,
              P : Spread<F, N>,
              Convolution<S, P> : Spread<F, N> {
            fn $fn_assign(&mut self, t : F) {
                self.weights.$fn_assign(t);
            }
        }

        impl<F, S, P, const N : usize>
        std::ops::$trait<F>
        for SensorGridSupportGenerator<F, S, P, N>
        where F : Float,
              S : Sensor<F, N>,
              P : Spread<F, N>,
              Convolution<S, P> : Spread<F, N> {
            type Output = SensorGridSupportGenerator<F, S, P, N>;
            fn $fn(mut self, t : F) -> Self::Output {
                std::ops::$trait_assign::$fn_assign(&mut self.weights, t);
                self
            }
        }

        impl<'a, F, S, P, const N : usize>
        std::ops::$trait<F>
        for &'a SensorGridSupportGenerator<F, S, P, N>
        where F : Float,
              S : Sensor<F, N>,
              P : Spread<F, N>,
              Convolution<S, P> : Spread<F, N> {
            type Output = SensorGridSupportGenerator<F, S, P, N>;
            fn $fn(self, t : F) -> Self::Output {
                SensorGridSupportGenerator{
                    base_sensor : self.base_sensor.clone(),
                    grid : self.grid,
                    weights : (&self.weights).$fn(t)
                }
            }
        }
    }
}

make_sensorgridsupportgenerator_scalarop_rhs!(Mul, mul, MulAssign, mul_assign);
make_sensorgridsupportgenerator_scalarop_rhs!(Div, div, DivAssign, div_assign);

macro_rules! make_sensorgridsupportgenerator_unaryop {
    ($trait:ident, $fn:ident) => {
        impl<F, S, P, const N : usize>
        std::ops::$trait
        for SensorGridSupportGenerator<F, S, P, N>
        where F : Float,
              S : Sensor<F, N>,
              P : Spread<F, N>,
              Convolution<S, P> : Spread<F, N> {
            type Output = SensorGridSupportGenerator<F, S, P, N>;
            fn $fn(mut self) -> Self::Output {
                self.weights = self.weights.$fn();
                self
            }
        }

        impl<'a, F, S, P, const N : usize>
        std::ops::$trait
        for &'a SensorGridSupportGenerator<F, S, P, N>
        where F : Float,
              S : Sensor<F, N>,
              P : Spread<F, N>,
              Convolution<S, P> : Spread<F, N> {
            type Output = SensorGridSupportGenerator<F, S, P, N>;
            fn $fn(self) -> Self::Output {
                SensorGridSupportGenerator{
                    base_sensor : self.base_sensor.clone(),
                    grid : self.grid,
                    weights : (&self.weights).$fn()
                }
            }
        }
    }
}

make_sensorgridsupportgenerator_unaryop!(Neg, neg);

impl<'a, F, S, P, BT, const N : usize> Mapping<DVector<F>>
for PreadjointHelper<'a, SensorGrid<F, S, P, BT, N>, RNDM<F,N>>
where F : Float,
      BT : SensorGridBT<F, S, P, N>,
      S : Sensor<F, N>,
      P : Spread<F, N>,
      Convolution<S, P> : Spread<F, N> + LocalAnalysis<F, Bounds<F>, N>,
{

    type Codomain = SensorGridBTFN<F, S, P, BT, N>;

    fn apply<I : Instance<DVector<F>>>(&self, x : I) -> Self::Codomain {
        let fwd = &self.forward_op;
        let generator = SensorGridSupportGenerator{
            base_sensor : fwd.base_sensor.clone(),
            grid : fwd.grid(),
            weights : x.own()
        };
        BTFN::new_refresh(&fwd.bt, generator)
    }
}

impl<'a, F, S, P, BT, const N : usize> Linear<DVector<F>>
for PreadjointHelper<'a, SensorGrid<F, S, P, BT, N>, RNDM<F,N>>
where F : Float,
      BT : SensorGridBT<F, S, P, N>,
      S : Sensor<F, N>,
      P : Spread<F, N>,
      Convolution<S, P> : Spread<F, N> + LocalAnalysis<F, Bounds<F>, N>,
{ }