src/forward_pdps.rs

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

use crate::domain::{AdmissibleDomain, FullSpace};
use crate::fb::*;
use crate::measures::merging::SpikeMerging;
use crate::measures::DiscreteMeasure;
use crate::plot::Plotter;
use crate::prox_penalty::{ProxPenalty, StepLengthBoundPair};
use crate::regularisation::RegTerm;
use crate::types::*;
use alg_tools::bounds::MinMaxMapping;
use alg_tools::convex::{Conjugable, Prox, Zero};
use alg_tools::direct_product::Pair;
use alg_tools::error::DynResult;
use alg_tools::euclidean::ClosedEuclidean;
use alg_tools::iterate::AlgIteratorFactory;
use alg_tools::linops::{
    BoundedLinear, IdOp, SimplyAdjointable, Space, StaticEuclideanOriginGenerator, ZeroOp, AXPY,
    GEMV,
};
use alg_tools::mapping::{DifferentiableMapping, Instance};
use alg_tools::nalgebra_support::ToNalgebraRealField;
use alg_tools::norms::{Norm, L2};
use anyhow::ensure;
use measures::{DeltaMeasure, Radon};
use numeric_literals::replace_float_literals;
use serde::{Deserialize, Serialize};
//use std::ops::Mul;

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

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

type MeasureZ<F, Z, U> = Pair<DiscreteMeasure<U, F>, Z>;

/// Iteratively solve the pointsource localisation with an additional variable
/// using primal-dual proximal splitting with a forward step.
///
/// The problem is
/// $$
///    \min_{μ, z}~ F(μ, z) + R(z) + H(K_z z) + Q(μ),
/// $$
/// where
///   * The data term $F$ is given in `f`,
///   * the measure (Radon or positivity-constrained Radon) regulariser in $Q$ is given in `reg`,
///   * the functions $R$ and $H$ are given in `fnR` and `fnH`, and
///   * the operator $K_z$ in `opKz`.
///
/// This is dualised to
/// $$
///    \min_{μ, z}\max_y~ F(μ, z) + R(z) + ⟨K_z z, y⟩ + Q(μ) - H^*(y).
/// $$
///
/// The algorithm is controlled by:
///   * the proximal penalty in `prox_penalty`.
///   * the initial iterates in `z`, `y`
///   * The configuration in `config`.
///   * The `iterator` that controls stopping and reporting.
/// Moreover, plotting is performed by `plotter`.
///
/// The step lengths 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$.
/// Since we are given “scalings” $τ_0$, $σ_{p,0}$, and $σ_{d,0}$ in `config`, we take
/// $σ_d=σ_{d,0}/‖K_z‖$, and $σ_p = σ_{p,0} / (L_z σ_d‖K_z‖)$. This satisfies the
/// part $[σ_p L_z + σ_pσ_d‖K_z‖^2] < 1$. Then with these cohices, we solve
/// $$
///     τ = τ_0 \frac{1 - σ_{p,0}}{(σ_d M (1-σ_p L_z) + (1 - σ_{p,0} L)}.
/// $$
#[replace_float_literals(F::cast_from(literal))]
pub fn pointsource_forward_pdps_pair_dom<
    F,
    U,
    I,
    S,
    Dat,
    Reg,
    P,
    Z,
    R,
    Y,
    /*KOpM, */ KOpZ,
    H,
    Plot,
