Cubed sphere interpolation

src/atlas/CMakeLists.txt

@@ -72,6 +72,7 @@ option/TransOptions.cc
 projection.h
 projection/Projection.cc
 projection/Projection.h
+projection/Jacobian.h
 projection/detail/CubedSphereEquiAnglProjection.cc
 projection/detail/CubedSphereEquiAnglProjection.h
 projection/detail/CubedSphereEquiDistProjection.cc
@@ -203,6 +204,8 @@ grid/detail/partitioner/MatchingMeshPartitioner.h
 grid/detail/partitioner/MatchingMeshPartitioner.cc
 grid/detail/partitioner/MatchingMeshPartitionerBruteForce.cc
 grid/detail/partitioner/MatchingMeshPartitionerBruteForce.h
+grid/detail/partitioner/MatchingMeshPartitionerCubedSphere.cc
+grid/detail/partitioner/MatchingMeshPartitionerCubedSphere.h
 grid/detail/partitioner/MatchingMeshPartitionerLonLatPolygon.cc
 grid/detail/partitioner/MatchingMeshPartitionerLonLatPolygon.h
 grid/detail/partitioner/MatchingMeshPartitionerSphericalPolygon.cc
@@ -559,6 +562,10 @@ interpolation/method/fe/FiniteElement.cc
 interpolation/method/fe/FiniteElement.h
 interpolation/method/bil/BilinearRemapping.cc
 interpolation/method/bil/BilinearRemapping.h
+interpolation/method/cubedsphere/CellFinder.cc
+interpolation/method/cubedsphere/CellFinder.h
+interpolation/method/cubedsphere/CubedSphereBilinear.cc
+interpolation/method/cubedsphere/CubedSphereBilinear.h
 interpolation/method/knn/GridBox.cc
 interpolation/method/knn/GridBox.h
 interpolation/method/knn/GridBoxAverage.cc

src/atlas/functionspace/NodeColumns.cc

@@ -358,6 +358,27 @@ void dispatch_haloExchange(Field& field, const parallel::HaloExchange& halo_exch
     }
     field.set_dirty(false);
 }
+
+template <int RANK>
+void dispatch_adjointHaloExchange(Field& field, const parallel::HaloExchange& halo_exchange, bool on_device) {
+    if (field.datatype() == array::DataType::kind<int>()) {
+        halo_exchange.template execute_adjoint<int, RANK>(field.array(), on_device);
+    }
+    else if (field.datatype() == array::DataType::kind<long>()) {
+        halo_exchange.template execute_adjoint<long, RANK>(field.array(), on_device);
+    }
+    else if (field.datatype() == array::DataType::kind<float>()) {
+        halo_exchange.template execute_adjoint<float, RANK>(field.array(), on_device);
+    }
+    else if (field.datatype() == array::DataType::kind<double>()) {
+        halo_exchange.template execute_adjoint<double, RANK>(field.array(), on_device);
+    }
+    else {
+        throw_Exception("datatype not supported", Here());
+    }
+    field.set_dirty(false);
+}
+
 }  // namespace
 
 void NodeColumns::haloExchange(const FieldSet& fieldset, bool on_device) const {
@@ -382,11 +403,40 @@ void NodeColumns::haloExchange(const FieldSet& fieldset, bool on_device) const {
     }
 }
 
+void NodeColumns::adjointHaloExchange(const FieldSet& fieldset, bool on_device) const {
+    for (idx_t f = 0; f < fieldset.size(); ++f) {
+        Field& field = const_cast<FieldSet&>(fieldset)[f];
+        switch (field.rank()) {
+            case 1:
+                dispatch_adjointHaloExchange<1>(field, halo_exchange(), on_device);
+                break;
+            case 2:
+                dispatch_adjointHaloExchange<2>(field, halo_exchange(), on_device);
+                break;
+            case 3:
+                dispatch_adjointHaloExchange<3>(field, halo_exchange(), on_device);
+                break;
+            case 4:
+                dispatch_adjointHaloExchange<4>(field, halo_exchange(), on_device);
+                break;
+            default:
+                throw_Exception("Rank not supported", Here());
+        }
+    }
+}
+
 void NodeColumns::haloExchange(const Field& field, bool on_device) const {
     FieldSet fieldset;
     fieldset.add(field);
     haloExchange(fieldset, on_device);
 }
