pyerrors/correlators.py: Allowing for None values in pe.Corr.prune

pyerrors/correlators.py

@@ -1405,13 +1405,15 @@ def prune(self, Ntrunc, tproj=3, t0proj=2, basematrix=None):
         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):

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])