Merge pull request #388 from astro-informatics/mm/ApproxNormedStochasticBenchmark

Use static variables to avoid recalculating the norm

Author: tkoskela

cpp/benchmarks/stochastic_algorithm.cc

@@ -120,10 +120,76 @@ BENCHMARK_DEFINE_F(StochasticAlgoFixture, ForwardBackward)(benchmark::State &sta
   }
 }
 
+BENCHMARK_DEFINE_F(StochasticAlgoFixture, ForwardBackwardApproxNorm)(benchmark::State &state) {
+  // This functor would be defined in Purify
+  std::function<std::shared_ptr<sopt::IterationState<Vector<t_complex>>>()> random_updater =
+      [this]() {
+        H5::H5Handler h5file(m_input_data_path, m_world);
+        utilities::vis_params uv_data = H5::stochread_visibility(h5file, m_N, false);
+        uv_data.units = utilities::vis_units::radians;
+        auto phi = factory::measurement_operator_factory<Vector<t_complex>>(
+            factory::distributed_measurement_operator::mpi_distribute_image, uv_data, m_imsizex,
+            m_imsizey, 1, 1, 2, kernels::kernel_from_string.at("kb"), 4, 4);
+
+        // declaration of static variables to avoid recalculating the normalisation
+        static auto const power_method_stuff = sopt::algorithm::power_method<Vector<t_complex>>(
+            *phi, 1000, 1e-5,
+            m_world.broadcast(Vector<t_complex>::Ones(m_imsizex * m_imsizey).eval()));
+
+        static const t_real op_norm = std::get<0>(power_method_stuff);
+
+        // set the normalisation of the new phi
+        phi->set_norm(op_norm);
+
+        return std::make_shared<sopt::IterationState<Vector<t_complex>>>(uv_data.vis, phi);
+      };
+
+  // wavelets
+  auto const wavelets = factory::wavelet_operator_factory<Vector<t_complex>>(
+      factory::distributed_wavelet_operator::serial, m_sara, m_imsizey, m_imsizex);
+
+  // algorithm
+  sopt::algorithm::ImagingForwardBackward<t_complex> fb(random_updater);
+  fb.itermax(state.range(2))
+      .step_size(m_beta * sqrt(2))
+      .sigma(m_sigma * sqrt(2))
+      .regulariser_strength(m_gamma)
+      .relative_variation(1e-3)
+      .residual_tolerance(0)
+      .tight_frame(true)
+      .obj_comm(m_world);
+
+  auto gp = std::make_shared<sopt::algorithm::L1GProximal<t_complex>>(false);
+  gp->l1_proximal_tolerance(1e-4)
+      .l1_proximal_nu(1)
+      .l1_proximal_itermax(50)
+      .l1_proximal_positivity_constraint(true)
+      .l1_proximal_real_constraint(true)
+      .Psi(*wavelets);
+  fb.g_function(gp);
+
+  PURIFY_INFO("Start iteration loop");
+
+  while (state.KeepRunning()) {
+    auto start = std::chrono::high_resolution_clock::now();
+    fb();
+    auto end = std::chrono::high_resolution_clock::now();
+    state.SetIterationTime(b_utilities::duration(start, end, m_world));
+  }
+}
+
 BENCHMARK_REGISTER_F(StochasticAlgoFixture, ForwardBackward)
     ->Args({128, 10000, 10})
     ->UseManualTime()
     ->MinTime(60.0)
     ->MinWarmUpTime(5.0)
     ->Repetitions(3)  //->ReportAggregatesOnly(true)
     ->Unit(benchmark::kMillisecond);
+
+BENCHMARK_REGISTER_F(StochasticAlgoFixture, ForwardBackwardApproxNorm)
+    ->Args({128, 10000, 10})
+    ->UseManualTime()
+    ->MinTime(60.0)
+    ->MinWarmUpTime(5.0)
+    ->Repetitions(3)  //->ReportAggregatesOnly(true)
+    ->Unit(benchmark::kMillisecond);