Add predict logic

CMakeLists.txt

@@ -51,7 +51,8 @@ target_link_libraries(
   predict PUBLIC xtensor xtensor-blas xsimd OpenMP::OpenMP_CXX
                  ${CASACORE_LIBRARIES} ${BLAS_LIBRARIES})
 
-target_include_directories(predict PRIVATE ${CASACORE_INCLUDE_DIRS})
+target_include_directories(predict PRIVATE ${CASACORE_INCLUDE_DIRS}
+                                           ${aocommon_SOURCE_DIR}/include)
 
 # clang-tidy
 find_program(
@@ -75,6 +76,8 @@ if(PREDICT_BUILD_BENCHMARK)
   add_subdirectory(benchmark)
 endif()
 
+find_package(EveryBeam REQUIRED)
+
 if(PREDICT_BUILD_TESTING)
   # Enable testing
   include(CTest)

include/BeamResponse.h

+#ifndef BEAM_RESPONSE_H_
+#define BEAM_RESPONSE_H_
+
+#include <Baseline.h>
+#include <EveryBeam/pointresponse/pointresponse.h>
+#include <xtensor/xtensor.hpp>
+using Baseline = dp3::base::Baseline;
+
+size_t ComputeBeam(const size_t n_channels, double time,
+                   const everybeam::vector3r_t &srcdir,
+                   const everybeam::telescope::Telescope *telescope,
+                   aocommon::MC2x2 *beam_values, bool invert,
+                   everybeam::CorrectionMode mode, std::mutex *mutex,
+                   const std::vector<size_t> &skip_station_indices);
+
+void ApplyBeamToDataAndAdd(const std::vector<Baseline> &baselines,
+                           const std::vector<double> &frequencies,
+                           size_t n_stations,
+                           const xt::xtensor<std::complex<double>, 3> &data,
+                           xt::xtensor<std::complex<double>, 3> &model_data,
+                           const aocommon::MC2x2 *beam_values);
+#endif // BEAM_RESPONSE_H_
\ No newline at end of file

src/BeamResponse.cpp

+// Routines to compute and apply the beams
+// Copyright (C) 2020 ASTRON (Netherlands Institute for Radio Astronomy)
+// SPDX-License-Identifier: GPL-3.0-or-later
+//
+// @author Tammo Jan Dijkema
+
+// ApplyBeam is templated on the type of the data, could be complex<double>
+// or complex<float>
+
+#include <Baseline.h>
+#include <BeamResponse.h>
+#include <aocommon/matrix2x2.h>
+
+#include <memory>
+#include <vector>
+
+using dp3::base::Baseline;
+
+void ApplyWeights(const aocommon::MC2x2F &gain_a,
+                  const aocommon::MC2x2F &gain_b, float *weight) {
+  float cov[4], normGainA[4], normGainB[4];
+  for (unsigned int i = 0; i < 4; ++i) {
+    cov[i] = 1. / weight[i];
+    normGainA[i] = std::norm(gain_a.Get(i));
+    normGainB[i] = std::norm(gain_b.Get(i));
+  }
+
+  weight[0] = cov[0] * (normGainA[0] * normGainB[0]) +
+              cov[1] * (normGainA[0] * normGainB[1]) +
+              cov[2] * (normGainA[1] * normGainB[0]) +
+              cov[3] * (normGainA[1] * normGainB[1]);
+  weight[0] = 1. / weight[0];
+
+  weight[1] = cov[0] * (normGainA[0] * normGainB[2]) +
+              cov[1] * (normGainA[0] * normGainB[3]) +
+              cov[2] * (normGainA[1] * normGainB[2]) +
+              cov[3] * (normGainA[1] * normGainB[3]);
+  weight[1] = 1. / weight[1];
+
+  weight[2] = cov[0] * (normGainA[2] * normGainB[0]) +
+              cov[1] * (normGainA[2] * normGainB[1]) +
+              cov[2] * (normGainA[3] * normGainB[0]) +
+              cov[3] * (normGainA[3] * normGainB[1]);
+  weight[2] = 1. / weight[2];
+
+  weight[3] = cov[0] * (normGainA[2] * normGainB[2]) +
+              cov[1] * (normGainA[2] * normGainB[3]) +
+              cov[2] * (normGainA[3] * normGainB[2]) +
+              cov[3] * (normGainA[3] * normGainB[3]);
+  weight[3] = 1. / weight[3];
+}
+
+template <typename T>
+void ApplyBeamToData(std::vector<dp3::base::Baseline> baselines,
+                     const size_t n_channels, const size_t n_baselines,
+                     const size_t n_stations, T *data0, float *weight0,
+                     aocommon::MC2x2 *beam_values, bool doUpdateWeights) {
+  /*
+    Applies the beam to each baseline and each frequency of the
+    model patch
+  */
+
+  // Apply beam for channel ch on all baselines
+  // For mode=ARRAY_FACTOR, too much work is done here
+  // because we know that r and l are diagonal
+  for (size_t bl = 0; bl < n_baselines; ++bl) {
+    // If the beam is the same for all stations (i.e. when n_stations = 1),
+    // all baselines will have the same beam values
+    size_t index_left =
+        (n_stations == 1 ? 0 : n_channels * baselines[bl].first);
+    size_t index_right =
+        (n_stations == 1 ? 0 : n_channels * baselines[bl].second);
+    for (size_t ch = 0; ch < n_channels; ++ch) {
+      T *data = data0 + bl * 4 * n_channels + ch * 4;
+      const aocommon::MC2x2F mat(data);
+
+      const aocommon::MC2x2F left(beam_values[index_left + ch]);
+      const aocommon::MC2x2F right(beam_values[index_right + ch]);
+      const aocommon::MC2x2F result = left * mat.MultiplyHerm(right);
+      result.AssignTo(data);
+      if (doUpdateWeights) {
+        ApplyWeights(left, right, weight0 + bl * 4 * n_channels + ch * 4);
+      }
+    }
+  }
+}
+
+std::vector<size_t> UniqueStationIndices(std::vector<Baseline> baselines) {
+  /*
+      Returns a vector of unique station indices from the baselines
+  */
+  std::set<size_t> unique_stations;
+  for (const auto &bl : baselines) {
+    unique_stations.insert(bl.first);
+    unique_stations.insert(bl.second);
+  }
+  return std::vector<size_t>(unique_stations.begin(), unique_stations.end());
+}
+
+size_t ComputeBeam(double time, const std::vector<Baseline> &baselines,
+                   const std::vector<double> &chan_freqs,
+                   const everybeam::vector3r_t &srcdir,
+                   const everybeam::telescope::Telescope *telescope,
+                   aocommon::MC2x2 *beam_values, bool invert,
+                   everybeam::CorrectionMode mode,
+                   const std::vector<size_t> &skip_station_indices) {
+  /*
+    Compute the beam values for each station in a specific direction
+    and store them into beam_values
+
+    For convenience it returns the number of stations the beam
+    was computed for.
+  */
+
+  std::unique_ptr<everybeam::pointresponse::PointResponse> point_response =
+      telescope->GetPointResponse(time);
+
+  const std::vector<size_t> station_indices = UniqueStationIndices(baselines);
+  const size_t n_stations = station_indices.size();
+  const size_t n_channels = chan_freqs.size();
+
+  // Apply the beam values of both stations to the ApplyBeamed data.
+  for (size_t ch = 0; ch < n_channels; ++ch) {
+    switch (mode) {
+    case everybeam::CorrectionMode::kFull:
+    case everybeam::CorrectionMode::kElement:
+      // Fill beam_values for channel ch
+      for (size_t st = 0; st < n_stations; ++st) {
+        if (std::find(skip_station_indices.begin(), skip_station_indices.end(),
+                      st) != skip_station_indices.end()) {
+          beam_values[n_channels * st + ch] = aocommon::MC2x2::Unity();
+        } else {
+          beam_values[n_channels * st + ch] = point_response->Response(
+              mode, station_indices[st], chan_freqs[ch], srcdir);
+          if (invert) {
+            // Terminate if the matrix is not invertible.
+            [[maybe_unused]] bool status =
+                beam_values[n_channels * st + ch].Invert();
+            assert(status);
+          }
+        }
+      }
+      break;
+    case everybeam::CorrectionMode::kArrayFactor: {
+      aocommon::MC2x2 af_tmp;
+      for (size_t st = 0; st < n_stations; ++st) {
+        if (std::find(skip_station_indices.begin(), skip_station_indices.end(),
+                      st) != skip_station_indices.end()) {
+          beam_values[n_channels * st + ch] = aocommon::MC2x2::Unity();
+        } else {
+          af_tmp = point_response->Response(mode, station_indices[st],
+                                            chan_freqs[ch], srcdir);
+
+          if (invert) {
+            af_tmp = aocommon::MC2x2(1.0 / af_tmp.Get(0), 0.0, 0.0,
+                                     1.0 / af_tmp.Get(3));
+          }
+          beam_values[n_channels * st + ch] = af_tmp;
+        }
+      }
+      break;
+    }
+    case everybeam::CorrectionMode::kNone: // this should not happen
+      for (size_t st = 0; st < n_stations; ++st) {
+        beam_values[n_channels * st + ch] = aocommon::MC2x2::Unity();
+      }
+      break;
+    }
+  }
+  return n_stations;
+}
+
+void ApplyBeamToDataAndAdd(const std::vector<Baseline> &baselines,
+                           const std::vector<double> &frequencies,
+                           size_t n_stations,
+                           const xt::xtensor<std::complex<double>, 3> &data,
+                           xt::xtensor<std::complex<double>, 3> &model_data,
+                           const aocommon::MC2x2 *beam_values) {
+  /*
+    Applies the beam to each baseline and each frequency of the
+    model patch and sum the contribution to the model data
+  */
+  const size_t n_channels = frequencies.size();
+  const size_t n_baselines = baselines.size();
+
+  // Apply beam for channel ch on all baselines
+  // For mode=ARRAY_FACTOR, too much work is done here
+  // because we know that r and l are diagonal
+  for (size_t bl = 0; bl < n_baselines; ++bl) {
+    // If the beam is the same for all stations (i.e. when n_stations = 1),
+    // all baselines will have the same beam values
+    const size_t index_left =
+        (n_stations == 1 ? 0 : n_channels * baselines[bl].first);
+    const size_t index_right =
+        (n_stations == 1 ? 0 : n_channels * baselines[bl].second);
+    for (size_t ch = 0; ch < n_channels; ++ch) {
+      std::complex<double> *model_data_pointer = &model_data(bl, ch, 0);
+      // Using (model_)data2x2 avoids aliases with the (model_)data arguments.
+      const aocommon::MC2x2F data2x2(&data(bl, ch, 0));
+      const aocommon::MC2x2F model_data2x2(model_data_pointer);
+
+      const aocommon::MC2x2F left(beam_values[index_left + ch]);
+      const aocommon::MC2x2F right(beam_values[index_right + ch]);
+      const aocommon::MC2x2F result =
+          model_data2x2 + left * data2x2.MultiplyHerm(right);
+      result.AssignTo(model_data_pointer);
+    }
+  }
+}
\ No newline at end of file

src/Predict.cpp

@@ -63,6 +63,7 @@ void Predict::run(PredictPlanExec &plan, const PointSourceCollection &sources,
 
     // TODO compute should be extended to accumulate per beam_id
     plan.Compute(sources, direction_buffer);
+
     xt::view(buffer, beam_id, xt::all(), xt::all(), xt::all()) =
         direction_buffer;
   }