A generic units class for pyiron

pyiron_base/generic/units.py

@@ -29,12 +29,12 @@ class PyironUnitRegistry:
 
     After instantiating, the `pint` units for different physical quantities can be registered as follows
 
-    >>> base_units = PyironUnitRegistry()
-    >>> base_units.add_quantity(quantity="energy", unit=pint_registry.eV, data_type=float)
+    >>> base_registry = PyironUnitRegistry()
+    >>> base_registry.add_quantity(quantity="energy", unit=pint_registry.eV, data_type=float)
 
     Labels corresponding to a particular physical quantity can also be registered
 
-    >>> base_units.add_labels(labels=["energy_tot", "energy_pot"], quantity="energy")
+    >>> base_registry.add_labels(labels=["energy_tot", "energy_pot"], quantity="energy")
 
     For more information on working with `pint`, see: https://pint.readthedocs.io/en/0.10.1/tutorial.html
 
@@ -63,7 +63,8 @@ def quantity_dict(self):
     @property
     def dtype_dict(self):
         """
-        A dictionary of the names of the different physical quantities to the corresponding datatype in which they are to be stored
+        A dictionary of the names of the different physical quantities to the corresponding datatype in which they are
+        to be stored
 
         Returns:
             dict
@@ -103,15 +104,17 @@ def add_labels(self, labels, quantity):
             labels (list/ndarray): List of labels
             quantity (str): Physical quantity associated with the labels
 
+        Raises:
+            KeyError: If quantity is not yet added with :method:`.add_quantity()`
+
         Note: `quantity` should already be a key of unit_dict
-Raises:
-    ValueError: if quantity is not yet added with :method:`.add_quantity()`.
+
         """
         for label in labels:
             if quantity in self.unit_dict.keys():
                 self._quantity_dict[label] = quantity
             else:
-                raise ValueError("Quantity {} is not defined. Use `add_quantity` to register the unit of this label")
+                raise KeyError("Quantity {} is not defined. Use `add_quantity` to register the unit of this label")
 
     def __getitem__(self, item):
         """
@@ -122,13 +125,16 @@ def __getitem__(self, item):
 
         Returns:
             pint.unit.Unit/pint.quantity.Quantity: The corresponding `pint` unit/quantity
+
+        Raises:
+            KeyError: If quantity is not yet added with :method:`.add_quantity()` or :method:`.add_labels()`
         """
         if item in self._unit_dict.keys():
             return self._unit_dict[item]
         elif item in self._quantity_dict.keys():
             return self._unit_dict[self._quantity_dict[item]]
         else:
-            raise ValueError("Quantity/label '{}' not registered in this unit registry".format(item))
+            raise KeyError("Quantity/label '{}' not registered in this unit registry".format(item))
 
     def get_dtype(self, quantity):
         """
@@ -139,13 +145,16 @@ def get_dtype(self, quantity):
 
         Returns:
             type: Corresponding data type
+
+        Raises:
+            KeyError: If quantity is not yet added with :method:`.add_quantity()` or :method:`.add_labels()`
         """
         if quantity in self._unit_dict.keys():
             return self._dtype_dict[quantity]
         elif quantity in self._quantity_dict.keys():
             return self._dtype_dict[self._quantity_dict[quantity]]
         else:
-            raise ValueError("Quantity/label '{}' not registered in this unit registry".format(quantity))
+            raise KeyError("Quantity/label '{}' not registered in this unit registry".format(quantity))
 
 
 class UnitConverter:
@@ -158,17 +167,17 @@ class UnitConverter:
 
     >>> import pint
     >>> pint_registry = pint.UnitRegistry()
-    >>> base_units = PyironUnitRegistry()
-    >>> base_units.add_quantity(quantity="energy", unit=pint_registry.eV)
-    >>> code_units = PyironUnitRegistry()
-    >>> code_units.add_quantity(quantity="energy",
-    >>>                         unit=1e3 * pint_registry.cal / (pint_registry.mol * pint_registry.N_A))
-    >>> unit_converter = UnitConverter(base_units=base_units, code_units=code_units)
+    >>> base = PyironUnitRegistry()
+    >>> base.add_quantity(quantity="energy", unit=pint_registry.eV)
+    >>> code = PyironUnitRegistry()
+    >>> code.add_quantity(quantity="energy",
+    >>>                         unit=pint_registry.kilocal / (pint_registry.mol * pint_registry.N_A))
+    >>> unit_converter = UnitConverter(base_registry=base, code_registry=code)
 
     The unit converter instance can then be used to obtain conversion factors between code and base units either as a
     `pint` quantity:
 
