From 4b0d45d9b1ab92a6c22867ade49c5bdc5af940d3 Mon Sep 17 00:00:00 2001
From: JordiManyer <jordi.manyer@monash.edu>
Date: Tue, 23 Jul 2024 17:08:53 +1000
Subject: [PATCH 01/15] Started adding sero-mean FESpace

---
 src/FESpaces.jl | 33 +++++++++++++++++++++++++++++++--
 1 file changed, 31 insertions(+), 2 deletions(-)

diff --git a/src/FESpaces.jl b/src/FESpaces.jl
index 4480c2f..ec354bc 100644
--- a/src/FESpaces.jl
+++ b/src/FESpaces.jl
@@ -356,13 +356,13 @@ function _EvaluationFunction(func,
 end
 
 function FESpaces.get_fe_basis(f::DistributedSingleFieldFESpace)
-  fields = map(get_fe_basis,f.spaces)
+  fields = map(get_fe_basis,local_views(f))
   trian  = get_triangulation(f)
   DistributedCellField(fields,trian)
 end
 
 function FESpaces.get_trial_fe_basis(f::DistributedSingleFieldFESpace)
-  fields = map(get_trial_fe_basis,f.spaces)
+  fields = map(get_trial_fe_basis,local_views(f))
   trian  = get_triangulation(f)
   DistributedCellField(fields,trian)
 end
@@ -702,3 +702,32 @@ function FESpaces.SparseMatrixAssembler(
   Tm = SparseMatrixCSC{T,Int}
   SparseMatrixAssembler(Tm,Tv,trial,test,par_strategy)
 end
+
+# ZeroMean FESpace
+
+#const DistributedZeroMeanFESpace = DistributedSingleFieldFESpace{}
+
+struct ZeroMeanCache{A}
+  vol_i::A
+  vol::Float64
+end
+
+function FESpaces.ZeroMeanFESpace(space::DistributedFESpace)#,dΩ::DistributedMeasure)
+  gids = get_free_dof_ids(space)
+  ranks = get_parts(space)
+  spaces = map(ranks,partition(gids),local_views(space)) do r, gids, lspace
+    fix_constant = isone(r) # Only main processor fixes the constant
+    dof_to_fix = Int(first(own_to_local(gids))) # Make sure it's an owned DoF
+    FESpaceWithConstantFixed(lspace,fix_constant,dof_to_fix)
+  end
+
+  #vol_i = assemble_vector(v->∫(v)*dΩ,space)
+  #vol = sum(vol_i)
+  #metadata = ZeroMeanCache(vol_i,vol)
+
+  trian = get_triangulation(space)
+  model = get_background_model(trian)
+  gids =  generate_gids(model,spaces)
+  vector_type = _find_vector_type(spaces,gids)
+  DistributedSingleFieldFESpace(spaces,gids,trian,vector_type)
+end

From 503e37ea7934eed477ea6d9c89d8fa4a81650404 Mon Sep 17 00:00:00 2001
From: JordiManyer <jordi.manyer@monash.edu>
Date: Wed, 24 Jul 2024 12:46:58 +1000
Subject: [PATCH 02/15] Added vtk kwargs

---
 src/Visualization.jl | 41 +++++++++++++++++++++++++++++++----------
 1 file changed, 31 insertions(+), 10 deletions(-)

diff --git a/src/Visualization.jl b/src/Visualization.jl
index 6a8b0d8..9f79d34 100644
--- a/src/Visualization.jl
+++ b/src/Visualization.jl
@@ -33,7 +33,8 @@ end
 function Visualization.visualization_data(
   model::DistributedDiscreteModel{Dc},
   filebase::AbstractString;
-  labels=get_face_labeling(model)) where Dc
+  labels=get_face_labeling(model)
+) where Dc
 
   cell_gids = get_cell_gids(model)
   vd = map(local_views(model),partition(cell_gids),labels.labels) do model,gids,labels
@@ -57,7 +58,8 @@ function Visualization.visualization_data(
   order=-1,
   nsubcells=-1,
   celldata=nothing,
-  cellfields=nothing)
+  cellfields=nothing
+)
 
   trians    = trian.trians
   cell_gids = get_cell_gids(trian.model)
@@ -145,8 +147,13 @@ end
 
 function Visualization.write_vtk_file(
   parts::AbstractArray,
-  grid::AbstractArray{<:Grid}, filebase; celldata, nodaldata)
-  pvtk = Visualization.create_vtk_file(parts,grid,filebase;celldata=celldata,nodaldata=nodaldata)
+  grid::AbstractArray{<:Grid}, filebase; celldata, nodaldata,
+  compress=false,append=true,ascii=false,vtkversion=:default
+  )
+  pvtk = Visualization.create_vtk_file(
+    parts,grid,filebase;celldata=celldata,nodaldata=nodaldata,
+    compress=compress,append=append,ascii=ascii,vtkversion=vtkversion
+  )
   map(vtk_save,pvtk)
 end
 
@@ -154,33 +161,47 @@ function Visualization.create_vtk_file(
   parts::AbstractArray,
   grid::AbstractArray{<:Grid}, 
   filebase; 
-  celldata, nodaldata)
+  celldata, nodaldata,
+  compress=false,append=true,ascii=false,vtkversion=:default
+)
   nparts = length(parts)
   map(parts,grid,celldata,nodaldata) do part,g,c,n
     Visualization.create_pvtk_file(
       g,filebase;
       part=part,nparts=nparts,
-      celldata=c,nodaldata=n)
+      celldata=c,nodaldata=n,
+      compress=compress,append=append,ascii=ascii,vtkversion=vtkversion
+    )
   end
 end
 
 const DistributedModelOrTriangulation = Union{DistributedDiscreteModel,DistributedTriangulation}
 
-function Visualization.writevtk(arg::DistributedModelOrTriangulation,args...;kwargs...)
+function Visualization.writevtk(
+  arg::DistributedModelOrTriangulation,args...;
+  compress=false,append=true,ascii=false,vtkversion=:default,kwargs...
+)
   parts=get_parts(arg)
   map(visualization_data(arg,args...;kwargs...)) do visdata
     write_vtk_file(
-    parts,visdata.grid,visdata.filebase,celldata=visdata.celldata,nodaldata=visdata.nodaldata)
+      parts,visdata.grid,visdata.filebase,celldata=visdata.celldata,nodaldata=visdata.nodaldata,
+      compress=compress, append=append, ascii=ascii, vtkversion=vtkversion
+    )
   end
 end
 
-function Visualization.createvtk(arg::DistributedModelOrTriangulation,args...;kwargs...)
+function Visualization.createvtk(
+  arg::DistributedModelOrTriangulation,args...;
+  compress=false,append=true,ascii=false,vtkversion=:default,kwargs...
+)
   v = visualization_data(arg,args...;kwargs...)
   parts=get_parts(arg)
   @notimplementedif length(v) != 1
   visdata = first(v)
   Visualization.create_vtk_file(
-    parts,visdata.grid,visdata.filebase,celldata=visdata.celldata,nodaldata=visdata.nodaldata)
+    parts,visdata.grid,visdata.filebase,celldata=visdata.celldata,nodaldata=visdata.nodaldata,
+    compress=compress, append=append, ascii=ascii, vtkversion=vtkversion
+  )
 end
 
 struct DistributedPvd{T<:AbstractArray}

From b363db40a44537bd954221e19f764dfb1aa5ccf4 Mon Sep 17 00:00:00 2001
From: JordiManyer <jordi.manyer@monash.edu>
Date: Tue, 6 Aug 2024 23:38:17 +1000
Subject: [PATCH 03/15] Distributed ZeroMeanFESpace works

---
 src/FESpaces.jl | 118 ++++++++++++++++++++++++++++++++++++++----------
 1 file changed, 93 insertions(+), 25 deletions(-)

diff --git a/src/FESpaces.jl b/src/FESpaces.jl
index ec354bc..f4896c0 100644
--- a/src/FESpaces.jl
+++ b/src/FESpaces.jl
@@ -252,7 +252,7 @@ end
 # FEFunction related
 """
 """
-struct DistributedFEFunctionData{T<:AbstractVector} <:GridapType
+struct DistributedFEFunctionData{T<:AbstractVector} <: GridapType
   free_values::T
 end
 
@@ -267,20 +267,24 @@ end
 # Single field related
 """
 """
-struct DistributedSingleFieldFESpace{A,B,C,D} <: DistributedFESpace
+struct DistributedSingleFieldFESpace{A,B,C,D,E} <: DistributedFESpace
   spaces::A
   gids::B
   trian::C
   vector_type::Type{D}
