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/run_card.toml

@@ -69,7 +69,7 @@ sde_strategy = %(phasespace.sde_strategy)s #options: diagrams, denominators
 decays = %(phasespace.decays)s # options: all, massive, none
 t_channel = %(phasespace.t_channel)s # options: propagator, rambo, chili
 flat_mode = %(phasespace.flat_mode)s # options: propagator, rambo, chili
-simplified_channel_count = %(phasespace.simplified_channel_count)s
+combine_channel_threshold = %(phasespace.combine_channel_threshold)s
 invariant_power = %(phasespace.invariant_power)s
 bw_cutoff = %(phasespace.bw_cutoff)s
 
@@ -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
@@ -136,3 +137,4 @@ batch_size_threshold = %(madnis.batch_size_threshold)s
 channel_grouping_mode = %(madnis.channel_grouping_mode)s # options: none, uniform, learned
 fixed_cwnet_fraction = %(madnis.fixed_cwnet_fraction)s
 softclip_threshold = %(madnis.softclip_threshold)s
+compressed_channel_weight_count = %(madnis.compressed_channel_weight_count)s

madgraph/various/banner.py

@@ -6425,11 +6425,19 @@ def __init__(self, *args, **opts):
     # ------------------------------------------------------------------
     # parameter declaration
     # ------------------------------------------------------------------
