src/prox_penalty/radon_squared.rs

/*!
Solver for the point source localisation problem using a simplified forward-backward splitting method.

Instead of the $𝒟$-norm of `fb.rs`, this uses a standard Radon norm for the proximal map.
*/

use numeric_literals::replace_float_literals;
use serde::{Serialize, Deserialize};
use nalgebra::DVector;

use alg_tools::iterate::{
    AlgIteratorIteration,
    AlgIterator
};
use alg_tools::norms::{L2, Norm};
use alg_tools::linops::Mapping;
use alg_tools::bisection_tree::{
    BTFN,
    Bounds,
    BTSearch,
    SupportGenerator,
    LocalAnalysis,
};
use alg_tools::mapping::RealMapping;
use alg_tools::nalgebra_support::ToNalgebraRealField;

use crate::types::*;
use crate::measures::{
    RNDM,
    DeltaMeasure,
    Radon,
};
use crate::measures::merging::SpikeMerging;
use crate::regularisation::RegTerm;
use crate::forward_model::{
    ForwardModel,
    AdjointProductBoundedBy
};
use super::{
    FBGenericConfig,
    ProxPenalty,
};

/// Radon-norm squared proximal penalty

#[derive(Copy,Clone,Serialize,Deserialize)]
pub struct RadonSquared;

#[replace_float_literals(F::cast_from(literal))]
impl<F, GA, BTA, S, Reg, const N : usize>
ProxPenalty<F, BTFN<F, GA, BTA, N>, Reg, N> for RadonSquared
where
    F : Float + ToNalgebraRealField,
    GA : SupportGenerator<F, N, SupportType = S, Id = usize> + Clone,
    BTA : BTSearch<F, N, Data=usize, Agg=Bounds<F>>,
    S: RealMapping<F, N> + LocalAnalysis<F, Bounds<F>, N>,
    Reg : RegTerm<F, N>,
    RNDM<F, N> : SpikeMerging<F>,
{
    type ReturnMapping = BTFN<F, GA, BTA, N>;

    fn insert_and_reweigh<I>(
        &self,
        μ : &mut RNDM<F, N>,
        τv : &mut BTFN<F, GA, BTA, N>,
        μ_base : &RNDM<F, N>,
        ν_delta: Option<&RNDM<F, N>>,
        τ : F,
        ε : F,
        config : &FBGenericConfig<F>,
        reg : &Reg,
        _state : &AlgIteratorIteration<I>,
        stats : &mut IterInfo<F, N>,
    ) -> (Option<Self::ReturnMapping>, bool)
    where
        I : AlgIterator
    {
        let mut y = μ_base.masses_dvector();

        assert!(μ_base.len() <= μ.len());
        
        'i_and_w: for i in 0..=1 {
            // Optimise weights
            if μ.len() > 0 {
                // Form finite-dimensional subproblem. The subproblem references to the original μ^k
                // from the beginning of the iteration are all contained in the immutable c and g.
                // TODO: observe negation of -τv after switch from minus_τv: finite-dimensional
                // problems have not yet been updated to sign change.
                let g̃ = DVector::from_iterator(μ.len(),
                                               μ.iter_locations()
                                                .map(|ζ| - F::to_nalgebra_mixed(τv.apply(ζ))));
                let mut x = μ.masses_dvector();
                y.extend(std::iter::repeat(0.0.to_nalgebra_mixed()).take(0.max(x.len()-y.len())));
                assert_eq!(y.len(), x.len());
                // Solve finite-dimensional subproblem.
                // TODO: This assumes that ν_delta has no common locations with μ-μ_base, to
                // ignore it.
                stats.inner_iters += reg.solve_findim_l1squared(&y, &g̃, τ, &mut x, ε, config);

                // Update masses of μ based on solution of finite-dimensional subproblem.
                μ.set_masses_dvector(&x);
            }

            if i>0 {
                // Simple debugging test to see if more inserts would be needed. Doesn't seem so.
                //let n = μ.dist_matching(μ_base);
                //println!("{:?}", reg.find_tolerance_violation_slack(τv, τ, ε, false, config, n));
                break 'i_and_w
            }
            
            // Calculate ‖μ - μ_base‖_ℳ
            // TODO: This assumes that ν_delta has no common locations with μ-μ_base.
            let n = μ.dist_matching(μ_base) + ν_delta.map_or(0.0, |ν| ν.norm(Radon));
        
            // Find a spike to insert, if needed.
            // This only check the overall tolerances, not tolerances on support of μ-μ_base or μ,
            // which are supposed to have been guaranteed by the finite-dimensional weight optimisation.
            match reg.find_tolerance_violation_slack(τv, τ, ε, false, config, n) {
                None => { break 'i_and_w },
                Some((ξ, _v_ξ, _in_bounds)) => {
                    // Weight is found out by running the finite-dimensional optimisation algorithm
                    // above
                    *μ += DeltaMeasure { x : ξ, α : 0.0 };
                    stats.inserted += 1;
                }
            };
        }

        (None, true)
    }

    fn merge_spikes(
        &self,
        μ : &mut RNDM<F, N>,
        τv : &mut BTFN<F, GA, BTA, N>,
        μ_base : &RNDM<F, N>,
        ν_delta: Option<&RNDM<F, N>>,
        τ : F,
        ε : F,
        config : &FBGenericConfig<F>,
        reg : &Reg,
        fitness : Option<impl Fn(&RNDM<F, N>) -> F>,
    ) -> usize
    {
        if config.fitness_merging {
            if let Some(f) = fitness {
                return μ.merge_spikes_fitness(config.merging, f, |&v| v)
                        .1
            }
        }
        μ.merge_spikes(config.merging, |μ_candidate| {
            // Important: μ_candidate's new points are afterwards,
            // and do not conflict with μ_base.
            // TODO: could simplify to requiring μ_base instead of μ_radon.
            // but may complicate with sliding base's exgtra points that need to be
            // after μ_candidate's extra points.
            // TODO: doesn't seem to work, maybe need to merge μ_base as well?
            // Although that doesn't seem to make sense.
            let μ_radon = match ν_delta {
                None => μ_candidate.sub_matching(μ_base),
                Some(ν) => μ_candidate.sub_matching(μ_base) - ν,
            };
            reg.verify_merge_candidate_radonsq(τv, μ_candidate, τ, ε, &config, &μ_radon)
            //let n = μ_candidate.dist_matching(μ_base);
            //reg.find_tolerance_violation_slack(τv, τ, ε, false, config, n).is_none()
        })
    }
}


impl<F, A, const N : usize> AdjointProductBoundedBy<RNDM<F, N>, RadonSquared>
for A
where
    F : Float,
    A : ForwardModel<RNDM<F, N>, F>
{
    type FloatType = F;

    fn adjoint_product_bound(&self, _ : &RadonSquared) -> Option<Self::FloatType> {
        self.opnorm_bound(Radon, L2).powi(2).into()
    }
}