renumber_partition and other minor changes

Author: fverdugo

CHANGELOG.md

@@ -26,7 +26,7 @@ The previous "monolithic" storage is not implemented anymore for the new version
 - Efficient re-construction of `PSparseMatrix` and `PVector` objects.
 - Functions `assemble` and `consistent` (allocating versions of `assemble!` and `consistent!` with a slightly different
 treatment of the ghost rows).
-- Function `consitent` for `PSparseMatrix`.
+- Function `consistent` for `PSparseMatrix`.
 - Functions `repartition` and `repartition!` used to change the data partition of `PSparseMatrix` and `PVector` objects.
 - Functions `psystem` and `psystem!` for generating a system matrix and vector at once.
 - Function `trivial_partition`.

src/PartitionedArrays.jl

@@ -66,6 +66,7 @@ export uniform_partition
 export variable_partition
 export partition_from_color
 export trivial_partition
+export renumber_partition
 export AbstractLocalIndices
 export OwnAndGhostIndices
 export LocalIndices
@@ -102,6 +103,12 @@ export partition
 export assembly_graph
 export assembly_neighbors
 export assembly_local_indices
+export map_local_to_global!
+export map_global_to_local!
+export map_ghost_to_global!
+export map_global_to_ghost!
+export map_own_to_global!
+export map_global_to_own!
 include("p_range.jl")
 
 export local_values
@@ -126,6 +133,7 @@ export repartition!
 include("p_vector.jl")
 
 export OldPSparseMatrix
+export SplitMatrix
 export PSparseMatrix
 export old_psparse
 export psparse

src/p_range.jl

@@ -771,6 +771,27 @@ function trivial_partition(ranks,n;destination=MAIN)
     partition_in_main
 end
 
+function renumber_partition(partition_in)
+    own_ids = map(own_to_global,partition_in)
+    if eltype(own_ids) <: BlockPartitionOwnToGlobal{1}
+        return partition_in
+    end
+    n_global = PartitionedArrays.getany(map(global_length,partition_in))
+    n_own = map(own_length,partition_in)
+    new_gids = variable_partition(n_own,n_global)
+    v = PVector{Vector{Int}}(undef,partition_in)
+    map(own_values(v),new_gids) do own_v, new_gids
+        own_v .= own_to_global(new_gids)
+    end
+    consistent!(v) |> wait
+    I = ghost_values(v)
+    I_owner = map(ghost_to_owner,partition_in)
+    new_ids2 = map(union_ghost,new_gids,I,I_owner)
+    perm = map(PartitionedArrays.local_permutation,partition_in)
+    partition_out = map(permute_indices,new_ids2,perm)
+    partition_out
+end
+
 function local_range(p,np,n,ghost=false,periodic=false)
     l = n ÷ np
     offset = l * (p-1)
@@ -1223,7 +1244,7 @@ assembly_cache(a::OwnAndGhostIndices) = a.assembly_cache
 local_permutation(a::OwnAndGhostIndices) = Int32(1):Int32(local_length(a))
 
 function replace_ghost(a::OwnAndGhostIndices,ghost::GhostIndices)
-    OwnAndGhostIndices(a.own,ghost)
+    OwnAndGhostIndices(a.own,ghost,a.global_to_owner)
 end
 
 function find_owner(indices,global_ids,::Type{<:OwnAndGhostIndices{T}}) where T

src/p_sparse_matrix.jl

@@ -1272,8 +1272,8 @@ function psparse_assemble_impl(
         blocks = split_matrix_blocks(own_own,own_ghost,ghost_own,ghost_ghost)
         values = split_matrix(blocks,local_permutation(rows_fa),local_permutation(rows_fa))
         nnz_own_own = nnz(own_own)
-        k_own_sa = precompute_nzindex(own_own,nz_own_own[1:2]...)
-        k_ghost_sa = precompute_nzindex(own_ghost,nz_own_ghost[1:2]...)
+        k_own_sa = precompute_nzindex(own_own,own_own_triplet[1:2]...)
+        k_ghost_sa = precompute_nzindex(own_ghost,own_ghost_triplet[1:2]...)
         for p in 1:length(I_rcv_own)
             i = I_rcv_own[p]
             j = J_rcv_own[p]

test/p_range_tests.jl

@@ -143,8 +143,10 @@ function p_range_tests(distribute)
            LocalIndices(n,part,[1,3,7,9,10],Int32[1,1,3,4,4])
        end
    end
+   @test renumber_partition(ids) !== ids
 
    ids = uniform_partition(parts,n)
+   @test renumber_partition(ids) === ids
    @test length(ids) == length(parts)
 
    gids = map(parts) do part
@@ -176,6 +178,7 @@ function p_range_tests(distribute)
        end
    end
    ids5 = variable_partition(a,sum(a))
+   @test renumber_partition(ids5) === ids5
    map(parts,ids5) do part, local_to_global
        if part == 1
            @test local_to_global == [1, 2, 3, 4]