add test for prediction tool

Author: TjarkMiener

ctlearn/conftest.py

@@ -3,8 +3,11 @@
 """
 
 import pytest
+import shutil
+
 from ctapipe.core import run_tool
 from ctapipe.utils import get_dataset_path
+from ctlearn.tools import TrainCTLearnModel
 
 @pytest.fixture(scope="session")
 def gamma_simtel_path():
@@ -33,13 +36,14 @@ def dl1_gamma_file(dl1_tmp_path, gamma_simtel_path):
     """
     from ctapipe.tools.process import ProcessorTool
 
+    allowed_tels = {7, 13, 15, 16, 17, 19}
     output = dl1_tmp_path / "gamma.dl1.h5"
-
     argv = [
         f"--input={gamma_simtel_path}",
         f"--output={output}",
         "--write-images",
         "--SimTelEventSource.focal_length_choice=EQUIVALENT",
+        f"--SimTelEventSource.allowed_tels={allowed_tels}",
     ]
     assert run_tool(ProcessorTool(), argv=argv, cwd=dl1_tmp_path) == 0
     return output
@@ -51,12 +55,14 @@ def dl1_proton_file(dl1_tmp_path, proton_simtel_path):
     """
     from ctapipe.tools.process import ProcessorTool
 
+    allowed_tels = {7, 13, 15, 16, 17, 19}
     output = dl1_tmp_path / "proton.dl1.h5"
     argv = [
         f"--input={proton_simtel_path}",
         f"--output={output}",
         "--write-images",
         "--SimTelEventSource.focal_length_choice=EQUIVALENT",
+        f"--SimTelEventSource.allowed_tels={allowed_tels}",
     ]
     assert run_tool(ProcessorTool(), argv=argv, cwd=dl1_tmp_path) == 0
     return output
@@ -77,4 +83,56 @@ def r1_gamma_file(r1_tmp_path, gamma_simtel_path):
         "--SimTelEventSource.focal_length_choice=EQUIVALENT",
     ]
     assert run_tool(ProcessorTool(), argv=argv, cwd=r1_tmp_path) == 0
-    return output
+    return output
+
+@pytest.fixture(scope="session")
+def ctlearn_trained_dl1_models(dl1_gamma_file, dl1_proton_file, tmp_path_factory):
+    """
+    Test training CTLearn model using the DL1 gamma and proton files for all reconstruction tasks.
+    Each test run gets its own isolated temp directories.
+    """
+    tmp_path = tmp_path_factory.mktemp("ctlearn_models")
+
+    # Temporary directories for signal and background
+    signal_dir = tmp_path / "gamma_dl1"
+    signal_dir.mkdir(parents=True, exist_ok=True)
+
+    background_dir = tmp_path / "proton_dl1"
+    background_dir.mkdir(parents=True, exist_ok=True)
+
+    # Hardcopy DL1 gamma file to the signal directory
+    shutil.copy(dl1_gamma_file, signal_dir)
+    # Hardcopy DL1 proton file to the background directory
+    shutil.copy(dl1_proton_file, background_dir)
+
+    ctlearn_trained_dl1_models = {}
+    for reco_task in ["type", "energy", "cameradirection"]:
+        # Output directory for trained model
+        output_dir = tmp_path / f"ctlearn_{reco_task}"
+        
+        # Build command-line arguments
+        argv = [
+            f"--signal={signal_dir}",
+            "--pattern-signal=*.dl1.h5",
+            f"--output={output_dir}",
+            f"--reco={reco_task}",
+            "--TrainCTLearnModel.n_epochs=1",
+            "--TrainCTLearnModel.batch_size=2",
+            "--DLImageReader.focal_length_choice=EQUIVALENT",
+        ]
+
+        # Include background only for classification task
+        if reco_task == "type":
+            argv.extend([
+                f"--background={background_dir}",
+                "--pattern-background=*.dl1.h5",
+                "--DLImageReader.enforce_subarray_equality=False",
+            ])
+
+        # Run training
+        assert run_tool(TrainCTLearnModel(), argv=argv, cwd=tmp_path) == 0
+
+        ctlearn_trained_dl1_models[reco_task] = output_dir / "ctlearn_model.keras"
+        # Check that the trained model exists
+        assert ctlearn_trained_dl1_models[reco_task].exists()
+    return ctlearn_trained_dl1_models

ctlearn/tools/tests/test_predict_model.py

@@ -0,0 +1,77 @@
+import shutil
+import numpy as np
+
+from ctapipe.core import run_tool
+from ctapipe.io import TableLoader
+from ctlearn.tools import MonoPredictCTLearnModel
+
+def test_predict_model(ctlearn_trained_dl1_models, dl1_gamma_file, tmp_path):
+    """
+    Test training CTLearn model using the DL1 gamma and proton files for all reconstruction tasks.
+    Each test run gets its own isolated temp directories.
+    """
+
+    model_dir = tmp_path / "trained_models"
+    model_dir.mkdir(parents=True, exist_ok=True)
+
+    dl2_dir = tmp_path / "dl2_output"
+    dl2_dir.mkdir(parents=True, exist_ok=True)
+
+    # Hardcopy the trained models to the model directory
+    for reco_task in ["type", "energy", "cameradirection"]:
+        shutil.copy(ctlearn_trained_dl1_models[f"{reco_task}"], model_dir / f"ctlearn_model_{reco_task}.keras")
+        model_file = model_dir / f"ctlearn_model_{reco_task}.keras"
+        assert model_file.exists(), f"Trained model file not found for {reco_task}"
+
+    # Build command-line arguments
+    output_file = dl2_dir / "gamma.dl2.h5"
+    argv = [
+        f"--input_url={dl1_gamma_file}",
+        f"--output={output_file}",
+        "--PredictCTLearnModel.batch_size=4",
+        "--DLImageReader.focal_length_choice=EQUIVALENT",
+    ]
+ 
+    # Run Prediction for energy and type together
+    assert run_tool(
+        MonoPredictCTLearnModel(),
+        argv = argv + [
+            f"--PredictCTLearnModel.load_type_model_from={model_dir}/ctlearn_model_type.keras",
+            f"--PredictCTLearnModel.load_energy_model_from={model_dir}/ctlearn_model_energy.keras",
+            "--use-HDF5Merger",
+            "--no-dl1-images",
+            "--no-true-images",
+        ],
+        cwd=tmp_path
+    ) == 0
+
+    assert run_tool(
+        MonoPredictCTLearnModel(),
+        argv= argv + [
+            f"--PredictCTLearnModel.load_cameradirection_model_from="
+            f"{model_dir}/ctlearn_model_cameradirection.keras",
+            "--no-use-HDF5Merger",
+        ],
+        cwd=tmp_path,
+    ) == 0
+
+
+    allowed_tels = [7, 13, 15, 16, 17, 19]
+    required_columns = [
+        "telescope_pointing_azimuth",
+        "telescope_pointing_altitude",
+        "CTLearn_alt",
+        "CTLearn_az",
+        "CTLearn_prediction",
+        "CTLearn_energy",
+    ]
+    # Check that the output DL2 file was created
+    assert output_file.exists(), "Output DL2 file not created"
+    # Check that the created DL2 file can be read with the TableLoader 
+    with TableLoader(output_file, pointing=True, focal_length_choice="EQUIVALENT") as loader:
+        events = loader.read_telescope_events_by_id(telescopes=allowed_tels)
+        for tel_id in allowed_tels:
+            assert len(events[tel_id]) > 0
+            for col in required_columns:
+                assert col in events[tel_id].colnames, f"{col} missing in DL2 file {output_file.name}"
+                assert events[tel_id][col][0] is not np.nan, f"{col} has NaN values in DL2 file {output_file.name}"

ctlearn/tools/tests/test_train_model.py

@@ -39,15 +39,11 @@ def test_train_ctlearn_model(reco_task, dl1_gamma_file, dl1_proton_file, tmp_pat
 
     # Include background only for classification task
     if reco_task == "type":
-        allowed_tels = {7, 13, 15, 16, 17, 19}
         argv.extend([
             f"--background={background_dir}",
             "--pattern-background=*.dl1.h5",
-            f"--DLImageReader.allowed_tels={allowed_tels}",
             "--DLImageReader.enforce_subarray_equality=False",
         ])
-    else:
-        argv.extend(["--DLImageReader.allowed_tel_types=LST_LST_LSTCam"])
 
     # Run training
     assert run_tool(TrainCTLearnModel(), argv=argv, cwd=tmp_path) == 0