Repository: cta-observatory/ctapipe (python)

Issue:
Write cross-validation results immediately after cross validation
**Please describe the use case that requires this feature.**
Currently, the `CrossValidator` keeps the results of the cross-validation, including the full result table in memory and writes it in `finish` method of the tools.

We should write the results immediately in the cross validation loop using `write_table(..., append=True)` to avoid keeping this information in memory unnecessarily.

**Describe the solution you'd like**

Write cross validation results out immediately after each step, not in `finish`.




Pull request:
Write out cross-validation results immediately
Fixes #2480

Patch:
diff --git a/ctapipe/reco/sklearn.py b/ctapipe/reco/sklearn.py
index bf85ad36..51928b6d 100644
--- a/ctapipe/reco/sklearn.py
+++ b/ctapipe/reco/sklearn.py
@@ -10,11 +10,12 @@ import astropy.units as u
 import joblib
 import numpy as np
 from astropy.coordinates import AltAz
-from astropy.table import QTable, Table, vstack
+from astropy.table import QTable, Table, hstack
 from astropy.utils.decorators import lazyproperty
 from sklearn.metrics import accuracy_score, r2_score, roc_auc_score
 from sklearn.model_selection import KFold, StratifiedKFold
 from sklearn.utils import all_estimators
+from tables import open_file
 from tqdm import tqdm
 from traitlets import TraitError, observe
 
@@ -28,7 +29,14 @@ from ..containers import (
     ReconstructedGeometryContainer,
 )
 from ..coordinates import TelescopeFrame
-from ..core import Component, FeatureGenerator, Provenance, QualityQuery, traits
+from ..core import (
+    Component,
+    FeatureGenerator,
+    Provenance,
+    QualityQuery,
+    ToolConfigurationError,
+    traits,
+)
 from ..io import write_table
 from .preprocessing import collect_features, table_to_X, telescope_to_horizontal
 from .reconstructor import ReconstructionProperty, Reconstructor
