using pint for units (#957)

* Proof of Concept: using pint for units

* Add delay context

* Fix mm alias

* Blessed units for backwards compatibility

* allow for all conversions, if kwarg is passed

* get preferred short name for wavenumber, fix spelling

* Big huge comment about 'blessed_units'

* Make labels nicer when no symbol is present

* clean up symbol logic, make it more readable, add theta for angles

* Do not include unit itself in valid conversion list

* Update docs

* Update api

* Fix code quality checks

* Exception raised by None units

* parentheses

* spelling

* UndefinedUnitError

* changelog

Co-authored-by: Blaise Thompson <blaise@untzag.com>

CHANGELOG.md

@@ -8,6 +8,9 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/).
 ### Added
 - `from_databroker` method to import Data objects from databroker catalogs
 
+### Changed
+- complete units overhaul, now using pint library
+
 ### Fixed
 - Avoid passing both `vmin/vmax` and `norm` to `pcolor*` methods
 

WrightTools/_dataset.py

@@ -306,9 +306,9 @@ def convert(self, destination_units):
             Units to convert into.
         """
         if not wt_units.is_valid_conversion(self.units, destination_units):
-            kind = wt_units.kind(self.units)
-            valid = list(wt_units.dicts[kind].keys())
-            raise wt_exceptions.UnitsError(valid, destination_units)
+            raise wt_exceptions.UnitsError(
+                wt_units.get_valid_conversions(self.units), destination_units
+            )
         if self.units is None:
             return
 

WrightTools/data/_axis.py

@@ -96,15 +96,15 @@ def label(self) -> str:
         label = self.expression.replace("_", "\\;")
         if self.units_kind:
             symbol = wt_units.get_symbol(self.units)
-            for v in self.variables:
-                vl = "%s_{%s}" % (symbol, v.label)
-                vl = vl.replace("_{}", "")  # label can be empty, no empty subscripts
-                label = label.replace(v.natural_name, vl)
-            units_dictionary = getattr(wt_units, self.units_kind)
-            label += r"\,"
-            label += r"\left("
-            label += units_dictionary[self.units][2]
-            label += r"\right)"
+            if symbol is not None:
+                for v in self.variables:
+                    vl = "%s_{%s}" % (symbol, v.label)
+                    vl = vl.replace("_{}", "")  # label can be empty, no empty subscripts
+
+                    label = label.replace(v.natural_name, vl)
+
+            label += fr"\,\left({wt_units.ureg.Unit(self.units):~}\right)"
+
         label = r"$\mathsf{%s}$" % label
         return label
 

WrightTools/data/_constant.py

@@ -63,20 +63,15 @@ def label(self) -> str:
         label = self.expression.replace("_", "\\;")
         if self.units_kind:
             symbol = wt_units.get_symbol(self.units)
-            for v in self.variables:
-                vl = "%s_{%s}" % (symbol, v.label)
-                vl = vl.replace("_{}", "")  # label can be empty, no empty subscripts
-                label = label.replace(v.natural_name, vl)
-                val = (
-                    round(self.value, self.round_spec)
-                    if self.round_spec is not None
-                    else self.value
-                )
+            if symbol is not None:
+                for v in self.variables:
+                    vl = "%s_{%s}" % (symbol, v.label)
+                    vl = vl.replace("_{}", "")  # label can be empty, no empty subscripts
+                    label = label.replace(v.natural_name, vl)
+        val = round(self.value, self.round_spec) if self.round_spec is not None else self.value
         label += r"\,=\,{}".format(format(val, self.format_spec))
         if self.units_kind:
-            units_dictionary = getattr(wt_units, self.units_kind)
-            label += r"\,"
-            label += units_dictionary[self.units][2]
+            label += fr"\,{wt_units.ureg.Unit(self.units):~}"
         label = r"$\mathsf{%s}$" % label
         return label
 

WrightTools/data/_data.py

@@ -776,11 +776,9 @@ def convert(self, destination_units, *, convert_variables=False, verbose=True):
             Convert a single axis object to compatable units. Call on an
             axis object in data.axes.
         """
