Repository: yt-project/unyt
Pull request #62 (2019-01-30 22:14:27)

give unit registries a unit system
this required a fair amount of refactoring to avoid a circular dependency generated by the fact that the unit systems currently store Unit objects. That's not actually necessary, the unit systems only need to store the units that make up the system and can generate Unit objects to return to the user when `__getitem__` is called.

Now that unit registries are associated with a unit system, the registry can be used to determine the appropriate unit system to use for a calculation. This allows the unit system to be controlled by an application by making use of a custom unit registry. Calls to in_base or get_base_equivalent will return data in the appropriately configured unit system associated with that registry.

I haven't yet added configuration for unit systems, but this at least makes it possible for application authors to have their own custom default unit system, as long as they make sure to always use their custom unit registry when they create new arrays.


Patch:
diff --git a/docs/usage.rst b/docs/usage.rst
index effa640..10260fb 100644
--- a/docs/usage.rst
+++ b/docs/usage.rst
@@ -759,7 +759,7 @@ In practice, the unit metadata for a unit object is contained in an instance of
   <unyt.unit_registry.UnitRegistry ...>
 
 All the unit objects in the :mod:`unyt` namespace make use of the default unit
-registry, importable as :data:`unyt.unit_registry.default_unit_registry`. This
+registry, importable as :data:`unyt.unit_object.default_unit_registry`. This
 registry object contains all of the real-world physical units that the
 :mod:`unyt` library ships with out of the box.
 