>(
    f: &Dat,
    reg: &Reg,
    prox_penalty: &P,
    config: &ForwardPDPSConfig<F>,
    iterator: I,
    mut plotter: Plot,
    (μ0, mut z, mut y): (Option<DiscreteMeasure<U, F>>, Z, Y),
    //opKμ : KOpM,
    opKz: &KOpZ,
    fnR: &R,
    fnH: &H,
    domain: &mut impl AdmissibleDomain<U>,
) -> DynResult<MeasureZ<F, Z, U>>
where
    F: Float + ToNalgebraRealField,
    U: Space + Clone,
    I: AlgIteratorFactory<IterInfo<F>>,
    Dat: DifferentiableMapping<MeasureZ<F, Z, U>, Codomain = F, DerivativeDomain = Pair<S, Z>>,
    S: ClosedMul<F> + MinMaxMapping<U, F>,
    for<'a> Pair<&'a P, &'a IdOp<Z>>: StepLengthBoundPair<F, Dat>,
    DiscreteMeasure<U, F>: SpikeMerging<F> + Norm<Radon, F>,
    Reg: RegTerm<U, F>,
    P: ProxPenalty<U, S, Reg, F>,
    KOpZ: BoundedLinear<Z, L2, L2, F, Codomain = Y>
        + GEMV<F, Z>
        + SimplyAdjointable<Z, Y, AdjointCodomain = Z>,
    KOpZ::SimpleAdjoint: GEMV<F, Y>,
    Y: ClosedEuclidean<F>,
    for<'b> &'b Y: Instance<Y>,
    Z: ClosedEuclidean<F> + AXPY + Clone,
    for<'b> &'b Z: Instance<Z>, // compiler overflow  + Mul<F, Output = Z>,
    R: Prox<Z, Codomain = F>,
    H: Conjugable<Y, F, Codomain = F>,
    for<'b> H::Conjugate<'b>: Prox<Y>,
    Plot: Plotter<P::ReturnMapping, S, DiscreteMeasure<U, F>>,
{
    // Check parameters
    ensure!(
        config.τ0 > 0.0 && config.σp0 > 0.0,
        //&& config.τ0 < 1.0 && config.σp0 < 1.0
        "Invalid step length parameters"
    );

    // Initialise iterates
    let mut μ = μ0.unwrap_or_else(|| DiscreteMeasure::new());

    // Set up parameters
    let bigM = 0.0; //opKμ.adjoint_product_bound(prox_penalty).unwrap().sqrt();
    let nKz = opKz.opnorm_bound(L2, L2)?;
    let is_fb = nKz == 0.0;
    let idOpZ = IdOp::new();
    let opKz_adj = opKz.adjoint();
    let (l, l_z) = Pair(prox_penalty, &idOpZ).step_length_bound_pair(&f)?;
    let l_z_prime = if l_z == 0.0 {
        1.0
    } else {
        ensure!(config.σp0 < 1.0);
        l_z
    };
    let l_prime = if l == 0.0 {
        1.0
    } else {
        ensure!(config.τ0 < 1.0);
        l
    };

    let (τ, σ_p, σ_d) = if is_fb {
        (config.τ0 / l_prime, config.σp0 / l_z_prime, 0.0)
    } else {
        ensure!(
            //config.τ0 < 1.0 &&
            config.σp0 < 1.0 && config.σd0 > 0.0 && config.σp0 * config.σd0 <= 1.0,
            "Invalid step length parameters"
        );
        let σ_d = config.σd0 / nKz;
        let σ_p = config.σp0 / (l_z_prime + config.σd0 * nKz);
        let φ = 1.0 - config.σp0;
        let a = 1.0 - σ_p * l_z_prime;
        let τ = config.τ0 * φ / (σ_d * bigM * a + φ * l_prime);
        (τ, σ_p, σ_d)
    };

    // 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 = |μ: &DiscreteMeasure<U, F>, v: &mut S, z: &Z, dz: &Z, ε, stats| {
        let value = f.apply(Pair(μ, z))
            + fnR.apply(z)
            + reg.apply(μ)
            + fnH.apply(/* opKμ.apply(μ) + */ opKz.apply(z));
        let discrepancy = reg.discrepancy(v, τ, ε, μ, &config.insertion);
        let aux_discrepancy = is_fb.then(|| z.dist2(fnR.prox(σ_p, Z::clone(dz) * (-σ_p) + z)));
        let joint_discrepancy = aux_discrepancy
            .map(|d| (discrepancy.powi(2) + τ * (2.0 / σ_p - l_z) * d.powi(2)).sqrt());
        IterInfo {
            value,
            n_spikes: μ.len(),
            discrepancy: Some(discrepancy),
            μ_norm: μ.norm(Radon),
            // compiler overflow with the necesssary trait bounds:
            //aux_discrepancy: is_fb.then(|| z.dist2(fnR.prox(σ_p, dz * (-σ_p) + z))),
            aux_discrepancy,
            joint_discrepancy,
            ε,
            // postprocessing: config.insertion.postprocessing.then(|| μ.clone()),
            ..stats
        }
    };
    let mut stats = IterInfo::new();

    let zy_step = |mut z, mut y, mut dz| {
        // Do z variable primal update
        opKz_adj.apply_add(&mut dz, &y);
        dz.axpy(1.0, &z, -σ_p);
        let z_new = fnR.prox(σ_p, dz);
        //let z_new = fnR.prox(σ_p, dz * (-σ_p) + &z);
        // Do dual update
        z.axpy(1.0 + ω, &z_new, -ω);
        //z = (z - &z_new) * (-ω) + &z_new;
        opKz.gemv(&mut y, σ_d, z, 1.0);
        let y = starH.prox(σ_d, y);
        (z_new, y)
    };

    // Calculate differential of smooth part
    let Pair(mut v, mut dz) = f.differential(Pair(&μ, &z));
    let mut τv = v * τ;

    // Run the algorithm
    for state in iterator.iter_init(|| full_stats(&μ, &mut τv, &z, &dz, ε, stats.clone())) {
        // Calculate initial transport
        let μ_base = μ.clone();
        let μ_base_len = μ_base.len();

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

        (z, y) = zy_step(z, y, dz);

        stats.inserted += μ.len() - μ_base.len();

        // Merge spikes.
        // This crucially expects the merge routine to be stable with respect to spike locations,
        // and not to performing any pruning. That is be to done below simultaneously for γ.
        let did_merge = config.insertion.if_merge_now(&state, |ins| {
            stats.merged += prox_penalty.merge_spikes(
                &mut μ,
                &mut τv,
                &μ_base,
                τ,
                ε,
                ins,
                &reg,
                Some(|μ̃: &DiscreteMeasure<U, F>| f.apply(Pair(μ̃, &z)) + reg.apply(μ̃)),
            );
            // Prune. Do not prune new insertions on first round.
            stats.pruned += prune_with_stats(&mut μ, μ_base_len);

            if config.insertion.random_insertions {
                if let Some(x) = domain.random_point() {
                    μ += DeltaMeasure { x, α: 0.0 };
                    stats.inserted += 1;
                }
            }

            for _ in 0..config.insertion.extra_weight_optimisation_steps {
                // Do extra weight optimisation step heuristic
                let Pair(v, mut _dz) = f.differential(Pair(&μ, &z));
                τv = v * τ;
                prox_penalty.reweigh(
                    &mut μ,
                    &mut τv,
                    τ,
                    ε,
                    &config.insertion,
                    &reg,
                    &state,
                    &mut stats,
                )?;

                // opKz_adj.apply_add(&mut dz, &y);
                // dz.axpy(1.0, &z, -σ_p);
                // let z_new = fnR.prox(σ_p, dz);
                // z.axpy(1.0 + ω, &z_new, -ω);
                // //z = (z - &z_new) * (-ω) + &z_new;
                // opKz.gemv(&mut y, σ_d, z, 1.0);
                // y = starH.prox(σ_d, y);
                // z = z_new;
            }
            Ok(())
        })?;

        if did_merge.is_none() {
            // Prune. Do not prune new insertions on first round.
            stats.pruned += prune_with_stats(&mut μ, μ_base_len);
        }

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

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

        // Update differential of smooth part
        Pair(v, dz) = f.differential(Pair(&μ, &z));
        τv = v * τ;

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

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

    let fit = |μ̃: &DiscreteMeasure<U, F>| {
        f.apply(Pair(μ̃, &z)) /*+ fnR.apply(z) + reg.apply(μ)*/
        + fnH.apply(/* opKμ.apply(&μ̃) + */ opKz.apply(&z))
        + reg.apply(μ̃)
    };

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

#[deprecated]
pub fn pointsource_forward_pdps_pair<
    F,
    U,
    I,
    S,
    Dat,
    Reg,
    P,
    Z,
    R,
    Y,
    /*KOpM, */ KOpZ,
    H,
    Plot,
>(
    f: &Dat,
    reg: &Reg,
    prox_penalty: &P,
    config: &ForwardPDPSConfig<F>,
    iterator: I,
    plotter: Plot,
    init: (Option<DiscreteMeasure<U, F>>, Z, Y),
    //opKμ : KOpM,
    opKz: &KOpZ,
    fnR: &R,
    fnH: &H,
) -> DynResult<MeasureZ<F, Z, U>>
where
    F: Float + ToNalgebraRealField,
    U: Space + Clone,
    //for<'a> &'a U: Instance<U>,
    I: AlgIteratorFactory<IterInfo<F>>,
    Dat: DifferentiableMapping<MeasureZ<F, Z, U>, Codomain = F, DerivativeDomain = Pair<S, Z>>,
    S: ClosedMul<F> + MinMaxMapping<U, F>,
    for<'a> Pair<&'a P, &'a IdOp<Z>>: StepLengthBoundPair<F, Dat>,
    //Pair<S, Z>: ClosedMul<F>,
    DiscreteMeasure<U, F>: SpikeMerging<F> + Norm<Radon, F>,
    Reg: RegTerm<U, F>,
    P: ProxPenalty<U, S, Reg, F>,
    KOpZ: BoundedLinear<Z, L2, L2, F, Codomain = Y>
        + GEMV<F, Z>
        + SimplyAdjointable<Z, Y, AdjointCodomain = Z>,
    KOpZ::SimpleAdjoint: GEMV<F, Y>,
    Y: ClosedEuclidean<F>,
    for<'b> &'b Y: Instance<Y>,
    Z: ClosedEuclidean<F> + AXPY + Clone, // clone not neede if Mul worked below
    for<'b> &'b Z: Instance<Z>,           // compiler overflow  + Mul<F, Output = Z>,
    R: Prox<Z, Codomain = F>,
    H: Conjugable<Y, F, Codomain = F>,
    for<'b> H::Conjugate<'b>: Prox<Y>,
    Plot: Plotter<P::ReturnMapping, S, DiscreteMeasure<U, F>>,
{
    pointsource_forward_pdps_pair_dom(
        f,
        reg,
        prox_penalty,
        config,
        iterator,
        plotter,
        init,
        opKz,
        fnR,
        fnH,
        &mut FullSpace,
    )
}

/// Iteratively solve the pointsource localisation with an additional variable
/// using forward-backward splitting.
///
/// The implementation uses [`pointsource_forward_pdps_pair`] with appropriate dummy
/// variables, operators, and functions.
#[replace_float_literals(F::cast_from(literal))]
pub fn pointsource_fb_pair_dom<F, U, I, S, Dat, Reg, P, Z, R, Plot>(
    f: &Dat,
    reg: &Reg,
    prox_penalty: &P,
    config: &FBConfig<F>,
    iterator: I,
    plotter: Plot,
    (μ0, z): (Option<DiscreteMeasure<U, F>>, Z),
    //opKμ : KOpM,
    fnR: &R,
    domain: &mut impl AdmissibleDomain<U>,
) -> DynResult<MeasureZ<F, Z, U>>
where
    F: Float + ToNalgebraRealField,
    U: Space + Clone,
    for<'a> &'a U: Instance<U>,
    I: AlgIteratorFactory<IterInfo<F>>,
    Dat: DifferentiableMapping<MeasureZ<F, Z, U>, Codomain = F, DerivativeDomain = Pair<S, Z>>,
    S: ClosedMul<F> + MinMaxMapping<U, F>,
    DiscreteMeasure<U, F>: SpikeMerging<F> + Norm<Radon, F>,
    Reg: RegTerm<U, F>,
    P: ProxPenalty<U, S, Reg, F>,
    for<'a> Pair<&'a P, &'a IdOp<Z>>: StepLengthBoundPair<F, Dat>,
    Z: ClosedEuclidean<F> + AXPY + Clone,
    for<'b> &'b Z: Instance<Z>, // compiler overflow + Mul<F, Output = Z>,
    R: Prox<Z, Codomain = F>,
    Plot: Plotter<P::ReturnMapping, S, DiscreteMeasure<U, F>>,
{
    let opKz: ZeroOp<Z, Loc<0, F>, _, _, F> =
        ZeroOp::new_dualisable(StaticEuclideanOriginGenerator, z.dual_origin());
    let fnH = Zero::new();
    // Convert config. We don't implement From (that could be done with the o2o crate), as σd0
    // needs to be chosen in a general case; for the problem of this fucntion, anything is valid.
    let &FBConfig { τ0, σp0, insertion } = config;
    let pdps_config = ForwardPDPSConfig { τ0, σp0, insertion, σd0: 0.0 };

    pointsource_forward_pdps_pair_dom(
        f,
        reg,
        prox_penalty,
        &pdps_config,
        iterator,
        plotter,
        (μ0, z, Loc([])),
        &opKz,
        fnR,
        &fnH,
        domain,
    )
}

#[deprecated]
#[replace_float_literals(F::cast_from(literal))]
pub fn pointsource_fb_pair<F, U, I, S, Dat, Reg, P, Z, R, Plot>(
    f: &Dat,
    reg: &Reg,
    prox_penalty: &P,
    config: &FBConfig<F>,
    iterator: I,
    plotter: Plot,
    init: (Option<DiscreteMeasure<U, F>>, Z),
    //opKμ : KOpM,
    fnR: &R,
) -> DynResult<MeasureZ<F, Z, U>>
where
    F: Float + ToNalgebraRealField,
    U: Space + Clone,
    for<'a> &'a U: Instance<U>,
    I: AlgIteratorFactory<IterInfo<F>>,
    Dat: DifferentiableMapping<MeasureZ<F, Z, U>, Codomain = F, DerivativeDomain = Pair<S, Z>>,
    S: ClosedMul<F> + MinMaxMapping<U, F>,
    DiscreteMeasure<U, F>: SpikeMerging<F> + Norm<Radon, F>,
    Reg: RegTerm<U, F>,
    P: ProxPenalty<U, S, Reg, F>,
    for<'a> Pair<&'a P, &'a IdOp<Z>>: StepLengthBoundPair<F, Dat>,
    Z: ClosedEuclidean<F> + AXPY + Clone,
    for<'b> &'b Z: Instance<Z>, // compiler overflow  + Mul<F, Output = Z>,
    R: Prox<Z, Codomain = F>,
    Plot: Plotter<P::ReturnMapping, S, DiscreteMeasure<U, F>>,
{
    pointsource_fb_pair_dom(
        f,
        reg,
        prox_penalty,
        config,
        iterator,
        plotter,
        init,
        fnR,
        &mut FullSpace,
    )
}