Merge pull request #107 from ven-k/vkb/fix_register_units

Allow External Unit Registration
SymbolicML · Feb 12, 2024 · 889d020 · 889d020
2 parents 9681769 + 92df2f4
commit 889d020
Show file tree

Hide file tree

Showing 8 changed files with 274 additions and 75 deletions.
diff --git a/docs/src/units.md b/docs/src/units.md
@@ -42,3 +42,11 @@ Units.T
 Units.L
 Units.bar
 ```
+
+## Custom Units
+
+You can define custom units with the `@register_unit` macro:
+
+```@docs
+@register_unit
+```
diff --git a/src/DynamicQuantities.jl b/src/DynamicQuantities.jl
@@ -9,10 +9,12 @@ export QuantityArray
 export DimensionError
 export ustrip, dimension
 export ulength, umass, utime, ucurrent, utemperature, uluminosity, uamount
-export uparse, @u_str, sym_uparse, @us_str, uexpand, uconvert
+export uparse, @u_str, sym_uparse, @us_str, uexpand, uconvert, @register_unit
+
 
 include("internal_utils.jl")
 include("fixed_rational.jl")
+include("write_once_read_many.jl")
 include("types.jl")
 include("utils.jl")
 include("math.jl")
@@ -22,6 +24,7 @@ include("constants.jl")
 include("uparse.jl")
 include("symbolic_dimensions.jl")
 include("complex.jl")
+include("register_units.jl")
 include("disambiguities.jl")
 
 include("deprecated.jl")
@@ -38,12 +41,11 @@ using .Units: UNIT_SYMBOLS
 let _units_import_expr = :(using .Units: m, g)
     append!(
         _units_import_expr.args[1].args,
-        map(s -> Expr(:(.), s), filter(s -> s ∉ (:m, :g), UNIT_SYMBOLS))
+        Expr(:(.), s) for s in UNIT_SYMBOLS if s ∉ (:m, :g)
     )
     eval(_units_import_expr)
 end
 
-
 function __init__()
     @require_extensions
 end

diff --git a/src/register_units.jl b/src/register_units.jl
@@ -0,0 +1,72 @@
+import .Units: UNIT_MAPPING, UNIT_SYMBOLS, UNIT_VALUES, _lazy_register_unit
+import .SymbolicUnits: update_external_symbolic_unit_value
+
+# Update the unit collections
+const UNIT_UPDATE_LOCK = Threads.SpinLock()
+
+function update_all_values(name_symbol, unit)
+    lock(UNIT_UPDATE_LOCK) do
+        push!(ALL_SYMBOLS, name_symbol)
+        push!(ALL_VALUES, unit)
+        i = lastindex(ALL_VALUES)
+        ALL_MAPPING[name_symbol] = i
+        UNIT_MAPPING[name_symbol] = i
+        update_external_symbolic_unit_value(name_symbol)
+    end
+end
+
+"""
+    @register_unit symbol value
+
+Register a new unit under the given symbol to have
+a particular value.
+
+# Example
+
+```julia
+julia> @register_unit MyVolt 1.5u"V"
+```
+
+This will register a new unit `MyVolt` with a value of `1.5u"V"`.
+You can then use this unit in your calculations:
+
+```julia
+julia> x = 20us"MyVolt^2"
+20.0 MyVolt²
+
+julia> y = 2.5us"A"
+2.5 A
+
+julia> x * y^2 |> uconvert(us"W^2")
+281.25 W²
+
+julia> x * y^2 |> uconvert(us"W^2") |> sqrt |> uexpand
+16.77050983124842 m² kg s⁻³
+```
+
+"""
+macro register_unit(symbol, value)
+    return esc(_register_unit(symbol, value))
+end
+
+function _register_unit(name::Symbol, value)
+    name_symbol = Meta.quot(name)
+    index = get(ALL_MAPPING, name, INDEX_TYPE(0))
+    if !iszero(index)
+        unit = ALL_VALUES[index]
+        # When a utility function to expand `value` to its final form becomes
+        # available, enable the following check. This will avoid throwing an error
+        # if user is trying to register an existing unit with matching values.
+        # unit.value != value && throw("Unit $name is already defined as $unit")
+        error("Unit `$name` is already defined as `$unit`")
+    end
+    reg_expr = _lazy_register_unit(name, value)
+    push!(
+        reg_expr.args,
+        quote
+            $update_all_values($name_symbol, $value)
+            nothing
+        end
+    )
+    return reg_expr
+end
diff --git a/src/symbolic_dimensions.jl b/src/symbolic_dimensions.jl
@@ -1,21 +1,17 @@
+import ..WriteOnceReadMany
 import .Units: UNIT_SYMBOLS, UNIT_MAPPING, UNIT_VALUES
 import .Constants: CONSTANT_SYMBOLS, CONSTANT_MAPPING, CONSTANT_VALUES
 
 
-const SYMBOL_CONFLICTS = intersect(UNIT_SYMBOLS, CONSTANT_SYMBOLS)
+disambiguate_constant_symbol(s) = s in UNIT_SYMBOLS ? Symbol(s, :_constant) : s
 
-disambiguate_symbol(s) = s in SYMBOL_CONFLICTS ? Symbol(s, :_constant) : s
-
-const INDEX_TYPE = UInt8
+const INDEX_TYPE = UInt16
 # Prefer units over constants:
 # For example, this means we can't have a symbolic Planck's constant,
 # as it is just "hours" (h), which is more common.
-const ALL_SYMBOLS = (
-    UNIT_SYMBOLS...,
-    disambiguate_symbol.(CONSTANT_SYMBOLS)...
-)
-const ALL_VALUES = (UNIT_VALUES..., CONSTANT_VALUES...)
-const ALL_MAPPING = NamedTuple{ALL_SYMBOLS}(INDEX_TYPE(1):INDEX_TYPE(length(ALL_SYMBOLS)))
+const ALL_SYMBOLS = WriteOnceReadMany([UNIT_SYMBOLS..., disambiguate_constant_symbol.(CONSTANT_SYMBOLS)...])
+const ALL_VALUES = WriteOnceReadMany([UNIT_VALUES..., CONSTANT_VALUES...])
+const ALL_MAPPING = WriteOnceReadMany(Dict(s => INDEX_TYPE(i) for (i, s) in enumerate(ALL_SYMBOLS)))
 
 """
     AbstractSymbolicDimensions{R} <: AbstractDimensions{R}
@@ -169,7 +165,7 @@ uexpand(q::QuantityArray) = uexpand.(q)
     uconvert(qout::UnionAbstractQuantity{<:Any, <:AbstractSymbolicDimensions}, q::UnionAbstractQuantity{<:Any, <:Dimensions})
 
 Convert a quantity `q` with base SI units to the symbolic units of `qout`, for `q` and `qout` with compatible units.
-Mathematically, the result has value `q / uexpand(qout)` and units `dimension(qout)`. 
+Mathematically, the result has value `q / uexpand(qout)` and units `dimension(qout)`.
 """
 function uconvert(qout::UnionAbstractQuantity{<:Any, <:SymbolicDimensions}, q::UnionAbstractQuantity{<:Any, <:Dimensions})
     @assert isone(ustrip(qout)) "You passed a quantity with a non-unit value to uconvert."
@@ -224,7 +220,7 @@ end
 """
     uconvert(qout::UnionAbstractQuantity{<:Any, <:AbstractSymbolicDimensions})
 
-Create a function that converts an input quantity `q` with base SI units to the symbolic units of `qout`, i.e 
+Create a function that converts an input quantity `q` with base SI units to the symbolic units of `qout`, i.e
 a function equivalent to `q -> uconvert(qout, q)`.
 """
 uconvert(qout::UnionAbstractQuantity{<:Any,<:AbstractSymbolicDimensions}) = Base.Fix1(uconvert, qout)
@@ -371,29 +367,30 @@ module SymbolicUnits
     import ..UNIT_SYMBOLS
     import ..CONSTANT_SYMBOLS
     import ..SymbolicDimensionsSingleton
-    import ...constructorof
-    import ...DEFAULT_SYMBOLIC_QUANTITY_TYPE
-    import ...DEFAULT_SYMBOLIC_QUANTITY_OUTPUT_TYPE
-    import ...DEFAULT_VALUE_TYPE
-    import ...DEFAULT_DIM_BASE_TYPE
+    import ..constructorof
+    import ..DEFAULT_SYMBOLIC_QUANTITY_TYPE
+    import ..DEFAULT_SYMBOLIC_QUANTITY_OUTPUT_TYPE
+    import ..DEFAULT_VALUE_TYPE
+    import ..DEFAULT_DIM_BASE_TYPE
+    import ..WriteOnceReadMany
 
     # Lazily create unit symbols (since there are so many)
     module Constants
         import ...CONSTANT_SYMBOLS
         import ...SymbolicDimensionsSingleton
         import ...constructorof
-        import ...disambiguate_symbol
-        import ....DEFAULT_SYMBOLIC_QUANTITY_TYPE
-        import ....DEFAULT_VALUE_TYPE
-        import ....DEFAULT_DIM_BASE_TYPE
+        import ...disambiguate_constant_symbol
+        import ...DEFAULT_SYMBOLIC_QUANTITY_TYPE
+        import ...DEFAULT_VALUE_TYPE
+        import ...DEFAULT_DIM_BASE_TYPE
 
         const _SYMBOLIC_CONSTANT_VALUES = DEFAULT_SYMBOLIC_QUANTITY_TYPE[]
 
         for unit in CONSTANT_SYMBOLS
             @eval begin
                 const $unit = constructorof(DEFAULT_SYMBOLIC_QUANTITY_TYPE)(
                     DEFAULT_VALUE_TYPE(1.0),
-                    SymbolicDimensionsSingleton{DEFAULT_DIM_BASE_TYPE}($(QuoteNode(disambiguate_symbol(unit))))
+                    SymbolicDimensionsSingleton{DEFAULT_DIM_BASE_TYPE}($(QuoteNode(disambiguate_constant_symbol(unit))))
                 )
                 push!(_SYMBOLIC_CONSTANT_VALUES, $unit)
             end
@@ -404,18 +401,30 @@ module SymbolicUnits
     import .Constants as SymbolicConstants
     import .Constants: SYMBOLIC_CONSTANT_VALUES
 
-    const _SYMBOLIC_UNIT_VALUES = DEFAULT_SYMBOLIC_QUANTITY_TYPE[]
-    for unit in UNIT_SYMBOLS
+    const SYMBOLIC_UNIT_VALUES = WriteOnceReadMany{Vector{DEFAULT_SYMBOLIC_QUANTITY_TYPE}}()
+
+    function update_symbolic_unit_values!(unit, symbolic_unit_values = SYMBOLIC_UNIT_VALUES)
         @eval begin
             const $unit = constructorof(DEFAULT_SYMBOLIC_QUANTITY_TYPE)(
                 DEFAULT_VALUE_TYPE(1.0),
                 SymbolicDimensionsSingleton{DEFAULT_DIM_BASE_TYPE}($(QuoteNode(unit)))
             )
-            push!(_SYMBOLIC_UNIT_VALUES, $unit)
+            push!($symbolic_unit_values, $unit)
         end
     end
-    const SYMBOLIC_UNIT_VALUES = Tuple(_SYMBOLIC_UNIT_VALUES)
 
+    update_symbolic_unit_values!(w::WriteOnceReadMany) = update_symbolic_unit_values!.(w._raw_data)
+    update_symbolic_unit_values!(UNIT_SYMBOLS)
+
+    # Non-eval version of `update_symbolic_unit_values!` for registering units in
+    # an external module.
+    function update_external_symbolic_unit_value(unit)
+        unit = constructorof(DEFAULT_SYMBOLIC_QUANTITY_TYPE)(
+            DEFAULT_VALUE_TYPE(1.0),
+            SymbolicDimensionsSingleton{DEFAULT_DIM_BASE_TYPE}(unit)
+        )
+        push!(SYMBOLIC_UNIT_VALUES, unit)
+    end
 
     """
         sym_uparse(raw_string::AbstractString)

diff --git a/src/units.jl b/src/units.jl
@@ -1,29 +1,30 @@
 module Units
 
+import ..WriteOnceReadMany
 import ..DEFAULT_DIM_TYPE
 import ..DEFAULT_VALUE_TYPE
 import ..DEFAULT_QUANTITY_TYPE
 
 @assert DEFAULT_VALUE_TYPE == Float64 "`units.jl` must be updated to support a different default value type."
 
-const _UNIT_SYMBOLS = Symbol[]
-const _UNIT_VALUES = DEFAULT_QUANTITY_TYPE[]
+const UNIT_SYMBOLS = WriteOnceReadMany{Vector{Symbol}}()
+const UNIT_VALUES = WriteOnceReadMany{Vector{DEFAULT_QUANTITY_TYPE}}()
 
-macro register_unit(name, value)
-    return esc(_register_unit(name, value))
+macro _lazy_register_unit(name, value)
+    return esc(_lazy_register_unit(name, value))
 end
 
 macro add_prefixes(base_unit, prefixes)
     @assert prefixes.head == :tuple
-    return esc(_add_prefixes(base_unit, prefixes.args, _register_unit))
+    return esc(_add_prefixes(base_unit, prefixes.args, _lazy_register_unit))
 end
 
-function _register_unit(name::Symbol, value)
-    s = string(name)
-    return quote
+function _lazy_register_unit(name::Symbol, value)
+    name_symbol = Meta.quot(name)
+    quote
         const $name = $value
-        push!(_UNIT_SYMBOLS, Symbol($s))
-        push!(_UNIT_VALUES, $name)
+        push!($UNIT_SYMBOLS, $name_symbol)
+        push!($UNIT_VALUES, $name)
     end
 end
 
@@ -42,13 +43,13 @@ function _add_prefixes(base_unit::Symbol, prefixes, register_function)
 end
 
 # SI base units
-@register_unit m DEFAULT_QUANTITY_TYPE(1.0, length=1)
-@register_unit g DEFAULT_QUANTITY_TYPE(1e-3, mass=1)
-@register_unit s DEFAULT_QUANTITY_TYPE(1.0, time=1)
-@register_unit A DEFAULT_QUANTITY_TYPE(1.0, current=1)
-@register_unit K DEFAULT_QUANTITY_TYPE(1.0, temperature=1)
-@register_unit cd DEFAULT_QUANTITY_TYPE(1.0, luminosity=1)
-@register_unit mol DEFAULT_QUANTITY_TYPE(1.0, amount=1)
+@_lazy_register_unit m DEFAULT_QUANTITY_TYPE(1.0, length=1)
+@_lazy_register_unit g DEFAULT_QUANTITY_TYPE(1e-3, mass=1)
+@_lazy_register_unit s DEFAULT_QUANTITY_TYPE(1.0, time=1)
+@_lazy_register_unit A DEFAULT_QUANTITY_TYPE(1.0, current=1)
+@_lazy_register_unit K DEFAULT_QUANTITY_TYPE(1.0, temperature=1)
+@_lazy_register_unit cd DEFAULT_QUANTITY_TYPE(1.0, luminosity=1)
+@_lazy_register_unit mol DEFAULT_QUANTITY_TYPE(1.0, amount=1)
 
 @add_prefixes m (f, p, n, μ, u, c, d, m, k, M, G)
 @add_prefixes g (p, n, μ, u, m, k)
@@ -88,17 +89,17 @@ end
 )
 
 # SI derived units
-@register_unit Hz inv(s)
-@register_unit N kg * m / s^2
-@register_unit Pa N / m^2
-@register_unit J N * m
-@register_unit W J / s
-@register_unit C A * s
-@register_unit V W / A
-@register_unit F C / V
-@register_unit Ω V / A
-@register_unit ohm Ω
-@register_unit T N / (A * m)
+@_lazy_register_unit Hz inv(s)
+@_lazy_register_unit N kg * m / s^2
+@_lazy_register_unit Pa N / m^2
+@_lazy_register_unit J N * m
+@_lazy_register_unit W J / s
+@_lazy_register_unit C A * s
+@_lazy_register_unit V W / A
+@_lazy_register_unit F C / V
+@_lazy_register_unit Ω V / A
+@_lazy_register_unit ohm Ω
+@_lazy_register_unit T N / (A * m)
 
 @add_prefixes Hz (n, μ, u, m, k, M, G)
 @add_prefixes N ()
@@ -156,17 +157,17 @@ end
 
 # Common assorted units
 ## Time
-@register_unit min 60 * s
-@register_unit minute min
-@register_unit h 60 * min
-@register_unit hr h
-@register_unit day 24 * h
-@register_unit d day
-@register_unit wk 7 * day
-@register_unit yr 365.25 * day
-@register_unit inch 2.54 * cm
-@register_unit ft 12 * inch
-@register_unit mi 5280 * ft
+@_lazy_register_unit min 60 * s
+@_lazy_register_unit minute min
+@_lazy_register_unit h 60 * min
+@_lazy_register_unit hr h
+@_lazy_register_unit day 24 * h
+@_lazy_register_unit d day
+@_lazy_register_unit wk 7 * day
+@_lazy_register_unit yr 365.25 * day
+@_lazy_register_unit inch 2.54 * cm
+@_lazy_register_unit ft 12 * inch
+@_lazy_register_unit mi 5280 * ft
 
 @add_prefixes min ()
 @add_prefixes minute ()
@@ -178,7 +179,7 @@ end
 @add_prefixes yr (k, M, G)
 
 ## Volume
-@register_unit L dm^3
+@_lazy_register_unit L dm^3
 
 @add_prefixes L (μ, u, m, c, d)
 
@@ -188,7 +189,7 @@ end
 )
 
 ## Pressure
-@register_unit bar 100 * kPa
+@_lazy_register_unit bar 100 * kPa
 
 @add_prefixes bar (m,)
 
@@ -203,9 +204,7 @@ end
 # Do not wish to define physical constants, as the number of symbols might lead to ambiguity.
 # The user should define these instead.
 
-"""A tuple of all possible unit symbols."""
-const UNIT_SYMBOLS = Tuple(_UNIT_SYMBOLS)
-const UNIT_VALUES = Tuple(_UNIT_VALUES)
-const UNIT_MAPPING = NamedTuple([s => i for (i, s) in enumerate(UNIT_SYMBOLS)])
+# Update `UNIT_MAPPING` with all internally defined unit symbols.
+const UNIT_MAPPING = WriteOnceReadMany(Dict(s => i for (i, s) in enumerate(UNIT_SYMBOLS)))
 
 end