diff --git a/tests/correlators_test.py b/tests/correlators_test.py index fc3528d2..cb8517fa 100644 --- a/tests/correlators_test.py +++ b/tests/correlators_test.py @@ -781,3 +781,27 @@ 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]) + with pytest.raises(ValueError): + 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])