+  metadata::E
   function DistributedSingleFieldFESpace(
     spaces::AbstractArray{<:SingleFieldFESpace},
     gids::PRange,
     trian::DistributedTriangulation,
-    vector_type::Type{D}) where D
+    vector_type::Type{D},
+    metadata = nothing
+  ) where D
     A = typeof(spaces)
     B = typeof(gids)
     C = typeof(trian)
-    new{A,B,C,D}(spaces,gids,trian,vector_type)
+    E = typeof(metadata)
+    new{A,B,C,D,E}(spaces,gids,trian,vector_type,metadata)
   end
 end
 
@@ -375,35 +379,35 @@ end
 
 function FESpaces.TrialFESpace(f::DistributedSingleFieldFESpace)
   spaces = map(TrialFESpace,f.spaces)
-  DistributedSingleFieldFESpace(spaces,f.gids,f.trian,f.vector_type)
+  DistributedSingleFieldFESpace(spaces,f.gids,f.trian,f.vector_type,f.metadata)
 end
 
 function FESpaces.TrialFESpace(f::DistributedSingleFieldFESpace,fun)
   spaces = map(f.spaces) do s
     TrialFESpace(s,fun)
   end
-  DistributedSingleFieldFESpace(spaces,f.gids,f.trian,f.vector_type)
+  DistributedSingleFieldFESpace(spaces,f.gids,f.trian,f.vector_type,f.metadata)
 end
 
 function FESpaces.TrialFESpace(fun,f::DistributedSingleFieldFESpace)
   spaces = map(f.spaces) do s
     TrialFESpace(fun,s)
   end
-  DistributedSingleFieldFESpace(spaces,f.gids,f.trian,f.vector_type)
+  DistributedSingleFieldFESpace(spaces,f.gids,f.trian,f.vector_type,f.metadata)
 end
 
 function FESpaces.TrialFESpace!(f::DistributedSingleFieldFESpace,fun)
   spaces = map(f.spaces) do s
     TrialFESpace!(s,fun)
   end
-  DistributedSingleFieldFESpace(spaces,f.gids,f.trian,f.vector_type)
+  DistributedSingleFieldFESpace(spaces,f.gids,f.trian,f.vector_type,f.metadata)
 end
 
 function FESpaces.HomogeneousTrialFESpace(f::DistributedSingleFieldFESpace)
   spaces = map(f.spaces) do s
     HomogeneousTrialFESpace(s)
   end
-  DistributedSingleFieldFESpace(spaces,f.gids,f.trian,f.vector_type)
+  DistributedSingleFieldFESpace(spaces,f.gids,f.trian,f.vector_type,f.metadata)
 end
 
 function generate_gids(
@@ -704,30 +708,94 @@ function FESpaces.SparseMatrixAssembler(
 end
 
 # ZeroMean FESpace
-
-#const DistributedZeroMeanFESpace = DistributedSingleFieldFESpace{}
-
-struct ZeroMeanCache{A}
-  vol_i::A
-  vol::Float64
+struct DistributedZeroMeanCache{A,B}
+  dvol::A
+  vol::B
 end
 
-function FESpaces.ZeroMeanFESpace(space::DistributedFESpace)#,dΩ::DistributedMeasure)
+const DistributedZeroMeanFESpace{A,B,C,D,E,F} = DistributedSingleFieldFESpace{A,B,C,D,DistributedZeroMeanCache{E,F}}
+
+function FESpaces.ZeroMeanFESpace(space::DistributedFESpace)
   gids = get_free_dof_ids(space)
-  ranks = get_parts(space)
-  spaces = map(ranks,partition(gids),local_views(space)) do r, gids, lspace
-    fix_constant = isone(r) # Only main processor fixes the constant
-    dof_to_fix = Int(first(own_to_local(gids))) # Make sure it's an owned DoF
-    FESpaceWithConstantFixed(lspace,fix_constant,dof_to_fix)
+
+  # We always fix the first gid
+  lid_to_fix = map(partition(gids)) do gids
+    Int(global_to_local(gids)[1]) # returns 0 if not found in the processor
   end
 
-  #vol_i = assemble_vector(v->∫(v)*dΩ,space)
-  #vol = sum(vol_i)
-  #metadata = ZeroMeanCache(vol_i,vol)
+  # Create local spaces
+  spaces = map(local_views(space),lid_to_fix) do lspace, lid_to_fix
+    fix_constant = !iszero(lid_to_fix)
+    FESpaceWithConstantFixed(lspace,fix_constant,lid_to_fix)
+  end
+
+  # Setup volume integration
+  order = 2 # TODO: make it a parameter
+  _vol, _dvol = map(local_views(space),partition(gids)) do lspace, gids
+    _trian = get_triangulation(lspace)
+    trian = remove_ghost_cells(_trian,gids)
+    dΩ = Measure(trian,2*order)
+    dvol = assemble_vector(v -> ∫(v)dΩ, lspace)
+    vol  = sum(dvol)
+    return vol, dvol
+  end |> tuple_of_arrays
+  vol  = reduce(+,_vol,init=zero(eltype(vol)))
+  dvol = PVector(_dvol,partition(gids))
+  metadata = DistributedZeroMeanCache(dvol,vol)
 
+  # Create the new global FESpace
   trian = get_triangulation(space)
   model = get_background_model(trian)
   gids =  generate_gids(model,spaces)
   vector_type = _find_vector_type(spaces,gids)
-  DistributedSingleFieldFESpace(spaces,gids,trian,vector_type)
+  return DistributedSingleFieldFESpace(spaces,gids,trian,vector_type,metadata)
+end
+
+function FESpaces.FEFunction(
+  f::DistributedZeroMeanFESpace,
+  free_values::AbstractVector,
+  isconsistent=false
+  )
+  dirichlet_values = get_dirichlet_values(f)
+  FEFunction(f,free_values,dirichlet_values,isconsistent)
+end
+
+function FESpaces.FEFunction(
+  f::DistributedZeroMeanFESpace,
+  free_values::AbstractVector,
+  dirichlet_values::AbstractArray{<:AbstractVector},
+  isconsistent=false
+)
+  free_values = change_ghost(free_values,f.gids,is_consistent=isconsistent,make_consistent=true)
+  
+  c = _compute_new_distributed_fixedval(
+    f,free_values,dirichlet_values
+  )
+  fv = free_values .+ c
+  dv = map(dirichlet_values) do dv
+    dv .+ c
+  end
+  
+  fields = map(FEFunction,f.spaces,partition(fv),dv)
+  trian = get_triangulation(f)
+  metadata = DistributedFEFunctionData(free_values)
+  DistributedCellField(fields,trian,metadata)
+end
+
+function _compute_new_distributed_fixedval(
+  f::DistributedZeroMeanFESpace,fv,dv
+)
+  dvol = f.metadata.dvol
+  vol  = f.metadata.vol
+  
+  c_i = map(local_views(f),partition(fv),dv,partition(dvol)) do space,fv,dv,dvol
+    if isa(FESpaces.ConstantApproach(space),FESpaces.FixConstant)
+      lid_to_fix = space.dof_to_fix
+      c = FESpaces._compute_new_fixedval(fv,dv,dvol,vol,lid_to_fix)
+    else
+      c = - dot(fv,dvol)/vol
+    end
+  end
+  c = reduce(+,c_i,init=zero(eltype(c_i)))
+  return c
 end

From 46e40fa3b40d80960406d2370c1e86d4ac086453 Mon Sep 17 00:00:00 2001
From: JordiManyer <jordi.manyer@monash.edu>
Date: Wed, 7 Aug 2024 08:56:59 +1000
Subject: [PATCH 04/15] Merged ZeroMeanFESpaces into FESpace factories

---
 src/FESpaces.jl | 108 ++++++++++++++++++++++++++++++++++--------------
 src/Geometry.jl |   6 +++
 2 files changed, 83 insertions(+), 31 deletions(-)

diff --git a/src/FESpaces.jl b/src/FESpaces.jl
index f4896c0..4c4b85f 100644
--- a/src/FESpaces.jl
+++ b/src/FESpaces.jl
@@ -412,13 +412,24 @@ end
 
 function generate_gids(
   model::DistributedDiscreteModel{Dc},
-  spaces::AbstractArray{<:SingleFieldFESpace}) where Dc
+  spaces::AbstractArray{<:SingleFieldFESpace}
+) where Dc
   cell_to_ldofs = map(get_cell_dof_ids,spaces)
   nldofs = map(num_free_dofs,spaces)
   cell_gids = get_cell_gids(model)
   generate_gids(cell_gids,cell_to_ldofs,nldofs)
 end
 
+function generate_gids(
+  trian::DistributedTriangulation{Dc},
+  spaces::AbstractArray{<:SingleFieldFESpace}
+) where Dc
+  cell_to_ldofs = map(get_cell_dof_ids,spaces)
+  nldofs = map(num_free_dofs,spaces)
+  cell_gids = generate_cell_gids(trian)
+  generate_gids(cell_gids,cell_to_ldofs,nldofs)
+end
+
 function FESpaces.interpolate(u,f::DistributedSingleFieldFESpace)
   free_values = zero_free_values(f)
   interpolate!(u,free_values,f)
@@ -484,27 +495,30 @@ end
 
 # Factories
 
-function FESpaces.FESpace(model::DistributedDiscreteModel,reffe;split_own_and_ghost=false,kwargs...)
+function FESpaces.FESpace(
+  model::DistributedDiscreteModel,reffe;split_own_and_ghost=false,constraint=nothing,kwargs...
+)
   spaces = map(local_views(model)) do m
     FESpace(m,reffe;kwargs...)
   end
   gids =  generate_gids(model,spaces)
   trian = DistributedTriangulation(map(get_triangulation,spaces),model)
   vector_type = _find_vector_type(spaces,gids;split_own_and_ghost=split_own_and_ghost)
-  DistributedSingleFieldFESpace(spaces,gids,trian,vector_type)
+  space = DistributedSingleFieldFESpace(spaces,gids,trian,vector_type)
+  return _add_distributed_constraint(space,reffe,constraint)
 end
 
-function FESpaces.FESpace(_trian::DistributedTriangulation,reffe;split_own_and_ghost=false,kwargs...)
+function FESpaces.FESpace(
+  _trian::DistributedTriangulation,reffe;split_own_and_ghost=false,constraint=nothing,kwargs...
+)
   trian = add_ghost_cells(_trian)
-  trian_gids = generate_cell_gids(trian)
-  spaces = map(trian.trians) do t
+  spaces = map(local_views(trian)) do t
     FESpace(t,reffe;kwargs...)
   end
-  cell_to_ldofs = map(get_cell_dof_ids,spaces)
-  nldofs = map(num_free_dofs,spaces)
-  gids = generate_gids(trian_gids,cell_to_ldofs,nldofs)
+  gids = generate_gids(trian,spaces)
   vector_type = _find_vector_type(spaces,gids;split_own_and_ghost=split_own_and_ghost)
-  DistributedSingleFieldFESpace(spaces,gids,trian,vector_type)
+  space = DistributedSingleFieldFESpace(spaces,gids,trian,vector_type)
+  return _add_distributed_constraint(space,reffe,constraint)
 end
 
 function _find_vector_type(spaces,gids;split_own_and_ghost=false)
@@ -522,6 +536,40 @@ function _find_vector_type(spaces,gids;split_own_and_ghost=false)
   return vector_type
 end
 
+# TODO: We would like to avoid this, but I cannot extract the maximal order 
+#       from the space itself...
+function _add_distributed_constraint(
+  F::DistributedFESpace,reffe::ReferenceFE,constraint
+)
+  order = get_order(reffe)
+  _add_distributed_constraint(F,order,constraint)
+end
+
+function _add_distributed_constraint(
+  F::DistributedFESpace,reffe::Tuple{<:ReferenceFEName,Any,Any},constraint
+)
+  args = reffe[2]
+  order = maximum(args[2])
+  _add_distributed_constraint(F,order,constraint)
+end
+
+function _add_distributed_constraint(F::DistributedFESpace,order::Integer,constraint)
+  if isnothing(constraint)
+    V = F
+  elseif constraint == :zeromean
+    _trian = get_triangulation(F)
+    trian = remove_ghost_cells(_trian,get_free_dof_ids(F))
+    dΩ = Measure(trian,order)
+    V = ZeroMeanFESpace(F,dΩ)
+  else
+    @unreachable """\n
+    The passed option constraint=$constraint is not valid.
+    Valid values for constraint: nothing, :zeromean
+    """
+  end
+  V
+end
+
 # Assembly
 
 function FESpaces.collect_cell_matrix(
@@ -659,7 +707,8 @@ function local_assembly_strategy(::FullyAssembledRows,rows,cols)
     identity,
     identity,
     row->rows_local_to_ghost[row]==0,
-    col->true)
+    col->true
+  )
 end
 
 # Assembler high level constructors
