Refactor Unitful.jl usage to use package extensions

Project.toml

@@ -64,7 +64,6 @@ SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 URIs = "5c2747f8-b7ea-4ff2-ba2e-563bfd36b1d4"
 UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
-Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [weakdeps]
 BifurcationKit = "0f109fa4-8a5d-4b75-95aa-f515264e7665"
@@ -74,6 +73,7 @@ FMI = "14a09403-18e3-468f-ad8a-74f8dda2d9ac"
 InfiniteOpt = "20393b10-9daf-11e9-18c9-8db751c92c57"
 LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 Pyomo = "0e8e1daf-01b5-4eba-a626-3897743a3816"
+Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [extensions]
 MTKBifurcationKitExt = "BifurcationKit"
@@ -83,6 +83,7 @@ MTKFMIExt = "FMI"
 MTKInfiniteOptExt = "InfiniteOpt"
 MTKLabelledArraysExt = "LabelledArrays"
 MTKPyomoDynamicOptExt = "Pyomo"
+ModelingToolkitUnitfulExt = "Unitful"
 
 [compat]
 ADTypes = "1.14.0"
@@ -165,7 +166,6 @@ SymbolicUtils = "3.26.1"
 Symbolics = "6.40"
 URIs = "1"
 UnPack = "0.1, 1.0"
-Unitful = "1.1"
 julia = "1.9"
 
 [extras]
@@ -205,6 +205,7 @@ StochasticDelayDiffEq = "29a0d76e-afc8-11e9-03a4-eda52ae4b960"
 StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [targets]
