diff --git a/madarrays/mad_array.py b/madarrays/mad_array.py index 208eb27465a94057e316684404182b016d032223..7764f43f45781ec964cd934bd10e3568798356dd 100644 --- a/madarrays/mad_array.py +++ b/madarrays/mad_array.py @@ -98,8 +98,15 @@ def _complex_masking_only(f): return decorated -UFUNC_RETURNING_MADARRAYS = ['add', 'subtract', 'multiply', 'true_divide', - 'floor_divide', 'floor', 'ceil', 'absolute'] +UFUNC_NOT_RETURNING_MADARRAYS = ['bitwise_and', 'bitwise_or', 'bitwise_xor', + 'invert', 'left_shift', 'right_shift', + 'greater', 'greater_equal', 'less', + 'less_equal', 'not_equal', 'equal', + 'logical_and', 'logical_or', 'logical_xor', + 'logical_not', 'maximum', 'minimum', 'fmax', + 'fmin', 'isfinite', 'isinf', 'isnan', 'isnat', + 'signbit', 'copysign', 'nextafter', 'spacing', + 'modf', 'frexp', 'fmod'] class MadArray(np.ndarray): @@ -341,7 +348,7 @@ class MadArray(np.ndarray): for result, output in zip(results, outputs): if output is None: if (method == '__call__' and - ufunc.__name__ in UFUNC_RETURNING_MADARRAYS): + ufunc.__name__ not in UFUNC_NOT_RETURNING_MADARRAYS): new_results.append(np.asarray(result).view(MadArray)) new_results[-1]._mask = mask new_results[-1]._complex_masking = complex_masking diff --git a/madarrays/tests/test_madarray.py b/madarrays/tests/test_madarray.py index ff08855a878676221e9d0f8d7aec7b8b9f82e9f0..29ad6412490b69ea05e1d30f9054e6e502ef7356 100644 --- a/madarrays/tests/test_madarray.py +++ b/madarrays/tests/test_madarray.py @@ -523,6 +523,39 @@ class TestMadArray: assert not isinstance(m, MadArray) np.testing.assert_equal(m, np.std(x)) + m = np.abs(ma) + assert isinstance(m, MadArray) + np.testing.assert_equal(m, np.abs(x)) + + m = np.sqrt(ma) + assert isinstance(m, MadArray) + np.testing.assert_equal(m, np.sqrt(x)) + + m = ma**2 + assert isinstance(m, MadArray) + np.testing.assert_equal(m, x**2) + + m = np.conj(ma) + assert isinstance(m, MadArray) + np.testing.assert_equal(m, np.conj(x)) + + if np.issubdtype(ma.dtype, np.floating): + m = np.floor(ma) + assert isinstance(m, MadArray) + np.testing.assert_equal(m, np.floor(x)) + + m = np.ceil(ma) + assert isinstance(m, MadArray) + np.testing.assert_equal(m, np.ceil(x)) + + a = ma < np.mean(ma) + assert not isinstance(a, MadArray) + np.testing.assert_equal(a, x < np.mean(ma)) + + a = ma >= np.mean(ma) + assert not isinstance(a, MadArray) + np.testing.assert_equal(a, x >= np.mean(ma)) + def test_eq_ne_numpy(self): ma = MadArray(self.x_float)