[med-svn] [python-mne] 26/52: ENH : adding numeric interface to SourceEstimate class

[Yaroslav Halchenko]

This is an automated email from the git hooks/post-receive script.

yoh pushed a commit to annotated tag v0.2
in repository python-mne.

commit f84a0540c65383273ebe58f9798503c2eadbfea9
Author: Alexandre Gramfort <alexandre.gramfort at inria.fr>
Date:   Mon Oct 3 17:59:20 2011 -0400

    ENH : adding numeric interface to SourceEstimate class
---
 mne/source_estimate.py            | 79 +++++++++++++++++++++++++++++++++++++++
 mne/tests/test_source_estimate.py | 27 +++++++++++++
 2 files changed, 106 insertions(+)

diff --git a/mne/source_estimate.py b/mne/source_estimate.py
index 7dc785b..4ec62e2 100644
--- a/mne/source_estimate.py
+++ b/mne/source_estimate.py
@@ -166,6 +166,85 @@ class SourceEstimate(object):
         self.data = self.data[:, mask]
         self.tmin = self.times[0]
 
+    def __add__(self, a):
+        stc = copy.deepcopy(self)
+        stc += a
+        return stc
+
+    def __iadd__(self, a):
+        if isinstance(a, SourceEstimate):
+            self.data += a.data
+        else:
+            self.data += a
+        return self
+
+    def __sub__(self, a):
+        stc = copy.deepcopy(self)
+        stc -= a
+        return stc
+
+    def __isub__(self, a):
+        if isinstance(a, SourceEstimate):
+            self.data -= a.data
+        else:
+            self.data -= a
+        return self
+
+    def __div__(self, a):
+        stc = copy.deepcopy(self)
+        stc /= a
+        return stc
+
+    def __idiv__(self, a):
+        if isinstance(a, SourceEstimate):
+            self.data /= a.data
+        else:
+            self.data /= a
+        return self
+
+    def __mul__(self, a):
+        stc = copy.deepcopy(self)
+        stc *= a
+        return stc
+
+    def __imul__(self, a):
+        if isinstance(a, SourceEstimate):
+            self.data *= a.data
+        else:
+            self.data *= a
+        return self
+
+    def __pow__(self, a):
+        stc = copy.deepcopy(self)
+        stc **= a
+        return stc
+
+    def __ipow__(self, a):
+        self.data **= a
+        return self
+
+    def __radd__(self, a):
+        return self + a
+
+    def __rsub__(self, a):
+        return self - a
+
+    def __rmul__(self, a):
+        return self * a
+
+    def __rdiv__(self, a):
+        return self / a
+
+    def __neg__(self):
+        stc = copy.deepcopy(self)
+        stc.data *= -1
+        return stc
+
+    def __pos__(self):
+        return self
+
+    def sqrt(self):
+        return self ** (0.5)
 
 ###############################################################################
 # Morphing
diff --git a/mne/tests/test_source_estimate.py b/mne/tests/test_source_estimate.py
index e24eed0..d10c433 100644
--- a/mne/tests/test_source_estimate.py
+++ b/mne/tests/test_source_estimate.py
@@ -31,6 +31,33 @@ def test_io_stc():
     assert_array_almost_equal(stc['tstep'], stc2['tstep'])
 
 
+def test_stc_arithmetic():
+    """Test arithmetic for STC files
+    """
+    fname = op.join(data_path, 'MEG', 'sample', 'sample_audvis-meg')
+    stc = SourceEstimate(fname)
+    data = stc.data.copy()
+
+    out = list()
+    for a in [data, stc]:
+        a = a + a * 3 + 3 * a - a ** 2 / 2
+
+        a += a
+        a -= a
+        a /=  2 * a
+        a *= -a
+
+        a += 2
+        a -= 1
+        a *= -1
+        a /= 2
+        a **= 3
+        out.append(a)
+
+    assert_array_equal(out[0], out[1].data)
+    assert_array_equal(stc.sqrt().data, np.sqrt(stc.data))
+
+
 def test_morph_data():
     """Test morphing of data
     """

-- 
Alioth's /usr/local/bin/git-commit-notice on /srv/git.debian.org/git/debian-med/python-mne.git