-    def add_toml_param(self, section, key, value, gridpack=False, **opts):
+    def add_toml_param(self, section, key, value, gridpack=False, auto=False, **opts):
         """Declare one fixed (typed) TOML parameter belonging to ``section``.
 
         ``gridpack=True`` marks the parameter as relevant during gridpack
-        execution; such params are written to ``grid_run_card.toml``."""
+        execution; such params are written to ``grid_run_card.toml``.
+
+        ``auto=True`` declares a numeric parameter whose default is
+        determined automatically: ``value`` fixes the accepted type (int/
+        float/...) and provides the internal placeholder default, but the
+        card reads/writes it as the string ``"auto"`` until the user sets an
+        explicit value of that type. This reuses the same ``auto_set``
+        machinery as typing ``auto`` for any other numeric parameter, so
+        type-checking for genuine numeric overrides is unaffected."""
         section = section.lower()
         key = key.lower()
         internal = '%s.%s' % (section, key)
@@ -6442,6 +6450,8 @@ def add_toml_param(self, section, key, value, gridpack=False, **opts):
             self.toml_sections[section].append(key)
         if gridpack:
             self.gridpack_params.add(internal)
+        if auto:
+            self.auto_set.add(internal)
 
     def default_setup(self):
         """Define every parameter of the default ``run_card.toml``."""
@@ -6526,8 +6536,8 @@ def default_setup(self):
         self.add_toml_param('vegas', 'max_batch_size', 32000)
 
         # -------------------------- [phasespace] ----------------------
-        self.add_toml_param('phasespace', 'mode', "multichannel",
-            allowed=['multichannel', 'flat', 'both'])
+        self.add_toml_param('phasespace', 'mode', "auto",
+            allowed=['auto', 'multichannel', 'flat', 'both'])
         self.add_toml_param('phasespace', 'sde_strategy', "diagrams",
             allowed=['diagrams', 'denominators'])
         self.add_toml_param('phasespace', 'decays', "all",
@@ -6536,23 +6546,23 @@ def default_setup(self):
             allowed=['propagator', 'rambo', 'chili'])
         self.add_toml_param('phasespace', 'flat_mode', "rambo",
             allowed=['propagator', 'rambo', 'chili'])
-        self.add_toml_param('phasespace', 'simplified_channel_count', 10)
+        self.add_toml_param('phasespace', 'combine_channel_threshold', 0.01)
         self.add_toml_param('phasespace', 'invariant_power', 0.7)
         self.add_toml_param('phasespace', 'bw_cutoff', 15)
 
         # ----------------------------- [madnis] -----------------------
-        self.add_toml_param('madnis', 'enable', False)
-        self.add_toml_param('madnis', 'flow_hidden_dim', 64)
-        self.add_toml_param('madnis', 'flow_layers', 3)
+        self.add_toml_param('madnis', 'enable', False, allowed=["auto", True, False], auto=True)
+        self.add_toml_param('madnis', 'flow_hidden_dim', 64, auto=True)
+        self.add_toml_param('madnis', 'flow_layers', 3, auto=True)
         self.add_toml_param('madnis', 'flow_spline_bins', 10)
         self.add_toml_param('madnis', 'flow_activation', "leaky_relu",
             allowed=['relu', 'leaky_relu', 'elu', 'gelu', 'sigmoid', 'softplus'])
         self.add_toml_param('madnis', 'flow_invert_spline', False)
-        self.add_toml_param('madnis', 'discrete_hidden_dim', 64)
+        self.add_toml_param('madnis', 'discrete_hidden_dim', 64, auto=True)
         self.add_toml_param('madnis', 'discrete_layers', 3)
         self.add_toml_param('madnis', 'discrete_activation', "leaky_relu",
             allowed=['relu', 'leaky_relu', 'elu', 'gelu', 'sigmoid', 'softplus'])
-        self.add_toml_param('madnis', 'cwnet_hidden_dim', 64)
+        self.add_toml_param('madnis', 'cwnet_hidden_dim', 64, auto=True)
         self.add_toml_param('madnis', 'cwnet_layers', 3)
         self.add_toml_param('madnis', 'cwnet_activation', "leaky_relu",
             allowed=['relu', 'leaky_relu', 'elu', 'gelu', 'sigmoid', 'softplus'])
@@ -6564,18 +6574,19 @@ def default_setup(self):
         self.add_toml_param('madnis', 'batch_size_per_channel', 128)
         self.add_toml_param('madnis', 'generator_target_size_factor', 32)
         self.add_toml_param('madnis', 'gpu_generator_batch_granularity', 1000)
-        self.add_toml_param('madnis', 'lr', 1e-3)
+        self.add_toml_param('madnis', 'lr', 3e-4, auto=True)
         self.add_toml_param('madnis', 'lr_decay', 0.01)
         self.add_toml_param('madnis', 'lr_max', 3e-3)
         self.add_toml_param('madnis', 'lr_scheduler', "cosine",
             allowed=['none', 'cosine'])
         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)
@@ -6588,6 +6599,7 @@ def default_setup(self):
             allowed=['none', 'uniform', 'learned'])
         self.add_toml_param('madnis', 'fixed_cwnet_fraction', 0.33)
         self.add_toml_param('madnis', 'softclip_threshold', 30.0)
+        self.add_toml_param('madnis', 'compressed_channel_weight_count', 50)
 
         # ----------------- dynamic (free-form) sections ---------------
         self.dynamic_sections['multiparticles'] = collections.OrderedDict([

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];
 }
@@ -366,6 +370,15 @@ Value apply_subchannel_weights(Value channel_weights_in, Value subchannel_weight
     return instruction("apply_subchannel_weights", {channel_weights_in, subchannel_weights, channel_indices, subchannel_indices})[0];
 }
 
+std::array<Value, 2> compress_channel_weights(Value channel_index, Value channel_weights, Value keep_count) {
+    auto output_vector = instruction("compress_channel_weights", {channel_index, channel_weights, keep_count});
+    return {output_vector[0], output_vector[1]};
+}
+
+Value restore_channel_weights(Value chan_weight_values, Value chan_weight_indices, Value full_count) {
+    return instruction("restore_channel_weights", {chan_weight_values, chan_weight_indices, full_count})[0];
+}
+
 Value pt_eta_phi_x(Value p_ext, Value x1, Value x2) {
     return instruction("pt_eta_phi_x", {p_ext, x1, x2})[0];
 }