206 lines
7.4 KiB
Python
206 lines
7.4 KiB
Python
"""
|
|
Unit system for physical quantities; include definition of constants.
|
|
"""
|
|
|
|
from typing import Dict as tDict, Set as tSet
|
|
|
|
from sympy.core.add import Add
|
|
from sympy.core.function import (Derivative, Function)
|
|
from sympy.core.mul import Mul
|
|
from sympy.core.power import Pow
|
|
from sympy.core.singleton import S
|
|
from sympy.physics.units.dimensions import _QuantityMapper
|
|
from sympy.physics.units.quantities import Quantity
|
|
|
|
from .dimensions import Dimension
|
|
|
|
|
|
class UnitSystem(_QuantityMapper):
|
|
"""
|
|
UnitSystem represents a coherent set of units.
|
|
|
|
A unit system is basically a dimension system with notions of scales. Many
|
|
of the methods are defined in the same way.
|
|
|
|
It is much better if all base units have a symbol.
|
|
"""
|
|
|
|
_unit_systems = {} # type: tDict[str, UnitSystem]
|
|
|
|
def __init__(self, base_units, units=(), name="", descr="", dimension_system=None, derived_units: tDict[Dimension, Quantity]={}):
|
|
|
|
UnitSystem._unit_systems[name] = self
|
|
|
|
self.name = name
|
|
self.descr = descr
|
|
|
|
self._base_units = base_units
|
|
self._dimension_system = dimension_system
|
|
self._units = tuple(set(base_units) | set(units))
|
|
self._base_units = tuple(base_units)
|
|
self._derived_units = derived_units
|
|
|
|
super().__init__()
|
|
|
|
def __str__(self):
|
|
"""
|
|
Return the name of the system.
|
|
|
|
If it does not exist, then it makes a list of symbols (or names) of
|
|
the base dimensions.
|
|
"""
|
|
|
|
if self.name != "":
|
|
return self.name
|
|
else:
|
|
return "UnitSystem((%s))" % ", ".join(
|
|
str(d) for d in self._base_units)
|
|
|
|
def __repr__(self):
|
|
return '<UnitSystem: %s>' % repr(self._base_units)
|
|
|
|
def extend(self, base, units=(), name="", description="", dimension_system=None, derived_units: tDict[Dimension, Quantity]={}):
|
|
"""Extend the current system into a new one.
|
|
|
|
Take the base and normal units of the current system to merge
|
|
them to the base and normal units given in argument.
|
|
If not provided, name and description are overridden by empty strings.
|
|
"""
|
|
|
|
base = self._base_units + tuple(base)
|
|
units = self._units + tuple(units)
|
|
|
|
return UnitSystem(base, units, name, description, dimension_system, {**self._derived_units, **derived_units})
|
|
|
|
def get_dimension_system(self):
|
|
return self._dimension_system
|
|
|
|
def get_quantity_dimension(self, unit):
|
|
qdm = self.get_dimension_system()._quantity_dimension_map
|
|
if unit in qdm:
|
|
return qdm[unit]
|
|
return super().get_quantity_dimension(unit)
|
|
|
|
def get_quantity_scale_factor(self, unit):
|
|
qsfm = self.get_dimension_system()._quantity_scale_factors
|
|
if unit in qsfm:
|
|
return qsfm[unit]
|
|
return super().get_quantity_scale_factor(unit)
|
|
|
|
@staticmethod
|
|
def get_unit_system(unit_system):
|
|
if isinstance(unit_system, UnitSystem):
|
|
return unit_system
|
|
|
|
if unit_system not in UnitSystem._unit_systems:
|
|
raise ValueError(
|
|
"Unit system is not supported. Currently"
|
|
"supported unit systems are {}".format(
|
|
", ".join(sorted(UnitSystem._unit_systems))
|
|
)
|
|
)
|
|
|
|
return UnitSystem._unit_systems[unit_system]
|
|
|
|
@staticmethod
|
|
def get_default_unit_system():
|
|
return UnitSystem._unit_systems["SI"]
|
|
|
|
@property
|
|
def dim(self):
|
|
"""
|
|
Give the dimension of the system.
|
|
|
|
That is return the number of units forming the basis.
|
|
"""
|
|
return len(self._base_units)
|
|
|
|
@property
|
|
def is_consistent(self):
|
|
"""
|
|
Check if the underlying dimension system is consistent.
|
|
"""
|
|
# test is performed in DimensionSystem
|
|
return self.get_dimension_system().is_consistent
|
|
|
|
@property
|
|
def derived_units(self) -> tDict[Dimension, Quantity]:
|
|
return self._derived_units
|
|
|
|
def get_dimensional_expr(self, expr):
|
|
from sympy.physics.units import Quantity
|
|
if isinstance(expr, Mul):
|
|
return Mul(*[self.get_dimensional_expr(i) for i in expr.args])
|
|
elif isinstance(expr, Pow):
|
|
return self.get_dimensional_expr(expr.base) ** expr.exp
|
|
elif isinstance(expr, Add):
|
|
return self.get_dimensional_expr(expr.args[0])
|
|
elif isinstance(expr, Derivative):
|
|
dim = self.get_dimensional_expr(expr.expr)
|
|
for independent, count in expr.variable_count:
|
|
dim /= self.get_dimensional_expr(independent)**count
|
|
return dim
|
|
elif isinstance(expr, Function):
|
|
args = [self.get_dimensional_expr(arg) for arg in expr.args]
|
|
if all(i == 1 for i in args):
|
|
return S.One
|
|
return expr.func(*args)
|
|
elif isinstance(expr, Quantity):
|
|
return self.get_quantity_dimension(expr).name
|
|
return S.One
|
|
|
|
def _collect_factor_and_dimension(self, expr):
|
|
"""
|
|
Return tuple with scale factor expression and dimension expression.
|
|
"""
|
|
from sympy.physics.units import Quantity
|
|
if isinstance(expr, Quantity):
|
|
return expr.scale_factor, expr.dimension
|
|
elif isinstance(expr, Mul):
|
|
factor = 1
|
|
dimension = Dimension(1)
|
|
for arg in expr.args:
|
|
arg_factor, arg_dim = self._collect_factor_and_dimension(arg)
|
|
factor *= arg_factor
|
|
dimension *= arg_dim
|
|
return factor, dimension
|
|
elif isinstance(expr, Pow):
|
|
factor, dim = self._collect_factor_and_dimension(expr.base)
|
|
exp_factor, exp_dim = self._collect_factor_and_dimension(expr.exp)
|
|
if self.get_dimension_system().is_dimensionless(exp_dim):
|
|
exp_dim = 1
|
|
return factor ** exp_factor, dim ** (exp_factor * exp_dim)
|
|
elif isinstance(expr, Add):
|
|
factor, dim = self._collect_factor_and_dimension(expr.args[0])
|
|
for addend in expr.args[1:]:
|
|
addend_factor, addend_dim = \
|
|
self._collect_factor_and_dimension(addend)
|
|
if not self.get_dimension_system().equivalent_dims(dim, addend_dim):
|
|
raise ValueError(
|
|
'Dimension of "{}" is {}, '
|
|
'but it should be {}'.format(
|
|
addend, addend_dim, dim))
|
|
factor += addend_factor
|
|
return factor, dim
|
|
elif isinstance(expr, Derivative):
|
|
factor, dim = self._collect_factor_and_dimension(expr.args[0])
|
|
for independent, count in expr.variable_count:
|
|
ifactor, idim = self._collect_factor_and_dimension(independent)
|
|
factor /= ifactor**count
|
|
dim /= idim**count
|
|
return factor, dim
|
|
elif isinstance(expr, Function):
|
|
fds = [self._collect_factor_and_dimension(arg) for arg in expr.args]
|
|
dims = [Dimension(1) if self.get_dimension_system().is_dimensionless(d[1]) else d[1] for d in fds]
|
|
return (expr.func(*(f[0] for f in fds)), *dims)
|
|
elif isinstance(expr, Dimension):
|
|
return S.One, expr
|
|
else:
|
|
return expr, Dimension(1)
|
|
|
|
def get_units_non_prefixed(self) -> tSet[Quantity]:
|
|
"""
|
|
Return the units of the system that do not have a prefix.
|
|
"""
|
|
return set(filter(lambda u: not u.is_prefixed and not u.is_physical_constant, self._units))
|