From ab452bb2075f130268e1c051e0c262f555181943 Mon Sep 17 00:00:00 2001
From: JordiManyer <jordi.manyer@monash.edu>
Date: Mon, 16 Dec 2024 21:10:50 +1100
Subject: [PATCH] Expanded ConstantFESpaces to work on sub-triangulations

---
 src/FESpaces/ConstantFESpaces.jl           | 51 ++++++++++++----------
 test/FESpacesTests/ConstantFESpaceTests.jl | 12 +++--
 2 files changed, 37 insertions(+), 26 deletions(-)

diff --git a/src/FESpaces/ConstantFESpaces.jl b/src/FESpaces/ConstantFESpaces.jl
index f7115f46c..71b7d26b9 100644
--- a/src/FESpaces/ConstantFESpaces.jl
+++ b/src/FESpaces/ConstantFESpaces.jl
@@ -5,45 +5,52 @@ struct ConstantFESpace <: SingleFieldFESpace
 end
 """
 struct ConstantFESpace{V,T,A,B,C} <: SingleFieldFESpace
-  model::DiscreteModel
+  trian::Triangulation
   cell_basis::A
   cell_dof_basis::B
   cell_dof_ids::C
 
-  function ConstantFESpace(model;
-                         vector_type::Type{V}=Vector{Float64},
-                         field_type::Type{T}=Float64) where {V,T}
-    function setup_cell_reffe(model::DiscreteModel,
-        reffe::Tuple{<:ReferenceFEName,Any,Any}; kwargs...)
-      basis, reffe_args,reffe_kwargs = reffe
-      cell_reffe = ReferenceFE(model,basis,reffe_args...;reffe_kwargs...)
-    end
-
-    reffe = ReferenceFE(lagrangian,T,0)
-    cell_reffe = setup_cell_reffe(model,reffe)
+  function ConstantFESpace(
+    model::DiscreteModel{Dc},
+    trian::Triangulation{Dc};
+    vector_type::Type{V}=Vector{Float64},
+    field_type::Type{T}=Float64
+  ) where {Dc,V,T}
+    @assert num_cells(model) == num_cells(trian)
+
+    basis, reffe_args, reffe_kwargs = ReferenceFE(lagrangian,T,0)
+    cell_reffe = ReferenceFE(model,basis,reffe_args...;reffe_kwargs...)
     cell_basis_array = lazy_map(get_shapefuns,cell_reffe)
 
     cell_basis = SingleFieldFEBasis(
-      cell_basis_array,
-      Triangulation(model),
-      TestBasis(),
-      ReferenceDomain())
-
-    cell_dof_basis_array = lazy_map(get_dof_basis,cell_reffe)
-    cell_dof_basis = CellDof(cell_dof_basis_array,Triangulation(model),ReferenceDomain())
+      cell_basis_array, trian, TestBasis(), ReferenceDomain()
+    )
+    cell_dof_basis = CellDof(
+      lazy_map(get_dof_basis,cell_reffe),trian,ReferenceDomain()
+    )
+    cell_dof_ids = Fill(Int32(1):Int32(num_indep_components(field_type)),num_cells(trian))
 
-    cell_dof_ids = Fill(Int32(1):Int32(num_indep_components(field_type)),num_cells(model))
     A = typeof(cell_basis)
     B = typeof(cell_dof_basis)
     C = typeof(cell_dof_ids)
-    new{V,T,A,B,C}(model, cell_basis, cell_dof_basis, cell_dof_ids)
+    new{V,T,A,B,C}(trian, cell_basis, cell_dof_basis, cell_dof_ids)
   end
 end
 
+function ConstantFESpace(model::DiscreteModel; kwargs...)
+  trian = Triangulation(model)
+  ConstantFESpace(model,trian; kwargs...)
+end
+
+function ConstantFESpace(trian::Triangulation; kwargs...)
+  model = get_active_model(trian)
+  ConstantFESpace(model,trian; kwargs...)
+end
+
 TrialFESpace(f::ConstantFESpace) = f
 
 # Delegated functions
-get_triangulation(f::ConstantFESpace) = Triangulation(f.model)
+get_triangulation(f::ConstantFESpace) = f.trian
 
 ConstraintStyle(::Type{<:ConstantFESpace}) = UnConstrained()
 
diff --git a/test/FESpacesTests/ConstantFESpaceTests.jl b/test/FESpacesTests/ConstantFESpaceTests.jl
index 980ee6124..cc81846f7 100644
--- a/test/FESpacesTests/ConstantFESpaceTests.jl
+++ b/test/FESpacesTests/ConstantFESpaceTests.jl
@@ -6,9 +6,9 @@ using Test
 domain = (0,1,0,1)
 partition = (4,4)
 model = CartesianDiscreteModel(domain,partition)
-Λ=ConstantFESpace(model)
+Λ = ConstantFESpace(model)
 Gridap.FESpaces.test_fe_space(Λ)
-M=TrialFESpace(Λ)
+M = TrialFESpace(Λ)
 
 order = 2
 u((x,y)) = (x+y)^order
@@ -29,13 +29,17 @@ uh = solve(op)
 @assert sum(∫((uh[1]-u)*(uh[1]-u))dΩ) < 1.0e-14
 abs(sum(∫(uh[2])dΩ)) < 1.0e-12
 
-Λ2=ConstantFESpace(model,field_type=VectorValue{2,Float64})
+Λ2 = ConstantFESpace(model,field_type=VectorValue{2,Float64})
 Gridap.FESpaces.test_fe_space(Λ2)
-M2=TrialFESpace(Λ2)
+M2 = TrialFESpace(Λ2)
 a2(μ,λ) = ∫(λ⋅μ)dΩ
 l2(λ) = ∫(VectorValue(0.0,0.0)⋅λ)dΩ
 op2 = AffineFEOperator(a2,l2,M2,Λ2)
 μ2h = solve(op2)
 @assert sum(∫(μ2h⋅μ2h)dΩ) < 1.0e-12
 
+trian = Triangulation(model,[1,2,3,4])
+Λ3 = ConstantFESpace(trian,field_type=VectorValue{2,Float64})
+Gridap.FESpaces.test_fe_space(Λ3)
+
 end # module