-        # get kind of units
-        units_kind = wt_units.kind(destination_units)
         # apply to all compatible axes
         for axis in self.axes:
-            if axis.units_kind == units_kind:
+            if wt_units.is_valid_conversion(axis.units, destination_units):
                 orig = axis.units
                 axis.convert(destination_units, convert_variables=convert_variables)
                 if verbose:
@@ -791,7 +789,7 @@ def convert(self, destination_units, *, convert_variables=False, verbose=True):
                     )
         # apply to all compatible constants
         for constant in self.constants:
-            if constant.units_kind == units_kind:
+            if wt_units.is_valid_conversion(constant.units, destination_units):
                 orig = constant.units
                 constant.convert(destination_units, convert_variables=convert_variables)
                 if verbose:
@@ -802,7 +800,7 @@ def convert(self, destination_units, *, convert_variables=False, verbose=True):
                     )
         if convert_variables:
             for var in self.variables:
-                if wt_units.kind(var.units) == units_kind:
+                if wt_units.is_valid_conversion(var.units, destination_units):
                     orig = var.units
                     var.convert(destination_units)
                     if verbose:

WrightTools/units.py

@@ -4,92 +4,92 @@
 # --- import --------------------------------------------------------------------------------------
 
 
-import collections
 import warnings
 
-import numpy as np
-import numexpr
+import pint
 
 
 # --- define --------------------------------------------------------------------------------------
 