-    >>> print(unit_converter.code_to_base("energy"))
+    >>> print(unit_converter.code_to_base_pint("energy"))
     0.043364104241800934 electron_volt
 
     or as a scalar:
@@ -179,34 +188,34 @@ class UnitConverter:
     Alternatively, the unit converter can also be used as decorators for functions that return an array scaled into
     appropriate units:
 
-    >>> @unit_converter(quantity="energy", conversion="code_to_base")
+    >>> @unit_converter.code_to_base(quantity="energy")
     ... def return_ones():
     ...    return np.ones(5)
     >>> print(return_ones())
-    [0.0433641 0.0433641 0.0433641 0.0433641 0.0433641]
+    [0.0433641 0.0433641 0.0433641 0.0433641 0.0433641
 
     The decorator can also be used to assign units for numpy arrays
     (for more info see https://pint.readthedocs.io/en/0.10.1/numpy.html)
 
-    >>> @unit_converter(quantity="energy", conversion="base_units")
+    >>> @unit_converter.base_units(quantity="energy")
     ... def return_ones_ev():
     ...     return np.ones(5)
     >>> print(return_ones_ev())
     [1.0 1.0 1.0 1.0 1.0] electron_volt
 
     """
 
-    def __init__(self, base_units, code_units):
+    def __init__(self, base_registry, code_registry):
         """
 
         Args:
-            base_units (:class:`pyiron_base.generic.units.PyironUnitRegistry`): Base unit registry
-            code_units (pyiron_base.generic.units.PyironUnitRegistry): Code specific unit registry
+            base_registry (:class:`pyiron_base.generic.units.PyironUnitRegistry`): Base unit registry
+            code_registry (:class:`pyiron_base.generic.units.PyironUnitRegistry`): Code specific unit registry
         """
-        self._base_units = base_units
-        self._code_units = code_units
+        self._base_registry = base_registry
+        self._code_registry = code_registry
 
-    def code_to_base(self, quantity):
+    def code_to_base_pint(self, quantity):
         """
         Get the conversion factor as a `pint` quantity from code to base units
 
@@ -216,9 +225,9 @@ def code_to_base(self, quantity):
         Returns:
             pint.quantity.Quantity: Conversion factor as a `pint` quantity
         """
-        return (1 * self._code_units[quantity]).to(self._base_units[quantity])
+        return (1 * self._code_registry[quantity]).to(self._base_registry[quantity])
 
-    def base_to_code(self, quantity):
+    def base_to_code_pint(self, quantity):
         """
         Get the conversion factor as a `pint` quantity from base to code units
 
@@ -228,7 +237,7 @@ def base_to_code(self, quantity):
         Returns:
             pint.quantity.Quantity: Conversion factor as a `pint` quantity
         """
-        return (1 * self._base_units[quantity]).to(self._code_units[quantity])
+        return (1 * self._base_registry[quantity]).to(self._code_registry[quantity])
 
     def code_to_base_value(self, quantity):
         """
@@ -240,7 +249,7 @@ def code_to_base_value(self, quantity):
         Returns:
             float: Conversion factor as a float
         """
-        return self.code_to_base(quantity).magnitude
+        return self.code_to_base_pint(quantity).magnitude
 
     def base_to_code_value(self, quantity):
         """
@@ -252,15 +261,20 @@ def base_to_code_value(self, quantity):
         Returns:
             float: Conversion factor as a float
         """
-        return self.base_to_code(quantity).magnitude
+        return self.base_to_code_pint(quantity).magnitude
 
     def __call__(self, conversion, quantity):
         """
         Function call operator used as a decorator for functions that return numpy array
 
         Args:
-            conversion (str):
-            quantity (str):
+            conversion (str): Conversion type which should be one of
+                'code_to_base' To multiply by the code to base units conversion factor
+                'base_to_code' To multiply by the base to code units conversion factor
+                'code_units' To assign code units to the nunpy array returned by the decorated function
+                'base_units' To assign base units to the nunpy array returned by the decorated function
+
+            quantity (str): Name of quantity
 
         Returns:
             function: Decorated function
@@ -270,32 +284,86 @@ def _decorate_to_base(function):
                 @functools.wraps(function)
                 def dec(*args, **kwargs):
                     return np.array(function(*args, **kwargs) * self.code_to_base_value(quantity),
-                                    dtype=self._base_units.get_dtype(quantity))
+                                    dtype=self._base_registry.get_dtype(quantity))
                 return dec
             return _decorate_to_base
         elif conversion == "base_to_code":
             def _decorate_to_code(function):
                 @functools.wraps(function)
                 def dec(*args, **kwargs):
                     return np.array(function(*args, **kwargs) * self.base_to_code_value(quantity),
-                                    dtype=self._code_units.get_dtype(quantity))
+                                    dtype=self._code_registry.get_dtype(quantity))
                 return dec
             return _decorate_to_code
         elif conversion == "base_units":
             def _decorate_base_units(function):
                 @functools.wraps(function)
                 def dec(*args, **kwargs):