+
+void NodeColumns::adjointHaloExchange(const Field& field, bool) const {
+    FieldSet fieldset;
+    fieldset.add(field);
+    adjointHaloExchange(fieldset);
+}
+
 const parallel::HaloExchange& NodeColumns::halo_exchange() const {
     if (halo_exchange_) {
         return *halo_exchange_;

src/atlas/functionspace/NodeColumns.h

@@ -85,6 +85,9 @@ class NodeColumns : public functionspace::FunctionSpaceImpl {
     void haloExchange(const Field&, bool on_device = false) const override;
     const parallel::HaloExchange& halo_exchange() const;
 
+    void adjointHaloExchange(const FieldSet&, bool on_device = false) const override;
+    void adjointHaloExchange(const Field&, bool on_device = false) const override;
+
     void gather(const FieldSet&, FieldSet&) const;
     void gather(const Field&, Field&) const;
     const parallel::GatherScatter& gather() const;

src/atlas/functionspace/detail/CubedSphereStructure.cc

@@ -121,16 +121,13 @@ idx_t CubedSphereStructure::index(idx_t t, idx_t i, idx_t j) const {
 }
 
 bool CubedSphereStructure::is_valid_index(idx_t t, idx_t i, idx_t j) const {
-
-
     // Check if t is in range.
     if (t < 0 || t > 5) {
         return false;
     }
 
     // Check if i and j are in range in index method.
-    if ( i < i_begin(t) || i >= i_end(t) ||
-         j < j_begin(t) || j >= j_end(t)) {
+    if (i < i_begin(t) || i >= i_end(t) || j < j_begin(t) || j >= j_end(t)) {
         return false;
     }
 
@@ -140,7 +137,6 @@ bool CubedSphereStructure::is_valid_index(idx_t t, idx_t i, idx_t j) const {
     }
 
     return true;
-
 }
 
 Field CubedSphereStructure::tij() const {

src/atlas/grid/CubedSphereGrid.h

@@ -263,7 +263,9 @@ class IterateTIJ {
 
 class CubedSphereGrid : public Grid {
 public:
-    using grid_t = grid::detail::grid::CubedSphere;
+    using grid_t                = grid::detail::grid::CubedSphere;
+    using CubedSphereProjection = projection::detail::CubedSphereProjectionBase;
+    using CubedSphereTiles      = grid::CubedSphereTiles;
 
 public:
     CubedSphereGrid();
@@ -296,16 +298,11 @@ class CubedSphereGrid : public Grid {
     inline int N() const { return grid_->N(); }
 
     /// @brief return tiles object.
-    inline atlas::grid::CubedSphereTiles tiles() const { return grid_->tiles(); }
+    inline CubedSphereTiles tiles() const { return grid_->tiles(); }
 
     /// @brief return cubed sphere projection object.
-    inline const projection::detail::CubedSphereProjectionBase& cubedSphereProjection() const {
-
-        const auto projPtr =
-            dynamic_cast<const projection::detail::CubedSphereProjectionBase*>
-            (projection().get());
-
-        return *projPtr;
+    inline const CubedSphereProjection& cubedSphereProjection() const {
+        return dynamic_cast<const CubedSphereProjection&>(*projection().get());
     };
 
     temporary::IterateTIJ tij() const { return temporary::IterateTIJ(*grid_); }

src/atlas/grid/Tiles.cc

@@ -12,6 +12,7 @@
 #include "atlas/grid/detail/tiles/FV3Tiles.h"
 #include "atlas/grid/detail/tiles/LFRicTiles.h"
 #include "atlas/grid/detail/tiles/Tiles.h"
+#include "atlas/projection/Jacobian.h"
 
 
 namespace atlas {
@@ -76,7 +77,7 @@ const PointXY& CubedSphereTiles::tileCentre(size_t t) const {
     return get()->tileCentre(t);
 }
 
-const Jacobian& CubedSphereTiles::tileJacobian(size_t t) const {
+const projection::Jacobian& CubedSphereTiles::tileJacobian(size_t t) const {
     return get()->tileJacobian(t);
 }
 

src/atlas/grid/Tiles.h

@@ -15,7 +15,6 @@
 #include <string>
 
 #include "atlas/library/config.h"
-#include "atlas/projection/detail/ProjectionImpl.h"
 #include "atlas/util/ObjectHandle.h"
 
 //---------------------------------------------------------------------------------------------------------------------
@@ -36,6 +35,10 @@ namespace util {
 class Config;
 }  // namespace util
 
+namespace projection {
+class Jacobian;
+}
+
 namespace grid {
 
 #ifndef DOXYGEN_SHOULD_SKIP_THIS
@@ -44,7 +47,6 @@ class CubedSphereTiles;
 }  // namespace detail
 #endif
 
-using Jacobian = projection::detail::ProjectionImpl::Jacobian;
 //---------------------------------------------------------------------------------------------------------------------
 
 class CubedSphereTiles : DOXYGEN_HIDE(public util::ObjectHandle<atlas::grid::detail::CubedSphereTiles>) {
@@ -92,7 +94,7 @@ class CubedSphereTiles : DOXYGEN_HIDE(public util::ObjectHandle<atlas::grid::det
 
     /// @brief Return the Jacobian of xy with respect to the curvilinear
     ///        coordinates of the tile.
-    const Jacobian& tileJacobian(size_t t) const;
+    const projection::Jacobian& tileJacobian(size_t t) const;
 
 private:
     /// Output to stream

src/atlas/grid/detail/partitioner/MatchingMeshPartitionerCubedSphere.cc

@@ -0,0 +1,49 @@
+/*
+ * (C) Crown Copyright 2022 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.
+ */
+
+#include "atlas/grid/Iterator.h"
+#include "atlas/grid/detail/partitioner/MatchingMeshPartitionerCubedSphere.h"
+#include "atlas/interpolation/method/cubedsphere/CellFinder.h"
+#include "atlas/parallel/mpi/mpi.h"
+
+namespace atlas {
+namespace grid {
+namespace detail {
+namespace partitioner {
+
+void MatchingMeshPartitionerCubedSphere::partition(const Grid& grid, int partitioning[]) const {
+
+    // Make cell finder from owned mesh cells.
+    const auto finder = interpolation::method::cubedsphere::CellFinder(prePartitionedMesh_);
+
+    // Loop over grid and set partioning[].
+    auto lonlatIt = grid.lonlat().begin();
+    for (gidx_t i = 0; i < grid.size(); ++i) {
+
+        // This is probably more expensive than it needs to be, as it performs
+        // a dry run of the cubedsphere interpolation method.
+        const auto lonlat = *lonlatIt;
+        partitioning[i] = finder.getCell(lonlat).isect ? mpi::rank() : -1;;
+        ++lonlatIt;
+    }
+
+    // AllReduce to get full partitioning array.
+    mpi::comm().allReduceInPlace(partitioning, grid.size(), eckit::mpi::Operation::MAX);
+    const auto misses = std::count_if(partitioning, partitioning + grid.size(), [](int elem){return elem < 0;});
+    if (misses > 0) {
+        throw_Exception(
+            "Could not find partition for target node (source "
+            "mesh does not contain all target grid points)\n" +
+            std::to_string(misses) + " misses.", Here());
+    }
+
+}
+
+} // namespace partitioner
+} // namespace detail
+} // namespace grid
+} // namespace atlas

src/atlas/grid/detail/partitioner/MatchingMeshPartitionerCubedSphere.h

@@ -0,0 +1,38 @@
+/*
+ * (C) Crown Copyright 2022 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 "atlas/grid/detail/partitioner/MatchingMeshPartitioner.h"
+
+namespace atlas {
+namespace grid {
+namespace detail {
+namespace partitioner {
+
+
+class MatchingMeshPartitionerCubedSphere : public MatchingMeshPartitioner {
+public:
+    static std::string static_type() { return "cubedsphere"; }
+
+public:
+    MatchingMeshPartitionerCubedSphere(): MatchingMeshPartitioner() {}
+    MatchingMeshPartitionerCubedSphere(const idx_t nb_partitions): MatchingMeshPartitioner(nb_partitions) {}
+    MatchingMeshPartitionerCubedSphere(const idx_t nb_partitions, const eckit::Parametrisation&):
+        MatchingMeshPartitioner(nb_partitions) {}
+    MatchingMeshPartitionerCubedSphere(const Mesh& mesh): MatchingMeshPartitioner(mesh) {}
+
+    using MatchingMeshPartitioner::partition;
+    virtual void partition(const Grid& grid, int partitioning[]) const;
+
+    virtual std::string type() const { return static_type(); }
+};
+
+}  // namespace partitioner
+}  // namespace detail
+}  // namespace grid
+}  // namespace atlas

src/atlas/grid/detail/partitioner/Partitioner.cc

@@ -29,6 +29,7 @@
 #include "atlas/grid/detail/partitioner/MatchingFunctionSpacePartitionerLonLatPolygon.h"
 #include "atlas/grid/detail/partitioner/MatchingMeshPartitioner.h"
 #include "atlas/grid/detail/partitioner/MatchingMeshPartitionerBruteForce.h"
+#include "atlas/grid/detail/partitioner/MatchingMeshPartitionerCubedSphere.h"
 #include "atlas/grid/detail/partitioner/MatchingMeshPartitionerLonLatPolygon.h"
 #include "atlas/grid/detail/partitioner/MatchingMeshPartitionerSphericalPolygon.h"
 #include "atlas/grid/detail/partitioner/RegularBandsPartitioner.h"
@@ -202,6 +203,9 @@ Partitioner* MatchingPartitionerFactory::build(const std::string& type, const Me
     else if (type == MatchingMeshPartitionerBruteForce::static_type()) {
         return new MatchingMeshPartitionerBruteForce(partitioned);
     }
+    else if (type == MatchingMeshPartitionerCubedSphere::static_type()) {
+        return new MatchingMeshPartitionerCubedSphere(partitioned);
+    }
     else {
         ATLAS_NOTIMPLEMENTED;
     }

src/atlas/grid/detail/tiles/FV3Tiles.cc

@@ -457,8 +457,9 @@ const PointXY& FV3CubedSphereTiles::tileCentre(size_t t) const {
     throw_NotImplemented("tileCentre not implemented for FV3Tiles", Here());
 }
 
-const Jacobian& FV3CubedSphereTiles::tileJacobian(size_t t) const {
-    throw_NotImplemented("tileJacobian not implemented for FV3Tiles", Here());;
+const FV3CubedSphereTiles::Jacobian& FV3CubedSphereTiles::tileJacobian(size_t t) const {
+    throw_NotImplemented("tileJacobian not implemented for FV3Tiles", Here());
+    ;
 }