Merge pull request #264 from MilesCranmer/compathelper/new_version/20…

…23-11-03-00-09-25-600-02903764031

CompatHelper: bump compat for DynamicQuantities to 0.10

Project.toml

@@ -35,9 +35,10 @@ SymbolicRegressionJSON3Ext = "JSON3"
 SymbolicRegressionSymbolicUtilsExt = "SymbolicUtils"
 
 [compat]
+Aqua = "0.7"
 Compat = "^4.2"
 DynamicExpressions = "0.13"
-DynamicQuantities = "0.7"
+DynamicQuantities = "0.10"
 JSON3 = "1"
 LineSearches = "7"
 LossFunctions = "0.10, 0.11"

src/DimensionalAnalysis.jl

@@ -2,7 +2,7 @@ module DimensionalAnalysisModule
 
 import DynamicExpressions: Node
 import DynamicQuantities:
-    Quantity, DimensionError, AbstractQuantity, dimension, ustrip, uparse, constructor_of
+    Quantity, DimensionError, AbstractQuantity, dimension, ustrip, uparse, constructorof
 import Tricks: static_hasmethod
 
 import ..CoreModule: Options, Dataset
@@ -75,14 +75,14 @@ for op in (:(Base.:+), :(Base.:-))
             return W($(op)(l.val, r.val), l.wildcard && r.wildcard, false)
         elseif l.wildcard && r.wildcard
             return W(
-                constructor_of(Q)($(op)(ustrip(l), ustrip(r)), typeof(dimension(l))),
+                constructorof(Q)($(op)(ustrip(l), ustrip(r)), typeof(dimension(l))),
                 true,
                 false,
             )
         elseif l.wildcard
-            return W($(op)(constructor_of(Q)(ustrip(l), dimension(r)), r.val), false, false)
+            return W($(op)(constructorof(Q)(ustrip(l), dimension(r)), r.val), false, false)
         elseif r.wildcard
-            return W($(op)(l.val, constructor_of(Q)(ustrip(r), dimension(l))), false, false)
+            return W($(op)(l.val, constructorof(Q)(ustrip(r), dimension(l))), false, false)
         else
             return W(one(Q), false, true)
         end

src/InterfaceDynamicQuantities.jl

@@ -1,14 +1,15 @@
 module InterfaceDynamicQuantitiesModule
 
 import DynamicQuantities:
+    UnionAbstractQuantity,
     AbstractDimensions,
-    AbstractQuantity,
     Dimensions,
     SymbolicDimensions,
     Quantity,
     dimension,
     uparse,
     sym_uparse,
+    dim_type,
     DEFAULT_DIM_BASE_TYPE
 
 """
@@ -21,7 +22,7 @@ argument is a function for parsing strings (in case a string is passed)
 """
 function get_units(args...)
     return error(
-        "Unit information must be passed as one of `AbstractDimensions`, `AbstractQuantity`, `AbstractString`, `Real`.",
+        "Unit information must be passed as one of `AbstractDimensions`, `AbstractQuantity`, `AbstractString`, `Function`.",
     )
 end
 function get_units(_, _, ::Nothing, ::Function)
@@ -43,6 +44,7 @@ end
 function get_units(::Type{T}, ::Type{D}, x::AbstractVector, f::Function) where {T,D}
     return Quantity{T,D}[get_units(T, D, xi, f) for xi in x]
 end
