Add a _mesh_local_subdomains cache

Author: roystgnr

include/mesh/mesh_base.h

@@ -1329,7 +1329,7 @@ class MeshBase : public ParallelObject
    * Recalculate any cached data after elements and nodes have been
    * repartitioned.
    */
-  virtual void update_post_partitioning () {}
+  virtual void update_post_partitioning ();
 
   /**
    * If false is passed in then this mesh will no longer be renumbered
@@ -1993,6 +1993,15 @@ class MeshBase : public ParallelObject
   const std::set<subdomain_id_type> & get_mesh_subdomains() const
   { libmesh_assert(this->is_prepared()); return _mesh_subdomains; }
 
+
+  /**
+   * \return The cached mesh subdomains. As long as the mesh is prepared, this
+   * should contain all the subdomain ids across processors. Relies on the mesh
+   * being prepared
+   */
+  const std::set<subdomain_id_type> & get_mesh_local_subdomains() const
+  { libmesh_assert(this->is_prepared()); return _mesh_local_subdomains; }
+
 #ifdef LIBMESH_ENABLE_PERIODIC
   /**
    * Register a pair of boundaries as disjoint neighbor boundary pairs.
@@ -2257,6 +2266,12 @@ class MeshBase : public ParallelObject
    */
   std::set<subdomain_id_type> _mesh_subdomains;
 
+  /**
+   * We also cache the subdomain ids of the elements owned by this
+   * processor.
+   */
+  std::set<subdomain_id_type> _mesh_local_subdomains;
+
   /**
    * Map from "element set code" to list of set ids to which that element
    * belongs (and vice-versa). Remarks:

src/mesh/distributed_mesh.C

@@ -1060,7 +1060,7 @@ void DistributedMesh::redistribute ()
 
 void DistributedMesh::update_post_partitioning ()
 {
-  // this->recalculate_n_partitions();
+  this->UnstructuredMesh::update_post_partitioning();
 
   // Partitioning changes our numbers of unpartitioned objects
   this->update_parallel_id_counts();

src/mesh/mesh_base.C

@@ -120,6 +120,7 @@ MeshBase::MeshBase (const MeshBase & other_mesh) :
   _elem_default_orders(other_mesh._elem_default_orders),
   _supported_nodal_order(other_mesh._supported_nodal_order),
   _mesh_subdomains(other_mesh._mesh_subdomains),
+  _mesh_local_subdomains(other_mesh._mesh_local_subdomains),
   _elemset_codes_inverse_map(other_mesh._elemset_codes_inverse_map),
   _all_elemset_ids(other_mesh._all_elemset_ids),
   _spatial_dimension(other_mesh._spatial_dimension),
@@ -209,6 +210,7 @@ MeshBase& MeshBase::operator= (MeshBase && other_mesh)
   _elem_default_orders = std::move(other_mesh.elem_default_orders());
   _supported_nodal_order = other_mesh.supported_nodal_order();
   _mesh_subdomains = other_mesh._mesh_subdomains;
+  _mesh_local_subdomains = other_mesh._mesh_local_subdomains;
   _elemset_codes = std::move(other_mesh._elemset_codes);
   _elemset_codes_inverse_map = std::move(other_mesh._elemset_codes_inverse_map);
   _all_elemset_ids = std::move(other_mesh._all_elemset_ids);
@@ -325,6 +327,8 @@ bool MeshBase::locally_equals (const MeshBase & other_mesh) const
     return false;
   if (_mesh_subdomains != other_mesh._mesh_subdomains)
     return false;
+  if (_mesh_local_subdomains != other_mesh._mesh_local_subdomains)
+    return false;
   if (_all_elemset_ids != other_mesh._all_elemset_ids)
     return false;
   if (_elem_integer_names != other_mesh._elem_integer_names)
@@ -1112,6 +1116,16 @@ void MeshBase::redistribute()
 
 
 
+void MeshBase::update_post_partitioning()
+{
+  _mesh_local_subdomains.clear();
+
+  for (const Elem * elem : this->active_local_element_ptr_range())
+    _mesh_local_subdomains.insert(elem->subdomain_id());
+}
+
+
+
 subdomain_id_type MeshBase::n_subdomains() const
 {
   // This requires an inspection on every processor
@@ -1860,13 +1874,16 @@ void MeshBase::cache_elem_data()
   _elem_dims.clear();
   _elem_default_orders.clear();
   _mesh_subdomains.clear();
+  _mesh_local_subdomains.clear();
   _supported_nodal_order = MAXIMUM;
 
   for (const auto & elem : this->active_element_ptr_range())
   {
     _elem_dims.insert(cast_int<unsigned char>(elem->dim()));
     _elem_default_orders.insert(elem->default_order());
     _mesh_subdomains.insert(elem->subdomain_id());
+    if (elem->processor_id() == this->processor_id())
+      _mesh_local_subdomains.insert(elem->subdomain_id());
     _supported_nodal_order =
       static_cast<Order>
         (std::min(static_cast<int>(_supported_nodal_order),
@@ -2220,6 +2237,7 @@ MeshBase::copy_cached_data(const MeshBase & other_mesh)
   this->_elem_default_orders = other_mesh._elem_default_orders;
   this->_supported_nodal_order = other_mesh._supported_nodal_order;
   this->_mesh_subdomains = other_mesh._mesh_subdomains;
+  this->_mesh_local_subdomains = other_mesh._mesh_local_subdomains;
 }
 
 
@@ -2474,6 +2492,7 @@ MeshBase::copy_constraint_rows(const SparseMatrix<T> & constraint_operator,
           (std::min(static_cast<int>(this->_supported_nodal_order),
                     static_cast<int>(added_elem->supported_nodal_order())));
       this->_mesh_subdomains.insert(new_sbd_id);
+      this->_mesh_local_subdomains.insert(new_sbd_id);
       node_to_elem_ptrs.emplace(n, std::make_pair(added_elem->id(), 0));
       existing_unconstrained_columns.emplace(j,n->id());