src/forward_model/sensor_grid.rs

/*!
Sensor grid forward model
*/

use super::{BasicCurvatureBoundEstimates, ForwardModel};
use crate::frank_wolfe::FindimQuadraticModel;
use crate::kernels::{AutoConvolution, BoundedBy, Convolution};
use crate::measures::{DiscreteMeasure, Radon, RNDM};
use crate::preadjoint_helper::PreadjointHelper;
use crate::seminorms::{ConvolutionOp, SimpleConvolutionKernel};
use crate::types::*;
use alg_tools::bisection_tree::*;
use alg_tools::bounds::Bounded;
use alg_tools::error::{DynError, DynResult};
use alg_tools::euclidean::Euclidean;
use alg_tools::instance::Instance;
use alg_tools::iter::{MapX, Mappable};
use alg_tools::lingrid::*;
pub use alg_tools::linops::*;
use alg_tools::mapping::{DifferentiableMapping, RealMapping};
use alg_tools::maputil::map2;
use alg_tools::nalgebra_support::ToNalgebraRealField;
use alg_tools::norms::{Linfinity, Norm, L1, L2};
use alg_tools::tabledump::write_csv;
use anyhow::anyhow;
use nalgebra::base::{DMatrix, DVector};
use numeric_literals::replace_float_literals;
use std::iter::Zip;
use std::ops::RangeFrom;

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

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

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

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

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

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

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

