propagationmatrix.h

Go to the documentation of this file.

/* Copyright (C) 2017
 * Richard Larsson <ric.larsson@gmail.com>
 * 
 * This program is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by the
 * Free Software Foundation; either version 2, or (at your option) any
 * later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307,
 * USA. */
 
#ifndef propagationmatrix_h
#define propagationmatrix_h
 
#include "matpack_complex.h"
#include "matpackIV.h"
 
template <int N>
Eigen::Matrix<Numeric, N, N> prop_matrix(const ConstVectorView& vec) {
#define a vec[0]
#define b vec[1]
#define c vec[2]
#define d vec[3]
#define u vec[N]
#define v vec[5]
#define w vec[6]
  static_assert (N>0 and N<5, "Bad size N");
  if constexpr (N == 1) {
    return Eigen::Matrix<Numeric, 1, 1>(a);
  } else if constexpr (N==2) {
    return (Eigen::Matrix2d() << a, b, b, a).finished();
  } else if constexpr (N==3) {
    return (Eigen::Matrix3d() << a, b, c, b, a, u, c, -u, a).finished();
  } else if constexpr (N==4) {
    return (Eigen::Matrix4d() << a, b, c, d, b, a, u, v, c, -u, a, w, d, -v, -w, a).finished();
  }
#undef a
#undef b
#undef c
#undef d
#undef u
#undef v
#undef w
}
 
template <int N>
Eigen::Matrix<Numeric, N, N> prop_matrix(const ConstMatrixView& m) {
  static_assert (N>0 and N<5, "Bad size N");
  if constexpr (N == 1) {
    return Eigen::Matrix<Numeric, 1, 1>(m(0, 0));
  } else if constexpr (N == 2) {
    return (Eigen::Matrix2d() << m(0, 0), m(0, 1),
                                 m(1, 0), m(1, 1)).finished();
  } else if constexpr (N == 3) {
    return (Eigen::Matrix3d() << m(0, 0), m(0, 1), m(0, 2),
                                 m(1, 0), m(1, 1), m(1, 2),
                                 m(2, 0), m(2, 1), m(2, 2)) .finished();
  } else if constexpr (N == 4) {
    return (Eigen::Matrix4d() << m(0, 0), m(0, 1), m(0, 2), m(0, 3),
                                 m(1, 0), m(1, 1), m(1, 2), m(1, 3),
                                 m(2, 0), m(2, 1), m(2, 2), m(2, 3),
                                 m(3, 0), m(3, 1), m(3, 2), m(3, 3)) .finished();
  }
}
 
template <int N>
Eigen::Matrix<Numeric, N, N> inv_prop_matrix(const ConstVectorView& vec) {
#define a vec[0]
#define b vec[1]
#define c vec[2]
#define d vec[3]
#define u vec[N]
#define v vec[5]
#define w vec[6]
  static_assert (N>0 and N<5, "Bad size N");
  if constexpr (N == 1) {
    return Eigen::Matrix<double, 1, 1>(1 / a);
  } else if constexpr (N==2) {
    return (Eigen::Matrix2d() << a, -b, -b, a).finished() / (a * a - b * b);
  } else if constexpr (N==3) {
    return (Eigen::Matrix3d() << a * a + u * u,
          -a * b - c * u,
          -a * c + b * u,
          -a * b + c * u,
          a * a - c * c,
          -a * u + b * c,
          -a * c - b * u,
          a * u + b * c,
          a * a - b * b)
             .finished() /
         (a * a * a - a * b * b - a * c * c + a * u * u);
  } else if constexpr (N==4) {
    return (Eigen::Matrix4d() << a * a * a + a * u * u + a * v * v + a * w * w,
          -a * a * b - a * c * u - a * d * v - b * w * w + c * v * w -
              d * u * w,
          -a * a * c + a * b * u - a * d * w + b * v * w - c * v * v +
              d * u * v,
          -a * a * d + a * b * v + a * c * w - b * u * w + c * u * v -
              d * u * u,
          -a * a * b + a * c * u + a * d * v - b * w * w + c * v * w -
              d * u * w,
          a * a * a - a * c * c - a * d * d + a * w * w,
          -a * a * u + a * b * c - a * v * w + b * d * w - c * d * v +
              d * d * u,
          -a * a * v + a * b * d + a * u * w - b * c * w + c * c * v -
              c * d * u,
          -a * a * c - a * b * u + a * d * w + b * v * w - c * v * v +
              d * u * v,
          a * a * u + a * b * c - a * v * w - b * d * w + c * d * v - d * d * u,
          a * a * a - a * b * b - a * d * d + a * v * v,
          -a * a * w + a * c * d - a * u * v + b * b * w - b * c * v +
              b * d * u,
          -a * a * d - a * b * v - a * c * w - b * u * w + c * u * v -
              d * u * u,
          a * a * v + a * b * d + a * u * w + b * c * w - c * c * v + c * d * u,
          a * a * w + a * c * d - a * u * v - b * b * w + b * c * v - b * d * u,
          a * a * a - a * b * b - a * c * c + a * u * u)
             .finished() /
         (a * a * a * a - a * a * b * b - a * a * c * c - a * a * d * d +
          a * a * u * u + a * a * v * v + a * a * w * w - b * b * w * w +
          2 * b * c * v * w - 2 * b * d * u * w - c * c * v * v +
          2 * c * d * u * v - d * d * u * u);
  }
#undef a
#undef b
#undef c
#undef d
#undef u
#undef v
#undef w
}
 