-                    return Q_(np.array(function(*args, **kwargs), dtype=self._base_units.get_dtype(quantity)),
-                              self._base_units[quantity])
+                    return Q_(np.array(function(*args, **kwargs), dtype=self._base_registry.get_dtype(quantity)),
+                              self._base_registry[quantity])
                 return dec
             return _decorate_base_units
         elif conversion == "code_units":
             def _decorate_code_units(function):
                 @functools.wraps(function)
                 def dec(*args, **kwargs):
-                    return Q_(np.array(function(*args, **kwargs), dtype=self._code_units.get_dtype(quantity)),
-                              self._code_units[quantity])
+                    return Q_(np.array(function(*args, **kwargs), dtype=self._code_registry.get_dtype(quantity)),
+                              self._code_registry[quantity])
                 return dec
             return _decorate_code_units
         else:
             raise ValueError("Conversion type {} not implemented!".format(conversion))
+
+    def code_to_base(self, quantity):
+        """
+        Decorator for functions that returns a numpy array. Multiples the function output by the code to base units
+        conversion factor
+
+        Args:
+            quantity (str):  Name of the quantity
+
+        Returns:
+            function: Decorated function
+
+        """
+        return self(quantity=quantity, conversion="code_to_base")
+
+    def base_to_code(self, quantity):
+        """
+        Decorator for functions that returns a numpy array. Multiples the function output by the base to code units
+        conversion factor
+
+        Args:
+            quantity (str):  Name of the quantity
+
+        Returns:
+            function: Decorated function
+
+        """
+        return self(quantity=quantity, conversion="base_to_code")
+
+    def code_units(self, quantity):
+        """
+        Decorator for functions that returns a numpy array. Assigns the code unit of the quantity to the function output
+
+        Args:
+            quantity (str):  Name of the quantity
+
+        Returns:
+            function: Decorated function
+
+        """
+        return self(quantity=quantity, conversion="code_units")
+
+    def base_units(self, quantity):
+        """
+        Decorator for functions that returns a numpy array. Assigns the base unit of the quantity to the function output
+
+        Args:
+            quantity (str):  Name of the quantity
+
+        Returns:
+            function: Decorated function
+
+        """
+        return self(quantity=quantity, conversion="base_units")

tests/generic/test_units.py

@@ -18,21 +18,26 @@ def test_units(self):
         self.assertRaises(ValueError, base_units.add_quantity, quantity="energy", unit="eV")
         dim_less_q = pint_registry.Quantity
         base_units.add_quantity(quantity="dimensionless_integer_quantity", unit=dim_less_q(1), data_type=int)
