diff --git a/pyerrors/correlators.py b/pyerrors/correlators.py index c39c7095..736bcb7a 100644 --- a/pyerrors/correlators.py +++ b/pyerrors/correlators.py @@ -1198,8 +1198,8 @@ class Corr: @property def real(self): def return_real(obs_OR_cobs): - if isinstance(obs_OR_cobs[0], CObs): - return obs_OR_cobs[0].real + if isinstance(obs_OR_cobs.flatten()[0], CObs): + return np.vectorize(lambda x: x.real)(obs_OR_cobs) else: return obs_OR_cobs @@ -1208,8 +1208,8 @@ class Corr: @property def imag(self): def return_imag(obs_OR_cobs): - if isinstance(obs_OR_cobs[0], CObs): - return obs_OR_cobs[0].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 49c950be..1ab31f12 100644 --- a/tests/correlators_test.py +++ b/tests/correlators_test.py @@ -536,7 +536,9 @@ def test_prune(): def test_complex_Corr(): o1 = pe.pseudo_Obs(1.0, 0.1, "test") - cobs = pe.CObs(o1, o1) + 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)]) + 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)])