Vegas+

examples/retracing.py

@@ -0,0 +1,42 @@
+"""
+    Retracing example in VegasFlowPlus
+"""
+
+from vegasflow import VegasFlowPlus, VegasFlow, PlainFlow
+from vegasflow.configflow import DTYPE, DTYPEINT, run_eager, float_me
+import time
+import numpy as np
+import tensorflow as tf
+
+# MC integration setup
+dim = 2
+ncalls = np.int32(1e3)
+n_iter = 5
+
+@tf.function(input_signature=[
+                   tf.TensorSpec(shape=[None,dim], dtype=DTYPE),
+                   tf.TensorSpec(shape=[], dtype=DTYPEINT),
+                  tf.TensorSpec(shape=[None], dtype=DTYPE)
+               ]
+           )
+def symgauss(xarr, n_dim=None, weight=None, **kwargs):
+    """symgauss test function"""
+    if n_dim is None:
+        n_dim = xarr.shape[-1]
+    a = tf.constant(0.1, dtype=DTYPE)
+    n100 = tf.cast(100 * n_dim, dtype=DTYPE)
+    pref = tf.pow(1.0 / a / np.sqrt(np.pi), float_me(n_dim))
+    coef = tf.reduce_sum(tf.range(n100 + 1))
+    coef += tf.reduce_sum(tf.square((xarr - 1.0 / 2.0) / a), axis=1)
+    coef -= (n100 + 1) * n100 / 2.0
+    return pref * tf.exp(-coef)
+
+
+
+if __name__ == "__main__":
+    """Testing several different integrations"""
+
+#     run_eager()
+    vegas_instance = VegasFlowPlus(dim, ncalls,adaptive=True)
+    vegas_instance.compile(symgauss)
+    vegas_instance.run_integration(n_iter)

src/vegasflow/__init__.py

@@ -1,8 +1,10 @@
 """Monte Carlo integration with Tensorflow"""
 
 from vegasflow.configflow import int_me, float_me, run_eager
+
 # Expose the main interfaces
 from vegasflow.vflow import VegasFlow, vegas_wrapper, vegas_sampler
 from vegasflow.plain import PlainFlow, plain_wrapper, plain_sampler
+from vegasflow.vflowplus import VegasFlowPlus, vegasflowplus_wrapper, vegasflowplus_sampler
 
-__version__ = "1.2.2"
+__version__ = "1.3.0"

src/vegasflow/configflow.py

@@ -9,11 +9,17 @@
 # Some global parameters
 BINS_MAX = 50
 ALPHA = 1.5
+BETA = 0.75  # Vegas+
 TECH_CUT = 1e-8
+
 # Set up the logistics of the integration
+# set the limits lower if hitting memory problems
+
 # Events Limit limits how many events are done in one single run of the event_loop
-# set it lower if hitting memory problems
 MAX_EVENTS_LIMIT = int(1e6)
+# Maximum number of evaluation per hypercube for VegasFlowPlus
+MAX_NEVAL_HCUBE = int(1e4)
+
 # Select the list of devices to look for
 DEFAULT_ACTIVE_DEVICES = ["GPU"]  # , 'CPU']
 
@@ -66,15 +72,15 @@
 else:
     DTYPE = tf.float64
     FMAX = tf.constant(np.finfo(np.float64).max, dtype=DTYPE)
-    logger.warning(f"PDFFLOW_FLOAT={_float_env} not understood, defaulting to 64 bits")
+    logger.warning(f"VEGASFLOW_FLOAT={_float_env} not understood, defaulting to 64 bits")
 
 if _int_env == "64":
     DTYPEINT = tf.int64
 elif _int_env == "32":
     DTYPEINT = tf.int32
 else:
     DTYPEINT = tf.int64
-    logger.warning(f"PDFFLOW_INT={_int_env} not understood, defaulting to 64 bits")
+    logger.warning(f"VEGASFLOW_INT={_int_env} not understood, defaulting to 64 bits")
 
 
 def run_eager(flag=True):

