Integrate with rtbm

examples/simgauss_tf.py

@@ -4,18 +4,20 @@
     Basic example using the vegas_wrapper helper
 """
 
-from vegasflow.configflow import DTYPE
+import tensorflow as tf
+from vegasflow.configflow import DTYPE, run_eager
+
 import time
 import numpy as np
-import tensorflow as tf
 from vegasflow.vflow import vegas_wrapper
 from vegasflow.plain import plain_wrapper 
+from vegasflow.rtbm import rtbm_wrapper
 
 
 # MC integration setup
 dim = 4
-ncalls = np.int32(1e5)
-n_iter = 5
+ncalls = np.int32(1e3)
+n_iter = 3
 
 
 @tf.function
@@ -34,9 +36,16 @@ def symgauss(xarr, n_dim=None, **kwargs):
 
 if __name__ == "__main__":
     """Testing several different integrations"""
+    print(f"RTBM MC, ncalls={ncalls}:")
+    start = time.time()
+    ncalls = ncalls
+    r = rtbm_wrapper(symgauss, dim, n_iter, ncalls)
+    end = time.time()
+    print(f"RTBM took: time (s): {end-start}")
+
     print(f"VEGAS MC, ncalls={ncalls}:")
     start = time.time()
-    ncalls = 10*ncalls
+    ncalls = ncalls
     r = vegas_wrapper(symgauss, dim, n_iter, ncalls)
     end = time.time()
     print(f"Vegas took: time (s): {end-start}")

examples/sin_tf.py

@@ -0,0 +1,70 @@
+"""
+    Example: basic integration
+
+    Basic example using the vegas_wrapper helper
+"""
+
+from vegasflow import run_eager, float_me
+import tensorflow as tf
+
+import time
+import numpy as np
+from vegasflow.vflow import vegas_wrapper
+from vegasflow.rtbm import RTBMFlow
+
+
+# MC integration setup
+dim = 2
+ncalls = np.int32(1e2)
+n_iter = 5
+n_hidden = 1
+tf_pi = float_me(np.pi)
+npeaks = 2.0
+
+
+@tf.function
+def sin_fun(xarr, **kwargs):
+    """symgauss test function"""
+    res = tf.pow(tf.sin(npeaks * xarr * tf_pi), 2)
+    return tf.reduce_prod(res, axis=1)
+
+
+integrand = sin_fun
+# from simgauss_tf import symgauss as integrand
+
+if __name__ == "__main__":
+    # Testing several different integrations
+    print(f"VEGAS MC, ncalls={ncalls}:")
+    start = time.time()
+    ncalls = ncalls
+    _, vegas_instance = vegas_wrapper(integrand, dim, n_iter, ncalls, return_instance=True)
+    vegas_instance.freeze_grid()
+    end = time.time()
+    print(f"Vegas took: time (s): {end-start}")
+
+    print(f"RTBM MC, ncalls={ncalls}:")
+    start = time.time()
+    rtbm = RTBMFlow(n_dim=dim, n_events=ncalls, train=True, n_hidden=n_hidden, generations=500)
+    rtbm.compile(integrand)
+    _ = rtbm.run_integration(n_iter)
+    rtbm.freeze()
+    end = time.time()
+    print(f"RTBM took: time (s): {end-start}")
+
+    parameters = rtbm._rtbm.get_parameters()
+    param_file = f"PARAMS_for_ndim={dim}_{npeaks}xpeaks_hidden={n_hidden}.npy"
+    np.save(param_file, parameters)
+    print(f" > Saved parameters to {param_file}")
+
+    print("Results with frozen grids")
+    r = vegas_instance.run_integration(5)
+    rt = rtbm.run_integration(5)
+    print(f"Result computed by Vegas: {r[0]:.5f} +- {r[1]:.6f}")
+    print(f"Result computed by RTBM:  {rt[0]:.5f} +- {rt[1]:.6f}")
+
+# Notes
+# For 1 and 2 dimensions is enough with n_hidden=1 to get a better per event error
+# For 3 dimensions we need to go to n_hidden=2
+# We might be actually overfitting big time because after a few iterations the integration stops being so good
+#   a good way of testing this would be plotting the distribution of points one gets after each training!
+#   or fitting wrongly anyway

examples/singletop_lo_tf.py

@@ -8,11 +8,12 @@
 from vegasflow.configflow import DTYPE
 import tensorflow as tf
 from vegasflow.vflow import vegas_wrapper
+from vegasflow.rtbm import rtbm_wrapper
 
 # MC integration setup
 dim = 3
 ncalls = np.int32(1e5)
-n_iter = 5
+n_iter = 4
 
 # Physics setup
 # top mass
@@ -272,6 +273,13 @@ def singletop(xarr, n_dim=None, **kwargs):
 
 if __name__ == "__main__":
     """Testing a basic integration"""
+
+    print(f"RTBM MC, ncalls={ncalls}:")
+    start = time.time()
+    r = rtbm_wrapper(singletop, dim, n_iter, ncalls)
+    end = time.time()
+    print(f"RTBM took: time (s): {end-start}")
+
     print(f"VEGAS MC, ncalls={ncalls}:")
     start = time.time()
     r = vegas_wrapper(singletop, dim, n_iter, ncalls)

src/vegasflow/configflow.py

@@ -7,7 +7,7 @@
 import numpy as np
 
 # Set TF to only log errors
-os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "1")
+os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3")
 import tensorflow as tf
 
 # uncomment this line for debugging to avoid compiling any tf.function

src/vegasflow/monte_carlo.py

@@ -57,27 +57,7 @@ def print_iteration(it, res, error, extra="", threshold=0.1):
         logger.info(f"Result for iteration {it}: {res:.4f} +/- {error:.4f}" + extra)
 
 
-def _accumulate(accumulators):
-    """Accumulate all the quantities in accumulators
-    The default accumulation is implemented for tensorflow tensors
-    as a sum of all partial results.
 
