@@ -8,7 +8,7 @@ import Base: map, map!, broadcast, copy, copyto!, _extrema_dims, _extrema_itr
8
8
9
9
using Base: front, tail, to_shape
10
10
using .. SparseArrays: SparseVector, SparseMatrixCSC, AbstractSparseVector, AbstractSparseMatrixCSC,
11
- AbstractSparseMatrix, AbstractSparseArray, indtype, nnz, nzrange,
11
+ AbstractSparseMatrix, AbstractSparseArray, indtype, nnz, nzrange, spzeros,
12
12
SparseVectorUnion, AdjOrTransSparseVectorUnion, nonzeroinds, nonzeros, rowvals, getcolptr, widelength
13
13
using Base. Broadcast: BroadcastStyle, Broadcasted, flatten
14
14
using LinearAlgebra
@@ -132,12 +132,17 @@ function trimstorage!(A::SparseVecOrMat, maxstored)
132
132
resize! (storedvals (A), maxstored)
133
133
return maxstored
134
134
end
135
+
135
136
function expandstorage! (A:: SparseVecOrMat , maxstored)
136
- length (storedinds (A)) < maxstored && resize! (storedinds (A), maxstored)
137
- length (storedvals (A)) < maxstored && resize! (storedvals (A), maxstored)
137
+ if length (storedinds (A)) < maxstored
138
+ resize! (storedinds (A), maxstored)
139
+ resize! (storedvals (A), maxstored)
140
+ end
138
141
return maxstored
139
142
end
140
143
144
+ _checkbuffers (S:: SparseMatrixCSC ) = (@assert length (getcolptr (S)) == size (S, 2 ) + 1 && getcolptr (S)[end ] - 1 == length (rowvals (S)) == length (nonzeros (S)); S)
145
+ _checkbuffers (S:: SparseVector ) = (@assert length (storedvals (S)) == length (storedinds (S)); S)
141
146
142
147
# (2) map[!] entry points
143
148
map (f:: Tf , A:: SparseVector ) where {Tf} = _noshapecheck_map (f, A)
@@ -181,7 +186,7 @@ copy(bc::SpBroadcasted1) = _noshapecheck_map(bc.f, bc.args[1])
181
186
storedvals (C)[1 ] = fofnoargs
182
187
broadcast! (f, view (storedvals (C), 2 : length (storedvals (C))))
183
188
end
184
- return C
189
+ return _checkbuffers (C)
185
190
end
186
191
187
192
function _diffshape_broadcast (f:: Tf , A:: SparseVecOrMat , Bs:: Vararg{SparseVecOrMat,N} ) where {Tf,N}
@@ -224,22 +229,13 @@ _maxnnzfrom(shape::NTuple{2}, A::AbstractSparseMatrixCSC) = nnz(A) * div(shape[1
224
229
@inline _unchecked_maxnnzbcres (shape, As... ) = _unchecked_maxnnzbcres (shape, As)
225
230
@inline _checked_maxnnzbcres (shape:: NTuple{1} , As... ) = shape[1 ] != 0 ? _unchecked_maxnnzbcres (shape, As) : 0
226
231
@inline _checked_maxnnzbcres (shape:: NTuple{2} , As... ) = shape[1 ] != 0 && shape[2 ] != 0 ? _unchecked_maxnnzbcres (shape, As) : 0
227
- @inline function _allocres (shape:: NTuple{1} , indextype, entrytype, maxnnz)
228
- storedinds = Vector {indextype} (undef, maxnnz)
229
- storedvals = Vector {entrytype} (undef, maxnnz)
230
- return SparseVector (shape... , storedinds, storedvals)
231
- end
232
- @inline function _allocres (shape:: NTuple{2} , indextype, entrytype, maxnnz)
233
- pointers = ones (indextype, shape[2 ] + 1 )
234
- storedinds = Vector {indextype} (undef, maxnnz)
235
- storedvals = Vector {entrytype} (undef, maxnnz)
236
- return SparseMatrixCSC (shape... , pointers, storedinds, storedvals)
237
- end
232
+ @inline _allocres (shape:: NTuple{1} , indextype, entrytype, maxnnz) = sizehint! (SparseVector (shape[1 ], indextype[], entrytype[]), maxnnz)
233
+ @inline _allocres (shape:: NTuple{2} , indextype, entrytype, maxnnz) = sizehint! (SparseMatrixCSC (shape... , fill (indextype (1 ), shape[2 ]+ 1 ), indextype[], entrytype[]), maxnnz)
238
234
239
235
# (4) _map_zeropres!/_map_notzeropres! specialized for a single sparse vector/matrix
240
236
" Stores only the nonzero entries of `map(f, Array(A))` in `C`."
241
237
function _map_zeropres! (f:: Tf , C:: SparseVecOrMat , A:: SparseVecOrMat ) where Tf
242
- spaceC:: Int = min ( length (storedinds (C)), length ( storedvals (C) ))
238
+ spaceC:: Int = length (nonzeros (C ))
243
239
Ck = 1
244
240
@inbounds for j in columns (C)
245
241
setcolptr! (C, j, Ck)
@@ -255,7 +251,7 @@ function _map_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat) where Tf
255
251
end
256
252
@inbounds setcolptr! (C, numcols (C) + 1 , Ck)
257
253
trimstorage! (C, Ck - 1 )
258
- return C
254
+ return _checkbuffers (C)
259
255
end
260
256
"""
261
257
Densifies `C`, storing `fillvalue` in place of each unstored entry in `A` and
@@ -274,7 +270,7 @@ function _map_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseVecOrMa
274
270
end
275
271
# NOTE: Combining the fill! above into the loop above to avoid multiple sweeps over /
276
272
# nonsequential access of storedvals(C) does not appear to improve performance.
277
- return C
273
+ return _checkbuffers (C)
278
274
end
279
275
# helper functions for these methods and some of those below
280
276
@inline _densecoloffsets (A:: SparseVector ) = 0
297
293
298
294
# (5) _map_zeropres!/_map_notzeropres! specialized for a pair of sparse vectors/matrices
299
295
function _map_zeropres! (f:: Tf , C:: SparseVecOrMat , A:: SparseVecOrMat , B:: SparseVecOrMat ) where Tf
300
- spaceC:: Int = min ( length (storedinds (C)), length ( storedvals (C) ))
296
+ spaceC:: Int = length (nonzeros (C ))
301
297
rowsentinelA = convert (indtype (A), numrows (C) + 1 )
302
298
rowsentinelB = convert (indtype (B), numrows (C) + 1 )
303
299
Ck = 1
@@ -336,7 +332,7 @@ function _map_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat, B::SparseVe
336
332
end
337
333
@inbounds setcolptr! (C, numcols (C) + 1 , Ck)
338
334
trimstorage! (C, Ck - 1 )
339
- return C
335
+ return _checkbuffers (C)
340
336
end
341
337
function _map_notzeropres! (f:: Tf , fillvalue, C:: SparseVecOrMat , A:: SparseVecOrMat , B:: SparseVecOrMat ) where Tf
342
338
# Build dense matrix structure in C, expanding storage if necessary
@@ -368,13 +364,13 @@ function _map_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseVecOrMa
368
364
Cx != fillvalue && (storedvals (C)[jo + Ci] = Cx)
369
365
end
370
366
end
371
- return C
367
+ return _checkbuffers (C)
372
368
end
373
369
374
370
375
371
# (6) _map_zeropres!/_map_notzeropres! for more than two sparse matrices / vectors
376
372
function _map_zeropres! (f:: Tf , C:: SparseVecOrMat , As:: Vararg{SparseVecOrMat,N} ) where {Tf,N}
377
- spaceC:: Int = min ( length (storedinds (C)), length ( storedvals (C) ))
373
+ spaceC:: Int = length (nonzeros (C ))
378
374
rowsentinel = numrows (C) + 1
379
375
Ck = 1
380
376
stopks = _colstartind_all (1 , As)
@@ -398,7 +394,7 @@ function _map_zeropres!(f::Tf, C::SparseVecOrMat, As::Vararg{SparseVecOrMat,N})
398
394
end
399
395
@inbounds setcolptr! (C, numcols (C) + 1 , Ck)
400
396
trimstorage! (C, Ck - 1 )
401
- return C
397
+ return _checkbuffers (C)
402
398
end
403
399
function _map_notzeropres! (f:: Tf , fillvalue, C:: SparseVecOrMat , As:: Vararg{SparseVecOrMat,N} ) where {Tf,N}
404
400
# Build dense matrix structure in C, expanding storage if necessary
@@ -421,7 +417,7 @@ function _map_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, As::Vararg{Spars
421
417
activerow = min (rows... )
422
418
end
423
419
end
424
- return C
420
+ return _checkbuffers (C)
425
421
end
426
422
427
423
# helper methods for map/map! methods just above
462
458
# (7) _broadcast_zeropres!/_broadcast_notzeropres! specialized for a single (input) sparse vector/matrix
463
459
function _broadcast_zeropres! (f:: Tf , C:: SparseVecOrMat , A:: SparseVecOrMat ) where Tf
464
460
isempty (C) && return _finishempty! (C)
465
- spaceC:: Int = min ( length (storedinds (C)), length ( storedvals (C) ))
461
+ spaceC:: Int = length (nonzeros (C ))
466
462
# C and A cannot have the same shape, as we directed that case to map in broadcast's
467
463
# entry point; here we need efficiently handle only heterogeneous C-A combinations where
468
464
# one or both of C and A has at least one singleton dimension.
@@ -509,7 +505,7 @@ function _broadcast_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat) where
509
505
end
510
506
@inbounds setcolptr! (C, numcols (C) + 1 , Ck)
511
507
trimstorage! (C, Ck - 1 )
512
- return C
508
+ return _checkbuffers (C)
513
509
end
514
510
function _broadcast_notzeropres! (f:: Tf , fillvalue, C:: SparseVecOrMat , A:: SparseVecOrMat ) where Tf
515
511
# For information on this code, see comments in similar code in _broadcast_zeropres! above
@@ -540,14 +536,14 @@ function _broadcast_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseV
540
536
end
541
537
end
542
538
end
543
- return C
539
+ return _checkbuffers (C)
544
540
end
545
541
546
542
547
543
# (8) _broadcast_zeropres!/_broadcast_notzeropres! specialized for a pair of (input) sparse vectors/matrices
548
544
function _broadcast_zeropres! (f:: Tf , C:: SparseVecOrMat , A:: SparseVecOrMat , B:: SparseVecOrMat ) where Tf
549
545
isempty (C) && return _finishempty! (C)
550
- spaceC:: Int = min ( length (storedinds (C)), length ( storedvals (C) ))
546
+ spaceC:: Int = length (nonzeros (C ))
551
547
rowsentinelA = convert (indtype (A), numrows (C) + 1 )
552
548
rowsentinelB = convert (indtype (B), numrows (C) + 1 )
553
549
# C, A, and B cannot all have the same shape, as we directed that case to map in broadcast's
@@ -711,7 +707,7 @@ function _broadcast_zeropres!(f::Tf, C::SparseVecOrMat, A::SparseVecOrMat, B::Sp
711
707
end
712
708
@inbounds setcolptr! (C, numcols (C) + 1 , Ck)
713
709
trimstorage! (C, Ck - 1 )
714
- return C
710
+ return _checkbuffers (C)
715
711
end
716
712
function _broadcast_notzeropres! (f:: Tf , fillvalue, C:: SparseVecOrMat , A:: SparseVecOrMat , B:: SparseVecOrMat ) where Tf
717
713
# For information on this code, see comments in similar code in _broadcast_zeropres! above
@@ -810,7 +806,7 @@ function _broadcast_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, A::SparseV
810
806
end
811
807
end
812
808
end
813
- return C
809
+ return _checkbuffers (C)
814
810
end
815
811
_finishempty! (C:: SparseVector ) = C
816
812
_finishempty! (C:: AbstractSparseMatrixCSC ) = (fill! (getcolptr (C), 1 ); C)
861
857
# (9) _broadcast_zeropres!/_broadcast_notzeropres! for more than two (input) sparse vectors/matrices
862
858
function _broadcast_zeropres! (f:: Tf , C:: SparseVecOrMat , As:: Vararg{SparseVecOrMat,N} ) where {Tf,N}
863
859
isempty (C) && return _finishempty! (C)
864
- spaceC:: Int = min ( length (storedinds (C)), length ( storedvals (C) ))
860
+ spaceC:: Int = length (nonzeros (C ))
865
861
expandsverts = _expandsvert_all (C, As)
866
862
expandshorzs = _expandshorz_all (C, As)
867
863
rowsentinel = numrows (C) + 1
@@ -909,7 +905,7 @@ function _broadcast_zeropres!(f::Tf, C::SparseVecOrMat, As::Vararg{SparseVecOrMa
909
905
end
910
906
@inbounds setcolptr! (C, numcols (C) + 1 , Ck)
911
907
trimstorage! (C, Ck - 1 )
912
- return C
908
+ return _checkbuffers (C)
913
909
end
914
910
function _broadcast_notzeropres! (f:: Tf , fillvalue, C:: SparseVecOrMat , As:: Vararg{SparseVecOrMat,N} ) where {Tf,N}
915
911
isempty (C) && return _finishempty! (C)
@@ -950,7 +946,7 @@ function _broadcast_notzeropres!(f::Tf, fillvalue, C::SparseVecOrMat, As::Vararg
950
946
end
951
947
end
952
948
end
953
- return C
949
+ return _checkbuffers (C)
954
950
end
955
951
956
952
# helper method for broadcast/broadcast! methods just above
0 commit comments