IDGPredict.h

Go to the documentation of this file.

// Copyright (C) 2023 ASTRON (Netherlands Institute for Radio Astronomy)
// SPDX-License-Identifier: GPL-3.0-or-later
 
#ifndef DP3_STEPS_IDGPREDICT_H_
#define DP3_STEPS_IDGPREDICT_H_
 
#include <complex>
#include <functional>
#include <string>
#include <vector>
#include <utility>
 
#ifdef HAVE_IDG
#include <idg-api.h>
#endif
 
#include <casacore/ms/MeasurementSets/MeasurementSet.h>
 
#include <EveryBeam/aterms/atermbase.h>
 
#include <schaapcommon/facets/facetimage.h>
#include <schaapcommon/facets/facet.h>
 
#include <aocommon/fits/fitsreader.h>
#include <aocommon/uvector.h>
 
#include "steps/Step.h"
 
#include "common/ParameterSet.h"
#include "common/Timer.h"
 
namespace dp3 {
namespace steps {
 
class IDGPredict : public ModelDataStep {
 public:
  IDGPredict(const common::ParameterSet& parset, const std::string& prefix,
             const std::vector<aocommon::FitsReader>& readers,
             std::vector<schaapcommon::facets::Facet>&& facets,
             const std::string& ds9_regions_file = "");
 
  IDGPredict(const common::ParameterSet&, const std::string& prefix);
 
  common::Fields getRequiredFields() const override { return kUvwField; }
 
  void updateInfo(const base::DPInfo& info) override;
 
  bool process(std::unique_ptr<base::DPBuffer> buffer) override;
 
  void finish() override;
 
  void show(std::ostream&) const override;
 
  void showTimings(std::ostream&, double duration) const override;
 
  void flush();
 
  std::vector<aocommon::xt::UTensor<std::complex<float>, 3>> Predict(
      size_t direction);
 
  bool IsStarted() const;
 
  base::Direction GetFirstDirection() const override;
 
  void SetBufferSize(size_t nTimesteps);
  size_t GetBufferSize() const;
 
  static std::vector<aocommon::FitsReader> GetReaders(
      const std::vector<std::string>& fits_model_files);
 
  static std::vector<schaapcommon::facets::Facet> GetFacets(
      const std::string& ds9_regions_file, const double ra, const double dec,
      const double pixel_size_x, const double pixel_size_y,
      const size_t full_width, const size_t full_height);
 
  static std::vector<schaapcommon::facets::Facet> GetFacets(
      const std::string& ds9_regions_file, const aocommon::FitsReader& reader);
 
#ifdef HAVE_IDG
 private:
  std::vector<aocommon::UVector<float>> GetModelImages();
 
  void StartIDG();
 
  void InitializeATerms(const casacore::MeasurementSet& ms);
 
  aocommon::UVector<std::complex<float>> GetAtermValues(size_t direction) const;
 
  std::vector<const double*> InitializeUVWs();
 
  std::vector<aocommon::xt::UTensor<std::complex<float>, 3>>
  ComputeVisibilities(size_t direction, const std::vector<const double*>& uvws,
                      std::complex<float>* term_data) const;
 
  void CorrectVisibilities(
      std::vector<aocommon::xt::UTensor<std::complex<float>, 3>>& result,
      const std::complex<float>* term_data);
 
  size_t GetAllocatableBuffers(size_t memory);
 
  size_t GetSubgridCount(size_t direction) const;
 
  std::string name_;
  common::ParameterSet parset_;
 
  std::vector<schaapcommon::facets::FacetImage> images_;
  std::vector<std::unique_ptr<idg::api::BufferSet>> buffersets_;
 
  struct FacetMetaData {
    FacetMetaData(double _dl, double _dm, double _dp)
        : dl(_dl), dm(_dm), dp(_dp) {}
    double dl, dm, dp;
  };
  std::vector<FacetMetaData> meta_data_;
 
  std::vector<std::unique_ptr<base::DPBuffer>> buffers_;
 
  double ref_frequency_;
  double pixel_size_x_, pixel_size_y_;
  std::vector<aocommon::FitsReader> readers_;
 
  size_t buffer_size_;  
 
  std::vector<std::size_t> ant1_;  
  std::vector<std::size_t> ant2_;  
 
  common::NSTimer timer_;
  double max_w_;
  double max_baseline_;
  std::vector<std::pair<double, double>> directions_;
  bool save_facets_;  
 
  // Required for aterms (beam)
  std::vector<std::unique_ptr<everybeam::aterms::ATermBase>> aterms_;
  mutable std::vector<aocommon::UVector<std::complex<float>>> aterm_values_;
#endif
};
 
}  // namespace steps
}  // namespace dp3
 
#endif  // header guard