template <class base>
class LazyScale {
 public:
 
  LazyScale(const base& t, const Numeric& x) : bas(t), scale(x) {}
 
  const base& bas;
 
  const Numeric& scale;
};
 
template<bool matrix>
constexpr Index need2stokes(Index nstokes_needed) {
  if constexpr (matrix) {
    if (nstokes_needed == 7) return 4;
    if (nstokes_needed == 4) return 3;
    if (nstokes_needed == 2) return 2;
    if (nstokes_needed == 1) return 1;
  }
 
  if (nstokes_needed < 5 and nstokes_needed > 0) return nstokes_needed;
 
  ARTS_ASSERT (false, "Cannot understand the input Stokes dimensions");
  return std::numeric_limits<Index>::max();
}
 
template<bool matrix>
constexpr Index stokes2need(Index nstokes) {
  if constexpr (matrix) {
      if (nstokes == 4) return 7;
      if (nstokes == 3) return 4;
      if (nstokes == 2) return 2;
      if (nstokes == 1) return 1;
  }
  
  if (nstokes < 5 and nstokes > 0) return nstokes;
  
  ARTS_ASSERT (false, "Cannot understand the input Stokes dimensions");
  return std::numeric_limits<Index>::max();
}
 
class PropagationMatrix {
 public:
  PropagationMatrix(const Index nr_frequencies = 0,
                    const Index stokes_dim = 1,
                    const Index nr_za = 1,
                    const Index nr_aa = 1,
                    const Numeric v = 0.0)
      : mfreqs(nr_frequencies),
        mstokes_dim(stokes_dim),
        mza(nr_za),
        maa(nr_aa) {
    ARTS_ASSERT(mstokes_dim < 5 and mstokes_dim > 0);
    mdata = Tensor4(maa, mza, mfreqs, NumberOfNeededVectors(), v);
  }
 
  PropagationMatrix(const PropagationMatrix& pm) = default;
 
  PropagationMatrix(PropagationMatrix&& pm) noexcept
      : mfreqs(pm.mfreqs),
        mstokes_dim(pm.mstokes_dim),
        mza(pm.mza),
        maa(pm.maa),
        mdata(std::move(pm.mdata)),
        mvectortype(pm.mvectortype) {}
 
  explicit PropagationMatrix(Tensor4 x)
      : mfreqs(x.nrows()),
        mza(x.npages()),
        maa(x.nbooks()),
        mdata(std::move(x)) {
    switch (x.ncols()) {
      case 7:
        mstokes_dim = 4;
        break;
      case 4:
        mstokes_dim = 3;
        break;
      case 2:
        mstokes_dim = 2;
        break;
      case 1:
        mstokes_dim = 1;
        break;
      default:
        ARTS_ASSERT(false, "Tensor4 not representative of PropagationMatrix");
    }
  }
 