+# TODO: Allow for AbstractQuantity output here
 
 """
     get_si_units(::Type{T}, units)
@@ -62,26 +64,69 @@ function get_sym_units(::Type{T}, units) where {T}
     return get_units(T, SymbolicDimensions{DEFAULT_DIM_BASE_TYPE}, units, sym_uparse)
 end
 
-#! format: off
 """
     get_dimensions_type(A, default_dimensions)
 
 Recursively finds the dimension type from an array, or,
 if no quantity is found, returns the default type.
 """
-function get_dimensions_type(A::AbstractArray, ::Type{D}) where {D}
-    @inbounds for i in eachindex(A)
-        # Look through columns for any dimensions (so we can return the correct type)
-        A[i] isa AbstractQuantity && return typeof(dimension(A[i]))
+function get_dimensions_type(A::AbstractArray, default::Type{D}) where {D}
+    i = findfirst(a -> isa(a, UnionAbstractQuantity), A)
+    if i === nothing
+        return D
+    else
+        return typeof(dimension(A[i]))
     end
-    return D
 end
-function get_dimensions_type(::AbstractArray{T}, ::Type{D}) where {D,T<:Number}
+function get_dimensions_type(
+    ::AbstractArray{Q}, default::Type
+) where {Q<:UnionAbstractQuantity}
+    return dim_type(Q)
+end
+function get_dimensions_type(_, default::Type{D}) where {D}
     return D
 end
-function get_dimensions_type(::AbstractArray{Q}, ::Type{D}) where {Dout,Q<:AbstractQuantity{<:Any,Dout},D}
-    return Dout
+
+# Shortcut for basic numeric types
+function get_dimensions_type(
+    ::AbstractArray{
+        <:Union{
+            Bool,
+            Int8,
+            UInt8,
+            Int16,
+            UInt16,
+            Int32,
+            UInt32,
+            Int64,
+            UInt64,
+            Int128,
+            UInt128,
+            Float16,
+            Float32,
+            Float64,
+            BigFloat,
+            BigInt,
+            ComplexF16,
+            ComplexF32,
+            ComplexF64,
+            Complex{BigFloat},
+            Rational{Int8},
+            Rational{UInt8},
+            Rational{Int16},
+            Rational{UInt16},
+            Rational{Int32},
+            Rational{UInt32},
+            Rational{Int64},
+            Rational{UInt64},
+            Rational{Int128},
+            Rational{UInt128},
+            Rational{BigInt},
+        },
+    },
+    default::Type{D},
+) where {D}
+    return D
 end
-#! format: on
 
 end

src/MLJInterface.jl

@@ -4,7 +4,8 @@ using Optim: Optim
 import MLJModelInterface as MMI
 import DynamicExpressions: eval_tree_array, string_tree, Node
 import DynamicQuantities:
-    AbstractQuantity,
+    QuantityArray,
+    UnionAbstractQuantity,
     AbstractDimensions,
     SymbolicDimensions,
     Quantity,
@@ -260,12 +261,12 @@ function validate_units(X_units, old_X_units)
 end
 
 # TODO: Test whether this conversion poses any issues in data normalization...
-function dimension_fallback(
-    q::Union{<:Quantity{T,<:AbstractDimensions}}, ::Type{D}
-) where {T,D}
+function dimension_with_fallback(q::UnionAbstractQuantity{T}, ::Type{D}) where {T,D}
     return dimension(convert(Quantity{T,D}, q))::D
 end
-dimension_fallback(_, ::Type{D}) where {D} = D()
+function dimension_with_fallback(_, ::Type{D}) where {D}
+    return D()
+end
 function prediction_warn()
     @warn "Evaluation failed either due to NaNs detected or due to unfinished search. Using 0s for prediction."
 end
@@ -307,25 +308,28 @@ function prediction_fallback(
     end
 end
 
-function unwrap_units_single(A::AbstractMatrix{T}, ::Type{D}) where {D,T<:Number}
-    return A, [D() for _ in eachrow(A)]
-end
+compat_ustrip(A::QuantityArray) = ustrip(A)
+compat_ustrip(A) = ustrip.(A)
+
+"""
+    unwrap_units_single(::AbstractArray, ::Type{<:AbstractDimensions})
+    
+Remove units from some features in a matrix, and return, as a tuple,
+(1) the matrix with stripped units, and (2) the dimensions for those features.
+"""
 function unwrap_units_single(A::AbstractMatrix, ::Type{D}) where {D}
-    for (i, row) in enumerate(eachrow(A))
-        allequal(Base.Fix2(dimension_fallback, D).(row)) ||
+    dims = D[dimension_with_fallback(first(row), D) for row in eachrow(A)]
+    @inbounds for (i, row) in enumerate(eachrow(A))
+        all(xi -> dimension_with_fallback(xi, D) == dims[i], row) ||
             error("Inconsistent units in feature $i of matrix.")
     end
-    dims = map(Base.Fix2(dimension_fallback, D) ∘ first, eachrow(A))
-    return stack([ustrip.(row) for row in eachrow(A)]; dims=1), dims
-end
-function unwrap_units_single(v::AbstractVector{T}, ::Type{D}) where {D,T<:Number}
-    return v, D()
+    return stack(compat_ustrip, eachrow(A); dims=1)::AbstractMatrix, dims
 end
 function unwrap_units_single(v::AbstractVector, ::Type{D}) where {D}
-    allequal(Base.Fix2(dimension_fallback, D).(v)) || error("Inconsistent units in vector.")
-    dims = dimension_fallback(first(v), D)
-    v = ustrip.(v)
-    return v, dims
+    dims = dimension_with_fallback(first(v), D)
+    all(xi -> dimension_with_fallback(xi, D) == dims, v) ||
+        error("Inconsistent units in vector.")
+    return compat_ustrip(v)::AbstractVector, dims
 end
 
 function MMI.fitted_params(m::AbstractSRRegressor, fitresult)

test/test_units.jl

@@ -1,9 +1,11 @@
 using SymbolicRegression
-using SymbolicRegression.InterfaceDynamicQuantitiesModule: get_units
+using SymbolicRegression.InterfaceDynamicQuantitiesModule: get_units, get_dimensions_type
+using SymbolicRegression.MLJInterfaceModule: unwrap_units_single
 using SymbolicRegression.DimensionalAnalysisModule:
     violates_dimensional_constraints, @maybe_return_call, WildcardQuantity
 import DynamicQuantities:
     DEFAULT_DIM_BASE_TYPE,
+    RealQuantity,
     Quantity,
     QuantityArray,
     SymbolicDimensions,
@@ -16,6 +18,7 @@ import DynamicQuantities:
     dimension
 using Test
 import MLJBase as MLJ
+import MLJModelInterface as MMI
 
 custom_op(x, y) = x + y
 
@@ -215,14 +218,18 @@ end
             @test dimension(ypred[begin]) == dimension(y[begin])
         end
 
-        # Multiple outputs:
+        # Multiple outputs, and with RealQuantity
         model = MultitargetSRRegressor(;
             binary_operators=[+, *],
             unary_operators=[sqrt, cbrt, abs],
             early_stop_condition=(loss, complexity) -> (loss < 1e-7 && complexity <= 8),
         )
         X = (; x1=randn(128), x2=randn(128))
-        y = (; a=(@. cbrt(ustrip(X.x1)) + sqrt(abs(ustrip(X.x2)))) .* u"kg", b=X.x1)
+        y = (;
+            a=(@. cbrt(ustrip(X.x1)) + sqrt(abs(ustrip(X.x2)))) .* RealQuantity(u"kg"),
+            b=X.x1,
+        )
+        @test typeof(y.a) <: AbstractArray{<:RealQuantity}
         mach = MLJ.machine(model, X, y)
         MLJ.fit!(mach)
         report = MLJ.report(mach)
@@ -239,7 +246,10 @@ end
         # Prediction should have same units:
         ypred = MLJ.predict(mach; rows=1:3)
         @test dimension(ypred.a[begin]) == dimension(y.a[begin])
-        @test typeof(ypred.a[begin]) == typeof(y.a[begin])
+        @test typeof(dimension(ypred.a[begin])) == typeof(dimension(y.a[begin]))
+        # TODO: Should return same quantity as input
+        @test typeof(ypred.a[begin]) <: Quantity
+        @test typeof(y.a[begin]) <: RealQuantity
         VERSION >= v"1.8" &&
             @eval @test(typeof(ypred.b[begin]) == typeof(y.b[begin]), broken = true)
     end
@@ -315,4 +325,28 @@ end
 
     # But method errors are safely caught
     @test test_return_call(+, 1.0, "1.0") === nothing
+
+    # Edge case
+    ## First, what happens if we just pass some data with quantities,
+    ## and some without?
+    data = (a=randn(3), b=fill(us"m", 3), c=fill(u"m/s", 3))
+    Xm_t = MMI.matrix(data; transpose=true)
+    @test typeof(Xm_t) <: Matrix{<:Quantity}
+    _, test_dims = unwrap_units_single(Xm_t, Dimensions)
+    @test test_dims == dimension.([u"1", u"m", u"m/s"])
+    @test test_dims != dimension.([u"m", u"m", u"m"])
+    @inferred unwrap_units_single(Xm_t, Dimensions)
+
+    ## Now, we force promotion to generic `Number` type:
+    data = (a=Number[randn(3)...], b=fill(us"m", 3), c=fill(u"m/s", 3))
+    Xm_t = MMI.matrix(data; transpose=true)
+    @test typeof(Xm_t) === Matrix{Number}
+    _, test_dims = unwrap_units_single(Xm_t, Dimensions)
+    @test test_dims == dimension.([u"1", u"m", u"m/s"])
+    @test_skip @inferred unwrap_units_single(Xm_t, Dimensions)
+
+    # Another edge case
+    ## Should be able to pull it out from array:
+    @test get_dimensions_type(Number[1.0, us"1"], Dimensions) <: SymbolicDimensions
+    @test get_dimensions_type(Number[1.0, 1.0], Dimensions) <: Dimensions
 end