Merge pull request #33 from SymbolicML/arrays-2

Create `QuantityArray <: AbstractArray`

Author: MilesCranmer

Project.toml

@@ -1,21 +1,23 @@
 name = "DynamicQuantities"
 uuid = "06fc5a27-2a28-4c7c-a15d-362465fb6821"
 authors = ["MilesCranmer <[email protected]> and contributors"]
-version = "0.6.4"
+version = "0.7.0"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
-LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
 Tricks = "410a4b4d-49e4-4fbc-ab6d-cb71b17b3775"
 
 [weakdeps]
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 ScientificTypes = "321657f4-b219-11e9-178b-2701a2544e81"
 ScientificTypesBase = "30f210dd-8aff-4c5f-94ba-8e64358c1161"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [extensions]
+DynamicQuantitiesLinearAlgebraExt = "LinearAlgebra"
 DynamicQuantitiesScientificTypesExt = ["ScientificTypes", "ScientificTypesBase"]
 DynamicQuantitiesUnitfulExt = "Unitful"
 
@@ -30,13 +32,15 @@ julia = "1.6"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Ratios = "c84ed2f1-dad5-54f0-aa8e-dbefe2724439"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SaferIntegers = "88634af6-177f-5301-88b8-7819386cfa38"
 ScientificTypes = "321657f4-b219-11e9-178b-2701a2544e81"
 ScientificTypesBase = "30f210dd-8aff-4c5f-94ba-8e64358c1161"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [targets]
-test = ["Test", "Aqua", "Ratios", "SaferIntegers", "SafeTestsets", "ScientificTypes", "ScientificTypesBase", "Unitful"]
+test = ["Test", "Aqua", "LinearAlgebra", "Ratios", "SaferIntegers", "SafeTestsets", "ScientificTypes", "ScientificTypesBase", "StaticArrays", "Unitful"]

benchmark/benchmarks.jl

@@ -3,15 +3,17 @@ using DynamicQuantities
 
 const SUITE = BenchmarkGroup()
 
