Refactoring of interpolation::method::SphericalVector and implementation of adjoint methods.

src/atlas/CMakeLists.txt

@@ -524,7 +524,7 @@ functionspace/detail/PointCloudInterface.cc
 functionspace/detail/CubedSphereStructure.h
 functionspace/detail/CubedSphereStructure.cc
 
-# for cubedsphere matching mesh partitioner
+# for cubedsphere matching mesh partitioner
 interpolation/method/cubedsphere/CellFinder.cc
 interpolation/method/cubedsphere/CellFinder.h
 interpolation/Vector2D.cc
@@ -541,8 +541,8 @@ interpolation/element/Triag3D.cc
 interpolation/element/Triag3D.h
 interpolation/method/Intersect.cc
 interpolation/method/Intersect.h
-interpolation/method/Ray.cc # For testing Quad
-interpolation/method/Ray.h  # For testing Quad
+interpolation/method/Ray.cc # For testing Quad
+interpolation/method/Ray.h  # For testing Quad
 
 # for BuildConvexHull3D
 
@@ -634,8 +634,11 @@ interpolation/method/knn/KNearestNeighboursBase.cc
 interpolation/method/knn/KNearestNeighboursBase.h
 interpolation/method/knn/NearestNeighbour.cc
 interpolation/method/knn/NearestNeighbour.h
-interpolation/method/sphericalvector/SphericalVector.h
+interpolation/method/sphericalvector/ComplexMatrixMultiply.h
+interpolation/method/sphericalvector/SparseMatrix.h
 interpolation/method/sphericalvector/SphericalVector.cc
+interpolation/method/sphericalvector/SphericalVector.h
+interpolation/method/sphericalvector/Types.h
 interpolation/method/structured/Cubic2D.cc
 interpolation/method/structured/Cubic2D.h
 interpolation/method/structured/Cubic3D.cc
@@ -868,7 +871,7 @@ if( NOT atlas_HAVE_ATLAS_FUNCTIONSPACE )
   unset( atlas_parallel_srcs )
   unset( atlas_output_srcs )
   unset( atlas_redistribution_srcs )
-  unset( atlas_linalg_srcs ) # only depends on array
+  unset( atlas_linalg_srcs ) # only depends on array
 endif()
 
 if( NOT atlas_HAVE_ATLAS_INTERPOLATION  )

src/atlas/interpolation/method/sphericalvector/ComplexMatrixMultiply.h

