Merge pull request #10638 from sinhrks/array_finalize

CLN: Make ufunc works for Index

Author: jreback

doc/source/whatsnew/v0.17.0.txt

@@ -309,6 +309,8 @@ Other enhancements
 
 - ``DataFrame.apply`` will return a Series of dicts if the passed function returns a dict and ``reduce=True`` (:issue:`8735`).
 
+- ``PeriodIndex`` now supports arithmetic with ``np.ndarray`` (:issue:`10638`)
+
 - ``concat`` will now use existing Series names if provided (:issue:`10698`).
 
   .. ipython:: python
@@ -333,6 +335,7 @@ Other enhancements
 
     pd.concat([foo, bar, baz], 1)
 
+
 .. _whatsnew_0170.api:
 
 .. _whatsnew_0170.api_breaking:
@@ -1005,3 +1008,5 @@ Bug Fixes
 - Bug when constructing ``DataFrame`` where passing a dictionary with only scalar values and specifying columns did not raise an error (:issue:`10856`)
 - Bug in ``.var()`` causing roundoff errors for highly similar values (:issue:`10242`)
 - Bug in ``DataFrame.plot(subplots=True)`` with duplicated columns outputs incorrect result (:issue:`10962`)
+- Bug in ``Index`` arithmetic may result in incorrect class (:issue:`10638`)
+

pandas/core/index.py

@@ -273,7 +273,12 @@ def __array_wrap__(self, result, context=None):
         """
         Gets called after a ufunc
         """
-        return self._shallow_copy(result)
+        if is_bool_dtype(result):
+            return result
+
+        attrs = self._get_attributes_dict()
+        attrs = self._maybe_update_attributes(attrs)
+        return Index(result, **attrs)
 
     @cache_readonly
     def dtype(self):
@@ -2809,6 +2814,10 @@ def invalid_op(self, other=None):
         cls.__abs__ = _make_invalid_op('__abs__')
         cls.__inv__ = _make_invalid_op('__inv__')
 
+    def _maybe_update_attributes(self, attrs):
+        """ Update Index attributes (e.g. freq) depending on op """
+        return attrs
+
     @classmethod
     def _add_numeric_methods(cls):
         """ add in numeric methods """
@@ -2849,7 +2858,9 @@ def _evaluate_numeric_binop(self, other):
                 if reversed:
                     values, other = other, values
 
-                return self._shallow_copy(op(values, other))
+                attrs = self._get_attributes_dict()
+                attrs = self._maybe_update_attributes(attrs)
+                return Index(op(values, other), **attrs)
 
             return _evaluate_numeric_binop
 
@@ -2861,8 +2872,9 @@ def _evaluate_numeric_unary(self):
                 if not self._is_numeric_dtype:
                     raise TypeError("cannot evaluate a numeric op {opstr} for type: {typ}".format(opstr=opstr,
                                                                                                   typ=type(self)))
-
-                return self._shallow_copy(op(self.values))
+                attrs = self._get_attributes_dict()
+                attrs = self._maybe_update_attributes(attrs)
+                return Index(op(self.values), **attrs)
 
             return _evaluate_numeric_unary
 

pandas/core/ops.py

@@ -613,7 +613,8 @@ def wrapper(left, right, name=name, na_op=na_op):
         else:
             # scalars
             if hasattr(lvalues, 'values') and not isinstance(lvalues, pd.DatetimeIndex):
-                 lvalues = lvalues.values
+                lvalues = lvalues.values
+
             return left._constructor(wrap_results(na_op(lvalues, rvalues)),
                                      index=left.index, name=left.name,
                                      dtype=dtype)

pandas/tests/test_index.py

@@ -509,6 +509,56 @@ def test_equals_op(self):
             tm.assert_numpy_array_equal(index_a == item, expected3)
             tm.assert_numpy_array_equal(series_a == item, expected3)
 