// 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<S, F, 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<N, F>,
    P: Spread<N, F>,
    Convolution<S, P>: Spread<N, F>,
    BT: SensorGridBT<F, S, P, N>,
{
    domain: Cube<N, F>,
    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<N, F>,
    P: Spread<N, F>,
    Convolution<S, P>: Spread<N, F> + 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<N, F>,
        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<N, F>,
    P: Spread<N, F>,
    Convolution<S, P>: Spread<N, F>,
{
    /// Return the grid of sensor locations.
    pub fn grid(&self) -> LinGrid<N, F> {
        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<N, F>) -> 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<N, F>> for SensorGrid<F, S, P, BT, N>
where
    F: Float,
    BT: SensorGridBT<F, S, P, N>,
    S: Sensor<N, F>,
    P: Spread<N, F>,
    Convolution<S, P>: Spread<N, F>,
{
    type Codomain = DVector<F>;

    #[inline]
    fn apply<I: Instance<RNDM<N, F>>>(&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<N, F>> for SensorGrid<F, S, P, BT, N>
where
    F: Float,
    BT: SensorGridBT<F, S, P, N>,
    S: Sensor<N, F>,
    P: Spread<N, F>,
    Convolution<S, P>: Spread<N, F>,
{
}

#[replace_float_literals(F::cast_from(literal))]
impl<F, S, P, BT, const N: usize> GEMV<F, RNDM<N, F>, DVector<F>> for SensorGrid<F, S, P, BT, N>
where
    F: Float,
    BT: SensorGridBT<F, S, P, N>,
    S: Sensor<N, F>,
    P: Spread<N, F>,
    Convolution<S, P>: Spread<N, F>,
{
    fn gemv<I: Instance<RNDM<N, F>>>(&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 {
            μ.eval_ref(|μr| {
                for δ in μr {
                    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<N, F>>>(&self, y: &mut DVector<F>, μ: I) {
        let grid = self.grid();
        μ.eval_ref(|μr| {
            for δ in μr {
                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<N, F>, Radon, L2, F>
    for SensorGrid<F, S, P, BT, N>
where
    F: Float,
    BT: SensorGridBT<F, S, P, N>,
    S: Sensor<N, F>,
    P: Spread<N, F>,
    Convolution<S, P>: Spread<N, F>,
{
    /// 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) -> DynResult<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();
        Ok(self.base_sensor.bounds().uniform() * n.sqrt())
    }
}

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

impl<F, S, P, BT, const N: usize> Preadjointable<RNDM<N, F>, DVector<F>>
    for SensorGrid<F, S, P, BT, N>
where
    F: Float,
    BT: SensorGridBT<F, S, P, N>,
    S: Sensor<N, F>,
    P: Spread<N, F>,
    Convolution<S, P>: Spread<N, F> + LocalAnalysis<F, BT::Agg, N>,
{
    type PreadjointCodomain = BTFN<F, SensorGridSupportGenerator<S, F, 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<N, F>,
      P : Spread<N, F>,
      Convolution<S, P> : Spread<N, F> + Lipschitz<L2, FloatType=F> + DifferentiableMapping<Loc<N, F>> + LocalAnalysis<F, BT::Agg, N>,
      for<'b> <Convolution<S, P> as DifferentiableMapping<Loc<N, F>>>::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)
    }
}
*/

#[replace_float_literals(F::cast_from(literal))]
impl<'a, F, S, P, BT, const N: usize> BasicCurvatureBoundEstimates<F> for SensorGrid<F, S, P, BT, N>
where
    F: Float + ToNalgebraRealField,
    BT: SensorGridBT<F, S, P, N>,
    S: Sensor<N, F>,
    P: Spread<N, F>,
    DVector<F>: Euclidean<F>,
    Convolution<S, P>: Spread<N, F>
        + Lipschitz<L2, FloatType = F>
        + DifferentiableMapping<Loc<N, F>>
        + LocalAnalysis<F, BT::Agg, N>,
    for<'b> <Convolution<S, P> as DifferentiableMapping<Loc<N, F>>>::Differential<'b>:
        Lipschitz<L2, FloatType = F>,
{
    fn basic_curvature_bound_components(&self) -> (DynResult<F>, DynResult<F>) {
        let n_ψ = self.max_overlapping();
        let ψ_diff_lip = self.base_sensor.diff_ref().lipschitz_factor(L2);
        let ψ_lip = self.base_sensor.lipschitz_factor(L2);
        let ℓ_F0 = ψ_diff_lip.map(|l| (2.0 * n_ψ).sqrt() * l);
        let Θ2 = ψ_lip.map(|l| 4.0 * n_ψ * l.powi(2));

        (ℓ_F0, Θ2)
    }
}

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

impl<F, S, P, const N: usize> SensorGridSupportGenerator<S, F, P, N>
where
    F: Float,
    S: Sensor<N, F>,
    P: Spread<N, F>,
    Convolution<S, P>: Spread<N, F>,
{
    #[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<N, F> for SensorGridSupportGenerator<S, F, P, N>
where
    F: Float,
    S: Sensor<N, F>,
    P: Spread<N, F>,
    Convolution<S, P>: Spread<N, F>,
{
    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<N, F>, 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<N, F>,
    P: Spread<N, F>,
    Convolution<S, P>: Spread<N, F> + 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<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<N, F>,
    P: Spread<N, F>,
    Convolution<S, P>: Spread<N, F> + LocalAnalysis<F, BT::Agg, N>,
{
    fn findim_quadratic_model(
        &self,
        μ: &DiscreteMeasure<Loc<N, F>, 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<'a, F, BT, S, P, K, const N: usize>
    BoundedLinear<RNDM<N, F>, &'a ConvolutionOp<F, K, BT, N>, L2, F> for SensorGrid<F, S, P, BT, N>
where
    F: Float + ToNalgebraRealField,
    BT: SensorGridBT<F, S, P, N>,
    S: Sensor<N, F>,
    P: Spread<N, F>,
    Convolution<S, P>: Spread<N, F>,
    K: SimpleConvolutionKernel<N, F>,
    AutoConvolution<P>: BoundedBy<F, K>,
    Weighted<Shift<K, N, F>, F>: LocalAnalysis<F, BT::Agg, N>,
{
    fn opnorm_bound(&self, seminorm: &'a ConvolutionOp<F, K, BT, N>, _: L2) -> DynResult<F> {
        // Sensors should not take on negative values to allow
        // A_*A to be upper bounded by a simple convolution of `spread`.
        // TODO: Do we really need this restriction?
        if self.sensor.bounds().lower() < 0.0 {
            return Err(anyhow!("Sensor not bounded from below by zero"));
        }

        // 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:
        Ok((l0 * l1).sqrt())
    }
}

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<S, F, P, N>
        where
            F: Float,
            S: Sensor<N, F>,
            P: Spread<N, F>,
            Convolution<S, P>: Spread<N, F>,
        {
            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<S, F, P, N>
        where
            F: Float,
            S: Sensor<N, F>,
            P: Spread<N, F>,
            Convolution<S, P>: Spread<N, F>,
        {
            type Output = SensorGridSupportGenerator<S, F, 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<S, F, P, N>
        where
            F: Float,
            S: Sensor<N, F>,
            P: Spread<N, F>,
            Convolution<S, P>: Spread<N, F>,
        {
            type Output = SensorGridSupportGenerator<S, F, 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<S, F, P, N>
        where
            F: Float,
            S: Sensor<N, F>,
            P: Spread<N, F>,
            Convolution<S, P>: Spread<N, F>,
        {
            type Output = SensorGridSupportGenerator<S, F, 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<S, F, P, N>
        where
            F: Float,
            S: Sensor<N, F>,
            P: Spread<N, F>,
            Convolution<S, P>: Spread<N, F>,
        {
            type Output = SensorGridSupportGenerator<S, F, 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<N, F>>
where
    F: Float,
    BT: SensorGridBT<F, S, P, N>,
    S: Sensor<N, F>,
    P: Spread<N, F>,
    Convolution<S, P>: Spread<N, F> + 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<N, F>>
where
    F: Float,
    BT: SensorGridBT<F, S, P, N>,
    S: Sensor<N, F>,
    P: Spread<N, F>,
    Convolution<S, P>: Spread<N, F> + LocalAnalysis<F, Bounds<F>, N>,
{
}