Added documentation

Author: JordiManyer

src/FESpaces.jl

@@ -161,6 +161,15 @@ function fetch_vector_ghost_values!(vector_partition,cache)
   assemble!((a,b)->b, vector_partition, cache) 
 end
 
+"""
+    generate_gids(cell_gids::PRange, cell_to_lids, nlids) -> PRange
+
+Given a set of cell global ids, a distributed array of local tables mapping cells to local dof ids,
+and a distributed array with the number of local dofs in each partition, this function 
+generates the global dof ids and returns them as a `PRange` of `LocalIndices`.
+
+Ignores negative local dof ids (usually used for Dirichlet dofs).
+"""
 function generate_gids(
   cell_range::PRange,
   cell_to_ldofs::AbstractArray{<:AbstractArray},
@@ -248,6 +257,13 @@ function generate_gids(
   return PRange(local_indices)
 end
 
+"""
+    generate_posneg_gids(cell_gids::PRange, cell_to_lposneg, nlpos, nlneg) -> (PRange,PRange)
+
+Similar to `generate_gids`, but also handles negative local dof ids. Returns two sets of 
+global ids: one for positive local ids and another for negative local ids.
+This can be used to generate simultaneously free and dirichlet dof global ids.
+"""
 function generate_posneg_gids(
   cell_range::PRange,
   cell_to_lposneg::AbstractArray{<:AbstractArray},
@@ -360,6 +376,16 @@ function generate_posneg_gids(
   return PRange(local_indices_pos), PRange(local_indices_neg)
 end
 
+"""
+    generate_gids_by_color(cell_gids::PRange, cell_to_lids, lid_to_color, ncolors)
+
+Similar to `generate_gids`, but uses a global partition given by `lid_to_color` to generate
+a different set of global ids per color. Returns:
+
+- a tuple with a `PRange` of `LocalIndices` per color.
+- a mapping `lid_to_clid` that maps local ids to local color ids
+- a mapping `color_to_clid_to_lid` that maps, for each color, local color ids to local ids.
+"""
 function generate_gids_by_color(
   cell_range::PRange, 
   cell_to_lids::AbstractArray{<:AbstractArray},
@@ -490,6 +516,13 @@ function generate_gids_by_color(
   return map(PRange,color_to_indices), lid_to_clid, color_to_clid_to_lid
 end
 
+"""
+    split_gids_by_color(gids::PRange, lid_to_color::AbstractArray, ncolors) -> NTuple{ncolors,PRange}
+
+Given a set of global ids and a mapping from local ids to colors, this function splits
+the global ids into different sets of global ids per color. Returns a tuple with a `PRange`
+of `LocalIndices` per color.
+"""
 function split_gids_by_color(
   gids::PRange, lid_to_color::AbstractArray, 
   ncolors = getany(reduction(max,map(maximum,lid_to_color);destination=:all))

test/FESpacesTests.jl

@@ -9,6 +9,44 @@ using Test
 
 u((x,y)) = x+y
 
+function gid_generation_tests(distribute,parts)
+  ranks = distribute(LinearIndices((prod(parts),)))
+
+  model = CartesianDiscreteModel(ranks,parts,(0,1,0,1),(8,4))
+  reffe = ReferenceFE(lagrangian,Float64,1)
+  cgids = GridapDistributed.get_cell_gids(model)
+
+  V0 = FESpace(model,reffe)
+  I0 = partition(get_free_dof_ids(V0))
+
+  V1 = FESpace(model,reffe;dirichlet_tags=["tag_1","tag_2","tag_3","tag_4","tag_5","tag_6"])
+  I1 = partition(get_free_dof_ids(V1))
+
+  fgids_1, dgids_1 = GridapDistributed.generate_posneg_gids(
+    cgids,get_cell_dof_ids(V1),map(num_free_dofs,local_views(V1)),map(num_dirichlet_dofs,local_views(V1))
+  )
+
+  # 1 -> free, 2 -> dirichlet
+  lid_to_color = map(local_views(V0),local_views(V1)) do V0, V1
+    lid_to_color = zeros(Int16,num_free_dofs(V0))
+    for (I,J) in zip(get_cell_dof_ids(V0),get_cell_dof_ids(V1))
+      lid_to_color[I] .= map(j -> ifelse(j > 0, 1, 2), J)
+    end
+    return lid_to_color
+  end
+
+  fgids_2, dgids_2 = GridapDistributed.split_gids_by_color(
+    get_free_dof_ids(V0), lid_to_color
+  )
+
+  (fgids_3, dgids_3), _, _ = GridapDistributed.generate_gids_by_color(
+    cgids, get_cell_dof_ids(V0), lid_to_color
+  )
+
+  @test fgids_1 == fgids_2 == fgids_3
+  @test dgids_1 == dgids_2 == dgids_3
+end
+
 function assemble_tests(das,dΩ,dΩass,U,V)
   # Assembly
   dv = get_fe_basis(V)