Smarter defaults for MadNIS

madgraph/iolibs/template_files/mg7/gridpack.py

@@ -171,7 +171,7 @@ def main() -> None:
     event_generator = ms.EventGenerator(
         contexts=contexts,
         channels=channel_generators,
-        status_file=os.path.join(run_path, "info.json"),
+        status_file=ms.StatusFile(os.path.join(run_path, "info.json")),
         config=config,
     )
 

madgraph/iolibs/template_files/mg7/madevent.py

@@ -205,6 +205,7 @@ def init_event_dir(self) -> None:
                 break
             except FileExistsError:
                 run_index += 1
+        self.status_file = ms.StatusFile(os.path.join(self.run_path, "info.json"))
 
     def init_context(self) -> None:
         device_names = self.run_card["run"]["devices"]
@@ -424,7 +425,7 @@ def build_event_generator(self, phasespaces: list[PhaseSpace]) -> ms.EventGenera
         event_generator = ms.EventGenerator(
             contexts=self.contexts,
             channels=channel_generators,
-            status_file=os.path.join(self.run_path, "info.json"),
+            status_file=self.status_file,
             config=self.event_generator_config,
         )
         unused_globals = (
@@ -502,9 +503,6 @@ def train_madnis(self) -> None:
         madnis_args = self.run_card["madnis"]
         if not madnis_args["enable"]:
             return
-        if madnis_args.get("old", False):
-            self.train_madnis_old()
-            return
 
         gen_args = self.run_card["generation"]
         run_args = self.run_card["run"]
@@ -528,6 +526,7 @@ def train_madnis(self) -> None:
         config.adam_beta1 = madnis_args["adam_beta1"]
         config.adam_beta2 = madnis_args["adam_beta2"]
         config.adam_eps = madnis_args["adam_eps"]
+        config.grad_clip_threshold = madnis_args["grad_clip_threshold"]
         config.buffer_capacity = madnis_args["buffer_capacity"]
         config.minimum_buffer_size = madnis_args["minimum_buffer_size"]
         config.buffered_steps = madnis_args["buffered_steps"]
@@ -559,6 +558,7 @@ def train_madnis(self) -> None:
             config=config,
             integrands=integrands,
             cwnets=cwnets,
+            status_file=self.status_file,
         )
         madnis_training.train()
         for phasespace, active_channels in zip(
@@ -572,49 +572,6 @@ def train_madnis(self) -> None:
             context.copy_globals_from(self.contexts[0])
         self.event_generator = self.build_event_generator(madnis_phasespaces)
 
-    def train_madnis_old(self) -> None:
-        madnis_args = self.run_card["madnis"]
-        if not madnis_args["enable"]:
-            return
-
-        if len(self.subprocesses) > 1:
-            self.madnis_lower_box = ms.PrettyBox(
-                "Subprocesses", len(self.subprocesses) + 1, [12, 12, 12, 0],
-            )
-            self.madnis_lower_box.set_row(0, ["Subprocess", "Loss", "Channels", "Batch"])
-            self.madnis_upper_box = ms.PrettyBox(
-                "MadNIS training", 2, [18, 0], self.madnis_lower_box.line_count
-            )
-            self.madnis_upper_box.set_column(0, ["Subprocesses:", "Run time:"])
-            self.madnis_upper_box.print_first()
-            self.madnis_lower_box.print_first()
-        else:
-            self.madnis_box = ms.PrettyBox(
-                "MadNIS training", 4, [18, 0]
-            )
-            self.madnis_box.set_column(0, ["Batch:", "Loss:", "Channels:", "Run time:"])
-            self.madnis_box.print_first()
-
-        self.last_update_time = 0
-        self.madnis_wall_time = time.time()
-        self.madnis_cpu_time = time.process_time()
-
-        madnis_phasespaces = []
-        for subproc, phasespace in zip(self.subprocesses, self.phasespaces):
-            phasespace = subproc.build_madnis(phasespace)
-            if len(self.subprocesses) > 1:
-                status_func = lambda *args: self.update_madnis_status_multi(
-                    subproc.subproc_id, *args
-                )
-            else:
-                status_func = self.update_madnis_status_single
-            subproc.train_madnis(phasespace, status_func)
-            madnis_phasespaces.append(phasespace)
-        self.phasespaces = madnis_phasespaces
-        for context in self.contexts[1:]:
-            context.copy_globals_from(self.contexts[0])
-        self.event_generator = self.build_event_generator(madnis_phasespaces)
-
     def update_madnis_status_single(
         self, batch: int, batch_target: int, loss: float, lr: float, channel_count: int
     ) -> None:

madgraph/iolibs/template_files/mg7/run_card.toml

@@ -121,6 +121,7 @@ lr_scheduler = %(madnis.lr_scheduler)s # options: none, cosine
 adam_beta1 = %(madnis.adam_beta1)s
 adam_beta2 = %(madnis.adam_beta2)s
 adam_eps = %(madnis.adam_eps)s
+grad_clip_threshold = %(madnis.grad_clip_threshold)s
 train_mcw = %(madnis.train_mcw)s
 buffer_capacity = %(madnis.buffer_capacity)s
 minimum_buffer_size = %(madnis.minimum_buffer_size)s

madgraph/various/banner.py

@@ -6572,10 +6572,11 @@ def default_setup(self):
         self.add_toml_param('madnis', 'adam_beta1', 0.9)
         self.add_toml_param('madnis', 'adam_beta2', 0.999)
         self.add_toml_param('madnis', 'adam_eps', 1e-8)
+        self.add_toml_param('madnis', 'grad_clip_threshold', 0.0)
         self.add_toml_param('madnis', 'train_mcw', True)
-        self.add_toml_param('madnis', 'buffer_capacity', 0)
-        self.add_toml_param('madnis', 'minimum_buffer_size', 50000)
-        self.add_toml_param('madnis', 'buffered_steps', 0)
+        self.add_toml_param('madnis', 'buffer_capacity', 100000)
+        self.add_toml_param('madnis', 'minimum_buffer_size', 10000)
+        self.add_toml_param('madnis', 'buffered_steps', 5)
         self.add_toml_param('madnis', 'buffer_unweighting_quantile', 0.99)
         self.add_toml_param('madnis', 'uniform_channel_ratio', 0.1)
         self.add_toml_param('madnis', 'integration_history_length', 100)

madspace/CMakeLists.txt

@@ -173,6 +173,7 @@ add_library(
     src/driver/event_generator.cpp
     src/driver/channel_generator.cpp
     src/driver/generator_data.cpp
+    src/driver/status_file.cpp
     src/driver/thread_pool.cpp
     src/driver/io.cpp
     src/driver/vegas_optimizer.cpp

madspace/include/madspace/compgraphs/function_builder_mixin.inc

@@ -87,6 +87,10 @@ Value reduce_sum_vector(Value in1) {
     return instruction("reduce_sum_vector", {in1})[0];
 }
 
+Value batch_reduce_sum(Value in1) {
+    return instruction("batch_reduce_sum", {in1})[0];
+}
+
 Value batch_reduce_mean(Value in1) {
     return instruction("batch_reduce_mean", {in1})[0];
 }