Add Grid getter to FunctionSpace

src/atlas/functionspace/CellColumns.cc

@@ -16,6 +16,7 @@
 
 #include "atlas/array/MakeView.h"
 #include "atlas/functionspace/CellColumns.h"
+#include "atlas/grid/UnstructuredGrid.h"
 #include "atlas/library/config.h"
 #include "atlas/mesh/HybridElements.h"
 #include "atlas/mesh/IsGhostNode.h"
@@ -283,6 +284,8 @@ CellColumns::CellColumns(const Mesh& mesh, const eckit::Configuration& config):
         nb_cells_ = mesh.cells().size();
     }
     ATLAS_ASSERT(nb_cells_);
+
+    if (mesh_.grid()) grid_ = mesh_.grid();
 }
 
 CellColumns::~CellColumns() = default;
@@ -570,6 +573,44 @@ const parallel::Checksum& CellColumns::checksum() const {
     return *checksum_;
 }
 
+const Grid& CellColumns::grid() const {
+    if (grid_) return grid_;
+
+    const auto& comm = mpi::comm(mpi_comm());
+    std::vector<PointXY> points;
+    if (comm.size() == 1) {
+        const auto xy = atlas::array::make_view<double, 2>(mesh_.nodes().xy());
+        for (auto i = 0; i < xy.shape(0); i++) {
+            points.push_back({xy(i, 0), xy(i, 1)});
+        }
+    } else {
+        std::vector<int> gidx;
+        std::vector<double> x, y;
+        const auto gidxView = array::make_view<gidx_t, 1>(global_index());
+        const auto ghostView = array::make_view<int, 1>(ghost());
+        const auto xy = atlas::array::make_view<double, 2>(mesh_.nodes().xy());
+        for (auto i = 0; i < xy.shape(0); i++) {
+            if (ghostView(i) == 0) {
+                gidx.push_back(gidxView(i));
+                x.push_back(xy(i, 0));
+                y.push_back(xy(i, 1));
+            }
+        }
+        eckit::mpi::Buffer<int> gidxBuffer(comm.size());
+        eckit::mpi::Buffer<double> xBuffer(comm.size());
+        eckit::mpi::Buffer<double> yBuffer(comm.size());
+        comm.allGatherv(gidx.begin(), gidx.end(), gidxBuffer);
+        comm.allGatherv(x.begin(), x.end(), xBuffer);
+        comm.allGatherv(y.begin(), y.end(), yBuffer);
+        points.reserve(gidxBuffer.buffer.size());
+        for (auto i : gidxBuffer.buffer) {
+            points[i - 1] = atlas::PointXY{xBuffer.buffer[i - 1], yBuffer.buffer[i - 1]};
+        }
+    }
+    grid_ = UnstructuredGrid(points);
+    return grid_;
+}
+
 Field CellColumns::lonlat() const {
     if (!mesh_.cells().has_field("lonlat")) {
         mesh::actions::Build2DCellCentres("lonlat")(const_cast<Mesh&>(mesh_));

src/atlas/functionspace/CellColumns.h

@@ -89,6 +89,8 @@ class CellColumns : public functionspace::FunctionSpaceImpl {
 
     idx_t size() const override { return nb_cells_; }
 
+    const Grid& grid() const override;
+
     Field lonlat() const override;
 
     Field remote_index() const override;
@@ -111,6 +113,7 @@ class CellColumns : public functionspace::FunctionSpaceImpl {
     size_t footprint() const override;
 
 private:                           // data
+    mutable Grid grid_;
     Mesh mesh_;                    // non-const because functionspace may modify mesh
     mesh::HybridElements& cells_;  // non-const because functionspace may modify mesh
     idx_t nb_levels_;

src/atlas/functionspace/EdgeColumns.cc

@@ -17,6 +17,7 @@
 #include "atlas/array/MakeView.h"
 #include "atlas/field/detail/FieldImpl.h"
 #include "atlas/functionspace/EdgeColumns.h"
+#include "atlas/grid/UnstructuredGrid.h"
 #include "atlas/library/config.h"
 #include "atlas/mesh/HybridElements.h"
 #include "atlas/mesh/IsGhostNode.h"
@@ -281,6 +282,8 @@ EdgeColumns::EdgeColumns(const Mesh& mesh, const eckit::Configuration& config):
         nb_edges_ = get_nb_edges_from_metadata();
     }
     ATLAS_ASSERT(nb_edges_);
+
+    if (mesh_.grid()) grid_ = mesh_.grid();
 }
 
 EdgeColumns::~EdgeColumns() = default;
@@ -560,6 +563,44 @@ const parallel::Checksum& EdgeColumns::checksum() const {
     return *checksum_;
 }
 
+const Grid& EdgeColumns::grid() const {
+    if (grid_) return grid_;
+
+    const auto& comm = mpi::comm(mpi_comm());
+    std::vector<PointXY> points;
+    if (comm.size() == 1) {
+        const auto xy = atlas::array::make_view<double, 2>(mesh_.nodes().xy());
+        for (auto i = 0; i < xy.shape(0); i++) {
+            points.push_back({xy(i, 0), xy(i, 1)});
+        }
+    } else {
+        std::vector<int> gidx;
+        std::vector<double> x, y;
+        const auto gidxView = array::make_view<gidx_t, 1>(global_index());
+        const auto ghostView = array::make_view<int, 1>(ghost());
+        const auto xy = atlas::array::make_view<double, 2>(mesh_.nodes().xy());
+        for (auto i = 0; i < xy.shape(0); i++) {
+            if (ghostView(i) == 0) {
+                gidx.push_back(gidxView(i));
+                x.push_back(xy(i, 0));
+                y.push_back(xy(i, 1));
+            }
+        }
+        eckit::mpi::Buffer<int> gidxBuffer(comm.size());
+        eckit::mpi::Buffer<double> xBuffer(comm.size());
+        eckit::mpi::Buffer<double> yBuffer(comm.size());
+        comm.allGatherv(gidx.begin(), gidx.end(), gidxBuffer);
+        comm.allGatherv(x.begin(), x.end(), xBuffer);
+        comm.allGatherv(y.begin(), y.end(), yBuffer);
+        points.reserve(gidxBuffer.buffer.size());
+        for (auto i : gidxBuffer.buffer) {
+            points[i - 1] = atlas::PointXY{xBuffer.buffer[i - 1], yBuffer.buffer[i - 1]};
+        }
+    }
+    grid_ = UnstructuredGrid(points);
+    return grid_;
+}
+
 Field EdgeColumns::lonlat() const {
     return edges().field("lonlat");
 }

src/atlas/functionspace/EdgeColumns.h

@@ -85,6 +85,8 @@ class EdgeColumns : public functionspace::FunctionSpaceImpl {
 
     idx_t size() const override { return nb_edges_; }
 
+    const Grid& grid() const override;
+
     Field lonlat() const override;
 
     Field global_index() const override;
@@ -105,6 +107,7 @@ class EdgeColumns : public functionspace::FunctionSpaceImpl {
     size_t footprint() const override;
 
 private:                           // data
+    mutable Grid grid_;
     Mesh mesh_;                    // non-const because functionspace may modify mesh
     mesh::HybridElements& edges_;  // non-const because functionspace may modify mesh
     idx_t nb_levels_;

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 {
@@ -30,6 +31,10 @@ size_t FunctionSpace::footprint() const {
     return get()->footprint();
 }
 
+const Grid& FunctionSpace::grid() const {
+    return get()->grid();
+}
+
 Field FunctionSpace::createField(const eckit::Configuration& config) const {
     return get()->createField(config);
 }

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;
@@ -85,6 +86,8 @@ class FunctionSpace : DOXYGEN_HIDE(public util::ObjectHandle<functionspace::Func
 
     idx_t size() const;
 
+    const Grid& grid() const;
+
     Field lonlat() const;
 
     Field ghost() const;

src/atlas/functionspace/NodeColumns.cc

@@ -18,6 +18,7 @@
 #include "atlas/field/FieldSet.h"
 #include "atlas/functionspace/NodeColumns.h"
 #include "atlas/grid/Grid.h"
+#include "atlas/grid/UnstructuredGrid.h"
 #include "atlas/library/config.h"
 #include "atlas/mesh/IsGhostNode.h"
 #include "atlas/mesh/Mesh.h"
@@ -223,6 +224,8 @@ NodeColumns::NodeColumns(Mesh mesh, const eckit::Configuration& config):
             nb_nodes_ = mesh_.nodes().size();
         }
     }
+
+    if (mesh_.grid()) grid_ = mesh_.grid();
 }
 
 NodeColumns::~NodeColumns() = default;
@@ -641,6 +644,43 @@ const parallel::Checksum& NodeColumns::checksum() const {
 //  return checksum(fieldset);
 //}
 
+const Grid& NodeColumns::grid() const {
+    if (grid_) return grid_;
+
+    const auto& comm = mpi::comm(mpi_comm());
+    std::vector<PointXY> points;
+    if (comm.size() == 1) {
+        const auto xy = atlas::array::make_view<double, 2>(mesh_.nodes().xy());
+        for (auto i = 0; i < xy.shape(0); i++) {
+            points.push_back({xy(i, 0), xy(i, 1)});
+        }
+    } else {
+        std::vector<int> gidx;
+        std::vector<double> x, y;
+        const auto gidxView = array::make_view<gidx_t, 1>(global_index());
+        const auto ghostView = array::make_view<int, 1>(ghost());
+        const auto xy = atlas::array::make_view<double, 2>(mesh_.nodes().xy());
+        for (auto i = 0; i < xy.shape(0); i++) {
+            if (ghostView(i) == 0) {
+                gidx.push_back(gidxView(i));
+                x.push_back(xy(i, 0));
+                y.push_back(xy(i, 1));
+            }
+        }
+        eckit::mpi::Buffer<int> gidxBuffer(comm.size());
+        eckit::mpi::Buffer<double> xBuffer(comm.size());
+        eckit::mpi::Buffer<double> yBuffer(comm.size());
+        comm.allGatherv(gidx.begin(), gidx.end(), gidxBuffer);
+        comm.allGatherv(x.begin(), x.end(), xBuffer);
+        comm.allGatherv(y.begin(), y.end(), yBuffer);
+        points.reserve(gidxBuffer.buffer.size());
+        for (auto i : gidxBuffer.buffer) {
+            points[i - 1] = atlas::PointXY{xBuffer.buffer[i - 1], yBuffer.buffer[i - 1]};
+        }
+    }
+    grid_ = UnstructuredGrid(points);
+    return grid_;
+}
 
 }  // namespace detail
 

src/atlas/functionspace/NodeColumns.h

@@ -262,6 +262,8 @@ class NodeColumns : public functionspace::FunctionSpaceImpl {
 
     idx_t nb_parts() const override { return mesh_.nb_parts(); }
 
+    const Grid& grid() const override;
+
     Field lonlat() const override { return nodes_.lonlat(); }
 
     Field ghost() const override { return nodes_.ghost(); }
@@ -293,6 +295,7 @@ class NodeColumns : public functionspace::FunctionSpaceImpl {
     virtual size_t footprint() const override { return 0; }
 
 private:                  // data
+    mutable Grid grid_;
     Mesh mesh_;           // non-const because functionspace may modify mesh
     mesh::Nodes& nodes_;  // non-const because functionspace may modify mesh
     mesh::Halo halo_;

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,48 @@ Field PointCloud::ghost() const {
     return ghost_;
 }
 
+const Grid& PointCloud::grid() const {
+    if (grid_) return grid_;
+
+    std::vector<PointXY> points;
+    points.reserve(size_global_);
+    if (nb_partitions_ == 1) {
+        for (const auto& point : iterate().xy()) {
+            points.push_back(point);
+        }
+    } else {
+        std::vector<int> gidx;
+        gidx.reserve(size_global_);
+        std::vector<double> x, y;
+        x.reserve(size_global_);
+        y.reserve(size_global_);
+        const auto gidxView = array::make_view<gidx_t, 1>(global_index_);
+        const auto ghostView = array::make_view<int, 1>(ghost_);
+        int i = 0;
+        for (const auto& point : iterate().xy()) {
+            if (ghostView(i) == 0) {
+                gidx.push_back(gidxView(i));
+                x.push_back(point.x());
+                y.push_back(point.y());
+            }
+            i++;
+        }
+        const auto& comm = mpi::comm(mpi_comm_);
+        eckit::mpi::Buffer<int> gidxBuffer(comm.size());
+        eckit::mpi::Buffer<double> xBuffer(comm.size());
+        eckit::mpi::Buffer<double> yBuffer(comm.size());
+        comm.allGatherv(gidx.begin(), gidx.end(), gidxBuffer);
+        comm.allGatherv(x.begin(), x.end(), xBuffer);
+        comm.allGatherv(y.begin(), y.end(), yBuffer);
+        for (auto i : gidxBuffer.buffer) {
+            points[i - 1] = atlas::PointXY{xBuffer.buffer[i - 1], yBuffer.buffer[i - 1]};
+        }
+    }
+    grid_ = UnstructuredGrid(points);
+    return grid_;
+}
+
+
 array::ArrayShape PointCloud::config_shape(const eckit::Configuration& config) const {
     idx_t _size  = size();
     bool global(false);

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"
@@ -57,6 +58,7 @@ class PointCloud : public functionspace::FunctionSpaceImpl {
     operator bool() const override { return true; }
     size_t footprint() const override { return sizeof(*this); }
     std::string distribution() const override;
+    const Grid& grid() const override;
     Field lonlat() const override { return lonlat_; }
     const Field& vertical() const { return vertical_; }
     Field ghost() const override;
@@ -170,6 +172,7 @@ class PointCloud : public functionspace::FunctionSpaceImpl {
     void create_remote_index() const;
 
 private:
+    mutable Grid grid_;
     Field lonlat_;
     Field vertical_;
     mutable Field ghost_;

src/atlas/functionspace/detail/BlockStructuredColumns.h

@@ -98,7 +98,7 @@ class BlockStructuredColumns : public FunctionSpaceImpl {
     idx_t nblks() const { return nblks_; }
 
     const Vertical& vertical() const { return structuredcolumns_->vertical(); }
-    const StructuredGrid& grid() const { return structuredcolumns_->grid(); }
+    const StructuredGrid& grid() const override { return structuredcolumns_->grid(); }
 
     idx_t levels() const { return structuredcolumns_->levels(); }
     Field lonlat() const override { return structuredcolumns_->lonlat(); }

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"
@@ -54,6 +55,10 @@ Field NoFunctionSpace::createField(const Field&, const eckit::Configuration&) co
     ATLAS_NOTIMPLEMENTED;
 }
 
+const Grid& FunctionSpaceImpl::grid() const {
+    ATLAS_NOTIMPLEMENTED;
+}
+
 Field FunctionSpaceImpl::lonlat() const {
     ATLAS_NOTIMPLEMENTED;
 }