@@ -850,7 +858,7 @@ class DispReconstructor(Reconstructor):
         Update n_jobs of all associated models.
         """
         if hasattr(self, "_models"):
-            for (disp, sign) in self._models.values():
+            for disp, sign in self._models.values():
                 disp.n_jobs = n_jobs.new
                 sign.n_jobs = n_jobs.new
 
@@ -878,9 +886,8 @@ class CrossValidator(Component):
         default_value=1337, help="Random seed for splitting the training data."
     ).tag(config=True)
 
-    def __init__(self, model_component, **kwargs):
+    def __init__(self, model_component, overwrite=False, **kwargs):
         super().__init__(**kwargs)
-        self.cv_predictions = {}
         self.model_component = model_component
         self.rng = np.random.default_rng(self.rng_seed)
 
@@ -898,6 +905,29 @@ class CrossValidator(Component):
                 "Unsupported Model of type %s supplied", self.model_component
             )
 
+        if self.output_path:
+            if self.output_path.exists():
+                if overwrite:
+                    self.log.warning("Overwriting %s", self.output_path)
+                else:
+                    raise ToolConfigurationError(
+                        f"Output path {self.output_path} exists, but overwrite=False"
+                    )
+
+            Provenance().add_output_file(self.output_path, role="ml-cross-validation")
+            self.h5file = open_file(self.output_path, mode="w")
+
+    def close(self):
+        """Close the output hdf5 file, if ``self.output_path`` is given."""
+        if self.output_path:
+            self.h5file.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        self.close()
+
     def __call__(self, telescope_type, table):
         """Perform cross validation for the given model."""
         if self.n_cross_validations == 0:
@@ -913,8 +943,6 @@ class CrossValidator(Component):
         )
 
         scores = defaultdict(list)
-        predictions = []
-
         kfold = self.split_data(
             n_splits=self.n_cross_validations,
             shuffle=True,
@@ -937,21 +965,27 @@ class CrossValidator(Component):
             train = table[train_indices]
             test = table[test_indices]
 
-            cv_prediction, truth, metrics = self.cross_validate(
-                telescope_type, train, test
-            )
-            predictions.append(
-                Table(
-                    {
-                        "cv_fold": np.full(len(truth), fold, dtype=np.uint8),
-                        "tel_type": [str(telescope_type)] * len(truth),
-                        "truth": truth,
+            cv_result, metrics = self.cross_validate(telescope_type, train, test)
+            if self.output_path:
+                results = Table(
+                    data={
+                        "cv_fold": np.full(len(cv_result), fold, dtype=np.uint8),
+                        "tel_type": [str(telescope_type)] * len(cv_result),
                         "true_energy": test["true_energy"],
                         "true_impact_distance": test["true_impact_distance"],
-                        **cv_prediction,
-                    }
+                    },
+                    descriptions={
+                        "cv_fold": "Cross validation iteration",
+                        "tel_type": "Telescope type",
+                    },
+                )
+                results = hstack([results, cv_result], join_type="exact")
+                write_table(
+                    results,
+                    self.h5file,
+                    f"/cv_predictions/{telescope_type}",
+                    append=True,
                 )
-            )
 
             for metric, value in metrics.items():
                 scores[metric].append(value)
@@ -965,7 +999,6 @@ class CrossValidator(Component):
                     cv_values.mean(),
                     cv_values.std(),
                 )
-        self.cv_predictions[telescope_type] = vstack(predictions)
 
     def _cross_validate_regressor(self, telescope_type, train, test):
         regressor = self.model_component
@@ -973,7 +1006,17 @@ class CrossValidator(Component):
         prediction, _ = regressor._predict(telescope_type, test)
         truth = test[regressor.target]
         r2 = r2_score(truth, prediction)
-        return {f"{regressor.prefix}_energy": prediction}, truth, {"R^2": r2}
+        result = Table(
+            data={
+                f"{regressor.prefix}_energy": prediction,
+                "truth": truth,
+            },
+            descriptions={
+                f"{regressor.prefix}_energy": "Predicted Energy",
+                "truth": "Simulated Energy",
+            },
+        )
+        return result, {"R^2": r2}
 
     def _cross_validate_classification(self, telescope_type, train, test):
         classifier = self.model_component
@@ -985,11 +1028,17 @@ class CrossValidator(Component):
             0,
         )
         roc_auc = roc_auc_score(truth, prediction)
-        return (
-            {f"{classifier.prefix}_prediction": prediction},
-            truth,
-            {"ROC AUC": roc_auc},
+        result = Table(
+            data={
+                f"{classifier.prefix}_prediction": prediction,
+                "truth": truth,
+            },
+            descriptions={
+                f"{classifier.prefix}_prediction": "Predicted gammaness score",
+                "truth": "Particle id (default is 1 for gammas)",
+            },
         )
+        return result, {"ROC AUC": roc_auc}
 
     def _cross_validate_disp(self, telescope_type, train, test):
         models = self.model_component
@@ -998,21 +1047,16 @@ class CrossValidator(Component):
         truth = test[models.target]
         r2 = r2_score(np.abs(truth), np.abs(disp))
         accuracy = accuracy_score(np.sign(truth), np.sign(disp))
-        prediction = {
-            f"{models.prefix}_parameter": disp,
-            f"{models.prefix}_sign_score": sign_score,
-        }
-        return prediction, truth, {"R^2": r2, "accuracy": accuracy}
-
-    def write(self, overwrite=False):
-        if self.n_cross_validations == 0:
-            return 0
-
-        if self.output_path.exists() and not overwrite:
-            raise IOError(f"Path {self.output_path} exists and overwrite=False")
-
-        Provenance().add_output_file(self.output_path, role="ml-cross-validation")
-        for tel_type, results in self.cv_predictions.items():
-            write_table(
-                results, self.output_path, f"/cv_predictions_{tel_type}", overwrite=True
-            )
+        result = Table(
+            data={
+                f"{models.prefix}_parameter": disp,
+                f"{models.prefix}_sign_score": sign_score,
+                "truth": truth,
+            },
+            descriptions={
+                f"{models.prefix}_parameter": "Predicted disp parameter",
+                f"{models.prefix}_sign_score": "Score for how certain the disp sign classification was",
+                "truth": "True disp parameter",
+            },
+        )
+        return result, {"R^2": r2, "accuracy": accuracy}
diff --git a/ctapipe/tools/train_disp_reconstructor.py b/ctapipe/tools/train_disp_reconstructor.py
index 405eb9f9..ac2813df 100644
--- a/ctapipe/tools/train_disp_reconstructor.py
+++ b/ctapipe/tools/train_disp_reconstructor.py
@@ -108,12 +108,15 @@ class TrainDispReconstructor(Tool):
             )
         )
         self.n_events.attach_subarray(self.loader.subarray)
-
         self.models = DispReconstructor(self.loader.subarray, parent=self)
 
-        self.cross_validate = CrossValidator(parent=self, model_component=self.models)
+        self.cross_validate = self.enter_context(
+            CrossValidator(
+                parent=self, model_component=self.models, overwrite=self.overwrite
+            )
+        )
         self.rng = np.random.default_rng(self.random_seed)
-        self.check_output(self.output_path, self.cross_validate.output_path)
+        self.check_output(self.output_path)
 
     def start(self):
         """
