From a9eb000677d7faf236a040598d9d78dfa249f02c Mon Sep 17 00:00:00 2001
From: Toby Dixon <toby.dixon.23@ucl.ac.uk>
Date: Mon, 6 May 2024 23:50:11 +0200
Subject: [PATCH] update to tests

---
 tests/evt/test_xtalk.py | 75 +++++++++++++++++++++++++++++++++++++++--
 1 file changed, 73 insertions(+), 2 deletions(-)

diff --git a/tests/evt/test_xtalk.py b/tests/evt/test_xtalk.py
index d8e6481f7..01cce4c98 100644
--- a/tests/evt/test_xtalk.py
+++ b/tests/evt/test_xtalk.py
@@ -1,13 +1,33 @@
+from pathlib import Path
+
 import numpy as np
+import pytest
+from lgdo import lh5
 
+from pygama.evt import utils
 from pygama.evt.modules import xtalk
 
+config_dir = Path(__file__).parent / "configs"
+
+
+@pytest.fixture(scope="module")
+def files_config(lgnd_test_data, tmptestdir):
+    tcm_path = "lh5/prod-ref-l200/generated/tier/tcm/phy/p03/r001/l200-p03-r001-phy-20230322T160139Z-tier_tcm.lh5"
+    outfile = f"{tmptestdir}/l200-p03-r001-phy-20230322T160139Z-tier_evt.lh5"
+
+    return {
+        "tcm": (lgnd_test_data.get_path(tcm_path), "hardware_tcm_1"),
+        "dsp": (lgnd_test_data.get_path(tcm_path.replace("tcm", "dsp")), "dsp", "ch{}"),
+        "hit": (lgnd_test_data.get_path(tcm_path.replace("tcm", "hit")), "hit", "ch{}"),
+        "evt": (outfile, "evt"),
+    }
+
 
-def test_xtalk_corrected_energy():
+def test_xtalk_corrected_energy(lgnd_test_data, files_config):
 
     energy = np.array([[1, 2, 3], [4, 5, 6], [2, 0, 1], [0, 1, 0]])
     matrix = np.array([[0, 0, 1], [1, 0, 2], [0, 2, 0]])
-    energy_corrected_zero_threshold = xtalk.xtalk_corrected_energy(
+    energy_corrected_zero_threshold = xtalk.xtalk_correct_energy_impl(
         energy, energy, matrix, None
     )
 
@@ -15,3 +35,54 @@ def test_xtalk_corrected_energy():
         energy_corrected_zero_threshold
         == (energy - np.array([[3, 7, 4], [6, 16, 10], [1, 4, 0], [0, 0, 2]]))
     )
+
+    # test with a 2.1 threshold
+    energy_corrected_two_threshold = xtalk.xtalk_correct_energy_impl(
+        energy, energy, matrix, 2.1
+    )
+    assert np.all(
+        energy_corrected_two_threshold
+        == (energy - np.array([[3, 6, 0], [6, 16, 10], [0, 0, 0], [0, 0, 0]]))
+    )
+
+
+def test_gather_energy(lgnd_test_data, files_config):
+    f = utils.make_files_config(files_config)
+    tcm = utils.TCMData(
+        id=lh5.read_as(f"/{f.tcm.group}/array_id", f.tcm.file, library="np"),
+        idx=lh5.read_as(f"/{f.tcm.group}/array_idx", f.tcm.file, library="np"),
+        cumulative_length=lh5.read_as(
+            f"/{f.tcm.group}/cumulative_length", f.tcm.file, library="np"
+        ),
+    )
+    energy = xtalk.gather_energy(
+        "hit.cuspEmax_ctc_cal", tcm, f, np.array([1084803, 1084804])
+    )
+    n_rows = np.max(tcm.idx)+1
+    assert isinstance(energy, np.ndarray)
+    assert energy.ndim == 2
+    assert np.shape(energy) == (n_rows,2)
+
+
+def test_filter_hits(lgnd_test_data, files_config):
+    f = utils.make_files_config(files_config)
+    tcm = utils.TCMData(
+        id=lh5.read_as(f"/{f.tcm.group}/array_id", f.tcm.file, library="np"),
+        idx=lh5.read_as(f"/{f.tcm.group}/array_idx", f.tcm.file, library="np"),
+        cumulative_length=lh5.read_as(
+            f"/{f.tcm.group}/cumulative_length", f.tcm.file, library="np"
+        ),
+    )
+    n_rows = np.max(tcm.idx) + 1
+    
+    filter = xtalk.filter_hits(
+        f,
+        tcm,
+        "hit.cuspEmax_ctc_cal>5",
+        np.zeros((n_rows, 2)),
+        np.array([1084803, 1084804]),
+    )
+
+    assert isinstance(filter, np.ndarray)
+    assert filter.ndim == 2
+    assert np.shape(filter)==(n_rows,2)