Merge pull request #74 from N3PDF/trace_option

Add a signature discovery and trace option

Author: scarlehoff

doc/source/examples.rst

@@ -174,7 +174,8 @@ by using the CFFI library as you can see in `this <https://github.com/N3PDF/vega
     # Now you can read up the compiled C code as a python library
     from _integrand_cffi import ffi, lib
     
-    def integrand(xarr, n_dim, **kwargs):
+    def integrand(xarr, **kwargs):
+        n_dim = xarr.shape[-1]
         result = np.empty(n_events, dtype=DTYPE.as_numpy_dtype)
         x_flat = xarr.numpy().flatten()
         

doc/source/how_to.rst

@@ -44,7 +44,7 @@ using the ``tf.function`` decorator.
 
     import tensorflow as tf
     
-    def example_integrand(xarr, n_dim=None, weight=None):
+    def example_integrand(xarr, weight=None):
       s = tf.reduce_sum(xarr, axis=1)
       result = tf.pow(0.1/s, 2)
       return result
@@ -76,17 +76,16 @@ the integration algorithm).
 
     import numpy as np
     
-    def example_integrand(xarr, n_dim=None, weight=None):
-      s = np.pow(xarr, axis=1)
+    def example_integrand(xarr, weight=None):
+      s = np.sum(xarr, axis=1)
       result = np.square(0.1/s)
       return result
       
     vegas_instance.compile(example_integrand, compilable=False)
 
-
-.. note:: Integrands must accept the keyword arguments ``n_dim`` and ``weight``, as the integrators
-   will try to pass those argument when convenient. It is however possible to construct integrands
-   that accept only the array of random number ``xarr``, see :ref:`simple-label`
+.. note:: Integrands must always accept as first argument the random number (``xarr``)
+  and can also accept the keyword argument ``weight``. The ``compile`` method of the integration
+   will try to find the most adequate signature in each situation.
 
 
 It is also possible to completely avoid compilation,
@@ -112,7 +111,7 @@ These functions are wrappers around ``tf.cast`` `🔗 <https://www.tensorflow.or
     
     constant = float_me(0.1)
     
-    def example_integrand(xarr, n_dim=None, weight=None):
+    def example_integrand(xarr, weight=None):
       s = tf.reduce_sum(xarr, axis=1)
       result = tf.pow(constant/s, 2)
       return result
@@ -141,42 +140,6 @@ The full list of integration algorithms and wrappers can be consulted at: :ref:`
 Tips and Tricks
 ===============
 
-.. _simple-label:
-
-Improving results by simplifying the integrand
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-In the above example the integrand receives the keyword arguments `n_dim` and `xjac`.
-Although these are useful for instance for filling histograms or dimensional-dependent integrands,
-these extra argument can harm the performance of the integration when they are not being used.
-
-It is possible to instantiate ``VegasFlow`` algorithms with ``simplify_signature``.
-In this case the integrand will only receive the array of random numbers and, in exchange for this
-loss of flexibility, the function will be retraced less often.
-For more details in what function retracing entails we direct you to the `TensorFlow documentation <https://www.tensorflow.org/api_docs/python/tf/function>`_.
-
-.. code-block:: python
-
-    from vegasflow import VegasFlow, float_me
-    import tensorflow as tf
-
-    def example_integrand(xarr):
-        c = float_me(0.1)
-        s = tf.reduce_sum(xarr)
-        result = tf.pow(c/s)
-        return result
-
-    dimensions = 3
-    n_calls = int(1e7)
-    # Create an instance of the VegasFlow class
-    vegas_instance = VegasFlow(dimensions, n_calls, simplify_signature = True)
-    # Compile the function to be integrated
-    vegas_instance.compile(example_integrand)
-    # Compute the result after a number of iterations
-    n_iter = 5
-    result = vegas_instance.run_integration(n_iter)
-    
-
 Seeding the random number generator
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
@@ -396,7 +359,7 @@ This is a crucial step (and the only fixed step) as this tensor will be accumula
         cummulator_tensor.assign(new_histograms)
 
     @tf.function
-    def integrand_example(xarr, n_dim=None, weight=fone):
+    def integrand_example(xarr, weight=fone):
         # some complicated calculation that generates 
         # a final_result and some histogram values:
         final_result = tf.constant(42, dtype=tf.float64)

doc/source/intalg.rst

@@ -35,7 +35,10 @@ Once that is generated it is possible to register an integrand by calling the ``
 .. code-block:: python
 
     def example_integrand(x, **kwargs):