@@ -192,9 +195,8 @@ class TrainDispReconstructor(Tool):
         self.log.info("Writing output")
         self.models.n_jobs = None
         self.models.write(self.output_path, overwrite=self.overwrite)
-        if self.cross_validate.output_path:
-            self.cross_validate.write(overwrite=self.overwrite)
         self.loader.close()
+        self.cross_validate.close()
 
 
 def main():
diff --git a/ctapipe/tools/train_energy_regressor.py b/ctapipe/tools/train_energy_regressor.py
index faf7c12d..2c3d8882 100644
--- a/ctapipe/tools/train_energy_regressor.py
+++ b/ctapipe/tools/train_energy_regressor.py
@@ -86,7 +86,7 @@ class TrainEnergyRegressor(Tool):
 
     def setup(self):
         """
-        Initialize components from config
+        Initialize components from config.
         """
         self.loader = self.enter_context(
             TableLoader(
@@ -99,14 +99,15 @@ class TrainEnergyRegressor(Tool):
             )
         )
         self.n_events.attach_subarray(self.loader.subarray)
-
         self.regressor = EnergyRegressor(self.loader.subarray, parent=self)
-        self.log.warning(f"{self.regressor._models}")
-        self.cross_validate = CrossValidator(
-            parent=self, model_component=self.regressor
+
+        self.cross_validate = self.enter_context(
+            CrossValidator(
+                parent=self, model_component=self.regressor, overwrite=self.overwrite
+            )
         )
         self.rng = np.random.default_rng(self.random_seed)
-        self.check_output(self.output_path, self.cross_validate.output_path)
+        self.check_output(self.output_path)
 
     def start(self):
         """
@@ -147,9 +148,8 @@ class TrainEnergyRegressor(Tool):
         self.log.info("Writing output")
         self.regressor.n_jobs = None
         self.regressor.write(self.output_path, overwrite=self.overwrite)
-        if self.cross_validate.output_path:
-            self.cross_validate.write(overwrite=self.overwrite)
         self.loader.close()
+        self.cross_validate.close()
 
 
 def main():
diff --git a/ctapipe/tools/train_particle_classifier.py b/ctapipe/tools/train_particle_classifier.py
index 14ca3183..c9fac65c 100644
--- a/ctapipe/tools/train_particle_classifier.py
+++ b/ctapipe/tools/train_particle_classifier.py
@@ -115,9 +115,8 @@ class TrainParticleClassifier(Tool):
 
     def setup(self):
         """
-        Initialize components from config
+        Initialize components from config.
         """
-
         self.signal_loader = self.enter_context(
             TableLoader(
                 parent=self,
@@ -129,7 +128,6 @@ class TrainParticleClassifier(Tool):
                 load_instrument=True,
             )
         )
-
         self.background_loader = self.enter_context(
             TableLoader(
                 parent=self,
@@ -141,20 +139,21 @@ class TrainParticleClassifier(Tool):
                 load_instrument=True,
             )
         )
-
         if self.signal_loader.subarray != self.background_loader.subarray:
             raise ValueError("Signal and background subarrays do not match")
 
         self.subarray = self.signal_loader.subarray
         self.n_signal.attach_subarray(self.subarray)
         self.n_background.attach_subarray(self.subarray)
-
         self.classifier = ParticleClassifier(subarray=self.subarray, parent=self)
-        self.rng = np.random.default_rng(self.random_seed)
-        self.cross_validate = CrossValidator(
-            parent=self, model_component=self.classifier
+
+        self.cross_validate = self.enter_context(
+            CrossValidator(
+                parent=self, model_component=self.classifier, overwrite=self.overwrite
+            )
         )
-        self.check_output(self.output_path, self.cross_validate.output_path)
+        self.rng = np.random.default_rng(self.random_seed)
+        self.check_output(self.output_path)
 
     def start(self):
         """
@@ -217,8 +216,7 @@ class TrainParticleClassifier(Tool):
         self.classifier.write(self.output_path, overwrite=self.overwrite)
         self.signal_loader.close()
         self.background_loader.close()
-        if self.cross_validate.output_path:
-            self.cross_validate.write(overwrite=self.overwrite)
+        self.cross_validate.close()
 
 
 def main():
diff --git a/docs/changes/2483.optimization.rst b/docs/changes/2483.optimization.rst
new file mode 100644
index 00000000..36ed94bc
--- /dev/null
+++ b/docs/changes/2483.optimization.rst
@@ -0,0 +1,1 @@
+Write cross validation results for each model out immediatly after validation to free up memory earlier.


Test patch:
diff --git a/ctapipe/tools/tests/test_train.py b/ctapipe/tools/tests/test_train.py
index 216e7e7a..abee895f 100644
--- a/ctapipe/tools/tests/test_train.py
+++ b/ctapipe/tools/tests/test_train.py
@@ -1,6 +1,6 @@
 import pytest
 
-from ctapipe.core import run_tool
+from ctapipe.core import ToolConfigurationError, run_tool
 from ctapipe.exceptions import TooFewEvents
 from ctapipe.utils.datasets import resource_file
 
@@ -86,6 +86,34 @@ def test_cross_validation_results(tmp_path, gamma_train_clf, proton_train_clf):
     assert ret == 0
     assert energy_cv_out_file.exists()
 
+    # test overwrite of cv results works
+    with pytest.raises(
+        ToolConfigurationError,
+        match=f"Output path {energy_cv_out_file} exists, but overwrite=False",
+    ):
+        run_tool(
+            tool,
+            argv=[
+                "--input=dataset://gamma_diffuse_dl2_train_small.dl2.h5",
+                f"--output={out_file}",
+                f"--config={config}",
+                f"--cv-output={energy_cv_out_file}",
+            ],
+            raises=True,
+        )
+
+    ret = run_tool(
+        tool,
+        argv=[
+            "--input=dataset://gamma_diffuse_dl2_train_small.dl2.h5",
+            f"--output={out_file}",
+            f"--config={config}",
+            f"--cv-output={energy_cv_out_file}",
+            "--overwrite",
+        ],
+    )
+    assert ret == 0
+
     tool = TrainParticleClassifier()
     config = resource_file("train_particle_classifier.yaml")
     out_file = tmp_path / "particle_classifier_.pkl"