src/vegasflow/monte_carlo.py

@@ -48,6 +48,7 @@
     DTYPEINT,
     TECH_CUT,
     float_me,
+    int_me,
 )
 
 
@@ -330,7 +331,7 @@ def set_distribute(self, queue_object):
         self.cluster = queue_object
         self.distribute = True
 
-    def run_event(self, **kwargs):
+    def run_event(self, tensorize_events=False, **kwargs):
         """
         Runs the Monte Carlo event. This corresponds to a number of calls
         decided by the `events_per_run` variable. The variable `acc` is exposed
@@ -358,7 +359,10 @@ def run_event(self, **kwargs):
             ncalls = min(events_left, self.events_per_run)
             pc += ncalls / self.n_events * 100
             percentages.append(pc)
-            events_to_do.append(ncalls)
+            if tensorize_events:
+                events_to_do.append(int_me(ncalls))
+            else:
+                events_to_do.append(ncalls)
             events_left -= self.events_per_run
 
         if self.devices:

src/vegasflow/plain.py

@@ -2,7 +2,7 @@
     Plain implementation of the plainest possible MonteCarlo
 """
 
-from vegasflow.configflow import DTYPE, fone, fzero
+from vegasflow.configflow import fone, fzero
 from vegasflow.monte_carlo import MonteCarloFlow, wrapper, sampler
 import tensorflow as tf
 
@@ -38,9 +38,10 @@ def _run_iteration(self):
 
 
 def plain_wrapper(*args, **kwargs):
-    """ Wrapper around PlainFlow """
+    """Wrapper around PlainFlow"""
     return wrapper(PlainFlow, *args, **kwargs)
 
+
 def plain_sampler(*args, **kwargs):
-    """ Wrapper sampler around PlainFlow """
+    """Wrapper sampler around PlainFlow"""
     return sampler(PlainFlow, *args, **kwargs)

src/vegasflow/tests/test_algs.py

@@ -13,6 +13,7 @@
 from vegasflow.configflow import DTYPE
 from vegasflow.vflow import VegasFlow
 from vegasflow.plain import PlainFlow
+from vegasflow.vflowplus import VegasFlowPlus
 from vegasflow import plain_sampler, vegas_sampler
 
 # Test setup
@@ -66,6 +67,7 @@ def check_is_one(result, sigmas=3):
 
 
 def test_VegasFlow():
+    """ Test VegasFlow class, importance sampling algorithm"""
     for mode in range(4):
         vegas_instance = instance_and_compile(VegasFlow, mode)
         _ = vegas_instance.run_integration(n_iter)
@@ -90,6 +92,7 @@ def test_VegasFlow():
 
 
 def test_VegasFlow_save_grid():
+    """ Test the grid saving feature of vegasflow """
     tmp_filename = tempfile.mktemp()
     vegas_instance = instance_and_compile(VegasFlow)
     # Run an iteration so the grid is not trivial
@@ -173,3 +176,17 @@ def test_rng_generation(n_events=100):
     _ = helper_rng_tester(p.generate_random_array, n_events)
     v = vegas_sampler(example_integrand, dim, n_events, training_steps=2)
     _ = helper_rng_tester(v, n_events)
+
+def test_VegasFlowPlus_ADAPTIVE_SAMPLING():
+    """ Test Vegasflow with Adaptive Sampling on (the default) """
+    for mode in range(4):
+        vflowplus_instance = instance_and_compile(VegasFlowPlus, mode)
+        result = vflowplus_instance.run_integration(n_iter)
+        check_is_one(result)
+
+def test_VegasFlowPlus_NOT_ADAPTIVE_SAMPLING():
+    """ Test Vegasflow with Adaptive Sampling off (non-default) """
+    vflowplus_instance = VegasFlowPlus(dim, ncalls, adaptive=False)
+    vflowplus_instance.compile(example_integrand)
+    result = vflowplus_instance.run_integration(n_iter)
+    check_is_one(result)

src/vegasflow/utils.py

@@ -6,9 +6,13 @@
 import tensorflow as tf
 
 
