Mercurial > repos > pointsource_algs / file diff

--- a/src/forward_pdps.rs	Thu Jul 16 19:51:51 2026 +0300
+++ b/src/forward_pdps.rs	Sun Jul 19 07:30:06 2026 +0200
@@ -197,15 +197,29 @@
     };
     let mut stats = IterInfo::new();
 
+    let zy_step = |mut z, mut y, mut z_tmp| {
+        // Do z variable primal update
+        opKz_adj.apply_add(&mut z_tmp, &y);
+        z_tmp.axpy(1.0, &z, -σ_p);
+        let z_new = fnR.prox(σ_p, z_tmp);
+        //let z_new = fnR.prox(σ_p, z_tmp * (-σ_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)
+    };
+
     // Run the algorithm
     for state in iterator.iter_init(|| full_stats(&μ, &z, ε, stats.clone())) {
         // Calculate initial transport
-        let Pair(v, mut z_tmp) = f.differential(Pair(&μ, &z));
+        let Pair(v, z_tmp) = f.differential(Pair(&μ, &z));
         let mut τv = v * τ;
         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(
+        prox_penalty.insert_and_reweigh(
             &mut μ,
             &mut τv,
             τ,
@@ -216,6 +230,8 @@
             &mut stats,
         )?;
 
+        (z, y) = zy_step(z, y, z_tmp);
+
         stats.inserted += μ.len() - μ_base.len();
 
         // Merge spikes.
@@ -241,17 +257,27 @@
         // Prune spikes with zero weight.
         stats.pruned += prune_with_stats(&mut μ);
 
-        // Do z variable primal update
-        opKz_adj.apply_add(&mut z_tmp, &y);
-        z_tmp.axpy(1.0, &z, -σ_p);
-        let z_new = fnR.prox(σ_p, z_tmp);
-        //let z_new = fnR.prox(σ_p, z_tmp * (-σ_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);
-        y = starH.prox(σ_d, y);
-        z = z_new;
+        for _ in 0..config.insertion.extra_weight_optimisation_steps {
+            // Do extra weight optimisation step heuristic
+            let Pair(v, _z_tmp) = f.differential(Pair(&μ, &z));
+            τv = v * τ;
+            prox_penalty.reweigh(
+                &mut μ,
+                &mut τv,
+                τ,
+                ε,
+                &config.insertion,
+                &reg,
+                &state,
+                &mut stats,
+            )?;
+            //(z, y) = zy_step(z, y, z_tmp);
+        }
+
+        if config.insertion.extra_weight_optimisation_steps > 0 {
+            // Prune spikes with zero weight.
+            stats.pruned += prune_with_stats(&mut μ);
+        }
 
         // Update step length parameters
         // let ω = pdpsconfig.acceleration.accelerate(&mut τ, &mut σ, γ);
@@ -261,7 +287,7 @@
         stats.this_iters += 1;
 
         state.if_verbose(|| {
-            plotter.plot_spikes(iter, maybe_d.as_ref(), Some(&τv), &μ);
+            plotter.plot_spikes(iter, None, Some(&τv), &μ);
             full_stats(&μ, &z, ε, std::mem::replace(&mut stats, IterInfo::new()))
         });