@@ -4,7 +4,7 @@ module HigherOrderFns
4
4
5
5
# This module provides higher order functions specialized for sparse arrays,
6
6
# particularly map[!]/broadcast[!] for SparseVectors and SparseMatrixCSCs at present.
7
- import Base: map, map!, broadcast, copy, copyto!
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,
@@ -29,6 +29,7 @@ using LinearAlgebra
29
29
# (11) Define broadcast[!] methods handling combinations of scalars, sparse vectors/matrices,
30
30
# structured matrices, and one- and two-dimensional Arrays.
31
31
# (12) Define map[!] methods handling combinations of sparse and structured matrices.
32
+ # (13) Define extrema methods optimized for sparse vectors/matrices.
32
33
33
34
34
35
# (0) BroadcastStyle rules and convenience types for dispatch
@@ -1154,4 +1155,58 @@ map(f::Tf, A::SparseOrStructuredMatrix, Bs::Vararg{SparseOrStructuredMatrix,N})
1154
1155
map! (f:: Tf , C:: AbstractSparseMatrixCSC , A:: SparseOrStructuredMatrix , Bs:: Vararg{SparseOrStructuredMatrix,N} ) where {Tf,N} =
1155
1156
(_checksameshape (C, A, Bs... ); _noshapecheck_map! (f, C, _sparsifystructured (A), map (_sparsifystructured, Bs)... ))
1156
1157
1158
+
1159
+ # (13) extrema methods optimized for sparse vectors/matrices.
1160
+ function _extrema_itr (f, A:: SparseVecOrMat )
1161
+ M = length (A)
1162
+ iszero (M) && throw (ArgumentError (" Sparse array must have at least one element." ))
1163
+ N = nnz (A)
1164
+ iszero (N) && return f (zero (eltype (A))), f (zero (eltype (A)))
1165
+ vmin, vmax = _extrema_itr (f, nonzeros (A))
1166
+ if N != M
1167
+ f0 = f (zero (eltype (A)))
1168
+ vmin = min (f0, vmin)
1169
+ vmax = max (f0, vmax)
1170
+ end
1171
+ vmin, vmax
1172
+ end
1173
+
1174
+ function _extrema_dims (f, x:: SparseVector{Tv, Ti} , dims) where {Tv, Ti}
1175
+ sz = zeros (1 )
1176
+ for d in dims
1177
+ sz[d] = 1
1178
+ end
1179
+ if sz == [1 ] && ! iszero (length (x))
1180
+ return [_extrema_itr (f, x)]
1181
+ end
1182
+ invoke (_extrema_dims, Tuple{Any, AbstractArray, Any}, f, x, dims)
1183
+ end
1184
+
1185
+ function _extrema_dims (f, A:: AbstractSparseMatrix{Tv, Ti} , dims) where {Tv, Ti}
1186
+ sz = zeros (2 )
1187
+ for d in dims
1188
+ sz[d] = 1
1189
+ end
1190
+ if sz == [1 , 0 ] && ! iszero (length (A))
1191
+ B = Array {Tuple{Tv,Tv}} (undef, 1 , size (A, 2 ))
1192
+ @inbounds for col_idx in 1 : size (A, 2 )
1193
+ col = @view A[:,col_idx]
1194
+ fx = (nnz (col) == size (A, 1 )) ? f (A[1 ,col_idx]) : f (zero (Tv))
1195
+ B[col_idx] = (fx, fx)
1196
+ for x in nonzeros (col)
1197
+ fx = f (x)
1198
+ if fx < B[col_idx][1 ]
1199
+ B[col_idx] = (fx, B[col_idx][2 ])
1200
+ elseif fx > B[col_idx][2 ]
1201
+ B[col_idx] = (B[col_idx][1 ], fx)
1202
+ end
1203
+ end
1204
+ end
1205
+ return B
1206
+ end
1207
+ invoke (_extrema_dims, Tuple{Any, AbstractArray, Any}, f, A, dims)
1208
+ end
1209
+
1210
+ _extrema_dims (f, A:: SparseVecOrMat , :: Colon ) = _extrema_itr (f, A)
1211
+
1157
1212
end
0 commit comments