  explicit PropagationMatrix(const ConstMatrixView& x, const bool& assume_fit = false)
      : mfreqs(1), mstokes_dim(x.ncols()), mza(1), maa(1) {
    ARTS_ASSERT(mstokes_dim < 5 and mstokes_dim > 0);
    mvectortype = false;
 
    if (not assume_fit) {
      ARTS_ASSERT(FittingShape(x),
        "Matrix not fit as propagation matrix");
    }
 
    mdata.resize(1, 1, 1, NumberOfNeededVectors());
 
    switch (mstokes_dim) {
      case 4:
        mdata(0, 0, 0, 5) = x(1, 3);
        mdata(0, 0, 0, 6) = x(2, 3);
        mdata(0, 0, 0, 3) = x(0, 3); /* FALLTHROUGH */
      case 3:
        mdata(0, 0, 0, mstokes_dim) = x(1, 2);
        mdata(0, 0, 0, 2) = x(0, 2); /* FALLTHROUGH */
      case 2:
        mdata(0, 0, 0, 1) = x(0, 1); /* FALLTHROUGH */
      case 1:
        mdata(0, 0, 0, 0) = x(0, 0); /* FALLTHROUGH */
    }
  };
 
  [[nodiscard]] Index StokesDimensions() const { return mstokes_dim; };
 
  [[nodiscard]] Index NumberOfFrequencies() const { return mfreqs; };
 
  [[nodiscard]] Index NumberOfZenithAngles() const { return mza; };
 
  [[nodiscard]] Index NumberOfAzimuthAngles() const { return maa; };
  
  [[nodiscard]] bool OK() const {return mdata.ncols() == NumberOfNeededVectors() and mdata.nrows() == mfreqs and mdata.npages() == mza and mdata.nbooks() == maa;}
 
  [[nodiscard]] bool IsEmpty() const { return not mfreqs or not mza or not maa; };
 
  [[nodiscard]] bool IsZero(const Index iv = 0,
                            const Index iz = 0,
                            const Index ia = 0) const {
    // FIXME: matpack does not do pointers in a clear manner
    // return std::any_of(mdata(ia, iz, iv, joker).begin(), mdata(ia, iz, iv, joker).end(), [](auto& x){return x not_eq 0;});
    for (auto& n : mdata(ia, iz, iv, joker))
      if (n not_eq 0) return false;
    return true;
  };
 
  [[nodiscard]] bool IsRotational(const Index iv = 0,
                                  const Index iz = 0,
                                  const Index ia = 0) const {
    if (mdata(ia, iz, iv, 0) == 0.0)
      return true;
    return false;
  };
 
  #pragma GCC diagnostic push
  #pragma GCC diagnostic ignored "-Wreturn-type"
  [[nodiscard]] Index NumberOfNeededVectors() const {
    if (not mvectortype) {
      return stokes2need<true>(mstokes_dim);
    }
   return stokes2need<false>(mstokes_dim);
  }
  #pragma GCC diagnostic pop
 
  Numeric operator()(const Index iv = 0,
                     const Index is1 = 0,
                     const Index is2 = 0,
                     const Index iz = 0,
                     const Index ia = 0) const;
 
  void AddFaraday(const Numeric& rot,
                  const Index iv = 0,
                  const Index iz = 0,
                  const Index ia = 0) {
    mdata(ia, iz, iv, mstokes_dim) += rot;
  }
 
  void SetFaraday(const Numeric& rot,
                  const Index iv = 0,
                  const Index iz = 0,
                  const Index ia = 0) {
    mdata(ia, iz, iv, mstokes_dim) = rot;
  }
 
  void MatrixAtPosition(MatrixView ret,
                        const Index iv = 0,
                        const Index iz = 0,
                        const Index ia = 0) const;
 
  PropagationMatrix& operator=(PropagationMatrix&& pm) noexcept {
    if (this != &pm) {
      mfreqs = pm.mfreqs;
      mstokes_dim = pm.mstokes_dim;
      mza = pm.mza;
      maa = pm.maa;
      mdata = std::move(pm.mdata);
      mvectortype = pm.mvectortype;
    }
    return *this;
  }
 