@@ -668,17 +717,18 @@ function FESpaces.SparseMatrixAssembler(
   local_vec_type,
   rows::PRange,
   cols::PRange,
-  par_strategy=SubAssembledRows())
-
+  par_strategy=SubAssembledRows()
+)
   assems = map(partition(rows),partition(cols)) do rows,cols
     local_strategy = local_assembly_strategy(par_strategy,rows,cols)
-    FESpaces.GenericSparseMatrixAssembler(SparseMatrixBuilder(local_mat_type),
-                                          ArrayBuilder(local_vec_type),
-                                          Base.OneTo(length(rows)),
-                                          Base.OneTo(length(cols)),
-                                          local_strategy)
+    FESpaces.GenericSparseMatrixAssembler(
+      SparseMatrixBuilder(local_mat_type),
+      ArrayBuilder(local_vec_type),
+      Base.OneTo(length(rows)),
+      Base.OneTo(length(cols)),
+      local_strategy
+    )
   end
-
   mat_builder = PSparseMatrixBuilderCOO(local_mat_type,par_strategy)
   vec_builder = PVectorBuilder(local_vec_type,par_strategy)
   return DistributedSparseMatrixAssembler(par_strategy,assems,mat_builder,vec_builder,rows,cols)
@@ -689,8 +739,8 @@ function FESpaces.SparseMatrixAssembler(
   local_vec_type,
   trial::DistributedFESpace,
   test::DistributedFESpace,
-  par_strategy=SubAssembledRows())
-
+  par_strategy=SubAssembledRows()
+)
   rows = get_free_dof_ids(test)
   cols = get_free_dof_ids(trial)
   SparseMatrixAssembler(local_mat_type,local_vec_type,rows,cols,par_strategy)
@@ -699,8 +749,8 @@ end
 function FESpaces.SparseMatrixAssembler(
   trial::DistributedFESpace,
   test::DistributedFESpace,
-  par_strategy=SubAssembledRows())
-
+  par_strategy=SubAssembledRows()
+)
   Tv = PartitionedArrays.getany(map(get_vector_type,local_views(trial)))
   T  = eltype(Tv)
   Tm = SparseMatrixCSC{T,Int}
@@ -715,7 +765,7 @@ end
 
 const DistributedZeroMeanFESpace{A,B,C,D,E,F} = DistributedSingleFieldFESpace{A,B,C,D,DistributedZeroMeanCache{E,F}}
 
-function FESpaces.ZeroMeanFESpace(space::DistributedFESpace)
+function FESpaces.ZeroMeanFESpace(space::DistributedFESpace,dΩ::DistributedMeasure)
   gids = get_free_dof_ids(space)
 
   # We always fix the first gid
@@ -730,11 +780,7 @@ function FESpaces.ZeroMeanFESpace(space::DistributedFESpace)
   end
 
   # Setup volume integration
-  order = 2 # TODO: make it a parameter
-  _vol, _dvol = map(local_views(space),partition(gids)) do lspace, gids
-    _trian = get_triangulation(lspace)
-    trian = remove_ghost_cells(_trian,gids)
-    dΩ = Measure(trian,2*order)
+  _vol, _dvol = map(local_views(space),local_views(dΩ)) do lspace, dΩ
     dvol = assemble_vector(v -> ∫(v)dΩ, lspace)
     vol  = sum(dvol)
     return vol, dvol
@@ -755,8 +801,8 @@ function FESpaces.FEFunction(
   f::DistributedZeroMeanFESpace,
   free_values::AbstractVector,
   isconsistent=false
-  )
-  dirichlet_values = get_dirichlet_values(f)
+)
+  dirichlet_values = get_dirichlet_dof_values(f)
   FEFunction(f,free_values,dirichlet_values,isconsistent)
 end
 
