From 5372644858b53e05b844ce54d48b7083c55583bb Mon Sep 17 00:00:00 2001
From: Francesc Verdugo <fverdugo@cimne.upc.edu>
Date: Wed, 15 Apr 2020 11:40:01 +0200
Subject: [PATCH] Added HomogeneousTrialFESpace and other enhancements in
 TrialFESpace

---
 src/FESpaces/FESpaces.jl                 |  2 ++
 src/FESpaces/TrialFESpaces.jl            | 28 ++++++++++++++++++------
 test/FESpacesTests/TrialFESpacesTests.jl | 10 ++++++++-
 3 files changed, 32 insertions(+), 8 deletions(-)

diff --git a/src/FESpaces/FESpaces.jl b/src/FESpaces/FESpaces.jl
index 460cc389e..3ca28154c 100644
--- a/src/FESpaces/FESpaces.jl
+++ b/src/FESpaces/FESpaces.jl
@@ -162,6 +162,8 @@ export is_a_fe_cell_basis
 
 export TrialFESpace
 export TrialFESpace!
+export HomogeneousTrialFESpace
+export HomogeneousTrialFESpace!
 export TestFESpace
 export compute_conforming_cell_dofs
 export SparseMatrixAssembler
diff --git a/src/FESpaces/TrialFESpaces.jl b/src/FESpaces/TrialFESpaces.jl
index a5c0d6c3a..2b6bef990 100644
--- a/src/FESpaces/TrialFESpaces.jl
+++ b/src/FESpaces/TrialFESpaces.jl
@@ -4,30 +4,34 @@ struct TrialFESpace{B} <: SingleFieldFESpace
   dirichlet_values::AbstractVector
   cell_basis::CellBasis
   constraint_style::Val{B}
+
+  function TrialFESpace(dirichlet_values::AbstractVector,space::SingleFieldFESpace)
+    cell_basis = _prepare_trial_cell_basis(space)
+    cs = constraint_style(space)
+    B = get_val_parameter(cs)
+    new{B}(space,dirichlet_values,cell_basis,cs)
+  end
 end
 
 """
 """
 function TrialFESpace(space::SingleFieldFESpace)
   dirichlet_values = get_dirichlet_values(space)
-  cell_basis = _prepare_trial_cell_basis(space)
-  TrialFESpace(space,dirichlet_values,cell_basis,constraint_style(space))
+  TrialFESpace(dirichlet_values,space)
 end
 
 """
 """
 function TrialFESpace(space::SingleFieldFESpace,objects)
   dirichlet_values = compute_dirichlet_values_for_tags(space,objects)
-  cell_basis = _prepare_trial_cell_basis(space)
-  TrialFESpace(space,dirichlet_values,cell_basis,constraint_style(space))
+  TrialFESpace(dirichlet_values,space)
 end
 
 """
 """
-function TrialFESpace(dir_values::AbstractVector,space::SingleFieldFESpace,objects)
+function TrialFESpace!(dir_values::AbstractVector,space::SingleFieldFESpace,objects)
   dir_values = compute_dirichlet_values_for_tags!(dir_values,space,objects)
-  cell_basis = _prepare_trial_cell_basis(space)
-  TrialFESpace(space,dir_values,cell_basis,constraint_style(space))
+  TrialFESpace(dir_values,space)
 end
 
 """
@@ -42,6 +46,16 @@ function TrialFESpace(space::TrialFESpace)
   space
 end
 
+function HomogeneousTrialFESpace(U::FESpace)
+  dirichlet_values = zero_dirichlet_values(U)
+  TrialFESpace(dirichlet_values,U)
+end
+
+function HomogeneousTrialFESpace!(dirichlet_values::AbstractVector,U::FESpace)
+  fill!(dirichlet_values,zero(eltype(dirichlet_values)))
+  TrialFESpace(dirichlet_values,U)
+end
+
 function  _prepare_trial_cell_basis(space)
   cb = get_cell_basis(space)
   a = get_array(cb)
diff --git a/test/FESpacesTests/TrialFESpacesTests.jl b/test/FESpacesTests/TrialFESpacesTests.jl
index af1c8c219..6aa452232 100644
--- a/test/FESpacesTests/TrialFESpacesTests.jl
+++ b/test/FESpacesTests/TrialFESpacesTests.jl
@@ -29,7 +29,7 @@ isa(V,SingleFieldFESpace)
 ud = compute_dirichlet_values_for_tags!(v,V,[4,3])
 @test all(ud .== v)
 test_single_field_fe_space(U)
-U = TrialFESpace(v,V,[4,3])
+U = TrialFESpace!(v,V,[4,3])
 
 
 matvecdata = ([],[],[])
@@ -46,6 +46,14 @@ TrialFESpace!(U,[1,2])
 cell_basis = get_cell_basis(U)
 @test is_trial(cell_basis)
 
+U0 = HomogeneousTrialFESpace(V)
+@test get_dirichlet_values(U0) == zeros(6)
+
+U0 = HomogeneousTrialFESpace!(v,V)
+@test v === get_dirichlet_values(U0)
+@test v == zeros(6)
+@test get_dirichlet_values(U0) == zeros(6)
+
 #trian = get_triangulation(model)
 #
 #using Gridap.Visualization