  PropagationMatrix& operator=(const LazyScale<PropagationMatrix>& lpms) {
    operator=(lpms.bas);
    mdata *= lpms.scale;
    return *this;
  }
 
  PropagationMatrix& operator=(const PropagationMatrix& other) = default;
 
  PropagationMatrix& operator=(const Numeric& x) {
    mdata = x;
    return *this;
  }
 
  void SetAtPosition(const PropagationMatrix& x,
                     const Index iv = 0,
                     const Index iz = 0,
                     const Index ia = 0) {
    mdata(ia, iz, iv, joker) = x.mdata(ia, iz, iv, joker);
  }
 
  void SetAtPosition(const ConstMatrixView& x,
                     const Index iv = 0,
                     const Index iz = 0,
                     const Index ia = 0);
 
  
  void SetAtPosition(const Numeric& x,
                     const Index iv = 0,
                     const Index iz = 0,
                     const Index ia = 0) {
    mdata(ia, iz, iv, joker) = x;
  }
 
  PropagationMatrix& operator/=(const PropagationMatrix& other) {
    mdata /= other.mdata;
    return *this;
  }
 
  PropagationMatrix& operator/=(const ConstVectorView& x) {
    for (Index i = 0; i < NumberOfNeededVectors(); i++) {
      for (Index j = 0; j < mza; j++) {
        for (Index k = 0; k < maa; k++) {
          mdata(k, j, joker, i) /= x;
        }
      }
    }
    return *this;
  }
 
  PropagationMatrix& operator/=(const Numeric& x) {
    mdata /= x;
    return *this;
  }
 
  void DivideAtPosition(const PropagationMatrix& x,
                        const Index iv = 0,
                        const Index iz = 0,
                        const Index ia = 0) {
    mdata(ia, iz, iv, joker) /= x.mdata(ia, iz, iv, joker);
  }
 
  void DivideAtPosition(const ConstMatrixView& x,
                        const Index iv = 0,
                        const Index iz = 0,
                        const Index ia = 0);
 
  void DivideAtPosition(const Numeric& x,
                        const Index iv = 0,
                        const Index iz = 0,
                        const Index ia = 0) {
    mdata(ia, iz, iv, joker) /= x;
  }
 
  PropagationMatrix& operator*=(const PropagationMatrix& other) {
    mdata *= other.mdata;
    return *this;
  }
 
  PropagationMatrix& operator*=(const ConstVectorView& x) {
    for (Index i = 0; i < NumberOfNeededVectors(); i++)
      for (Index j = 0; j < mza; j++)
        for (Index k = 0; k < maa; k++) mdata(k, j, joker, i) *= x;
    return *this;
  }
 
  PropagationMatrix& operator*=(const Numeric& x) {
    mdata *= x;
    return *this;
  }
 
  void MultiplyAtPosition(const PropagationMatrix& x,
                          const Index iv = 0,
                          const Index iz = 0,
                          const Index ia = 0) {
    mdata(ia, iz, iv, joker) *= x.mdata(ia, iz, iv, joker);
  }
 
  void MultiplyAtPosition(const ConstMatrixView& x,
                          const Index iv = 0,
                          const Index iz = 0,
                          const Index ia = 0);
 
  void MultiplyAtPosition(const Numeric& x,
                          const Index iv = 0,
                          const Index iz = 0,
                          const Index ia = 0) {
    mdata(ia, iz, iv, joker) *= x;
  }
 
  PropagationMatrix& operator+=(const PropagationMatrix& other) {
    mdata += other.mdata;
    return *this;
  }
 
  PropagationMatrix& operator+=(const LazyScale<PropagationMatrix>& lpms) {
    MultiplyAndAdd(lpms.scale, lpms.bas);
    return *this;
  }
 
  PropagationMatrix& operator+=(const ConstVectorView& x) {
    for (Index i = 0; i < NumberOfNeededVectors(); i++)
      for (Index j = 0; j < mza; j++)
        for (Index k = 0; k < maa; k++) mdata(k, j, joker, i) += x;
    return *this;
  }
 
  PropagationMatrix& operator+=(const Numeric& x) {
    mdata += x;
    return *this;
  }
 
