Added pvector_from_split_blocks

Author: fverdugo

src/PartitionedArrays.jl

@@ -137,6 +137,7 @@ export find_local_indices
 export SplitVector
 export split_vector
 export split_vector_blocks
+export pvector_from_split_blocks
 include("p_vector.jl")
 
 export SplitMatrix

src/p_vector.jl

@@ -129,21 +129,29 @@ function split_vector_blocks(own::A,ghost::A) where A
     SplitVectorBlocks(own,ghost)
 end
 
-struct SplitVector{A,T} <: AbstractVector{T}
+struct SplitVector{A,B,T} <: AbstractVector{T}
     blocks::SplitVectorBlocks{A}
-    permutation::UnitRange{Int32}
+    permutation::B
     function SplitVector(
         blocks::SplitVectorBlocks{A},permutation) where A
         T = eltype(blocks.own)
-        perm = convert(UnitRange{Int32},permutation)
-        new{A,T}(blocks,perm)
+        B = typeof(permutation)
+        new{A,B,T}(blocks,permutation)
     end
 end
 
-function split_vector(blocks::SplitVectorBlocks,permutation::UnitRange)
+function split_vector(blocks::SplitVectorBlocks,permutation)
     SplitVector(blocks,permutation)
 end
 
+function split_vector(
+    own::AbstractVector,
+    ghost::AbstractVector,
+    permutation)
+    blocks = split_vector_blocks(own,ghost)
+    split_vector(blocks,permutation)
+end
+
 Base.IndexStyle(::Type{<:SplitVector}) = IndexLinear()
 Base.size(a::SplitVector) = (length(a.blocks.own)+length(a.blocks.ghost),)
 function Base.getindex(a::SplitVector,local_id::Int)
@@ -974,6 +982,12 @@ function pvector!(C,V,cache)
     end
 end
 
+function pvector_from_split_blocks(own,ghost,row_partition)
+    perms = map(local_permutation,row_partition)
+    values = map(split_vector,own,ghost,perms)
+    PVector(values,row_partition)
+end
+
 function old_pvector!(I,V,index_partition;kwargs...)
     old_pvector!(default_local_values,I,V,index_partition;kwargs...)
 end

test/p_vector_tests.jl

@@ -57,6 +57,9 @@ function p_vector_tests(distribute)
 
     @test a == copy(a)
 
+    ac = pvector_from_split_blocks(own_values(aa),ghost_values(aa),row_partition)
+    @test aa == ac
+
     n = 10
     I,V = map(rank) do rank
         Random.seed!(rank)