Commit 9d3c57af authored by Ronan Hamon's avatar Ronan Hamon

Update ufunc in MadArray and Waveform

parent f4cba663
......@@ -98,6 +98,10 @@ def _complex_masking_only(f):
return decorated
UFUNC_RETURNING_MADARRAYS = ['add', 'subtract', 'multiply', 'true_divide',
'floor_divide', 'floor', 'ceil', 'absolute']
class MadArray(np.ndarray):
"""Subclass of numpy.ndarray to handle data with missing elements.
......@@ -351,7 +355,7 @@ class MadArray(np.ndarray):
new_results = []
for result, output in zip(results, outputs):
if output is None:
if not ufunc.__name__.startswith('is'):
if ufunc.__name__ in UFUNC_RETURNING_MADARRAYS:
new_results.append(np.asarray(result).view(MadArray))
new_results[-1]._mask = mask
new_results[-1]._complex_masking = complex_masking
......
......@@ -772,21 +772,9 @@ class Waveform(MadArray):
def copy(self):
return Waveform(self)
def __add__(self, other):
if isinstance(other, Waveform) and self.fs != other.fs:
errmsg = 'Waveforms do not have the same fs: {} and {}'
raise ValueError(errmsg.format(self.fs, other.fs))
return Waveform(super().__add__(other), fs=self.fs)
def __sub__(self, ma):
return self.__add__(-ma)
def __eq__(self, other):
if isinstance(other, Waveform) and self.fs != other.fs:
errmsg = 'Waveforms do not have the same fs: {} and {}'
raise ValueError(errmsg.format(self.fs, other.fs))
# Test the compatability between Waveform
_ = self + other
return super().__eq__(other)
def is_equal(self, other):
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment