Integrate calculateLLSTolerance into CreateLLSSolver.

Use explicit m, n, nrhs when calling CreateLLSSolver.

ddecal/gain_solvers/DiagonalSolver.cc

@@ -212,13 +212,12 @@ void DiagonalSolver::PerformIteration(const SolverBuffer& solver_buffer,
 
   // The matrices have been filled; compute the linear solution
   // for each antenna.
-  const size_t m = n_antennas_ * 2 * n_times * cur_channel_block_size;
-  const size_t n = n_directions_;
+  const size_t m = NAntennas() * n_times * cur_channel_block_size * 2;
+  const size_t n = NDirections();
   const size_t nrhs = 1;
   std::unique_ptr<LLSSolver> solver =
-      LLSSolver::Make(lls_solver_type_, m, n, nrhs);
-  solver->SetTolerance(
-      calculateLLSTolerance(iterationfraction, solverprecision));
+      CreateLLSSolver(m, n, nrhs, iterationfraction, solverprecision);
+
   for (size_t ant = 0; ant != n_antennas_; ++ant) {
     for (size_t pol = 0; pol != 2; ++pol) {
       // solve x^H in [g C] x^H  = v

ddecal/gain_solvers/ScalarSolver.cc

@@ -24,71 +24,66 @@ ScalarSolver::SolveResult ScalarSolver::Solve(
     const SolverBuffer& solver_buffer,
     std::vector<std::vector<DComplex>>& solutions, double time,
     std::ostream* stat_stream) {
+  assert(solutions.size() == NChannelBlocks());
+
   const size_t n_times = solver_buffer.NTimes();
 
-  for (size_t i = 0; i != constraints_.size(); ++i)
-    constraints_[i]->PrepareIteration(false, 0, false);
+  for (Constraint* c : GetConstraints()) c->PrepareIteration(false, 0, false);
 
-  std::vector<std::vector<DComplex>> next_solutions(n_channel_blocks_);
+  std::vector<std::vector<DComplex>> next_solutions(NChannelBlocks());
 
   SolveResult result;
-#ifndef NDEBUG
-  if (solutions.size() != n_channel_blocks_) {
-    std::cout << "Error: 'solutions' parameter does not have the right shape\n";
-    result.iterations = 0;
-    return result;
-  }
-#endif
-
-  result.results.resize(constraints_.size());
+  result.results.resize(GetConstraints().size());
 
   // Model matrix ant x [N x D] and visibility vector ant x [N x 1],
   // for each channelblock
   // The following loop allocates all structures
-  std::vector<std::vector<Matrix>> g_times_cs(n_channel_blocks_);
-  std::vector<std::vector<Matrix>> vs(n_channel_blocks_);
-  for (size_t chBlock = 0; chBlock != n_channel_blocks_; ++chBlock) {
-    next_solutions[chBlock].resize(n_directions_ * n_antennas_);
-    const size_t channel_index_start =
-                     chBlock * n_channels_ / n_channel_blocks_,
-                 channel_index_end =
-                     (chBlock + 1) * n_channels_ / n_channel_blocks_,
-                 cur_channel_block_size =
+  std::vector<std::vector<Matrix>> g_times_cs(NChannelBlocks());
+  std::vector<std::vector<Matrix>> vs(NChannelBlocks());
+  for (size_t chBlock = 0; chBlock != NChannelBlocks(); ++chBlock) {
+    next_solutions[chBlock].resize(NDirections() * NAntennas());
+    const size_t channel_index_start = FirstChannel(chBlock);
+    const size_t channel_index_end = FirstChannel(chBlock + 1);
+    const size_t cur_channel_block_size =
         channel_index_end - channel_index_start;
-    g_times_cs[chBlock].resize(n_antennas_);
-    vs[chBlock].resize(n_antennas_);
+    g_times_cs[chBlock].reserve(NAntennas());
+    vs[chBlock].reserve(NAntennas());
 
-    for (size_t ant = 0; ant != n_antennas_; ++ant) {
+    for (size_t ant = 0; ant != NAntennas(); ++ant) {
       // Model matrix [N x D] and visibility vector [N x 1]
       // Also space for the auto correlation is reserved, but they will be set
       // to 0.
-      size_t m = n_antennas_ * n_times * cur_channel_block_size * 4,
-             n = n_directions_, nrhs = 1;
-      g_times_cs[chBlock][ant] = Matrix(m, n);
-      vs[chBlock][ant] = Matrix(std::max(m, n), nrhs);
+      const size_t m = NAntennas() * n_times * cur_channel_block_size * 4;
+      const size_t n = NDirections();
+      const size_t nrhs = 1;
+      g_times_cs[chBlock].emplace_back(m, n);
+      vs[chBlock].emplace_back(std::max(m, n), nrhs);
     }
   }
 
   ///
   /// Start iterating
   ///
-  size_t iteration = 0, constrained_iterations = 0;
-  bool has_converged = false, has_previously_converged = false,
-       constraints_satisfied = false, has_stalled = false;
+  size_t iteration = 0;
+  size_t constrained_iterations = 0;
+  bool has_converged = false;
+  bool has_previously_converged = false;
+  bool constraints_satisfied = false;
+  bool has_stalled = false;
 
   std::vector<double> step_magnitudes;
-  step_magnitudes.reserve(max_iterations_);
+  step_magnitudes.reserve(GetMaxIterations());
 
   double avg_squared_diff = 1.0E4;
 
   do {
     MakeSolutionsFinite1Pol(solutions);
 
-    ParallelFor<size_t> loop(n_threads_);
-    loop.Run(0, n_channel_blocks_, [&](size_t chBlock, size_t /*thread*/) {
+    ParallelFor<size_t> loop(GetNThreads());
+    loop.Run(0, NChannelBlocks(), [&](size_t chBlock, size_t /*thread*/) {
       PerformIteration(solver_buffer, chBlock, g_times_cs[chBlock], vs[chBlock],
                        solutions[chBlock], next_solutions[chBlock],
-                       (double)(iteration + 1) / max_iterations_,
+                       double(iteration + 1) / GetMaxIterations(),
                        avg_squared_diff);
     });
 
@@ -100,16 +95,15 @@ ScalarSolver::SolveResult ScalarSolver::Solve(
 
     constraints_satisfied = true;
 
-    for (size_t i = 0; i != constraints_.size(); ++i) {
+    for (size_t i = 0; i != GetConstraints().size(); ++i) {
+      Constraint* c = GetConstraints()[i];
       // PrepareIteration() might change Satisfied(), and since we always want
       // to iterate at least once more when a constraint is not yet satisfied,
       // we evaluate Satisfied() before preparing.
-      constraints_satisfied =
-          constraints_[i]->Satisfied() && constraints_satisfied;
-      constraints_[i]->PrepareIteration(has_previously_converged, iteration,
-                                        iteration + 1 >= max_iterations_);
-      result.results[i] =
-          constraints_[i]->Apply(next_solutions, time, stat_stream);
+      constraints_satisfied = c->Satisfied() && constraints_satisfied;
+      c->PrepareIteration(has_previously_converged, iteration,
+                          iteration + 1 >= GetMaxIterations());
+      result.results[i] = c->Apply(next_solutions, time, stat_stream);
     }
 
     if (!constraints_satisfied) constrained_iterations = iteration + 1;
@@ -147,50 +141,46 @@ void ScalarSolver::PerformIteration(const SolverBuffer& solver_buffer,
                                     std::vector<DComplex>& next_solutions,
                                     double iterationfraction,
                                     double solverprecision) {
-  for (size_t ant = 0; ant != n_antennas_; ++ant) {
+  for (size_t ant = 0; ant != NAntennas(); ++ant) {
     g_times_cs[ant].SetZero();
     vs[ant].SetZero();
   }
 
-  const size_t channel_index_start =
-      channel_block_index * n_channels_ / n_channel_blocks_;
-  const size_t channel_index_end =
-      (channel_block_index + 1) * n_channels_ / n_channel_blocks_;
+  const size_t channel_index_start = FirstChannel(channel_block_index);
+  const size_t channel_index_end = FirstChannel(channel_block_index + 1);
   const size_t cur_channel_block_size = channel_index_end - channel_index_start;
   const size_t n_times = solver_buffer.NTimes();
 
   // The following loop fills the matrices for all antennas
   for (size_t time_index = 0; time_index != n_times; ++time_index) {
-    std::vector<const Complex*> model_ptrs(n_directions_);
-    for (size_t baseline = 0; baseline != ant1_.size(); ++baseline) {
-      size_t antenna1 = ant1_[baseline];
-      size_t antenna2 = ant2_[baseline];
+    std::vector<const Complex*> model_ptrs(NDirections());
+    for (size_t baseline = 0; baseline != NBaselines(); ++baseline) {
+      size_t antenna1 = AntennaIndex1(baseline);
+      size_t antenna2 = AntennaIndex2(baseline);
       if (antenna1 != antenna2) {
         Matrix& g_times_c1 = g_times_cs[antenna1];
         Matrix& v1 = vs[antenna1];
         Matrix& g_times_c2 = g_times_cs[antenna2];
         Matrix& v2 = vs[antenna2];
-        for (size_t d = 0; d != n_directions_; ++d) {
+        for (size_t d = 0; d != NDirections(); ++d) {
           model_ptrs[d] = solver_buffer.ModelDataPointer(
               time_index, d, baseline, channel_index_start);
         }
         const Complex* data_ptr = solver_buffer.DataPointer(
             time_index, baseline, channel_index_start);
         const size_t p1_to_p2[4] = {0, 2, 1, 3};
-        for (size_t ch = channel_index_start; ch != channel_index_end; ++ch) {
-          const size_t data_index1 = ch - channel_index_start +
-                                     (time_index + antenna1 * n_times) *
-                                         cur_channel_block_size,
-                       data_index2 = ch - channel_index_start +
-                                     (time_index + antenna2 * n_times) *
-                                         cur_channel_block_size;
+        for (size_t ch = 0; ch < cur_channel_block_size; ++ch) {
+          const size_t data_index1 =
+              ch + (time_index + antenna1 * n_times) * cur_channel_block_size;
+          const size_t data_index2 =
+              ch + (time_index + antenna2 * n_times) * cur_channel_block_size;
           for (size_t p1 = 0; p1 != 4; ++p1) {
             size_t p2 = p1_to_p2[p1];
-            for (size_t d = 0; d != n_directions_; ++d) {
+            for (size_t d = 0; d != NDirections(); ++d) {
               std::complex<double> predicted = *model_ptrs[d];
 
-              size_t sol_index1 = antenna1 * n_directions_ + d;
-              size_t sol_index2 = antenna2 * n_directions_ + d;
+              size_t sol_index1 = antenna1 * NDirections() + d;
+              size_t sol_index2 = antenna2 * NDirections() + d;
               g_times_c2(data_index1 * 4 + p1, d) = std::conj(
                   solutions[sol_index1] * predicted);  // using a* b* = (ab)*
               g_times_c1(data_index2 * 4 + p2, d) =
@@ -210,28 +200,27 @@ void ScalarSolver::PerformIteration(const SolverBuffer& solver_buffer,
 
   // The matrices have been filled; compute the linear solution
   // for each antenna.
-  size_t m = n_antennas_ * n_times * cur_channel_block_size * 4;
-  size_t n = n_directions_, nrhs = 1;
-
+  const size_t m = NAntennas() * n_times * cur_channel_block_size * 4;
+  const size_t n = NDirections();
+  const size_t nrhs = 1;
   std::unique_ptr<LLSSolver> solver =
-      LLSSolver::Make(lls_solver_type_, m, n, nrhs);
-  solver->SetTolerance(
-      calculateLLSTolerance(iterationfraction, solverprecision));
-  for (size_t ant = 0; ant != n_antennas_; ++ant) {
+      CreateLLSSolver(m, n, nrhs, iterationfraction, solverprecision);
+
+  for (size_t ant = 0; ant != NAntennas(); ++ant) {
     // solve x^H in [g C] x^H  = v
-    std::vector<Complex> x0(n_directions_);
-    for (size_t d = 0; d != n_directions_; ++d) {
-      x0[d] = solutions[(ant * n_directions_ + d)];
+    std::vector<Complex> x0(NDirections());
+    for (size_t d = 0; d != NDirections(); ++d) {
+      x0[d] = solutions[(ant * NDirections() + d)];
     }
     bool success =
         solver->Solve(g_times_cs[ant].data(), vs[ant].data(), x0.data());
     Matrix& x = vs[ant];
     if (success && x(0, 0) != Complex(0.0, 0.0)) {
-      for (size_t d = 0; d != n_directions_; ++d)
-        next_solutions[ant * n_directions_ + d] = x(d, 0);
+      for (size_t d = 0; d != NDirections(); ++d)
+        next_solutions[ant * NDirections() + d] = x(d, 0);
     } else {
-      for (size_t d = 0; d != n_directions_; ++d)
-        next_solutions[ant * n_directions_ + d] =
+      for (size_t d = 0; d != NDirections(); ++d)
+        next_solutions[ant * NDirections() + d] =
             std::numeric_limits<double>::quiet_NaN();
     }
   }

ddecal/gain_solvers/SolverBase.cc

@@ -244,12 +244,10 @@ void SolverBase::MakeSolutionsFinite4Pol(
 }
 
 void SolverBase::GetTimings(std::ostream& os, double duration) const {
-  if (!constraints_.empty()) {
-    for (auto& constraint : constraints_) {
+  for (Constraint* constraint : constraints_) {
     constraint->GetTimings(os, duration);
   }
 }
-}
 
 void SolverBase::SetLLSSolverType(const LLSSolverType solver,
                                   const std::pair<double, double> tolerances) {
@@ -258,19 +256,22 @@ void SolverBase::SetLLSSolverType(const LLSSolverType solver,
   lls_max_tolerance_ = tolerances.second;
 }
 
-double SolverBase::calculateLLSTolerance(const double iteration_fraction,
-                                         const double solver_precision) const {
-  double result;
-  if (lls_max_tolerance_ == lls_min_tolerance_) {
-    result = lls_max_tolerance_;
-  } else {
-    result = std::min(
-        solver_precision / 10.0,
-        std::min(lls_max_tolerance_,
-                 lls_min_tolerance_ / (iteration_fraction * iteration_fraction *
-                                       iteration_fraction)));
-  }
-  return result;
+std::unique_ptr<LLSSolver> SolverBase::CreateLLSSolver(
+    const size_t m, const size_t n, const size_t nrhs,
+    const double iteration_fraction, const double solver_precision) const {
+  std::unique_ptr<LLSSolver> solver =
+      LLSSolver::Make(lls_solver_type_, m, n, nrhs);
+
+  double tolerance = lls_max_tolerance_;
+  if (lls_max_tolerance_ != lls_min_tolerance_) {
+    const double iteration_fraction3 =
+        iteration_fraction * iteration_fraction * iteration_fraction;
+    tolerance = std::min({solver_precision / 10.0, lls_max_tolerance_,
+                          lls_min_tolerance_ / iteration_fraction3});
+  }
+  solver->SetTolerance(tolerance);
+
+  return solver;
 }
 
 }  // namespace base

ddecal/gain_solvers/SolverBase.h

@@ -25,17 +25,19 @@ class SolverBase {
   typedef std::complex<double> DComplex;
   typedef std::complex<float> Complex;
 
-  class Matrix : public std::vector<Complex> {
+  class Matrix final {
    public:
-    Matrix() : columns_(0) {}
+    Matrix() : data_(), columns_(0) {}
     Matrix(size_t columns, size_t rows)
-        : std::vector<Complex>(columns * rows, 0.0), columns_(columns) {}
-    void SetZero() { assign(size(), Complex(0.0, 0.0)); }
+        : data_(columns * rows, 0.0), columns_(columns) {}
+    void SetZero() { data_.assign(data_.size(), Complex(0.0, 0.0)); }
     Complex& operator()(size_t column, size_t row) {
-      return (*this)[column + row * columns_];
+      return data_[column + row * columns_];
     }
+    Complex* data() { return data_.data(); }
 
    private:
+    std::vector<Complex> data_;
     size_t columns_;
   };
 
@@ -148,6 +150,7 @@ class SolverBase {
    * The solving is parallelized over channel blocks.
    */
   void SetNThreads(size_t n_threads) { n_threads_ = n_threads; }
+  size_t GetNThreads() const { return n_threads_; }
 
   /**
    * Output timing information to a stream.
@@ -188,9 +191,6 @@ class SolverBase {
     return std::isfinite(val.real()) && std::isfinite(val.imag());
   }
 
-  double calculateLLSTolerance(double iteration_fraction,
-                               double solver_precision) const;
-
   bool ReachedStoppingCriterion(
       size_t iteration, bool has_converged, bool constraints_satisfied,
       const std::vector<double>& step_magnitudes) const {
@@ -204,6 +204,33 @@ class SolverBase {
     return iteration >= min_iterations_ && is_ready;
   }
 
+  size_t NAntennas() const { return n_antennas_; }
+  size_t NDirections() const { return n_directions_; }
+  size_t NChannels() const { return n_channels_; }
+  size_t NChannelBlocks() const { return n_channel_blocks_; }
+  size_t NBaselines() const { return ant1_.size(); }
+  int AntennaIndex1(size_t baseline) const { return ant1_[baseline]; }
+  int AntennaIndex2(size_t baseline) const { return ant2_[baseline]; }
+  const std::vector<Constraint*>& GetConstraints() const {
+    return constraints_;
+  }
+
+  /**
+   * @param block A channel block index, less than NChannelBlocks().
+   * @return The index of the first channel in the given channel block.
+   */
+  size_t FirstChannel(size_t block) const {
+    return block * n_channels_ / n_channel_blocks_;
+  }
+
+  /**
+   * Create an LLSSolver with the given matrix dimensions.
+   * Set the tolerance using 'iteration_fraction' and 'solver_precision'.
+   */
+  std::unique_ptr<LLSSolver> CreateLLSSolver(size_t m, size_t n, size_t nrhs,
+                                             double iteration_fraction,
+                                             double solver_precision) const;
+
   size_t n_antennas_;
   size_t n_directions_;
   size_t n_channels_;
@@ -224,6 +251,7 @@ class SolverBase {
   bool phase_only_;
   std::vector<Constraint*>
       constraints_;  // Does not own the Constraint objects.
+ private:
   LLSSolverType lls_solver_type_;
   double lls_min_tolerance_;
   double lls_max_tolerance_;