src/forward_pdps.rs

/*!
Solver for the point source localisation problem using a
primal-dual proximal splitting with a forward step.
*/

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

use alg_tools::iterate::AlgIteratorFactory;
use alg_tools::euclidean::Euclidean;
use alg_tools::mapping::{Mapping, DifferentiableRealMapping, Instance};
use alg_tools::norms::Norm;
use alg_tools::direct_product::Pair;
use alg_tools::nalgebra_support::ToNalgebraRealField;
use alg_tools::linops::{
    BoundedLinear, AXPY, GEMV, Adjointable, IdOp,
};
use alg_tools::convex::{Conjugable, Prox};
use alg_tools::norms::{L2, PairNorm};

use crate::types::*;
use crate::measures::{DiscreteMeasure, Radon, RNDM};
use crate::measures::merging::SpikeMerging;
use crate::forward_model::{
    ForwardModel,
    AdjointProductPairBoundedBy,
};
use crate::plot::{
    SeqPlotter,
    Plotting,
    PlotLookup
};
use crate::fb::*;
use crate::regularisation::RegTerm;
use crate::dataterm::calculate_residual;

/// Settings for [`pointsource_forward_pdps_pair`].
#[derive(Clone, Copy, Eq, PartialEq, Serialize, Deserialize, Debug)]
#[serde(default)]
pub struct ForwardPDPSConfig<F : Float> {
    /// Primal step length scaling.
    pub τ0 : F,
    /// Primal step length scaling.
    pub σp0 : F,
    /// Dual step length scaling.
    pub σd0 : F,
    /// Generic parameters
    pub insertion : FBGenericConfig<F>,
}

#[replace_float_literals(F::cast_from(literal))]
impl<F : Float> Default for ForwardPDPSConfig<F> {
    fn default() -> Self {
        let τ0 = 0.99;
        ForwardPDPSConfig {
            τ0,
            σd0 : 0.1,
            σp0 : 0.99,
            insertion : Default::default()
        }
    }
}

type MeasureZ<F, Z, const N : usize> = Pair<RNDM<F, N>, Z>;

/// Iteratively solve the pointsource localisation with an additional variable
/// using primal-dual proximal splitting with a forward step.
#[replace_float_literals(F::cast_from(literal))]
pub fn pointsource_forward_pdps_pair<
    F, I, A, S, Reg, P, Z, R, Y, /*KOpM, */ KOpZ, H, const N : usize
