From 16334831369ec13ae684bb022df4f7dc82cb2ef5 Mon Sep 17 00:00:00 2001
From: Willem Deconinck <willem.deconinck@ecmwf.int>
Date: Tue, 19 Sep 2023 15:06:27 +0200
Subject: [PATCH] Implement PointCloud::nb_partitions()

---
 src/atlas/functionspace/PointCloud.cc | 1 +
 src/atlas/functionspace/PointCloud.h  | 4 +++-
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/atlas/functionspace/PointCloud.cc b/src/atlas/functionspace/PointCloud.cc
index 25594c800..797b9e43c 100644
--- 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
index 2f763ee1b..c28a7e774 100644
--- 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();