Skip to content

Commit

Permalink
Accelerate extrema via mapreduce
Browse files Browse the repository at this point in the history
also fix for NaN and missing
  • Loading branch information
N5N3 committed Dec 30, 2021
1 parent ab4e036 commit b468dcd
Showing 1 changed file with 21 additions and 33 deletions.
54 changes: 21 additions & 33 deletions base/multidimensional.jl
Original file line number Diff line number Diff line change
Expand Up @@ -1746,41 +1746,29 @@ of `A`.
This method requires Julia 1.2 or later.
"""
extrema(f, A::AbstractArray; dims=:) = _extrema_dims(f, A, dims)

_extrema_dims(f, A::AbstractArray, ::Colon) = _extrema_itr(f, A)

function _extrema_dims(f, A::AbstractArray, dims)
sz = size(A)
for d in dims
sz = setindex(sz, 1, d)
end
T = promote_op(f, eltype(A))
B = Array{Tuple{T,T}}(undef, sz...)
return extrema!(f, B, A)
end

@noinline function extrema!(f, B, A)
require_one_based_indexing(B, A)
sA = size(A)
sB = size(B)
for I in CartesianIndices(sB)
fAI = f(A[I])
B[I] = (fAI, fAI)
end
Bmax = CartesianIndex(sB)
@inbounds @simd for I in CartesianIndices(sA)
J = min(Bmax,I)
BJ = B[J]
fAI = f(A[I])
if fAI < BJ[1]
B[J] = (fAI, BJ[2])
elseif fAI > BJ[2]
B[J] = (BJ[1], fAI)
end
_extrema_dims(f, A::AbstractArray, dims) = mapreduce(x -> (fx = f(x); (fx, fx)), _extrema_op, A; dims)
extrema!(B, A) = extrema!(identity, B, A)
extrema!(f, B, A) = mapreduce!(x -> (fx = f(x); (fx, fx)), _extrema_op, B, A)
function _extrema_op((a, b), (c, d))
# if a(c) is unordered, then b(d) === a(c)
ismissing(c) && return c, d # missing is stronger than nan
isunordered(a) && return a, b
isunordered(c) && return c, d
ifelse(c < a, c, a), ifelse(b < d, d, b) # @fastmath min(a,c), max(b,d)
end
function reducedim_init(f, ::typeof(_extrema_op), A::AbstractArray, region)
ri = reduced_indices(A, region)
any(i -> isempty(axes(A, i)), region) && _empty_reduce_error()
A1 = view(A, ri...)
IT = eltype(A)
if missing isa IT
RT = Base._return_type(f, Tuple{nonmissingtype(IT)})
T = Union{RT,Tuple{Missing,Missing}}
else
T = Base._return_type(f, Tuple{IT})
end
return B
map!(f, similar(A1, T), A1)
end
extrema!(B, A) = extrema!(identity, B, A)

# Show for pairs() with Cartesian indices. Needs to be here rather than show.jl for bootstrap order
function Base.showarg(io::IO, r::Iterators.Pairs{<:Integer, <:Any, <:Any, T}, toplevel) where T <: Union{AbstractVector, Tuple}
Expand Down

0 comments on commit b468dcd

Please sign in to comment.