--- /dev/null Thu Jan 01 00:00:00 1970 +0000
+++ b/src/forward_pdps.rs Tue Dec 31 09:25:45 2024 -0500
@@ -0,0 +1,268 @@
+/*!
+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::sets::Cube;
+use alg_tools::loc::Loc;
+use alg_tools::mapping::{Mapping, Instance};
+use alg_tools::norms::Norm;
+use alg_tools::direct_product::Pair;
+use alg_tools::bisection_tree::{
+ BTFN,
+ PreBTFN,
+ Bounds,
+ BTNodeLookup,
+ BTNode,
+ BTSearch,
+ P2Minimise,
+ SupportGenerator,
+ LocalAnalysis,
+ //Bounded,
+};
+use alg_tools::mapping::RealMapping;
+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, Linfinity, PairNorm};
+
+use crate::types::*;
+use crate::measures::{DiscreteMeasure, Radon, RNDM};
+use crate::measures::merging::SpikeMerging;
+use crate::forward_model::{
+ ForwardModel,
+ AdjointProductPairBoundedBy,
+};
+use crate::seminorms::DiscreteMeasureOp;
+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<
+ 'a, F, I, A, GA, 𝒟, BTA, BT𝒟, G𝒟, S, K, Reg, Z, R, Y, /*KOpM, */ KOpZ, H, const N : usize
+>(
+ opA : &'a A,
+ b : &A::Observable,
+ reg : Reg,
+ op𝒟 : &'a 𝒟,
+ 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>>,
+ for<'b> &'b A::Observable : std::ops::Neg<Output=A::Observable> + Instance<A::Observable>,
+ GA : SupportGenerator<F, N, SupportType = S, Id = usize> + Clone,
+ A : ForwardModel<
+ MeasureZ<F, Z, N>,
+ F,
+ PairNorm<Radon, L2, L2>,
+ PreadjointCodomain = Pair<BTFN<F, GA, BTA, N>, Z>,
+ >
+ + AdjointProductPairBoundedBy<MeasureZ<F, Z, N>, 𝒟, IdOp<Z>, FloatType=F>,
+ BTA : BTSearch<F, N, Data=usize, Agg=Bounds<F>>,
+ G𝒟 : SupportGenerator<F, N, SupportType = K, Id = usize> + Clone,
+ 𝒟 : DiscreteMeasureOp<Loc<F, N>, F, PreCodomain = PreBTFN<F, G𝒟, N>,
+ Codomain = BTFN<F, G𝒟, BT𝒟, N>>,
+ BT𝒟 : BTSearch<F, N, Data=usize, Agg=Bounds<F>>,
+ S: RealMapping<F, N> + LocalAnalysis<F, Bounds<F>, N>,
+ K: RealMapping<F, N> + LocalAnalysis<F, Bounds<F>, N>,
+ BTNodeLookup: BTNode<F, usize, Bounds<F>, N>,
+ Cube<F, N>: P2Minimise<Loc<F, N>, F>,
+ PlotLookup : Plotting<N>,
+ RNDM<F, N> : SpikeMerging<F>,
+ Reg : RegTerm<F, 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 op𝒟norm = op𝒟.opnorm_bound(Radon, Linfinity);
+ let bigM = 0.0; //opKμ.adjoint_product_bound(&op𝒟).unwrap().sqrt();
+ let nKz = opKz.opnorm_bound(L2, L2);
+ let opIdZ = IdOp::new();
+ let (l, l_z) = opA.adjoint_product_pair_bound(&op𝒟, &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(τ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 (d, _within_tolerances) = insert_and_reweigh(
+ &mut μ, &τv, &μ_base, None,
+ op𝒟, op𝒟norm,
+ τ, ε, &config.insertion,
+ ®, &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, Some(&d), 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)
+}