ecmwf · tom-j-h · Mar 13, 2025
diff --git a/src/atlas/functionspace/CellColumns.cc b/src/atlas/functionspace/CellColumns.cc
@@ -593,6 +593,10 @@ Field CellColumns::partition() const {
     return mesh_.cells().partition();
 }
 
+Grid CellColumns::get_grid_copy() const {
+    return mesh_.grid();
+}
+
 //------------------------------------------------------------------------------
 //------------------------------------------------------------------------------
 //------------------------------------------------------------------------------

diff --git a/src/atlas/functionspace/CellColumns.h b/src/atlas/functionspace/CellColumns.h
@@ -99,6 +99,8 @@ class CellColumns : public functionspace::FunctionSpaceImpl {
 
     Field partition() const override;
 
+    Grid get_grid_copy() const override;
+
     std::string mpi_comm() const override { return mesh_.mpi_comm(); }
 
 private:  // methods

diff --git a/src/atlas/functionspace/EdgeColumns.cc b/src/atlas/functionspace/EdgeColumns.cc
@@ -576,6 +576,10 @@ Field EdgeColumns::partition() const {
     return edges().partition();
 }
 
+Grid EdgeColumns::get_grid_copy() const {
+    return mesh().grid();
+}
+
 //------------------------------------------------------------------------------
 //------------------------------------------------------------------------------
 //------------------------------------------------------------------------------

diff --git a/src/atlas/functionspace/EdgeColumns.h b/src/atlas/functionspace/EdgeColumns.h
@@ -93,6 +93,8 @@ class EdgeColumns : public functionspace::FunctionSpaceImpl {
 
     Field partition() const override;
 
+    Grid get_grid_copy() const override;
+
     std::string mpi_comm() const override { return mesh_.mpi_comm(); }
 
 private:  // methods

diff --git a/src/atlas/functionspace/FunctionSpace.cc b/src/atlas/functionspace/FunctionSpace.cc
@@ -11,6 +11,7 @@
 
 #include "atlas/functionspace/FunctionSpace.h"
 #include "atlas/field/Field.h"
+#include "atlas/grid.h"
 #include "atlas/functionspace/detail/FunctionSpaceImpl.h"
 
 namespace atlas {
@@ -86,6 +87,10 @@ Field FunctionSpace::partition() const {
     return get()->partition();
 }
 
+Grid FunctionSpace::get_grid_copy() const {
+    return get()->get_grid_copy();
+}
+
 void FunctionSpace::haloExchange(const FieldSet& fields, bool on_device) const {
     get()->haloExchange(fields, on_device);
 }

diff --git a/src/atlas/functionspace/FunctionSpace.h b/src/atlas/functionspace/FunctionSpace.h
@@ -22,6 +22,7 @@ class Configuration;
 namespace atlas {
 class Field;
 class FieldSet;
+class Grid;
 class Projection;
 namespace functionspace {
 class FunctionSpaceImpl;
@@ -95,6 +96,8 @@ class FunctionSpace : DOXYGEN_HIDE(public util::ObjectHandle<functionspace::Func
 
     Field partition() const;
 
+    Grid get_grid_copy() const;
+
     const parallel::GatherScatter& gather() const;
     const parallel::GatherScatter& scatter() const;
 

diff --git a/src/atlas/functionspace/NodeColumns.h b/src/atlas/functionspace/NodeColumns.h
@@ -272,6 +272,8 @@ class NodeColumns : public functionspace::FunctionSpaceImpl {
 
     Field partition() const override { return nodes_.partition(); }
 
+    Grid get_grid_copy() const override { return mesh_.grid(); }
+
     const util::PartitionPolygon& polygon(idx_t halo = 0) const override { return mesh_.polygon(halo); }
 
     const util::PartitionPolygons& polygons() const override { return mesh_.polygons(); }

diff --git a/src/atlas/functionspace/PointCloud.cc b/src/atlas/functionspace/PointCloud.cc
@@ -20,6 +20,7 @@
 #include "atlas/grid/Distribution.h"
 #include "atlas/grid/Grid.h"
 #include "atlas/grid/Iterator.h"
+#include "atlas/grid/UnstructuredGrid.h"
 #include "atlas/option/Options.h"
 #include "atlas/parallel/HaloExchange.h"
 #include "atlas/parallel/mpi/mpi.h"
@@ -156,6 +157,7 @@ PointCloud::PointCloud(const Grid& grid, const grid::Distribution& distribution,
     auto& comm = mpi::comm(mpi_comm_);
     double halo_radius;
     config.get("halo_radius", halo_radius = 0.);
+    grid_ = grid;
 
     part_ = comm.rank();
 
@@ -278,6 +280,40 @@ Field PointCloud::ghost() const {
     return ghost_;
 }
 
+Grid PointCloud::get_grid_copy() const {
+    if (grid_) return grid_;
+
+    if (nb_partitions_ == 1) {
+        std::vector<PointXY> points;
+        points.reserve(size_global_);
+        for (const auto& point : iterate().xy()) {
+            points.push_back(point);
+        }
+        return UnstructuredGrid(points);
+    } else {
+        std::vector<double> points_x(size_global_, 0.0);
+        std::vector<double> points_y(size_global_, 0.0);
+        const auto gidx = array::make_view<gidx_t, 1>(global_index_);
+        const auto ghost = array::make_view<int, 1>(ghost_);
+        int i = 0;
+        for (const auto& point : iterate().xy()) {
+            if (ghost(i) != 0) continue;
+            points_x[gidx(i) - 1] = point.x();
+            points_y[gidx(i) - 1] = point.y();
+            i++;
+        }
+        const auto& comm = mpi::comm(mpi_comm_);
+        comm.allReduceInPlace(points_x.begin(), points_x.end(), eckit::mpi::sum());
+        comm.allReduceInPlace(points_y.begin(), points_y.end(), eckit::mpi::sum());
+        std::vector<PointXY> points;
+        points.reserve(size_global_);
+        for (int i = 0; i < points_x.size(); i++) {
+            points.emplace_back(PointXY{points_x[i], points_y[i]});
+        }
+        return UnstructuredGrid(points);
+    }
+}
+
 array::ArrayShape PointCloud::config_shape(const eckit::Configuration& config) const {
     idx_t _size  = size();
     bool global(false);

diff --git a/src/atlas/functionspace/PointCloud.h b/src/atlas/functionspace/PointCloud.h
@@ -19,6 +19,7 @@
 #include "atlas/field/FieldSet.h"
 #include "atlas/functionspace/FunctionSpace.h"
 #include "atlas/functionspace/detail/FunctionSpaceImpl.h"
+#include "atlas/grid/Grid.h"
 #include "atlas/parallel/HaloExchange.h"
 #include "atlas/parallel/GatherScatter.h"
 #include "atlas/runtime/Exception.h"
@@ -63,6 +64,7 @@ class PointCloud : public functionspace::FunctionSpaceImpl {
     Field remote_index() const override { return remote_index_; }
     Field global_index() const override { return global_index_; }
     Field partition() const override { return partition_; }
+    Grid get_grid_copy() const override;
     idx_t size() const override { return lonlat_.shape(0); }
     idx_t part() const override { return part_; }
     idx_t nb_parts() const override { return nb_partitions_; }
@@ -170,6 +172,7 @@ class PointCloud : public functionspace::FunctionSpaceImpl {
     void create_remote_index() const;
 
 private:
+    Grid grid_;
     Field lonlat_;
     Field vertical_;
     mutable Field ghost_;

diff --git a/src/atlas/functionspace/detail/BlockStructuredColumns.h b/src/atlas/functionspace/detail/BlockStructuredColumns.h
@@ -105,6 +105,7 @@ class BlockStructuredColumns : public FunctionSpaceImpl {
     Field xy() const { return structuredcolumns_->xy(); }
     Field z() const { return structuredcolumns_->z(); }
     Field partition() const override { return structuredcolumns_->partition(); }
+    Grid get_grid_copy() const override { return structuredcolumns_->get_grid_copy(); }
     Field global_index() const override { return structuredcolumns_->global_index(); }
     Field remote_index() const override { return structuredcolumns_->remote_index(); }
     Field index_i() const { return structuredcolumns_->index_i(); }

diff --git a/src/atlas/functionspace/detail/FunctionSpaceImpl.cc b/src/atlas/functionspace/detail/FunctionSpaceImpl.cc
@@ -10,6 +10,7 @@
 
 #include "FunctionSpaceImpl.h"
 #include "atlas/field/Field.h"
+#include "atlas/grid.h"
 #include "atlas/option/Options.h"
 #include "atlas/runtime/Exception.h"
 #include "atlas/util/Metadata.h"
@@ -70,6 +71,10 @@ Field FunctionSpaceImpl::partition() const {
     ATLAS_NOTIMPLEMENTED;
 }
 
+Grid FunctionSpaceImpl::get_grid_copy() const {
+    ATLAS_NOTIMPLEMENTED;
+}
+
 Field FunctionSpaceImpl::global_index() const {
     ATLAS_NOTIMPLEMENTED;
 }

diff --git a/src/atlas/functionspace/detail/FunctionSpaceImpl.h b/src/atlas/functionspace/detail/FunctionSpaceImpl.h
@@ -25,6 +25,7 @@ class Configuration;
 namespace atlas {
 class FieldSet;
 class Field;
+class Grid;
 class Projection;
 namespace util {
 class Metadata;
@@ -107,6 +108,8 @@ class FunctionSpaceImpl : public util::Object {
 
     virtual atlas::Field partition() const;
 
+    virtual atlas::Grid get_grid_copy() const;
+
     virtual atlas::Field global_index() const;
 
     virtual const util::PartitionPolygons& polygons() const;

diff --git a/src/atlas/functionspace/detail/StructuredColumns.h b/src/atlas/functionspace/detail/StructuredColumns.h
@@ -143,6 +143,7 @@ class StructuredColumns : public FunctionSpaceImpl {
     Field xy() const { return field_xy_; }
     Field z() const { return field_z_; }
     Field partition() const override { return field_partition_; }
+    Grid get_grid_copy() const override { return *grid_; }
     Field global_index() const override { return field_global_index_; }
     Field remote_index() const override {
         if (not field_remote_index_) {

diff --git a/src/tests/functionspace/test_pointcloud.cc b/src/tests/functionspace/test_pointcloud.cc
@@ -13,6 +13,7 @@
 #include "atlas/array.h"
 #include "atlas/field.h"
 #include "atlas/functionspace/PointCloud.h"
+#include "atlas/grid/UnstructuredGrid.h"
 #include "atlas/option.h"
 #include "atlas/parallel/mpi/mpi.h"
 
@@ -161,6 +162,24 @@ CASE("test_createField") {
 
 //-----------------------------------------------------------------------------
 
+CASE("test_get_grid_copy") {
+    UnstructuredGrid grid(ref_xy());
+
+    FunctionSpace p1;
+    p1 = functionspace::PointCloud(grid);  // stores copy of Grid
+
+    FunctionSpace p2;
+    p2 = functionspace::PointCloud(ref_xy());  // no Grid
+
+    Grid gridP1 = p1.get_grid_copy();  // returns copy
+    Grid gridP2 = p2.get_grid_copy();  // creates via iterator
+
+    EXPECT(gridP1 == grid);
+    EXPECT(gridP1 == gridP2);
+}
+
+//-----------------------------------------------------------------------------
+
 }  // namespace test
 }  // namespace atlas