Merge branch 'master' into grid

rafaqz · Oct 13, 2019 · 6be2505 · 6be2505
2 parents 2e29fb5 + f9eeb37
commit 6be2505
Show file tree

Hide file tree

Showing 10 changed files with 176 additions and 127 deletions.
diff --git a/Project.toml b/Project.toml
@@ -12,8 +12,10 @@ julia = "1"
 
 [extras]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [targets]
-test = ["BenchmarkTools", "Test", "Unitful"]
+test = ["BenchmarkTools", "Test", "SparseArrays", "Statistics", "Unitful"]
diff --git a/src/array.jl b/src/array.jl
@@ -2,56 +2,60 @@ abstract type AbstractDimensionalArray{T,N,D<:Tuple} <: AbstractArray{T,N} end
 
 const AbDimArray = AbstractDimensionalArray
 
-dims(a::AbDimArray) = a.dims
-label(a::AbDimArray) = ""
+const StandardIndices = Union{AbstractArray,Colon,Integer}
+
+dims(A::AbDimArray) = A.dims
+label(A::AbDimArray) = ""
 
 # Array interface
-Base.size(a::AbDimArray) = size(parent(a))
-Base.iterate(a::AbDimArray, args...) = iterate(parent(a), args...)
-Base.show(io::IO, a::AbDimArray) = begin
-    printstyled(io, "\n", label(a), ": "; color=:red)
-    show(io, typeof(a))
-    show(io, parent(a))
+Base.size(A::AbDimArray) = size(parent(A))
+Base.iterate(A::AbDimArray, args...) = iterate(parent(A), args...)
+Base.show(io::IO, A::AbDimArray) = begin
+    printstyled(io, "\n", label(A), ": "; color=:red)
+    show(io, typeof(A))
+    show(io, parent(A))
     printstyled(io, "\n\ndims:\n"; color=:magenta)
-    show(io, dims(a))
-    show(io, refdims(a))
+    show(io, dims(A))
+    show(io, refdims(A))
     printstyled(io, "\n\nmetadata:\n"; color=:cyan)
 end
 
-Base.@propagate_inbounds Base.getindex(a::AbDimArray, I::Vararg{<:Integer}) =
-    getindex(parent(a), I...)
-Base.@propagate_inbounds Base.getindex(a::AbDimArray, I::Vararg{<:Union{AbstractArray,Colon,Integer}}) =
-    rebuildsliced(a, getindex(parent(a), I...), I)
+Base.@propagate_inbounds Base.getindex(A::AbDimArray, I::Vararg{<:Integer}) =
+    getindex(parent(A), I...)
+Base.@propagate_inbounds Base.getindex(A::AbDimArray, I::Vararg{<:StandardIndices}) =
+    rebuildsliced(A, getindex(parent(A), I...), I)
 
-Base.@propagate_inbounds Base.view(a::AbDimArray, I::Vararg{<:Union{AbstractArray,Colon,Integer}}) =
-    rebuildsliced(a, view(parent(a), I...), I)
+Base.@propagate_inbounds Base.view(A::AbDimArray, I::Vararg{<:StandardIndices}) =
+    rebuildsliced(A, view(parent(A), I...), I)
 
-Base.convert(::Type{Array{T,N}}, a::AbDimArray{T,N}) where {T,N} = 
-    convert(Array{T,N}, parent(a))
+Base.convert(::Type{Array{T,N}}, A::AbDimArray{T,N}) where {T,N} = 
+    convert(Array{T,N}, parent(A))
 
 Base.copy!(dst::AbDimArray, src::AbDimArray) = copy!(parent(src), parent(dst))
 
 # Similar. TODO this need a rethink. How do we know what the new dims are?
 Base.BroadcastStyle(::Type{<:AbDimArray}) = Broadcast.ArrayStyle{AbDimArray}()
 # Need to cover a few type signatures to avoid ambiguity with base