+    def test_numpy_ufuncs(self):
+        # test ufuncs of numpy 1.9.2. see:
+        # http://docs.scipy.org/doc/numpy/reference/ufuncs.html
+
+        # some functions are skipped because it may return different result
+        # for unicode input depending on numpy version
+
+        for name, idx in compat.iteritems(self.indices):
+            for func in [np.exp, np.exp2, np.expm1, np.log, np.log2, np.log10,
+                         np.log1p, np.sqrt, np.sin, np.cos,
+                         np.tan, np.arcsin, np.arccos, np.arctan,
+                         np.sinh, np.cosh, np.tanh, np.arcsinh, np.arccosh,
+                         np.arctanh, np.deg2rad, np.rad2deg]:
+                if isinstance(idx, pd.tseries.base.DatetimeIndexOpsMixin):
+                    # raise TypeError or ValueError (PeriodIndex)
+                    # PeriodIndex behavior should be changed in future version
+                    with tm.assertRaises(Exception):
+                        func(idx)
+                elif isinstance(idx, (Float64Index, Int64Index)):
+                    # coerces to float (e.g. np.sin)
+                    result = func(idx)
+                    exp = Index(func(idx.values), name=idx.name)
+                    self.assert_index_equal(result, exp)
+                    self.assertIsInstance(result, pd.Float64Index)
+                else:
+                    # raise AttributeError or TypeError
+                    if len(idx) == 0:
+                        continue
+                    else:
+                        with tm.assertRaises(Exception):
+                            func(idx)
+
+            for func in [np.isfinite, np.isinf, np.isnan, np.signbit]:
+                if isinstance(idx, pd.tseries.base.DatetimeIndexOpsMixin):
+                    # raise TypeError or ValueError (PeriodIndex)
+                    with tm.assertRaises(Exception):
+                        func(idx)
+                elif isinstance(idx, (Float64Index, Int64Index)):
+                    # results in bool array
+                    result = func(idx)
+                    exp = func(idx.values)
+                    self.assertIsInstance(result, np.ndarray)
+                    tm.assertNotIsInstance(result, Index)
+                else:
+                    if len(idx) == 0:
+                        continue
+                    else:
+                        with tm.assertRaises(Exception):
+                            func(idx)
+
 
 class TestIndex(Base, tm.TestCase):
     _holder = Index
@@ -2848,6 +2898,41 @@ def test_slice_keep_name(self):
         idx = Int64Index([1, 2], name='asdf')
         self.assertEqual(idx.name, idx[1:].name)
 
+    def test_ufunc_coercions(self):
+        idx = pd.Int64Index([1, 2, 3, 4, 5], name='x')
+
+        result = np.sqrt(idx)
+        tm.assertIsInstance(result, Float64Index)
+        exp = pd.Float64Index(np.sqrt(np.array([1, 2, 3, 4, 5])), name='x')
+        tm.assert_index_equal(result, exp)
+
+        result = np.divide(idx, 2.)
+        tm.assertIsInstance(result, Float64Index)
+        exp = pd.Float64Index([0.5, 1., 1.5, 2., 2.5], name='x')
+        tm.assert_index_equal(result, exp)
+
+        # _evaluate_numeric_binop
+        result = idx + 2.
+        tm.assertIsInstance(result, Float64Index)
+        exp = pd.Float64Index([3., 4., 5., 6., 7.], name='x')
+        tm.assert_index_equal(result, exp)
+
+        result = idx - 2.
+        tm.assertIsInstance(result, Float64Index)
+        exp = pd.Float64Index([-1., 0., 1., 2., 3.], name='x')
+        tm.assert_index_equal(result, exp)
+
+        result = idx * 1.
+        tm.assertIsInstance(result, Float64Index)
+        exp = pd.Float64Index([1., 2., 3., 4., 5.], name='x')
+        tm.assert_index_equal(result, exp)
+
+        result = idx / 2.
+        tm.assertIsInstance(result, Float64Index)
+        exp = pd.Float64Index([0.5, 1., 1.5, 2., 2.5], name='x')
+        tm.assert_index_equal(result, exp)
+
+
 class DatetimeLike(Base):
 
     def test_str(self):
@@ -3101,7 +3186,9 @@ def test_get_loc(self):
                                      tolerance=timedelta(1)), 1)
         with tm.assertRaisesRegexp(ValueError, 'must be convertible'):
             idx.get_loc('2000-01-10', method='nearest', tolerance='foo')