-    Parameters
-    ----------
-        `accumulators`: list of tensorflow tensors
-
-    Returns
-    -------
-        `results`: `sum` for each element of the accumulators
-
-    Function not compiled
-    """
-    results = []
-    len_acc = len(accumulators[0])
-    for i in range(len_acc):
-        total = tf.reduce_sum([acc[i] for acc in accumulators], axis=0)
-        results.append(total)
-    return results
 
 
 class MonteCarloFlow(ABC):
@@ -99,13 +79,15 @@ class MonteCarloFlow(ABC):
 
     def __init__(
         self,
-        n_dim,
-        n_events,
+        n_dim=None,
+        n_events=None,
         events_limit=MAX_EVENTS_LIMIT,
         list_devices=DEFAULT_ACTIVE_DEVICES,
         verbose=True,
         simplify_signature=False,
     ):
+        if not ( isinstance(n_dim, int) and n_dim > 0 ):
+            raise ValueError(f"Can't integrate an integrand with non-positive or non-integer dimensionality, received: {n_dim}")
         # Save some parameters
         self.n_dim = n_dim
         self.integrand = None
@@ -132,6 +114,15 @@ def __init__(
             self.pool = joblib.Parallel(n_jobs=len(devices), prefer="threads")
         else:
             self.devices = None
+        self._train = False
+
+    @property
+    def train(self):
+        return self._train
+
+    @train.setter
+    def train(self, new_val):
+        self._train = new_val
 
     @property
     def events_per_run(self):
@@ -162,6 +153,29 @@ def history(self):
         """
         return self._history
 
+    @staticmethod
+    def _accumulate(accumulators):
+        """Accumulate all the quantities in accumulators
+        The default accumulation is implemented for tensorflow tensors
+        as a sum of all partial results.
+
+        Parameters
+        ----------
+            `accumulators`: list of tensorflow tensors
+
+        Returns
+        -------
+            `results`: `sum` for each element of the accumulators
+
+        Function not compiled
+        """
+        results = []
+        len_acc = len(accumulators[0])
+        for i in range(len_acc):
+            total = tf.reduce_sum([acc[i] for acc in accumulators], axis=0)
+            results.append(total)
+        return results
+
     def generate_random_array(self, n_events):
         """Generate a 2D array of (n_events, n_dim) points
         Parameters
@@ -314,7 +328,7 @@ def run_event(self, **kwargs):
             # Generate the client to control the distribution using the cluster variable
             client = Client(cluster)
             accumulators_future = client.map(self.device_run, events_to_do, percentages)
-            result_future = client.submit(_accumulate, accumulators_future)
+            result_future = client.submit(self._accumulate, accumulators_future)
             result = result_future.result()
             # Liberate the client
             client.close()
@@ -325,7 +339,7 @@ def run_event(self, **kwargs):
             for ncalls, pc in zip(events_to_do, percentages):
                 res = self.device_run(ncalls, sent_pc=pc, **kwargs)
                 accumulators.append(res)
-        return _accumulate(accumulators)
+        return self._accumulate(accumulators)
 
     def compile(self, integrand, compilable=True):
         """Receives an integrand, prepares it for integration
@@ -466,7 +480,7 @@ def run_integration(self, n_iter, log_time=True, histograms=None):
         return final_result.numpy(), sigma
 
 
-def wrapper(integrator_class, integrand, n_dim, n_iter, total_n_events, compilable=True):
+def wrapper(integrator_class, integrand, n_dim, n_iter, total_n_events, compilable=True, return_instance=False):
     """Convenience wrapper
 
     Parameters
@@ -482,6 +496,10 @@ def wrapper(integrator_class, integrand, n_dim, n_iter, total_n_events, compilab
         `final_result`: integral value
         `sigma`: monte carlo error
     """
-    mc_instance = integrator_class(n_dim, total_n_events)
+    mc_instance = integrator_class(n_dim = n_dim, n_events=total_n_events)
     mc_instance.compile(integrand, compilable=compilable)
-    return mc_instance.run_integration(n_iter)
+    if return_instance:
+        r = mc_instance.run_integration(n_iter)
+        return r, mc_instance
+    else:
+        return mc_instance.run_integration(n_iter)