-Base.similar(a::AbDimArray, ::Type{T}, I::Vararg{<:Integer}) where T =
-    rebuildsliced(a, similar(parent(a), T, I...), I)
-
-Base.similar(a::AbDimArray) = rebuild(a, similar(parent(a)))
-Base.similar(a::AbDimArray, ::Type{T}) where T = rebuild(a, similar(parent(a), T))
-Base.similar(a::AbDimArray, ::Type{T}, ::Tuple{Int64,Vararg{Int64}}) where T = 
-    rebuild(a, similar(parent(a), T))
-Base.similar(a::AbDimArray, ::Type{T}, I::Tuple{Union{Integer,OneTo},Vararg{Union{Integer,OneTo},N}}) where {T,N} =
-    rebuildsliced(a, similar(parent(a), T, I...), I)
-Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{AbDimArray}}, ::Type{ElType}) where ElType =
-    rebuildsliced(find_dimensional(bc), similar(Array{ElType}, axes(bc)), axes(bc))
+Base.similar(A::AbDimArray, ::Type{T}, I::Vararg{<:Integer}) where T =
+    rebuildsliced(A, similar(parent(A), T, I...), I)
+
+Base.similar(A::AbDimArray) = rebuild(A, similar(parent(A)))
+Base.similar(A::AbDimArray, ::Type{T}) where T = rebuild(A, similar(parent(A), T))
+Base.similar(A::AbDimArray, ::Type{T}, ::Tuple{Int64,Vararg{Int64}}) where T = 
+    rebuild(A, similar(parent(A), T))
+Base.similar(A::AbDimArray, ::Type{T}, I::Tuple{Union{Integer,OneTo},Vararg{Union{Integer,OneTo},N}}) where {T,N} =
+    rebuildsliced(A, similar(parent(A), T, I...), I)
+Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{AbDimArray}}, ::Type{ElType}) where ElType = begin
+    A = find_dimensional(bc)
+    rebuildsliced(A, similar(Array{ElType}, axes(bc)), axes(bc))
+end
 
 @inline find_dimensional(bc::Base.Broadcast.Broadcasted) = find_dimensional(bc.args)
 @inline find_dimensional(ext::Base.Broadcast.Extruded) = find_dimensional(ext.x)
 @inline find_dimensional(args::Tuple{}) = error("dimensional array not found")
 @inline find_dimensional(args::Tuple) = find_dimensional(find_dimensional(args[1]), tail(args))
 @inline find_dimensional(x) = x
-@inline find_dimensional(a::AbDimArray, rest) = a
+@inline find_dimensional(A::AbDimArray, rest) = A
 @inline find_dimensional(::Any, rest) = find_dimensional(rest)
 
 
@@ -65,18 +69,18 @@ struct DimensionalArray{T,N,D<:Tuple,R<:Tuple,A<:AbstractArray{T,N}} <: Abstract
     dims::D
     refdims::R
 end
-DimensionalArray(a::AbstractArray, dims; refdims=()) = 
-    DimensionalArray(a, formatdims(a, dims), refdims)
+DimensionalArray(A::AbstractArray, dims; refdims=()) = 
+    DimensionalArray(A, formatdims(A, dims), refdims)
 
 # Getters
-refdims(a::DimensionalArray) = a.refdims
+refdims(A::DimensionalArray) = A.refdims
 
 # DimensionalArray interface
-@inline rebuild(a::DimensionalArray, data, dims, refdims) = 
+@inline rebuild(A::DimensionalArray, data, dims, refdims) = 
     DimensionalArray(data, dims, refdims)
 
 # Array interface (AbstractDimensionalArray takes care of everything else)
-Base.parent(a::DimensionalArray) = a.data
+Base.parent(A::DimensionalArray) = A.data
 
-Base.@propagate_inbounds Base.setindex!(a::DimensionalArray, x, I::Vararg{<:Union{AbstractArray,Colon,Real}}) =
-    setindex!(parent(a), x, I...)
+Base.@propagate_inbounds Base.setindex!(A::DimensionalArray, x, I::Vararg{StandardIndices}) =
+    setindex!(parent(A), x, I...)
diff --git a/src/dimension.jl b/src/dimension.jl
@@ -46,7 +46,7 @@ shortname(d::AbDim) = shortname(typeof(d))
 shortname(d::Type{<:AbDim}) = name(d)
 units(dim::AbDim) = metadata(dim) == nothing ? "" : get(metadata(dim), :units, "")
 
-bounds(a, args...) = bounds(dims(a), args...)
+bounds(A, args...) = bounds(dims(A), args...)
 bounds(dims::AbDimTuple, lookupdims::Tuple) = bounds(dims[[dimnum(dims, lookupdims)...]]...)
 bounds(dims::AbDimTuple, dim::DimOrDimType) = bounds(dims[dimnum(dims, dim)])
 bounds(dims::AbDimTuple) = (bounds(dims[1]), bounds(tail(dims))...)