-        return x
+        y = 0.0
+        for d in range(dims):
+            y += x[:,d]
+        return y
 
     vegas_instance.compile(example_integrand)
 
@@ -44,6 +47,7 @@ After each iteration the grid will be refined, producing more points (and hence
 
 .. code-block:: python
 
+    n_iter = 3
     result = vegas_instance.run_integration(n_iter)
 
 The output variable, in this example named ``result``, is a tuple variable where the first element is the result of the integration while the second element is the error of the integration.
@@ -103,7 +107,7 @@ The usage pattern is similar to :ref:`vegas-label`.
 .. code-block:: python
 
     from vegasflow import PlainFlow
-    plain_instance = PlainFlow(dimensions, ncalls)
+    plain_instance = PlainFlow(dims, n_calls)
     plain_instance.compile(example_integrand)
     plain_instance.run_integration(n_iter)
 

examples/asian_options_tf.py

@@ -29,7 +29,7 @@
 
 
 @tf.function
-def example_integrand(xarr, n_dim=None, **kwargs):
+def example_integrand(xarr, **kwargs):
     """Asian options test function"""
     sum1 = tf.reduce_sum(ndtri(xarr), axis=1)
     a = S0 * tf.exp((r-sigma2/2) + sigma*sqrtdt*sum1)

examples/basic_example.py

@@ -1,7 +1,7 @@
 """
     Example: basic integration
 
-    Very basic example with a simple integrand, uses simplify_signature
+    Very basic example with a simple integrand
 """
 
 from vegasflow import VegasFlow, float_me, run_eager
@@ -32,8 +32,11 @@ def symgauss(xarr):
     """Testing several different integrations"""
     print(f"VEGAS MC, ncalls={ncalls}:")
     start = time.time()
-    vegas_instance = VegasFlow(dim, ncalls, simplify_signature=True, events_limit=int(7e4))
+    vegas_instance = VegasFlow(dim, ncalls, events_limit=int(7e4))
     vegas_instance.compile(symgauss)
     result = vegas_instance.run_integration(n_iter)
     end = time.time()
     print(f"Vegas took: time (s): {end-start}")
+    print("Change the number of events and freeze the grid...")
+    vegas_instance.freeze_grid()
+    vegas_instance.run_integration(n_iter)

examples/drellyan_lo_tf.py

@@ -240,7 +240,7 @@ def build_luminosity(x1, x2):
 
 
 @tf.function
-def drellyan(xarr, **kwargs):
+def drellyan(xarr, weight=None, **kwargs):
     """Single-top (t-channel) at LO"""
     psw, flux, p0, p1, p2, p3, x1, x2 = make_event(xarr)
     wgts = evaluate_matrix_element_square(p0, p1, p2, p3)

examples/example_eager.py

@@ -20,10 +20,9 @@
 
 
 @tf.function
-def symgauss_sigmoid(xarr, n_dim=None, **kwargs):
+def symgauss_sigmoid(xarr, **kwargs):
     """symgauss test function"""
-    if n_dim is None:
-        n_dim = xarr.shape[-1]
+    n_dim = xarr.shape[-1]
     a = 0.1
     pref = pow(1.0 / a / np.sqrt(np.pi), n_dim)
     coef = np.sum(np.arange(1, 101))

examples/example_pineappl.py

@@ -3,15 +3,15 @@
 import random as rn
 from multiprocessing.pool import ThreadPool as Pool
 from functools import partial
+from vegasflow.configflow import DTYPE, MAX_EVENTS_LIMIT, run_eager
+run_eager(True)
 from pdfflow.pflow import mkPDF
 import pineappl
 from vegasflow.utils import generate_condition_function
 from vegasflow.vflow import VegasFlow
-from vegasflow.configflow import DTYPE, MAX_EVENTS_LIMIT
 import time
 import numpy as np
 import tensorflow as tf
-tf.config.run_functions_eagerly(True)
 
 
 parser = argparse.ArgumentParser()
@@ -82,7 +82,7 @@ def fill(grid, x1, x2, q2, yll, weight):
     grid.fill_array(x1, x2, q2, zeros, yll, zeros, weight)
 
 
-def fill_grid(xarr, n_dim=None, **kwargs):
+def fill_grid(xarr, weight=1, **kwargs):
     s, t, u, x1, x2, jacobian = hadronic_pspgen(xarr, 10.0, 7000.0)
 
     ptl = tf.sqrt((t * u / s))
@@ -102,18 +102,18 @@ def fill_grid(xarr, n_dim=None, **kwargs):
     t_6 = mll <= 120.0
     full_mask, indices = cuts(t_1, t_2, t_3, t_4, t_5, t_6)
 
-    weight = tf.boolean_mask(jacobian * int_photo(s, u, t), full_mask, axis=0)
+    wgt = tf.boolean_mask(jacobian * int_photo(s, u, t), full_mask, axis=0)
     x1 = tf.boolean_mask(x1, full_mask, axis=0)
     x2 = tf.boolean_mask(x2, full_mask, axis=0)
     yll = tf.boolean_mask(yll, full_mask, axis=0)
-    vweight = weight * tf.boolean_mask(kwargs.get('weight'), full_mask, axis=0)
+    vweight = wgt  * tf.boolean_mask(weight, full_mask, axis=0)
 
     if kwargs.get('fill_pineappl'):
         q2 = 90.0 * 90.0 * tf.ones(weight.shape, dtype=tf.float64)
         kwargs.get('pool').apply_async(fill, [kwargs.get('grid'), x1.numpy(), x2.numpy(),
                                               q2.numpy(), tf.abs(yll).numpy(), vweight.numpy()])
 
-    return tf.scatter_nd(indices, weight, shape=xarr.shape[0:1])
+    return tf.scatter_nd(indices, wgt, shape=xarr.shape[0:1])
 
 
 if __name__ == "__main__":

examples/histogram_ex.py

@@ -19,41 +19,34 @@
 
 
 def generate_integrand(cummulator_tensor):
-    """ 
+    """
     This function will generate an integrand function
     which will already hold a reference to the tensor to accumulate
     """
 
     @tf.function
     def histogram_collector(results, variables):
-        """ This function will receive a tensor (result)
-        and the variables corresponding to those integrand results 
-        In the example integrand below, these corresponds to 
+        """This function will receive a tensor (result)
+        and the variables corresponding to those integrand results
+        In the example integrand below, these corresponds to
             `final_result` and `histogram_values` respectively.
         `current_histograms` instead is the current value of the histogram
-        which will be overwritten """
+        which will be overwritten"""
         # Fill a histogram with HISTO_BINS (2) bins, (0 to 0.5, 0.5 to 1)
         # First generate the indices with TF
-        indices = tf.histogram_fixed_width_bins(
-            variables, [fzero, fone], nbins=HISTO_BINS
-        )
+        indices = tf.histogram_fixed_width_bins(variables, [fzero, fone], nbins=HISTO_BINS)
         t_indices = tf.transpose(indices)
         # Then consume the results with the utility we provide
         partial_hist = consume_array_into_indices(results, t_indices, HISTO_BINS)
         # Then update the results of current_histograms
         new_histograms = partial_hist + current_histograms
         cummulator_tensor.assign(new_histograms)
 
-    def integrand_example(xarr, n_dim=None, weight=fone):
-        """ Example of function which saves histograms """
-        if n_dim is None:
-            n_dim = xarr.shape[-1]
-        if n_dim < hst_dim:
-            raise ValueError(
-                f"The number of dimensions has to be greater than {hst_dim} for this example"
-            )
+    def integrand_example(xarr, weight=fone):
+        """Example of function which saves histograms"""
+        n_dim = xarr.shape[-1]
         a = tf.constant(0.1, dtype=DTYPE)
-        n100 = tf.cast(100 * n_dim, dtype=DTYPE)
+        n100 = tf.cast(100 * dim, dtype=DTYPE)
         pref = tf.pow(1.0 / a / np.sqrt(np.pi), n_dim)
         coef = tf.reduce_sum(tf.range(n100 + 1))
         coef += tf.reduce_sum(tf.square((xarr - 1.0 / 2.0) / a), axis=1)
@@ -69,6 +62,10 @@ def integrand_example(xarr, n_dim=None, weight=fone):
 
 if __name__ == "__main__":
     """Testing histogram generation"""
+    if dim < hst_dim:
+        raise ValueError(
+            f"The number of dimensions has to be greater than {hst_dim} for this example"
+        )
     print(f"Plain MC, ncalls={ncalls}:")
     start = time.time()
     # First we create the tensor in which to accumulate the histogra

examples/simgauss_cffi.py

@@ -49,9 +49,8 @@
 
 from _symgauss_cffi import ffi, lib
 
-def symgauss(xarr, n_dim=None, **kwargs):
-    if n_dim is None:
-        n_dim = xarr.shape[-1]
+def symgauss(xarr, **kwargs):
+    n_dim = xarr.shape[-1]
     n_events = xarr.shape[0]
 
     res = np.empty(n_events, dtype = DTYPE.as_numpy_dtype)