-
-# units are stored in dictionaries of like kind. format:
-#     unit : to native, from native, units_symbol, units_label
-
-# angle units (native: rad)
-angle = {"rad": ["x", "x", r"rad"], "deg": ["x/57.2958", "57.2958*x", r"deg"]}
-
-# delay units (native: fs)
-fs_per_mm = 3336.0
-delay = {
-    "fs": ["x", "x", r"fs"],
-    "ps": ["x*1e3", "x/1e3", r"ps"],
-    "ns": ["x*1e6", "x/1e6", r"ns"],
-    "mm_delay": ["x*2*fs_per_mm", "x/(2*fs_per_mm)", r"mm"],
-}
-
-# energy units (native: nm)
-energy = {
-    "nm": ["x", "x", r"nm"],
-    "wn": ["1e7/x", "1e7/x", r"cm^{-1}"],
-    "eV": ["1240./x", "1240./x", r"eV"],
-    "meV": ["1240000./x", "1240000./x", r"meV"],
-    "Hz": ["2.99792458e17/x", "2.99792458e17/x", r"Hz"],
-    "THz": ["2.99792458e5/x", "2.99792458e5/x", r"THz"],
-    "GHz": ["2.99792458e8/x", "2.99792458e8/x", r"GHz"],
-}
-
-# fluence units (native: uJ per sq. cm)
-fluence = {"uJ per sq. cm": ["x", "x", r"\frac{\mu J}{cm^{2}}"]}
-
-# optical density units (native: od)
-od = {"mOD": ["1e3*x", "x/1e3", r"mOD"], "OD": ["x", "x", r"OD"]}
-
-# position units (native: mm)
-position = {
-    "nm_p": ["x/1e6", "1e6*x", r"nm"],
-    "um": ["x/1000.", "1000.*x", r"\mu m"],
-    "mm": ["x", "x", r"mm"],
-    "cm": ["10.*x", "x/10.", r"cm"],
-    "in": ["x*25.4", "x/25.4", r"in"],
-}
-
-# pulse width units (native: FWHM)
-pulse_width = {"FWHM": ["x", "x", r"FWHM"]}
-
-# temperature units (native: K)
-temperature = {
-    "K": ["x", "x", r"K"],
-    "deg_C": ["x+273.15", "x-273.15", r"^\circ C"],
-    "deg_F": ["(x+459.67)*5/9", "x*9/5-456.67", r"^\circ F"],
-    "deg_R": ["x*5/9", "x*9/5", r"^\circ R"],
-}
-
-# time units (native: s)
-time = {
-    "fs_t": ["x/1e15", "x*1e15", r"fs"],
-    "ps_t": ["x/1e12", "x*1e12", r"ps"],
-    "ns_t": ["x/1e9", "x*1e9", r"ns"],
-    "us_t": ["x/1e6", "x*1e6", r"\mu s"],
-    "ms_t": ["x/1000.", "x*1000.", r"ms"],
-    "s_t": ["x", "x", r"s"],
-    "m_t": ["x*60.", "x/60.", r"m"],
-    "h_t": ["x*3600.", "x/3600.", r"h"],
-    "d_t": ["x*86400.", "x/86400.", r"d"],
-}
-
-dicts = collections.OrderedDict()
-dicts["angle"] = angle
-dicts["delay"] = delay
-dicts["energy"] = energy
-dicts["time"] = time
-dicts["position"] = position
-dicts["pulse_width"] = pulse_width
-dicts["fluence"] = fluence
-dicts["od"] = od
-dicts["temperature"] = temperature
-
+# Thise "blessed" units are here primarily for backwards compatibility, in particular
+# to enable the behavior of `data.convert` which will convert freely between the energy units
+# but does not go to time (where delay will)
+# Since both of these context can convert to [length] units, they are interconvertible, but we
+# do not want them to automatically do so.
+# This list is (at creation time) purely reflective of historical units supported pre pint
+# There is nothing preventing other units from being used and converted to, only to enable
+# expected behavior
+# 2021-01-29 KFS
+blessed_units = (
+    # angle
+    "rad",
+    "deg",
+    # delay
+    "fs",
+    "ps",
+    "ns",
+    "mm_delay",
+    # energy
+    "nm",
+    "wn",
+    "eV",
+    "meV",
+    "Hz",
+    "THz",
+    "GHz",
+    # optical density
+    "mOD",
+    # position
+    "nm_p",
+    "um",
+    "mm",
+    "cm",
+    "in",
+    # absolute temperature
+    "K",
+    "deg_C",
+    "deg_F",
+    "deg_R",
+    # time
+    "fs_t",
+    "ps_t",
+    "ns_t",
+    "us_t",
+    "ns_t",
+    "s_t",
+    "m_t",
+    "h_t",
+    "d_t",
+)
+
+ureg = pint.UnitRegistry()
+ureg.define("[fluence] = [energy] / [area]")
+
+ureg.define("OD = [] ")
+
+ureg.define("wavenumber = 1 / cm = cm^{-1} = wn")
+
+
+# Aliases for backwards compatability
+ureg.define("@alias s = s_t")
+ureg.define("@alias min = m_t")
+ureg.define("@alias hour = h_t")
+ureg.define("@alias d = d_t")
+
+ureg.define("@alias degC = deg_C")
+ureg.define("@alias degF = deg_F")
+ureg.define("@alias degR = deg_R")
+
+ureg.define("@alias m = m_delay")
+
+delay = pint.Context("delay", defaults={"n": 1, "num_pass": 2})
+delay.add_transformation(
+    "[length]", "[time]", lambda ureg, x, n=1, num_pass=2: num_pass * x / ureg.speed_of_light * n
+)
+delay.add_transformation(
+    "[time]", "[length]", lambda ureg, x, n=1, num_pass=2: x / num_pass * ureg.speed_of_light / n
+)
+ureg.enable_contexts("spectroscopy", delay)
 
 # --- functions -----------------------------------------------------------------------------------
 
