Merge pull request #120 from fverdugo/partition_from_colors

Partition from colors

CHANGELOG.md

@@ -5,12 +5,18 @@ 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
+
+### Added 
+
+- Function `partition_from_color`.
+
 ## [0.3.3] - 2023-08-09
 
 ### Added 
 
-- Added an MPI ibarrier-based (supposedly scalable) algorithm to find rcv neighbours in a sparse all-to-all communication graph given the snd neighbors. We left the previous non-scalable algorithm as default (based on gather-scatter) until we have experimental evidence on the relative performance and scalability of the former with respect to the latter and for which core ranges.
-- Added a new kwarg `discover_cols=true` to the `psparse!` constructor, which allows the user to skip column index discovery.
+- MPI ibarrier-based (supposedly scalable) algorithm to find rcv neighbours in a sparse all-to-all communication graph given the snd neighbors. We left the previous non-scalable algorithm as default (based on gather-scatter) until we have experimental evidence on the relative performance and scalability of the former with respect to the latter and for which core ranges.
+- New kwarg `discover_cols=true` to the `psparse!` constructor, which allows the user to skip column index discovery.
 
 ### Fixed
 

docs/src/reference/partition.md

@@ -5,6 +5,7 @@
 ```@docs
 uniform_partition
 variable_partition
+partition_from_color
 ```
 
 ## AbstractLocalIndices

src/PartitionedArrays.jl

@@ -61,6 +61,7 @@ include("mpi_array.jl")
 export PRange
 export uniform_partition
 export variable_partition
+export partition_from_color
 export AbstractLocalIndices
 export OwnAndGhostIndices
 export LocalIndices

src/p_range.jl

@@ -696,6 +696,43 @@ function variable_partition(
     indices
 end
 
+"""
+    partition_from_color(ranks,global_to_color;multicast=false,source=MAIN)
+
+Build an arbitrary 1d partition by defining the parts via the argument `global_to_color` (see below).
+The output is a vector of vectors containing the indices in each component of
+the partition. The `eltype` of the result implements the [`AbstractLocalIndices`](@ref)
+interface.
+
+# Arguments
+
+- `ranks`: Array containing the distribution of ranks.
+- `global_to_color`: If `multicast==false`,  `global_to_color[gid]` contains the part id that owns the global id `gid`. If `multicast==true`, then   `global_to_color[source][gid]` contains the part id that owns the global id `gid`.
+
+# Key-word arguments
+- `multicast=false`
+- `source=MAIN`
+
+This function is useful when generating a partition using a graph partitioner such as METIS.
+The argument `global_to_color` is the usual output of such tools.
+"""
+function partition_from_color(ranks,global_to_color;multicast=false,source=MAIN)
+    if multicast == true
+        global_to_owner = getany(emit(global_to_color;source))
+    else
+        global_to_owner = global_to_color
+    end
+    map(ranks) do rank
+        nglobal = length(global_to_owner)
+        own_to_global = findall(owner->owner==rank,global_to_owner)
+            ghost_to_global = Int[]
+        ghost_to_owner = Int32[]
+        own = OwnIndices(nglobal,rank,own_to_global)
+        ghost = GhostIndices(nglobal,ghost_to_global,ghost_to_owner)
+        OwnAndGhostIndices(own,ghost,global_to_owner)
+    end
+end
+
 function local_range(p,np,n,ghost=false,periodic=false)
     l = n ÷ np
     offset = l * (p-1)
@@ -1119,23 +1156,28 @@ Local indices are defined by concatenating own and ghost ones.
 
 - `own::OwnIndices`: Container for the own indices.
 - `ghost::GhostIndices`: Container for the ghost indices.
+- `global_to_owner`: [optional: it can be `nothing`] Vector containing the owner of each global id.
 
 # Supertype hierarchy
 
-    OwnAndGhostIndices <: AbstractLocalIndices
+    OwnAndGhostIndices{A} <: AbstractLocalIndices
+
+where `A=typeof(global_to_owner)`.
 
 """
-struct OwnAndGhostIndices <: AbstractLocalIndices
+struct OwnAndGhostIndices{A} <: AbstractLocalIndices
     own::OwnIndices
     ghost::GhostIndices
+    global_to_owner::A
     assembly_cache::AssemblyCache
     @doc """
-        OwnAndGhostIndices(own::OwnIndices,ghost::GhostIndices)
+        OwnAndGhostIndices(own::OwnIndices,ghost::GhostIndices,global_to_owner=nothing)
 
-    Build an instance of [`OwnAndGhostIndices`](@ref) from the underlying properties `own` and `ghost`.
+    Build an instance of [`OwnAndGhostIndices`](@ref) from the underlying properties `own`, `ghost`, and `global_to_owner`.
     """
-    function OwnAndGhostIndices(own::OwnIndices,ghost::GhostIndices)
-        new(own,ghost,AssemblyCache())
+    function OwnAndGhostIndices(own::OwnIndices,ghost::GhostIndices,global_to_owner=nothing)
+        A = typeof(global_to_owner)
+        new{A}(own,ghost,global_to_owner,AssemblyCache())
     end
 end
 assembly_cache(a::OwnAndGhostIndices) = a.assembly_cache
@@ -1146,6 +1188,15 @@ function replace_ghost(a::OwnAndGhostIndices,ghost::GhostIndices)
     OwnAndGhostIndices(a.own,ghost)
 end
 
+function find_owner(indices,global_ids,::Type{<:OwnAndGhostIndices{T}}) where T
+    if T == Nothing
+        error("Not enough data to perform this operation without communciation")
+    end
+    map(indices,global_ids) do indices,global_ids
+        indices.global_to_owner[global_ids]
+    end
+end
+
 part_id(a::OwnAndGhostIndices) = a.own.owner
 
 function own_to_global(a::OwnAndGhostIndices)

test/p_range_tests.jl

@@ -245,4 +245,51 @@ function p_range_tests(distribute)
 
    display(PRange(ids4))
 
+   n = 21
+   np = 4
+   rank = distribute(LinearIndices((4,)))
+   global_to_color = zeros(Int,n)
+   for p in 1:np
+       global_to_color[local_range(p,np,n)] .= p
+   end
+   ids_color = partition_from_color(rank,global_to_color)
+   ids_uniform = uniform_partition(rank,n)
+   map(ids_color,ids_uniform) do ids1,ids2
+       @test ids1 == ids2
+   end
+
+   global_to_color = map(rank) do rank
+       if rank == MAIN
+           my_global_to_color = zeros(Int,n)
+           for p in 1:np
+               my_global_to_color[local_range(p,np,n)] .= p
+           end
+       else
+           my_global_to_color = Int[]
+       end
+       my_global_to_color
+   end
+   ids_color = partition_from_color(rank,global_to_color;multicast=true)
+   ids_color = partition_from_color(rank,global_to_color;multicast=true,source=MAIN)
+   map(ids_color,ids_uniform) do ids1,ids2
+       @test ids1 == ids2
+   end
+
+   gids = map(rank) do part
+       if part == 1
+           [1,4,6]
+       elseif part == 2
+           [3,1,2,8]
+       elseif part == 3
+           [1,9,6]
+       else
+           [3,2,8,10]
+       end
+   end
+   owners1 = find_owner(ids_color,gids)
+   owners2 = find_owner(ids_uniform,gids)
+   map(owners1,owners2) do ids1,ids2
+       @test ids1 == ids2
+   end
+
 end