Add support for UFloat in PintArray (#139)

CHANGES

@@ -14,6 +14,7 @@ pint-pandas Changelog
 - Fixed bug preventing adding to empty slices of PintArrays #69
 - Notebook updated to show plotting #116 #9 #43
 - Tests reorganised #131
+- Allow UFloat as type of magnitude supported in PintArray
 
 0.2 (2021-03-23)
 ----------------

pint_pandas/pint_array.py

@@ -19,8 +19,14 @@
 from pandas.core.arrays.base import ExtensionOpsMixin
 from pandas.core.indexers import check_array_indexer
 from pint import compat, errors
-from pint.quantity import _Quantity
-from pint.unit import _Unit
+from pint.compat import HAS_UNCERTAINTIES
+from pint.facets.plain.quantity import PlainQuantity as _Quantity
+from pint.facets.plain.unit import PlainUnit as _Unit
+
+if HAS_UNCERTAINTIES:
+    from uncertainties import UFloat, ufloat
+    from uncertainties import unumpy as unp
+    _ufloat_nan = ufloat(np.nan, 0)
 
 
 class PintType(ExtensionDtype):
@@ -131,7 +137,10 @@ def name(self):
 
     @property
     def na_value(self):
-        return self.ureg.Quantity(np.nan, self.units)
+        if HAS_UNCERTAINTIES:
+            return self.ureg.Quantity(_ufloat_nan, self.units)
+        else:
+            return self.ureg.Quantity(np.nan, self.units)
 
     def __hash__(self):
         # make myself hashable
@@ -216,13 +225,31 @@ def __init__(self, values, dtype=None, copy=False):
         if not isinstance(values, np.ndarray):
             values = np.array(values, copy=copy)
             copy = False
-        if not np.issubdtype(values.dtype, np.floating):
-            warnings.warn(
-                f"pint-pandas does not support magnitudes of {values.dtype}. Converting magnitudes to float.",
-                category=RuntimeWarning,
-            )
-            values = values.astype(float)
-            copy = False
+        if HAS_UNCERTAINTIES:
+            if np.issubdtype(values.dtype, np.floating):
+                pass
+            elif not all([isinstance(v, UFloat) for v in values]):
+                warnings.warn(
+                    f"pint-pandas does not support certain magnitudes of {values.dtype}. Converting magnitudes to ufloat.",
+                    category=RuntimeWarning,
+                )
+                for i in range(len(values)):
+                    # List comprehensions are great, but they are not np.arrays!
+                    if not isinstance(values[i], UFloat):
+                        if np.isnan(values[i]):
+                            values[i] = _ufloat_nan
+                        else:
+                            values[i] = ufloat(values[i], 0)
+                    elif unp.isnan(values[i]):
+                        values[i] = _ufloat_nan
+                copy = False
+        elif not np.issubdtype(values.dtype, np.floating):
+                warnings.warn(
+                    f"pint-pandas does not support magnitudes of {values.dtype}. Converting magnitudes to float.",
+                    category=RuntimeWarning,
+                )
+                values = values.astype(float)
+                copy = False
         if copy:
             values = values.copy()
         self._data = values
@@ -269,7 +296,7 @@ def __getitem__(self, item):
         item : scalar or PintArray
         """
         if is_integer(item):
-            return self._data[item] * self.units
+            return self._Q(self._data[item], self.units)
 
         item = check_array_indexer(self, item)
 
@@ -336,6 +363,12 @@ def isna(self):
         -------
         missing : np.array
         """
+        if HAS_UNCERTAINTIES:
+            # GH https://github.com/lebigot/uncertainties/issues/164
+            if isinstance(self._data, np.ndarray) and len(self._data)==0:
+                # True or False doesn't matter--we just need the value for the type
+                return np.full((0), True)
+            return unp.isnan(self._data)
         return np.isnan(self._data)
 
     def astype(self, dtype, copy=True):
@@ -483,8 +516,17 @@ def _from_sequence(cls, scalars, dtype=None, copy=False):
             dtype = PintType(master_scalar.units)
 
         def quantify_nan(item):
-            if type(item) is float:
-                return item * dtype.units
+            if HAS_UNCERTAINTIES:
+                if type(item) is UFloat:
+                    return item * dtype.units
+                if type(item) is float:
+                    if np.isnan(item):
+                        return _ufloat_nan * dtype.units
+                    else:
+                        return UFloat(item, 0) * dtype.units
+            else:
+                if type(item) is float:
+                    return item * dtype.units
             return item
 
         if isinstance(master_scalar, _Quantity):
@@ -504,6 +546,8 @@ def _from_factorized(cls, values, original):
 
     def _values_for_factorize(self):
         arr = self._data
+        if HAS_UNCERTAINTIES:
+            return arr, _ufloat_nan
         return arr, np.NaN
 
     def value_counts(self, dropna=True):
@@ -531,7 +575,10 @@ def value_counts(self, dropna=True):
         # compute counts on the data with no nans
         data = self._data
         if dropna:
-            data = data[~np.isnan(data)]
+            if HAS_UNCERTAINTIES:
+                data = data[~unp.isnan(data)]
+            else:
+                data = data[~np.isnan(data)]
 
         data_list = data.tolist()
         index = list(set(data))

pint_pandas/testsuite/test_issues.py

@@ -7,6 +7,15 @@
 from pandas.tests.extension.base.base import BaseExtensionTests
 from pint.testsuite import helpers
 
+try:
+    import uncertainties.unumpy as unp
+    from uncertainties import ufloat, UFloat
+    HAS_UNCERTAINTIES = True
+except ImportError:
+    unp = np
+    ufloat = Ufloat = None
+    HAS_UNCERTAINTIES = False
+
 from pint_pandas import PintArray, PintType
 
 ureg = PintType.ureg
@@ -15,9 +24,9 @@
 class TestIssue21(BaseExtensionTests):
     @pytest.mark.filterwarnings("ignore::RuntimeWarning")
     def test_offset_concat(self):
-        q_a = ureg.Quantity(np.arange(5), ureg.Unit("degC"))
-        q_b = ureg.Quantity(np.arange(6), ureg.Unit("degC"))
-        q_a_ = np.append(q_a, np.nan)
+        q_a = ureg.Quantity(np.arange(5)+ufloat(0,0), ureg.Unit("degC"))
+        q_b = ureg.Quantity(np.arange(6)+ufloat(0,0), ureg.Unit("degC"))
+        q_a_ = np.append(q_a, ufloat(np.nan, 0))
 
         a = pd.Series(PintArray(q_a))
         b = pd.Series(PintArray(q_b))
@@ -107,3 +116,25 @@ def test_issue_88():
     q_mm = ureg.Quantity([1000, 2000], "mm")
     b = PintArray(q_mm, "m")
     helpers.assert_quantity_almost_equal(q_m, b.quantity)
+
+def test_issue_139():
+    from pint.compat import HAS_UNCERTAINTIES
+    assert(HAS_UNCERTAINTIES)
+    from uncertainties import ufloat
+    from uncertainties import unumpy as unp
+
+    q1 = 1.234
+    q2 = 5.678
+    q_nan = np.nan
+
+    u1 = ufloat(1, 0)
+    u2 = ufloat(3, 0)
+    u_nan = ufloat(np.nan, 0.0)
+    u_plus_or_minus_nan = ufloat(0.0, np.nan)
+    u_nan_plus_or_minus_nan = ufloat(np.nan, np.nan)
+
+    a_m = PintArray([q1, u1, q2, u2, q_nan, u_nan, u_plus_or_minus_nan, u_nan_plus_or_minus_nan], ureg.m)
+    a_cm = a_m.astype('pint[cm]')
+    assert np.all(a_m[0:4] == a_cm[0:4])
+    for x, y in zip(a_m[4:], a_cm[4:]):
+        assert unp.isnan(x) == unp.isnan(y)