From 95633733a48b3eae143adf585eeef36b4e1177eb Mon Sep 17 00:00:00 2001
From: JordiManyer <jordi.manyer@monash.edu>
Date: Wed, 12 Feb 2025 15:46:06 +1100
Subject: [PATCH] Added more cases to add_ghost_cells that tries to reuse the
 same triangulation if possible

---
 src/Geometry.jl | 77 ++++++++++++++++++++++++++-----------------------
 1 file changed, 41 insertions(+), 36 deletions(-)

diff --git a/src/Geometry.jl b/src/Geometry.jl
index 026ea93..f69d003 100644
--- a/src/Geometry.jl
+++ b/src/Geometry.jl
@@ -741,33 +741,39 @@ function add_ghost_cells(
   dmodel::DistributedDiscreteModel{Dm},
   dtrian::DistributedTriangulation{Dt}
 ) where {Dm,Dt}
-  covers_all_faces = _covers_all_faces(dmodel,dtrian)
-  if (covers_all_faces)
-    trians = map(local_views(dmodel)) do model
-      Triangulation(ReferenceFE{Dt},model)
-    end
-    return DistributedTriangulation(trians,dmodel)
-  else
-    mcell_intrian = map(local_views(dmodel),local_views(dtrian)) do model, trian
-      glue = get_glue(trian,Val(Dt))
-      @assert isa(glue,FaceToFaceGlue)
-      nmcells = num_faces(model,Dt)
-      mcell_intrian = fill(false,nmcells)
-      tcell_to_mcell = glue.tface_to_mface
-      mcell_intrian[tcell_to_mcell] .= true
-      mcell_intrian
-    end
-    gids = get_face_gids(dmodel,Dt)
 
-    cache=fetch_vector_ghost_values_cache(mcell_intrian,partition(gids))
-    fetch_vector_ghost_values!(mcell_intrian,cache) |> wait
+  tface_to_mface = map(local_views(dtrian)) do trian
+    glue = get_glue(trian,Val(Dt))
+    @assert isa(glue,FaceToFaceGlue)
+    glue.tface_to_mface
+  end
+
+  # Case 1: All model faces are already in the triangulation
+  covers_all_faces = reduce(&,map(x -> isa(x,IdentityVector), tface_to_mface),init=true)
+  covers_all_faces && return dtrian
 
-    dreffes = map(local_views(dmodel)) do model
-      ReferenceFE{Dt}
-    end
-    trians = map(Triangulation,dreffes,local_views(dmodel),mcell_intrian)
-    return DistributedTriangulation(trians,dmodel)
+  # Otherwise: Add ghost cells to triangulation
+  mface_intrian = map(local_views(dmodel),tface_to_mface) do model, tface_to_mface
+    mface_intrian = fill(false,num_faces(model,Dt))
+    mface_intrian[tface_to_mface] .= true
+    mface_intrian
   end
+  consistent!(PVector(mface_intrian,partition(get_face_gids(dmodel,Dt)))) |> wait
+
+  # Case 2: New triangulation contains all model faces
+  covers_all_faces = reduce(&,map(all,mface_intrian),init=true)
+  covers_all_faces && return Triangulation(with_ghost,ReferenceFE{Dt},dmodel)
+
+  # Case 3: Original triangulation already had the ghost cells
+  new_tface_to_mface = map(findall,mface_intrian)
+  had_ghost = reduce(&,map(==,tface_to_mface,new_tface_to_mface),init=true)
+  had_ghost && return dtrian
+
+  # Case 4: Otherwise, create a new triangulation with the ghost cells
+  trians = map(local_views(dmodel),new_tface_to_mface) do model, new_tface_to_mface
+    Triangulation(ReferenceFE{Dt},model,new_tface_to_mface)
+  end
+  return DistributedTriangulation(trians,dmodel)
 end
 
 function generate_cell_gids(dtrian::DistributedTriangulation)
@@ -778,11 +784,11 @@ end
 function generate_cell_gids(dmodel::DistributedDiscreteModel{Dm},
                             dtrian::DistributedTriangulation{Dt}) where {Dm,Dt}
 
+  mgids = get_face_gids(dmodel,Dt)
   covers_all_faces = _covers_all_faces(dmodel,dtrian)
   if (covers_all_faces)
-    get_face_gids(dmodel,Dt)
+    tgids = mgids
   else
-    mgids = get_face_gids(dmodel,Dt)
     # count number owned cells
     notcells, tcell_to_mcell = map(
       local_views(dmodel),local_views(dtrian),PArrays.partition(mgids)) do model,trian,partition
@@ -802,7 +808,7 @@ function generate_cell_gids(dmodel::DistributedDiscreteModel{Dm},
 
     # Assign global cell ids to owned cells
     mcell_to_gtcell = map(
-      first_gtcell,tcell_to_mcell,PArrays.partition(mgids)) do first_gtcell,tcell_to_mcell,partition
+      first_gtcell,tcell_to_mcell,partition(mgids)) do first_gtcell,tcell_to_mcell,partition
       mcell_to_gtcell = zeros(Int,local_length(partition))
       loc_to_owner = local_to_owner(partition)
       part = part_id(partition)
@@ -816,22 +822,21 @@ function generate_cell_gids(dmodel::DistributedDiscreteModel{Dm},
       mcell_to_gtcell
     end
 
-    cache = fetch_vector_ghost_values_cache(mcell_to_gtcell,PArrays.partition(mgids))
+    cache = fetch_vector_ghost_values_cache(mcell_to_gtcell,partition(mgids))
     fetch_vector_ghost_values!(mcell_to_gtcell,cache) |> wait
 
     # Prepare new partition
     ngtcells = reduction(+,notcells,destination=:all,init=zero(eltype(notcells)))
-    partition = map(ngtcells,
-                    mcell_to_gtcell,
-                    tcell_to_mcell,
-                    PArrays.partition(mgids)) do ngtcells,mcell_to_gtcell,tcell_to_mcell,partition
+    indices = map(
+      ngtcells,mcell_to_gtcell,tcell_to_mcell,partition(mgids)
+    ) do ngtcells,mcell_to_gtcell,tcell_to_mcell,partition
       tcell_to_gtcell = mcell_to_gtcell[tcell_to_mcell]
-      lid_to_owner = local_to_owner(partition)
+      lid_to_owner  = local_to_owner(partition)
       tcell_to_part = lid_to_owner[tcell_to_mcell]
       LocalIndices(ngtcells,part_id(partition),tcell_to_gtcell,tcell_to_part)
     end
-    _find_neighbours!(partition, PArrays.partition(mgids))
-    gids = PRange(partition)
-    gids
+    _find_neighbours!(indices, partition(mgids))
+    tgids = PRange(indices)
   end
+  return tgids
 end