Merge pull request #135 from fjosw/fix/complex_Corr

Fix/complex corr

Author: fjosw

pyerrors/correlators.py

@@ -969,6 +969,8 @@ def __repr__(self, print_range=None):
             content_string += "Description: " + self.tag + "\n"
         if self.N != 1:
             return content_string
+        if isinstance(self[0], CObs):
+            return content_string
 
         if print_range[1]:
             print_range[1] += 1
@@ -1136,8 +1138,10 @@ def _apply_func_to_corr(self, func):
         for t in range(self.T):
             if _check_for_none(self, newcontent[t]):
                 continue
-            if np.isnan(np.sum(newcontent[t]).value):
-                newcontent[t] = None
+            tmp_sum = np.sum(newcontent[t])
+            if hasattr(tmp_sum, "value"):
+                if np.isnan(tmp_sum.value):
+                    newcontent[t] = None
         if all([item is None for item in newcontent]):
             raise Exception('Operation returns undefined correlator')
         return Corr(newcontent)
@@ -1194,8 +1198,8 @@ def __rtruediv__(self, y):
     @property
     def real(self):
         def return_real(obs_OR_cobs):
-            if isinstance(obs_OR_cobs, CObs):
-                return obs_OR_cobs.real
+            if isinstance(obs_OR_cobs.flatten()[0], CObs):
+                return np.vectorize(lambda x: x.real)(obs_OR_cobs)
             else:
                 return obs_OR_cobs
 
@@ -1204,8 +1208,8 @@ def return_real(obs_OR_cobs):
     @property
     def imag(self):
         def return_imag(obs_OR_cobs):
-            if isinstance(obs_OR_cobs, CObs):
-                return obs_OR_cobs.imag
+            if isinstance(obs_OR_cobs.flatten()[0], CObs):
+                return np.vectorize(lambda x: x.imag)(obs_OR_cobs)
             else:
                 return obs_OR_cobs * 0  # So it stays the right type
 

tests/correlators_test.py

@@ -532,3 +532,13 @@ def test_prune():
     with pytest.raises(Exception):
         corr_mat.prune(3)
         corr_mat.prune(4)
+
+
+def test_complex_Corr():
+    o1 = pe.pseudo_Obs(1.0, 0.1, "test")
+    cobs = pe.CObs(o1, -o1)
+    ccorr = pe.Corr([cobs, cobs, cobs])
+    assert np.all([ccorr.imag[i] == -ccorr.real[i] for i in range(ccorr.T)])
+    print(ccorr)
+    mcorr = pe.Corr(np.array([[ccorr, ccorr], [ccorr, ccorr]]))
+    assert np.all([mcorr.imag[i] == -mcorr.real[i] for i in range(mcorr.T)])