|
1 /*! |
|
2 Solver for the point source localisation problem using a |
|
3 primal-dual proximal splitting with a forward step. |
|
4 */ |
|
5 |
|
6 use numeric_literals::replace_float_literals; |
|
7 use serde::{Serialize, Deserialize}; |
|
8 |
|
9 use alg_tools::iterate::AlgIteratorFactory; |
|
10 use alg_tools::euclidean::Euclidean; |
|
11 use alg_tools::mapping::{Mapping, DifferentiableRealMapping, Instance}; |
|
12 use alg_tools::norms::Norm; |
|
13 use alg_tools::direct_product::Pair; |
|
14 use alg_tools::nalgebra_support::ToNalgebraRealField; |
|
15 use alg_tools::linops::{ |
|
16 BoundedLinear, AXPY, GEMV, Adjointable, IdOp, |
|
17 }; |
|
18 use alg_tools::convex::{Conjugable, Prox}; |
|
19 use alg_tools::norms::{L2, PairNorm}; |
|
20 |
|
21 use crate::types::*; |
|
22 use crate::measures::{DiscreteMeasure, Radon, RNDM}; |
|
23 use crate::measures::merging::SpikeMerging; |
|
24 use crate::forward_model::{ |
|
25 ForwardModel, |
|
26 AdjointProductPairBoundedBy, |
|
27 }; |
|
28 use crate::plot::{ |
|
29 SeqPlotter, |
|
30 Plotting, |
|
31 PlotLookup |
|
32 }; |
|
33 use crate::fb::*; |
|
34 use crate::regularisation::RegTerm; |
|
35 use crate::dataterm::calculate_residual; |
|
36 |
|
37 /// Settings for [`pointsource_forward_pdps_pair`]. |
|
38 #[derive(Clone, Copy, Eq, PartialEq, Serialize, Deserialize, Debug)] |
|
39 #[serde(default)] |
|
40 pub struct ForwardPDPSConfig<F : Float> { |
|
41 /// Primal step length scaling. |
|
42 pub τ0 : F, |
|
43 /// Primal step length scaling. |
|
44 pub σp0 : F, |
|
45 /// Dual step length scaling. |
|
46 pub σd0 : F, |
|
47 /// Generic parameters |
|
48 pub insertion : FBGenericConfig<F>, |
|
49 } |
|
50 |
|
51 #[replace_float_literals(F::cast_from(literal))] |
|
52 impl<F : Float> Default for ForwardPDPSConfig<F> { |
|
53 fn default() -> Self { |
|
54 ForwardPDPSConfig { |
|
55 τ0 : 0.99, |
|
56 σd0 : 0.05, |
|
57 σp0 : 0.99, |
|
58 insertion : Default::default() |
|
59 } |
|
60 } |
|
61 } |
|
62 |
|
63 type MeasureZ<F, Z, const N : usize> = Pair<RNDM<F, N>, Z>; |
|
64 |
|
65 /// Iteratively solve the pointsource localisation with an additional variable |
|
66 /// using primal-dual proximal splitting with a forward step. |
|
67 #[replace_float_literals(F::cast_from(literal))] |
|
68 pub fn pointsource_forward_pdps_pair< |
|
69 F, I, A, S, Reg, P, Z, R, Y, /*KOpM, */ KOpZ, H, const N : usize |
|
70 >( |
|
71 opA : &A, |
|
72 b : &A::Observable, |
|
73 reg : Reg, |
|
74 prox_penalty : &P, |
|
75 config : &ForwardPDPSConfig<F>, |
|
76 iterator : I, |
|
77 mut plotter : SeqPlotter<F, N>, |
|
78 //opKμ : KOpM, |
|
79 opKz : &KOpZ, |
|
80 fnR : &R, |
|
81 fnH : &H, |
|
82 mut z : Z, |
|
83 mut y : Y, |
|
84 ) -> MeasureZ<F, Z, N> |
|
85 where |
|
86 F : Float + ToNalgebraRealField, |
|
87 I : AlgIteratorFactory<IterInfo<F, N>>, |
|
88 A : ForwardModel< |
|
89 MeasureZ<F, Z, N>, |
|
90 F, |
|
91 PairNorm<Radon, L2, L2>, |
|
92 PreadjointCodomain = Pair<S, Z>, |
|
93 > |
|
94 + AdjointProductPairBoundedBy<MeasureZ<F, Z, N>, P, IdOp<Z>, FloatType=F>, |
|
95 S: DifferentiableRealMapping<F, N>, |
|
96 for<'b> &'b A::Observable : std::ops::Neg<Output=A::Observable> + Instance<A::Observable>, |
|
97 PlotLookup : Plotting<N>, |
|
98 RNDM<F, N> : SpikeMerging<F>, |
|
99 Reg : RegTerm<F, N>, |
|
100 P : ProxPenalty<F, S, Reg, N>, |
|
101 KOpZ : BoundedLinear<Z, L2, L2, F, Codomain=Y> |
|
102 + GEMV<F, Z> |
|
103 + Adjointable<Z, Y, AdjointCodomain = Z>, |
|
104 for<'b> KOpZ::Adjoint<'b> : GEMV<F, Y>, |
|
105 Y : AXPY<F> + Euclidean<F, Output=Y> + Clone + ClosedAdd, |
|
106 for<'b> &'b Y : Instance<Y>, |
|
107 Z : AXPY<F, Owned=Z> + Euclidean<F, Output=Z> + Clone + Norm<F, L2>, |
|
108 for<'b> &'b Z : Instance<Z>, |
|
109 R : Prox<Z, Codomain=F>, |
|
110 H : Conjugable<Y, F, Codomain=F>, |
|
111 for<'b> H::Conjugate<'b> : Prox<Y>, |
|
112 { |
|
113 |
|
114 // Check parameters |
|
115 assert!(config.τ0 > 0.0 && |
|
116 config.τ0 < 1.0 && |
|
117 config.σp0 > 0.0 && |
|
118 config.σp0 < 1.0 && |
|
119 config.σd0 > 0.0 && |
|
120 config.σp0 * config.σd0 <= 1.0, |
|
121 "Invalid step length parameters"); |
|
122 |
|
123 // Initialise iterates |
|
124 let mut μ = DiscreteMeasure::new(); |
|
125 let mut residual = calculate_residual(Pair(&μ, &z), opA, b); |
|
126 |
|
127 // Set up parameters |
|
128 let bigM = 0.0; //opKμ.adjoint_product_bound(prox_penalty).unwrap().sqrt(); |
|
129 let nKz = opKz.opnorm_bound(L2, L2); |
|
130 let opIdZ = IdOp::new(); |
|
131 let (l, l_z) = opA.adjoint_product_pair_bound(prox_penalty, &opIdZ).unwrap(); |
|
132 // We need to satisfy |
|
133 // |
|
134 // τσ_dM(1-σ_p L_z)/(1 - τ L) + [σ_p L_z + σ_pσ_d‖K_z‖^2] < 1 |
|
135 // ^^^^^^^^^^^^^^^^^^^^^^^^^ |
|
136 // with 1 > σ_p L_z and 1 > τ L. |
|
137 // |
|
138 // To do so, we first solve σ_p and σ_d from standard PDPS step length condition |
|
139 // ^^^^^ < 1. then we solve τ from the rest. |
|
140 let σ_d = config.σd0 / nKz; |
|
141 let σ_p = config.σp0 / (l_z + config.σd0 * nKz); |
|
142 // Observe that = 1 - ^^^^^^^^^^^^^^^^^^^^^ = 1 - σ_{p,0} |
|
143 // We get the condition τσ_d M (1-σ_p L_z) < (1-σ_{p,0})*(1-τ L) |
|
144 // ⟺ τ [ σ_d M (1-σ_p L_z) + (1-σ_{p,0}) L ] < (1-σ_{p,0}) |
|
145 let φ = 1.0 - config.σp0; |
|
146 let a = 1.0 - σ_p * l_z; |
|
147 let τ = config.τ0 * φ / ( σ_d * bigM * a + φ * l ); |
|
148 // Acceleration is not currently supported |
|
149 // let γ = dataterm.factor_of_strong_convexity(); |
|
150 let ω = 1.0; |
|
151 |
|
152 // We multiply tolerance by τ for FB since our subproblems depending on tolerances are scaled |
|
153 // by τ compared to the conditional gradient approach. |
|
154 let tolerance = config.insertion.tolerance * τ * reg.tolerance_scaling(); |
|
155 let mut ε = tolerance.initial(); |
|
156 |
|
157 let starH = fnH.conjugate(); |
|
158 |
|
159 // Statistics |
|
160 let full_stats = |residual : &A::Observable, μ : &RNDM<F, N>, z : &Z, ε, stats| IterInfo { |
|
161 value : residual.norm2_squared_div2() + fnR.apply(z) |
|
162 + reg.apply(μ) + fnH.apply(/* opKμ.apply(μ) + */ opKz.apply(z)), |
|
163 n_spikes : μ.len(), |
|
164 ε, |
|
165 // postprocessing: config.insertion.postprocessing.then(|| μ.clone()), |
|
166 .. stats |
|
167 }; |
|
168 let mut stats = IterInfo::new(); |
|
169 |
|
170 // Run the algorithm |
|
171 for state in iterator.iter_init(|| full_stats(&residual, &μ, &z, ε, stats.clone())) { |
|
172 // Calculate initial transport |
|
173 let Pair(mut τv, τz) = opA.preadjoint().apply(residual * τ); |
|
174 let μ_base = μ.clone(); |
|
175 |
|
176 // Construct μ^{k+1} by solving finite-dimensional subproblems and insert new spikes. |
|
177 let (maybe_d, _within_tolerances) = prox_penalty.insert_and_reweigh( |
|
178 &mut μ, &mut τv, &μ_base, None, |
|
179 τ, ε, &config.insertion, |
|
180 ®, &state, &mut stats, |
|
181 ); |
|
182 |
|
183 // Merge spikes. |
|
184 // This crucially expects the merge routine to be stable with respect to spike locations, |
|
185 // and not to performing any pruning. That is be to done below simultaneously for γ. |
|
186 // Merge spikes. |
|
187 // This crucially expects the merge routine to be stable with respect to spike locations, |
|
188 // and not to performing any pruning. That is be to done below simultaneously for γ. |
|
189 let ins = &config.insertion; |
|
190 if ins.merge_now(&state) { |
|
191 stats.merged += prox_penalty.merge_spikes_no_fitness( |
|
192 &mut μ, &mut τv, &μ_base, None, τ, ε, ins, ®, |
|
193 //Some(|μ̃ : &RNDM<F, N>| calculate_residual(Pair(μ̃, &z), opA, b).norm2_squared_div2()), |
|
194 ); |
|
195 } |
|
196 |
|
197 // Prune spikes with zero weight. |
|
198 stats.pruned += prune_with_stats(&mut μ); |
|
199 |
|
200 // Do z variable primal update |
|
201 let mut z_new = τz; |
|
202 opKz.adjoint().gemv(&mut z_new, -σ_p, &y, -σ_p/τ); |
|
203 z_new = fnR.prox(σ_p, z_new + &z); |
|
204 // Do dual update |
|
205 // opKμ.gemv(&mut y, σ_d*(1.0 + ω), &μ, 1.0); // y = y + σ_d K[(1+ω)(μ,z)^{k+1}] |
|
206 opKz.gemv(&mut y, σ_d*(1.0 + ω), &z_new, 1.0); |
|
207 // opKμ.gemv(&mut y, -σ_d*ω, μ_base, 1.0);// y = y + σ_d K[(1+ω)(μ,z)^{k+1} - ω (μ,z)^k]-b |
|
208 opKz.gemv(&mut y, -σ_d*ω, z, 1.0);// y = y + σ_d K[(1+ω)(μ,z)^{k+1} - ω (μ,z)^k]-b |
|
209 y = starH.prox(σ_d, y); |
|
210 z = z_new; |
|
211 |
|
212 // Update residual |
|
213 residual = calculate_residual(Pair(&μ, &z), opA, b); |
|
214 |
|
215 // Update step length parameters |
|
216 // let ω = pdpsconfig.acceleration.accelerate(&mut τ, &mut σ, γ); |
|
217 |
|
218 // Give statistics if requested |
|
219 let iter = state.iteration(); |
|
220 stats.this_iters += 1; |
|
221 |
|
222 state.if_verbose(|| { |
|
223 plotter.plot_spikes(iter, maybe_d.as_ref(), Some(&τv), &μ); |
|
224 full_stats(&residual, &μ, &z, ε, std::mem::replace(&mut stats, IterInfo::new())) |
|
225 }); |
|
226 |
|
227 // Update main tolerance for next iteration |
|
228 ε = tolerance.update(ε, iter); |
|
229 } |
|
230 |
|
231 let fit = |μ̃ : &RNDM<F, N>| { |
|
232 (opA.apply(Pair(μ̃, &z))-b).norm2_squared_div2() |
|
233 //+ fnR.apply(z) + reg.apply(μ) |
|
234 + fnH.apply(/* opKμ.apply(&μ̃) + */ opKz.apply(&z)) |
|
235 }; |
|
236 |
|
237 μ.merge_spikes_fitness(config.insertion.final_merging_method(), fit, |&v| v); |
|
238 μ.prune(); |
|
239 Pair(μ, z) |
|
240 } |