+        self.assertIsInstance(base_units.quantity_dict, dict)
+        self.assertIsInstance(base_units.unit_dict, dict)
+        self.assertEqual(base_units.unit_dict["energy"], pint_registry.eV)
+        self.assertIsInstance(base_units.dtype_dict, dict)
+        self.assertIsInstance(base_units.dtype_dict["energy"], float.__class__)
         code_units = PyironUnitRegistry()
         # Define unit kJ/mol
         code_units.add_quantity(quantity="energy",
                                 unit=pint_registry.kilocal / (pint_registry.mol * pint_registry.N_A))
         code_units.add_labels(labels=["energy_tot", "energy_pot"], quantity="energy")
         # Raise Error for undefined quantity
-        self.assertRaises(ValueError, code_units.add_labels, labels=["mean_forces"], quantity="force")
+        self.assertRaises(KeyError, code_units.add_labels, labels=["mean_forces"], quantity="force")
         self.assertTrue(code_units["energy"], code_units["energy_tot"])
         self.assertTrue(code_units["energy"], code_units["energy_pot"])
         # Define converter
-        unit_converter = UnitConverter(base_units=base_units, code_units=code_units)
+        unit_converter = UnitConverter(base_registry=base_units, code_registry=code_units)
         self.assertAlmostEqual(round(unit_converter.code_to_base_value("energy"), 3), 0.043)
         # Raise error if quantity not defined in any of the unit registries
-        self.assertRaises(ValueError, unit_converter.code_to_base_value, "dimensionless_integer_quantity")
-        self.assertRaises(ValueError, code_units.get_dtype, "dimensionless_integer_quantity")
+        self.assertRaises(KeyError, unit_converter.code_to_base_value, "dimensionless_integer_quantity")
+        self.assertRaises(KeyError, code_units.get_dtype, "dimensionless_integer_quantity")
         # Define dimensionless quantity in the code units registry
         code_units.add_quantity(quantity="dimensionless_integer_quantity", unit=dim_less_q(1), data_type=int)
         self.assertIsInstance(code_units.get_dtype("dimensionless_integer_quantity"), int.__class__)
@@ -42,22 +47,22 @@ def test_units(self):
                                * unit_converter.base_to_code_value("energy"), 1)
 
         # Use decorator to convert units
-        @unit_converter(quantity="energy", conversion="code_to_base")
+        @unit_converter.code_to_base(quantity="energy")
         def return_ones_base():
             return np.ones(10)
 
-        @unit_converter(quantity="energy", conversion="base_to_code")
+        @unit_converter.base_to_code(quantity="energy")
         def return_ones_code():
             return np.ones(10)
 
-        @unit_converter(quantity="energy", conversion="base_units")
+        @unit_converter.base_units(quantity="energy")
         def return_ones_ev():
             return np.ones(10)
 
-        @unit_converter(quantity="energy", conversion="code_units")
+        @unit_converter.code_units(quantity="energy")
         def return_ones_kj_mol():
             return np.ones(10)
-        print(return_ones_ev())
+
         self.assertEqual(1 * return_ones_kj_mol().units, 1 * code_units["energy"])
         self.assertEqual(1 * return_ones_ev().units, 1 * base_units["energy"])
         self.assertTrue(np.allclose(return_ones_base(), np.ones(10) * 0.0433641))
@@ -66,11 +71,11 @@ def return_ones_kj_mol():
         # Define dimensionally incorrect units
         code_units.add_quantity(quantity="energy",
                                 unit=pint_registry.N * pint_registry.metre ** 2)
-        unit_converter = UnitConverter(base_units=base_units, code_units=code_units)
+        unit_converter = UnitConverter(base_registry=base_units, code_registry=code_units)
         # Check if dimensionality error raised
         self.assertRaises(pint.DimensionalityError, unit_converter.code_to_base_value, "energy")
         # Try SI units
         code_units.add_quantity(quantity="energy",
                                 unit=pint_registry.N * pint_registry.metre)
-        unit_converter = UnitConverter(base_units=base_units, code_units=code_units)
+        unit_converter = UnitConverter(base_registry=base_units, code_registry=code_units)
         self.assertAlmostEqual(round(unit_converter.code_to_base_value("energy") / 1e18, 3), 6.242)