From 272dc6f9492cce113171f712d4a87954b771173d Mon Sep 17 00:00:00 2001
From: JordiManyer <jordi.manyer@monash.edu>
Date: Tue, 1 Oct 2024 10:46:32 +1000
Subject: [PATCH] Added local-constant fespaces

---
 src/FESpaces.jl               | 18 ++++++++++--
 test/ConstantFESpacesTests.jl | 52 +++++++++++++++--------------------
 2 files changed, 37 insertions(+), 33 deletions(-)

diff --git a/src/FESpaces.jl b/src/FESpaces.jl
index 5bb0dff..43c7d8d 100644
--- a/src/FESpaces.jl
+++ b/src/FESpaces.jl
@@ -870,17 +870,29 @@ function _compute_new_distributed_fixedval(
   return c
 end
 
-function FESpaces.ConstantFESpace(model::DistributedDiscreteModel;kwargs...)
-  @warn "ConstantFESpace is NOT scalable in parallel. For testing purposes only."
+function FESpaces.ConstantFESpace(
+  model::DistributedDiscreteModel;
+  constraint_type=:global,kwargs...
+)
+  @assert constraint_type ∈ [:global,:local]
+  if constraint_type == :global
+    @warn "ConstantFESpace is NOT scalable in parallel. For testing purposes only."
+  end
+
   spaces = map(local_views(model)) do model
     ConstantFESpace(model;kwargs...)
   end
 
   # Single dof, owned by processor 1 (ghost for all other processors)
+  nranks = length(spaces)
   cell_gids = get_cell_gids(model)
   indices = map(partition(cell_gids)) do cell_indices
     me = part_id(cell_indices)
-    LocalIndices(1,me,Int[1],Int32[1])
+    if constraint_type == :global
+      LocalIndices(1,me,Int[1],Int32[1])
+    else
+      LocalIndices(nranks,me,Int[me],Int32[me])
+    end
   end
   gids = PRange(indices)
 
diff --git a/test/ConstantFESpacesTests.jl b/test/ConstantFESpacesTests.jl
index 8a66568..7623981 100644
--- a/test/ConstantFESpacesTests.jl
+++ b/test/ConstantFESpacesTests.jl
@@ -1,46 +1,38 @@
-
+using Test
 using Gridap
 using GridapDistributed, PartitionedArrays
 
 using Gridap.FESpaces, Gridap.Arrays, Gridap.Algebra
 
-np = (2,2)
-ranks = with_mpi() do distribute
-  distribute(LinearIndices((prod(np),)))
-end
+function main(distribute,np)
+  @assert prod(np) == 4
+  ranks = distribute(LinearIndices((prod(np),)))
 
-model = CartesianDiscreteModel(ranks,np,(0,1,0,1),(10,10))
+  model = CartesianDiscreteModel(ranks,np,(0,1,0,1),(10,10))
 
-Ω = Triangulation(model)
-dΩ = Measure(Ω,2)
+  Ω = Triangulation(model)
+  dΩ = Measure(Ω,2)
 
-reffe = ReferenceFE(lagrangian, Float64, 1)
-U = FESpace(model,reffe)
-V = ConstantFESpace(model)
-X = MultiFieldFESpace([U,V])
+  reffe = ReferenceFE(lagrangian, Float64, 1)
+  U = FESpace(model,reffe)
+  V = ConstantFESpace(model)
+  X = MultiFieldFESpace([U,V])
 
-gids = get_free_dof_ids(V)
-local_to_global(gids)
-local_to_owner(gids)
-ghost_to_local(gids)
-own_to_local(gids)
+  gids = get_free_dof_ids(V)
+  @test length(gids) == 1
 
-uh = zero(V)
-sum(∫(uh)dΩ)
+  a((u,λ),(v,μ)) = ∫(u*v + λ*v + μ*u)dΩ 
+  A = assemble_matrix(a,X,X)
 
-map(ranks) do r
-  println("flag 1 from ", r)
-end
+  V2 = ConstantFESpace(model;constraint_type=:local)
+  X2 = MultiFieldFESpace([U,V2])
 
-a((u,λ),(v,μ)) = ∫(u*v + λ*v + μ*u)dΩ 
-A = assemble_matrix(a,X,X)
+  gids = get_free_dof_ids(V2)
+  @test length(gids) == 4
 
-map(ranks) do r
-  println("flag 2 from ", r)
+  A2 = assemble_matrix(a,X2,X2)
 end
 
-x = allocate_in_domain(A)
-
-map(ranks) do r
-  println("done from ", r)
+with_debug() do distribute
+  main(distribute,(2,2))
 end