Commit 1d1192df authored by Markus Scheidgen's avatar Markus Scheidgen

optimized pint unit conversions

parent 75a663a7
......@@ -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__":
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment