From 35b18198f256b1df444b62b6d1e9ec3df954a49b Mon Sep 17 00:00:00 2001
From: Martin Kronbichler <kronbichler@lnm.mw.tum.de>
Date: Sun, 4 Jul 2021 22:48:05 +0200
Subject: [PATCH] Add ParticleAccessor::get_local_index

---
 include/deal.II/particles/particle_accessor.h | 27 +++++++++++++++++++
 include/deal.II/particles/particle_handler.h  | 11 ++++++++
 include/deal.II/particles/property_pool.h     | 18 ++++++++++++-
 source/particles/particle_handler.cc          | 16 +++++++++--
 4 files changed, 69 insertions(+), 3 deletions(-)

diff --git a/include/deal.II/particles/particle_accessor.h b/include/deal.II/particles/particle_accessor.h
index b90fe459e4..70103e3754 100644
--- a/include/deal.II/particles/particle_accessor.h
+++ b/include/deal.II/particles/particle_accessor.h
@@ -130,6 +130,22 @@ namespace Particles
     types::particle_index
     get_id() const;
 
+    /**
+     * Return a particle ID number local to each MPI process. This number
+     * enables the direct array access (similar to
+     * LinearAlgebra::distributed::Vector::local_element()) to quantities used
+     * for local computations. Use
+     * ParticleHandler::get_max_local_particle_index() to query the largest
+     * index for suitable array sizes.
+     *
+     * @note The number returned by this function is not stable between calls
+     * of ParticleHandler::sort_particles_into_subdomains_and_cells() and
+     * therefore it cannot be used to track quantities across those calls. Use
+     * particle properties for storing persistent information.
+     */
+    types::particle_index
+    get_local_index() const;
+
     /**
      * Return whether this particle has a valid property pool and a valid
      * handle to properties.
@@ -579,6 +595,17 @@ namespace Particles
 
 
 
+  template <int dim, int spacedim>
+  inline types::particle_index
+  ParticleAccessor<dim, spacedim>::get_local_index() const
+  {
+    Assert(state() == IteratorState::valid, ExcInternalError());
+
+    return get_handle();
+  }
+
+
+
   template <int dim, int spacedim>
   inline void
   ParticleAccessor<dim, spacedim>::set_id(const types::particle_index &new_id)
diff --git a/include/deal.II/particles/particle_handler.h b/include/deal.II/particles/particle_handler.h
index d3ed4cab45..2fe3ec54a3 100644
--- a/include/deal.II/particles/particle_handler.h
+++ b/include/deal.II/particles/particle_handler.h
@@ -660,6 +660,17 @@ namespace Particles
     unsigned int
     n_properties_per_particle() const;
 
+    /**
+     * Return the maximal local index (in MPI-local index space) returned by
+     * ParticleAccessor::get_local_index(). More precisely, the returned value
+     * is one larger than the largest index that will be returned, which makes
+     * this value useful for resizing vectors. This number can be larger than
+     * locally_owned_particle_ids().n_elements(), because local indices can
+     * contain gaps in the range 0 to get_max_local_particle_index().
+     */
+    types::particle_index
+    get_max_local_particle_index() const;
+
     /**
      * Return a reference to the property pool that owns all particle
      * properties, and organizes them physically.
diff --git a/include/deal.II/particles/property_pool.h b/include/deal.II/particles/property_pool.h
index 944d2ca643..576effac94 100644
--- a/include/deal.II/particles/property_pool.h
+++ b/include/deal.II/particles/property_pool.h
@@ -137,7 +137,7 @@ namespace Particles
 
     /**
      * Return a new handle that allows a particle to store information such as
-     * properties and locations. This also allocated memory in this PropertyPool
+     * properties and locations. This also allocates memory in this PropertyPool
      * variable.
      */
     Handle
@@ -211,6 +211,13 @@ namespace Particles
     unsigned int
     n_properties_per_slot() const;
 
+    /**
+     * Return the total number of slots in the pool, including both registered
+     * and unregistered ones.
+     */
+    unsigned int
+    n_slots() const;
+
     /**
      * Return how many slots are currently registered in the pool.
      */
@@ -440,6 +447,15 @@ namespace Particles
   }
 
 
+
+  template <int dim, int spacedim>
+  inline unsigned int
+  PropertyPool<dim, spacedim>::n_slots() const
+  {
+    return locations.size();
+  }
+
+
 } // namespace Particles
 
 DEAL_II_NAMESPACE_CLOSE
diff --git a/source/particles/particle_handler.cc b/source/particles/particle_handler.cc
index 5426e9e6bd..2f12ce3e91 100644
--- a/source/particles/particle_handler.cc
+++ b/source/particles/particle_handler.cc
@@ -919,15 +919,27 @@ namespace Particles
   IndexSet
   ParticleHandler<dim, spacedim>::locally_owned_particle_ids() const
   {
-    IndexSet set(get_next_free_particle_index());
+    IndexSet                           set(get_next_free_particle_index());
+    std::vector<types::particle_index> indices;
+    indices.reserve(n_locally_owned_particles());
     for (const auto &p : *this)
-      set.add_index(p.get_id());
+      indices.push_back(p.get_id());
+    set.add_indices(indices.begin(), indices.end());
     set.compress();
     return set;
   }
 
 
 
+  template <int dim, int spacedim>
+  types::particle_index
+  ParticleHandler<dim, spacedim>::get_max_local_particle_index() const
+  {
+    return property_pool->n_slots();
+  }
+
+
+
   template <int dim, int spacedim>
   void
   ParticleHandler<dim, spacedim>::get_particle_positions(
-- 
2.39.5