Add unit system using pint

# Conflicts:
#	est/core/process/process.py
#	est/core/process/pymca/test/test_io.py
#	est/core/process/test/test_roi.py
#	est/core/test/test_types.py
#	est/io/io.py
#	est/setup.py

# Conflicts:
#	est/core/types.py
#	requirements.txt

est/core/io.py

@@ -33,6 +33,7 @@ from est.io import read_xas, write_xas, get_xasproc
 from est.core.types import XASObject
 from silx.io.url import DataUrl
 from est.core.types import Dim
+from est.units import ur, convert_to
 import h5py
 import logging
 
@@ -45,7 +46,8 @@ DEFAULT_CHANNEL_PATH = "/data/NXdata/Channel"
 DEFAULT_CONF_PATH = "/configuration"
 
 
-def read(spectra_url, channel_url, config_url=None, dimensions=None):
+def read(spectra_url, channel_url, config_url=None, dimensions=None,
+         energy_unit=ur.eV):
     """
 
     :param DataUrl spectra_url: data url to the spectra
@@ -68,10 +70,11 @@ def read(spectra_url, channel_url, config_url=None, dimensions=None):
         channel_url=channel_url,
         config_url=config_url,
         dimensions=dimensions_,
+        energy_unit=energy_unit,
     )
 
 
-def read_frm_file(file_path):
+def read_frm_file(file_path, energy_unit=ur.eV):
     """
 
     :param str file_path: path to the file containing the spectra. Must ba a
@@ -80,24 +83,26 @@ def read_frm_file(file_path):
     :rtype: XASObject
     """
     reader = XASReader()
-    return reader.read_from_file(file_path=file_path)
+    return reader.read_from_file(file_path=file_path, energy=energy_unit)
 
 
 class XASReader(object):
     """Simple reader of a xas file"""
 
     @staticmethod
-    def read_frm_url(spectra_url, channel_url, dimensions=None, config_url=None):
+    def read_frm_url(spectra_url, channel_url, dimensions=None,
+                     config_url=None, energy_unit=ur.eV):
         sp, en, conf = read_xas(
             spectra_url=spectra_url,
             channel_url=channel_url,
             config_url=config_url,
             dimensions=dimensions,
+            energy_unit=energy_unit,
         )
         return XASObject(spectra=sp, energy=en, configuration=conf)
 
     @staticmethod
-    def read_from_file(file_path):
+    def read_from_file(file_path, energy_unit=ur.eV):
         """
 
         :param str file_path:
@@ -107,11 +112,13 @@ class XASReader(object):
             return XASReader.read_frm_url(
                 spectra_url=DataUrl(file_path=file_path, scheme="PyMca"),
                 channel_url=DataUrl(file_path=file_path, scheme="PyMca"),
+                energy_unit=energy_unit,
             )
         elif file_path.endswith(".xmu"):
             return XASReader.read_frm_url(
                 spectra_url=DataUrl(file_path=file_path, scheme="larch"),
                 channel_url=DataUrl(file_path=file_path, scheme="larch"),
+                energy_unit=energy_unit,
             )
         elif h5py.is_hdf5(file_path):
             return XASReader.read_frm_url(
@@ -124,6 +131,7 @@ class XASReader(object):
                 config_url=DataUrl(
                     file_path=file_path, scheme="silx", data_path="configuration"
                 ),
+                energy_unit=energy_unit,
             )
         else:
             raise ValueError("file type not managed, unable to load")

est/core/process/process.py

@@ -32,6 +32,7 @@ __date__ = "07/08/2019"
 from .progress import Progress
 from est.core.types import XASObject
 import logging
+import pint
 
 _logger = logging.getLogger(__name__)
 
@@ -83,7 +84,7 @@ class Process(object):
         if _xas_obj.n_spectrum > 0:
             for spectrum in _xas_obj.spectra:
                 assert isinstance(spectrum, Spectrum)
-            assert isinstance(spectrum.energy, numpy.ndarray)
+            assert isinstance(spectrum.energy, (numpy.ndarray, pint.Quantity))
             assert isinstance(spectrum.mu, numpy.ndarray)
         assert isinstance(_xas_obj, XASObject)
         return _xas_obj

est/core/process/pymca/test/test_io.py

@@ -92,7 +92,8 @@ class TestNxWriting(unittest.TestCase):
             self.assertTrue("absorbed_beam" in hdf["scan1"].keys())
             self.assertTrue("monochromator" in hdf["scan1"].keys())
 
