Add would_stop_epoch to json

Author: achiefa

n3fit/src/n3fit/io/writer.py

@@ -279,18 +279,11 @@ 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()
@@ -310,6 +303,11 @@ def _write_metadata_json(self, i, out_path):
             # Note: the 2 arguments below are the same for all replicas, unless run separately
             timing=self.timings,
             stop_epoch=self.stopping_object.stop_epoch,
+            would_stop_epoch=(
+                self.stopping_object.would_stop_epoch
+                if self.stopping_object._dont_stop
+                else self.stopping_object.stop_epoch
+            ),
         )
 
         with open(out_path, "w", encoding="utf-8") as fs:
@@ -354,6 +352,7 @@ def jsonfit(
     true_chi2,
     stop_epoch,
     timing,
+    would_stop_epoch,
 ):
     """Generates a dictionary containing all relevant metadata for the fit
 
@@ -379,6 +378,8 @@ def jsonfit(
             epoch at which the stopping stopped (not the one for the best fit!)
         timing: dict
             dictionary of the timing of the different events that happened
+        would_stop_epoch: int
+            epoch at which the stopping would have stopped if it were not set to "dont_stop"
     """
     all_info = {}
     # Generate preprocessing information
@@ -393,6 +394,7 @@ def jsonfit(
     all_info["arc_lengths"] = arc_lengths
     all_info["integrability"] = integrability_numbers
     all_info["timing"] = timing
+    all_info["would_stop_epoch"] = would_stop_epoch
     # Versioning info
     all_info["version"] = version()
     return all_info

n3fit/src/n3fit/stopping.py

@@ -487,8 +487,7 @@ def make_stop(self):
             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()
+        self._restore_best_weights()
 
     @property
     def would_stop_epoch(self):