-@tf.function(input_signature=[tf.TensorSpec(shape=[None], dtype=DTYPE),
-    tf.TensorSpec(shape=[None, None], dtype=DTYPEINT),
-    tf.TensorSpec(shape=[], dtype=DTYPEINT)]) 
+@tf.function(
+    input_signature=[
+        tf.TensorSpec(shape=[None], dtype=DTYPE),
+        tf.TensorSpec(shape=[None, None], dtype=DTYPEINT),
+        tf.TensorSpec(shape=[], dtype=DTYPEINT),
+    ]
+)
 def consume_array_into_indices(input_arr, indices, result_size):
     """
     Accumulate the input tensor `input_arr` into an output tensor of
@@ -38,6 +42,7 @@ def consume_array_into_indices(input_arr, indices, result_size):
     final_result = tf.reduce_sum(res_tmp, axis=1)
     return final_result
 
+
 def py_consume_array_into_indices(input_arr, indices, result_size):
     """
     Python interface wrapper for ``consume_array_into_indices``.
@@ -46,8 +51,8 @@ def py_consume_array_into_indices(input_arr, indices, result_size):
     return consume_array_into_indices(float_me(input_arr), int_me(indices), int_me(result_size))
 
 
-def generate_condition_function(n_mask, condition='and'):
-    """ Generates a function that takes a number of masks
+def generate_condition_function(n_mask, condition="and"):
+    """Generates a function that takes a number of masks
     and returns a combination of all n_masks for the given condition.
 
     It is possible to pass a list of allowed conditions, in that case
@@ -76,7 +81,7 @@ def generate_condition_function(n_mask, condition='and'):
     full_mask=<tf.Tensor: shape=(3,), dtype=bool, numpy=array([ True, False, False])>
     indices=<tf.Tensor: shape=(1, 1), dtype=int32, numpy=array([[0]], dtype=int32)>
 
-    
+
 
     Parameters
     ----------
@@ -90,16 +95,14 @@ def generate_condition_function(n_mask, condition='and'):
     `condition_to_idx`: function
         function(*masks) -> full_mask, true indices
     """
-    allowed_conditions = {
-            'and': tf.math.logical_and,
-            'or' : tf.math.logical_or
-            }
+    allowed_conditions = {"and": tf.math.logical_and, "or": tf.math.logical_or}
     allo = list(allowed_conditions.keys())
 
     # Check that the user is not asking for anything weird
     if n_mask < 2:
-        raise ValueError(f"At least two masks needed to generate a wrapper")
-    elif isinstance(condition, str):
+        raise ValueError("At least two masks needed to generate a wrapper")
+
+    if isinstance(condition, str):
         if condition not in allowed_conditions:
             raise ValueError(f"Wrong condition {condition}, allowed values are {allo}")
         is_list = False
@@ -112,7 +115,7 @@ def generate_condition_function(n_mask, condition='and'):
         is_list = True
 
     def py_condition(*masks):
-        """ Receives a list of conditions and returns a result mask
+        """Receives a list of conditions and returns a result mask
         and the list of indices in which the result mask is True
 
         Returns
@@ -125,15 +128,15 @@ def py_condition(*masks):
         if is_list:
             res = masks[0]
             for i, cond in enumerate(condition):
-                res = allowed_conditions[cond](res, masks[i+1])
-        elif condition == 'and':
+                res = allowed_conditions[cond](res, masks[i + 1])
+        elif condition == "and":
             res = tf.math.reduce_all(masks, axis=0)
-        elif condition == 'or':
+        elif condition == "or":
             res = tf.math.reduce_any(masks, axis=0)
         indices = int_me(tf.where(res))
         return res, indices
 
-    signature = n_mask*[tf.TensorSpec(shape=[None], dtype=tf.bool)]
+    signature = n_mask * [tf.TensorSpec(shape=[None], dtype=tf.bool)]
 
     condition_to_idx = tf.function(py_condition, input_signature=signature)
-    return condition_to_idx
+    return condition_to_idx