>(
    opA : &A,
    b : &A::Observable,
    reg : Reg,
    prox_penalty : &P,
    config : &ForwardPDPSConfig<F>,
    iterator : I,
    mut plotter : SeqPlotter<F, N>,
    //opKμ : KOpM,
    opKz : &KOpZ,
    fnR : &R,
    fnH : &H,
    mut z : Z,
    mut y : Y,
) -> MeasureZ<F, Z, N>
where
    F : Float + ToNalgebraRealField,
    I : AlgIteratorFactory<IterInfo<F, N>>,
    A : ForwardModel<
            MeasureZ<F, Z, N>,
            F,
            PairNorm<Radon, L2, L2>,
            PreadjointCodomain = Pair<S, Z>,
        >
        + AdjointProductPairBoundedBy<MeasureZ<F, Z, N>, P, IdOp<Z>, FloatType=F>,
    S: DifferentiableRealMapping<F, N>,
    for<'b> &'b A::Observable : std::ops::Neg<Output=A::Observable> + Instance<A::Observable>,
    PlotLookup : Plotting<N>,
    RNDM<F, N> : SpikeMerging<F>,
    Reg : RegTerm<F, N>,
    P : ProxPenalty<F, S, Reg, N>,
    KOpZ : BoundedLinear<Z, L2, L2, F, Codomain=Y>
        + GEMV<F, Z>
        + Adjointable<Z, Y, AdjointCodomain = Z>,
    for<'b> KOpZ::Adjoint<'b> : GEMV<F, Y>,
    Y : AXPY<F> + Euclidean<F, Output=Y> + Clone + ClosedAdd,
    for<'b> &'b Y : Instance<Y>,
    Z : AXPY<F, Owned=Z> + Euclidean<F, Output=Z> + Clone + Norm<F, L2>,
    for<'b> &'b Z : Instance<Z>,
    R : Prox<Z, Codomain=F>,
    H : Conjugable<Y, F, Codomain=F>,
    for<'b> H::Conjugate<'b> : Prox<Y>,
{

    // Check parameters
    assert!(config.τ0 > 0.0 &&
            config.τ0 < 1.0 &&
            config.σp0 > 0.0 &&
            config.σp0 < 1.0 &&
            config.σd0 > 0.0 &&
            config.σp0 * config.σd0 <= 1.0,
            "Invalid step length parameters");

    // Initialise iterates
    let mut μ = DiscreteMeasure::new();
    let mut residual = calculate_residual(Pair(&μ, &z), opA, b);

    // Set up parameters
    let bigM = 0.0; //opKμ.adjoint_product_bound(prox_penalty).unwrap().sqrt();
    let nKz = opKz.opnorm_bound(L2, L2);
    let opIdZ = IdOp::new();
    let (l, l_z) = opA.adjoint_product_pair_bound(prox_penalty, &opIdZ).unwrap();
    // We need to satisfy
    //
    //     τσ_dM(1-σ_p L_z)/(1 - τ L) + [σ_p L_z + σ_pσ_d‖K_z‖^2] < 1
    //                                  ^^^^^^^^^^^^^^^^^^^^^^^^^
    // with 1 > σ_p L_z and 1 > τ L.
    //
    // To do so, we first solve σ_p and σ_d from standard PDPS step length condition
    // ^^^^^ < 1. then we solve τ from  the rest.
    let σ_d = config.σd0 / nKz;
    let σ_p = config.σp0 / (l_z + config.σd0 * nKz);
    // Observe that = 1 - ^^^^^^^^^^^^^^^^^^^^^ = 1 - σ_{p,0}
    // We get the condition τσ_d M (1-σ_p L_z) < (1-σ_{p,0})*(1-τ L)
    // ⟺ τ [ σ_d M (1-σ_p L_z) + (1-σ_{p,0}) L ] < (1-σ_{p,0})
    let φ = 1.0 - config.σp0;
    let a = 1.0 - σ_p * l_z;
    let τ = config.τ0 * φ / ( σ_d * bigM * a + φ * l );
    // Acceleration is not currently supported
    // let γ = dataterm.factor_of_strong_convexity();
    let ω = 1.0;

    // We multiply tolerance by τ for FB since our subproblems depending on tolerances are scaled
    // by τ compared to the conditional gradient approach.
    let tolerance = config.insertion.tolerance * τ * reg.tolerance_scaling();
    let mut ε = tolerance.initial();

    let starH = fnH.conjugate();

    // Statistics
    let full_stats = |residual : &A::Observable, μ : &RNDM<F, N>, z : &Z, ε, stats| IterInfo {
        value : residual.norm2_squared_div2() + fnR.apply(z)
                + reg.apply(μ) + fnH.apply(/* opKμ.apply(μ) + */ opKz.apply(z)),
        n_spikes : μ.len(),
        ε,
        // postprocessing: config.insertion.postprocessing.then(|| μ.clone()),
        .. stats
    };
    let mut stats = IterInfo::new();

    // Run the algorithm
    for state in iterator.iter_init(|| full_stats(&residual, &μ, &z, ε, stats.clone())) {
        // Calculate initial transport
        let Pair(mut τv, τz) = opA.preadjoint().apply(residual * τ);
        let z_base = z.clone();
        let μ_base = μ.clone();

        // Construct μ^{k+1} by solving finite-dimensional subproblems and insert new spikes.
        let (maybe_d, _within_tolerances) = prox_penalty.insert_and_reweigh(
            &mut μ, &mut τv, &μ_base, None,
            τ, ε, &config.insertion,
            &reg, &state, &mut stats,
        );

        // // Merge spikes.
        // // This expects the prune below to prune γ.
        // // TODO: This may not work correctly in all cases.
        // let ins = &config.insertion;
        // if ins.merge_now(&state) {
        //     if let SpikeMergingMethod::None = ins.merging {
        //     } else {
        //         stats.merged += μ.merge_spikes(ins.merging, |μ_candidate| {
        //             let ν = μ_candidate.sub_matching(&γ1)-&μ_base_minus_γ0;
        //             let mut d = &τv̆ + op𝒟.preapply(ν);
        //             reg.verify_merge_candidate(&mut d, μ_candidate, τ, ε, ins)
        //         });
        //     }
        // }

        // Prune spikes with zero weight.
        stats.pruned += prune_with_stats(&mut μ);

        // Do z variable primal update
        z.axpy(-σ_p/τ, τz, 1.0); // TODO: simplify nasty factors
        opKz.adjoint().gemv(&mut z, -σ_p, &y, 1.0);
        z = fnR.prox(σ_p, z);
        // Do dual update
        // opKμ.gemv(&mut y, σ_d*(1.0 + ω), &μ, 1.0);    // y = y + σ_d K[(1+ω)(μ,z)^{k+1}]
        opKz.gemv(&mut y, σ_d*(1.0 + ω), &z, 1.0);
        // opKμ.gemv(&mut y, -σ_d*ω, μ_base, 1.0);// y = y + σ_d K[(1+ω)(μ,z)^{k+1} - ω (μ,z)^k]-b
        opKz.gemv(&mut y, -σ_d*ω, z_base, 1.0);// y = y + σ_d K[(1+ω)(μ,z)^{k+1} - ω (μ,z)^k]-b
        y = starH.prox(σ_d, y);

        // Update residual
        residual = calculate_residual(Pair(&μ, &z), opA, b);

        // Update step length parameters
        // let ω = pdpsconfig.acceleration.accelerate(&mut τ, &mut σ, γ);

        // Give statistics if requested
        let iter = state.iteration();
        stats.this_iters += 1;

        state.if_verbose(|| {
            plotter.plot_spikes(iter, maybe_d.as_ref(), Some(&τv), &μ);
            full_stats(&residual, &μ, &z, ε, std::mem::replace(&mut stats, IterInfo::new()))
        });

        // Update main tolerance for next iteration
        ε = tolerance.update(ε, iter);
    }

    let fit = |μ̃ : &RNDM<F, N>| {
        (opA.apply(Pair(μ̃, &z))-b).norm2_squared_div2()
        //+ fnR.apply(z) + reg.apply(μ)
        + fnH.apply(/* opKμ.apply(&μ̃) + */ opKz.apply(&z))
    };

    μ.merge_spikes_fitness(config.insertion.merging, fit, |&v| v);
    μ.prune();
    Pair(μ, z)
}