  void AddAtPosition(const PropagationMatrix& x,
                     const Index iv = 0,
                     const Index iz = 0,
                     const Index ia = 0) {
    mdata(ia, iz, iv, joker) += x.mdata(ia, iz, iv, joker);
  }
 
  void AddAtPosition(const ConstMatrixView& x,
                     const Index iv = 0,
                     const Index iz = 0,
                     const Index ia = 0);
 
  void AddAtPosition(const Numeric& x,
                     const Index iv = 0,
                     const Index iz = 0,
                     const Index ia = 0) {
    mdata(ia, iz, iv, joker) += x;
  }
 
  PropagationMatrix& operator-=(const PropagationMatrix& other) {
    mdata -= other.mdata;
    return *this;
  }
 
  PropagationMatrix& operator-=(const ConstVectorView& x) {
    for (Index i = 0; i < NumberOfNeededVectors(); i++) {
      for (Index j = 0; j < mza; j++) {
        for (Index k = 0; k < maa; k++) {
          mdata(k, j, joker, i) -= x;
        }
      }
    }
    return *this;
  }
 
  PropagationMatrix& operator-=(const Numeric& x) {
    mdata -= x;
    return *this;
  }
 
  void RemoveAtPosition(const PropagationMatrix& x,
                        const Index iv = 0,
                        const Index iz = 0,
                        const Index ia = 0) {
    mdata(ia, iz, iv, joker) -= x.mdata(ia, iz, iv, joker);
  }
 
  void RemoveAtPosition(const ConstMatrixView& x,
                        const Index iv = 0,
                        const Index iz = 0,
                        const Index ia = 0);
 
  void RemoveAtPosition(const Numeric& x,
                        const Index iv = 0,
                        const Index iz = 0,
                        const Index ia = 0) {
    mdata(ia, iz, iv, joker) -= x;
  }
 
  void AddAbsorptionVectorAtPosition(const ConstVectorView& x,
                                     const Index iv = 0,
                                     const Index iz = 0,
                                     const Index ia = 0) {
    for (Index i = 0; i < mstokes_dim; i++) mdata(ia, iz, iv, i) += x[i];
  }
 
  void AddAverageAtPosition(const ConstMatrixView& mat1,
                            const ConstMatrixView& mat2,
                            const Index iv = 0,
                            const Index iz = 0,
                            const Index ia = 0);
 
  void MultiplyAndAdd(const Numeric x, const PropagationMatrix& y);
 
  void MatrixInverseAtPosition(MatrixView ret,
                               const Index iv = 0,
                               const Index iz = 0,
                               const Index ia = 0) const;
 
  [[nodiscard]] bool FittingShape(const ConstMatrixView& x) const;
 
  void GetTensor3(Tensor3View tensor3, const Index iz = 0, const Index ia = 0);
 
  VectorView Kjj(const Index iz = 0, const Index ia = 0) {
    return mdata(ia, iz, joker, 0);
  }
 
  VectorView K12(const Index iz = 0, const Index ia = 0) {
    return mdata(ia, iz, joker, 1);
  }
 
  VectorView K13(const Index iz = 0, const Index ia = 0) {
    return mdata(ia, iz, joker, 2);
  }
 
  VectorView K14(const Index iz = 0, const Index ia = 0) {
    return mdata(ia, iz, joker, 3);
  }
 
  VectorView K23(const Index iz = 0, const Index ia = 0) {
    return mdata(ia, iz, joker, mstokes_dim);
  }
 
  VectorView K24(const Index iz = 0, const Index ia = 0) {
    return mdata(ia, iz, joker, 5);
  }
 
  VectorView K34(const Index iz = 0, const Index ia = 0) {
    return mdata(ia, iz, joker, 6);
  }
 
  [[nodiscard]] ConstVectorView Kjj(const Index iz = 0, const Index ia = 0) const {
    return mdata(ia, iz, joker, 0);
  }
 
  [[nodiscard]] ConstVectorView K12(const Index iz = 0, const Index ia = 0) const {
    return mdata(ia, iz, joker, 1);
  }
 