-        with tm.assertRaisesRegexp(ValueError, 'different freq'):
+
+        msg = 'Input has different freq from PeriodIndex\\(freq=D\\)'
+        with tm.assertRaisesRegexp(ValueError, msg):
             idx.get_loc('2000-01-10', method='nearest', tolerance='1 hour')
         with tm.assertRaises(KeyError):
             idx.get_loc('2000-01-10', method='nearest', tolerance='1 day')
@@ -3119,7 +3206,8 @@ def test_get_indexer(self):
             idx.get_indexer(target, 'nearest', tolerance='1 hour'),
             [0, -1, 1])
 
-        with self.assertRaisesRegexp(ValueError, 'different freq'):
+        msg = 'Input has different freq from PeriodIndex\\(freq=H\\)'
+        with self.assertRaisesRegexp(ValueError, msg):
             idx.get_indexer(target, 'nearest', tolerance='1 minute')
 
         tm.assert_numpy_array_equal(
@@ -3215,6 +3303,44 @@ def test_numeric_compat(self):
     def test_pickle_compat_construction(self):
         pass
 
+    def test_ufunc_coercions(self):
+        # normal ops are also tested in tseries/test_timedeltas.py
+        idx = TimedeltaIndex(['2H', '4H', '6H', '8H', '10H'],
+                                freq='2H', name='x')
+
+        for result in [idx * 2, np.multiply(idx, 2)]:
+            tm.assertIsInstance(result, TimedeltaIndex)
+            exp = TimedeltaIndex(['4H', '8H', '12H', '16H', '20H'],
+                                 freq='4H', name='x')
+            tm.assert_index_equal(result, exp)
+            self.assertEqual(result.freq, '4H')
+
+        for result in [idx / 2, np.divide(idx, 2)]:
+            tm.assertIsInstance(result, TimedeltaIndex)
+            exp = TimedeltaIndex(['1H', '2H', '3H', '4H', '5H'],
+                                 freq='H', name='x')
+            tm.assert_index_equal(result, exp)
+            self.assertEqual(result.freq, 'H')
+
+        idx = TimedeltaIndex(['2H', '4H', '6H', '8H', '10H'],
+                                freq='2H', name='x')
+        for result in [ - idx, np.negative(idx)]:
+            tm.assertIsInstance(result, TimedeltaIndex)
+            exp = TimedeltaIndex(['-2H', '-4H', '-6H', '-8H', '-10H'],
+                                 freq='-2H', name='x')
+            tm.assert_index_equal(result, exp)
+            self.assertEqual(result.freq, None)
+
+        idx = TimedeltaIndex(['-2H', '-1H', '0H', '1H', '2H'],
+                                freq='H', name='x')
+        for result in [ abs(idx), np.absolute(idx)]:
+            tm.assertIsInstance(result, TimedeltaIndex)
+            exp = TimedeltaIndex(['2H', '1H', '0H', '1H', '2H'],
+                                 freq=None, name='x')
+            tm.assert_index_equal(result, exp)
+            self.assertEqual(result.freq, None)
+
+
 class TestMultiIndex(Base, tm.TestCase):
     _holder = MultiIndex
     _multiprocess_can_split_ = True

pandas/tseries/index.py

@@ -1077,15 +1077,6 @@ def _fast_union(self, other):
                               end=max(left_end, right_end),
                               freq=left.offset)
 
