unitsystem.py 6.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194
  1. """
  2. Unit system for physical quantities; include definition of constants.
  3. """
  4. from typing import Dict as tDict
  5. from sympy.core.add import Add
  6. from sympy.core.function import (Derivative, Function)
  7. from sympy.core.mul import Mul
  8. from sympy.core.power import Pow
  9. from sympy.core.singleton import S
  10. from sympy.physics.units.dimensions import _QuantityMapper
  11. from .dimensions import Dimension
  12. class UnitSystem(_QuantityMapper):
  13. """
  14. UnitSystem represents a coherent set of units.
  15. A unit system is basically a dimension system with notions of scales. Many
  16. of the methods are defined in the same way.
  17. It is much better if all base units have a symbol.
  18. """
  19. _unit_systems = {} # type: tDict[str, UnitSystem]
  20. def __init__(self, base_units, units=(), name="", descr="", dimension_system=None):
  21. UnitSystem._unit_systems[name] = self
  22. self.name = name
  23. self.descr = descr
  24. self._base_units = base_units
  25. self._dimension_system = dimension_system
  26. self._units = tuple(set(base_units) | set(units))
  27. self._base_units = tuple(base_units)
  28. super().__init__()
  29. def __str__(self):
  30. """
  31. Return the name of the system.
  32. If it does not exist, then it makes a list of symbols (or names) of
  33. the base dimensions.
  34. """
  35. if self.name != "":
  36. return self.name
  37. else:
  38. return "UnitSystem((%s))" % ", ".join(
  39. str(d) for d in self._base_units)
  40. def __repr__(self):
  41. return '<UnitSystem: %s>' % repr(self._base_units)
  42. def extend(self, base, units=(), name="", description="", dimension_system=None):
  43. """Extend the current system into a new one.
  44. Take the base and normal units of the current system to merge
  45. them to the base and normal units given in argument.
  46. If not provided, name and description are overridden by empty strings.
  47. """
  48. base = self._base_units + tuple(base)
  49. units = self._units + tuple(units)
  50. return UnitSystem(base, units, name, description, dimension_system)
  51. def get_dimension_system(self):
  52. return self._dimension_system
  53. def get_quantity_dimension(self, unit):
  54. qdm = self.get_dimension_system()._quantity_dimension_map
  55. if unit in qdm:
  56. return qdm[unit]
  57. return super().get_quantity_dimension(unit)
  58. def get_quantity_scale_factor(self, unit):
  59. qsfm = self.get_dimension_system()._quantity_scale_factors
  60. if unit in qsfm:
  61. return qsfm[unit]
  62. return super().get_quantity_scale_factor(unit)
  63. @staticmethod
  64. def get_unit_system(unit_system):
  65. if isinstance(unit_system, UnitSystem):
  66. return unit_system
  67. if unit_system not in UnitSystem._unit_systems:
  68. raise ValueError(
  69. "Unit system is not supported. Currently"
  70. "supported unit systems are {}".format(
  71. ", ".join(sorted(UnitSystem._unit_systems))
  72. )
  73. )
  74. return UnitSystem._unit_systems[unit_system]
  75. @staticmethod
  76. def get_default_unit_system():
  77. return UnitSystem._unit_systems["SI"]
  78. @property
  79. def dim(self):
  80. """
  81. Give the dimension of the system.
  82. That is return the number of units forming the basis.
  83. """
  84. return len(self._base_units)
  85. @property
  86. def is_consistent(self):
  87. """
  88. Check if the underlying dimension system is consistent.
  89. """
  90. # test is performed in DimensionSystem
  91. return self.get_dimension_system().is_consistent
  92. def get_dimensional_expr(self, expr):
  93. from sympy.physics.units import Quantity
  94. if isinstance(expr, Mul):
  95. return Mul(*[self.get_dimensional_expr(i) for i in expr.args])
  96. elif isinstance(expr, Pow):
  97. return self.get_dimensional_expr(expr.base) ** expr.exp
  98. elif isinstance(expr, Add):
  99. return self.get_dimensional_expr(expr.args[0])
  100. elif isinstance(expr, Derivative):
  101. dim = self.get_dimensional_expr(expr.expr)
  102. for independent, count in expr.variable_count:
  103. dim /= self.get_dimensional_expr(independent)**count
  104. return dim
  105. elif isinstance(expr, Function):
  106. args = [self.get_dimensional_expr(arg) for arg in expr.args]
  107. if all(i == 1 for i in args):
  108. return S.One
  109. return expr.func(*args)
  110. elif isinstance(expr, Quantity):
  111. return self.get_quantity_dimension(expr).name
  112. return S.One
  113. def _collect_factor_and_dimension(self, expr):
  114. """
  115. Return tuple with scale factor expression and dimension expression.
  116. """
  117. from sympy.physics.units import Quantity
  118. if isinstance(expr, Quantity):
  119. return expr.scale_factor, expr.dimension
  120. elif isinstance(expr, Mul):
  121. factor = 1
  122. dimension = Dimension(1)
  123. for arg in expr.args:
  124. arg_factor, arg_dim = self._collect_factor_and_dimension(arg)
  125. factor *= arg_factor
  126. dimension *= arg_dim
  127. return factor, dimension
  128. elif isinstance(expr, Pow):
  129. factor, dim = self._collect_factor_and_dimension(expr.base)
  130. exp_factor, exp_dim = self._collect_factor_and_dimension(expr.exp)
  131. if self.get_dimension_system().is_dimensionless(exp_dim):
  132. exp_dim = 1
  133. return factor ** exp_factor, dim ** (exp_factor * exp_dim)
  134. elif isinstance(expr, Add):
  135. factor, dim = self._collect_factor_and_dimension(expr.args[0])
  136. for addend in expr.args[1:]:
  137. addend_factor, addend_dim = \
  138. self._collect_factor_and_dimension(addend)
  139. if dim != addend_dim:
  140. raise ValueError(
  141. 'Dimension of "{}" is {}, '
  142. 'but it should be {}'.format(
  143. addend, addend_dim, dim))
  144. factor += addend_factor
  145. return factor, dim
  146. elif isinstance(expr, Derivative):
  147. factor, dim = self._collect_factor_and_dimension(expr.args[0])
  148. for independent, count in expr.variable_count:
  149. ifactor, idim = self._collect_factor_and_dimension(independent)
  150. factor /= ifactor**count
  151. dim /= idim**count
  152. return factor, dim
  153. elif isinstance(expr, Function):
  154. fds = [self._collect_factor_and_dimension(
  155. arg) for arg in expr.args]
  156. return (expr.func(*(f[0] for f in fds)),
  157. expr.func(*(d[1] for d in fds)))
  158. elif isinstance(expr, Dimension):
  159. return S.One, expr
  160. else:
  161. return expr, Dimension(1)