Implement PointCloud::nb_partitions()

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

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/src/atlas/functionspace/PointCloud.cc b/src/atlas/functionspace/PointCloud.cc
@@ -118,6 +118,7 @@ PointCloud::PointCloud(const Grid& grid, const grid::Partitioner& _partitioner,
         partitioner = grid::Partitioner("equal_regions");
     }
 
+    nb_partitions_ = partitioner.nb_partitions();
     auto distribution = partitioner.partition(grid);
     auto size_owned = distribution.nb_pts()[part];
     size_owned_ = size_owned;

diff --git a/src/atlas/functionspace/PointCloud.h b/src/atlas/functionspace/PointCloud.h
@@ -60,7 +60,8 @@ 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_; }
-    virtual idx_t size() const override { return lonlat_.shape(0); }
+    idx_t size() const override { return lonlat_.shape(0); }
+    idx_t nb_partitions() const override { return nb_partitions_; }
 
     using FunctionSpaceImpl::createField;
     Field createField(const eckit::Configuration&) const override;
@@ -165,6 +166,7 @@ class PointCloud : public functionspace::FunctionSpaceImpl {
     idx_t max_glb_idx_{0};
     std::unique_ptr<parallel::HaloExchange> halo_exchange_;
     idx_t levels_{0};
+    idx_t nb_partitions_{1};
 
     void setupHaloExchange();