diff --git a/pyerrors/correlators.py b/pyerrors/correlators.py index 453beb8b..736bcb7a 100644 --- a/pyerrors/correlators.py +++ b/pyerrors/correlators.py @@ -969,6 +969,8 @@ class Corr: 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 @@ class Corr: 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 @@ class Corr: @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 @@ class Corr: @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 diff --git a/tests/correlators_test.py b/tests/correlators_test.py index 49332254..1ab31f12 100644 --- a/tests/correlators_test.py +++ b/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)])