diff --git a/Project.toml b/Project.toml index bbec3e0..a31dbe7 100644 --- a/Project.toml +++ b/Project.toml @@ -26,6 +26,10 @@ Zarr = "0a941bbe-ad1d-11e8-39d9-ab76183a1d99" [weakdeps] Dagger = "d58978e5-989f-55fb-8d15-ea34adc7bf54" +# TODO remove once JuliaIO/DiskArrays.jl#315 is merged and released +[sources] +DiskArrays = {url = "https://github.com/JuliaIO/DiskArrays.jl", rev = "as/backend_fallbacks"} + [extensions] DaggerExt = "Dagger" diff --git a/src/diskarrayresults.jl b/src/diskarrayresults.jl index 91bb083..87eb9c1 100644 --- a/src/diskarrayresults.jl +++ b/src/diskarrayresults.jl @@ -2,7 +2,7 @@ export results_as_diskarrays using DiskArrays: AbstractDiskArray, RegularChunks using OffsetArrays: OffsetArray -struct GMWOPResult{T,N,G<:GMDWop,CS,ISPEC} <: AbstractEngineArray{T,N} +struct GMWOPResult{T,N,G<:GMDWop,CS,ISPEC} <: AbstractDiskArray{T,N} op::G ires::Val{ISPEC} chunksize::CS diff --git a/src/enginearrays.jl b/src/enginearrays.jl index 09fd6b0..506edb8 100644 --- a/src/enginearrays.jl +++ b/src/enginearrays.jl @@ -2,30 +2,9 @@ using DiskArrays using OnlineStats: OnlineStats using Statistics export engine, compute, compute! -abstract type AbstractEngineArray{T,N} <: AbstractDiskArray{T,N} end -struct EngineArray{T,N,P} <: AbstractEngineArray{T,N} - parent::P - bcdims::Base.RefValue{NTuple{N,Int}} -end -EngineArray(p) = EngineArray{eltype(p),ndims(p),typeof(p)}(p,Ref(ntuple(identity,ndims(p)))) -unengine(a::EngineArray) = a.parent -unengine(a) = a -engine(p;bcdims=ntuple(identity,ndims(p))) = EngineArray{eltype(p),ndims(p),typeof(p)}(p,Ref(bcdims)) -engine(p::EngineArray;bcdims=ntuple(identity,ndims(p))) = EngineArray{eltype(p),ndims(p),typeof(p)}(p.parent,bcdims) -function engine(p::AbstractEngineArray;bcdims=nothing) - if bcdims === nothing - p - else - bcdims=ntuple(identity,ndims(p)) - EngineArray{eltype(p),ndims(p),typeof(p)}(p,Ref(bcdims)) - end -end -Base.parent(p::EngineArray) = p.parent -DiskArrays.readblock!(a::EngineArray,xout,r::OrdinalRange...) = DiskArrays.readblock!(a.parent,xout,r...) -DiskArrays.writeblock!(a::EngineArray,xout,r::OrdinalRange...) = DiskArrays.writeblock!(a.parent,xout,r...) -Base.size(a::EngineArray) = size(a.parent) -DiskArrays.eachchunk(a::EngineArray) = DiskArrays.eachchunk(a.parent) -bcdims(a::EngineArray) = a.bcdims[] + +# The backend type is owned by DiskArrays, we only add methods for it +import DiskArrays: DiskArrayEngineBackend bcdims(p) = ntuple(identity,ndims(p)) function collect_bcdims(A) @@ -55,7 +34,29 @@ function collect_bcdims(A) last.(oc), first.(oc) end -function Base.mapreduce(f, op, A::AbstractEngineArray...; dims=:, init = nothing, fin = identity) +function DiskArrays.diskarrays_mapreduce_impl(f, op, a, dims, init, ::DiskArrayEngineBackend) + res = mapreduce_engine(f, op, a; dims, init=neutral_element(f, op, a, init)) + # Like in Base, a reduction over all dimensions returns a scalar + dims === Colon() ? only(compute(res)) : res +end + +# The output buffers are initialized with `init`, so without one we need a neutral element of `op` +neutral_element(f, op, a, init) = init +function neutral_element(f, op, a, ::Base._InitialValue) + T = Base.promote_op(f, eltype(a)) + op === max && return typemin(T) + op === min && return typemax(T) + Base.reduce_empty(op, T) +end + +struct _DAEReducer{F,R} <: Function + f::F + r::R +end +(d::_DAEReducer{F,R})(xout, xin...) where {F,R} = xout[] = d.r(xout[], d.f(xin...)) + + +function mapreduce_engine(f, op, A...; dims=:, init=nothing, fin=identity) s, bcd = collect_bcdims(A) nd = maximum(bcd) ia = map(A) do ar @@ -68,7 +69,7 @@ function Base.mapreduce(f, op, A::AbstractEngineArray...; dims=:, init = nothing fill(1,sloop) end end - InputArray(unengine,windows = lw, dimsmap = bcdims(ar)) + InputArray(ar, windows=lw, dimsmap=bcdims(ar)) end # tf = Base.promote_op(f,Base.nonmissingtype.(eltype.(A))...) # top = Base.promote_op(op,tf,tf) @@ -77,13 +78,13 @@ function Base.mapreduce(f, op, A::AbstractEngineArray...; dims=:, init = nothing # end tf = Base.promote_op(f,eltype.(A)...) top = Base.promote_op(op,tf,tf) - func = create_userfunction(f,top,red=op,init=init,buftype=top,finalize=fin) + func = create_userfunction(_DAEReducer(f, op), top, red=op, init=init, buftype=top, finalize=fin, is_mutating=true, allow_threads=false) if dims === Colon() dims = ntuple(identity,nd) end windows = map(ntuple(identity,nd)) do idim if idim in dims - Repeated(s[idim]) + Repeated(1, s[idim]) else 1:s[idim] end @@ -108,25 +109,28 @@ extrred(x,y) = min(first(x),first(y)),max(last(x),last(y)) wrap_reduction(a) = a wrap_reduction(a::OnlineStats.OnlineStat) = OnlineStats.value(a) -for (f,red,i) in ( - (:(Base.maximum),:max,:init_max), - (:(Base.minimum),:min,:init_min), - (:(Base.extrema),:extrred,:init_ex), - (:(Base.sum),:(+),:init_sum), - ) +for func in (:maximum, :minimum, :sum, :extrema, :median, :mean) + fname = Symbol("diskarrays_$(func)_impl") + # The direct aggregator broadcasts its result into the output, which does not work for the Tuple from `extrema` + default_strategy = QuoteNode(func === :extrema ? :reduce : :auto) eval(quote - $f(a::AbstractEngineArray;dims=:,skipmissing=false) = $f(identity,a;dims,skipmissing) - function $f(ff::Base.Callable,a::AbstractEngineArray;dims=:,skipmissing=false) - red = skipmissing ? missred($red) : $red - init = skipmissing ? missing : $i - mapreduce(ff,red,a;dims,init=$(i)(a)) - end - - -end) + # `a::AbstractDiskArray` avoids an ambiguity with the `::ComputeBackend` fallback in DiskArrays + function DiskArrays.$(fname)(ff, a::A, ::DiskArrayEngineBackend; dims=:, skipmissing=false, strategy=$default_strategy) where {A<:AbstractDiskArray} + if dims === Colon() + dimspec = ntuple(i->i=>nothing, ndims(a)) + res = aggregate_diskarray(a, $func, dimspec; skipmissing, preproc=ff, strategy) + + only(compute(res)) + else + dimspec = ntuple(i->i=>nothing, dims) + aggregate_diskarray(a, $func, dimspec; skipmissing, preproc=ff, strategy) + end + end + end) end -Base.mapslices(f,A::AbstractEngineArray...;dims,outchunks=nothing) = mapslices_engine(f,A;dims,outchunks) +diskarrays_mapslices_impl(f, A...; dims, outchunks=nothing) = + mapslices_engine(Base.Fix1(Statistics.median, f), A; dims, outchunks) function mapslices_engine(f,A...;dims,outchunks=nothing) s, bcd = collect_bcdims(A) @@ -173,14 +177,14 @@ end get_infered_types(f,inputtypes) = Base.promote_op(f,inputtypes...) -Statistics.median(a::AbstractEngineArray;dims=()) = mapslices_engine(median,a;dims) +diskarrays_median_impl(f, a; dims=()) = mapslices_engine(median, a; dims) struct EngineStyle{N} <: Base.Broadcast.AbstractArrayStyle{N} end using DiskArrays: ChunkStyle using Base.Broadcast: DefaultArrayStyle -Base.BroadcastStyle(::Type{<:AbstractEngineArray{<:Any,N}}) where N = EngineStyle{N}() +diskarrays_broadcaststyle(::Type{T}, ::DiskArrayEngineBackend) where T = EngineStyle{ndims(T)}() Base.BroadcastStyle(::EngineStyle{N}, ::EngineStyle{M}) where {N,M} = EngineStyle{max(N, M)}() function Base.BroadcastStyle(::EngineStyle{N}, ::DefaultArrayStyle{M}) where {N,M} return EngineStyle{max(N, M)}() diff --git a/src/util/aggregate.jl b/src/util/aggregate.jl index b69e588..0df33c3 100644 --- a/src/util/aggregate.jl +++ b/src/util/aggregate.jl @@ -83,11 +83,11 @@ This will never happen if you use a do-block. ```julia aggregate_diskarray(a,mean,(2=>nothing,3=>4,4=>[1,1,1,2,2,2,3,3,3,4,4,4,4,4,5,5,5])) """ -function aggregate_diskarray(a, f, dimspec; skipmissing=false, strategy=:auto,outchunks=nothing) +function aggregate_diskarray(a, f, dimspec; skipmissing=false, strategy=:auto, outchunks=nothing, preproc=identity) hasmissings = Missing <: eltype(a) if strategy == :reduce || ((isa(f,DataType) || isa(f,UnionAll)) && f <: OnlineStat) - agg = ReduceAggregator(disk_onlinestat(f)) + agg = ReduceAggregator(disk_onlinestat(f, preproc)) op = gmwop_for_aggregator(agg,dimspec,a;outchunks) results_as_diskarrays(op)[1] elseif strategy == :direct @@ -95,7 +95,7 @@ function aggregate_diskarray(a, f, dimspec; skipmissing=false, strategy=:auto,ou if hasmissings rett = Union{rett,Missing} end - agg = DirectAggregator(create_userfunction(f,rett)) + agg = DirectAggregator(create_userfunction(Base.Fix1(f, preproc), rett)) op = gmwop_for_aggregator(agg,dimspec,a;outchunks) results_as_diskarrays(op)[1] elseif strategy == :auto @@ -103,8 +103,8 @@ function aggregate_diskarray(a, f, dimspec; skipmissing=false, strategy=:auto,ou if hasmissings rett = Union{rett,Missing} end - agg1 = DirectAggregator(create_userfunction(f,rett)) - agg2 = ReduceAggregator(disk_onlinestat(f)) + agg1 = DirectAggregator(create_userfunction(Base.Fix1(f, preproc), rett)) + agg2 = ReduceAggregator(disk_onlinestat(f, preproc)) op1 = gmwop_for_aggregator(agg1,dimspec,a;outchunks) p1 = optimize_loopranges(op1,5e8) diff --git a/src/util/onlinestats_helpers.jl b/src/util/onlinestats_helpers.jl index 0172e34..1bc04a6 100644 --- a/src/util/onlinestats_helpers.jl +++ b/src/util/onlinestats_helpers.jl @@ -44,8 +44,9 @@ has_onlineversion(f) = f in keys(func_to_online) const func_to_online = Dict([ mean => (OnlineStats.Mean,Union{Float64,Missing}), sum => (OnlineStats.Sum,Union{Float64,Missing}), - extrema => (OnlineStats.Extrema, Union{Tuple{Float64,Float64},Missing}), - maximum => (DerivedOnlineStat{OnlineStats.Extrema,v -> last(OnlineStats.value(v)),OnlineStats.fit!,()}, Union{Float64,Missing}), - minimum => (DerivedOnlineStat{OnlineStats.Extrema,v -> first(OnlineStats.value(v)),OnlineStats.fit!,()}, Union{Float64,Missing}), + # `value(::Extrema)` is a NamedTuple `(min, max, nmin, nmax)`, so use the accessors instead of `first`/`last` + extrema => (DerivedOnlineStat{OnlineStats.Extrema,extrema,OnlineStats.fit!,()}, Union{Tuple{Float64,Float64},Missing}), + maximum => (DerivedOnlineStat{OnlineStats.Extrema,maximum,OnlineStats.fit!,()}, Union{Float64,Missing}), + minimum => (DerivedOnlineStat{OnlineStats.Extrema,minimum,OnlineStats.fit!,()}, Union{Float64,Missing}), median => (DerivedOnlineStat{OnlineStats.ExpandingHist,OnlineStats.median,OnlineStats.fit!,(200,)}, Union{Float64,Missing}) ]) diff --git a/test/runtests.jl b/test/runtests.jl index 8556bfc..480f023 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -9,4 +9,5 @@ include("test_distribute.jl") include("test_restart.jl") include("test_graph.jl") include("test_mergeops.jl") -include("test_interpolation.jl") \ No newline at end of file +include("test_interpolation.jl") +include("test_compute_backend.jl") \ No newline at end of file diff --git a/test/test_compute_backend.jl b/test/test_compute_backend.jl new file mode 100644 index 0000000..4d3b97d --- /dev/null +++ b/test/test_compute_backend.jl @@ -0,0 +1,154 @@ +using DiskArrayEngine +using DiskArrays: AbstractDiskArray, DiskArrayEngineBackend, withbackend +using DiskArrays.TestTypes: AccessCountDiskArray +using Statistics +using Test + +import DiskArrayEngine as DAE + +# ── Helpers ───────────────────────────────────────────────────────────────── + +function make_arrays(data; chunksize=size(data)) + mat = data + # Select the backend per array, so the tests do not depend on the backend preference + da = withbackend(AccessCountDiskArray(data; chunksize=chunksize), DiskArrayEngineBackend()) + return (materialized=mat, disk=da) +end + +# ── Scalar reductions ─────────────────────────────────────────────────────── + +function test_scalar_reductions(data; chunksize=ntuple(i -> max(1, size(data)[i] ÷ 2), ndims(data))) + @testset "DiskArrayEngine: scalar reductions" begin + mat, da = make_arrays(data; chunksize) + + @testset "sum" begin + @test sum(da) ≈ sum(mat) + @test sum(identity, da) ≈ sum(mat) + @test sum(x -> 2x, da) ≈ sum(x -> 2x, mat) + end + + @testset "prod" begin + @test prod(da) ≈ prod(mat) + @test prod(identity, da) ≈ prod(mat) + @test prod(x -> 2x, da) ≈ prod(x -> 2x, mat) + end + + @testset "minimum / maximum" begin + @test minimum(da) ≈ minimum(mat) + @test minimum(identity, da) ≈ minimum(mat) + @test minimum(x -> abs(x), da) ≈ minimum(x -> abs(x), mat) + @test maximum(da) ≈ maximum(mat) + @test maximum(identity, da) ≈ maximum(mat) + @test maximum(x -> abs(x), da) ≈ maximum(x -> abs(x), mat) + # The OnlineStats path, `:auto` mostly picks the direct one for these small arrays + @test minimum(da; strategy=:reduce) ≈ minimum(mat) + @test maximum(da; strategy=:reduce) ≈ maximum(mat) + end + + @testset "extrema" begin + @test extrema(da) == extrema(mat) + @test extrema(identity, da) == extrema(mat) + @test extrema(x -> abs(x), da) == extrema(x -> abs(x), mat) + end + + @testset "count" begin + @test count(x -> x > 0, da) == count(x -> x > 0, mat) + end + + @testset "mean" begin + @test mean(da) ≈ mean(mat) + @test mean(identity, da) ≈ mean(mat) + @test mean(x -> 2x, da) ≈ mean(x -> 2x, mat) + end + + @testset "median" begin + @test median(da) ≈ median(mat) + end + end +end + +# ── Lazy result verification ──────────────────────────────────────────────── + +function test_lazy_results(data; chunksize=ntuple(i -> max(1, size(data)[i] ÷ 2), ndims(data))) + @testset "DiskArrayEngine: lazy results" begin + da = make_arrays(data; chunksize).disk + + # sum/mean with dims≠: should return lazy GMWOPResult (DAE's impl + # calls aggregate_diskarray which returns a lazy array when dims is set). + # The DefaultBackend fallback would return a materialized Matrix, so + # checking the type proves the DAE impl path is exercised. + s = sum(da, dims=1) + @test s isa DAE.GMWOPResult + @test Array(s) ≈ sum(data, dims=1) + + m = mean(da, dims=1) + @test m isa DAE.GMWOPResult + @test Array(m) ≈ mean(data, dims=1) + end +end + +# ── mapreduce ─────────────────────────────────────────────────────────────── + +function test_mapreduce(data; chunksize=ntuple(i -> max(1, size(data)[i] ÷ 2), ndims(data))) + @testset "DiskArrayEngine: mapreduce" begin + mat, da = make_arrays(data; chunksize) + + @testset "mapreduce (no dims, no init)" begin + @test mapreduce(x -> 2x, +, da) ≈ mapreduce(x -> 2x, +, mat) + # `op` must be associative, chunks are reduced in an arbitrary order + @test mapreduce(abs, max, da) ≈ mapreduce(abs, max, mat) + end + + @testset "mapreduce (dims=)" begin + @test Array(mapreduce(x -> 2x, +, da; dims=1)) ≈ mapreduce(x -> 2x, +, mat; dims=1) + @test Array(mapreduce(x -> 2x, +, da; dims=(1, 2))) ≈ mapreduce(x -> 2x, +, mat; dims=(1, 2)) + end + + @testset "mapreduce (init)" begin + @test mapreduce(identity, +, da; init=0) ≈ mapreduce(identity, +, mat; init=0) + @test mapreduce(identity, *, da; init=1) ≈ mapreduce(identity, *, mat; init=1) + # An Int `init` with float data and `dims` is an InexactError in Base as well + @test Array(mapreduce(identity, +, da; dims=1, init=0.0)) ≈ mapreduce(identity, +, mat; dims=1, init=0.0) + end + + @testset "mapreducedim!" begin + # Reduce over the last and over the first dimension + redsize(d) = ntuple(i -> i == d ? 1 : size(mat, i), ndims(mat)) + R = zeros(redsize(ndims(mat))) + Base.mapreducedim!(x -> 2x, +, R, da) + @test R ≈ Base.mapreducedim!(x -> 2x, +, zero(R), mat) + R2 = zeros(redsize(1)) + Base.mapreducedim!(x -> x^2, +, R2, da) + @test R2 ≈ Base.mapreducedim!(x -> x^2, +, zero(R2), mat) + end + + @testset "mapfoldl (no init)" begin + @test mapfoldl(x -> 2x, +, da) ≈ mapfoldl(x -> 2x, +, mat) + end + + @testset "mapfoldl (with init)" begin + @test mapfoldl(identity, +, da; init=0) ≈ mapfoldl(identity, +, mat; init=0) + @test mapfoldl(identity, *, da; init=1) ≈ mapfoldl(identity, *, mat; init=1) + end + end +end + +# ── Run tests over multiple array shapes ──────────────────────────────────── + +@testset "DiskArrayEngine: scalar reductions" begin + for data in (randn(5, 4, 2), randn(10), randn(10, 20), randn(3, 3, 3)) + test_scalar_reductions(data) + end +end + +@testset "DiskArrayEngine: lazy results" begin + for data in (randn(5, 4, 2), randn(10), randn(10, 20), randn(3, 3, 3)) + test_lazy_results(data) + end +end + +@testset "DiskArrayEngine: mapreduce" begin + for data in (randn(5, 4, 2), randn(10), randn(10, 20), randn(3, 3, 3)) + test_mapreduce(data) + end +end diff --git a/test/test_interpolation.jl b/test/test_interpolation.jl index e6f60c2..f2521f4 100644 --- a/test/test_interpolation.jl +++ b/test/test_interpolation.jl @@ -26,8 +26,7 @@ end @test size(agg_mean) == (1,5,6) @test agg_mean[:,:,:] == mean(a, dims=1) agg_max = aggregate_diskarray(a, maximum, (2=>nothing,), strategy=:reduce) - # This gives all ones for some reason - @test_broken agg_max[:,:,:] == maximum(a, dims=2) + @test agg_max[:,:,:] == maximum(a, dims=2) agg_sec = aggregate_diskarray(a, mean, (2=>2,)) # This should work but currently throws a bounds error @test_throws BoundsError agg_sec[:,:,:]