fix for ode units

Author: ChrisRackauckas

README.md

@@ -21,7 +21,18 @@ However, broadcast is overloaded to loop in an efficient manner, meaning that
 do not match types. A full array interface is included for completeness, which
 allows this array type to be used in place of a standard array in places where
 such a type stable broadcast may be needed. One example is in heterogeneous
-differential equations for [DifferentialEquations.jl](https://github.com/JuliaDiffEq/DifferentialEquations.jl). 
+differential equations for [DifferentialEquations.jl](https://github.com/JuliaDiffEq/DifferentialEquations.jl).
+
+An `ArrayPartition` acts like a single array. `A[i]` indexes through the first
+array, then the second, etc. all linearly. But `A.x` is where the arrays are stored.
+Thus for
+
+```julia
+using RecursiveArrayTools
+A = ArrayPartition(y,z)
+```
+
+We would have `A.x[1]==y` and `A.x[2]==z`. Broadcasting like `f.(A)` is efficient.
 
 ### Functions
 

src/array_partition.jl

@@ -11,14 +11,45 @@ function ArrayPartition{T}(x,::Type{Val{T}}=Val{false})
 end
 Base.similar(A::ArrayPartition) = ArrayPartition((similar(x) for x in A.x)...)
 Base.similar(A::ArrayPartition,dims::Tuple) = ArrayPartition((similar(x,dim) for (x,dim) in zip(A.x,dims))...)
+Base.similar(A::ArrayPartition,T,dims::Tuple) = ArrayPartition((similar(x,T,dim) for (x,dim) in zip(A.x,dims))...)
 Base.copy(A::ArrayPartition) = Base.similar(A)
 Base.zeros(A::ArrayPartition) = ArrayPartition((zeros(x) for x in A.x)...)
 
+# Special to work with units
+function Base.ones(A::ArrayPartition)
+  B = similar(A::ArrayPartition)
+  for i in eachindex(A.x)
+    B.x[i] .= eltype(A.x[i])(one(first(A.x[i])))
+  end
+  B
+end
+
+Base.:+(A::ArrayPartition, B::ArrayPartition) = ArrayPartition((x .+ y for (x,y) in zip(A.x,B.x))...)
+Base.:+(A::Number, B::ArrayPartition) = ArrayPartition((A .+ x for x in B.x)...)
+Base.:+(A::ArrayPartition, B::Number) = ArrayPartition((B .+ x for x in A.x)...)
+Base.:-(A::ArrayPartition, B::ArrayPartition) = ArrayPartition((x .- y for (x,y) in zip(A.x,B.x))...)
+Base.:-(A::Number, B::ArrayPartition) = ArrayPartition((A .- x for x in B.x)...)
+Base.:-(A::ArrayPartition, B::Number) = ArrayPartition((x .- B for x in A.x)...)
 Base.:*(A::Number, B::ArrayPartition) = ArrayPartition((A .* x for x in B.x)...)
 Base.:*(A::ArrayPartition, B::Number) = ArrayPartition((x .* B for x in A.x)...)
 Base.:/(A::ArrayPartition, B::Number) = ArrayPartition((x ./ B for x in A.x)...)
 Base.:\(A::Number, B::ArrayPartition) = ArrayPartition((x ./ A for x in B.x)...)
 
+if VERSION < v"0.6-"
+  Base.:.+(A::ArrayPartition, B::ArrayPartition) = ArrayPartition((x .+ y for (x,y) in zip(A.x,B.x))...)
+  Base.:.+(A::Number, B::ArrayPartition) = ArrayPartition((A .+ x for x in B.x)...)
+  Base.:.+(A::ArrayPartition, B::Number) = ArrayPartition((B .+ x for x in A.x)...)
+  Base.:.-(A::ArrayPartition, B::ArrayPartition) = ArrayPartition((x .- y for (x,y) in zip(A.x,B.x))...)
+  Base.:.-(A::Number, B::ArrayPartition) = ArrayPartition((A .- x for x in B.x)...)
+  Base.:.-(A::ArrayPartition, B::Number) = ArrayPartition((x .- B for x in A.x)...)
+  Base.:.*(A::ArrayPartition, B::ArrayPartition) = ArrayPartition((x .* y for (x,y) in zip(A.x,B.x))...)
+  Base.:.*(A::Number, B::ArrayPartition) = ArrayPartition((A .* x for x in B.x)...)
+  Base.:.*(A::ArrayPartition, B::Number) = ArrayPartition((x .* B for x in A.x)...)
+  Base.:./(A::ArrayPartition, B::ArrayPartition) = ArrayPartition((x ./ y for (x,y) in zip(A.x,B.x))...)
+  Base.:./(A::ArrayPartition, B::Number) = ArrayPartition((x ./ B for x in A.x)...)
+  Base.:.\(A::Number, B::ArrayPartition) = ArrayPartition((x ./ A for x in B.x)...)
+end
+
 @inline function Base.getindex( A::ArrayPartition,i::Int)
   @boundscheck i > length(A) && throw(BoundsError("Index out of bounds"))
   @inbounds for j in 1:length(A.x)
@@ -52,9 +83,9 @@ recursive_one(A::ArrayPartition) = recursive_one(first(A.x))
 Base.zero(A::ArrayPartition) = zero(first(A.x))
 Base.first(A::ArrayPartition) = first(A.x)
 
-Base.start(A::ArrayPartition) = chain(A.x...)
-Base.next(iter::ArrayPartition,state) = next(state,state)
-Base.done(iter::ArrayPartition,state) = done(state,state)
+Base.start(A::ArrayPartition) = start(chain(A.x...))
+Base.next(A::ArrayPartition,state) = next(chain(A.x...),state)
+Base.done(A::ArrayPartition,state) = done(chain(A.x...),state)
 
 Base.length(A::ArrayPartition) = sum((length(x) for x in A.x))
 Base.size(A::ArrayPartition) = (length(A),)
@@ -72,15 +103,12 @@ add_idxs{T<:ArrayPartition}(::Type{T},expr) = :($(expr).x[i])
 end
 
 @generated function Base.broadcast(f,B::Union{Number,ArrayPartition}...)
-  exs = ((add_idxs(B[i],:(B[$i])) for i in eachindex(B))...)
   arr_idx = 0
   for (i,b) in enumerate(B)
     if b <: ArrayPartition
       arr_idx = i
       break
     end
   end
-  :(for i in eachindex(B[$arr_idx].x)
-    broadcast(f,$(exs...))
-  end)
+  :(A = similar(B[$arr_idx]); broadcast!(f,A,B...); A)
 end