Moved pvector stuff to the right place

fverdugo · Jan 10, 2024 · 5f9bc5a · 5f9bc5a
1 parent 34d85a8
commit 5f9bc5a
Show file tree

Hide file tree

Showing 3 changed files with 127 additions and 123 deletions.
diff --git a/src/PartitionedArrays.jl b/src/PartitionedArrays.jl
@@ -114,6 +114,10 @@ export pones
 export prand
 export prandn
 export consistent!
+export assemble
+export consistent
+export repartition
+export repartition!
 include("p_vector.jl")
 
 export PSparseMatrix
@@ -126,8 +130,6 @@ export compress_values
 export compress_values!
 export subassemble
 export subassemble!
-export assemble
-export consistent
 export Disassembled
 export Subassembled
 export Assembled
@@ -136,8 +138,6 @@ export own_ghost_values
 export ghost_own_values
 export own_own_values
 export ghost_ghost_values
-export repartition
-export repartition!
 export psystem
 export psystem!
 export pvector_new

diff --git a/src/p_sparse_matrix.jl b/src/p_sparse_matrix.jl
@@ -1611,65 +1611,6 @@ function IterativeSolvers.zerox(A::PSparseMatrixNew,b::PVector)
     return x
 end
 
-function repartition_cache(v::PVector,new_partition)
-    rows_da = map(remove_ghost,new_partition)
-    row_partition = partition(axes(v,1))
-    I = map(collect∘own_to_global,row_partition)
-    V = own_values(v)
-    I_owner = find_owner(rows_da,I)
-    rows_sa = map(union_ghost,rows_da,I,I_owner)
-    map(map_global_to_local!,I,rows_sa)
-    v_sa = similar(v,PRange(rows_sa))
-    (;I,v_sa)
-end
-
-function repartition(v::PVector,new_partition;reuse=Val(false))
-    w = similar(v,PRange(new_partition))
-    cache = repartition_cache(v,new_partition)
-    t = repartition!(w,v,cache)
-    @async begin
-        wait(t)
-        if val_parameter(reuse) == true
-            w, cache
-        else
-            w
-        end
-    end
-end
-
-function repartition!(
-    w::PVector,v::PVector,cache=repartition_cache(v,partition(axes(w,1)));
-    reversed=false
-    )
-    new_partition = partition(axes(w,1))
-    old_partition = partition(axes(v,1))
-    I = cache.I
-    v_sa = cache.v_sa
-    if ! reversed
-        V = own_values(v)
-        fill!(v_sa,0)
-        map(setindex!,partition(v_sa),V,I)
-        t = assemble!(v_sa)
-        return @async begin
-            wait(t)
-            w .= v_sa
-            w
-        end
-    else
-        v_sa .= v
-        t = consistent!(v_sa)
-        return @async begin
-            wait(t)
-            map(partition(v_sa),partition(w),I) do v_sa,w,I
-                for k in 1:length(I)
-                    w[k] = v_sa[I[k]]
-                end
-            end
-            w
-        end
-    end
-end
-
 function repartition(A::PSparseMatrixNew,new_rows,new_cols;reuse=Val(false))
     function prepare_triplets(A_own_own,A_own_ghost,A_rows,A_cols)
         I1,J1,V1 = findnz(A_own_own)
@@ -1875,66 +1816,6 @@ function psystem!(A,b,V,V2,cache)
     end
 end
 
-function assemble(v::PVector,rows=map(remove_ghost,partition(axes(v,1)));reuse=Val(false))
-    @boundscheck @assert matching_own_indices(axes(v,1),PRange(rows))
-    # TODO this is just a reference implementation
-    # for the moment.
-    # The construction of v2 can (should) be avoided
-    w = similar(v,PRange(rows))
-    v2 = copy(v)
-    t = assemble!(v2)
-    @async begin
-        wait(t)
-        w .= v2
-        if val_parameter(reuse)
-            cache = v2
-            w,cache
-        else
-            w
-        end
-    end
-end
-
-function assemble!(w::PVector,v::PVector,cache)
-    # TODO this is just a reference implementation
-    # for the moment.
-    # The construction of v2 can (should) be avoided
-    v2 = cache
-    map(copy!,local_values(v2),local_values(v))
-    t = assemble!(v2)
-    @async begin
-        wait(t)
-        w .= v2
-        w
-    end
-end
-
-function consistent(v::PVector,rows;reuse=Val(false))
-    # TODO this is just a reference implementation
-    # for the moment. It can be optimized
-    @boundscheck @assert matching_own_indices(axes(v,1),PRange(rows))
-    w = similar(v,PRange(rows))
-    w .= v
-    t = consistent!(w)
-    @async begin
-        wait(t)
-        if val_parameter(reuse)
-            w,nothing
-        else
-            w
-        end
-    end
-end
-
-function consistent!(w::PVector,v::PVector,cache)
-    w .= v
-    t = consistent!(w)
-    @async begin
-        wait(t)
-        w
-    end
-end
-
 # Not efficient, just for convenience and debugging purposes
 function Base.:\(a::PSparseMatrixNew,b::PVector)
     m,n = size(a)

