Merge pull request #132 from fverdugo/allow_negative_indices

Allow non positive indices (cont.)

src/p_range.jl

@@ -208,6 +208,9 @@ function filter_ghost(indices,gids,owners)
     n_new_ghost = 0
     global_to_ghost_indices = global_to_ghost(indices)
     for (global_i,owner) in zip(gids,owners)
+        if global_i < 1
+            continue
+        end
         if owner != part_owner
             ghost_i = global_to_ghost_indices[global_i]
             if ghost_i == 0 && !(global_i in set)
@@ -221,6 +224,9 @@ function filter_ghost(indices,gids,owners)
     new_ghost_i = 0
     set = Set{Int}()
     for (global_i,owner) in zip(gids,owners)
+        if global_i < 1
+            continue
+        end
         if owner != part_owner
             ghost_i = global_to_ghost_indices[global_i]
             if ghost_i == 0 && !(global_i in set)
@@ -293,6 +299,9 @@ function map_x_to_y!(x_to_y,I,indices)
     local_to_global_indices = x_to_y(indices)
     for k in 1:length(I)
         Ik = I[k]
+        if Ik < 1
+            continue
+        end
         I[k] = local_to_global_indices[Ik]
     end
     I
@@ -1251,9 +1260,23 @@ 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]
+    map(global_ids,indices) do global_ids, indices
+        map_global_to_owner(global_ids,indices.global_to_owner)
+    end
+end
+
+function map_global_to_owner(I,global_to_owner)
+    Ti = eltype(global_to_owner)
+    owners = Vector{Ti}(undef,length(I))
+    for k in 1:length(I)
+        i = I[k]
+        if i<1
+            owners[k] = zero(Ti)
+            continue
+        end
+        owners[k] = global_to_owner[i]
     end
+    owners
 end
 
 part_id(a::OwnAndGhostIndices) = a.own.owner
@@ -1567,7 +1590,7 @@ function find_owner(indices,global_ids,::Type{<:LocalIndicesWithConstantBlockSiz
             start
         end
         global_to_owner = BlockPartitionGlobalToOwner(start)
-        global_to_owner[global_ids]
+        map_global_to_owner(global_ids,global_to_owner)
     end
 end
 
@@ -1600,7 +1623,7 @@ function find_owner(indices,global_ids,::Type{<:LocalIndicesWithVariableBlockSiz
             start = vcat(initial,[n+1])
         end
         global_to_owner = BlockPartitionGlobalToOwner(start)
-        global_to_owner[global_ids]
+        map_global_to_owner(global_ids,global_to_owner)
     end
 end
 

src/p_sparse_matrix.jl

@@ -1301,18 +1301,6 @@ function psparse_assemble_impl(
         k_rcv = JaggedArray(k_rcv_data,I_rcv.ptrs)
         (;I_rcv,J_rcv,V_rcv,k_rcv,parts_rcv)
     end
-    function setup_touched_col_ids(A,cache_rcv,cols_sa)
-        J_rcv_data = cache_rcv.J_rcv.data
-        l1 = nnz(A.own_ghost)
-        l2 = length(J_rcv_data)
-        J_aux = zeros(Int,l1+l2)
-        ghost_to_global_col = ghost_to_global(cols_sa)
-        for (p,(_,j,_)) in enumerate(nziterator(A.own_ghost))
-            J_own_ghost[p] = ghost_to_global_col[j]
-        end
-        J_aux[l1.+(1:l2)] = J_rcv_data
-        J_aux
-    end
     function setup_own_triplets(A,cache_rcv,rows_sa,cols_sa)
         nz_own_own = findnz(A.blocks.own_own)
         nz_own_ghost = findnz(A.blocks.own_ghost)

src/p_vector.jl

@@ -555,11 +555,24 @@ function dense_vector(I,V,n)
     T = eltype(V)
     a = zeros(T,n)
     for (i,v) in zip(I,V)
+        if i < 1
+            continue
+        end
         a[i] += v
     end
     a
 end
 
+function dense_vector!(A,K,V)
+    fill!(A,0)
+    for (k,v) in zip(K,V)
+        if k < 1
+            continue
+        end
+        A[k] += v
+    end
+end
+
 function pvector(I,V,rows;kwargs...)
     pvector(dense_vector,I,V,rows;kwargs...)
 end
@@ -665,16 +678,10 @@ end
     pvector!(B::PVector,V,cache)
 """
 function pvector!(B,V,cache)
-    function update!(A,K,V)
-        fill!(A,0)
-        for (k,v) in zip(K,V)
-            A[k] += v
-        end
-    end
     (A,cacheB,assemble,assembled,K) = cache
     rows_sa = partition(axes(A,1))
     values_sa = partition(A)
-    map(update!,values_sa,K,V)
+    map(dense_vector!,values_sa,K,V)
     if !assembled && assemble
         t = PartitionedArrays.assemble!(B,A,cacheB)
     else

test/p_sparse_matrix_tests.jl

@@ -206,11 +206,11 @@ function p_sparse_matrix_tests(distribute)
         if part == 1
             [1,2,1,2,2], [2,6,1,2,1], [1.0,2.0,30.0,10.0,1.0]
         elseif part == 2
-            [3,3,4,6], [3,9,4,2], [10.0,2.0,30.0,2.0]
+            [3,3,4,6,0], [3,9,4,2,0], [10.0,2.0,30.0,2.0,2.0]
         elseif part == 3
             [5,5,6,7], [5,6,6,7], [10.0,2.0,30.0,1.0]
         else
-            [9,9,8,10,6], [9,3,8,10,5], [10.0,2.0,30.0,50.0,2.0]
+            [9,9,8,10,6,-1], [9,3,8,10,5,1], [10.0,2.0,30.0,50.0,2.0,1.0]
         end
     end |> tuple_of_arrays