Merge pull request #554 from shoyer/real-and-imag

Fixes for complex numbers

Author: shoyer

doc/api-hidden.rst

@@ -40,7 +40,9 @@
    Dataset.clip
    Dataset.conj
    Dataset.conjugate
+   Dataset.imag
    Dataset.round
+   Dataset.real
    Dataset.T
 
    DataArray.ndim
@@ -80,8 +82,10 @@
    DataArray.clip
    DataArray.conj
    DataArray.conjugate
+   DataArray.imag
    DataArray.searchsorted
    DataArray.round
+   DataArray.real
    DataArray.T
 
    ufuncs.angle

doc/api.rst

@@ -139,7 +139,9 @@ Computation
 :py:attr:`~Dataset.clip`
 :py:attr:`~Dataset.conj`
 :py:attr:`~Dataset.conjugate`
+:py:attr:`~Dataset.imag`
 :py:attr:`~Dataset.round`
+:py:attr:`~Dataset.real`
 :py:attr:`~Dataset.T`
 
 **Grouped operations**:
@@ -253,8 +255,10 @@ Computation
 :py:attr:`~DataArray.clip`
 :py:attr:`~DataArray.conj`
 :py:attr:`~DataArray.conjugate`
+:py:attr:`~DataArray.imag`
 :py:attr:`~DataArray.searchsorted`
 :py:attr:`~DataArray.round`
+:py:attr:`~DataArray.real`
 :py:attr:`~DataArray.T`
 
 **Grouped operations**:

doc/whats-new.rst

@@ -19,7 +19,19 @@ API Changes
   :py:meth:`~xray.DataArray.plot` was changed to provide more compatibility
   with matplotlib's `contour` and `contourf` functions (:issue:`538`).
   Now discrete lists of colors should be specified using `colors` keyword,
-  rather than `cmap`. 
+  rather than `cmap`.
+
+Enhancements
+~~~~~~~~~~~~
+
+- Add :py:attr:`~xray.Dataset.real` and :py:attr:`~xray.Dataset.imag`
+  attributes to Dataset and DataArray (:issue:`553`).
+
+Bug fixes
+~~~~~~~~~
+
+- Aggregation functions now correctly skip ``NaN`` for data for ``complex128``
+  dtype (:issue:`554`).
 
 v0.6.0 (21 August 2015)
 -----------------------

xray/core/dataarray.py

@@ -1171,5 +1171,13 @@ def diff(self, dim, n=1, label='upper'):
         ds = self._dataset.diff(n=n, dim=dim, label=label)
         return self._with_replaced_dataset(ds)
 
+    @property
+    def real(self):
+        return self._with_replaced_dataset(self._dataset.real)
+
+    @property
+    def imag(self):
+        return self._with_replaced_dataset(self._dataset.imag)
+
 # priority most be higher than Variable to properly work with binary ufuncs
 ops.inject_all_ops_and_reduce_methods(DataArray, priority=60)

xray/core/dataset.py

@@ -1855,12 +1855,14 @@ def from_dataframe(cls, dataframe):
         return obj
 
     @staticmethod
-    def _unary_op(f):
+    def _unary_op(f, keep_attrs=False):
         @functools.wraps(f)
         def func(self, *args, **kwargs):
             ds = self.coords.to_dataset()
             for k in self.data_vars:
                 ds._variables[k] = f(self._variables[k], *args, **kwargs)
+            if keep_attrs:
+                ds._attrs = self._attrs
             return ds
         return func
 
@@ -2019,5 +2021,13 @@ def diff(self, dim, n=1, label='upper'):
         else:
             return difference
 
+    @property
+    def real(self):
+        return self._unary_op(lambda x: x.real, keep_attrs=True)(self)
+
+    @property
+    def imag(self):
+        return self._unary_op(lambda x: x.imag, keep_attrs=True)(self)
+
 
 ops.inject_all_ops_and_reduce_methods(Dataset, array_only=False)

xray/core/ops.py

@@ -292,8 +292,8 @@ def f(values, axis=None, skipna=None, **kwargs):
         if coerce_strings and values.dtype.kind in 'SU':
             values = values.astype(object)
 
-        if skipna or (skipna is None and values.dtype.kind == 'f'):
-            if values.dtype.kind not in ['i', 'f']:
+        if skipna or (skipna is None and values.dtype.kind in 'cf'):
+            if values.dtype.kind not in ['i', 'f', 'c']:
                 raise NotImplementedError(
                     'skipna=True not yet implemented for %s with dtype %s'
                     % (name, values.dtype))

xray/core/variable.py

@@ -751,6 +751,14 @@ def identical(self, other):
         except (TypeError, AttributeError):
             return False
 
+    @property
+    def real(self):
+        return type(self)(self.dims, self.data.real, self._attrs)
+
+    @property
+    def imag(self):
+        return type(self)(self.dims, self.data.imag, self._attrs)
+
     def __array_wrap__(self, obj, context=None):
         return Variable(self.dims, obj)
 

xray/test/test_dataarray.py

@@ -1460,3 +1460,8 @@ def test_dataarray_diff_n1(self):
                              [da['x'].values, da['y'].values[1:]],
                              ['x', 'y'])
         self.assertDataArrayEqual(expected, actual)
+
+    def test_real_and_imag(self):
+        array = DataArray(1 + 2j)
+        self.assertDataArrayIdentical(array.real, DataArray(1))
+        self.assertDataArrayIdentical(array.imag, DataArray(2))

xray/test/test_dataset.py

@@ -2115,3 +2115,13 @@ def test_dataset_diff_exception_label_str(self):
         ds = create_test_data(seed=1)
         with self.assertRaisesRegexp(ValueError, '\'label\' argument has to'):
             ds.diff('dim2', label='raise_me')
+
+    def test_real_and_imag(self):
+        attrs = {'foo': 'bar'}
+        ds = Dataset({'x': ((), 1 + 2j, attrs)}, attrs=attrs)
+
+        expected_re = Dataset({'x': ((), 1, attrs)}, attrs=attrs)
+        self.assertDatasetIdentical(ds.real, expected_re)
+
+        expected_im = Dataset({'x': ((), 2, attrs)}, attrs=attrs)
+        self.assertDatasetIdentical(ds.imag, expected_im)

xray/test/test_variable.py

@@ -371,6 +371,23 @@ def test_copy(self):
                                   source_ndarray(w.values))
         self.assertVariableIdentical(v, copy(v))
 
+    def test_real_and_imag(self):
+        v = self.cls('x', np.arange(3) - 1j * np.arange(3), {'foo': 'bar'})
+        expected_re = self.cls('x', np.arange(3), {'foo': 'bar'})
+        self.assertVariableIdentical(v.real, expected_re)
+
+        expected_im = self.cls('x', -np.arange(3), {'foo': 'bar'})
+        self.assertVariableIdentical(v.imag, expected_im)
+
+        expected_abs = self.cls('x', np.sqrt(2 * np.arange(3) ** 2))
+        self.assertVariableAllClose(abs(v), expected_abs)
+
+    def test_aggregate_complex(self):
+        # should skip NaNs
+        v = self.cls('x', [1, 2j, np.nan])
+        expected = Variable((), 0.5 + 1j)
+        self.assertVariableAllClose(v.mean(), expected)
+
 
 class TestVariable(TestCase, VariableSubclassTestCases):
     cls = staticmethod(Variable)