Spectrum.h

// Spectrum.h: Spectral fit for a given source model
//
// Copyright (C) 2020 ASTRON (Netherlands Institute for Radio Astronomy)
// SPDX-License-Identifier: Apache2.0

/// \file
/// \brief Spectral fit for a given source model

#ifndef SPECTRUM_H_
#define SPECTRUM_H_

#include <Stokes.h>
#include <StokesVector.h>
#include <casacore/casa/BasicSL/Constants.h>
#include <complex>
#include <xtensor/xcomplex.hpp>
#include <xtensor/xshape.hpp>
#include <xtensor/xtensor.hpp>
#include <xtensor/xview.hpp>

using Stokes = dp3::base::Stokes;
enum CrossCorrelation { XX = 0, XY = 1, YX = 2, YY = 3 };

class Spectrum {
public:
  Spectrum()
      : reference_flux_(), reference_frequency_(0.0), polarization_angle_(0.0),
        polarization_factor_(0.0), rotation_measure_(0.0),
        has_rotation_measure_(false), has_logarithmic_spectral_index_(false) {}

  void Evaluate(const xt::xtensor<double, 1> &frequencies,
                StokesVector &result) const {
    result.I.resize(frequencies.size());
    result.Q.resize(frequencies.size());
    result.U.resize(frequencies.size());
    result.V.resize(frequencies.size());

    double I0 = reference_flux_.I;
    double V0 = reference_flux_.V;

    if (spectral_terms_.size() > 0 && has_logarithmic_spectral_index_) {
      for (size_t f_idx = 0; f_idx < frequencies.size(); f_idx++) {
        const double spectral_term = EvaluateLogSpectrum(
            frequencies[f_idx], reference_frequency_, spectral_terms_);
        result.I[f_idx] = I0 * spectral_term;
        result.V[f_idx] = V0 * spectral_term;
      }
    } else if (spectral_terms_.size()) {
      for (size_t f_idx = 0; f_idx < frequencies.size(); f_idx++) {
        const double spectral_term = EvaluateSpectrum(
            frequencies[f_idx], reference_frequency_, spectral_terms_);
        result.I[f_idx] = I0 + spectral_term;
        result.V[f_idx] = V0 + spectral_term;
      }
    } else {
      for (size_t f_idx = 0; f_idx < frequencies.size(); f_idx++) {
        result.I[f_idx] = I0;
        result.V[f_idx] = V0;
      }
    }

    if (has_rotation_measure_) {
      // This way of computing is faster then using the xtensor extension
      for (size_t f_idx = 0; f_idx < frequencies.size(); f_idx++) {
        const double lambda = casacore::C::c / frequencies[f_idx];
        const double chi =
            2.0 * (polarization_angle_ + (rotation_measure_ * lambda * lambda));
        const double stokesQU = result.I[f_idx] * polarization_factor_;
        result.Q[f_idx] = stokesQU * cos(chi);
        result.U[f_idx] = stokesQU * sin(chi);
      }
    }
  }

  void SetSpectralTerms(double refFreq, bool isLogarithmicSpectralIndex,
                        const xt::xtensor<double, 1> &terms = {}) {
    reference_frequency_ = refFreq;
    has_logarithmic_spectral_index_ = isLogarithmicSpectralIndex;
    spectral_terms_ = terms;
  }

  const xt::xtensor<double, 1> &GetSpectralTerms() const {
    return spectral_terms_;
  }

  void ClearSpectralTerms() { spectral_terms_ = xt::xtensor<double, 1>(); }

  double GetReferenceFrequency() const { return reference_frequency_; }
  bool HasLogarithmicSpectralIndex() const {
    return has_logarithmic_spectral_index_;
  }
  bool HasSpectralTerms() const { return spectral_terms_.size() > 0; }
  bool HasRotationMeasure() const { return has_rotation_measure_; }
  void SetPolarization(double angle, double factor) {
    polarization_angle_ = angle;
    polarization_factor_ = factor;
  }

  double GetPolarizationAngle() const { return polarization_angle_; }
  double GetPolarizationFactor() const { return polarization_factor_; }

  void SetRotationMeasure(double rm, bool has_rm = true) {
    rotation_measure_ = rm;
    has_rotation_measure_ = has_rm;
  }

  double GetRotationMeasure() const { return rotation_measure_; }

  void ClearRotationMeasure() {
    rotation_measure_ = 0.0;
    has_rotation_measure_ = false;
  }

  void SetReferenceFlux(const Stokes &flux) { reference_flux_ = flux; }

