Merge pull request #97 from SymbolicML/fix-map-part-2

More improvements to `map` and `similar` on `QuantityArray`

Author: MilesCranmer

src/arrays.jl

@@ -196,6 +196,9 @@ end
 
 Base.similar(A::QuantityArray) = QuantityArray(similar(ustrip(A)), dimension(A), quantity_type(A))
 Base.similar(A::QuantityArray, ::Type{S}) where {S} = QuantityArray(similar(ustrip(A), S), dimension(A), quantity_type(A))
+for (type, _, _) in ABSTRACT_QUANTITY_TYPES
+    @eval Base.similar(A::QuantityArray, ::Type{S}) where {S<:$type} = QuantityArray(similar(ustrip(A), value_type(S)), dimension(A), S)
+end
 
 # Unfortunately this mess of `similar` is required to avoid ambiguous methods.
 # c.f. base/abstractarray.jl
@@ -211,8 +214,10 @@ end
 
 # `_similar_for` in Base does not account for changed dimensions, so
 # we need to overload it for QuantityArray.
-Base._similar_for(c::QuantityArray, ::Type{T}, itr, ::Base.HasShape, axs) where {T} =
+Base._similar_for(c::QuantityArray, ::Type{T}, itr, ::Base.HasShape, axs) where {T<:UnionAbstractQuantity} =
     QuantityArray(similar(ustrip(c), value_type(T), axs), dimension(materialize_first(itr))::dim_type(T), T)
+Base._similar_for(c::QuantityArray, ::Type{T}, itr, ::Base.HasShape, axs) where {T} =
+    similar(ustrip(c), T, axs)
 
 # These methods are not yet implemented, but the default implementation is dangerous,
 # as it may cause a stack overflow, so we raise a more helpful error instead.
@@ -223,8 +228,10 @@ Base._similar_for(::QuantityArray, ::Type{T}, _, ::Base.HasLength, ::Integer) wh
 
 # In earlier Julia, `Base._similar_for` has different signatures.
 @static if hasmethod(Base._similar_for, Tuple{Array,Type,Any,Base.HasShape})
-    @eval Base._similar_for(c::QuantityArray, ::Type{T}, itr, ::Base.HasShape) where {T} =
+    @eval Base._similar_for(c::QuantityArray, ::Type{T}, itr, ::Base.HasShape) where {T<:UnionAbstractQuantity} =
         QuantityArray(similar(ustrip(c), value_type(T), axes(itr)), dimension(materialize_first(itr))::dim_type(T), T)
+    @eval Base._similar_for(c::QuantityArray, ::Type{T}, itr, ::Base.HasShape) where {T} =
+        similar(ustrip(c), T, axes(itr))
 end
 @static if hasmethod(Base._similar_for, Tuple{Array,Type,Any,Base.HasLength})
     @eval Base._similar_for(::QuantityArray, ::Type{T}, _, ::Base.HasLength) where {T} =

test/unittests.jl

@@ -983,6 +983,14 @@ end
             @test dimension(new_qa) == dimension(qa)
             @test isa(ustrip(new_qa), Array{Float32,2})
 
+            if Q !== GenericQuantity
+                new_qa = similar(qa, typeof(GenericQuantity{Float16}(u"km/s")))
+                @test eltype(new_qa) <: GenericQuantity{Float16}
+                @test dim_type(new_qa) == dim_type(qa)
+                @test dimension(new_qa) == dimension(qa)
+                @test isa(ustrip(new_qa), Array{Float16,2})
+            end
+
             new_qa = similar(qa, axes(ones(6, 8)))
             @test size(new_qa) == (6, 8)
             @test eltype(new_qa) <: Q{Float64}
@@ -1104,6 +1112,10 @@ end
             @test prod(qa) == 8.0u"m^3"
             @inferred prod(qa)
 
+            # Map to non-quantity output:
+            @test map(x -> ustrip(x), qa) == fill(2.0, 3)
+            @test map(x -> cos(x/dimension(x)), qa) == fill(cos(2.0), 3)
+
             # Test that we can use a function that returns a different type
             if Q === RealQuantity
                 qa = fill(RealQuantity(2.0u"m"), 3)
@@ -1237,7 +1249,7 @@ end
             # There is no easy way to test whether it actually ran,
             # so we create a fake array type that has a custom `sizehint!`
             # which tells us it actually ran.
-            @eval begin
+            isdefined(Main, :MyCustomArray) || @eval begin
                 mutable struct MyCustomArray{T,N} <: AbstractArray{T,N}
                     data::Array{T,N}
                     sizehint_called::Bool