@@ -72,32 +72,32 @@ end
 
 # AbstractArray methods where dims are the dispatch argument
 
-@inline rebuildsliced(a, data, I) = rebuild(a, data, slicedims(a, I)...)
+@inline rebuildsliced(A, data, I) = rebuild(A, data, slicedims(A, I)...)
 
-Base.@propagate_inbounds Base.getindex(a::AbstractArray, dims::Vararg{<:AbDim{<:Number}}) =
-    getindex(a, dims2indices(a, dims)...)
-Base.@propagate_inbounds Base.getindex(a::AbstractArray, dims::Vararg{<:AbstractDimension}) = 
-    getindex(a, dims2indices(a, dims)...)
+Base.@propagate_inbounds Base.getindex(A::AbstractArray, dims::Vararg{<:AbDim{<:Number}}) =
+    getindex(A, dims2indices(A, dims)...)
+Base.@propagate_inbounds Base.getindex(A::AbstractArray, dims::Vararg{<:AbstractDimension}) = 
+    getindex(A, dims2indices(A, dims)...)
 
-Base.@propagate_inbounds Base.setindex!(a::AbstractArray, x, dims::Vararg{<:AbstractDimension}) =
-    setindex!(a, x, dims2indices(a, dims)...)
+Base.@propagate_inbounds Base.setindex!(A::AbstractArray, x, dims::Vararg{<:AbstractDimension}) =
+    setindex!(A, x, dims2indices(A, dims)...)
 
-Base.@propagate_inbounds Base.view(a::AbstractArray, dims::Vararg{<:AbstractDimension}) = 
-    view(a, dims2indices(a, dims)...)
+Base.@propagate_inbounds Base.view(A::AbstractArray, dims::Vararg{<:AbstractDimension}) = 
+    view(A, dims2indices(A, dims)...)
 
-@inline Base.axes(a::AbstractArray, dims::DimOrDimType) = axes(a, dimnum(a, dims))
-@inline Base.size(a::AbstractArray, dims::DimOrDimType) = size(a, dimnum(a, dims))
+@inline Base.axes(A::AbstractArray, dims::DimOrDimType) = axes(A, dimnum(A, dims))
+@inline Base.size(A::AbstractArray, dims::DimOrDimType) = size(A, dimnum(A, dims))
 
 
 #= SplitApplyCombine methods?
 Should allow groupby using dims lookup to make this worth the dependency
 Like group by time/lattitude/height band etc.
 
-SplitApplyCombine.splitdims(a::AbstractArray, dims::AllDimensions) =
-    SplitApplyCombine.splitdims(a, dimnum(a, dims))
+SplitApplyCombine.splitdims(A::AbstractArray, dims::AllDimensions) =
+    SplitApplyCombine.splitdims(A, dimnum(A, dims))
 
-SplitApplyCombine.splitdimsview(a::AbstractArray, dims::AllDimensions) =
-    SplitApplyCombine.splitdimsview(a, dimnum(a, dims))
+SplitApplyCombine.splitdimsview(A::AbstractArray, dims::AllDimensions) =
+    SplitApplyCombine.splitdimsview(A, dimnum(A, dims))
 =#
 
 

diff --git a/src/methods.jl b/src/methods.jl
@@ -1,57 +1,62 @@
 
-# Methods where dims are an argument. This is an experimental approach
-# that might need to change.
-# targeting underscore _methods so we can use dispatch on the dims arg
-
-# TODO hanbdle rebuild in the array dispatch, the dimensions are probably wrong in some cases. 
 for (mod, fname) in ((:Base, :sum), (:Base, :prod), (:Base, :maximum), (:Base, :minimum), (:Statistics, :mean))
     _fname = Symbol('_', fname)
     @eval begin
