Remove stopping and add would stop

Author: achiefa

n3fit/src/n3fit/io/writer.py

@@ -279,11 +279,18 @@ def write_data(self, save_path, fitname, weights_name):
             replica_path.mkdir(exist_ok=True, parents=True)
 
             self._write_chi2s(replica_path / "chi2exps.log")
+            self._write_would_stop_epoch(replica_path / "would_stop_epoch.txt")
             self._write_metadata_json(i, replica_path / f"{fitname}.json")
             self._export_pdf_grid(i, replica_path / f"{fitname}.exportgrid")
             if weights_name:
                 self._write_weights(i, replica_path / f"{weights_name}")
 
+    def _write_would_stop_epoch(self, out_path):
+        epoch = self.stopping_object.would_stop_epoch
+        with open(out_path, "w", encoding="utf-8") as f:
+            f.write(str(epoch) if epoch is not None else "None")
+            f.write("\n")
+
     def _write_chi2s(self, out_path):
         # Note: same for all replicas, unless run separately
         chi2_log = self.stopping_object.chi2exps_json()

n3fit/src/n3fit/model_trainer.py

@@ -116,6 +116,7 @@ def __init__(
         save_checkpoints=False,
         replica_path=None,
         checkpoint_freq=100,
+        dont_stop=False,
     ):
         """
         Parameters
@@ -163,6 +164,8 @@ def __init__(
                 root path for all replicas.
             checkpoint_freq: int
                 frequency (in epochs) at which to save checkpoints. Only relevant if `save_checkpoints` is True.
+            dont_stop: bool
+                whether to disable the stopping mechanism, i.e. to run for all epochs regardless of the validation chi2
         """
         # Save all input information
         self.exp_info = list(exp_info)
@@ -179,6 +182,7 @@ def __init__(
         self.lux_params = lux_params
         self.replicas = replicas
         self.experiments_data = experiments_data
+        self.dont_stop = dont_stop
 
         # Checkpointing options
         self.save_checkpoints = save_checkpoints
@@ -1035,6 +1039,7 @@ def hyperparametrizable(self, params):
                 stopping_patience=stopping_epochs,
                 threshold_positivity=threshold_pos,
                 threshold_chi2=threshold_chi2,
+                dont_stop=self.dont_stop,
             )
 
             # Compile each of the models with the right parameters

n3fit/src/n3fit/performfit.py

@@ -44,6 +44,7 @@ def performfit(
     parallel_models=True,
     save_checkpoints=False,
     checkpoint_freq=100,
+    dont_stop=False,
 ):
     """
     This action will (upon having read a validcard) process a full PDF fit
@@ -204,6 +205,7 @@ def performfit(
             save_checkpoints=save_checkpoints,
             replica_path=replica_path,
             checkpoint_freq=checkpoint_freq,
+            dont_stop=dont_stop,
         )
 
         # This is just to give a descriptive name to the fit function

n3fit/src/n3fit/stopping.py

@@ -345,6 +345,7 @@ def __init__(
 
         self._dont_stop = dont_stop
         self._stop_now = False
+        self._would_stop_epoch = None
         self.stopping_patience = stopping_patience
         self.total_epochs = total_epochs
 
@@ -481,7 +482,20 @@ def make_stop(self):
         and reload the history to the point of the best model if any
         """
         self._stop_now = True
-        self._restore_best_weights()
+        if self._would_stop_epoch is None:
+            # final_epoch is the last registered epoch (0-indexed); +1 to match stop_epoch convention
+            self._would_stop_epoch = (
+                -1 if self._history.final_epoch is None else self._history.final_epoch + 1
+            )
+        if not self._dont_stop:
+            self._restore_best_weights()
+
+    @property
+    def would_stop_epoch(self):
+        """Epoch at which early stopping would have triggered.
+        Returns None if stopping never triggered (fit converged within total_epochs).
+        When dont_stop=False this equals stop_epoch."""
+        return self._would_stop_epoch
 
     def _restore_best_weights(self):
         for i_replica, weights in enumerate(self._best_weights):