-    def __array_finalize__(self, obj):
-        if self.ndim == 0:  # pragma: no cover
-            return self.item()
-
-        self.offset = getattr(obj, 'offset', None)
-        self.tz = getattr(obj, 'tz', None)
-        self.name = getattr(obj, 'name', None)
-        self._reset_identity()
-
     def __iter__(self):
         """
         Return an iterator over the boxed values

pandas/tseries/period.py

@@ -19,7 +19,8 @@
 import pandas.core.common as com
 from pandas.core.common import (isnull, _INT64_DTYPE, _maybe_box,
                                 _values_from_object, ABCSeries,
-                                is_integer, is_float, is_object_dtype)
+                                is_integer, is_float, is_object_dtype,
+                                is_float_dtype)
 from pandas import compat
 from pandas.util.decorators import cache_readonly
 
@@ -307,6 +308,30 @@ def __contains__(self, key):
             return False
         return key.ordinal in self._engine
 
+    def __array_wrap__(self, result, context=None):
+        """
+        Gets called after a ufunc. Needs additional handling as
+        PeriodIndex stores internal data as int dtype
+
+        Replace this to __numpy_ufunc__ in future version
+        """
+        if isinstance(context, tuple) and len(context) > 0:
+            func = context[0]
+            if (func is np.add):
+                return self._add_delta(context[1][1])
+            elif (func is np.subtract):
+                return self._add_delta(-context[1][1])
+            elif isinstance(func, np.ufunc):
+                if 'M->M' not in func.types:
+                    msg = "ufunc '{0}' not supported for the PeriodIndex"
+                    # This should be TypeError, but TypeError cannot be raised
+                    # from here because numpy catches.
+                    raise ValueError(msg.format(func.__name__))
+
+        if com.is_bool_dtype(result):
+            return result
+        return PeriodIndex(result, freq=self.freq, name=self.name)
+
     @property
     def _box_func(self):
         return lambda x: Period._from_ordinal(ordinal=x, freq=self.freq)
@@ -522,7 +547,18 @@ def _maybe_convert_timedelta(self, other):
             base = frequencies.get_base_alias(freqstr)
             if base == self.freq.rule_code:
                 return other.n
-        raise ValueError("Input has different freq from PeriodIndex(freq={0})".format(self.freq))
+        elif isinstance(other, np.ndarray):
+            if com.is_integer_dtype(other):
+                return other
+            elif com.is_timedelta64_dtype(other):
+                offset = frequencies.to_offset(self.freq)
+                if isinstance(offset, offsets.Tick):
+                    nanos = tslib._delta_to_nanoseconds(other)
+                    offset_nanos = tslib._delta_to_nanoseconds(offset)
+                    if (nanos % offset_nanos).all() == 0:
+                        return nanos // offset_nanos
+        msg = "Input has different freq from PeriodIndex(freq={0})"
+        raise ValueError(msg.format(self.freqstr))
 
     def _add_delta(self, other):
         ordinal_delta = self._maybe_convert_timedelta(other)
@@ -775,14 +811,6 @@ def _format_native_types(self, na_rep=u('NaT'), date_format=None, **kwargs):
         values[imask] = np.array([formatter(dt) for dt in values[imask]])
         return values
 
-    def __array_finalize__(self, obj):
-        if not self.ndim:  # pragma: no cover
-            return self.item()
-
-        self.freq = getattr(obj, 'freq', None)
-        self.name = getattr(obj, 'name', None)
-        self._reset_identity()
-
     def take(self, indices, axis=0):
         """
         Analogous to ndarray.take

pandas/tseries/tdi.py

@@ -278,6 +278,14 @@ def __setstate__(self, state):
             raise Exception("invalid pickle state")
     _unpickle_compat = __setstate__
 
+    def _maybe_update_attributes(self, attrs):
+        """ Update Index attributes (e.g. freq) depending on op """
+        freq = attrs.get('freq', None)
+        if freq is not None:
+            # no need to infer if freq is None
+            attrs['freq'] = 'infer'
+        return attrs
+
     def _add_delta(self, delta):
         if isinstance(delta, (Tick, timedelta, np.timedelta64)):
             new_values = self._add_delta_td(delta)
@@ -560,14 +568,6 @@ def _fast_union(self, other):
         else:
             return left
 
-    def __array_finalize__(self, obj):
-        if self.ndim == 0:  # pragma: no cover
-            return self.item()
-
-        self.name = getattr(obj, 'name', None)
-        self.freq = getattr(obj, 'freq', None)
-        self._reset_identity()
-
     def _wrap_union_result(self, other, result):
         name = self.name if self.name == other.name else None
         return self._simple_new(result, name=name, freq=None)