diff --git a/src/Geometry.jl b/src/Geometry.jl
index 3b9b4dc..c0152a6 100644
--- a/src/Geometry.jl
+++ b/src/Geometry.jl
@@ -597,6 +597,12 @@ function filter_cells_when_needed(
   remove_ghost_cells(trian,cell_gids)
 end
 
+function remove_ghost_cells(trian::DistributedTriangulation,gids)
+  trians = map(remove_ghost_cells,local_views(trian),partition(gids))
+  model  = get_background_model(trian)
+  return DistributedTriangulation(trians,model)
+end
+
 function remove_ghost_cells(trian::Triangulation,gids)
   model = get_background_model(trian)
   Dt    = num_cell_dims(trian)

From 18b000094b2b083e02a2561dab380de7f964a8fe Mon Sep 17 00:00:00 2001
From: JordiManyer <jordi.manyer@monash.edu>
Date: Thu, 8 Aug 2024 09:47:56 +1000
Subject: [PATCH 05/15] Better multifield implementation for interpolate and
 FEFunction, that avoids repeating work and works better with ZeroMeanFESpaces

---
 src/FESpaces.jl   |  57 +++++++++++-----
 src/MultiField.jl | 168 +++++++++++++++++++++-------------------------
 2 files changed, 117 insertions(+), 108 deletions(-)

diff --git a/src/FESpaces.jl b/src/FESpaces.jl
index 4c4b85f..1d79d75 100644
--- a/src/FESpaces.jl
+++ b/src/FESpaces.jl
@@ -558,7 +558,8 @@ function _add_distributed_constraint(F::DistributedFESpace,order::Integer,constr
     V = F
   elseif constraint == :zeromean
     _trian = get_triangulation(F)
-    trian = remove_ghost_cells(_trian,get_free_dof_ids(F))
+    model = get_background_model(_trian)
+    trian = remove_ghost_cells(_trian,get_cell_gids(model))
     dΩ = Measure(trian,order)
     V = ZeroMeanFESpace(F,dΩ)
   else
@@ -765,12 +766,14 @@ end
 
 const DistributedZeroMeanFESpace{A,B,C,D,E,F} = DistributedSingleFieldFESpace{A,B,C,D,DistributedZeroMeanCache{E,F}}
 
-function FESpaces.ZeroMeanFESpace(space::DistributedFESpace,dΩ::DistributedMeasure)
+function FESpaces.FESpaceWithConstantFixed(
+  space::DistributedSingleFieldFESpace, 
+  gid_to_fix::Int = num_free_dofs(space)
+)
+  # Find the gid within the processors
   gids = get_free_dof_ids(space)
-
-  # We always fix the first gid
   lid_to_fix = map(partition(gids)) do gids
-    Int(global_to_local(gids)[1]) # returns 0 if not found in the processor
+    Int(global_to_local(gids)[gid_to_fix]) # returns 0 if not found in the processor
   end
 
   # Create local spaces
@@ -779,22 +782,29 @@ function FESpaces.ZeroMeanFESpace(space::DistributedFESpace,dΩ::DistributedMeas
     FESpaceWithConstantFixed(lspace,fix_constant,lid_to_fix)
   end
 
+  trian = get_triangulation(space)
+  model = get_background_model(trian)
+  gids  =  generate_gids(model,spaces)
+  vector_type = _find_vector_type(spaces,gids)
+  return DistributedSingleFieldFESpace(spaces,gids,trian,vector_type)
+end
+
+function FESpaces.ZeroMeanFESpace(space::DistributedSingleFieldFESpace,dΩ::DistributedMeasure)
+  # Create underlying space
+  _space = FESpaceWithConstantFixed(space,num_free_dofs(space))
+
   # Setup volume integration
-  _vol, _dvol = map(local_views(space),local_views(dΩ)) do lspace, dΩ
+  _vol, dvol = map(local_views(space),local_views(dΩ)) do lspace, dΩ
     dvol = assemble_vector(v -> ∫(v)dΩ, lspace)
     vol  = sum(dvol)
     return vol, dvol
   end |> tuple_of_arrays
   vol  = reduce(+,_vol,init=zero(eltype(vol)))
-  dvol = PVector(_dvol,partition(gids))
   metadata = DistributedZeroMeanCache(dvol,vol)
 
-  # Create the new global FESpace
-  trian = get_triangulation(space)
-  model = get_background_model(trian)
-  gids =  generate_gids(model,spaces)
-  vector_type = _find_vector_type(spaces,gids)
-  return DistributedSingleFieldFESpace(spaces,gids,trian,vector_type,metadata)
+  return DistributedSingleFieldFESpace(
+    _space.spaces,_space.gids,_space.trian,_space.vector_type,metadata
+  )
 end
 
 function FESpaces.FEFunction(
@@ -817,30 +827,45 @@ function FESpaces.FEFunction(
   c = _compute_new_distributed_fixedval(
     f,free_values,dirichlet_values
   )
-  fv = free_values .+ c
+  fv = free_values .+ c # TODO: Do we need to copy, or can we just modify? 
   dv = map(dirichlet_values) do dv
     dv .+ c
   end
   
   fields = map(FEFunction,f.spaces,partition(fv),dv)
   trian = get_triangulation(f)
-  metadata = DistributedFEFunctionData(free_values)
+  metadata = DistributedFEFunctionData(fv)
   DistributedCellField(fields,trian,metadata)
 end
 
+# This is required, otherwise we end up calling `FEFunction` with a fixed value of zero, 
+# which does not properly interpolate the function provided. 
+# With this change, we are interpolating in the unconstrained space and then
+# substracting the mean.
+function FESpaces.interpolate!(u,free_values::AbstractVector,f::DistributedZeroMeanFESpace)
+  dirichlet_values = get_dirichlet_dof_values(f)
+  interpolate_everywhere!(u,free_values,dirichlet_values,f)
+end
+function FESpaces.interpolate!(u::DistributedCellField,free_values::AbstractVector,f::DistributedZeroMeanFESpace)
+  dirichlet_values = get_dirichlet_dof_values(f)
+  interpolate_everywhere!(u,free_values,dirichlet_values,f)
+end
+
 function _compute_new_distributed_fixedval(
   f::DistributedZeroMeanFESpace,fv,dv
 )
   dvol = f.metadata.dvol
   vol  = f.metadata.vol
   
-  c_i = map(local_views(f),partition(fv),dv,partition(dvol)) do space,fv,dv,dvol
+  c_i = map(local_views(f),partition(fv),dv,dvol) do space,fv,dv,dvol
     if isa(FESpaces.ConstantApproach(space),FESpaces.FixConstant)
       lid_to_fix = space.dof_to_fix
       c = FESpaces._compute_new_fixedval(fv,dv,dvol,vol,lid_to_fix)
     else
       c = - dot(fv,dvol)/vol
     end
+    println("c = $c")
+    c
   end
   c = reduce(+,c_i,init=zero(eltype(c_i)))
   return c
diff --git a/src/MultiField.jl b/src/MultiField.jl
index e8927d9..e0ffbbb 100644
--- a/src/MultiField.jl
+++ b/src/MultiField.jl
@@ -103,10 +103,10 @@ end
 function MultiField.restrict_to_field(
   f::DistributedMultiFieldFESpace,free_values::AbstractVector,field::Integer
 )
-  values = map(local_views(f),partition(free_values)) do u,x
-    restrict_to_field(u,x,field)
+  values = map(local_views(f),partition(free_values)) do u,fv
+    restrict_to_field(u,fv,field)
   end
-  gids = f.field_fe_space[field].gids
+  gids = get_free_dof_ids(f[field])
   PVector(values,partition(gids))
 end
 
@@ -118,26 +118,22 @@ function FESpaces.get_dirichlet_dof_values(f::DistributedMultiFieldFESpace)
   return map(get_dirichlet_dof_values,f.field_fe_space)
 end
 
-function FESpaces.FEFunction(f::DistributedMultiFieldFESpace,x::AbstractVector,isconsistent=false)
-  free_values  = change_ghost(x,f.gids;is_consistent=isconsistent,make_consistent=true)
-  part_fe_fun  = map(FEFunction,f.part_fe_space,partition(free_values))
-  field_fe_fun = DistributedSingleFieldFEFunction[]
-  for i in 1:num_fields(f)
-    free_values_i = restrict_to_field(f,free_values,i)
-    fe_space_i = f.field_fe_space[i]
-    fe_fun_i = FEFunction(fe_space_i,free_values_i,true)
-    push!(field_fe_fun,fe_fun_i)
-  end
-  DistributedMultiFieldFEFunction(field_fe_fun,part_fe_fun,free_values)
+function FESpaces.FEFunction(
+  f::DistributedMultiFieldFESpace,free_values::AbstractVector,isconsistent=false
+)
+  dirichlet_values = get_dirichlet_dof_values(f)
+  return FEFunction(f,free_values,dirichlet_values,isconsistent)
 end
 
 function FESpaces.FEFunction(
-  f::DistributedMultiFieldFESpace,x::AbstractVector,
-  dirichlet_values::AbstractArray{<:AbstractVector},isconsistent=false
+  f::DistributedMultiFieldFESpace,
+  _free_values::AbstractVector,
+  dirichlet_values::AbstractArray{<:AbstractVector},
+  isconsistent=false
 )
-  free_values  = GridapDistributed.change_ghost(x,f.gids;is_consistent=isconsistent,make_consistent=true)
-  part_dirvals = to_parray_of_arrays(dirichlet_values)
-  part_fe_fun  = map(FEFunction,f.part_fe_space,partition(free_values),part_dirvals)
+  free_values = change_ghost(_free_values,f.gids;is_consistent=isconsistent,make_consistent=true)
+
+  # Create distributed single field functions
   field_fe_fun = DistributedSingleFieldFEFunction[]
   for i in 1:num_fields(f)
     free_values_i = restrict_to_field(f,free_values,i)
@@ -146,12 +142,24 @@ function FESpaces.FEFunction(
     fe_fun_i = FEFunction(fe_space_i,free_values_i,dirichlet_values_i,true)
     push!(field_fe_fun,fe_fun_i)
   end
+
+  # Retrieve the local multifield views
+  part_sf_fe_funs = map(local_views,field_fe_fun)
+  part_fe_fun = map(local_views(f),partition(free_values),part_sf_fe_funs...) do space,fv,part_sf_fe_funs...
+    MultiFieldFEFunction(fv,space,[part_sf_fe_funs...])
+  end
+
   DistributedMultiFieldFEFunction(field_fe_fun,part_fe_fun,free_values)
 end
 
-function FESpaces.EvaluationFunction(f::DistributedMultiFieldFESpace,x::AbstractVector,isconsistent=false)
-  free_values  = change_ghost(x,f.gids;is_consistent=isconsistent,make_consistent=true)
-  part_fe_fun  = map(EvaluationFunction,f.part_fe_space,partition(free_values))
+function FESpaces.EvaluationFunction(
+  f::DistributedMultiFieldFESpace,
+  _free_values::AbstractVector,
+  isconsistent=false
+)
+  free_values = change_ghost(_free_values,f.gids;is_consistent=isconsistent,make_consistent=true)
+  
+  # Create distributed single field functions
   field_fe_fun = DistributedSingleFieldFEFunction[]
   for i in 1:num_fields(f)
     free_values_i = restrict_to_field(f,free_values,i)
@@ -159,99 +167,74 @@ function FESpaces.EvaluationFunction(f::DistributedMultiFieldFESpace,x::Abstract
     fe_fun_i = EvaluationFunction(fe_space_i,free_values_i)
     push!(field_fe_fun,fe_fun_i)
   end
+
+  # Retrieve the local multifield views
+  part_sf_fe_funs = map(local_views,field_fe_fun)
+  part_fe_fun = map(local_views(f),partition(free_values),part_sf_fe_funs...) do space,fv,part_sf_fe_funs...
+    MultiFieldFEFunction(fv,space,[part_sf_fe_funs...])
+  end
+
   DistributedMultiFieldFEFunction(field_fe_fun,part_fe_fun,free_values)
 end
 
-function FESpaces.interpolate(objects,fe::DistributedMultiFieldFESpace)
-  free_values = zero_free_values(fe)
-  interpolate!(objects,free_values,fe)
+function FESpaces.interpolate(objects,space::DistributedMultiFieldFESpace)
+  free_values = zero_free_values(space)
+  interpolate!(objects,free_values,space)
 end
 
-function FESpaces.interpolate!(objects,free_values::AbstractVector,fe::DistributedMultiFieldFESpace)
-  part_fe_fun = map(partition(free_values),local_views(fe)) do x,f
-    interpolate!(objects,x,f)
-  end
+function FESpaces.interpolate!(objects,free_values::AbstractVector,space::DistributedMultiFieldFESpace)
+  msg = "free_values and FESpace have incompatible index partitions."
+  @check partition(axes(free_values,1)) === partition(space.gids) msg
+
+  # Interpolate each field
   field_fe_fun = DistributedSingleFieldFEFunction[]
-  for i in 1:num_fields(fe)
-    free_values_i = restrict_to_field(fe,free_values,i)
-    fe_space_i = fe.field_fe_space[i]
-    fe_fun_i = FEFunction(fe_space_i,free_values_i)
+  for i in 1:num_fields(space)
+    free_values_i = restrict_to_field(space,free_values,i)
+    fe_space_i = space.field_fe_space[i]
+    fe_fun_i = interpolate!(objects[i], free_values_i, fe_space_i)
     push!(field_fe_fun,fe_fun_i)
   end
-  DistributedMultiFieldFEFunction(field_fe_fun,part_fe_fun,free_values)
-end
 
-function Gridap.FESpaces.interpolate!(
-  objects::Union{<:DistributedMultiFieldCellField,<:DistributedCellField},
-  free_values::AbstractVector,
-  fe::DistributedMultiFieldFESpace
-)
-  part_fe_fun = map(local_views(objects),partition(free_values),local_views(fe)) do objects,x,f
-    interpolate!(objects,x,f)
+  # Retrieve the local multifield views
+  part_sf_fe_funs = map(local_views,field_fe_fun)
+  part_fe_fun = map(local_views(space),partition(free_values),part_sf_fe_funs...) do space,fv,part_sf_fe_funs...
+    MultiFieldFEFunction(fv,space,[part_sf_fe_funs...])
   end
-  field_fe_fun = GridapDistributed.DistributedSingleFieldFEFunction[]
-  for i in 1:num_fields(fe)
-    free_values_i = Gridap.MultiField.restrict_to_field(fe,free_values,i)
-    fe_space_i = fe.field_fe_space[i]
-    fe_fun_i = FEFunction(fe_space_i,free_values_i)
-    push!(field_fe_fun,fe_fun_i)
-  end
-  GridapDistributed.DistributedMultiFieldFEFunction(field_fe_fun,part_fe_fun,free_values)
+
+  DistributedMultiFieldFEFunction(field_fe_fun,part_fe_fun,free_values)
 end
 
 function FESpaces.interpolate_everywhere(objects,fe::DistributedMultiFieldFESpace)
   free_values = zero_free_values(fe)
-  part_fe_fun = map(partition(free_values),local_views(fe)) do x,f
-    interpolate!(objects,x,f)
-  end
-  field_fe_fun = DistributedSingleFieldFEFunction[]
-  for i in 1:num_fields(fe)
-    free_values_i = restrict_to_field(fe,free_values,i)
-    fe_space_i = fe.field_fe_space[i]
-    dirichlet_values_i = zero_dirichlet_values(fe_space_i)
-    fe_fun_i = interpolate_everywhere!(objects[i], free_values_i,dirichlet_values_i,fe_space_i)
-    push!(field_fe_fun,fe_fun_i)
-  end
-  DistributedMultiFieldFEFunction(field_fe_fun,part_fe_fun,free_values)
+  dirichlet_values = zero_dirichlet_values(fe)
+  return interpolate_everywhere!(objects,free_values,dirichlet_values,fe)
 end
 
 function FESpaces.interpolate_everywhere!(
-  objects,free_values::AbstractVector,
+  objects,
+  free_values::AbstractVector,
   dirichlet_values::Vector{AbstractArray{<:AbstractVector}},
-  fe::DistributedMultiFieldFESpace)
-  msg = "free_values and fe have incompatible index partitions."
-  @check partition(axes(free_values,1)) === partition(fe.gids) msg
+  space::DistributedMultiFieldFESpace
+)
+  msg = "free_values and FESpace have incompatible index partitions."
+  @check partition(axes(free_values,1)) === partition(space.gids) msg
   
-  part_fe_fun = map(partition(free_values),local_views(fe)) do x,f
-    interpolate!(objects,x,f)
-  end
+  # Interpolate each field
   field_fe_fun = DistributedSingleFieldFEFunction[]
-  for i in 1:num_fields(fe)
-    free_values_i = restrict_to_field(fe,free_values,i)
+  for i in 1:num_fields(space)
+    free_values_i = restrict_to_field(space,free_values,i)
     dirichlet_values_i = dirichlet_values[i]
-    fe_space_i = fe.field_fe_space[i]
-    fe_fun_i = interpolate_everywhere!(objects[i], free_values_i,dirichlet_values_i,fe_space_i)
+    fe_space_i = space.field_fe_space[i]
+    fe_fun_i = interpolate_everywhere!(objects[i], free_values_i, dirichlet_values_i,fe_space_i)
     push!(field_fe_fun,fe_fun_i)
   end
-  DistributedMultiFieldFEFunction(field_fe_fun,part_fe_fun,free_values)
-end
 
-function FESpaces.interpolate_everywhere(
-  objects::Vector{<:DistributedCellField},fe::DistributedMultiFieldFESpace)
-  local_objects = map(local_views,objects)
-  local_spaces = local_views(fe)
-  part_fe_fun = map(local_spaces,local_objects...) do f,o...
-    interpolate_everywhere(o,f)
-  end
-  free_values = zero_free_values(fe)
-  field_fe_fun = DistributedSingleFieldFEFunction[]
-  for i in 1:num_fields(fe)
-    free_values_i = restrict_to_field(fe,free_values,i)
-    fe_space_i = fe.field_fe_space[i]
-    dirichlet_values_i = get_dirichlet_dof_values(fe_space_i)
-    fe_fun_i = interpolate_everywhere!(objects[i], free_values_i,dirichlet_values_i,fe_space_i)
-    push!(field_fe_fun,fe_fun_i)
+  # Retrieve the local multifield views
+  part_sf_fe_funs = map(local_views,field_fe_fun)
+  part_fe_fun = map(local_views(space),partition(free_values),part_sf_fe_funs...) do space,fv,part_sf_fe_funs...
+    MultiFieldFEFunction(fv,space,[part_sf_fe_funs...])
   end
+
   DistributedMultiFieldFEFunction(field_fe_fun,part_fe_fun,free_values)
 end
 
@@ -278,7 +261,7 @@ const DistributedMultiFieldFEBasis{A} = DistributedMultiFieldCellField{A,<:Abstr
 function FESpaces.get_fe_basis(f::DistributedMultiFieldFESpace)
   part_mbasis = map(get_fe_basis,f.part_fe_space)
   field_fe_basis = map(1:num_fields(f)) do i
-    space_i  = f.field_fe_space[i]
+    space_i = f.field_fe_space[i]
     basis_i = map(b->b[i],part_mbasis)
     DistributedCellField(basis_i,get_triangulation(space_i))
   end
@@ -288,7 +271,7 @@ end
 function FESpaces.get_trial_fe_basis(f::DistributedMultiFieldFESpace)
   part_mbasis = map(get_trial_fe_basis,f.part_fe_space)
   field_fe_basis = map(1:num_fields(f)) do i
-    space_i  = f.field_fe_space[i]
+    space_i = f.field_fe_space[i]
     basis_i = map(b->b[i],part_mbasis)
     DistributedCellField(basis_i,get_triangulation(space_i))
   end
@@ -328,6 +311,7 @@ function generate_multi_field_gids(
   end
   f_frange = map(get_free_dof_ids,f_dspace)
   gids = generate_multi_field_gids(f_p_flid_lid,f_frange)
+  return gids
 end
 
 function generate_multi_field_gids(

From 64363809e321852afdabddbeb37fe1493680ae49 Mon Sep 17 00:00:00 2001
From: JordiManyer <jordi.manyer@monash.edu>
Date: Thu, 8 Aug 2024 23:29:11 +1000
Subject: [PATCH 06/15] Added tests

---
 src/FESpaces.jl                          |  1 -
 test/TestApp/src/TestApp.jl              |  1 +
 test/ZeroMeanFESpacesTests.jl            | 82 ++++++++++++++++++++++++
 test/mpi/runtests_np4_body.jl            |  3 +
 test/sequential/ZeroMeanFESpacesTests.jl | 10 +++
 test/sequential/runtests.jl              |  2 +
 6 files changed, 98 insertions(+), 1 deletion(-)
 create mode 100644 test/ZeroMeanFESpacesTests.jl
 create mode 100644 test/sequential/ZeroMeanFESpacesTests.jl

diff --git a/src/FESpaces.jl b/src/FESpaces.jl
index 1d79d75..c2b0947 100644
--- a/src/FESpaces.jl
+++ b/src/FESpaces.jl
@@ -864,7 +864,6 @@ function _compute_new_distributed_fixedval(
     else
       c = - dot(fv,dvol)/vol
     end
-    println("c = $c")
     c
   end
   c = reduce(+,c_i,init=zero(eltype(c_i)))
diff --git a/test/TestApp/src/TestApp.jl b/test/TestApp/src/TestApp.jl
index 8f84789..dce204b 100644
--- a/test/TestApp/src/TestApp.jl
+++ b/test/TestApp/src/TestApp.jl
@@ -9,6 +9,7 @@ module TestApp
   include("../../SurfaceCouplingTests.jl")
   include("../../TransientDistributedCellFieldTests.jl")
   include("../../TransientMultiFieldDistributedCellFieldTests.jl")
+  include("../../ZeroMeanFESpacesTests.jl")
   include("../../HeatEquationTests.jl")
   include("../../StokesOpenBoundaryTests.jl")
   include("../../AdaptivityTests.jl")
diff --git a/test/ZeroMeanFESpacesTests.jl b/test/ZeroMeanFESpacesTests.jl
new file mode 100644
index 0000000..a1c309f
--- /dev/null
+++ b/test/ZeroMeanFESpacesTests.jl
@@ -0,0 +1,82 @@
+module ZeroMeanFESpacesTests
+
+using Gridap
+using Gridap.FESpaces, Gridap.MultiField, Gridap.Algebra
+using GridapDistributed
+using PartitionedArrays
+using Test
+
+function main(distribute, parts)
+  ranks = distribute(LinearIndices((prod(parts),)))
+  model = CartesianDiscreteModel(ranks,parts,(0,1,0,1),(4,4))
+
+  Ω = Triangulation(model)
+  dΩ = Measure(Ω,4)
+
+  u_ex(x) = VectorValue(x[2],-x[1])
+  p_ex(x) = x[1] + 2*x[2]
+  p_mean = sum(∫(p_ex)*dΩ)
+
+  order = 2
+  reffe_u = ReferenceFE(lagrangian,VectorValue{2,Float64},order)
+  reffe_p = ReferenceFE(lagrangian,Float64,order-1;space=:P)
+
+  V  = FESpace(model,reffe_u;dirichlet_tags="boundary")
+  U  = TrialFESpace(V,u_ex)
+  Q  = FESpace(model,reffe_p;conformity=:L2)
+  Q0 = FESpace(model,reffe_p;conformity=:L2,constraint=:zeromean)
+  @test isa(Q0,GridapDistributed.DistributedZeroMeanFESpace)
+
+  X = MultiFieldFESpace([U,Q0])
+  Y = MultiFieldFESpace([V,Q0])
+
+  ph_i = interpolate(p_ex, Q0)
+  @test abs(sum(∫(ph_i)*dΩ)) < 1.0e-10
+  @test abs(sum(∫(ph_i - p_ex + p_mean)*dΩ)) < 1.0e-10
+
+  # Stokes
+
+  a((u,p),(v,q)) = ∫(∇(u)⊙∇(v) - (∇⋅v)*p - q*(∇⋅u))dΩ
+  l((v,q)) = a((u_ex,p_ex),(v,q))
+
+  op = AffineFEOperator(a,l,X,Y)
+  uh, ph = solve(op);
+
+  l2_error(u,v) = sqrt(sum(∫((u-v)⋅(u-v))*dΩ))
+  @test l2_error(uh,u_ex) < 1.0e-10
+  @test abs(sum(∫(ph)*dΩ)) < 1.0e-10
+  @test l2_error(ph,ph_i) < 1.0e-10
+
+  b(u,q) = ∫(q*(∇⋅u))dΩ
+  B = assemble_vector(q -> b(uh,q),Q)
+  B0 = assemble_vector(q -> b(uh,q),Q0)
+  @test abs(sum(B)) < 1.0e-10
+  @test abs(sum(B0)) < 1.0e-10
+
+  # Navier-Stokes
+
+  uh_ex = interpolate(u_ex, U)
+  @test l2_error(uh,uh_ex) < 1.0e-10
+  conv(u,∇u) = (∇u')⋅u
+  dconv(du,∇du,u,∇u) = conv(u,∇du)+conv(du,∇u)
+  c(u,v) = ∫(v⊙(conv∘(u,∇(u))))dΩ
+  dc(u,du,dv) = ∫(dv⊙(dconv∘(du,∇(du),u,∇(u))))dΩ
+  jac((u,p),(du,dp),(dv,dq)) = a((du,dp),(dv,dq)) + dc(u,du,dv)
+  res((u,p),(dv,dq)) = a((u,p),(dv,dq)) + c(u,dv) - a((u_ex,p_ex),(dv,dq)) - c(uh_ex,dv)
+
+  op = FEOperator(res,jac,X,Y)
+  uh, ph = solve(op);
+
+  @test l2_error(uh,u_ex) < 1.0e-10
+  @test abs(sum(∫(ph)*dΩ)) < 1.0e-10
+  @test l2_error(ph,ph_i) < 1.0e-10
+
+  B = assemble_vector(q -> b(uh,q),Q)
+  B0 = assemble_vector(q -> b(uh,q),Q0)
+  @test abs(sum(B)) < 1.0e-10
+  @test abs(sum(B0)) < 1.0e-10
+
+  return true
+end
+
+end
\ No newline at end of file
diff --git a/test/mpi/runtests_np4_body.jl b/test/mpi/runtests_np4_body.jl
index 720179a..355f662 100644
--- a/test/mpi/runtests_np4_body.jl
+++ b/test/mpi/runtests_np4_body.jl
@@ -28,6 +28,9 @@ function all_tests(distribute,parts)
   TestApp.TransientMultiFieldDistributedCellFieldTests.main(distribute,parts)
   PArrays.toc!(t,"TransientMultiFieldDistributedCellField")
 
+  TestApp.ZeroMeanFESpacesTests.main(distribute,parts)
+  PArrays.toc!(t,"ZeroMeanFESpaces")
+
   TestApp.PeriodicBCsTests.main(distribute,parts)
   PArrays.toc!(t,"PeriodicBCs")
 
diff --git a/test/sequential/ZeroMeanFESpacesTests.jl b/test/sequential/ZeroMeanFESpacesTests.jl
new file mode 100644
index 0000000..38fcc6e
--- /dev/null
+++ b/test/sequential/ZeroMeanFESpacesTests.jl
@@ -0,0 +1,10 @@
+module ZeroMeanFESpacesTestsSeq
+
+using PartitionedArrays
+include("../ZeroMeanFESpacesTests.jl")
+
+with_debug() do distribute
+  ZeroMeanFESpacesTests.main(distribute,(2,2))
+end
+
+end # module
diff --git a/test/sequential/runtests.jl b/test/sequential/runtests.jl
index d9285ff..7dabeb8 100644
--- a/test/sequential/runtests.jl
+++ b/test/sequential/runtests.jl
@@ -30,6 +30,8 @@ end
   include("TransientMultiFieldDistributedCellFieldTests.jl")
 end
 
+@time @testset "ZeroMeanFESpaces" begin include("ZeroMeanFESpacesTests.jl") end
+
 @time @testset "HeatEquation" begin include("HeatEquationTests.jl") end
 
 @time @testset "StokesOpenBoundary" begin include("StokesOpenBoundaryTests.jl") end

From d13a548869666e09b0d04ec1e85643b3edf71fee Mon Sep 17 00:00:00 2001
From: JordiManyer <jordi.manyer@monash.edu>
Date: Fri, 9 Aug 2024 23:01:41 +1000
Subject: [PATCH 07/15] Minor

---
 src/MultiField.jl | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/MultiField.jl b/src/MultiField.jl
index e0ffbbb..71f2f4d 100644
--- a/src/MultiField.jl
+++ b/src/MultiField.jl
@@ -7,7 +7,8 @@ struct DistributedMultiFieldCellField{A,B,C} <: CellField
   function DistributedMultiFieldCellField(
     field_fe_fun::AbstractVector{<:DistributedCellField},
     part_fe_fun::AbstractArray{<:CellField},
-    metadata=nothing)
+    metadata=nothing
+  )
     A = typeof(field_fe_fun)
     B = typeof(part_fe_fun)
     C = typeof(metadata)
@@ -45,7 +46,8 @@ const DistributedMultiFieldFEFunction{A,B,T} = DistributedMultiFieldCellField{A,
 function DistributedMultiFieldFEFunction(
   field_fe_fun::AbstractVector{<:DistributedSingleFieldFEFunction},
   part_fe_fun::AbstractArray{<:MultiFieldFEFunction},
-  free_values::AbstractVector)
+  free_values::AbstractVector
+)
   metadata = DistributedFEFunctionData(free_values)
   DistributedMultiFieldCellField(field_fe_fun,part_fe_fun,metadata)
 end

From 23bfc2f57a7e30161c924aaf48158486f205302a Mon Sep 17 00:00:00 2001
From: JordiManyer <jordi.manyer@monash.edu>
Date: Wed, 14 Aug 2024 09:34:18 +1000
Subject: [PATCH 08/15] Updated NEWS

---
 NEWS.md | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/NEWS.md b/NEWS.md
index 65ba24d..329a887 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Added
+
+- Added kwargs for VTK encoding options. Since PR[#156](https://github.com/gridap/GridapDistributed.jl/pull/156).
+
 ### Fixed
 
 - Fixed distributed interpolators for Vector-Valued FESpaces. Since PR[#152](https://github.com/gridap/GridapDistributed.jl/pull/152).

From bb27af55d3be599865a424026914c3437f971626 Mon Sep 17 00:00:00 2001
From: JordiManyer <jordi.manyer@monash.edu>
Date: Wed, 14 Aug 2024 09:51:09 +1000
Subject: [PATCH 09/15] Fixed tests

---
 test/MultiFieldTests.jl | 24 +++++++++++++-----------
 1 file changed, 13 insertions(+), 11 deletions(-)

diff --git a/test/MultiFieldTests.jl b/test/MultiFieldTests.jl
index 9a0475a..db05635 100644
--- a/test/MultiFieldTests.jl
+++ b/test/MultiFieldTests.jl
@@ -17,6 +17,7 @@ function main(distribute, parts, mfs)
 
   domain = (0,4,0,4)
   cells = (4,4)
+  model = CartesianDiscreteModel(domain,cells)
   model = CartesianDiscreteModel(ranks,parts,domain,cells)
   Ω = Triangulation(model)
 
@@ -26,18 +27,20 @@ function main(distribute, parts, mfs)
   reffe_p = ReferenceFE(lagrangian,Float64,k-1,space=:P)
 
   u((x,y)) = VectorValue((x+y)^2,(x-y)^2)
-  p((x,y)) = x+y
+  p̂((x,y)) = x+y
+  p̂_mean = sum(∫(p̂)dΩ)/sum(∫(1)dΩ)
+  p((x,y)) = p̂((x,y)) - p̂_mean
   f(x) = - Δ(u,x) + ∇(p,x)
   g(x) = tr(∇(u,x))
 
   V = TestFESpace(model,reffe_u,dirichlet_tags="boundary")
   Q = TestFESpace(model,reffe_p,constraint=:zeromean)
   U = TrialFESpace(V,u)
-  P = TrialFESpace(Q,p)
+  P = TrialFESpace(Q)
 
   VxQ = MultiFieldFESpace([V,Q];style=mfs)
-  UxP = MultiFieldFESpace([U,P];style=mfs) # This generates again the global numbering
   UxP = TrialFESpace([u,p],VxQ) # This reuses the one computed
+  UxP = MultiFieldFESpace([U,P];style=mfs) # This generates again the global numbering
   @test length(UxP) == 2
 
   uh, ph = interpolate([u,p],UxP)
@@ -53,20 +56,19 @@ function main(distribute, parts, mfs)
     uh, ph = solve(solver,op)
     @test l2_error(u,uh,dΩ) < 1.0e-9
     @test l2_error(p,ph,dΩ) < 1.0e-9
-
-    writevtk(Ω,"Ω",nsubcells=10,cellfields=["uh"=>uh,"ph"=>ph])
+    #writevtk(Ω,"Ω",nsubcells=10,cellfields=["uh"=>uh,"ph"=>ph])
   end
 
   A  = get_matrix(op)
   xh = interpolate([u,p],UxP)
   x  = GridapDistributed.change_ghost(get_free_dof_values(xh),axes(A,2))
-  uh1, ph1 = FESpaces.EvaluationFunction(UxP,x)
-  uh2, ph2 = FEFunction(UxP,x)
+  uh, ph = xh
+
+  @test l2_error(u,uh,dΩ) < 1.0e-9
+  @test l2_error(p,ph,dΩ) < 1.0e-9
 
-  @test l2_error(u,uh1,dΩ) < 1.0e-9
-  @test l2_error(p,ph1,dΩ) < 1.0e-9
-  @test l2_error(u,uh2,dΩ) < 1.0e-9
-  @test l2_error(p,ph2,dΩ) < 1.0e-9
+  x_p = partition(get_free_dof_values(ph))
+  x_p2 = partition(get_free_dof_values(ph2))
 
   a1(x,y) = ∫(x⋅y)dΩ
   a2((u,p),(v,q)) = ∫(u⋅v + p⋅q)dΩ

From 5f6cd9af12f518288d811b4a1a77f1b819b4c142 Mon Sep 17 00:00:00 2001
From: JordiManyer <jordi.manyer@monash.edu>
Date: Wed, 14 Aug 2024 09:52:36 +1000
Subject: [PATCH 10/15] Minor

---
 test/MultiFieldTests.jl | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/test/MultiFieldTests.jl b/test/MultiFieldTests.jl
index db05635..3e70c2e 100644
--- a/test/MultiFieldTests.jl
+++ b/test/MultiFieldTests.jl
@@ -67,9 +67,6 @@ function main(distribute, parts, mfs)
   @test l2_error(u,uh,dΩ) < 1.0e-9
   @test l2_error(p,ph,dΩ) < 1.0e-9
 
-  x_p = partition(get_free_dof_values(ph))
-  x_p2 = partition(get_free_dof_values(ph2))
-
   a1(x,y) = ∫(x⋅y)dΩ
   a2((u,p),(v,q)) = ∫(u⋅v + p⋅q)dΩ
   A1 = assemble_matrix(a1,UxP,UxP)

From 0f51e0685305a7719ac93d4d29d8ee35807ece93 Mon Sep 17 00:00:00 2001
From: JordiManyer <jordi.manyer@monash.edu>
Date: Wed, 14 Aug 2024 10:33:33 +1000
Subject: [PATCH 11/15] Added proper comparisions between BlockPRanges

---
 src/BlockPartitionedArrays.jl | 15 +++++++++++++++
 src/MultiField.jl             |  4 ++--
 test/MultiFieldTests.jl       |  1 -
 3 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/src/BlockPartitionedArrays.jl b/src/BlockPartitionedArrays.jl
index 25e4346..aa83bb2 100644
--- a/src/BlockPartitionedArrays.jl
+++ b/src/BlockPartitionedArrays.jl
@@ -25,6 +25,21 @@ function Base.getindex(a::BlockPRange,inds::Block{1})
   a.ranges[inds.n...]
 end
 
+function PartitionedArrays.matching_local_indices(a::BlockPRange,b::BlockPRange)
+  c = map(PartitionedArrays.matching_local_indices,blocks(a),blocks(b))
+  reduce(&,c,init=true)
+end
+
+function PartitionedArrays.matching_own_indices(a::BlockPRange,b::BlockPRange)
+  c = map(PartitionedArrays.matching_own_indices,blocks(a),blocks(b))
+  reduce(&,c,init=true)
+end
+
+function PartitionedArrays.matching_ghost_indices(a::BlockPRange,b::BlockPRange)
+  c = map(PartitionedArrays.matching_ghost_indices,blocks(a),blocks(b))
+  reduce(&,c,init=true)
+end
+
 """
   struct BlockPArray{V,T,N,A,B} <: BlockArrays.AbstractBlockArray{T,N}
 """
diff --git a/src/MultiField.jl b/src/MultiField.jl
index 71f2f4d..9987558 100644
--- a/src/MultiField.jl
+++ b/src/MultiField.jl
@@ -186,7 +186,7 @@ end
 
 function FESpaces.interpolate!(objects,free_values::AbstractVector,space::DistributedMultiFieldFESpace)
   msg = "free_values and FESpace have incompatible index partitions."
-  @check partition(axes(free_values,1)) === partition(space.gids) msg
+  @check PartitionedArrays.matching_local_indices(axes(free_values,1),get_free_dof_ids(space)) msg
 
   # Interpolate each field
   field_fe_fun = DistributedSingleFieldFEFunction[]
@@ -219,7 +219,7 @@ function FESpaces.interpolate_everywhere!(
   space::DistributedMultiFieldFESpace
 )
   msg = "free_values and FESpace have incompatible index partitions."
-  @check partition(axes(free_values,1)) === partition(space.gids) msg
+  @check PartitionedArrays.matching_local_indices(axes(free_values,1),get_free_dof_ids(space)) msg
   
   # Interpolate each field
   field_fe_fun = DistributedSingleFieldFEFunction[]
diff --git a/test/MultiFieldTests.jl b/test/MultiFieldTests.jl
index 3e70c2e..d38ff70 100644
--- a/test/MultiFieldTests.jl
+++ b/test/MultiFieldTests.jl
@@ -17,7 +17,6 @@ function main(distribute, parts, mfs)
 
   domain = (0,4,0,4)
   cells = (4,4)
-  model = CartesianDiscreteModel(domain,cells)
   model = CartesianDiscreteModel(ranks,parts,domain,cells)
   Ω = Triangulation(model)
 

From 080534ba12b628408a5a7772fac14cdfd333fce8 Mon Sep 17 00:00:00 2001
From: JordiManyer <jordi.manyer@monash.edu>
Date: Wed, 14 Aug 2024 10:53:26 +1000
Subject: [PATCH 12/15] Minor

---
 src/BlockPartitionedArrays.jl | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/src/BlockPartitionedArrays.jl b/src/BlockPartitionedArrays.jl
index aa83bb2..3debc3e 100644
--- a/src/BlockPartitionedArrays.jl
+++ b/src/BlockPartitionedArrays.jl
@@ -351,6 +351,9 @@ function LinearAlgebra.dot(x::BlockPVector,y::BlockPVector)
 end
 
 function LinearAlgebra.norm(v::BlockPVector,p::Real=2)
+  if p == 2
+    return sqrt(dot(v,v))
+  end
   block_norms = map(vi->norm(vi,p),blocks(v))
   return sum(block_norms.^p)^(1/p)
 end

From d02f19a33ca79f936b389a961aa3faac7940b878 Mon Sep 17 00:00:00 2001
From: JordiManyer <jordi.manyer@monash.edu>
Date: Wed, 14 Aug 2024 10:56:15 +1000
Subject: [PATCH 13/15] Minor

---
 src/BlockPartitionedArrays.jl | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/BlockPartitionedArrays.jl b/src/BlockPartitionedArrays.jl
index 3debc3e..41e3f6d 100644
--- a/src/BlockPartitionedArrays.jl
+++ b/src/BlockPartitionedArrays.jl
@@ -352,7 +352,9 @@ end
 
 function LinearAlgebra.norm(v::BlockPVector,p::Real=2)
   if p == 2
-    return sqrt(dot(v,v))
+    # More accurate, I think, given the fact we are not 
+    # repeating the sqrt(square(sqrt...)) process in every block and every processor
+    return sqrt(dot(v,v)) 
   end
   block_norms = map(vi->norm(vi,p),blocks(v))
   return sum(block_norms.^p)^(1/p)

From 9567fe0090aed6ee4526448fc761e8d1be8f6dae Mon Sep 17 00:00:00 2001
From: JordiManyer <jordi.manyer@monash.edu>
Date: Wed, 14 Aug 2024 12:21:44 +1000
Subject: [PATCH 14/15] Minor

---
 src/MultiField.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/MultiField.jl b/src/MultiField.jl
index 9987558..689a062 100644
--- a/src/MultiField.jl
+++ b/src/MultiField.jl
@@ -215,7 +215,7 @@ end
 function FESpaces.interpolate_everywhere!(
   objects,
   free_values::AbstractVector,
-  dirichlet_values::Vector{AbstractArray{<:AbstractVector}},
+  dirichlet_values::Vector{<:AbstractArray{<:AbstractVector}},
   space::DistributedMultiFieldFESpace
 )
   msg = "free_values and FESpace have incompatible index partitions."

From 0ec889d4592cbbd07bc00e35802ddfd5e13e36d1 Mon Sep 17 00:00:00 2001
From: JordiManyer <jordi.manyer@monash.edu>
Date: Wed, 14 Aug 2024 17:43:53 +1000
Subject: [PATCH 15/15] Bumped release

---
 NEWS.md      | 3 ++-
 Project.toml | 2 +-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/NEWS.md b/NEWS.md
index 329a887..c69eb42 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -5,11 +5,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
-## Unreleased
+## [0.4.4] 2024-08-14
 
 ### Added
 
 - Added kwargs for VTK encoding options. Since PR[#156](https://github.com/gridap/GridapDistributed.jl/pull/156).
+- Reimplemented distributed ZeroMeanFESpaces. Since PR[#155](https://github.com/gridap/GridapDistributed.jl/pull/155).
 
 ### Fixed
 
diff --git a/Project.toml b/Project.toml
index 00162ac..1dc2bd6 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "GridapDistributed"
 uuid = "f9701e48-63b3-45aa-9a63-9bc6c271f355"
 authors = ["S. Badia <santiago.badia@monash.edu>", "A. F. Martin <alberto.f.martin@anu.edu.au>", "F. Verdugo <f.verdugo.rojano@vu.nl>"]
-version = "0.4.3"
+version = "0.4.4"
 
 [deps]
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"