  const Stokes &GetReferenceFlux() const { return reference_flux_; }

  void EvaluateCrossCorrelations(
      const xt::xtensor<double, 1> &frequencies,
      xt::xtensor<std::complex<double>, 2> &result) const {
    xt::xtensor<double, 1>::shape_type kSpectralCorrectionShape{
        frequencies.size()};
    xt::xtensor<double, 2>::shape_type kRotationCoefficientsShape{
        2, frequencies.size()};
    xt::xtensor<double, 1> spectral_correction =
        xt::zeros<double>(kSpectralCorrectionShape);
    xt::xtensor<double, 2> rotation_coefficients =
        xt::zeros<double>(kRotationCoefficientsShape);

    result.resize({4, frequencies.size()});

    if (spectral_terms_.size() > 0 && has_logarithmic_spectral_index_) {
      for (size_t f_idx = 0; f_idx < frequencies.size(); f_idx++) {
        spectral_correction[f_idx] = EvaluateLogSpectrum(
            frequencies[f_idx], reference_frequency_, spectral_terms_);
      }
    } else if (spectral_terms_.size() > 0) {
      for (size_t f_idx = 0; f_idx < frequencies.size(); f_idx++) {
        spectral_correction[f_idx] = EvaluateSpectrum(
            frequencies[f_idx], reference_frequency_, spectral_terms_);
      }
    }

    if (has_rotation_measure_) {
      for (size_t f_idx = 0; f_idx < frequencies.size(); f_idx++) {
        const double kLambda = casacore::C::c / frequencies[f_idx];
        const double kChi = 2.0 * (polarization_angle_ +
                                   (rotation_measure_ * kLambda * kLambda));
        const double stokesQU = polarization_factor_;
        rotation_coefficients(0, f_idx) = stokesQU * cos(kChi);
        rotation_coefficients(1, f_idx) = stokesQU * sin(kChi);
      }
    }

    xt::xtensor<double, 1> I;
    xt::xtensor<double, 1> V;

    if (has_logarithmic_spectral_index_) {
      I = spectral_correction * reference_flux_.I;
      V = spectral_correction * reference_flux_.V;
    } else {
      I = spectral_correction + reference_flux_.I;
      V = spectral_correction + reference_flux_.V;
    }
    auto Q = I * xt::view(rotation_coefficients, 0, xt::all());
    auto U = I * xt::view(rotation_coefficients, 1, xt::all());

    xt::real(xt::view(result, CrossCorrelation::XX, xt::all())) = I + Q;
    xt::imag(xt::view(result, CrossCorrelation::XX, xt::all())) = 0.0;
    xt::real(xt::view(result, CrossCorrelation::XY, xt::all())) = U;
    xt::imag(xt::view(result, CrossCorrelation::XY, xt::all())) = V;
    xt::real(xt::view(result, CrossCorrelation::YX, xt::all())) = U;
    xt::imag(xt::view(result, CrossCorrelation::YX, xt::all())) = -V;
    xt::real(xt::view(result, CrossCorrelation::YY, xt::all())) = I - Q;
    xt::imag(xt::view(result, CrossCorrelation::YY, xt::all())) = 0.0;
  }

private:
  inline double
  EvaluatePolynomial(double x, const xt::xtensor<double, 1> &parameters) const {
    double partial = 0.0;
    for (auto it = spectral_terms_.rbegin(), end = spectral_terms_.rend();
         it != end; ++it) {
      partial = std::fma(partial, x, *it);
    }
    return partial;
  }
  inline double
  EvaluateLogPolynomial(double x,
                        const xt::xtensor<double, 1> &parameters) const {
    const double base = log10(x);
    return EvaluatePolynomial(base, parameters);
  }

  inline double
  EvaluateLogSpectrum(const double frequency, const double reference_frequency,
                      const xt::xtensor<double, 1> &parameters) const {
    double x = frequency / reference_frequency;
    return pow(10, x * EvaluateLogPolynomial(x, parameters));
  }
  inline double
  EvaluateSpectrum(const double frequency, const double reference_frequency,
                   const xt::xtensor<double, 1> &parameters) const {
    const double x = frequency / reference_frequency - 1.0;

    return EvaluatePolynomial(x, parameters) * x;
  }

  Stokes reference_flux_;

  xt::xtensor<double, 1> spectral_terms_;
  double reference_frequency_;
  bool has_logarithmic_spectral_index_;

  double polarization_angle_;
  double polarization_factor_;

  double rotation_measure_;
  bool has_rotation_measure_;
};

#endif // SPECTRUM_H_