  [[nodiscard]] ConstVectorView K13(const Index iz = 0, const Index ia = 0) const {
    return mdata(ia, iz, joker, 2);
  }
 
  [[nodiscard]] ConstVectorView K14(const Index iz = 0, const Index ia = 0) const {
    return mdata(ia, iz, joker, 3);
  }
 
  [[nodiscard]] ConstVectorView K23(const Index iz = 0, const Index ia = 0) const {
    return mdata(ia, iz, joker, mstokes_dim);
  }
 
  [[nodiscard]] ConstVectorView K24(const Index iz = 0, const Index ia = 0) const {
    return mdata(ia, iz, joker, 5);
  }
 
  [[nodiscard]] ConstVectorView K34(const Index iz = 0, const Index ia = 0) const {
    return mdata(ia, iz, joker, 6);
  }
 
  void SetZero() { mdata = 0.0; }
 
  Tensor4& Data() { return mdata; }
 
  [[nodiscard]] const Tensor4& Data() const { return mdata; }
 
  void LeftMultiplyAtPosition(MatrixView out,
                              const ConstMatrixView& in,
                              const Index iv = 0,
                              const Index iz = 0,
                              const Index ia = 0) const;
 
 
  void RightMultiplyAtPosition(MatrixView out,
                               const ConstMatrixView& in,
                               const Index iv = 0,
                               const Index iz = 0,
                               const Index ia = 0) const;
 
 protected:
  Index mfreqs, mstokes_dim;
  Index mza, maa;
  Tensor4 mdata;
  bool mvectortype{false};
};
 
using ArrayOfPropagationMatrix = Array<PropagationMatrix>;
using ArrayOfArrayOfPropagationMatrix = Array<ArrayOfPropagationMatrix>;
 
void compute_transmission_matrix(Tensor3View T,
                                 const Numeric& r,
                                 const PropagationMatrix& upper_level,
                                 const PropagationMatrix& lower_level,
                                 const Index iz = 0,
                                 const Index ia = 0);
 
 
void compute_transmission_matrix_from_averaged_matrix_at_frequency(
    MatrixView T,
    const Numeric& r,
    const PropagationMatrix& averaged_propagation_matrix,
    const Index iv,
    const Index iz = 0,
    const Index ia = 0);
 
void compute_transmission_matrix_and_derivative(
    Tensor3View T,
    Tensor4View dT_upper_level,
    Tensor4View dT_lower_level,
    const Numeric& r,
    const PropagationMatrix& upper_level,
    const PropagationMatrix& lower_level,
    const Array<PropagationMatrix>& dprop_mat_upper_level,
    const Array<PropagationMatrix>& dprop_mat_lower_level,
    const Numeric& dr_dTu = 0.0,
    const Numeric& dr_dTl = 0.0,
    const Index it = -1,
    const Index iz = 0,
    const Index ia = 0);
 
std::ostream& operator<<(std::ostream& os, const PropagationMatrix& pm);
 
std::ostream& operator<<(std::ostream& os, const ArrayOfPropagationMatrix& apm);
 
std::ostream& operator<<(std::ostream& os,
                         const ArrayOfArrayOfPropagationMatrix& aapm);
 
class StokesVector final : public PropagationMatrix {
 public:
 
  StokesVector(const Index nr_frequencies = 0,
               const Index stokes_dim = 1,
               const Index nr_za = 1,
               const Index nr_aa = 1,
               const Numeric& v = 0.0) {
    mvectortype = true;
    mfreqs = nr_frequencies;
    mstokes_dim = stokes_dim;
    mza = nr_za;
    maa = nr_aa;
    ARTS_ASSERT(mstokes_dim < 5 and mstokes_dim > 0);
    mdata = Tensor4(maa, mza, mfreqs, mstokes_dim, v);
  };
 
  explicit StokesVector(Tensor4 x) {
    mfreqs = x.nrows();
    mstokes_dim = x.ncols();
    mza = x.npages();
    maa = x.nbooks();
    mdata = std::move(x);
    mvectortype = true;
    ARTS_ASSERT (not (mstokes_dim > 4 or mstokes_dim < 1),
      "Tensor4 is bad for StokesVector");
  }
 