-test = ["AmplNLWriter", "BenchmarkTools", "BoundaryValueDiffEqMIRK", "BoundaryValueDiffEqAscher", "ControlSystemsBase", "DataInterpolations", "DelayDiffEq", "NonlinearSolve", "ForwardDiff", "Ipopt", "Ipopt_jll", "ModelingToolkitStandardLibrary", "Optimization", "OptimizationOptimJL", "OptimizationMOI", "OrdinaryDiffEq", "OrdinaryDiffEqCore", "OrdinaryDiffEqDefault", "REPL", "Random", "ReferenceTests", "SafeTestsets", "StableRNGs", "Statistics", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials", "StochasticDelayDiffEq", "Pkg", "JET", "OrdinaryDiffEqNonlinearSolve", "Logging", "OptimizationBase", "LinearSolve"]
+test = ["AmplNLWriter", "BenchmarkTools", "BoundaryValueDiffEqMIRK", "BoundaryValueDiffEqAscher", "ControlSystemsBase", "DataInterpolations", "DelayDiffEq", "NonlinearSolve", "ForwardDiff", "Ipopt", "Ipopt_jll", "ModelingToolkitStandardLibrary", "Optimization", "OptimizationOptimJL", "OptimizationMOI", "OrdinaryDiffEq", "OrdinaryDiffEqCore", "OrdinaryDiffEqDefault", "REPL", "Random", "ReferenceTests", "SafeTestsets", "StableRNGs", "Statistics", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials", "StochasticDelayDiffEq", "Pkg", "JET", "OrdinaryDiffEqNonlinearSolve", "Logging", "OptimizationBase", "LinearSolve", "Unitful"]

ext/ModelingToolkitUnitfulExt.jl

@@ -0,0 +1,103 @@
+module ModelingToolkitUnitfulExt
+
+using ModelingToolkit, Symbolics, SciMLBase, Unitful, RecursiveArrayTools
+using ModelingToolkit: ValidationError, Connection, instream, JumpType, VariableUnit,
+                         get_systems, Conditional, Comparison, Integral, Differential
+using JumpProcesses: MassActionJump, ConstantRateJump, VariableRateJump
+using Symbolics: Symbolic, value, issym, isadd, ismul, ispow, iscall, operation, arguments, getmetadata
+
+using Unitful
+using SciMLBase
+
+# Import necessary types and functions from ModelingToolkit
+import ModelingToolkit: ValidationError, _get_unittype, get_unit, screen_unit, 
+                        equivalent, _is_dimension_error, convert_units, check_units
+
+const MT = ModelingToolkit
+
+# Add Unitful-specific unit type detection
+function MT._get_unittype(u::Unitful.Unitlike)
+    return Val(:Unitful)
+end
+
+MT._oneunit(x::Unitful.FreeUnits) = 1 * x
+
+# Base operations for mixing Symbolic and Unitful
+Base.:*(x::Union{MT.Num, Symbolic}, y::Unitful.AbstractQuantity) = x * y
+Base.:/(x::Union{MT.Num, Symbolic}, y::Unitful.AbstractQuantity) = x / y
+
+# Unitful-specific get_unit method
+function MT.get_unit(x::Unitful.Quantity)
+    return screen_unit(Unitful.unit(x))
+end
+
+# Unitful-specific screen_unit method
+function MT.screen_unit(result::Unitful.Unitlike)
+    result isa Unitful.ScalarUnits ||
+        throw(ValidationError("Non-scalar units such as $result are not supported. Use a scalar unit instead."))
+    result == Unitful.u"°" &&
+        throw(ValidationError("Degrees are not supported. Use radians instead."))
+    return result
+end
+
+# Unitful-specific equivalence check
+function MT.equivalent(x::Unitful.Unitlike, y::Unitful.Unitlike)
+    return isequal(1 * x, 1 * y)
+end
+
+# Mixed equivalence checks
+MT.equivalent(x::Unitful.Unitlike, y) = isequal(1 * x, y)
+MT.equivalent(x, y::Unitful.Unitlike) = isequal(x, 1 * y)
+
+# The safe_get_unit function stays in the main package and already handles DQ.DimensionError
+# We just need to make sure it can handle Unitful.DimensionError too
+# This will be handled by the main function's MethodError catch
+
+# Unitful-specific dimension error detection for model parsing
+MT._is_dimension_error(e::Unitful.DimensionError) = true
+
+# Unitful-specific convert_units methods for model parsing
+function MT.convert_units(varunits::Unitful.FreeUnits, value)
+    Unitful.ustrip(varunits, value)
+end
+
+MT.convert_units(::Unitful.FreeUnits, value::MT.NoValue) = MT.NO_VALUE
+
+function MT.convert_units(varunits::Unitful.FreeUnits, value::AbstractArray{T}) where {T}
+    Unitful.ustrip.(varunits, value)
+end
+
+MT.convert_units(::Unitful.FreeUnits, value::MT.Num) = value
+
+# Unitful-specific check_units method
+function MT.check_units(::Val{:Unitful}, eqs...)
+    # Use the main package's validate function
+    MT.validate(eqs...) ||
+        throw(ValidationError("Some equations had invalid units. See warnings for details."))
+end
+
+# Define Unitful time variables (moved from main module)
+const t_unitful = let
+    MT.only(MT.@independent_variables t [unit = Unitful.u"s"])
+end
+const D_unitful = MT.Differential(t_unitful)
+
+
+"""
+Throw exception on invalid unit types, otherwise return argument.
+"""
+function screen_unit(result::Unitful.AbstractQuantity)
+    result isa Unitful.Unitlike ||
+        throw(ValidationError("Unit must be a subtype of Unitful.Unitlike, not $(typeof(result))."))
+    result isa Unitful.ScalarUnits ||
+        throw(ValidationError("Non-scalar units such as $result are not supported. Use a scalar unit instead."))
+    result == Unitful.u"°" &&
+        throw(ValidationError("Degrees are not supported. Use radians instead."))
+    result
+end
+
+const unitless = Unitful.unit(1)
+
+
+
+end # module UnitfulUnitCheck

src/ModelingToolkit.jl

@@ -30,7 +30,7 @@ using InteractiveUtils
 using JumpProcesses
 using DataStructures
 using Base.Threads
-using Latexify, Unitful, ArrayInterface
+using Latexify, ArrayInterface
 using Setfield, ConstructionBase
 import Libdl
 using DocStringExtensions
@@ -93,7 +93,7 @@ export independent_variables, unknowns, observables, parameters, full_parameters
 @reexport using UnPack
 RuntimeGeneratedFunctions.init(@__MODULE__)
 
-import DynamicQuantities, Unitful
+import DynamicQuantities
 const DQ = DynamicQuantities
 
 import DifferentiationInterface as DI
@@ -232,15 +232,11 @@ include("deprecations.jl")
 const t_nounits = let
     only(@independent_variables t)
 end
-const t_unitful = let
-    only(@independent_variables t [unit = Unitful.u"s"])
-end
 const t = let
     only(@independent_variables t [unit = DQ.u"s"])
 end
 
 const D_nounits = Differential(t_nounits)
-const D_unitful = Differential(t_unitful)
 const D = Differential(t)
 
 export ODEFunction, convert_system_indepvar,

src/systems/model_parsing.jl

@@ -890,17 +890,7 @@ function convert_units(
         DynamicQuantities.SymbolicUnits.as_quantity(varunits), value))
 end
 
-function convert_units(varunits::Unitful.FreeUnits, value)
-    Unitful.ustrip(varunits, value)
-end
-
-convert_units(::Unitful.FreeUnits, value::NoValue) = NO_VALUE
-
-function convert_units(varunits::Unitful.FreeUnits, value::AbstractArray{T}) where {T}
-    Unitful.ustrip.(varunits, value)
-end
 
-convert_units(::Unitful.FreeUnits, value::Num) = value
 
 convert_units(::DynamicQuantities.Quantity, value::Num) = value
 
@@ -919,8 +909,7 @@ function parse_variable_arg(dict, mod, arg, varclass, kwargs, where_types)
                     try
                         $setdefault($vv, $convert_units($unit, $name))
                     catch e
-                        if isa(e, $(DynamicQuantities.DimensionError)) ||
-                           isa(e, $(Unitful.DimensionError))
+                        if $_is_dimension_error(e)
                             error("Unable to convert units for \'" * string(:($$vv)) * "\'")
                         elseif isa(e, MethodError)
                             error("No or invalid units provided for \'" * string(:($$vv)) *

src/systems/unit_check.jl

@@ -19,14 +19,13 @@ function __get_literal_unit(x)
     u = getmetadata(v, VariableUnit, nothing)
     u isa DQ.AbstractQuantity ? screen_unit(u) : u
 end
+# Extensible unit type detection - extensions can add methods for specific unit types
+_get_unittype(u) = nothing
+_get_unittype(u::DQ.AbstractQuantity) = Val(:DynamicQuantities)
+
 function __get_scalar_unit_type(v)
     u = __get_literal_unit(v)
-    if u isa DQ.AbstractQuantity
-        return Val(:DynamicQuantities)
-    elseif u isa Unitful.Unitlike
-        return Val(:Unitful)
-    end
-    return nothing
+    return _get_unittype(u)
 end
 function __get_unit_type(vs′...)
     for vs in vs′
@@ -44,40 +43,45 @@ function __get_unit_type(vs′...)
     return nothing
 end
 
-function screen_unit(result)
-    if result isa DQ.AbstractQuantity
-        d = DQ.dimension(result)
-        if d isa DQ.Dimensions
-            return result
-        elseif d isa DQ.SymbolicDimensions
-            return DQ.uexpand(oneunit(result))
-        else
-            throw(ValidationError("$result doesn't have a recognized unit"))
-        end
+
+function screen_unit(result::DQ.AbstractQuantity)
+    d = DQ.dimension(result)
+    if d isa DQ.Dimensions
+        return result
+    elseif d isa DQ.SymbolicDimensions
+        return DQ.uexpand(oneunit(result))
     else
-        throw(ValidationError("$result doesn't have any unit."))
+        throw(ValidationError("$result doesn't have a recognized unit"))
     end
 end
 
+function screen_unit(result)
+    throw(ValidationError("$result doesn't have any unit."))
+end
+
 const unitless = DQ.Quantity(1.0)
 get_literal_unit(x) = screen_unit(something(__get_literal_unit(x), unitless))
 
+_oneunit(x) = oneunit(x)
+
 """
 Find the unit of a symbolic item.
 """
 get_unit(x::Real) = unitless
 get_unit(x::DQ.AbstractQuantity) = screen_unit(x)
 get_unit(x::AbstractArray) = map(get_unit, x)
 get_unit(x::Num) = get_unit(unwrap(x))
-get_unit(x::Symbolics.Arr) = get_unit(unwrap(x))
+function get_unit(x::Union{Symbolics.ArrayOp, Symbolics.Arr, Symbolics.CallWithMetadata})
+    get_literal_unit(x)
+end
 get_unit(op::Differential, args) = get_unit(args[1]) / get_unit(op.x)
 get_unit(op::Difference, args) = get_unit(args[1]) / get_unit(op.t)
 get_unit(op::typeof(getindex), args) = get_unit(args[1])
 get_unit(x::SciMLBase.NullParameters) = unitless
 get_unit(op::typeof(instream), args) = get_unit(args[1])
 
 function get_unit(op, args) # Fallback
-    result = oneunit(op(get_unit.(args)...))
+    result = _oneunit(op(get_unit.(args)...))
     try
         get_unit(result)
     catch
@@ -104,7 +108,10 @@ function get_unit(op::Integral, args)
     return get_unit(args[1]) * unit
 end
 
-equivalent(x, y) = isequal(x, y)
+"""
+Test unit equivalence.
+"""
+equivalent(x, y) = isequal(1 * x, 1 * y)
 function get_unit(op::Conditional, args)
     terms = get_unit.(args)
     terms[1] == unitless ||
@@ -167,6 +174,10 @@ function get_unit(x::Symbolic)
     end
 end
 
+# Dimension error detection function - extensible for different unit systems
+_is_dimension_error(e) = false  # Default fallback
+_is_dimension_error(e::DQ.DimensionError) = true
+
 """
 Get unit of term, returning nothing & showing warning instead of throwing errors.
 """
@@ -175,8 +186,8 @@ function safe_get_unit(term, info)
     try
         side = get_unit(term)
     catch err
-        if err isa DQ.DimensionError
-            @warn("$info: $(err.x) and $(err.y) are not dimensionally compatible.")
+        if _is_dimension_error(err)
+            @warn("$info: dimension error occurred.")
         elseif err isa ValidationError
             @warn(info*err.message)
         elseif err isa MethodError