Merge pull request #121 from JordiManyer/bugfix-similar-psparse

Bugfix: `Base.similar` for `PSparseMatrix`

Author: fverdugo

CHANGELOG.md

@@ -10,6 +10,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added 
 
 - Function `partition_from_color`.
+- `Base.copyto!` and `Base.copy!` for `PSparseMatrix`.
+
+### Fixed
+
+- Bugfix: `Base.similar` methods for `PSparseMatrix` not working.
 
 ## [0.3.3] - 2023-08-09
 

src/p_sparse_matrix.jl

@@ -397,15 +397,31 @@ function Base.similar(a::PSparseMatrix,::Type{T},inds::Tuple{<:PRange,<:PRange})
     matrix_partition = map(partition(a),partition(rows),partition(cols)) do values, row_indices, col_indices
         allocate_local_values(values,T,row_indices,col_indices)
     end
-    PSparseMatrix(values,row_partition,col_partition)
+    PSparseMatrix(matrix_partition,partition(rows),partition(cols))
 end
 
 function Base.similar(::Type{<:PSparseMatrix{V}},inds::Tuple{<:PRange,<:PRange}) where V
     rows,cols = inds
-    matrix_partition = map(partition(a),partition(rows),partition(cols)) do values, row_indices, col_indices
+    matrix_partition = map(partition(rows),partition(cols)) do row_indices, col_indices
         allocate_local_values(V,row_indices,col_indices)
     end
-    PSparseMatrix(matrix_partition,row_partition,col_partition)
+    PSparseMatrix(matrix_partition,partition(rows),partition(cols))
+end
+
+function Base.copy!(a::PSparseMatrix,b::PSparseMatrix)
+    @assert size(a) == size(b)
+    copyto!(a,b)
+end
+
+function Base.copyto!(a::PSparseMatrix,b::PSparseMatrix)
+    if partition(axes(a,1)) === partition(axes(b,1)) && partition(axes(a,2)) === partition(axes(b,2))
+        map(copy!,partition(a),partition(b))
+    elseif matching_own_indices(axes(a,1),axes(b,1)) && matching_own_indices(axes(a,2),axes(b,2))
+        map(copy!,own_values(a),own_values(b))
+    else
+        error("Trying to copy a PSparseMatrix into another one with a different data layout. This case is not implemented yet. It would require communications.")
+    end
+    a
 end
 
 function LinearAlgebra.fillstored!(a::PSparseMatrix,v)

test/p_sparse_matrix_tests.jl

@@ -60,6 +60,11 @@ function p_sparse_matrix_tests(distribute)
       @test all( values .== 6 )
     end
 
+    _A = similar(A)
+    _A = similar(A,eltype(A),axes(A))
+    #_A = similar(typeof(A),axes(A)) # This should work, but fails down the line in SparseArrays.jl
+    copy!(_A,A)
+
     LinearAlgebra.fillstored!(A,1.0)
     fill!(x,3.0)
     mul!(b,A,x)