  explicit StokesVector(const ConstVectorView& x) {
    mfreqs = 1;
    mstokes_dim = x.nelem();
    mza = 1;
    maa = 1;
    ARTS_ASSERT(mstokes_dim < 5 and mstokes_dim > 0);
    mvectortype = true;
    mdata.resize(1, 1, 1, mstokes_dim);
    for (Index i = 0; i < mstokes_dim; i++) mdata(0, 0, 0, i) = x[i];
  };
 
  StokesVector(const StokesVector& a,
               const StokesVector& b,
               const Numeric& scale = 0.5) {
    mfreqs = a.NumberOfFrequencies();
    mstokes_dim = a.StokesDimensions();
    mza = a.NumberOfZenithAngles();
    maa = a.NumberOfAzimuthAngles();
    mdata.resize(maa, mza, mfreqs, mstokes_dim);
    mvectortype = true;
 
    for (Index i = 0; i < maa; i++)
      for (Index j = 0; j < mza; j++)
        for (Index k = 0; k < mfreqs; k++) {
          switch (mstokes_dim) {
            case 4:
              mdata(i, j, k, 3) = (a.mdata(i, j, k, 3) + b.mdata(i, j, k, 3)) *
                                  scale; /* FALLTHROUGH */
            case 3:
              mdata(i, j, k, 2) = (a.mdata(i, j, k, 2) + b.mdata(i, j, k, 2)) *
                                  scale; /* FALLTHROUGH */
            case 2:
              mdata(i, j, k, 1) = (a.mdata(i, j, k, 1) + b.mdata(i, j, k, 1)) *
                                  scale; /* FALLTHROUGH */
            case 1:
              mdata(i, j, k, 0) = (a.mdata(i, j, k, 0) + b.mdata(i, j, k, 0)) *
                                  scale; /* FALLTHROUGH */
          }
        }
  };
 
  [[nodiscard]] Index NumberOfNeededVectors() const { return mstokes_dim; }
 
  StokesVector& operator+=(const PropagationMatrix& x) {
    mdata += x.Data()(joker, joker, joker, Range(0, mstokes_dim, 1));
    return *this;
  }
 
  StokesVector& operator+=(const LazyScale<PropagationMatrix>& lpms) {
    MultiplyAndAdd(lpms.scale, lpms.bas);
    return *this;
  }
 
  StokesVector& operator=(const PropagationMatrix& x) {
    mstokes_dim = x.StokesDimensions();
    mfreqs = x.NumberOfFrequencies();
    mza = x.NumberOfZenithAngles();
    maa = x.NumberOfAzimuthAngles();
    mdata.resize(maa, mza, mfreqs, mstokes_dim);
    mdata(joker, joker, joker, Range(0, mstokes_dim, 1)) = x.Data()(joker, joker, joker, Range(0, mstokes_dim, 1));
    return *this;
  }
 
  StokesVector& operator=(const LazyScale<PropagationMatrix>& lpms) {
    operator=(lpms.bas);
    mdata *= lpms.scale;
    return *this;
  }
 
  StokesVector& operator=(const Numeric& x) {
    mdata = x;
    return *this;
  }
 
  void MultiplyAndAdd(const Numeric x, const PropagationMatrix& y) {
    ARTS_ASSERT(mstokes_dim == y.StokesDimensions());
    ARTS_ASSERT(mfreqs == y.NumberOfFrequencies());
    ARTS_ASSERT(mza == y.NumberOfZenithAngles());
    ARTS_ASSERT(maa == y.NumberOfAzimuthAngles());
 
    const Tensor4& data = y.Data();
 
    for (Index i = 0; i < maa; i++) {
      for (Index j = 0; j < mza; j++) {
        for (Index k = 0; k < mfreqs; k++) {
          switch (mstokes_dim) {
            case 4:
              mdata(i, j, k, 3) += x * data(i, j, k, 3); /* FALLTHROUGH */
            case 3:
              mdata(i, j, k, 2) += x * data(i, j, k, 2); /* FALLTHROUGH */
            case 2:
              mdata(i, j, k, 1) += x * data(i, j, k, 1); /* FALLTHROUGH */
            case 1:
              mdata(i, j, k, 0) += x * data(i, j, k, 0); /* FALLTHROUGH */
          }
        }
      }
    }
  }
 
