fix: Corr.roll method now also works for correlator matrices by

explicitly specifying the axis.

Co-authored-by: Matteo Di Carlo <matteo.dicarlo93@gmail.com>
This commit is contained in:
Fabian Joswig 2023-07-13 18:40:47 +01:00
parent bbbacc7bb4
commit 4467a2777f
No known key found for this signature in database
2 changed files with 23 additions and 1 deletions

View file

@ -417,7 +417,7 @@ class Corr:
dt : int
number of timeslices
"""
return Corr(list(np.roll(np.array(self.content, dtype=object), dt)))
return Corr(list(np.roll(np.array(self.content, dtype=object), dt, axis=0)))
def reverse(self):
"""Reverse the time ordering of the Corr"""

View file

@ -666,3 +666,25 @@ def test_matrix_trace():
corr = pe.Corr([mat] * 4)
for el in corr.trace():
assert el == 0
def test_corr_roll():
T = 4
rn = lambda : np.random.normal(0.5, 0.1)
ll = []
for i in range(T):
re = pe.pseudo_Obs(rn(), rn(), "test")
im = pe.pseudo_Obs(rn(), rn(), "test")
ll.append(pe.CObs(re, im))
# Rolling by T should produce the same correlator
corr = pe.Corr(ll)
tt = corr - corr.roll(T)
for el in tt:
assert np.all(el == 0)
mcorr = pe.Corr(np.array([[corr, corr + 0.1], [corr - 0.1, 2 * corr]]))
tt = mcorr.roll(T) - mcorr
for el in tt:
assert np.all(el == 0)