| |
1 /*! |
| |
2 Solver for the point source localisation problem using a sliding |
| |
3 primal-dual proximal splitting method. |
| |
4 */ |
| |
5 |
| |
6 use numeric_literals::replace_float_literals; |
| |
7 use serde::{Serialize, Deserialize}; |
| |
8 //use colored::Colorize; |
| |
9 //use nalgebra::{DVector, DMatrix}; |
| |
10 use std::iter::Iterator; |
| |
11 |
| |
12 use alg_tools::iterate::AlgIteratorFactory; |
| |
13 use alg_tools::euclidean::Euclidean; |
| |
14 use alg_tools::sets::Cube; |
| |
15 use alg_tools::loc::Loc; |
| |
16 use alg_tools::mapping::{Mapping, DifferentiableRealMapping, Instance}; |
| |
17 use alg_tools::norms::Norm; |
| |
18 use alg_tools::direct_product::Pair; |
| |
19 use alg_tools::bisection_tree::{ |
| |
20 BTFN, |
| |
21 PreBTFN, |
| |
22 Bounds, |
| |
23 BTNodeLookup, |
| |
24 BTNode, |
| |
25 BTSearch, |
| |
26 P2Minimise, |
| |
27 SupportGenerator, |
| |
28 LocalAnalysis, |
| |
29 //Bounded, |
| |
30 }; |
| |
31 use alg_tools::mapping::RealMapping; |
| |
32 use alg_tools::nalgebra_support::ToNalgebraRealField; |
| |
33 use alg_tools::linops::{ |
| |
34 BoundedLinear, AXPY, GEMV, Adjointable, IdOp, |
| |
35 }; |
| |
36 use alg_tools::convex::{Conjugable, Prox}; |
| |
37 use alg_tools::norms::{L2, Linfinity, PairNorm}; |
| |
38 |
| |
39 use crate::types::*; |
| |
40 use crate::measures::{DiscreteMeasure, Radon, RNDM}; |
| |
41 use crate::measures::merging::SpikeMerging; |
| |
42 use crate::forward_model::{ |
| |
43 ForwardModel, |
| |
44 AdjointProductPairBoundedBy, |
| |
45 LipschitzValues, |
| |
46 }; |
| |
47 // use crate::transport::TransportLipschitz; |
| |
48 use crate::seminorms::DiscreteMeasureOp; |
| |
49 //use crate::tolerance::Tolerance; |
| |
50 use crate::plot::{ |
| |
51 SeqPlotter, |
| |
52 Plotting, |
| |
53 PlotLookup |
| |
54 }; |
| |
55 use crate::fb::*; |
| |
56 use crate::regularisation::SlidingRegTerm; |
| |
57 // use crate::dataterm::L2Squared; |
| |
58 use crate::sliding_fb::{ |
| |
59 TransportConfig, |
| |
60 TransportStepLength, |
| |
61 initial_transport, |
| |
62 aposteriori_transport, |
| |
63 }; |
| |
64 use crate::dataterm::{calculate_residual, calculate_residual2}; |
| |
65 |
| |
66 /// Settings for [`pointsource_sliding_pdps_pair`]. |
| |
67 #[derive(Clone, Copy, Eq, PartialEq, Serialize, Deserialize, Debug)] |
| |
68 #[serde(default)] |
| |
69 pub struct SlidingPDPSConfig<F : Float> { |
| |
70 /// Primal step length scaling. |
| |
71 pub τ0 : F, |
| |
72 /// Primal step length scaling. |
| |
73 pub σp0 : F, |
| |
74 /// Dual step length scaling. |
| |
75 pub σd0 : F, |
| |
76 /// Transport parameters |
| |
77 pub transport : TransportConfig<F>, |
| |
78 /// Generic parameters |
| |
79 pub insertion : FBGenericConfig<F>, |
| |
80 } |
| |
81 |
| |
82 #[replace_float_literals(F::cast_from(literal))] |
| |
83 impl<F : Float> Default for SlidingPDPSConfig<F> { |
| |
84 fn default() -> Self { |
| |
85 let τ0 = 0.99; |
| |
86 SlidingPDPSConfig { |
| |
87 τ0, |
| |
88 σd0 : 0.1, |
| |
89 σp0 : 0.99, |
| |
90 transport : Default::default(), |
| |
91 insertion : Default::default() |
| |
92 } |
| |
93 } |
| |
94 } |
| |
95 |
| |
96 type MeasureZ<F, Z, const N : usize> = Pair<RNDM<F, N>, Z>; |
| |
97 |
| |
98 /// Iteratively solve the pointsource localisation with an additional variable |
| |
99 /// using sliding primal-dual proximal splitting |
| |
100 /// |
| |
101 /// The parametrisation is as for [`crate::forward_pdps::pointsource_forward_pdps_pair`]. |
| |
102 #[replace_float_literals(F::cast_from(literal))] |
| |
103 pub fn pointsource_sliding_pdps_pair< |
| |
104 'a, F, I, A, GA, 𝒟, BTA, BT𝒟, G𝒟, S, K, Reg, Z, R, Y, /*KOpM, */ KOpZ, H, const N : usize |
| |
105 >( |
| |
106 opA : &'a A, |
| |
107 b : &A::Observable, |
| |
108 reg : Reg, |
| |
109 op𝒟 : &'a 𝒟, |
| |
110 config : &SlidingPDPSConfig<F>, |
| |
111 iterator : I, |
| |
112 mut plotter : SeqPlotter<F, N>, |
| |
113 //opKμ : KOpM, |
| |
114 opKz : &KOpZ, |
| |
115 fnR : &R, |
| |
116 fnH : &H, |
| |
117 mut z : Z, |
| |
118 mut y : Y, |
| |
119 ) -> MeasureZ<F, Z, N> |
| |
120 where |
| |
121 F : Float + ToNalgebraRealField, |
| |
122 I : AlgIteratorFactory<IterInfo<F, N>>, |
| |
123 for<'b> &'b A::Observable : std::ops::Neg<Output=A::Observable> + Instance<A::Observable>, |
| |
124 for<'b> A::Preadjoint<'b> : LipschitzValues<FloatType=F>, |
| |
125 BTFN<F, GA, BTA, N> : DifferentiableRealMapping<F, N>, |
| |
126 GA : SupportGenerator<F, N, SupportType = S, Id = usize> + Clone, |
| |
127 A : ForwardModel< |
| |
128 MeasureZ<F, Z, N>, |
| |
129 F, |
| |
130 PairNorm<Radon, L2, L2>, |
| |
131 PreadjointCodomain = Pair<BTFN<F, GA, BTA, N>, Z>, |
| |
132 > |
| |
133 + AdjointProductPairBoundedBy<MeasureZ<F, Z, N>, 𝒟, IdOp<Z>, FloatType=F>, |
| |
134 BTA : BTSearch<F, N, Data=usize, Agg=Bounds<F>>, |
| |
135 G𝒟 : SupportGenerator<F, N, SupportType = K, Id = usize> + Clone, |
| |
136 𝒟 : DiscreteMeasureOp<Loc<F, N>, F, PreCodomain = PreBTFN<F, G𝒟, N>, |
| |
137 Codomain = BTFN<F, G𝒟, BT𝒟, N>>, |
| |
138 BT𝒟 : BTSearch<F, N, Data=usize, Agg=Bounds<F>>, |
| |
139 S: RealMapping<F, N> + LocalAnalysis<F, Bounds<F>, N> |
| |
140 + DifferentiableRealMapping<F, N>, |
| |
141 K: RealMapping<F, N> + LocalAnalysis<F, Bounds<F>, N>, |
| |
142 //+ Differentiable<Loc<F, N>, Derivative=Loc<F,N>>, |
| |
143 BTNodeLookup: BTNode<F, usize, Bounds<F>, N>, |
| |
144 Cube<F, N>: P2Minimise<Loc<F, N>, F>, |
| |
145 PlotLookup : Plotting<N>, |
| |
146 RNDM<F, N> : SpikeMerging<F>, |
| |
147 Reg : SlidingRegTerm<F, N>, |
| |
148 // KOpM : Linear<RNDM<F, N>, Codomain=Y> |
| |
149 // + GEMV<F, RNDM<F, N>> |
| |
150 // + Preadjointable< |
| |
151 // RNDM<F, N>, Y, |
| |
152 // PreadjointCodomain = BTFN<F, GA, BTA, N>, |
| |
153 // > |
| |
154 // + TransportLipschitz<L2Squared, FloatType=F> |
| |
155 // + AdjointProductBoundedBy<RNDM<F, N>, 𝒟, FloatType=F>, |
| |
156 // for<'b> KOpM::Preadjoint<'b> : GEMV<F, Y>, |
| |
157 // Since Z is Hilbert, we may just as well use adjoints for K_z. |
| |
158 KOpZ : BoundedLinear<Z, L2, L2, F, Codomain=Y> |
| |
159 + GEMV<F, Z> |
| |
160 + Adjointable<Z, Y, AdjointCodomain = Z>, |
| |
161 for<'b> KOpZ::Adjoint<'b> : GEMV<F, Y>, |
| |
162 Y : AXPY<F> + Euclidean<F, Output=Y> + Clone + ClosedAdd, |
| |
163 for<'b> &'b Y : Instance<Y>, |
| |
164 Z : AXPY<F, Owned=Z> + Euclidean<F, Output=Z> + Clone + Norm<F, L2>, |
| |
165 for<'b> &'b Z : Instance<Z>, |
| |
166 R : Prox<Z, Codomain=F>, |
| |
167 H : Conjugable<Y, F, Codomain=F>, |
| |
168 for<'b> H::Conjugate<'b> : Prox<Y>, |
| |
169 { |
| |
170 |
| |
171 // Check parameters |
| |
172 assert!(config.τ0 > 0.0 && |
| |
173 config.τ0 < 1.0 && |
| |
174 config.σp0 > 0.0 && |
| |
175 config.σp0 < 1.0 && |
| |
176 config.σd0 > 0.0 && |
| |
177 config.σp0 * config.σd0 <= 1.0, |
| |
178 "Invalid step length parameters"); |
| |
179 config.transport.check(); |
| |
180 |
| |
181 // Initialise iterates |
| |
182 let mut μ = DiscreteMeasure::new(); |
| |
183 let mut γ1 = DiscreteMeasure::new(); |
| |
184 let mut residual = calculate_residual(Pair(&μ, &z), opA, b); |
| |
185 let zero_z = z.similar_origin(); |
| |
186 |
| |
187 // Set up parameters |
| |
188 let op𝒟norm = op𝒟.opnorm_bound(Radon, Linfinity); |
| |
189 // TODO: maybe this PairNorm doesn't make sense here? |
| |
190 let opAnorm = opA.opnorm_bound(PairNorm(Radon, L2, L2), L2); |
| |
191 let bigθ = 0.0; //opKμ.transport_lipschitz_factor(L2Squared); |
| |
192 let bigM = 0.0; //opKμ.adjoint_product_bound(&op𝒟).unwrap().sqrt(); |
| |
193 let nKz = opKz.opnorm_bound(L2, L2); |
| |
194 let ℓ = 0.0; |
| |
195 let opIdZ = IdOp::new(); |
| |
196 let (l, l_z) = opA.adjoint_product_pair_bound(&op𝒟, &opIdZ).unwrap(); |
| |
197 // We need to satisfy |
| |
198 // |
| |
199 // τσ_dM(1-σ_p L_z)/(1 - τ L) + [σ_p L_z + σ_pσ_d‖K_z‖^2] < 1 |
| |
200 // ^^^^^^^^^^^^^^^^^^^^^^^^^ |
| |
201 // with 1 > σ_p L_z and 1 > τ L. |
| |
202 // |
| |
203 // To do so, we first solve σ_p and σ_d from standard PDPS step length condition |
| |
204 // ^^^^^ < 1. then we solve τ from the rest. |
| |
205 let σ_d = config.σd0 / nKz; |
| |
206 let σ_p = config.σp0 / (l_z + config.σd0 * nKz); |
| |
207 // Observe that = 1 - ^^^^^^^^^^^^^^^^^^^^^ = 1 - σ_{p,0} |
| |
208 // We get the condition τσ_d M (1-σ_p L_z) < (1-σ_{p,0})*(1-τ L) |
| |
209 // ⟺ τ [ σ_d M (1-σ_p L_z) + (1-σ_{p,0}) L ] < (1-σ_{p,0}) |
| |
210 let φ = 1.0 - config.σp0; |
| |
211 let a = 1.0 - σ_p * l_z; |
| |
212 let τ = config.τ0 * φ / ( σ_d * bigM * a + φ * l ); |
| |
213 let ψ = 1.0 - τ * l; |
| |
214 let β = σ_p * config.σd0 * nKz / a; // σ_p * σ_d * (nKz * nK_z) / a; |
| |
215 assert!(β < 1.0); |
| |
216 // Now we need κ‖K_μ(π_♯^1 - π_♯^0)γ‖^2 ≤ (1/θ - τ[ℓ_v + ℓ]) ∫ c_2 dγ for κ defined as: |
| |
217 let κ = σ_d * ψ / ((1.0 - β) * ψ - τ * σ_d * bigM); |
| |
218 // The factor two in the manuscript disappears due to the definition of 𝚹 being |
| |
219 // for ‖x-y‖₂² instead of c_2(x, y)=‖x-y‖₂²/2. |
| |
220 let calculate_θ = |ℓ_v, max_transport| { |
| |
221 config.transport.θ0 / (τ*(ℓ + ℓ_v) + κ * bigθ * max_transport) |
| |
222 }; |
| |
223 let mut θ_or_adaptive = match opA.preadjoint().value_diff_unit_lipschitz_factor() { |
| |
224 // We only estimate w (the uniform Lipschitz for of v), if we also estimate ℓ_v |
| |
225 // (the uniform Lipschitz factor of ∇v). |
| |
226 // We assume that the residual is decreasing. |
| |
227 Some(ℓ_v0) => TransportStepLength::AdaptiveMax{ |
| |
228 l: ℓ_v0 * b.norm2(), |
| |
229 max_transport : 0.0, |
| |
230 g : calculate_θ |
| |
231 }, |
| |
232 None => TransportStepLength::FullyAdaptive{ |
| |
233 l : 0.0, |
| |
234 max_transport : 0.0, |
| |
235 g : calculate_θ |
| |
236 }, |
| |
237 }; |
| |
238 // Acceleration is not currently supported |
| |
239 // let γ = dataterm.factor_of_strong_convexity(); |
| |
240 let ω = 1.0; |
| |
241 |
| |
242 // We multiply tolerance by τ for FB since our subproblems depending on tolerances are scaled |
| |
243 // by τ compared to the conditional gradient approach. |
| |
244 let tolerance = config.insertion.tolerance * τ * reg.tolerance_scaling(); |
| |
245 let mut ε = tolerance.initial(); |
| |
246 |
| |
247 let starH = fnH.conjugate(); |
| |
248 |
| |
249 // Statistics |
| |
250 let full_stats = |residual : &A::Observable, μ : &RNDM<F, N>, z : &Z, ε, stats| IterInfo { |
| |
251 value : residual.norm2_squared_div2() + fnR.apply(z) |
| |
252 + reg.apply(μ) + fnH.apply(/* opKμ.apply(μ) + */ opKz.apply(z)), |
| |
253 n_spikes : μ.len(), |
| |
254 ε, |
| |
255 // postprocessing: config.insertion.postprocessing.then(|| μ.clone()), |
| |
256 .. stats |
| |
257 }; |
| |
258 let mut stats = IterInfo::new(); |
| |
259 |
| |
260 // Run the algorithm |
| |
261 for state in iterator.iter_init(|| full_stats(&residual, &μ, &z, ε, stats.clone())) { |
| |
262 // Calculate initial transport |
| |
263 let Pair(v, _) = opA.preadjoint().apply(&residual); |
| |
264 //opKμ.preadjoint().apply_add(&mut v, y); |
| |
265 let z_base = z.clone(); |
| |
266 // We want to proceed as in Example 4.12 but with v and v̆ as in §5. |
| |
267 // With A(ν, z) = A_μ ν + A_z z, following Example 5.1, we have |
| |
268 // P_ℳ[F'(ν, z) + Ξ(ν, z, y)]= A_ν^*[A_ν ν + A_z z] + K_μ ν = A_ν^*A(ν, z) + K_μ ν, |
| |
269 // where A_ν^* becomes a multiplier. |
| |
270 // This is much easier with K_μ = 0, which is the only reason why are enforcing it. |
| |
271 // TODO: Write a version of initial_transport that can deal with K_μ ≠ 0. |
| |
272 |
| |
273 let (μ_base_masses, mut μ_base_minus_γ0) = initial_transport( |
| |
274 &mut γ1, &mut μ, |ν| opA.apply(Pair(ν, &z)), |
| |
275 ε, τ, &mut θ_or_adaptive, opAnorm, |
| |
276 v, &config.transport, |
| |
277 ); |
| |
278 |
| |
279 // Solve finite-dimensional subproblem several times until the dual variable for the |
| |
280 // regularisation term conforms to the assumptions made for the transport above. |
| |
281 let (d, _within_tolerances, Pair(τv̆, τz̆)) = 'adapt_transport: loop { |
| |
282 // Calculate τv̆ = τA_*(A[μ_transported + μ_transported_base]-b) |
| |
283 let residual_μ̆ = calculate_residual2(Pair(&γ1, &z), |
| |
284 Pair(&μ_base_minus_γ0, &zero_z), |
| |
285 opA, b); |
| |
286 let Pair(τv̆, τz) = opA.preadjoint().apply(residual_μ̆ * τ); |
| |
287 // opKμ.preadjoint().gemv(&mut τv̆, τ, y, 1.0); |
| |
288 |
| |
289 // Construct μ^{k+1} by solving finite-dimensional subproblems and insert new spikes. |
| |
290 let (d, within_tolerances) = insert_and_reweigh( |
| |
291 &mut μ, &τv̆, &γ1, Some(&μ_base_minus_γ0), |
| |
292 op𝒟, op𝒟norm, |
| |
293 τ, ε, &config.insertion, |
| |
294 ®, &state, &mut stats, |
| |
295 ); |
| |
296 |
| |
297 // A posteriori transport adaptation. |
| |
298 // TODO: this does not properly treat v^{k+1} - v̆^k that depends on z^{k+1}! |
| |
299 if aposteriori_transport( |
| |
300 &mut γ1, &mut μ, &mut μ_base_minus_γ0, &μ_base_masses, |
| |
301 ε, &config.transport |
| |
302 ) { |
| |
303 break 'adapt_transport (d, within_tolerances, Pair(τv̆, τz)) |
| |
304 } |
| |
305 }; |
| |
306 |
| |
307 stats.untransported_fraction = Some({ |
| |
308 assert_eq!(μ_base_masses.len(), γ1.len()); |
| |
309 let (a, b) = stats.untransported_fraction.unwrap_or((0.0, 0.0)); |
| |
310 let source = μ_base_masses.iter().map(|v| v.abs()).sum(); |
| |
311 (a + μ_base_minus_γ0.norm(Radon), b + source) |
| |
312 }); |
| |
313 stats.transport_error = Some({ |
| |
314 assert_eq!(μ_base_masses.len(), γ1.len()); |
| |
315 let (a, b) = stats.transport_error.unwrap_or((0.0, 0.0)); |
| |
316 (a + μ.dist_matching(&γ1), b + γ1.norm(Radon)) |
| |
317 }); |
| |
318 |
| |
319 // // Merge spikes. |
| |
320 // // This expects the prune below to prune γ. |
| |
321 // // TODO: This may not work correctly in all cases. |
| |
322 // let ins = &config.insertion; |
| |
323 // if ins.merge_now(&state) { |
| |
324 // if let SpikeMergingMethod::None = ins.merging { |
| |
325 // } else { |
| |
326 // stats.merged += μ.merge_spikes(ins.merging, |μ_candidate| { |
| |
327 // let ν = μ_candidate.sub_matching(&γ1)-&μ_base_minus_γ0; |
| |
328 // let mut d = &τv̆ + op𝒟.preapply(ν); |
| |
329 // reg.verify_merge_candidate(&mut d, μ_candidate, τ, ε, ins) |
| |
330 // }); |
| |
331 // } |
| |
332 // } |
| |
333 |
| |
334 // Prune spikes with zero weight. To maintain correct ordering between μ and γ1, also the |
| |
335 // latter needs to be pruned when μ is. |
| |
336 // TODO: This could do with a two-vector Vec::retain to avoid copies. |
| |
337 let μ_new = DiscreteMeasure::from_iter(μ.iter_spikes().filter(|δ| δ.α != F::ZERO).cloned()); |
| |
338 if μ_new.len() != μ.len() { |
| |
339 let mut μ_iter = μ.iter_spikes(); |
| |
340 γ1.prune_by(|_| μ_iter.next().unwrap().α != F::ZERO); |
| |
341 stats.pruned += μ.len() - μ_new.len(); |
| |
342 μ = μ_new; |
| |
343 } |
| |
344 |
| |
345 // Do z variable primal update |
| |
346 z.axpy(-σ_p/τ, τz̆, 1.0); // TODO: simplify nasty factors |
| |
347 opKz.adjoint().gemv(&mut z, -σ_p, &y, 1.0); |
| |
348 z = fnR.prox(σ_p, z); |
| |
349 // Do dual update |
| |
350 // opKμ.gemv(&mut y, σ_d*(1.0 + ω), &μ, 1.0); // y = y + σ_d K[(1+ω)(μ,z)^{k+1}] |
| |
351 opKz.gemv(&mut y, σ_d*(1.0 + ω), &z, 1.0); |
| |
352 // opKμ.gemv(&mut y, -σ_d*ω, μ_base, 1.0);// y = y + σ_d K[(1+ω)(μ,z)^{k+1} - ω (μ,z)^k]-b |
| |
353 opKz.gemv(&mut y, -σ_d*ω, z_base, 1.0);// y = y + σ_d K[(1+ω)(μ,z)^{k+1} - ω (μ,z)^k]-b |
| |
354 y = starH.prox(σ_d, y); |
| |
355 |
| |
356 // Update residual |
| |
357 residual = calculate_residual(Pair(&μ, &z), opA, b); |
| |
358 |
| |
359 // Update step length parameters |
| |
360 // let ω = pdpsconfig.acceleration.accelerate(&mut τ, &mut σ, γ); |
| |
361 |
| |
362 // Give statistics if requested |
| |
363 let iter = state.iteration(); |
| |
364 stats.this_iters += 1; |
| |
365 |
| |
366 state.if_verbose(|| { |
| |
367 plotter.plot_spikes(iter, Some(&d), Some(&τv̆), &μ); |
| |
368 full_stats(&residual, &μ, &z, ε, std::mem::replace(&mut stats, IterInfo::new())) |
| |
369 }); |
| |
370 |
| |
371 // Update main tolerance for next iteration |
| |
372 ε = tolerance.update(ε, iter); |
| |
373 } |
| |
374 |
| |
375 let fit = |μ̃ : &RNDM<F, N>| { |
| |
376 (opA.apply(Pair(μ̃, &z))-b).norm2_squared_div2() |
| |
377 //+ fnR.apply(z) + reg.apply(μ) |
| |
378 + fnH.apply(/* opKμ.apply(&μ̃) + */ opKz.apply(&z)) |
| |
379 }; |
| |
380 |
| |
381 μ.merge_spikes_fitness(config.insertion.merging, fit, |&v| v); |
| |
382 μ.prune(); |
| |
383 Pair(μ, z) |
| |
384 } |