From dbd5c09a833043f5a406e1dd234e965f2341877c Mon Sep 17 00:00:00 2001
From: Marc Paterno <paterno@fnal.gov>
Date: Fri, 21 Jun 2024 15:02:49 -0500
Subject: [PATCH] Create test and fix for issue 427

---
 firecrown/likelihood/two_point.py             |  2 +-
 tests/conftest.py                             | 22 +++++++++-
 .../gauss_family/statistic/test_two_point.py  | 42 ++++++++++++++++++-
 3 files changed, 63 insertions(+), 3 deletions(-)

diff --git a/firecrown/likelihood/two_point.py b/firecrown/likelihood/two_point.py
index c7266115..042785cc 100644
--- a/firecrown/likelihood/two_point.py
+++ b/firecrown/likelihood/two_point.py
@@ -620,7 +620,6 @@ def compute_theory_vector_real_space(self, tools: ModelingTools) -> TheoryVector
         assert self.thetas is not None
         assert self.ells_for_xi is not None
 
-        self.cells = {}
         print(self.source0.parameter_prefix, self.source1.parameter_prefix)
         cells_for_xi = self.compute_cells(
             self.ells_for_xi, scale0, scale1, tools, tracers0, tracers1
@@ -724,6 +723,7 @@ def compute_cells(
         tracers1: Sequence[Tracer],
     ) -> npt.NDArray[np.float64]:
         """Compute the power spectrum for the given ells and tracers."""
+        self.cells = {}
         for tracer0 in tracers0:
             for tracer1 in tracers1:
                 pk_name = f"{tracer0.field}:{tracer1.field}"
diff --git a/tests/conftest.py b/tests/conftest.py
index d74b2ea1..6c637bb6 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -513,7 +513,7 @@ def fixture_sacc_galaxy_xis_lens0_lens0_no_data():
 
 @pytest.fixture(name="sacc_galaxy_cells_src0_src0_window")
 def fixture_sacc_galaxy_cells_src0_src0_window():
-    """Fixture for a SACC data without window functions."""
+    """Fixture for a SACC data with a window function."""
     sacc_data = sacc.Sacc()
 
     z = np.linspace(0, 1.0, 256) + 0.05
@@ -534,3 +534,23 @@ def fixture_sacc_galaxy_cells_src0_src0_window():
     sacc_data.add_covariance(cov)
 
     return sacc_data, z, dndz
+
+
+@pytest.fixture(name="sacc_galaxy_cells_src0_src0_no_window")
+def fixture_sacc_galaxy_cells_src0_src0_no_window():
+    """Fixture for a SACC data without a window function."""
+    sacc_data = sacc.Sacc()
+
+    z = np.linspace(0, 1.0, 256) + 0.05
+    ells = np.unique(np.logspace(1, 3, 10).astype(np.int64))
+
+    dndz = np.exp(-0.5 * (z - 0.5) ** 2 / 0.05 / 0.05)
+    sacc_data.add_tracer("NZ", "src0", z, dndz)
+
+    Cells = np.random.normal(size=ells.shape[0])
+    sacc_data.add_ell_cl("galaxy_shear_cl_ee", "src0", "src0", ells, Cells)
+
+    cov = np.diag(np.ones_like(Cells) * 0.01)
+    sacc_data.add_covariance(cov)
+
+    return sacc_data, z, dndz
diff --git a/tests/likelihood/gauss_family/statistic/test_two_point.py b/tests/likelihood/gauss_family/statistic/test_two_point.py
index 6931dca1..f530e1f0 100644
--- a/tests/likelihood/gauss_family/statistic/test_two_point.py
+++ b/tests/likelihood/gauss_family/statistic/test_two_point.py
@@ -78,6 +78,8 @@ def test_tracer_names():
 
 
 def test_two_point_src0_src0_window(sacc_galaxy_cells_src0_src0_window):
+    """This test also makes sure that TwoPoint theory calculations are
+    repeatable."""
     sacc_data, _, _ = sacc_galaxy_cells_src0_src0_window
 
     src0 = WeakLensing(sacc_tracer="src0")
@@ -93,8 +95,46 @@ def test_two_point_src0_src0_window(sacc_galaxy_cells_src0_src0_window):
     assert statistic.window.ells_for_interpolation is not None
     assert all(np.isfinite(statistic.window.ells_for_interpolation))
 
+    statistic.reset()
     statistic.update(ParamsMap())
-    statistic.compute_theory_vector(tools)
+    tools.update(ParamsMap())
+    result1 = statistic.compute_theory_vector(tools)
+    assert all(np.isfinite(result1))
+
+    statistic.reset()
+    statistic.update(ParamsMap())
+    tools.update(ParamsMap())
+    result2 = statistic.compute_theory_vector(tools)
+    assert np.array_equal(result1, result2)
+
+
+def test_two_point_src0_src0_no_window(sacc_galaxy_cells_src0_src0_no_window):
+    """This test also makes sure that TwoPoint theory calculations are
+    repeatable."""
+    sacc_data, _, _ = sacc_galaxy_cells_src0_src0_no_window
+
+    src0 = WeakLensing(sacc_tracer="src0")
+
+    statistic = TwoPoint("galaxy_shear_cl_ee", src0, src0)
+    statistic.read(sacc_data)
+
+    tools = ModelingTools()
+    tools.update(ParamsMap())
+    tools.prepare(pyccl.CosmologyVanillaLCDM())
+
+    assert statistic.window is None
+
+    statistic.reset()
+    statistic.update(ParamsMap())
+    tools.update(ParamsMap())
+    result1 = statistic.compute_theory_vector(tools)
+    assert all(np.isfinite(result1))
+
+    statistic.reset()
+    statistic.update(ParamsMap())
+    tools.update(ParamsMap())
+    result2 = statistic.compute_theory_vector(tools)
+    assert np.array_equal(result1, result2)
 
 
 def test_two_point_src0_src0_no_data_lin(sacc_galaxy_cells_src0_src0_no_data):