100 lines
2.7 KiB
Python
100 lines
2.7 KiB
Python
import functools
|
|
|
|
import sympy as sp
|
|
import sympy.physics.units as spu
|
|
|
|
####################
|
|
# - Useful Methods
|
|
####################
|
|
def uses_units(expression: sp.Expr) -> bool:
|
|
"""Checks if an expression uses any units (`Quantity`)."""
|
|
|
|
for arg in sp.preorder_traversal(expression):
|
|
if isinstance(arg, spu.Quantity):
|
|
return True
|
|
return False
|
|
|
|
# Function to return a set containing all units used in the expression
|
|
def get_units(expression: sp.Expr):
|
|
"""Gets all the units of an expression (as `Quantity`)."""
|
|
|
|
return {
|
|
arg
|
|
for arg in sp.preorder_traversal(expression)
|
|
if isinstance(arg, spu.Quantity)
|
|
}
|
|
|
|
####################
|
|
# - Time
|
|
####################
|
|
femtosecond = fs = spu.Quantity("femtosecond", abbrev="fs")
|
|
femtosecond.set_global_relative_scale_factor(spu.femto, spu.second)
|
|
|
|
|
|
####################
|
|
# - Force
|
|
####################
|
|
# Newton
|
|
nanonewton = nN = spu.Quantity("nanonewton", abbrev="nN")
|
|
nanonewton.set_global_relative_scale_factor(spu.nano, spu.newton)
|
|
|
|
micronewton = uN = spu.Quantity("micronewton", abbrev="μN")
|
|
micronewton.set_global_relative_scale_factor(spu.micro, spu.newton)
|
|
|
|
millinewton = mN = spu.Quantity("micronewton", abbrev="mN")
|
|
micronewton.set_global_relative_scale_factor(spu.milli, spu.newton)
|
|
|
|
####################
|
|
# - Frequency
|
|
####################
|
|
# Hertz
|
|
kilohertz = kHz = spu.Quantity("kilohertz", abbrev="kHz")
|
|
kilohertz.set_global_relative_scale_factor(spu.kilo, spu.hertz)
|
|
|
|
megahertz = MHz = spu.Quantity("megahertz", abbrev="MHz")
|
|
kilohertz.set_global_relative_scale_factor(spu.kilo, spu.hertz)
|
|
|
|
gigahertz = GHz = spu.Quantity("gigahertz", abbrev="GHz")
|
|
gigahertz.set_global_relative_scale_factor(spu.giga, spu.hertz)
|
|
|
|
terahertz = THz = spu.Quantity("terahertz", abbrev="THz")
|
|
terahertz.set_global_relative_scale_factor(spu.tera, spu.hertz)
|
|
|
|
petahertz = PHz = spu.Quantity("petahertz", abbrev="PHz")
|
|
petahertz.set_global_relative_scale_factor(spu.peta, spu.hertz)
|
|
|
|
exahertz = EHz = spu.Quantity("exahertz", abbrev="EHz")
|
|
exahertz.set_global_relative_scale_factor(spu.exa, spu.hertz)
|
|
|
|
####################
|
|
# - Sympy Expression Typing
|
|
####################
|
|
ALL_UNIT_SYMBOLS = {
|
|
unit.abbrev: unit
|
|
for unit in spu.__dict__.values()
|
|
if isinstance(unit, spu.Quantity)
|
|
} | {
|
|
unit.abbrev: unit
|
|
for unit in globals().values()
|
|
if isinstance(unit, spu.Quantity)
|
|
}
|
|
|
|
@functools.lru_cache(maxsize=1024)
|
|
def parse_abbrev_symbols_to_units(expr: sp.Basic) -> sp.Basic:
|
|
print("IN ABBREV", expr)
|
|
return expr.subs(ALL_UNIT_SYMBOLS)
|
|
|
|
#def has_units(expr: sp.Expr):
|
|
# return any(
|
|
# symbol in ALL_UNIT_SYMBOLS
|
|
# for symbol in expr.atoms(sp.Symbol)
|
|
# )
|
|
#def is_exactly_expressed_as_unit(expr: sp.Expr, unit) -> bool:
|
|
# #try:
|
|
# converted_expr = expr / unit
|
|
#
|
|
# return (
|
|
# converted_expr.is_number
|
|
# and not converted_expr.has(spu.Quantity)
|
|
# )
|