  VectorView VectorAtPosition(const Index iv = 0,
                              const Index iz = 0,
                              const Index ia = 0) {
    return mdata(ia, iz, iv, joker);
  }
 
  [[nodiscard]] ConstVectorView VectorAtPosition(const Index iv = 0,
                                                 const Index iz = 0,
                                                 const Index ia = 0) const {
    return mdata(ia, iz, iv, joker);
  }
 
  void SetAtPosition(const ConstVectorView& x,
                     const Index iv = 0,
                     const Index iz = 0,
                     const Index ia = 0) {
    mdata(ia, iz, iv, joker) = x;
  }
 
  void AddAverageAtPosition(const ConstVectorView& vec1,
                            const ConstVectorView& vec2,
                            const Index iv = 0,
                            const Index iz = 0,
                            const Index ia = 0) {
    switch (mstokes_dim) {
      case 4:
        mdata(ia, iz, iv, 3) += (vec1[3] + vec2[3]) * 0.5; /* FALLTHROUGH */
      case 3:
        mdata(ia, iz, iv, 2) += (vec1[2] + vec2[2]) * 0.5; /* FALLTHROUGH */
      case 2:
        mdata(ia, iz, iv, 1) += (vec1[1] + vec2[1]) * 0.5; /* FALLTHROUGH */
      case 1:
        mdata(ia, iz, iv, 0) += (vec1[0] + vec2[0]) * 0.5; /* FALLTHROUGH */
    }
  }
 
  [[nodiscard]] bool IsPolarized(const Index iv = 0,
                                 const Index iz = 0,
                                 const Index ia = 0) const {
    switch (mstokes_dim) {
      case 4:
        if (K14(iz, ia)[iv] not_eq 0.0) return true; /* FALLTHROUGH */
      case 3:
        if (K13(iz, ia)[iv] not_eq 0.0) return true; /* FALLTHROUGH */
      case 2:
        if (K12(iz, ia)[iv] not_eq 0.0) return true; /* FALLTHROUGH */
    }
    return false;
  }
 
  [[nodiscard]] bool IsUnpolarized(const Index iv = 0,
                                   const Index iz = 0,
                                   const Index ia = 0) const {
    return not IsPolarized(iv, iz, ia);
  }
  
  [[nodiscard]] bool allZeroes() const {
    for (Index i=0; i<maa; i++)
      for (Index j=0; j<mza; j++)
        for (Index k=0; k<mfreqs; k++)
          for (Index m=0; m<mstokes_dim; m++)
            if (mdata(i, j, k, m) not_eq 0)
              return false;
    return true;
  }
};
 
using ArrayOfStokesVector = Array<StokesVector>;
using ArrayOfArrayOfStokesVector = Array<ArrayOfStokesVector>;
using ArrayOfArrayOfArrayOfStokesVector = Array<ArrayOfArrayOfStokesVector>;
 
std::ostream& operator<<(std::ostream& os, const StokesVector& pm);
std::ostream& operator<<(std::ostream& os, const ArrayOfStokesVector& apm);
std::ostream& operator<<(std::ostream& os,
                         const ArrayOfArrayOfStokesVector& aapm);
 
inline LazyScale<PropagationMatrix> operator*(const PropagationMatrix& pm,
                                              const Numeric& x) {
  return LazyScale<PropagationMatrix>(pm, x);
}
 
inline LazyScale<PropagationMatrix> operator*(const Numeric& x,
                                              const PropagationMatrix& pm) {
  return LazyScale<PropagationMatrix>(pm, x);
}
 
 
template <class T>
bool bad_propmat(const Array<T>& main,
                 const Vector& f_grid,
                 const Index sd = 1) noexcept {
  const Index nf = f_grid.nelem();
  for (auto& var : main) {
    const bool bad_stokes = sd > var.StokesDimensions();
    const bool bad_freq = nf not_eq var.NumberOfFrequencies();
    if (bad_freq or bad_stokes) return true;
  }
 
  return false;
}
 
#endif  //propagationmatrix_h