Added/modified code required in order to have MPI-parallel Stokes test

working with global FE spaces

Author: amartinhuertas

src/MultiFieldDistributedFESpaces.jl

@@ -1,138 +1,183 @@
 struct MultiFieldDistributedFESpace{V} <: DistributedFESpace{V}
-   vector_type         :: Type{V}
-   distributed_spaces  :: Vector{<:DistributedFESpace}
-   spaces              :: DistributedData{<:MultiFieldFESpace}
-   gids                :: DistributedIndexSet
- end
+    vector_type::Type{V}
+    distributed_spaces::Vector{<:DistributedFESpace}
+    spaces::DistributedData{<:MultiFieldFESpace}
+    gids::DistributedIndexSet
+end
 
- function Gridap.MultiFieldFESpace(test_space::MultiFieldDistributedFESpace{V},
+function Gridap.MultiFieldFESpace(test_space::MultiFieldDistributedFESpace{V},
                                    trial_spaces::Vector{<:DistributedFESpace{V}}) where V
-  spaces = DistributedData(trial_spaces...) do part, spaces_and_gids...
-    MultiFieldFESpace([s[1] for s in spaces_and_gids])
-  end
-  MultiFieldDistributedFESpace(V,trial_spaces,spaces,test_space.gids)
- end
+    spaces = DistributedData(trial_spaces...) do part, spaces_and_gids...
+        MultiFieldFESpace([s[1] for s in spaces_and_gids])
+    end
+    MultiFieldDistributedFESpace(V, trial_spaces, spaces, test_space.gids)
+end
 
 
-function Gridap.FESpaces.FEFunction(dV::MultiFieldDistributedFESpace{T},x) where {T}
-  _gen_multifield_distributed_fe_function(dV,x,FEFunction)
+function Gridap.FESpaces.FEFunction(dV::MultiFieldDistributedFESpace{T}, x) where {T}
+    _gen_multifield_distributed_fe_function(dV, x, FEFunction)
 end
 
 
-function _gen_multifield_distributed_fe_function(dV::MultiFieldDistributedFESpace{T},x,f) where {T}
-  single_fe_functions = DistributedFEFunction{T}[]
-  for (field, U) in enumerate(dV.distributed_spaces)
-    free_values_i = restrict_to_field(dV,x,field)
-    uhi = f(U,free_values_i)
-    push!(single_fe_functions,uhi)
-  end
+function _gen_multifield_distributed_fe_function(dV::MultiFieldDistributedFESpace{T}, x, f) where {T}
+    single_fe_functions = DistributedFEFunction{T}[]
+    for (field, U) in enumerate(dV.distributed_spaces)
+        free_values_i = restrict_to_field(dV, x, field)
+        uhi = f(U, free_values_i)
+        push!(single_fe_functions, uhi)
+    end
 
-  funs = DistributedData(get_comm(dV.distributed_spaces[1]),
+    funs = DistributedData(get_comm(dV.distributed_spaces[1]),
                          dV.spaces, single_fe_functions...,) do part, V, fe_functions...
-      mfv=zero_free_values(V)
-      current=1
-      for fun in fe_functions
-        fv = get_free_values(fun)
-        for i=1:length(fv)
-          mfv[current]=fv[i]
-          current=current+1
+        mfv = zero_free_values(V)
+        current = 1
+        for fun in fe_functions
+            fv = get_free_values(fun)
+            for i = 1:length(fv)
+                mfv[current] = fv[i]
+                current = current + 1
+            end
         end
-      end
-      f(V,mfv)
-  end
-  multifield_fe_function=DistributedFEFunction(funs,x,dV)
-  MultiFieldDistributedFEFunction{T}(single_fe_functions,
+        f(V, mfv)
+    end
+    multifield_fe_function = DistributedFEFunction(funs, x, dV)
+    MultiFieldDistributedFEFunction{T}(single_fe_functions,
                                   multifield_fe_function,
                                   dV)
 end
 
 
-function restrict_to_field(dV::MultiFieldDistributedFESpace,x::Vector,field)
-  @assert isa(dV.gids,SequentialDistributedIndexSet)
-  ngids = dV.distributed_spaces[field].gids.ngids
-  xi = Vector{eltype(x)}(undef,ngids)
-  do_on_parts(dV.spaces,dV.gids,xi,x,dV.distributed_spaces...) do part, mfspace, mfgids, xi, x,fspaces_and_gids...
-    offset=0
-    for i=1:field-1
-      fspace = fspaces_and_gids[i][1]
-      offset = offset + num_free_dofs(fspace)
+function restrict_to_field(dV::MultiFieldDistributedFESpace, x::Vector, field)
+    @assert isa(dV.gids, SequentialDistributedIndexSet)
+
+    xi = Gridap.Algebra.allocate_vector(Vector{eltype(x)},
+                                      dV.distributed_spaces[field].gids)
+
+    do_on_parts(dV.spaces, dV.gids, xi, x, dV.distributed_spaces...) do part, mfspace, mfgids, xi, x, fspaces_and_gids...
+        offset = 0
+        for i = 1:field - 1
+            fspace = fspaces_and_gids[i][1]
+            offset = offset + num_free_dofs(fspace)
+        end
+        fspace = fspaces_and_gids[field][1]
+        fgids  = fspaces_and_gids[field][2]
+        for i = 1:num_free_dofs(fspace)
+            if fgids.lid_to_owner[i] == part
+                xi[fgids.lid_to_gid[i]] = x[mfgids.lid_to_gid[offset + i]]
+            end
+        end
     end
-    fspace = fspaces_and_gids[field][1]
-    fgids  = fspaces_and_gids[field][2]
-    for i=1:num_free_dofs(fspace)
-      if fgids.lid_to_owner[i] == part
-        xi[fgids.lid_to_gid[i]]=x[mfgids.lid_to_gid[offset+i]]
-      end
+    xi
+end
+
+function restrict_to_field(dV::MultiFieldDistributedFESpace, x::PETSc.Vec{Float64}, field)
+    fgids  = dV.distributed_spaces[field].gids
+    mfgids = dV.gids
+    @assert isa(fgids, MPIPETScDistributedIndexSet)
+
+    xi = Gridap.Algebra.allocate_vector(PETSc.Vec{Float64},
+                                      dV.distributed_spaces[field].gids)
+
+    comm = get_comm(fgids)
+    part = get_part(comm)
+
+    fis_gids = [ PETSc.PetscInt(fgids.lid_to_gid_petsc[i] - 1)
+                  for i = 1:length(fgids.lid_to_gid_petsc)
+                     if fgids.parts.part.lid_to_owner[i] == part ]
+
+    mfis_gids = Vector{PETSc.PetscInt}(undef,length(fis_gids))
+
+    do_on_parts(dV.gids, dV.distributed_spaces...) do part, lmfgids, fspaces_and_gids...
+        offset = 0
+        for i = 1:field - 1
+            fspace = fspaces_and_gids[i][1]
+            offset = offset + num_free_dofs(fspace)
+        end
+        fspace = fspaces_and_gids[field][1]
+        current=1
+        for i = 1:num_free_dofs(fspace)
+            if lmfgids.lid_to_owner[offset+i] == part
+                mfis_gids[current]=mfgids.lid_to_gid_petsc[offset+i]-1
+                current=current+1
+            end
+        end
     end
-  end
-  xi
+
+    fis  = PETSc.IS_(Float64, fis_gids; comm=comm.comm)
+    mfis = PETSc.IS_(Float64, mfis_gids; comm=comm.comm)
+
+    vscatter = PETSc.VecScatter(x, mfis, xi, fis)
+    scatter!(vscatter,x,xi)
+    xi
 end
 
-function Gridap.FESpaces.EvaluationFunction(dV::MultiFieldDistributedFESpace,x)
-  _gen_multifield_distributed_fe_function(dV,x,EvaluationFEFunction)
+
+
+function Gridap.FESpaces.EvaluationFunction(dV::MultiFieldDistributedFESpace, x)
+    _gen_multifield_distributed_fe_function(dV, x, EvaluationFEFunction)
 end
 
 
 function Gridap.MultiFieldFESpace(model::DistributedDiscreteModel,
                                   distributed_spaces::Vector{<:DistributedFESpace{V}}) where V
 
-  spaces = DistributedData(distributed_spaces...) do part, spaces_and_gids...
-    MultiFieldFESpace([s[1] for s in spaces_and_gids])
-  end
-
-  function init_lid_to_owner(part,lspace,spaces_and_gids...)
-    nlids = num_free_dofs(lspace)
-    lid_to_owner = zeros(Int,nlids)
-    current_lid = 1
-    for current_field_space_gids in spaces_and_gids
-      gids = current_field_space_gids[2]
-      for i=1:length(gids.lid_to_owner)
-        lid_to_owner[current_lid]=gids.lid_to_owner[i]
-        current_lid += 1
-      end
+    spaces = DistributedData(distributed_spaces...) do part, spaces_and_gids...
+        MultiFieldFESpace([s[1] for s in spaces_and_gids])
+    end
+
+    function init_lid_to_owner(part, lspace, spaces_and_gids...)
+        nlids = num_free_dofs(lspace)
+        lid_to_owner = zeros(Int, nlids)
+        current_lid = 1
+        for current_field_space_gids in spaces_and_gids
+            gids = current_field_space_gids[2]
+            for i = 1:length(gids.lid_to_owner)
+                lid_to_owner[current_lid] = gids.lid_to_owner[i]
+                current_lid += 1
+            end
+        end
+        lid_to_owner
     end
-    lid_to_owner
-  end
 
-  comm = get_comm(model)
+    comm = get_comm(model)
 
-  part_to_lid_to_owner = DistributedData{Vector{Int}}(init_lid_to_owner,
+    part_to_lid_to_owner = DistributedData{Vector{Int}}(init_lid_to_owner,
                                                       comm,
                                                       spaces,
                                                       distributed_spaces...)
 
-  offsets, ngids = _compute_offsets_and_ngids(part_to_lid_to_owner)
+    offsets, ngids = _compute_offsets_and_ngids(part_to_lid_to_owner)
 
-  num_dofs_x_cell = compute_num_dofs_x_cell(comm,spaces)
+    num_dofs_x_cell = compute_num_dofs_x_cell(comm, spaces)
 
-  part_to_lid_to_gid = _compute_part_to_lid_to_gid(model,
+    part_to_lid_to_gid = _compute_part_to_lid_to_gid(model,
                                                    spaces,
                                                    num_dofs_x_cell,
                                                    part_to_lid_to_owner,
                                                    offsets)
 
-  function init_free_gids(part, lid_to_gid, lid_to_owner, ngids)
-    IndexSet(ngids, lid_to_gid, lid_to_owner)
-  end
+    function init_free_gids(part, lid_to_gid, lid_to_owner, ngids)
+        IndexSet(ngids, lid_to_gid, lid_to_owner)
+    end
 
-  gids = DistributedIndexSet(init_free_gids,
+    gids = DistributedIndexSet(init_free_gids,
                              comm,
                              ngids,
                              part_to_lid_to_gid,
                              part_to_lid_to_owner,
                              ngids)
 
-  MultiFieldDistributedFESpace(V,
+    MultiFieldDistributedFESpace(V,
                                distributed_spaces,
                                spaces,
                                gids)
 end
 
 # FE Function
 struct MultiFieldDistributedFEFunction{T}
-  single_fe_functions    :: Vector{DistributedFEFunction{T}}
-  multifield_fe_function :: DistributedFEFunction{T}
-  space                  :: MultiFieldDistributedFESpace{T}
+    single_fe_functions::Vector{DistributedFEFunction{T}}
+    multifield_fe_function::DistributedFEFunction{T}
+    space::MultiFieldDistributedFESpace{T}
 end
 
 Gridap.FESpaces.FEFunctionStyle(::Type{MultiFieldDistributedFEFunction}) = Val{true}()

test/MPIPETScDistributedStokesTests.jl

@@ -48,35 +48,33 @@ function run(comm,assembly_strategy::AbstractString, global_dofs::Bool)
     add_tag_from_tags!(labels,"neumann",[5])
   end
 
-  # Build local and global test FE spaces
-  spaces = DistributedData(comm, model) do part, (model,gids)
-    labels = get_face_labeling(model)
-    V = TestFESpace(
-           reffe=:QLagrangian,
-           conformity=:H1,
-           valuetype=VectorValue{2,Float64},
-           model=model,
-           labels=labels,
-           order=2,
-           dirichlet_tags=["diri0","diri1"])
-    Q = TestFESpace(
-            reffe=:PLagrangian,
-            conformity=:L2,
-            valuetype=Float64,
-            model=model,
-            order=1) #,
-            #constraint=:zeromean)
-    MultiFieldFESpace([V,Q])
-  end
-  Y=GridapDistributed.DistributedFESpaceFromLocalFESpaces(vector_type,model,spaces)
 
-  # Build local and global trial FE spaces
-  trialspaces = DistributedData(comm, Y) do part, (y,gids)
-    U=TrialFESpace(y.spaces[1],[u,u])
-    P=TrialFESpace(y.spaces[2])
-    MultiFieldFESpace([U,P])
-  end
-  X=GridapDistributed.DistributedFESpaceFromLocalFESpaces(vector_type,trialspaces,Y.gids)
+  # FE Spaces
+  order = 2
+  V = FESpace(
+        vector_type,
+        valuetype = VectorValue{2,Float64},
+        reffe=:QLagrangian,
+        order = order,
+        model = model,
+        conformity = :H1,
+        dirichlet_tags=["diri0","diri1"],
+      )
+
+  Q = FESpace(
+       vector_type,
+       reffe=:PLagrangian,
+       conformity=:L2,
+       valuetype=Float64,
+       model=model,
+       order=order-1)
+       #constraint=:zeromean)
+
+  Y=MultiFieldFESpace(model,[V,Q])
+
+  U=TrialFESpace(V,[u,u])
+  P=TrialFESpace(Q)
+  X=MultiFieldFESpace(Y,[U,P])
 
   if (assembly_strategy == "RowsComputedLocally")
     strategy = RowsComputedLocally(Y; global_dofs=global_dofs)