Use cached data for n_subdomains()

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
@@ -1492,6 +1492,8 @@ class MeshBase : public ParallelObject
    * materials in a solid mechanics application, or regions where different
    * physical processes are important.  The subdomain mapping is independent
    * from the parallel decomposition.
+   *
+   * This relies on the mesh cached element data being prepared.
    */
   subdomain_id_type n_subdomains () const;
 
@@ -1501,6 +1503,8 @@ class MeshBase : public ParallelObject
    * materials in a solid mechanics application, or regions where different
    * physical processes are important.  The subdomain mapping is independent
    * from the parallel decomposition.
+   *
+   * This relies on the mesh cached element data being prepared.
    */
   subdomain_id_type n_local_subdomains () const;
 
@@ -1993,6 +1997,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 +2270,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)
@@ -1076,26 +1080,12 @@ void MeshBase::remove_ghosting_functor(GhostingFunctor & ghosting_functor)
 
 void MeshBase::subdomain_ids (std::set<subdomain_id_type> & ids, const bool global /* = true */) const
 {
-  // This requires an inspection on every processor
-  if (global)
-    parallel_object_only();
-
-  ids.clear();
-
-  for (const auto & elem : this->active_local_element_ptr_range())
-    ids.insert(elem->subdomain_id());
+  libmesh_assert(this->preparation().has_cached_elem_data);
 
   if (global)
-    {
-      // Only include the unpartitioned elements if the user requests the global IDs.
-      // In the case of the local subdomain IDs, it doesn't make sense to include the
-      // unpartitioned elements because said elements do not have a sense of locality.
-      for (const auto & elem : this->active_unpartitioned_element_ptr_range())
-        ids.insert(elem->subdomain_id());
-
-      // Some subdomains may only live on other processors
-      this->comm().set_union(ids);
-    }
+    ids = this->get_mesh_subdomains();
+  else
+    ids = this->get_mesh_local_subdomains();
 }
 
 
@@ -1112,16 +1102,19 @@ void MeshBase::redistribute()
 
 
 
-subdomain_id_type MeshBase::n_subdomains() const
+void MeshBase::update_post_partitioning()
 {
-  // This requires an inspection on every processor
-  parallel_object_only();
+  _mesh_local_subdomains.clear();
 
-  std::set<subdomain_id_type> ids;
+  for (const Elem * elem : this->active_local_element_ptr_range())
+    _mesh_local_subdomains.insert(elem->subdomain_id());
+}
 
-  this->subdomain_ids (ids);
 
-  return cast_int<subdomain_id_type>(ids.size());
+
+subdomain_id_type MeshBase::n_subdomains() const
+{
+  return cast_int<subdomain_id_type>(this->get_mesh_subdomains().size());
 }
 
 
@@ -1860,13 +1853,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 +2216,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 +2471,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());