-SUITE["creation"] = let s = BenchmarkGroup()
+SUITE["Quantity"] = BenchmarkGroup()
+
+SUITE["Quantity"]["creation"] = let s = BenchmarkGroup()
     s["Quantity(x)"] = @benchmarkable Quantity(x) setup = (x = randn()) evals = 1000
     s["Quantity(x, length=y)"] = @benchmarkable Quantity(x, length=y) setup = (x = randn(); y = rand(1:5)) evals = 1000
     s
 end
 
 default() = Quantity(rand(), length=rand(1:5), mass=rand(1:5) // 2)
 
-SUITE["with_numbers"] = let s = BenchmarkGroup()
+SUITE["Quantity"]["with_numbers"] = let s = BenchmarkGroup()
     f1(x, i) = x^i
     s["^int"] = @benchmarkable $f1(x, i) setup = (x = default(); i = rand(1:5)) evals = 1000
     f2(x, y) = x * y
@@ -21,7 +23,7 @@ SUITE["with_numbers"] = let s = BenchmarkGroup()
     s
 end
 
-SUITE["with_self"] = let s = BenchmarkGroup()
+SUITE["Quantity"]["with_self"] = let s = BenchmarkGroup()
     f4(x) = inv(x)
     s["inv"] = @benchmarkable $f4(x) setup = (x = default()) evals = 1000
     f7(x) = ustrip(x)
@@ -31,10 +33,31 @@ SUITE["with_self"] = let s = BenchmarkGroup()
     s
 end
 
-SUITE["with_quantity"] = let s = BenchmarkGroup()
+SUITE["Quantity"]["with_quantity"] = let s = BenchmarkGroup()
     f5(x, y) = x / y
     s["/y"] = @benchmarkable $f5(x, y) setup = (x = default(); y = default()) evals = 1000
     f6(x, y) = x + y
     s["+y"] = @benchmarkable $f6(x, y) setup = (x = default(); y = x + rand() * x) evals = 1000
     s
 end
+
+if @isdefined QuantityArray
+    SUITE["QuantityArray"] = BenchmarkGroup()
+
+    SUITE["QuantityArray"]["broadcasting"] = let s = BenchmarkGroup()
+        N = 10000
+        f9(x) = x^2
+        s["x^2_normal_array"] = @benchmarkable $f9.(arr) setup = (arr = randn($N))
+        s["x^2_quantity_array"] = @benchmarkable $f9.(arr) setup = (arr = QuantityArray(randn($N), u"km/s"))
+        s["x^2_array_of_quantities"] = @benchmarkable $f9.(arr) setup = (arr = randn($N) .* u"km/s")
+        f10(x) = x^4
+        s["x^4_normal_array"] = @benchmarkable $f10.(arr) setup = (arr = randn($N))
+        s["x^4_quantity_array"] = @benchmarkable $f10.(arr) setup = (arr = QuantityArray(randn($N), u"km/s"))
+        s["x^4_array_of_quantities"] = @benchmarkable $f10.(arr) setup = (arr = randn($N) .* u"km/s")
+        f11(x) = x^4 * 0.9 - x * x / 0.3 * x * 0.9 * x
+        s["multi_normal_array"] = @benchmarkable $f11.(arr) setup = (arr = randn($N))
+        s["multi_quantity_array"] = @benchmarkable $f11.(arr) setup = (arr = QuantityArray(randn($N), u"km/s"))
+        s["multi_array_of_quantities"] = @benchmarkable $f11.(arr) setup = (arr = randn($N) .* u"km/s")
+        s
+    end
+end

docs/src/types.md

@@ -17,8 +17,16 @@ AbstractQuantity
 Note also that the `Quantity` object can take a custom `AbstractDimensions`
 as input, so there is often no need to subtype `AbstractQuantity` separately.
 
+## Symbolic dimensions
+
 Another type which subtypes `AbstractDimensions` is `SymbolicDimensions`:
 
 ```@docs
 SymbolicDimensions
 ```
+
+## Arrays
+
+```@docs
+QuantityArray
+```

ext/DynamicQuantitiesLinearAlgebraExt.jl

@@ -0,0 +1,8 @@
+module DynamicQuantitiesLinearAlgebraExt
+
+import LinearAlgebra: norm
+import DynamicQuantities: AbstractQuantity, ustrip, dimension, new_quantity
+
+norm(q::AbstractQuantity, p::Real=2) = new_quantity(typeof(q), norm(ustrip(q), p), dimension(q))
+
+end

src/DynamicQuantities.jl

@@ -1,8 +1,18 @@
 module DynamicQuantities
 
+import TOML: parsefile
+
+const PACKAGE_VERSION = try
+    let project = parsefile(joinpath(pkgdir(@__MODULE__), "Project.toml"))
+        VersionNumber(project["version"])
+    end
+catch
+    VersionNumber(0, 0, 0)
+end
+
 export Units, Constants
 export AbstractQuantity, AbstractDimensions
-export Quantity, Dimensions, SymbolicDimensions, DimensionError
+export Quantity, Dimensions, SymbolicDimensions, QuantityArray, DimensionError
 export ustrip, dimension
 export ulength, umass, utime, ucurrent, utemperature, uluminosity, uamount
 export uparse, @u_str, sym_uparse, @us_str, expand_units
@@ -11,6 +21,7 @@ include("fixed_rational.jl")
 include("types.jl")
 include("utils.jl")
 include("math.jl")
+include("arrays.jl")
 include("units.jl")
 include("constants.jl")
 include("uparse.jl")

src/arrays.jl

@@ -0,0 +1,214 @@
+import Compat: allequal
+
+"""
+    QuantityArray{T,N,D<:AbstractDimensions,Q<:AbstractQuantity,V<:AbstractArray}
+
+An array of quantities with value `value` of type `V` and dimensions `dimensions` of type `D`
+(which are shared across all elements of the array). This is a subtype of `AbstractArray{Q,N}`,
+and so can be used in most places where a normal array would be used, including broadcasting operations.
+
+# Fields
+
+- `value`: The underlying array of values. Access with `ustrip(a)`.
+- `dimensions`: The dimensions of the array. Access with `dimension(a)`.
+
+# Constructors
+
+- `QuantityArray(value::AbstractArray, dimensions::AbstractDimensions)`: Create a `QuantityArray` with value `value` and dimensions `dimensions`.
+- `QuantityArray(value::AbstractArray, quantity::Quantity)`: Create a `QuantityArray` with value `value` and dimensions inferred
+   with `dimension(quantity)`. This is so that you can easily create an array with the units module, like so:
+   ```julia
+   julia> A = QuantityArray(randn(32), 1u"m")
+   ```
+- `QuantityArray(v::AbstractArray{<:AbstractQuantity})`: Create a `QuantityArray` from an array of quantities. This means the following
+  syntax works:
+  ```julia
+  julia> A = QuantityArray(randn(32) .* 1u"km/s")
+  ```
+"""
+struct QuantityArray{T,N,D<:AbstractDimensions,Q<:AbstractQuantity{T,D},V<:AbstractArray{T,N}} <: AbstractArray{Q,N}
+    value::V
+    dimensions::D
+
+    function QuantityArray(v::_V, d::_D, ::Type{_Q}) where {_T,_N,_D<:AbstractDimensions,_Q<:AbstractQuantity,_V<:AbstractArray{_T,_N}}
+        Q_out = constructor_of(_Q){_T,_D}
+        return new{_T,_N,_D,Q_out,_V}(v, d)
+    end
+end
+
+# Construct with a Quantity (easier, as you can use the units):
+QuantityArray(v::AbstractArray; kws...) = QuantityArray(v, DEFAULT_DIM_TYPE(; kws...))
+QuantityArray(v::AbstractArray, d::AbstractDimensions) = QuantityArray(v, d, Quantity)
+QuantityArray(v::AbstractArray, q::AbstractQuantity) = QuantityArray(v .* ustrip(q), dimension(q), typeof(q))
+QuantityArray(v::QA) where {Q<:AbstractQuantity,QA<:AbstractArray{Q}} =
+    let
+        allequal(dimension.(v)) || throw(DimensionError(first(v), v))
+        QuantityArray(ustrip.(v), dimension(first(v)), Q)
+    end
+
+function Base.promote_rule(::Type{QA1}, ::Type{QA2}) where {QA1<:QuantityArray,QA2<:QuantityArray}
+    D = promote_type(dim_type.((QA1, QA2))...)
+    Q = promote_type(quantity_type.((QA1, QA2))...)
+    T = promote_type(value_type.((QA1, QA2))...)
+    V = promote_type(array_type.((QA1, QA2))...)
+    N = ndims(QA1)
+
+    @assert(
+        N == ndims(QA2),
+        "Cannot promote quantity arrays with different dimensions."
+    )
+    @assert(
+        Q <: AbstractQuantity{T,D} && V <: AbstractArray{T},
+        "Incompatible promotion rules between\n    $(QA1)\nand\n    $(QA2)\nPlease convert to a common quantity type first."
+    )
+
+    return QuantityArray{T,N,D,Q,V}
+end
+
+function Base.convert(::Type{QA}, A::QA) where {QA<:QuantityArray}
+    return A
+end
+function Base.convert(::Type{QA1}, A::QA2) where {QA1<:QuantityArray,QA2<:QuantityArray}
+    Q = quantity_type(QA1)
+    V = array_type(QA1)
+    N = ndims(QA1)
+
+    raw_array = Base.Fix1(convert, Q).(A)
+    output = QuantityArray(convert(constructor_of(V){Q,N}, raw_array))
+    # TODO: This will mess with static arrays
+
+    return output::QA1
+end
+
+@inline ustrip(A::QuantityArray) = A.value
+@inline dimension(A::QuantityArray) = A.dimensions
+
+array_type(::Type{<:QuantityArray{T,N,D,Q,V}}) where {T,N,D,Q,V} = V
+array_type(A::QuantityArray) = array_type(typeof(A))
+
+quantity_type(::Type{<:QuantityArray{T,N,D,Q}}) where {T,N,D,Q} = Q
+quantity_type(A::QuantityArray) = quantity_type(typeof(A))
+
+dim_type(::Type{<:QuantityArray{T,N,D}}) where {T,N,D} = D
+dim_type(A::QuantityArray) = dim_type(typeof(A))
+
+value_type(::Type{<:AbstractQuantity{T}}) where {T} = T
+value_type(::Type{<:QuantityArray{T}}) where {T} = T
+value_type(A::Union{<:QuantityArray,<:AbstractQuantity}) = value_type(typeof(A))
+
+# One field:
+for f in (:size, :length, :axes)
+    @eval Base.$f(A::QuantityArray) = $f(ustrip(A))
+end
+
+function Base.getindex(A::QuantityArray, i...)
+    output_value = getindex(ustrip(A), i...)
+    if isa(output_value, AbstractArray)
+        return QuantityArray(output_value, dimension(A), quantity_type(A))
+    else
+        return new_quantity(quantity_type(A), output_value, dimension(A))
+    end
+end
+function Base.setindex!(A::QuantityArray{T,N,D,Q}, v::Q, i...) where {T,N,D,Q<:AbstractQuantity}
+    dimension(A) == dimension(v) || throw(DimensionError(A, v))
+    return unsafe_setindex!(A, v, i...)
+end
+function Base.setindex!(A::QuantityArray{T,N,D,Q}, v::AbstractQuantity, i...) where {T,N,D,Q<:AbstractQuantity}
+    return setindex!(A, convert(Q, v), i...)
+end
+
+unsafe_setindex!(A, v, i...) = setindex!(ustrip(A), ustrip(v), i...)
+
+Base.IndexStyle(::Type{Q}) where {Q<:QuantityArray} = IndexStyle(array_type(Q))
+
+
+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))
+
+# Unfortunately this mess of `similar` is required to avoid ambiguous methods.
+# c.f. base/abstractarray.jl
+for dim_type in (:(Dims), :(Tuple{Union{Integer,Base.OneTo},Vararg{Union{Integer,Base.OneTo}}}), :(Tuple{Integer, Vararg{Integer}}))
+    @eval Base.similar(A::QuantityArray, dims::$dim_type) = QuantityArray(similar(ustrip(A), dims), dimension(A), quantity_type(A))
+    @eval Base.similar(A::QuantityArray, ::Type{S}, dims::$dim_type) where {S} = QuantityArray(similar(ustrip(A), S, dims), dimension(A), quantity_type(A))
+end
+
+Base.BroadcastStyle(::Type{QA}) where {QA<:QuantityArray} = Broadcast.ArrayStyle{QA}()
+
+function Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{QA}}, ::Type{ElType}) where {QA<:QuantityArray,ElType<:AbstractQuantity}
+    T = value_type(ElType)
+    output_array = similar(bc, T)
+    first_output::ElType = materialize_first(bc)
+    return QuantityArray(output_array, dimension(first_output)::dim_type(ElType), ElType)
+end
+function Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{QuantityArray{T,N,D,Q,V}}}, ::Type{ElType}) where {T,N,D,Q,V<:Array{T,N},ElType}
+    return similar(Array{ElType}, axes(bc))
+end
+function Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{QuantityArray{T,N,D,Q,V}}}, ::Type{ElType}) where {T,N,D,Q,V,ElType}
+    # To deal with things like StaticArrays, we need to rely on
+    # only `similar(::Type{ArrayType}, axes)`. We can't specify the
+    # element type in `similar` if we only give the array type.
+    # TODO: However, this results in a redundant allocation.
+    return (_ -> zero(ElType)).(similar(V, axes(bc)))
+end
+
+# Basically, we want to solve a single element to find the output dimension.
+# Then we can put results in the output `QuantityArray`.
+materialize_first(bc::Base.Broadcast.Broadcasted) = bc.f(materialize_first.(bc.args)...)
+
+# Base cases
+materialize_first(q::AbstractQuantity{<:AbstractArray}) = new_quantity(typeof(q), first(ustrip(q)), dimension(q))
+materialize_first(q::AbstractQuantity) = q
+materialize_first(q::QuantityArray) = first(q)
+materialize_first(q::AbstractArray{Q}) where {Q<:AbstractQuantity} = first(q)
+
+# Derived calls
+materialize_first(r::Base.RefValue) = materialize_first(r.x)
+materialize_first(x::Base.Broadcast.Extruded) = materialize_first(x.x)
+materialize_first(args::Tuple) = materialize_first(first(args))
+materialize_first(args::AbstractArray) =
+    let
+        length(args) >= 1 || error("Unexpected broadcast format. Please submit a bug report.")
+        materialize_first(args[begin])
+    end
+materialize_first(::Tuple{}) = error("Unexpected broadcast format. Please submit a bug report.")
+
+# Everything else:
+materialize_first(x) = x
+
+function _print_array_type(io::IO, ::Type{QA}) where {QA<:QuantityArray}
+    return print(io, "QuantityArray(::", array_type(QA), ", ::", quantity_type(QA), ")")
+end
+Base.showarg(io::IO, v::QuantityArray, _) = _print_array_type(io, typeof(v))
+Base.show(io::IO, ::MIME"text/plain", ::Type{QA}) where {QA<:QuantityArray} = _print_array_type(io, QA)
+
+# Other array operations:
+Base.copy(A::QuantityArray) = QuantityArray(copy(ustrip(A)), copy(dimension(A)), quantity_type(A))
+for f in (:cat, :hcat, :vcat)
+    preamble = quote
+        allequal(dimension.(A)) || throw(DimensionError(A[begin], A[begin+1:end]))
+        A = promote(A...)
+        dimensions = dimension(A[begin])
+        Q = quantity_type(A[begin])
+    end
+    if f == :cat
+        @eval function Base.$f(A::QuantityArray...; dims)
+            $preamble
+            return QuantityArray($f(ustrip.(A)...; dims), dimensions, Q)
+        end
+    else
+        @eval function Base.$f(A::QuantityArray...)
+            $preamble
+            return QuantityArray($f(ustrip.(A)...), dimensions, Q)
+        end
+    end
+end
+Base.fill(x::AbstractQuantity, dims::Dims...) = QuantityArray(fill(ustrip(x), dims...), dimension(x), typeof(x))
+Base.fill(x::AbstractQuantity, t::Tuple{}) = QuantityArray(fill(ustrip(x), t), dimension(x), typeof(x))
+
+ulength(q::QuantityArray) = ulength(dimension(q))
+umass(q::QuantityArray) = umass(dimension(q))
+utime(q::QuantityArray) = utime(dimension(q))
+ucurrent(q::QuantityArray) = ucurrent(dimension(q))
+utemperature(q::QuantityArray) = utemperature(dimension(q))
+uluminosity(q::QuantityArray) = uluminosity(dimension(q))
+uamount(q::QuantityArray) = uamount(dimension(q))