@@ -0,0 +1,222 @@
+/*
+ * (C) Crown Copyright 2024 Met Office
+ *
+ * This software is licensed under the terms of the Apache Licence Version 2.0
+ * which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+ */
+
+#pragma once
+
+#include <array>
+#include <iomanip>
+#include <limits>
+#include <sstream>
+#include <tuple>
+#include <type_traits>
+
+#include "atlas/array/ArrayView.h"
+#include "atlas/array/Range.h"
+#include "atlas/array/helpers/ArrayForEach.h"
+#include "atlas/interpolation/method/sphericalvector/Types.h"
+#include "atlas/parallel/omp/omp.h"
+
+// I don't yet fully trust NVHPC. Turning off OpenMP until we can figure out why
+// it's failing CI.
+#ifdef __NVCOMPILER
+#warning turning off OpenMP for nvhpc build
+#undef atlas_omp_parallel_for
+#define atlas_omp_parallel_for for
+#endif
+
+namespace atlas {
+namespace interpolation {
+namespace method {
+namespace detail {
+
+struct TwoVectorTag {};
+struct ThreeVectorTag {};
+constexpr auto twoVector = TwoVectorTag{};
+constexpr auto threeVector = ThreeVectorTag{};
+
+template <typename VectorTag>
+constexpr bool isVectorTag() {
+  return std::is_same_v<VectorTag, TwoVectorTag> ||
+         std::is_same_v<VectorTag, ThreeVectorTag>;
+}
+
+/// @brief   Helper class to perform complex matrix multiplications
+///
+/// @details Performs matrix multiplication between fields of 2-vectors or
+///          3-vectors. Fields must have Rank >= 2. Here, the assumption is
+///          that Dim = 0 is the horizontal dimension, and Dim = (Rank - 1) is
+///          the vector element dimension.
+template <bool InitialiseTarget>
+class ComplexMatrixMultiply {
+ public:
+  ComplexMatrixMultiply() = default;
+
+  /// @brief   Construct object from sparse matrices.
+  ///
+  /// @details complexWeights is a SparseMatrix of weights. realWeights is a
+  ///          SparseMatrix containing the magnitudes of the elements of
+  ///          complexWeights.
+  ComplexMatrixMultiply(const ComplexMatrix& complexWeights,
+                        const RealMatrix& realWeights)
+      : complexWeightsPtr_{&complexWeights}, realWeightsPtr_{&realWeights} {
+    if constexpr (checkMatrices) {
+      ATLAS_ASSERT(complexWeightsPtr_->rows() == realWeightsPtr_->rows());
+      ATLAS_ASSERT(complexWeightsPtr_->cols() == realWeightsPtr_->cols());
+      ATLAS_ASSERT(complexWeightsPtr_->nonZeros() ==
+                   realWeightsPtr_->nonZeros());
+      for (auto rowIndex = Index{0}; rowIndex < complexWeightsPtr_->rows();
+           ++rowIndex) {
+        for (auto [complexRowIter, realRowIter] = rowIters(rowIndex);
+             complexRowIter; ++complexRowIter, ++realRowIter) {
+          ATLAS_ASSERT(realRowIter);
+          ATLAS_ASSERT(realRowIter.row() == complexRowIter.row());
+          ATLAS_ASSERT(realRowIter.col() == complexRowIter.col());
+          // tinyNum ~= 2.3e-13 for double.
+          constexpr auto tinyNum = 1024 * std::numeric_limits<Real>::epsilon();
+          const auto complexMagnitude = std::abs(complexRowIter.value());
+          const auto realValue = realRowIter.value();
+          const auto error = std::abs(complexMagnitude - realValue);
+
+          const auto printError = [&]() {
+            auto strStream = std::ostringstream{};
+            strStream << "Error complex components: ";
+            strStream << std::setprecision(19) << error;
+            return strStream.str();
+          };
+          ATLAS_ASSERT_MSG(error < tinyNum, printError());
+        }
+      }
+    }
+  }
+
+  /// @brief   Apply complex matrix vector multiplication.
+  ///
+  /// @details Multiply weights by the elements in sourceView to give
+  ///          elements in targetView. If VectorType == TwoVectorTag,
+  ///          complexWeights are applied to the horizontal elements of
+  ///          sourceView. If VectorType == ThreeVectorTag, then realWeights
+  ///          are additionally applied to the vertical elements of sourceView.
+  template <typename Value, int Rank, typename VectorType,
+            typename = std::enable_if_t<isVectorTag<VectorType>()>>
+  void apply(const array::ArrayView<const Value, Rank>& sourceView,
+             array::ArrayView<Value, Rank>& targetView, VectorType) const {
+    if constexpr (std::is_same_v<VectorType, TwoVectorTag>) {
+      applyTwoVector(sourceView, targetView);
+    }
+
+    if constexpr (std::is_same_v<VectorType, ThreeVectorTag>) {
+      applyThreeVector(sourceView, targetView);
+    }
+
+    return;
+  }
+
+ private:
+#ifdef __NVCOMPILER
+#warning explicitly checking weight matrices for nvhpc build
+  static constexpr bool checkMatrices = true;
+#else
+  static constexpr bool checkMatrices = ATLAS_BUILD_TYPE_DEBUG;
+#endif
+
+  /// @brief Apply 2-vector MatMul.
+  template <typename Value, int Rank>
+  void applyTwoVector(const array::ArrayView<const Value, Rank>& sourceView,
+                      array::ArrayView<Value, Rank>& targetView) const {
+    // We could probably optimise contiguous arrays using
+    // reinterpret_cast<std::complex<double>*>(view.data()). This is fine
+    // according to the C++ standard!
+    atlas_omp_parallel_for(auto rowIndex = Index{0};
+                           rowIndex < complexWeightsPtr_->rows(); ++rowIndex) {
+      auto targetSlice = sliceColumn(targetView, rowIndex);
+      if constexpr (InitialiseTarget) {
+        targetSlice.assign(0.);
+      }
+
+      for (auto complexRowIter = complexWeightsPtr_->rowIter(rowIndex);
+           complexRowIter; ++complexRowIter) {
+        const auto colIndex = complexRowIter.col();
+        const auto complexWeight = complexRowIter.value();
+        const auto sourceSlice = sliceColumn(sourceView, colIndex);
+
+        array::helpers::arrayForEachDim(
+            slicedColumnDims<Rank>{}, std::tie(sourceSlice, targetSlice),
+            [&](auto&& sourceElem, auto&& targetElem) {
+              const auto targetVector =
+                  complexWeight * Complex(sourceElem(0), sourceElem(1));
+              targetElem(0) += targetVector.real();
+              targetElem(1) += targetVector.imag();
+            });
+      }
+    }
+  }
+
+  /// @brief Apply 3-vector MatMul.
+  template <typename Value, int Rank>
+  void applyThreeVector(const array::ArrayView<const Value, Rank>& sourceView,
+                        array::ArrayView<Value, Rank>& targetView) const {
+    atlas_omp_parallel_for(auto rowIndex = Index{0};
+                           rowIndex < complexWeightsPtr_->rows(); ++rowIndex) {
+      auto targetSlice = sliceColumn(targetView, rowIndex);
+      if constexpr (InitialiseTarget) {
+        targetSlice.assign(0.);
+      }
+
+      for (auto [complexRowIter, realRowIter] = rowIters(rowIndex);
+           complexRowIter; ++complexRowIter, ++realRowIter) {
+        const auto colIndex = complexRowIter.col();
+        const auto complexWeight = complexRowIter.value();
+        const auto realWeight = realRowIter.value();
+        const auto sourceSlice = sliceColumn(sourceView, colIndex);
+
+        array::helpers::arrayForEachDim(
+            slicedColumnDims<Rank>{}, std::tie(sourceSlice, targetSlice),
+            [&](auto&& sourceElem, auto&& targetElem) {
+              const auto targetVector =
+                  complexWeight * Complex(sourceElem(0), sourceElem(1));
+              targetElem(0) += targetVector.real();
+              targetElem(1) += targetVector.imag();
+              targetElem(2) += realWeight * sourceElem(2);
+            });
+      }
+    }
+  }
+
+  /// @brief Return a pair of complex and real row iterators
+  std::pair<ComplexMatrix::RowIter, RealMatrix::RowIter> rowIters(
+      Index rowIndex) const {
+    return std::make_pair(complexWeightsPtr_->rowIter(rowIndex),
+                          realWeightsPtr_->rowIter(rowIndex));
+  }
+
+  /// @brief Makes the slice arrayView.slice(index, Range::all()...).
+  template <typename View, typename Index>
+  static auto sliceColumn(View& arrayView, Index index) {
+    constexpr auto Rank = std::decay_t<View>::rank();
+    using RangeAll = decltype(array::Range::all());
+
+    const auto slicerArgs = std::tuple_cat(std::make_tuple(index),
+                                           std::array<RangeAll, Rank - 1>{});
+    const auto slicer = [&](auto&&... args) {
+      return arrayView.slice(args...);
+    };
+
+    return std::apply(slicer, slicerArgs);
+  }
+
+  /// @brief Defines a sequence of iteration Dims for a sliced column.
+  template <int Rank>
+  using slicedColumnDims = std::make_integer_sequence<int, Rank - 2>;
+
+  const ComplexMatrix* complexWeightsPtr_{};
+  const RealMatrix* realWeightsPtr_{};
+};
+
+}  // namespace detail
+}  // namespace method
+}  // namespace interpolation
+}  // namespace atlas

