Source code for yt.units.unit_registry

"""
A registry for units that can be added to and modified.


"""

#-----------------------------------------------------------------------------
# Copyright (c) 2013, yt Development Team.
#
# Distributed under the terms of the Modified BSD License.
#
# The full license is in the file COPYING.txt, distributed with this software.
#-----------------------------------------------------------------------------

import json
import re

from distutils.version import LooseVersion
from yt.units.unit_lookup_table import \
    default_unit_symbol_lut
from yt.utilities.lib.fnv_hash import fnv_hash
from yt.extern import six
from sympy import \
    sympify, \
    srepr, \
    __version__ as sympy_version

SYMPY_VERSION = LooseVersion(sympy_version)

[docs]def positive_symbol_replacer(match): return match.group().replace(')\')', ')\', positive=True)')
[docs]class SymbolNotFoundError(Exception): pass
[docs]class UnitParseError(Exception): pass
[docs]class UnitRegistry: """A registry for unit symbols""" def __init__(self, add_default_symbols=True, lut=None): if lut: self.lut = lut else: self.lut = {} self.unit_objs = {} if add_default_symbols: self.lut.update(default_unit_symbol_lut) def __getitem__(self, key): return self.lut[key] def __contains__(self, item): return item in self.lut _unit_system_id = None @property def unit_system_id(self): """ This is a unique identifier for the unit registry created from a FNV hash. It is needed to register a dataset's code unit system in the unit system registry. """ if self._unit_system_id is None: hash_data = bytearray() for k, v in self.lut.items(): hash_data.extend(k.encode('ascii')) hash_data.extend(repr(v).encode('ascii')) self._unit_system_id = "code_%d" % fnv_hash(hash_data) return self._unit_system_id
[docs] def add(self, symbol, base_value, dimensions, tex_repr=None, offset=None): """ Add a symbol to this registry. """ from yt.units.unit_object import validate_dimensions # Validate if not isinstance(base_value, float): raise UnitParseError("base_value (%s) must be a float, got a %s." % (base_value, type(base_value))) if offset is not None: if not isinstance(offset, float): raise UnitParseError( "offset value (%s) must be a float, got a %s." % (offset, type(offset))) else: offset = 0.0 validate_dimensions(dimensions) if tex_repr is None: # make educated guess that will look nice in most cases tex_repr = r"\rm{" + symbol.replace('_', '\ ') + "}" # Add to lut self.lut.update({symbol: (base_value, dimensions, offset, tex_repr)})
[docs] def remove(self, symbol): """ Remove the entry for the unit matching `symbol`. """ if symbol not in self.lut: raise SymbolNotFoundError( "Tried to remove the symbol '%s', but it does not exist" \ "in this registry." % symbol) del self.lut[symbol] if symbol in self.unit_objs: del self.unit_objs[symbol]
[docs] def modify(self, symbol, base_value): """ Change the base value of a unit symbol. Useful for adjusting code units after parsing parameters. """ if symbol not in self.lut: raise SymbolNotFoundError( "Tried to modify the symbol '%s', but it does not exist" \ "in this registry." % symbol) if hasattr(base_value, "in_base"): new_dimensions = base_value.units.dimensions base_value = base_value.in_base('cgs-ampere') base_value = base_value.value else: new_dimensions = self.lut[symbol][1] self.lut[symbol] = ((float(base_value), new_dimensions) + self.lut[symbol][2:]) if symbol in self.unit_objs: del self.unit_objs[symbol]
[docs] def keys(self): """ Print out the units contained in the lookup table. """ return self.lut.keys()
[docs] def to_json(self): """ Returns a json-serialized version of the unit registry """ sanitized_lut = {} for k, v in six.iteritems(self.lut): san_v = list(v) repr_dims = srepr(v[1]) if SYMPY_VERSION < LooseVersion("1.0.0"): # see https://github.com/sympy/sympy/issues/6131 repr_dims = re.sub("Symbol\('\([a-z_]*\)'\)", positive_symbol_replacer, repr_dims) san_v[1] = repr_dims sanitized_lut[k] = tuple(san_v) return json.dumps(sanitized_lut)
@classmethod
[docs] def from_json(cls, json_text): """ Returns a UnitRegistry object from a json-serialized unit registry """ data = json.loads(json_text) lut = {} for k, v in six.iteritems(data): unsan_v = list(v) unsan_v[1] = sympify(v[1]) lut[k] = tuple(unsan_v) return cls(lut=lut, add_default_symbols=False)