-        @inline ($mod.$_fname)(a::AbstractArray{T,N}, dims::AllDimensions) where {T,N} =
-            ($mod.$_fname)(a, dimnum(a, dims))
-        @inline ($mod.$_fname)(f, a::AbstractArray{T,N}, dims::AllDimensions) where {T,N} =
-            ($mod.$_fname)(f, a, dimnum(a, dims))
+        @inline ($mod.$_fname)(A::AbstractArray, dims::AllDimensions) =
+            rebuild(A, ($mod.$_fname)(parent(A), dimnum(A, dims)), reducedims(A, dims))
+        @inline ($mod.$_fname)(f, A::AbstractArray, dims::AllDimensions) =
+            rebuild(A, ($mod.$_fname)(f, parent(A), dimnum(A, dims)), reducedims(A, dims))
+        @inline ($mod.$_fname)(A::AbDimArray, dims::Union{Int,Base.Dims}) =
+            rebuild(A, ($mod.$_fname)(parent(A), dims), reducedims(A, dims))
+        @inline ($mod.$_fname)(f, A::AbDimArray, dims::Union{Int,Base.Dims}) =
+            rebuild(A, ($mod.$_fname)(f, parent(A), dims), reducedims(A, dims))
     end
 end
 
 for fname in (:std, :var)
     _fname = Symbol('_', fname)
     @eval begin
-        @inline (Statistics.$_fname)(a::AbstractArray, corrected::Bool, mean, dims::AllDimensions) =
-            (Statistics.$_fname)(a, corrected, mean, dimnum(a, dims))
+        @inline (Statistics.$_fname)(A::AbstractArray, corrected::Bool, mean, dims::AllDimensions) =
+            rebuild(A, (Statistics.$_fname)(A, corrected, mean, dimnum(A, dims)), reducedims(A, dims))
+        @inline (Statistics.$_fname)(A::AbDimArray, corrected::Bool, mean, dims::Union{Int,Base.Dims}) =
+            rebuild(A, (Statistics.$_fname)(parent(A), corrected, mean, dims), reducedims(A, dims))
     end
 end
 
-Statistics._median(a::AbstractArray, dims::AllDimensions) =
-    Statistics._median(a, dimnum(a, dims))
+Statistics._median(A::AbstractArray, dims::AllDimensions) =
+    rebuild(A, Statistics._median(parent(A), dimnum(A, dims)), reducedims(A, dims))
+Statistics._median(A::AbDimArray, dims::Union{Int,Base.Dims}) =
+    rebuild(A, Statistics._median(parent(A), dims), reducedims(A, dims))
+
 Base._mapreduce_dim(f, op, nt::NamedTuple{(),<:Tuple}, A::AbstractArray, dims::AllDimensions) =
-    Base._mapreduce_dim(f, op, nt, A, dimnum(A, dims))
+    rebuild(A, Base._mapreduce_dim(f, op, nt, parent(A), dimnum(A, dims)), reducedims(A, dims))
+Base._mapreduce_dim(f, op, nt::NamedTuple{(),<:Tuple}, A::AbDimArray, dims::Union{Int,Base.Dims}) =
+    rebuild(A, Base._mapreduce_dim(f, op, nt, parent(A), dimnum(A, dims)), reducedims(A, dims))
 # Unfortunately Base/accumulate.jl kwargs methods all force dims to be Integer.
 # accumulate wont work unless that is relaxed, or we copy half of the file here.
 Base._accumulate!(op, B, A, dims::AllDimensions, init::Union{Nothing, Some}) =
     Base._accumulate!(op, B, A, dimnum(A, dims), init)
 
-Base._dropdims(a::AbstractArray, dim::Union{AbDim,Type{<:AbDim}}) = 
-    rebuildsliced(a, Base._dropdims(a, dimnum(a, dim)), dims2indices(a, basetype(dim)(1)))
-Base._dropdims(a::AbstractArray, dims::AbDimTuple) = 
-    rebuildsliced(a, Base._dropdims(a, dimnum(a, dims)), 
-                  dims2indices(a, Tuple((basetype(d)(1) for d in dims))))
-
+Base._dropdims(A::AbstractArray, dim::Union{AbDim,Type{<:AbDim}}) = 
+    rebuildsliced(A, Base._dropdims(A, dimnum(A, dim)), dims2indices(A, basetype(dim)(1)))
+Base._dropdims(A::AbstractArray, dims::AbDimTuple) = 
+    rebuildsliced(A, Base._dropdims(A, dimnum(A, dims)), 
+                  dims2indices(A, Tuple((basetype(d)(1) for d in dims))))
 
 
 # TODO cov, cor mapslices, eachslice, reverse, sort and sort! need _methods without kwargs in base so
 # we can dispatch on dims. Instead we dispatch on array type for now, which means
 # these aren't usefull unless you inherit from AbDimArray.
 