@@ -111,23 +111,9 @@ def converter(val, current_unit, destination_unit):
     number
         Converted value.
     """
-    x = val
-    for dic in dicts.values():
-        if current_unit in dic.keys() and destination_unit in dic.keys():
-            try:
-                native = numexpr.evaluate(dic[current_unit][0], {"x": x})
-            except ZeroDivisionError:
-                native = np.inf
-            x = native  # noqa: F841
-            try:
-                out = numexpr.evaluate(dic[destination_unit][1], {"x": x})
-            except ZeroDivisionError:
-                out = np.inf
-            return out
-    # if all dictionaries fail
-    if current_unit is None and destination_unit is None:
-        pass
-    else:
+    try:
+        val = ureg.Quantity(val, current_unit).to(destination_unit).magnitude
+    except (pint.errors.DimensionalityError, pint.errors.UndefinedUnitError, AttributeError):
         warnings.warn(
             "conversion {0} to {1} not valid: returning input".format(
                 current_unit, destination_unit
@@ -152,48 +138,50 @@ def get_symbol(units) -> str:
     string
         LaTeX formatted symbol.
     """
-    if kind(units) == "energy":
-        d = {}
-        d["nm"] = r"\lambda"
-        d["wn"] = r"\bar\nu"
-        d["eV"] = r"\hslash\omega"
-        d["Hz"] = r"f"
-        d["THz"] = r"f"
-        d["GHz"] = r"f"
-        return d.get(units, "E")
-    elif kind(units) == "delay":
+    quantity = ureg.Quantity(1, ureg[units])
+    if quantity.check("[length]"):
+        return r"\lambda"
+    elif quantity.check("1 / [length]"):
+        return r"\bar\nu"
+    elif quantity.check("[energy]"):
+        return r"\hslash\omega"
+    elif quantity.check("1 / [time]"):
+        return "f"
+    elif quantity.check("[time]"):
         return r"\tau"
-    elif kind(units) == "fluence":
+    elif quantity.check("[fluence]"):
         return r"\mathcal{F}"
-    elif kind(units) == "pulse_width":
-        return r"\sigma"
-    elif kind(units) == "temperature":
-        return r"T"
+    elif quantity.check("[temperature]"):
+        return "T"
+    elif ureg[units] in (ureg.deg, ureg.radian):
+        return r"\omega"
     else:
-        return kind(units)
+        return None
 
 
-def get_valid_conversions(units) -> tuple:
+def get_valid_conversions(units, options=blessed_units) -> tuple:
+    return tuple(i for i in options if is_valid_conversion(units, i) and units != i)
+
+
+def is_valid_conversion(a, b, blessed=True) -> bool:
+    if a is None:
+        return b is None
+    if blessed and a in blessed_units and b in blessed_units:
+        blessed_energy_units = {"nm", "wn", "eV", "meV", "Hz", "THz", "GHz"}
+        if a in blessed_energy_units:
+            return b in blessed_energy_units
+        blessed_delay_units = {"fs", "ps", "ns", "mm_delay"}
+        if a in blessed_delay_units:
+            return b in blessed_delay_units
+        return ureg.Unit(a).dimensionality == ureg.Unit(b).dimensionality
     try:
-        valid = list(dicts[kind(units)])
-    except KeyError:
-        return ()
-    valid.remove(units)
-    return tuple(valid)
-
-
-def is_valid_conversion(a, b) -> bool:
-    for dic in dicts.values():
-        if a in dic.keys() and b in dic.keys():
-            return True
-    if a is None and b is None:
-        return True
-    else:
+        return ureg.Unit(a).is_compatible_with(b, "spectroscopy")
+    except pint.UndefinedUnitError:
         return False
 
 
 def kind(units):
-    """Find the kind of given units.
+    """Find the dimensionality of given units.
 
     Parameters
     ----------
@@ -205,6 +193,6 @@ def kind(units):
     string
         The kind of the given units. If no match is found, returns None.
     """
-    for k, v in dicts.items():
-        if units in v.keys():
-            return k
+    if units is None:
+        return None
+    return str(ureg.Unit(units).dimensionality)