diff --git a/src/p_vector.jl b/src/p_vector.jl
@@ -843,3 +843,126 @@ for M in Distances.metrics
     end
 end
 
+# New stuff
+
+function assemble(v::PVector,rows=map(remove_ghost,partition(axes(v,1)));reuse=Val(false))
+    @boundscheck @assert matching_own_indices(axes(v,1),PRange(rows))
+    # TODO this is just a reference implementation
+    # for the moment.
+    # The construction of v2 can (should) be avoided
+    w = similar(v,PRange(rows))
+    v2 = copy(v)
+    t = assemble!(v2)
+    @async begin
+        wait(t)
+        w .= v2
+        if val_parameter(reuse)
+            cache = v2
+            w,cache
+        else
+            w
+        end
+    end
+end
+
+function assemble!(w::PVector,v::PVector,cache)
+    # TODO this is just a reference implementation
+    # for the moment.
+    # The construction of v2 can (should) be avoided
+    v2 = cache
+    map(copy!,local_values(v2),local_values(v))
+    t = assemble!(v2)
+    @async begin
+        wait(t)
+        w .= v2
+        w
+    end
+end
+
+function consistent(v::PVector,rows;reuse=Val(false))
+    # TODO this is just a reference implementation
+    # for the moment. It can be optimized
+    @boundscheck @assert matching_own_indices(axes(v,1),PRange(rows))
+    w = similar(v,PRange(rows))
+    w .= v
+    t = consistent!(w)
+    @async begin
+        wait(t)
+        if val_parameter(reuse)
+            w,nothing
+        else
+            w
+        end
+    end
+end
+
+function consistent!(w::PVector,v::PVector,cache)
+    w .= v
+    t = consistent!(w)
+    @async begin
+        wait(t)
+        w
+    end
+end
+
+
+function repartition_cache(v::PVector,new_partition)
+    rows_da = map(remove_ghost,new_partition)
+    row_partition = partition(axes(v,1))
+    I = map(collect∘own_to_global,row_partition)
+    V = own_values(v)
+    I_owner = find_owner(rows_da,I)
+    rows_sa = map(union_ghost,rows_da,I,I_owner)
+    map(map_global_to_local!,I,rows_sa)
+    v_sa = similar(v,PRange(rows_sa))
+    (;I,v_sa)
+end
+
+function repartition(v::PVector,new_partition;reuse=Val(false))
+    w = similar(v,PRange(new_partition))
+    cache = repartition_cache(v,new_partition)
+    t = repartition!(w,v,cache)
+    @async begin
+        wait(t)
+        if val_parameter(reuse) == true
+            w, cache
+        else
+            w
+        end
+    end
+end
+
+function repartition!(
+    w::PVector,v::PVector,cache=repartition_cache(v,partition(axes(w,1)));
+    reversed=false
+    )
+    new_partition = partition(axes(w,1))
+    old_partition = partition(axes(v,1))
+    I = cache.I
+    v_sa = cache.v_sa
+    if ! reversed
+        V = own_values(v)
+        fill!(v_sa,0)
+        map(setindex!,partition(v_sa),V,I)
+        t = assemble!(v_sa)
+        return @async begin
+            wait(t)
+            w .= v_sa
+            w
+        end
+    else
+        v_sa .= v
+        t = consistent!(v_sa)
+        return @async begin
+            wait(t)
+            map(partition(v_sa),partition(w),I) do v_sa,w,I
+                for k in 1:length(I)
+                    w[k] = v_sa[I[k]]
+                end
+            end
+            w
+        end
+    end
+end
+
+