def derive_unit(expr, name=None): """Derive SI-unit from an expression, omitting scale factors.""" from essm.variables import Variable from essm.variables.utils import extract_variables from sympy.physics.units import Dimension from sympy.physics.units.dimensions import dimsys_SI variables = extract_variables(expr) for var1 in variables: q1 = Quantity('q_' + str(var1)) q1.set_dimension( Dimension(Quantity.get_dimensional_expr(var1.definition.unit))) q1.set_scale_factor(var1.definition.unit) expr = expr.xreplace({var1: q1}) dim = Dimension(Quantity.get_dimensional_expr(expr)) return functools.reduce( operator.mul, (SI_DIMENSIONS[d]**p for d, p in dimsys_SI.get_dimensional_dependencies(dim).items()), 1)
def __new__(cls, name, parents, dct): """Build and register new variable.""" if '__registry__' not in dct: unit = dct.pop('unit', S.One) if unit == 1: unit = S.One definition = dct.pop('expr', None) dct.setdefault('name', name) dct.setdefault('domain', 'real') dct.setdefault('latex_name', dct['name']) dct.setdefault('unit', unit) instance = super(VariableMeta, cls).__new__(cls, name, parents, dct) # Variable with definition expression. if definition is not None: definition = build_instance_expression(instance, definition) derived_unit = derive_unit(definition, name=name) if unit == S.One: unit = derived_unit # only if unit is None instance.expr, instance.unit = definition, derived_unit if unit != instance.unit: raise ValueError( 'Invalid expression units {0} should be {1}'.format( instance.unit, unit ) ) expr = BaseVariable( instance, dct['name'], abbrev=dct['latex_name'], dimension=Dimension(Quantity.get_dimensional_expr(unit)), scale_factor=unit or S.One, ) instance[expr] = instance # Store definition as variable expression. if definition is not None: instance.__expressions__[expr] = definition # Store default variable only if it is defined. if 'default' in dct: instance.__defaults__[expr] = dct['default'] # Store unit for each variable: instance.__units__[expr] = instance.unit return expr return super(VariableMeta, cls).__new__(cls, name, parents, dct)
def get_dimensional_expr(expr): """Return dimensions of expression.""" if isinstance(expr, Mul): return Mul(*[Variable.get_dimensional_expr(i) for i in expr.args]) elif isinstance(expr, Pow): return Variable.get_dimensional_expr(expr.base) ** expr.exp elif isinstance(expr, Add): return Variable.get_dimensional_expr(expr.args[0]) elif isinstance(expr, Derivative): dim = Variable.get_dimensional_expr(expr.expr) for independent, count in expr.variable_count: dim /= Variable.get_dimensional_expr(independent)**count return dim elif isinstance(expr, Function): args = [Variable.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 expr.dimension.name elif isinstance(expr, BaseVariable): return Quantity.get_dimensional_expr(expr.definition.unit) return S.One
import sympy.physics.units as u from sympy.physics.units import Dimension, Quantity, find_unit from sympy.physics.units.systems import SI joule = u.joule kelvin = u.kelvin kilogram = u.kilogram meter = u.meter mole = u.mole pascal = u.pascal second = u.second watt = u.watt SI_DIMENSIONS = { str(Quantity.get_dimensional_expr(d)): d for d in SI._base_units } def markdown(unit): """Return markdown representation of a unit.""" from sympy.printing import StrPrinter from operator import itemgetter # displays short units (m instead of meter) StrPrinter._print_Quantity = lambda self, expr: str(expr.abbrev) if unit.is_Pow: item = unit.args return '{0}$^{{{1}}}$'.format(item[0], item[1]) if unit.is_Mul: str1 = ''