-Base.mapslices(f, a::AbDimArray; dims=1, kwargs...) = begin
-    dimnums = dimnum(a, dims)
-    data = mapslices(f, parent(a); dims=dimnums, kwargs...)
-    rebuildsliced(a, data, dims2indices(a, reducedims(DimensionalData.dims(a, dimnums))))
+Base.mapslices(f, A::AbDimArray; dims=1, kwargs...) = begin
+    dimnums = dimnum(A, dims)
+    data = mapslices(f, parent(A); dims=dimnums, kwargs...)
+    rebuild(A, data, reducedims(A, DimensionalData.dims(A, dimnums)))
 end
 
 # This is copied from base as we can't efficiently wrap this function
 # through the kwarg with a rebuild in the generator. Doing it this way 
-# wierdly makes it 2x faster to use a dim than an integer.
+# wierdly makes it faster to use a dim than an integer.
 if VERSION > v"1.1-"
     Base.eachslice(A::AbDimArray; dims=1, kwargs...) = begin
         if dims isa Tuple && length(dims) == 1 
@@ -65,23 +70,24 @@ if VERSION > v"1.1-"
 end
 
 for fname in (:cor, :cov)
-    @eval Statistics.$fname(a::AbDimArray{T,2}; dims=1, kwargs...) where T = begin
-        newdata = Statistics.$fname(parent(a); dims=dimnum(a, dims), kwargs...)
-        I = dims2indices(a, dims, 1)
-        newdims, newrefdims = slicedims(a, I)
-        rebuild(a, newdata, (newdims[1], newdims[1]), newrefdims)
+    @eval Statistics.$fname(A::AbDimArray{T,2}; dims=1, kwargs...) where T = begin
+        newdata = Statistics.$fname(parent(A); dims=dimnum(A, dims), kwargs...)
+        I = dims2indices(A, dims, 1)
+        newdims, newrefdims = slicedims(A, I)
+        rebuild(A, newdata, (newdims[1], newdims[1]), newrefdims)
     end
 end
 
-@inline Base.reverse(a::AbDimArray{T,N}; dims=1) where {T,N} = begin
-    dnum = dimnum(a, dims)
+@inline Base.reverse(A::AbDimArray{T,N}; dims=1) where {T,N} = begin
+    dnum = dimnum(A, dims)
     # Reverse the dimension. TODO: make this type stable
-    newdims = revdims(DimensionalData.dims(a), dnum)
+    newdims = revdims(DimensionalData.dims(A), dnum)
     # Reverse the data
-    newdata = reverse(parent(a); dims=dnum)
-    rebuild(a, newdata, newdims, refdims(a))
+    newdata = reverse(parent(A); dims=dnum)
+    rebuild(A, newdata, newdims, refdims(A))
 end
 
+# TODO change order after reverse
 @inline revdims(dimstorev::Tuple, dnum) = begin
     dim = dimstorev[end]
     if length(dimstorev) == dnum 
@@ -93,10 +99,10 @@ end
 
 for fname in [:permutedims, :transpose, :adjoint]
     @eval begin
-        @inline Base.$fname(a::AbDimArray{T,2}) where T =
-            rebuild(a, $fname(parent(a)), reverse(dims(a)), refdims(a))
+        @inline Base.$fname(A::AbDimArray{T,2}) where T =
+            rebuild(A, $fname(parent(A)), reverse(dims(A)), refdims(A))
     end
 end
 
-Base.permutedims(a::AbDimArray{T,N}, perm) where {T,N} = 
-    rebuild(a, permutedims(parent(a), dimnum(a, perm)), permutedims(dims(a), perm))
+Base.permutedims(A::AbDimArray{T,N}, perm) where {T,N} = 
+    rebuild(A, permutedims(parent(A), dimnum(A, perm)), permutedims(dims(A), perm))
diff --git a/src/order.jl b/src/order.jl
@@ -22,3 +22,7 @@ Trait indicating that the array or dimension is in the reverse order.
 Selector lookup or plotting will be reversed.
 """
 struct Reverse end
+
+Base.reverse(::Reverse) = Forward()
+Base.reverse(::Forward) = Reverse()
+Base.reverse(o::Order) = Order(revese(dimorder(o)), revese(arrayorder(o)))