diff --git a/common/python/nomadcore/unit_conversion/unit_conversion.py b/common/python/nomadcore/unit_conversion/unit_conversion.py index 1dc52f16deb9f0acf74a10e49eb10e37cd0eadea..dfe24285ab1b0990e3962b6ac89286579fe5bf4b 100644 --- a/common/python/nomadcore/unit_conversion/unit_conversion.py +++ b/common/python/nomadcore/unit_conversion/unit_conversion.py @@ -8,10 +8,10 @@ Any new units and constants can be added to the text files "units.txt" and from __future__ import print_function from builtins import str from builtins import object -import os import re import logging from pint import UnitRegistry +import cachetools logger = logging.getLogger(__name__) # disable warnings from pint logging.getLogger("pint").setLevel(logging.ERROR) @@ -20,19 +20,13 @@ logging.getLogger("pint").setLevel(logging.ERROR) ureg = UnitRegistry() ureg.define('forceAu = hartree / bohr') ureg.define('inversecm = 1.9864475e-23 * joule') -_ureg_cache = dict() +@cachetools.cached(cache={}) def ureg_cached(unit): - if unit in _ureg_cache: - return _ureg_cache[unit] - else: - unit_def = ureg(unit) - _ureg_cache[unit] = unit_def - return unit_def + return ureg(unit) -#=============================================================================== def register_userdefined_quantity(quantity, units, value=1): """Registers a user defined quantity, valid until redefined. The value should be equal to value using units, with value defaulting to 1 @@ -45,7 +39,6 @@ def register_userdefined_quantity(quantity, units, value=1): ureg.define(quantity + ' = ' + str(value) + " * " + units) -#=============================================================================== def convert_unit(value, unit, target_unit=None): """Converts the given value from the given units to the target units. For examples see the bottom section. @@ -63,6 +56,9 @@ def convert_unit(value, unit, target_unit=None): The given value in the target units. returned as the same data type as the original values. """ + factor = get_multiplicative_factor(unit, target_unit) + if factor is not None: + return value * factor # Check that the unit is valid unit_def = ureg_cached(unit) @@ -76,19 +72,40 @@ def convert_unit(value, unit, target_unit=None): pint_value = Q_(value, unit_def) converted_value = pint_value.to_base_units() # Base units are defined in the "units.txt" file, and they are the SI units. return converted_value.magnitude + else: # Check that the given target unit is valid target_unit_def = ureg_cached(target_unit) if not target_unit_def: logger.error("Undefined target unit given. Cannot do the conversion") return + Q_ = ureg.Quantity pint_value = Q_(value, unit_def) converted_value = pint_value.to(target_unit_def) return converted_value.magnitude -#=============================================================================== +@cachetools.cached(cache={}) +def get_multiplicative_factor(unit, target_unit=None): + ''' + An optimization for unit conversion. In 99% of the cases this can produce a scaling + factor to convert from unit to target_unit (or base units) and the + results can be cached. + ''' + try: + unit_factor, _ = ureg.get_base_units(ureg_cached(unit), check_nonmult=True) + if target_unit is not None: + target_factor, _ = ureg.get_base_units(ureg_cached(target_unit), check_nonmult=True) + else: + target_factor = 1 + + return unit_factor / target_factor + + except Exception: + return None + + def convert_unit_function_immediate(unit, target_unit=None): """Returns a function that converts scalar floats from unit to target_unit All units need to be already known. @@ -119,7 +136,6 @@ def convert_unit_function_immediate(unit, target_unit=None): return lambda x: convert_unit(x, unit, target_unit) -#=============================================================================== class LazyF(object): """helper class for lazy evaluation of conversion function""" def __init__(self, unit, target_unit): @@ -135,7 +151,6 @@ class LazyF(object): return self.f(x) -#=============================================================================== def convert_unit_function(unit, target_unit=None): """Returns a function that converts scalar floats from unit to target_unit if any of the unit are user defined (usr*), then the conversion is done lazily @@ -162,7 +177,6 @@ def convert_unit_function(unit, target_unit=None): return convert_unit_function_immediate(unit, target_unit) -#=============================================================================== # Testing if __name__ == "__main__":