From 24e07296024740e2e7aadcd1c4378d70377bfdb1 Mon Sep 17 00:00:00 2001 From: mancini <mancini@astron.nl> Date: Tue, 20 May 2025 15:52:31 +0200 Subject: [PATCH] Add predict logic --- CMakeLists.txt | 5 +- include/BeamResponse.h | 22 +++++ src/BeamResponse.cpp | 209 +++++++++++++++++++++++++++++++++++++++++ src/Predict.cpp | 1 + 4 files changed, 236 insertions(+), 1 deletion(-) create mode 100644 include/BeamResponse.h create mode 100644 src/BeamResponse.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index ce66efe..72bc4d4 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/include/BeamResponse.h b/include/BeamResponse.h new file mode 100644 index 0000000..9677da3 --- /dev/null +++ b/include/BeamResponse.h @@ -0,0 +1,22 @@ +#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 diff --git a/src/BeamResponse.cpp b/src/BeamResponse.cpp new file mode 100644 index 0000000..1fb87a7 --- /dev/null +++ b/src/BeamResponse.cpp @@ -0,0 +1,209 @@ +// 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 diff --git a/src/Predict.cpp b/src/Predict.cpp index 1ed984d..2eea9f7 100644 --- a/src/Predict.cpp +++ b/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; } -- GitLab