SolveData.h

Go to the documentation of this file.

// Copyright (C) 2021 ASTRON (Netherlands Institute for Radio Astronomy)
// SPDX-License-Identifier: GPL-3.0-or-later
 
#ifndef DDECAL_SOLVE_DATA_H
#define DDECAL_SOLVE_DATA_H
 
#include <cstddef>
#include <numeric>
#include <vector>
 
#include <aocommon/matrix2x2.h>
#include <aocommon/staticfor.h>
#include <xtensor/containers/xtensor.hpp>
 
#include "base/DPBuffer.h"
 
namespace dp3 {
namespace ddecal {
 
class BdaSolverBuffer;
 
template <typename MatrixType = aocommon::MC2x2F>
class SolveData {
 public:
  class ChannelBlockData {
   public:
    using const_iterator = typename std::vector<MatrixType>::const_iterator;
 
    void Resize(size_t n_visibilities, size_t n_directions) {
      data_.resize(n_visibilities);
      model_data_.resize({n_directions, n_visibilities});
      antenna_indices_.resize(n_visibilities);
      n_solutions_.resize(n_directions);
      solution_map_.resize({n_directions, n_visibilities});
    }
    void ResizeWeights(size_t n_visibilities) {
      constexpr size_t kNCorrelations = 4;
      weights_.resize({n_visibilities, kNCorrelations});
    }
    size_t NDirections() const { return model_data_.shape(0); }
    size_t NVisibilities() const { return data_.size(); }
    size_t NAntennaVisibilities(size_t antenna_index) const {
      return antenna_visibility_counts_[antenna_index];
    }
    size_t NSolutionVisibilities(size_t direction_index,
                                 uint32_t solution_index) const;
 
    uint32_t Antenna1Index(size_t visibility_index) const {
      return antenna_indices_[visibility_index].first;
    }
    uint32_t Antenna2Index(size_t visibility_index) const {
      return antenna_indices_[visibility_index].second;
    }
    uint32_t SolutionIndex(size_t direction_index,
                           size_t visibility_index) const {
      return solution_map_(direction_index, visibility_index);
    }
 
    const uint32_t* SolutionMapData() const { return solution_map_.data(); }
 
    const float& Weight(size_t index) const { return weights_(index, 0); }
    const MatrixType& Visibility(size_t index) const { return data_[index]; }
    const MatrixType& ModelVisibility(size_t direction, size_t index) const {
      return model_data_(direction, index);
    }
 
    const_iterator DataBegin() const { return data_.begin(); }
    const_iterator DataEnd() const { return data_.end(); }
 
    uint32_t NSolutionsForDirection(size_t direction_index) const {
      return n_solutions_[direction_index];
    }
 
    uint32_t NSubSolutions() const {
      return std::accumulate(n_solutions_.begin(), n_solutions_.end(),
                             uint32_t(0));
    }
 
   private:
    friend class SolveData<MatrixType>;
 
    std::vector<MatrixType> data_;
    // weights_(i, pol) contains the weight for data_[i][pol]. The vector will
    // be left empty when the algorithm does not need the weights.
    xt::xtensor<float, 2> weights_;
    // model_data_(d, i) is the model data for direction d, element i
    xt::xtensor<MatrixType, 2> model_data_;
    // Element i contains the first and second antenna corresponding with
    // data_[i] and model_data_(d, i). Using uint32_t instead of size_t reduces
    // memory usage and improves cache/memory performance (same holds for
    // n_solutions_ and solution_map_).
    std::vector<std::pair<uint32_t, uint32_t>> antenna_indices_;
    std::vector<size_t> antenna_visibility_counts_;
    std::vector<uint32_t> n_solutions_;
    xt::xtensor<uint32_t, 2> solution_map_;
  };
 
  SolveData(const std::vector<base::DPBuffer>& buffers,
            const std::vector<std::string>& direction_names,
            size_t n_channel_blocks, size_t n_antennas,
            const std::vector<size_t>& n_solutions_per_direction,
            const std::vector<int>& antennas1,
            const std::vector<int>& antennas2);
 
  SolveData(const BdaSolverBuffer& buffer, size_t n_channel_blocks,
            size_t n_antennas,
            const std::vector<size_t>& n_solutions_per_direction,
            const std::vector<int>& antennas1,
            const std::vector<int>& antennas2, bool with_weights);
 
  size_t NChannelBlocks() const { return channel_blocks_.size(); }
 
  const ChannelBlockData& ChannelBlock(size_t i) const {
    return channel_blocks_[i];
  }
 
  std::vector<std::vector<double>> GetSolutionWeights() const;
 
 private:
  void CountAntennaVisibilities();
 
  std::vector<ChannelBlockData> channel_blocks_;
  size_t n_antennas_;
};
 
using FullSolveData = SolveData<aocommon::MC2x2F>;
using DuoSolveData = SolveData<aocommon::MC2x2FDiag>;
using UniSolveData = SolveData<std::complex<float>>;
 
template <bool Add, typename MatrixType>
void DiagonalAddOrSubtractDirection(
    const typename SolveData<MatrixType>::ChannelBlockData& cb_data,
    std::vector<MatrixType>& v_residual, size_t direction, size_t n_solutions,
    const std::vector<std::complex<double>>& solutions, size_t max_threads) {
  using DComplex = std::complex<double>;
  using Complex = std::complex<float>;
  const size_t n_visibilities = cb_data.NVisibilities();
  aocommon::RunConstrainedStaticFor<size_t>(
      0, n_visibilities, max_threads,
      [&](size_t start_vis_index, size_t end_vis_index) {
        for (size_t vis_index = start_vis_index; vis_index != end_vis_index;
             ++vis_index) {
          const uint32_t antenna_1 = cb_data.Antenna1Index(vis_index);
          const uint32_t antenna_2 = cb_data.Antenna2Index(vis_index);
          const uint32_t solution_index =
              cb_data.SolutionIndex(direction, vis_index);
          const DComplex* solution_1 =
              &solutions[(antenna_1 * n_solutions + solution_index) * 2];
          const DComplex* solution_2 =
              &solutions[(antenna_2 * n_solutions + solution_index) * 2];
 
          MatrixType& data = v_residual[vis_index];
          const MatrixType& model =
              cb_data.ModelVisibility(direction, vis_index);
          // Convert first to single precision to make calculation easier
          const aocommon::MC2x2FDiag solution_a(
              static_cast<Complex>(solution_1[0]),
              static_cast<Complex>(solution_1[1]));
          const aocommon::MC2x2FDiag solution_b(
              static_cast<Complex>(solution_2[0]),
              static_cast<Complex>(solution_2[1]));
          const MatrixType contribution(solution_a * model *
                                        HermTranspose(solution_b));
          if constexpr (Add)
            data += contribution;
          else
            data -= contribution;
        }
      });
}
 
extern template class SolveData<std::complex<float>>;
extern template class SolveData<aocommon::MC2x2F>;
extern template class SolveData<aocommon::MC2x2FDiag>;
 
}  // namespace ddecal
}  // namespace dp3
 
#endif