diff --git a/unyt/array.py b/unyt/array.py
index 661a363..5500911 100644
--- a/unyt/array.py
+++ b/unyt/array.py
@@ -126,13 +126,8 @@ from unyt.exceptions import (
 from unyt.equivalencies import equivalence_registry
 from unyt._on_demand_imports import _astropy, _pint
 from unyt._pint_conversions import convert_pint_units
-from unyt.unit_object import (
-    _check_em_conversion,
-    _em_conversion,
-    _sanitize_unit_system,
-    Unit,
-)
-from unyt.unit_registry import UnitRegistry
+from unyt.unit_object import _check_em_conversion, _em_conversion, Unit
+from unyt.unit_registry import _sanitize_unit_system, UnitRegistry
 
 NULL_UNIT = Unit()
 POWER_SIGN_MAPPING = {multiply: 1, divide: -1}
@@ -657,7 +652,7 @@ class unyt_array(np.ndarray):
         else:
             self.convert_to_equivalent(units, equivalence, **kwargs)
 
-    def convert_to_base(self, unit_system="mks", equivalence=None, **kwargs):
+    def convert_to_base(self, unit_system=None, equivalence=None, **kwargs):
         """
         Convert the array in-place to the equivalent base units in
         the specified unit system.
@@ -669,7 +664,7 @@ class unyt_array(np.ndarray):
         ----------
         unit_system : string, optional
             The unit system to be used in the conversion. If not specified,
-            the default base units of mks are used.
+            the configured base units are used (defaults to MKS).
         equivalence : string, optional
             The equivalence you wish to use. To see which equivalencies
             are supported for this object, try the ``list_equivalencies``
@@ -927,7 +922,7 @@ class unyt_array(np.ndarray):
         else:
             return v
 
-    def in_base(self, unit_system="mks"):
+    def in_base(self, unit_system=None):
         """
         Creates a copy of this array with the data in the specified unit
         system, and returns it in that system's base units.
@@ -936,7 +931,7 @@ class unyt_array(np.ndarray):
         ----------
         unit_system : string, optional
             The unit system to be used in the conversion. If not specified,
-            the default base units of mks are used.
+            the configured default base units of are used (defaults to MKS).
 
         Examples
         --------
diff --git a/unyt/unit_object.py b/unyt/unit_object.py
index b07d2c5..7afadb6 100644
--- a/unyt/unit_object.py
+++ b/unyt/unit_object.py
@@ -66,10 +66,10 @@ from unyt._physical_ratios import speed_of_light_cm_per_s
 from unyt.unit_registry import (
     default_unit_registry,
     _lookup_unit_symbol,
-    _split_prefix,
     UnitRegistry,
     UnitParseError,
 )
+from unyt.unit_systems import _split_prefix
 
 sympy_one = sympify(1)
 
@@ -82,16 +82,6 @@ global_dict = {
 }
 
 
-def _sanitize_unit_system(unit_system, obj):
-    from unyt.unit_systems import unit_system_registry
-
-    if hasattr(unit_system, "unit_registry"):
-        unit_system = unit_system.unit_registry.unit_system_id
-    elif unit_system == "code":
-        unit_system = obj.registry.unit_system_id
-    return unit_system_registry[str(unit_system)]
-
-
 def _auto_positive_symbol(tokens, local_dict, global_dict):
     """
     Inserts calls to ``Symbol`` for undefined variables.
@@ -677,7 +667,7 @@ class Unit(object):
         old_dims = self.dimensions
         return old_dims in this_equiv._dims
 
-    def get_base_equivalent(self, unit_system="mks"):
+    def get_base_equivalent(self, unit_system=None):
         """Create and return dimensionally-equivalent units in a specified base.
 
         >>> from unyt import g, cm
@@ -686,6 +676,8 @@ class Unit(object):
         >>> (g/cm**3).get_base_equivalent('solar')
         Mearth/AU**3
         """
+        from unyt.unit_registry import _sanitize_unit_system
+
         unit_system = _sanitize_unit_system(unit_system, self)
         try:
             conv_data = _check_em_conversion(
@@ -1097,26 +1089,6 @@ def _validate_dimensions(dimensions):
         raise UnitParseError("Bad dimensionality expression '%s'." % dimensions)
 
 
-def _get_system_unit_string(dimensions, base_units):
-    # The dimensions of a unit object is the product of the base dimensions.
-    # Use sympy to factor the dimensions into base CGS unit symbols.
-    units = []
-    my_dims = dimensions.expand()
-    if my_dims is dimensionless:
-        return ""
-    if my_dims in base_units:
-        return base_units[my_dims]
-    for factor in my_dims.as_ordered_factors():
-        dim = list(factor.free_symbols)[0]
-        unit_string = str(base_units[dim])
-        if factor.is_Pow:
-            power_string = "**(%s)" % factor.as_base_exp()[1]
-        else:
-            power_string = ""
-        units.append("(%s)%s" % (unit_string, power_string))
-    return " * ".join(units)
-
-
 def define_unit(
     symbol, value, tex_repr=None, offset=None, prefixable=False, registry=None
 ):
diff --git a/unyt/unit_registry.py b/unyt/unit_registry.py
index 119be29..7011c06 100644
--- a/unyt/unit_registry.py
+++ b/unyt/unit_registry.py
@@ -21,23 +21,42 @@ from unyt._unit_lookup_table import (
     unit_prefixes,
     latex_prefixes,
 )
+from unyt.unit_systems import mks_unit_system, _split_prefix, unit_system_registry
 from hashlib import md5
 import six
 from sympy import sympify, srepr
 
 
+def _sanitize_unit_system(unit_system, obj):
+
+    if unit_system is None:
+        try:
+            unit_system = obj.units.registry.unit_system
+        except AttributeError:
+            unit_system = mks_unit_system
+    if hasattr(unit_system, "name"):
+        return unit_system_registry[unit_system.name]
+    elif hasattr(unit_system, "unit_registry"):
+        unit_system = unit_system.unit_registry.unit_system_id
+    elif unit_system == "code":
+        unit_system = obj.registry.unit_system_id
+    return unit_system_registry[str(unit_system)]
+
+
 class UnitRegistry:
     """A registry for unit symbols"""
 
     _unit_system_id = None
 
-    def __init__(self, add_default_symbols=True, lut=None):
+    def __init__(self, add_default_symbols=True, lut=None, unit_system=None):
         self._unit_object_cache = {}
         if lut:
             self.lut = lut
         else:
             self.lut = {}
 
+        self.unit_system = _sanitize_unit_system(unit_system, None)
+
         if add_default_symbols:
             self.lut.update(default_unit_symbol_lut)
 
@@ -45,9 +64,13 @@ class UnitRegistry:
         try:
             ret = self.lut[key]
         except KeyError:
-            raise SymbolNotFoundError(
-                "The symbol '%s' does not exist in this registry." % key
-            )
+            try:
+                _lookup_unit_symbol(str(key), self.lut)
+                ret = self.lut[key]
+            except UnitParseError:
+                raise SymbolNotFoundError(
+                    "The symbol '%s' does not exist in this registry." % key
+                )
         return ret
 
     def __contains__(self, item):
@@ -246,26 +269,8 @@ class UnitRegistry:
         return equiv
 
 
-def _split_prefix(symbol_str, unit_symbol_lut):
-    possible_prefix = symbol_str[0]
-
-    if symbol_str[:2] == "da":
-        possible_prefix = "da"
-
-    if possible_prefix in unit_prefixes:
-        # the first character could be a prefix, check the rest of the symbol
-        symbol_wo_pref = symbol_str[1:]
-
-        # deca is the only prefix with length 2
-        if symbol_str[:2] == "da":
-            symbol_wo_pref = symbol_str[2:]
-            possible_prefix = "da"
-
-        prefixable_units = [u for u in unit_symbol_lut if unit_symbol_lut[u][4]]
-
-        if symbol_wo_pref in unit_symbol_lut and symbol_wo_pref in prefixable_units:
-            return possible_prefix, symbol_wo_pref
-    return "", symbol_str
+#: The default unit registry
+default_unit_registry = UnitRegistry()
 
 
 def _lookup_unit_symbol(symbol_str, unit_symbol_lut):
@@ -326,7 +331,3 @@ def _lookup_unit_symbol(symbol_str, unit_symbol_lut):
     raise UnitParseError(
         "Could not find unit symbol '%s' in the provided " "symbols." % symbol_str
     )
-
-
-#: The default unit registry
-default_unit_registry = UnitRegistry()
diff --git a/unyt/unit_systems.py b/unyt/unit_systems.py
index a05ba13..f1d50e0 100644
--- a/unyt/unit_systems.py
+++ b/unyt/unit_systems.py
@@ -15,7 +15,51 @@ from collections import OrderedDict
 from six import string_types
 from unyt import dimensions
 from unyt.exceptions import MissingMKSCurrent, IllDefinedUnitSystem
-from unyt.unit_object import Unit, _get_system_unit_string
+from unyt._unit_lookup_table import (
+    default_unit_symbol_lut as default_lut,
+    unit_prefixes,
+)
+
+
+def _split_prefix(symbol_str, unit_symbol_lut):
+    possible_prefix = symbol_str[0]
+
+    if symbol_str[:2] == "da":
+        possible_prefix = "da"
+
+    if possible_prefix in unit_prefixes:
+        # the first character could be a prefix, check the rest of the symbol
+        symbol_wo_pref = symbol_str[1:]
+
+        # deca is the only prefix with length 2
+        if symbol_str[:2] == "da":
+            symbol_wo_pref = symbol_str[2:]
+            possible_prefix = "da"
+
+        prefixable_units = [u for u in unit_symbol_lut if unit_symbol_lut[u][4]]
+
+        if symbol_wo_pref in unit_symbol_lut and symbol_wo_pref in prefixable_units:
+            return possible_prefix, symbol_wo_pref
+    return "", symbol_str
+
+
+def _get_system_unit_string(dims, base_units):
+    # The dimensions of a unit object is the product of the base dimensions.
+    # Use sympy to factor the dimensions into base CGS unit symbols.
+    units = []
+    my_dims = dims.expand()
+    if my_dims is dimensions.dimensionless:
+        return ""
+    for factor in my_dims.as_ordered_factors():
+        dim = list(factor.free_symbols)[0]
+        unit_string = str(base_units[dim])
+        if factor.is_Pow:
+            power_string = "**(%s)" % factor.as_base_exp()[1]
+        else:
+            power_string = ""
+        units.append("(%s)%s" % (unit_string, power_string))
+    return " * ".join(units)
+
 
 unit_system_registry = {}
 
@@ -67,31 +111,27 @@ class UnitSystem(object):
         registry=None,
     ):
         self.registry = registry
-        if current_mks_unit is not None:
-            current_mks_unit = Unit(current_mks_unit, registry=self.registry)
         self.units_map = OrderedDict(
             [
-                (dimensions.length, Unit(length_unit, registry=self.registry)),
-                (dimensions.mass, Unit(mass_unit, registry=self.registry)),
-                (dimensions.time, Unit(time_unit, registry=self.registry)),
-                (
-                    dimensions.temperature,
-                    Unit(temperature_unit, registry=self.registry),
-                ),
-                (dimensions.angle, Unit(angle_unit, registry=self.registry)),
+                (dimensions.length, length_unit),
+                (dimensions.mass, mass_unit),
+                (dimensions.time, time_unit),
+                (dimensions.temperature, temperature_unit),
+                (dimensions.angle, angle_unit),
                 (dimensions.current_mks, current_mks_unit),
-                (
-                    dimensions.luminous_intensity,
-                    Unit(luminous_intensity_unit, registry=self.registry),
-                ),
+                (dimensions.luminous_intensity, luminous_intensity_unit),
             ]
         )
         for dimension, unit in self.units_map.items():
             # CGS sets the current_mks unit to none, so catch it here
-            if unit is None:
+            if unit is None and dimension is dimensions.current_mks:
                 continue
-            if unit.dimensions is not dimension:
+            if self.registry is not None and self.registry[unit][1] is not dimension:
                 raise IllDefinedUnitSystem(self.units_map)
+            elif self.registry is None:
+                bu = _split_prefix(unit, default_lut)[1]
+                if default_lut[bu][1] is not dimension:
+                    raise IllDefinedUnitSystem(self.units_map)
         self._dims = [
             "length",
             "mass",
@@ -107,15 +147,18 @@ class UnitSystem(object):
         self.name = name
 
     def __getitem__(self, key):
+        from unyt.unit_object import Unit
+
         if isinstance(key, string_types):
             key = getattr(dimensions, key)
         um = self.units_map
-        if key not in um or um[key] is None or um[key].dimensions is not key:
+        if key not in um or um[key] is None:
             if cmks in key.free_symbols and self.units_map[cmks] is None:
                 raise MissingMKSCurrent(self.name)
             units = _get_system_unit_string(key, self.units_map)
-            self.units_map[key] = Unit(units, registry=self.registry)
-        return self.units_map[key]
+            self.units_map[key] = units
+            return Unit(units, registry=self.registry)
+        return Unit(self.units_map[key], registry=self.registry)
 
     def __setitem__(self, key, value):
         if isinstance(key, string_types):
@@ -124,7 +167,7 @@ class UnitSystem(object):
             key = getattr(dimensions, key)
         if self.units_map[cmks] is None and cmks in key.free_symbols:
             raise MissingMKSCurrent(self.name)
-        self.units_map[key] = Unit(value, registry=self.registry)
+        self.units_map[key] = value
 
     def __str__(self):
         return self.name


Test patch:
diff --git a/unyt/tests/test_unit_systems.py b/unyt/tests/test_unit_systems.py
index 44bb431..546a0a9 100644
--- a/unyt/tests/test_unit_systems.py
+++ b/unyt/tests/test_unit_systems.py
@@ -59,6 +59,8 @@ def test_unit_system_id():
 def test_bad_unit_system():
     with pytest.raises(IllDefinedUnitSystem):
         UnitSystem("atomic", "nm", "fs", "nK", "rad")
+    with pytest.raises(IllDefinedUnitSystem):
+        UnitSystem("atomic", "nm", "fs", "nK", "rad", registry=UnitRegistry())
 
 
 def test_mks_current():
diff --git a/unyt/tests/test_units.py b/unyt/tests/test_units.py
index 6779f80..5d84582 100644
--- a/unyt/tests/test_units.py
+++ b/unyt/tests/test_units.py
@@ -44,7 +44,7 @@ from unyt.dimensions import (
 )
 from unyt.exceptions import InvalidUnitOperation, UnitsNotReducible, UnitConversionError
 from unyt.unit_object import default_unit_registry, Unit, UnitParseError
-from unyt.unit_systems import UnitSystem
+from unyt.unit_systems import cgs_unit_system, UnitSystem
 from unyt._unit_lookup_table import default_unit_symbol_lut, unit_prefixes
 import unyt.unit_symbols as unit_symbols
 from unyt._physical_ratios import (
@@ -470,6 +470,19 @@ def test_base_equivalent():
     with pytest.raises(UnitConversionError):
         u1.get_conversion_factor(Unit("degF"))
 
+    reg = UnitRegistry(unit_system=cgs_unit_system)
+
+    u = Unit("kg", registry=reg)
+
+    assert u.get_base_equivalent() == Unit("g")
+
+    u = Unit("kg")
+
+    assert u.get_base_equivalent() == Unit("kg")
+
+    u = Unit("A")
+    assert u.get_base_equivalent(unit_system="mks") == Unit("A")
+
 
 def test_temperature_offsets():
     u1 = Unit("degC")