From 1d9205ce79f08f2d96376dd7ed10ee482e8dd2e4 Mon Sep 17 00:00:00 2001
From: JordiManyer <jordi.manyer@monash.edu>
Date: Tue, 25 Jun 2024 15:06:17 +1000
Subject: [PATCH] Bugfix: Nonconforming block assembly

---
 src/Algebra.jl                           | 30 +++++++++--------
 test/BlockSparseMatrixAssemblersTests.jl | 42 +++++++++++++++++++++---
 2 files changed, 55 insertions(+), 17 deletions(-)

diff --git a/src/Algebra.jl b/src/Algebra.jl
index cd1a163..57bcbac 100644
--- a/src/Algebra.jl
+++ b/src/Algebra.jl
@@ -1129,25 +1129,29 @@ function _setup_prange(dofs_gids_prange::PRange,gids;ghost=true,owners=nothing,k
   end
 end
 
-function _setup_prange(dofs_gids_prange::AbstractVector{<:PRange},
-                       gids::AbstractMatrix;
-                       ax=:rows,ghost=true,owners=nothing)
+function _setup_prange(
+  dofs_gids_prange::AbstractVector{<:PRange},
+  gids::AbstractMatrix;
+  ax=:rows,ghost=true,owners=nothing
+)
   @check ax ∈ (:rows,:cols)
   block_ids = LinearIndices(dofs_gids_prange)
+  pvcat(x) = map(xi -> vcat(xi...), to_parray_of_arrays(x))
 
-  gids_ax_slice, _owners = map(block_ids,dofs_gids_prange) do id,prange
-    gids_ax_slice = (ax == :rows) ? gids[id,:] : gids[:,id]
-    _owners = nothing
-    if ghost
-      gids_ax_slice = map(x -> union(x...), to_parray_of_arrays(gids_ax_slice))
-      if !isa(owners,Nothing) # Recompute owners for the union
-        _owners = get_gid_owners(gids_ax_slice,prange)
-      end
+  gids_union, owners_union = map(block_ids,dofs_gids_prange) do id, prange
+    gids_slice = (ax == :rows) ? gids[id,:] : gids[:,id]
+    gids_union = pvcat(gids_slice)
+
+    owners_union = nothing
+    if !isnothing(owners)
+      owners_slice = (ax == :rows) ? owners[id,:] : owners[:,id]
+      owners_union = pvcat(owners_slice)
     end
-    return gids_ax_slice, _owners
+
+    return gids_union, owners_union
   end |> tuple_of_arrays
   
-  return map((p,g,o) -> _setup_prange(p,g;ghost=ghost,owners=o),dofs_gids_prange,gids_ax_slice,_owners)
+  return map((p,g,o) -> _setup_prange(p,g;ghost=ghost,owners=o),dofs_gids_prange,gids_union,owners_union)
 end
 
 # Create PRange for the rows of the linear system
diff --git a/test/BlockSparseMatrixAssemblersTests.jl b/test/BlockSparseMatrixAssemblersTests.jl
index c9c3384..82bafc9 100644
--- a/test/BlockSparseMatrixAssemblersTests.jl
+++ b/test/BlockSparseMatrixAssemblersTests.jl
@@ -9,6 +9,11 @@ using GridapDistributed
 using PartitionedArrays
 using GridapDistributed: BlockPVector, BlockPMatrix
 
+get_edge_measures(Ω::Triangulation,dΩ) = sqrt∘CellField(get_array(∫(1)dΩ),Ω)
+function get_edge_measures(Ω::GridapDistributed.DistributedTriangulation,dΩ)
+  return sqrt∘CellField(map(get_array,local_views(∫(1)*dΩ)),Ω)
+end
+
 function is_same_vector(x::BlockPVector,y::PVector,Ub,U)
   y_fespace = GridapDistributed.change_ghost(y,U.gids)
   x_fespace = GridapDistributed.change_ghost(x,Ub.gids)
@@ -33,7 +38,7 @@ function is_same_matrix(Ab::BlockPMatrix,A::PSparseMatrix,Xb,X)
   return is_same_vector(yb,y,Xb,X)
 end
 
-function _main(n_spaces,mfs,weakform,Ω,dΩ,U,V)
+function _main(n_spaces,mfs,weakform,U,V)
   biform, liform = weakform
 
   # Normal assembly 
@@ -94,13 +99,14 @@ function main(distribute,parts)
   ranks = distribute(LinearIndices((prod(parts),)))
 
   model = CartesianDiscreteModel(ranks,parts,(0,1,0,1),(8,8))
-  Ω = Triangulation(model)
 
+  # Conforming tests
   sol(x) = sum(x)
   reffe  = LagrangianRefFE(Float64,QUAD,1)
-  V = FESpace(Ω, reffe; dirichlet_tags="boundary")
+  V = FESpace(model, reffe; dirichlet_tags="boundary")
   U = TrialFESpace(sol,V)
 
+  Ω = Triangulation(model)
   dΩ = Measure(Ω, 2)
   biform2((u1,u2),(v1,v2)) = ∫(∇(u1)⋅∇(v1) + u2⋅v2 + u1⋅v2)*dΩ
   liform2((v1,v2)) = ∫(v1 + v2)*dΩ
@@ -109,9 +115,37 @@ function main(distribute,parts)
 
   for (n_spaces,weakform) in zip([2,3],[(biform2,liform2),(biform3,liform3)])
     for mfs in [BlockMultiFieldStyle(),BlockMultiFieldStyle(2,(1,n_spaces-1))]
-      _main(n_spaces,mfs,weakform,Ω,dΩ,U,V)
+      _main(n_spaces,mfs,weakform,U,V)
     end
   end
+
+  # Non-conforming tests (tests whether we can fetch neighbors from non-ghosts)
+  reffe = ReferenceFE(raviart_thomas,Float64,0)
+  D = FESpace(model,reffe; dirichlet_tags="boundary")
+
+  β_U = 100.0
+  Γ = Boundary(model)
+  Λ = Skeleton(model)
+  n_Γ = get_normal_vector(Γ)
+  n_Λ = get_normal_vector(Λ)
+  dΓ = Measure(Γ,2)
+  dΛ = Measure(Λ,2)
+  h_e_Λ = get_edge_measures(Λ,dΛ)
+  h_e_Γ = get_edge_measures(Γ,dΓ)
+  lap_dg(u,v) = ∫(∇(u)⊙∇(v))dΩ -
+                ∫(jump(v⊗n_Λ)⊙(mean(∇(u))))dΛ -
+                ∫(mean(∇(v))⊙jump(u⊗n_Λ))dΛ -
+                ∫(v⋅(∇(u)⋅n_Γ))dΓ -
+                ∫((∇(v)⋅n_Γ)⋅u)dΓ +
+                ∫(jump(v⊗n_Λ)⊙((β_U/h_e_Λ*jump(u⊗n_Λ))))dΛ +
+                ∫(v⋅(β_U/h_e_Γ*u))dΓ
+
+  f = VectorValue(1.0,1.0)
+  biform4((u1,u2),(v1,v2)) = lap_dg(u1,v1) + lap_dg(u1,v1) + ∫(u2⋅v2)*dΩ
+  liform4((v1,v2)) = ∫(v1⋅f)*dΩ
+  _main(2,BlockMultiFieldStyle(),(biform4,liform4),D,D)
 end
 
+main(DebugArray, (2,2))
+
 end
\ No newline at end of file