diff --git a/pyerrors/correlators.py b/pyerrors/correlators.py index 0375155f..fb14d6d1 100644 --- a/pyerrors/correlators.py +++ b/pyerrors/correlators.py @@ -1405,13 +1405,15 @@ class Corr: tmpmat = np.empty((Ntrunc, Ntrunc), dtype=object) rmat = [] for t in range(basematrix.T): - for i in range(Ntrunc): - for j in range(Ntrunc): - tmpmat[i][j] = evecs[i].T @ self[t] @ evecs[j] - rmat.append(np.copy(tmpmat)) + if self.content[t] is None: + rmat.append(None) + else: + for i in range(Ntrunc): + for j in range(Ntrunc): + tmpmat[i][j] = evecs[i].T @ self[t] @ evecs[j] + rmat.append(np.copy(tmpmat)) - newcontent = [None if (self.content[t] is None) else rmat[t] for t in range(self.T)] - return Corr(newcontent) + return Corr(rmat) def _sort_vectors(vec_set_in, ts): diff --git a/tests/correlators_test.py b/tests/correlators_test.py index fc3528d2..cd629275 100644 --- a/tests/correlators_test.py +++ b/tests/correlators_test.py @@ -781,3 +781,26 @@ def test_complex_add_and_mul(): cc += 2j cc = cc * 4j cc.real + cc.imag + + +def test_prune_with_Nones(): + N = 3 + T = 10 + + front_padding = 1 + back_padding = T // 2 + + Ntrunc = N - 1 + t0proj = 2 + tproj = 3 + + corr_content = np.array([[[pe.pseudo_Obs((i+j+1)**(-t), .01, "None_prune_test") for i in range(N)] for j in range(N)] for t in range(T // 2 - front_padding)]) + unpadded_corr = pe.Corr(corr_content) + padded_corr = pe.Corr(corr_content, padding=[front_padding, back_padding]) + + tmp_corr = unpadded_corr.prune(Ntrunc, t0proj=t0proj-front_padding, tproj=tproj-front_padding) + pruned_then_padded = pe.Corr(tmp_corr.content, padding=[front_padding, back_padding]) + padded_then_pruned = padded_corr.prune(Ntrunc, t0proj=t0proj, tproj=tproj) + + for t in range(T): + assert np.all(pruned_then_padded.content[t] == padded_then_pruned.content[t])