From e192d8c46e3a30090760247acca2b7d46ba1ec64 Mon Sep 17 00:00:00 2001
From: Lukas Beiske <lukas.beiske@tu-dortmund.de>
Date: Tue, 1 Oct 2024 17:12:58 +0200
Subject: [PATCH] Add test for cut optimization tool

---
 .../tests/test_optimize_event_selection.py    | 63 +++++++++++++++++++
 1 file changed, 63 insertions(+)
 create mode 100644 src/ctapipe/tools/tests/test_optimize_event_selection.py

diff --git a/src/ctapipe/tools/tests/test_optimize_event_selection.py b/src/ctapipe/tools/tests/test_optimize_event_selection.py
new file mode 100644
index 00000000000..d8f42c506c0
--- /dev/null
+++ b/src/ctapipe/tools/tests/test_optimize_event_selection.py
@@ -0,0 +1,63 @@
+import json
+
+import astropy.units as u
+import pytest
+from astropy.table import QTable
+
+from ctapipe.core import run_tool
+
+
+@pytest.mark.parametrize("point_like", (True, False))
+def test_cuts_optimization(
+    gamma_diffuse_full_reco_file,
+    proton_full_reco_file,
+    irf_events_loader_test_config,
+    tmp_path,
+    point_like,
+):
+    from ctapipe.irf import (
+        OptimizationResult,
+        OptimizationResultStore,
+        ResultValidRange,
+    )
+    from ctapipe.tools.optimize_event_selection import IrfEventSelector
+
+    output_path = tmp_path / "cuts.fits"
+    config_path = tmp_path / "config.json"
+    with config_path.open("w") as f:
+        json.dump(irf_events_loader_test_config, f)
+
+    argv = [
+        f"--gamma-file={gamma_diffuse_full_reco_file}",
+        f"--proton-file={proton_full_reco_file}",
+        # Use diffuse gammas weighted to electron spectrum as stand-in
+        f"--electron-file={gamma_diffuse_full_reco_file}",
+        f"--output={output_path}",
+        f"--config={config_path}",
+    ]
+    if not point_like:
+        argv.append("--full-enclosure")
+
+    ret = run_tool(
+        IrfEventSelector(),
+        argv=argv,
+    )
+    assert ret == 0
+
+    result = OptimizationResultStore().read(output_path)
+    assert isinstance(result, OptimizationResult)
+    assert isinstance(result.valid_energy, ResultValidRange)
+    assert isinstance(result.valid_offset, ResultValidRange)
+    assert isinstance(result.gh_cuts, QTable)
+    assert result.gh_cuts.meta["CLFNAME"] == "ExtraTreesClassifier"
+    assert "cut" in result.gh_cuts.colnames
+    if point_like:
+        assert isinstance(result.theta_cuts, QTable)
+        assert "cut" in result.theta_cuts.colnames
+
+    for c in ["low", "center", "high"]:
+        assert c in result.gh_cuts.colnames
+        assert result.gh_cuts[c].unit == u.TeV
+        if point_like:
+            assert c in result.theta_cuts.colnames
+            assert result.theta_cuts[c].unit == u.TeV