Skipping negative indices in psparse

Author: fverdugo

src/PartitionedArrays.jl

@@ -20,6 +20,8 @@ export nziterator
 export nzindex
 export compresscoo
 export indextype
+export sparse_matrix
+export sparse_matrix!
 include("sparse_utils.jl")
 
 export linear_indices

src/p_sparse_matrix.jl

@@ -1043,7 +1043,7 @@ function psparse(f,row_partition,col_partition;assembled)
 end
 
 function psparse(I,J,V,rows,cols;kwargs...)
-    psparse(sparse,I,J,V,rows,cols;kwargs...)
+    psparse(sparse_matrix,I,J,V,rows,cols;kwargs...)
 end
 
 """
@@ -1063,6 +1063,7 @@ function psparse(f,I,J,V,rows,cols;
         restore_ids = true,
         assembly_neighbors_options_rows = (;),
         assembly_neighbors_options_cols = (;),
+        assembled_rows = nothing,
         reuse=Val(false)
     )
 
@@ -1126,6 +1127,9 @@ function psparse(f,I,J,V,rows,cols;
     elseif subassembled
         rows_sa = rows
         cols_sa = cols
+        if assembled_rows == nothing
+            assembled_rows = map(remove_ghost,rows_sa)
+        end
         if indices === :global
             map(map_global_to_local!,I,rows_sa)
             map(map_global_to_local!,J,cols_sa)
@@ -1145,7 +1149,7 @@ function psparse(f,I,J,V,rows,cols;
             B,cacheB = A,nothing
         end
         if assemble
-            t = PartitionedArrays.assemble(B,rows;reuse=true,assembly_neighbors_options_cols)
+            t = PartitionedArrays.assemble(B,assembled_rows;reuse=true,assembly_neighbors_options_cols)
         else
             t = @async B,cacheB
         end
@@ -1196,7 +1200,7 @@ function psparse!(C,V,cache)
     rows_sa = partition(axes(A,1))
     cols_sa = partition(axes(A,2))
     values_sa = partition(A)
-    map(setcoofast!,values_sa,V,K)
+    map(sparse_matrix!,values_sa,V,K)
     if split_format
         split_format!(B,A,cacheB)
     end
@@ -1905,6 +1909,7 @@ function psystem(I,J,V,I2,V2,rows,cols;
         restore_ids = true,
         assembly_neighbors_options_rows = (;),
         assembly_neighbors_options_cols = (;),
+        assembled_rows = nothing,
         reuse=Val(false)
     )
 
@@ -1913,13 +1918,18 @@ function psystem(I,J,V,I2,V2,rows,cols;
     # It can be optimized to exploit the fact
     # that we want to generate a matrix and a vector
 
+    if assembled_rows === nothing && subassembled
+        assembled_rows = map(remove_ghost,rows)
+    end
+
     t1 = psparse(I,J,V,rows,cols;
             subassembled,
             assembled,
             assemble,
             restore_ids,
             assembly_neighbors_options_rows,
             assembly_neighbors_options_cols,
+            assembled_rows,
             reuse=true)
 
     t2 = pvector(I2,V2,rows;
@@ -1928,6 +1938,7 @@ function psystem(I,J,V,I2,V2,rows,cols;
             assemble,
             restore_ids,
             assembly_neighbors_options_rows,
+            assembled_rows,
             reuse=true)
 
     @async begin

src/p_vector.jl

@@ -579,6 +579,7 @@ function pvector(f,I,V,rows;
         restore_ids = true,
         indices = :global,
         reuse=Val(false),
+        assembled_rows = nothing,
         assembly_neighbors_options_rows = (;)
     )
 
@@ -609,6 +610,9 @@ function pvector(f,I,V,rows;
             t = @async A, nothing
         end
     elseif subassembled
+        if assembled_rows === nothing
+            assembled_rows = map(remove_ghost,rows)
+        end
         rows_sa = rows
         if indices === :global
             map(map_global_to_local!,I,rows_sa)
@@ -622,7 +626,7 @@ function pvector(f,I,V,rows;
         end
         A = PVector(values_sa,rows_sa)
         if assemble
-            t = PartitionedArrays.assemble(A,rows;reuse=true)
+            t = PartitionedArrays.assemble(A,assembled_rows;reuse=true)
         else
             t = @async A, nothing
         end

src/sparse_utils.jl

@@ -333,21 +333,82 @@ Base.@propagate_inbounds Base.getindex(v::EltypeVector{T},i::Integer) where T =
 Base.@propagate_inbounds Base.setindex!(v::EltypeVector,w,i::Integer) = (v.parent[i] = w)
 Base.IndexStyle(::Type{<:EltypeVector{T,V}}) where {T,V} = IndexStyle(V)
 
+# TODO deprecated
+function setcoofast!(A,V,K)
+    sparse_matrix!(A,V,K)
+end
+
+struct FilteredCooVector{F,A,B,C,T} <: AbstractVector{T}
+    f::F
+    I::A
+    J::B
+    V::C
+    function FilteredCooVector(f::F,I::A,J::B,V::C) where {F,A,B,C}
+        T = eltype(C)
+        new{F,A,B,C,T}(f,I,J,V)
+    end
+end
+Base.size(a::FilteredCooVector) = size(a.V)
+Base.IndexStyle(::Type{<:FilteredCooVector}) = IndexLinear()
+Base.@propagate_inbounds function Base.getindex(a::FilteredCooVector,k::Int)
+    i = a.I[k]
+    j = a.J[k]
+    v = a.V[k]
+    if i < 1 || j < 1
+        return a.f(v)
+    end
+    v
+end
+
+function sparse_matrix(I,J,V,m,n;kwargs...)
+    sparse_matrix(sparse,I,J,V,m,n;kwargs...)
+end
+function sparse_matrix(f,I,J,V,m,n;reuse=Val(false),skip_out_of_bounds=true)
+    if !skip_out_of_bounds
+        I2 = I
+        J2 = J
+        V2 = V
+    elseif m*n == 0
+        Ti = eltype(I)
+        T = eltype(V)
+        I2 = Ti[]
+        J2 = Ti[]
+        V2 = Tv[]
+    else
+        I2 = FilteredCooVector(one,I,J,I)
+        J2 = FilteredCooVector(one,I,J,J)
+        V2 = FilteredCooVector(zero,I,J,V)
+    end
+    A = f(I2,J2,V2,m,n)
+    if val_parameter(reuse)
+        K = precompute_nzindex(A,I,J)
+        return A,K
+    end
+    A
+end
+
 function precompute_nzindex(A,I,J)
     K = zeros(Int32,length(I))
     for (p,(i,j)) in enumerate(zip(I,J))
+        if i < 1 || j < 1
+            continue
+        end
         K[p] = nzindex(A,i,j)
     end
     K
 end
 
-function setcoofast!(A,V,K)
-    LinearAlgebra.fillstored!(A,0)
+function sparse_matrix!(A,V,K;reset=true)
+    if reset
+        LinearAlgebra.fillstored!(A,0)
+    end
     A_nz = nonzeros(A)
     for (k,v) in zip(K,V)
+        if k < 1
+            continue
+        end
         A_nz[k] += v
     end
     A
 end
 
-