src/atlas/interpolation/method/sphericalvector/SparseMatrix.h

@@ -0,0 +1,114 @@
+/*
+ * (C) Crown Copyright 2024 Met Office
+ *
+ * This software is licensed under the terms of the Apache Licence Version 2.0
+ * which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+ */
+
+#pragma once
+
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "atlas/library/defines.h"
+#if ATLAS_HAVE_EIGEN
+
+#ifdef __NVCOMPILER
+#define EIGEN_MAX_ALIGN_BYTES 0
+#define EIGEN_UNALIGNED_VECTORIZE 0
+#define EIGEN_DONT_VECTORIZE
+#define EIGEN_DONT_PARALLELIZE
+#endif
+
+#include <Eigen/Sparse>
+#endif
+
+#include "atlas/runtime/Exception.h"
+#include "eckit/log/CodeLocation.h"
+
+namespace atlas {
+namespace interpolation {
+namespace method {
+namespace detail {
+
+#if ATLAS_HAVE_EIGEN
+/// @brief   Wrapper class for Eigen sparse matrix
+///
+/// @details Eigen sparse matrix wrapper. Allows preprocessor disabling of class
+///          if Eigen library is not present.
+template <typename Value>
+class SparseMatrix {
+ public:
+  using Index = int;
+  using EigenMatrix = Eigen::SparseMatrix<Value, Eigen::RowMajor, Index>;
+  using Triplet = Eigen::Triplet<Value, Index>;
+  using Triplets = std::vector<Triplet>;
+  using RowIter = typename EigenMatrix::InnerIterator;
+
+  SparseMatrix() = default;
+
+  SparseMatrix(Index nRows, Index nCols, const Triplets& triplets)
+      : eigenMatrix_(nRows, nCols) {
+    eigenMatrix_.setFromTriplets(triplets.begin(), triplets.end());
+  }
+
+  Index nonZeros() const { return eigenMatrix_.nonZeros(); }
+  Index rows() const { return eigenMatrix_.rows(); }
+  Index cols() const { return eigenMatrix_.cols(); }
+  RowIter rowIter(Index rowIndex) const {
+    return RowIter(eigenMatrix_, rowIndex);
+  }
+  SparseMatrix<Value> adjoint() const {
+    auto adjointMatrix = SparseMatrix{};
+    adjointMatrix.eigenMatrix_ = eigenMatrix_.adjoint();
+    return adjointMatrix;
+  }
+
+ private:
+  EigenMatrix eigenMatrix_{};
+};
+#else
+
+template <typename Value>
+class SparseMatrix {
+ public:
+  using Index = int;
+
+  class Triplet {
+   public:
+    template <typename... Args>
+    constexpr Triplet(const Args&... args) {}
+  };
+  using Triplets = std::vector<Triplet>;
+
+  class RowIter {
+   public:
+    template <typename... Args>
+    constexpr RowIter(const Args&... args) {}
+    constexpr Index row() const { return Index{}; }
+    constexpr Index col() const { return Index{}; }
+    constexpr Value value() const { return Value{}; }
+    constexpr operator bool() const { return false; }
+    constexpr RowIter& operator++() { return *this; }
+  };
+
+  constexpr SparseMatrix() = default;
+  template <typename... Args>
+  SparseMatrix(const Args&... args) {
+    throw_Exception("Atlas has been compiled without Eigen", Here());
+  }
+  constexpr Index nonZeros() const { return Index{}; }
+  constexpr Index rows() const { return Index{}; }
+  constexpr Index cols() const { return Index{}; }
+  constexpr RowIter rowIter(Index rowIndex) const { return RowIter{}; }
+  constexpr SparseMatrix<Value> adjoint() const {
+    return SparseMatrix<Value>{};
+  }
+};
+#endif
+
+}  // namespace detail
+}  // namespace method
+}  // namespace interpolation
+}  // namespace atlas