Add FunctionSpace::partition() field

ecmwf · Sep 19, 2023 · b87aa4a · b87aa4a
1 parent 4f9b29c
commit b87aa4a
Show file tree

Hide file tree

Showing 12 changed files with 29 additions and 2 deletions.
diff --git a/src/atlas/functionspace/CellColumns.cc b/src/atlas/functionspace/CellColumns.cc
@@ -587,6 +587,10 @@ Field CellColumns::ghost() const {
     return mesh_.cells().field("ghost");
 }
 
+Field CellColumns::partition() const {
+    return mesh_.cells().partition();
+}
+
 //------------------------------------------------------------------------------
 //------------------------------------------------------------------------------
 //------------------------------------------------------------------------------

diff --git a/src/atlas/functionspace/CellColumns.h b/src/atlas/functionspace/CellColumns.h
@@ -97,6 +97,8 @@ class CellColumns : public functionspace::FunctionSpaceImpl {
 
     Field ghost() const override;
 
+    Field partition() const override;
+
 private:  // methods
     idx_t config_size(const eckit::Configuration& config) const;
     array::DataType config_datatype(const eckit::Configuration&) const;

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

diff --git a/src/atlas/functionspace/EdgeColumns.h b/src/atlas/functionspace/EdgeColumns.h
@@ -91,6 +91,8 @@ class EdgeColumns : public functionspace::FunctionSpaceImpl {
 
     Field remote_index() const override;
 
+    Field partition() const override;
+
 private:  // methods
     idx_t config_size(const eckit::Configuration& config) const;
     array::DataType config_datatype(const eckit::Configuration&) const;

diff --git a/src/atlas/functionspace/FunctionSpace.cc b/src/atlas/functionspace/FunctionSpace.cc
@@ -78,6 +78,10 @@ Field FunctionSpace::remote_index() const {
     return get()->remote_index();
 }
 
+Field FunctionSpace::partition() const {
+    return get()->partition();
+}
+
 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
@@ -91,6 +91,8 @@ class FunctionSpace : DOXYGEN_HIDE(public util::ObjectHandle<functionspace::Func
 
     Field remote_index() const;
 
+    Field partition() 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
@@ -268,6 +268,8 @@ class NodeColumns : public functionspace::FunctionSpaceImpl {
 
     Field remote_index() const override { return nodes_.remote_index(); }
 
+    Field partition() const override { return nodes_.partition(); }
+
     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.h b/src/atlas/functionspace/PointCloud.h
@@ -59,6 +59,7 @@ class PointCloud : public functionspace::FunctionSpaceImpl {
     Field ghost() const override;
     Field remote_index() const override { return remote_index_; }
     Field global_index() const override { return global_index_; }
+    Field partition() const override { return partition_; }
     virtual idx_t size() const override { return lonlat_.shape(0); }
 
     using FunctionSpaceImpl::createField;

diff --git a/src/atlas/functionspace/detail/BlockStructuredColumns.h b/src/atlas/functionspace/detail/BlockStructuredColumns.h
@@ -104,7 +104,7 @@ class BlockStructuredColumns : public FunctionSpaceImpl {
     Field lonlat() const override { return structuredcolumns_->lonlat(); }
     Field xy() const { return structuredcolumns_->xy(); }
     Field z() const { return structuredcolumns_->z(); }
-    Field partition() const { return structuredcolumns_->partition(); }
+    Field partition() const override { return structuredcolumns_->partition(); }
     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
@@ -65,6 +65,10 @@ Field FunctionSpaceImpl::remote_index() const {
     ATLAS_NOTIMPLEMENTED;
 }
 
+Field FunctionSpaceImpl::partition() 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
@@ -103,6 +103,8 @@ class FunctionSpaceImpl : public util::Object {
 
     virtual atlas::Field remote_index() const;
 
+    virtual atlas::Field partition() 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
@@ -142,7 +142,7 @@ class StructuredColumns : public FunctionSpaceImpl {
     Field lonlat() const override { return field_xy_; }
     Field xy() const { return field_xy_; }
     Field z() const { return field_z_; }
-    Field partition() const { return field_partition_; }
+    Field partition() const override { return field_partition_; }
     Field global_index() const override { return field_global_index_; }
     Field remote_index() const override {
         if (not field_remote_index_) {