-        loaded_xas_obj = XASObject.from_file(self.h5_file, configuration_path=None)
+        loaded_xas_obj = XASObject.from_file(self.h5_file,
+                                             configuration_path=None)
         numpy.testing.assert_allclose(loaded_xas_obj.energy, self.xas_obj.energy)
         numpy.testing.assert_allclose(
             loaded_xas_obj.absorbed_beam(), self.xas_obj.absorbed_beam()

est/core/process/roi.py

@@ -31,6 +31,7 @@ from est.core.types import XASObject, Spectrum
 from .process import Process
 import logging
 import numpy
+import pint
 
 _logger = logging.getLogger(__name__)
 
@@ -120,9 +121,8 @@ class ROIProcess(Process):
         # TODO: create a spectra object which deal automatically with
         # the following. Should be part of spectra object
         volumes = {}
-
         for key in xas_obj.spectra_keys():
-            if isinstance(xas_obj.spectra[0][key], numpy.ndarray):
+            if isinstance(xas_obj.spectra[0][key], (numpy.ndarray, pint.Quantity)):
                 # there is no processing for the _larch_grp_members case
                 if key == "_larch_grp_members":
                     continue
@@ -178,6 +178,7 @@ class ROIProcess(Process):
         """
         assert xas_obj is not None
         _xas_obj = self.getXasObject(xas_obj=xas_obj)
+
         # existing roi is priority. This is the case if called from pushworkflow
         # for example.
         if self._roi is not None:

est/core/test/test_io.py

@@ -32,6 +32,7 @@ from est.core.types import Dim
 from est.core.io import XASReader
 from silx.io.url import DataUrl
 from est.core.types import XASObject
+from est.units import ur
 import numpy
 import os
 import tempfile
@@ -126,7 +127,8 @@ class TestSpectraDimensions(unittest.TestCase):
         self.assertTrue(isinstance(xas_obj, XASObject))
         self.assertTrue(xas_obj.n_spectrum == x_dim * y_dim)
         numpy.testing.assert_array_equal(xas_obj.spectra[1].mu, spectra[:, 0, 1])
-        numpy.testing.assert_array_equal(xas_obj.spectra[2].energy, channel)
+        numpy.testing.assert_array_equal(xas_obj.spectra[2].energy.m,
+                                         (channel * ur.eV).m)
 
 
 def suite():

est/core/test/test_types.py

@@ -37,7 +37,9 @@ from est.core.types import Spectrum, XASObject, Dim
 from est.core.utils import spectra as spectra_utils
 from est.core.io import read as read_xas
 from silx.io.url import DataUrl
+from est.units import ur
 import json
+import pint
 import silx.io.utils
 
 try:
@@ -67,7 +69,7 @@ class TestSpectrum(unittest.TestCase):
         energy = numpy.arange(10, 20)
         mu = numpy.arange(10)
         spectrum = Spectrum(energy=energy, mu=mu)
-        numpy.testing.assert_array_equal(spectrum.energy, energy)
+        numpy.testing.assert_array_equal(spectrum.energy.m, (energy * ur.eV).m)
         numpy.testing.assert_array_equal(spectrum.mu, mu)
         mu_2 = numpy.arange(30, 40)
         spectrum["Mu"] = mu_2
@@ -184,7 +186,7 @@ class TestXASObjectSerialization(unittest.TestCase):
         # make sure we find a comparable xas object from it
         xas_obj_2 = XASObject.from_dict(dict_xas_obj)
 
-        numpy.testing.assert_array_equal(xas_obj.energy, xas_obj_2.energy)
+        numpy.testing.assert_array_equal(xas_obj.energy.m, xas_obj_2.energy.m)
         self.assertEqual(xas_obj, xas_obj_2)
 
         # simple test without the process_details
@@ -204,6 +206,7 @@ class TestXASObjectSerialization(unittest.TestCase):
             energy_url=self.url_energy,
             spectra_url=self.spectra_path,
         )
+        self.assertTrue(isinstance(xas_obj.energy, pint.Quantity))
         # if no h5 file defined, should fail to copy it to a dictionary
         with self.assertRaises(ValueError):
             xas_obj.to_dict()
@@ -215,8 +218,10 @@ class TestXASObjectSerialization(unittest.TestCase):
         json.dumps(dict_xas_obj)
         # make sure we find a comparable xas object from it
         xas_obj_2 = XASObject.from_dict(dict_xas_obj)
+        self.assertTrue(isinstance(xas_obj_2.energy, pint.Quantity))
 
-        numpy.testing.assert_array_equal(xas_obj.energy, xas_obj_2.energy)
+        self.assertEqual(xas_obj.energy.units, xas_obj_2.energy.units)
+        numpy.testing.assert_array_equal(xas_obj.energy.m, xas_obj_2.energy.m)
         self.assertEqual(xas_obj, xas_obj_2)
 
         # simple test without the process_details

est/core/types.py

@@ -39,7 +39,8 @@ import os
 import shutil
 from silx.utils.enum import Enum
 from silx.io.utils import h5py_read_dataset
-
+from est.units import units, ur, convert_to as convert_unit_to
+import pint
 
 _logger = logging.getLogger(__name__)
 
@@ -178,6 +179,8 @@ class XASObject(object):
     @spectra.setter
     def spectra(self, energy_spectra):
         energy, spectra, dim1, dim2 = energy_spectra
+        if isinstance(energy, numpy.ndarray):
+            energy = energy * ur.eV
         if spectra is None:
             self.__spectra = []
             self.__energy = energy
@@ -210,7 +213,10 @@ class XASObject(object):
                     assert isinstance(spectrum, Spectrum)
                     self.add_spectrum(spectrum)
             self._updateSpectraIndexes()
-        self.energy = energy
+        if energy is not None:
+            self.energy = convert_unit_to(energy, ur.eV)
+        else:
+            self.energy = None
 
     def _setSpectra(self, spectra):
         self.__spectra = spectra
@@ -260,7 +266,12 @@ class XASObject(object):
 
     @energy.setter
     def energy(self, energy):
-        self.__energy = energy
+        if isinstance(energy, numpy.ndarray):
+            energy = energy * ur.eV
+        if energy is None:
+            self.__energy = None
+        else:
+            self.__energy = convert_unit_to(energy, ur.eV)
         if len(self.__spectra) > 0:
             if len(self.__spectra[0].energy) != len(energy):
                 _logger.warning("spectra and energy have incoherent dimension")
@@ -306,7 +317,7 @@ class XASObject(object):
                 entry = "/".join((self.entry, "est_saving_pt", "channel"))
                 if entry in h5f:
                     del h5f[entry]
-                h5f[entry] = "None" if self.energy is None else self.energy
+                h5f[entry] = "None" if self.energy is None else convert_unit_to(self.energy, ur.eV).m
 
     def to_dict(self, with_process_details=True):
         """convert the XAS object to a dict
@@ -343,7 +354,6 @@ class XASObject(object):
                 return DataUrl(
                     file_path=self.linked_h5_file, data_path=data_path, scheme="silx"
                 ).path()
-
         self._create_saving_pt()
         spectra_ = get_spectra_and_processing()
         res = {
@@ -408,6 +418,8 @@ class XASObject(object):
                 except Exception:
                     raise KeyError("fail to access to", key)
                 else:
+                    if isinstance(value, pint.Quantity):
+                        value = convert_unit_to(value, ur.eV).m
                     if _has_larch and isinstance(value, larch.symboltable.Group):
                         _logger.info("pass larch details, not managed for now")
                         continue
@@ -438,7 +450,7 @@ class XASObject(object):
             if isinstance(spectra, str):
                 spectra = load_data(data_url=DataUrl(path=spectra), name="spectra")
             # if come from a list of spectrum
-            elif not isinstance(spectra, numpy.ndarray):
+            elif not isinstance(spectra, (numpy.ndarray, pint.Quantity)):
                 new_spectra = []
                 for spectrum in spectra:
                     assert isinstance(spectrum, dict)
@@ -542,7 +554,8 @@ class XASObject(object):
     def __eq__(self, other):
         return (
             isinstance(other, XASObject)
-            and numpy.array_equal(self.energy, other.energy)  # noqa
+            and ((self.energy is None and other.energy is None) or  # noqa
+                 numpy.array_equal(self.energy.m, other.energy.m))  # noqa
             and self.dim1 == other.dim1  # noqa
             and self.dim2 == other.dim2  # noqa
             and self.configuration == other.configuration  # noqa
@@ -762,7 +775,9 @@ class Spectrum(_Spectrum_Base):
     def __init__(self, energy=None, mu=None, x=None, y=None):
         super().__init__()
         if energy is not None:
-            assert isinstance(energy, numpy.ndarray)
+            assert isinstance(energy, (numpy.ndarray, pint.Quantity))
+            if isinstance(energy, numpy.ndarray):
+                energy = energy * ur.eV
 
         self.__x = x
         self.__y = y
@@ -801,8 +816,9 @@ class Spectrum(_Spectrum_Base):
         return self.__energy
 
     @energy.setter
+    @units(energy='eV')
     def energy(self, energy):
-        assert isinstance(energy, numpy.ndarray) or energy is None
+        assert isinstance(energy, (numpy.ndarray, pint.Quantity)) or energy is None
         self.__energy = energy
 
     @property
@@ -863,7 +879,7 @@ class Spectrum(_Spectrum_Base):
 
     @normalized_energy.setter
     def normalized_energy(self, energy):
-        assert isinstance(energy, numpy.ndarray) or energy is None
+        assert isinstance(energy, (numpy.ndarray, pint.Quantity)) or energy is None
         self.__normalized_energy = energy
 
     @property
@@ -994,7 +1010,7 @@ class Spectrum(_Spectrum_Base):
             self._X_POS_KEY: self.x,
             self._Y_POS_KEY: self.y,
             self._MU_KEY: self.mu,
-            self._ENERGY_KEY: self.energy,
+            self._ENERGY_KEY: convert_unit_to(self.energy, ur.eV).m,
             self._FT_KEY: self.ft.to_dict(),
             self._NORMALIZED_MU_KEY: "None"
             if self.normalized_mu is None

est/io/io.py

@@ -35,6 +35,7 @@ from silx.io import utils
 from silx.io.dictdump import dicttoh5, h5todict
 from silx.io.url import DataUrl
 from silx.utils.enum import Enum
+from est.units import ur, convert_to
 
 try:
     from est.io.utils.pymca import read_spectrum as pymca_read_spectrum
@@ -148,7 +149,8 @@ def load_data(data_url, name):
             _logger.warning("invalid url for", name, ",  will not load it")
 
 
-def read_xas(spectra_url, channel_url, dimensions=None, config_url=None):
+def read_xas(spectra_url, channel_url, dimensions=None, config_url=None,
+             energy_unit=ur.eV):
     """
     Read the given spectra url and the config url if any
 
@@ -218,7 +220,7 @@ def read_xas(spectra_url, channel_url, dimensions=None, config_url=None):
         )
         raise ValueError(err)
 
-    return (spectra, energy, configuration)
+    return (spectra, energy * energy_unit, configuration)
 
 
 def write_xas_proc(
@@ -352,6 +354,7 @@ def write_xas(
         nx_entry.attrs["NX_class"] = "NXentry"
 
         # store energy
+<<<<<<< HEAD
         nx_monochromator = nx_entry.require_group("monochromator")
         nx_monochromator.attrs["NX_class"] = "NXmonochromator"
         if overwrite and "energy" in nx_monochromator:
@@ -359,6 +362,16 @@ def write_xas(
         nx_monochromator["energy"] = energy
         nx_monochromator["energy"].attrs["interpretation"] = "spectrum"
         nx_monochromator["energy"].attrs["NX_class"] = "NXdata"
+=======
+        nx_monochromator = nx_entry.require_group('monochromator')
+        nx_monochromator.attrs['NX_class'] = "NXmonochromator"
+        if overwrite and 'energy' in nx_monochromator:
+            del nx_monochromator['energy']
+        nx_monochromator['energy'] = convert_to(energy, ur.eV).m
+        nx_monochromator['energy'].attrs['interpretation'] = 'spectrum'
+        nx_monochromator['energy'].attrs['NX_class'] = "NXdata"
+        nx_monochromator['energy'].attrs['unit'] = "eV"
+>>>>>>> 91be990... Add unit system using pint
 
         # store absorbed beam
         nx_absorbed_beam = nx_entry.require_group("absorbed_beam")

est/io/utils/pymca.py

@@ -28,9 +28,10 @@ __license__ = "MIT"
 __date__ = "06/11/2019"
 
 from PyMca5.PyMcaIO import specfilewrapper as specfile
+from est.units import ur
 
 
-def read_spectrum(spec_file):
+def read_spectrum(spec_file, energy_unit=ur.eV):
     """
 
     :param spec_file: path to the spec file containing the spectra definition
@@ -66,4 +67,6 @@ def read_spectrum(spec_file):
             else:
                 energy = data[0, :]
                 mu = data[1, :]
+    if energy is not None:
+        energy = energy * energy_unit
     return energy, mu

est/setup.py

@@ -39,6 +39,7 @@ def configuration(parent_package="", top_path=None):
     config.add_subpackage("io")
     config.add_subpackage("resources")
     config.add_subpackage("test")
+    config.add_subpackage("units")
 
     return config
 

est/units/__init__.py

+# -*- coding: utf-8 -*-
+#
+# This file is part of the bliss project
+#
+# Copyright (c) 2015-2020 Beamline Control Unit, ESRF
+# Distributed under the GNU LGPLv3. See LICENSE for more info.
+
+"""Physical units for engineering and science
+
+Implementation based on :mod:`pint`.
+
+Usage::
+
+    >>> from bliss.physics.units import ur
+
+    >>> # use ur as a pint UnitRegistry
+    >>> mass = 0.1*ur.mg
+    >>> E = mass * ur.c**2
+    >>> print( E.to(ur.kJ) )
+    8987551.78737 kilojoule
+
+    >>> # decorate your methods to ensure proper units
+    >>> @ur.units(mass='kg', result='J')
+    ... def energy(mass):
+    ...     return mass * ur.c**2
+
+    >>> # passing a Quantity will give you back a Quantity
+    >>> print( energy(mass) )
+    8987551787.37 joule
+
+    >>> # passing a float will assume it is in the proper (kg)
+    >>> # and will give you back a float in the corresponding unit (J)
+    >>> print( energy(0.1e-6) )
+    8987551787.37
+
+
+Advanced usage
+--------------
+
+Use another UnitRegistry
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+Bliss unit system creates a default UnitRegistry (ur).
+
+If you want to interact with another library which also uses pint it is
+important that both libraries use the same UnitRegistry.
+
+This library allows you to change the active UnitRegistry. You should do
+this as soon as possible in the code of your application::
+
+    import pint
+    ureg = UnitRegistry()
+
+    from bliss.physics import units
+    units.ur = ureg
+
+"""
+
+from functools import wraps
+from inspect import getfullargspec
+
+import pint
+
+__all__ = ["ur", "units"]
+
+#: unit registry
+ur = pint.UnitRegistry()
+
+
+def is_quantity(arg):
+    """Return whether the given argument is a quantity."""
+    return isinstance(arg, ur.Quantity)
+
+
+def to_unit(arg):
+    """Permissively cast the given argument into a unit."""
+    return ur.Unit(arg) if arg else None
+
+
+def values_to_units(dct):
+    """Cast the values of the given dict into units"""
+    return {k: to_unit(v) for k, v in dct.items()}
+
+
+def convert_to(arg, unit):
+    """Permissively convert the given argument into the given unit.
+
+    The argument can either be a number or a quantity.
+    """
+    if not unit:
+        return arg
+    if arg is None:
+        return None
+    else:
+        return arg.to(unit) if is_quantity(arg) else arg * unit
+
+
+def units(**kwarg_units):
+    """
+    Use as a decorator to protect your function against unit errors.
+
+    Each keyword argument must be an argument of the function. The
+    value is the unit (string or Unit) in which your function argument
+    should be called with.
+
+    An extra argument *result* should provide the Unit for the expected
+    return value. And yes, you cannot have a function which has an
+    argument called *result* but that is just good naming!
+
+    Missing arguments will be ignored.
+
+    Example::
+
+        from bliss.physics.units import ur, units
+
+        @units(mass='kg', result=ur.J)
+        def energy(mass):
+            return mass * ur.c**2
+
+    When you call a decorated function, it will return a Quantity if the
+    *result* is a Unit and at least one of the arguments is a Quantity. If none
+    of the arguments is a Quantity the result is a float with a value in the
+    units specified by *result*
+    """
+    result_unit = to_unit(kwarg_units.pop("result", None))
+    kwarg_units = values_to_units(kwarg_units)
+
+    def decorator(func):
+        arg_spec = getfullargspec(func).args
+        if not set(arg_spec).issuperset(kwarg_units):
+            raise TypeError("units argument names differ from function argument names")
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            # Everything is a kwargs
+            kwargs.update(zip(arg_spec, args))
+            # Check for quantity-free use case
+            all_magnitude = all(
+                not is_quantity(value)
+                for key, value in kwargs.items()
+                if key in kwarg_units
+            )
+            # Kwargs conversion
+            kwargs = {
+                key: convert_to(value, kwarg_units.get(key))
+                for key, value in kwargs.items()
+            }
+            # Call the actual func
+            result = func(**kwargs)
+            if not result_unit:
+                return result
+            # Safety check
+            if not is_quantity(result):
+                raise TypeError("Function {!r} did not return a quantity".format(func))
+            # Convert the result and return magnitude or quantity
+            result = convert_to(result, result_unit)
+            return result.magnitude if all_magnitude else result
+
+        return wrapper
+
+    return decorator

requirements.txt

 silx>=0.14b
 pypushflow>=0.2.0b
-h5py>=3.1
\ No newline at end of file
+h5py>=3.1
+Pint