Repository: N3PDF/vegasflow

Issue:
Specifying integration limits
Hi vegasflow developers!

Can I specify integration (upper and lower) limits in vegasflow? I searched the documentation but didn't found a way.

Discussion:
Mulliken: Thanks for making this change @scarlehoff!

Patch:
diff --git a/doc/source/how_to.rst b/doc/source/how_to.rst
index 488d6cc..3bc787d 100644
--- a/doc/source/how_to.rst
+++ b/doc/source/how_to.rst
@@ -189,6 +189,22 @@ The full list of integration algorithms and wrappers can be consulted at: :ref:`
 Tips and Tricks
 ===============
 
+Changing the integration limits
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+By default ``VegasFlow`` provides random number only in the 0 to 1 range (and so all integrals are expected to be integrals from 0 to 1).
+But it is possible to choose any other ranges by passing to the initializer of the algorithm the ``xmin`` and ``xman`` variables.
+
+Note that if any limit is to be changed all ``xmin`` and ``xmax`` must be provided:
+
+.. code-block:: python
+
+    from vegasflow import VegasFlow
+
+    dimensions = 2
+    vegas_instance = VegasFlow(dimensions, n_calls, xmin=[0, -4], xmax=[1, 10])
+
+
 Seeding the random number generator
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
diff --git a/examples/simgauss_tf.py b/examples/simgauss_tf.py
index 4232a00..78fa172 100644
--- a/examples/simgauss_tf.py
+++ b/examples/simgauss_tf.py
@@ -9,6 +9,7 @@ import time
 import numpy as np
 import tensorflow as tf
 from vegasflow.vflow import vegas_wrapper
+from vegasflow import PlainFlow
 from vegasflow.plain import plain_wrapper 
 
 
@@ -35,13 +36,6 @@ if __name__ == "__main__":
     """Testing several different integrations"""
     print(f"VEGAS MC, ncalls={ncalls}:")
     start = time.time()
-    ncalls = 10*ncalls
     r = vegas_wrapper(symgauss, dim, n_iter, ncalls)
     end = time.time()
     print(f"Vegas took: time (s): {end-start}")
-
-#     print(f"Plain MC, ncalls={ncalls}:")
-#     start = time.time()
-#     r = plain_wrapper(symgauss, dim, n_iter, ncalls)
-#     end = time.time()
-#     print(f"Plain took: time (s): {end-start}")
diff --git a/src/vegasflow/monte_carlo.py b/src/vegasflow/monte_carlo.py
index f8bc328..029c8f3 100644
--- a/src/vegasflow/monte_carlo.py
+++ b/src/vegasflow/monte_carlo.py
@@ -49,6 +49,8 @@ from vegasflow.configflow import (
     TECH_CUT,
     float_me,
     int_me,
+    fone,
+    fzero,
 )
 
 
@@ -114,6 +116,8 @@ class MonteCarloFlow(ABC):
         events_limit=MAX_EVENTS_LIMIT,
         list_devices=DEFAULT_ACTIVE_DEVICES,  # pylint: disable=dangerous-default-value
         verbose=True,
+        xmin=None,
+        xmax=None,
         **kwargs,
     ):
         if "simplify_signature" in kwargs:
@@ -151,6 +155,24 @@ class MonteCarloFlow(ABC):
         else:
             self.devices = None
 
+        if xmin is not None or xmax is not None:
+            # If the ranges are provided, check that they are correct
+            if xmin is None or xmax is None:
+                raise ValueError(
+                    "Both xmin and xmax must be provided if the integration limits are to change"
+                )
+            if not (len(xmin) == len(xmax) == n_dim):
+                raise ValueError("The integration limits must be given for all dimensions")
+            self._xmin = float_me(xmin)
+            self._xdelta = float_me(xmax) - float_me(xmin)
+            if any(self._xdelta < 0.0):
+                raise ValueError(f"No xmin ({xmin}) can be bigger than xmax ({xmax})")
+            self._xdeltajac = tf.reduce_prod(self._xdelta)
+        else:
+            self._xmin = None
+            self._xdelta = None
+            self._xdeltajac = None
+
     # Note:
     # The number of events to run in a single iteration is `n_events`
     # while the total number of events to be run per step (so, for instance, per GPU call)
@@ -203,7 +225,7 @@ class MonteCarloFlow(ABC):
         """The default jacobian is 1 / total number of events"""
         return float_me([1.0 / self.n_events])
 
-    def generate_random_array(self, n_events):
+    def generate_random_array(self, n_events, *args):
         """External interface for the generation of random
         points as a 2D array of (n_events, n_dim).
         It calls the internal version of ``_generate_random_array``
@@ -216,16 +238,14 @@ class MonteCarloFlow(ABC):
         Returns
         -------
             `rnds`: array of (n_events, n_dim) random points
-            `idx` : index associated to each random point
             `p(x)` : p(x) associated to the random points
         """
-        rnds, idx, xjac_raw = self._generate_random_array(n_events)
-        # returns a p(x) corresponding to the number of events
-        # the algorithm was trained with, reweight
-        xjac = xjac_raw / self.xjac / n_events
-        return rnds, idx, xjac
+        rnds, xjac_raw, *extra = self._generate_random_array(n_events, *args)
+        # Since the n_events of this method might not be the "evaluation" value, reweight
+        xjac = xjac_raw / (self.xjac * n_events)
+        return rnds, xjac
 
-    def _generate_random_array(self, n_events):
+    def _generate_random_array(self, n_events, *args):
         """Generate a 2D array of (n_events, n_dim) points
         For the weight of the given point, this function is considered
         as part of an integration with ``self.n_events`` calls.
@@ -240,11 +260,18 @@ class MonteCarloFlow(ABC):
             `idx` : index associated to each random point
             `wgt` : wgt associated to the random point
         """
-        rnds = tf.random.uniform(
+        rnds_raw = tf.random.uniform(
             (n_events, self.n_dim), minval=TECH_CUT, maxval=1.0 - TECH_CUT, dtype=DTYPE
         )
-        idx = 0
-        return rnds, idx, self.xjac
+        # Now allow for the algorithm to produce the random numbers for the integration
+        rnds, wgts_raw, *extra = self._digest_random_generation(rnds_raw, *args)
+
+        wgts = wgts_raw * self.xjac
+        if self._xdelta is not None:
+            # Now apply integration limits
+            rnds = self._xmin + rnds * self._xdelta
+            wgts *= self._xdeltajac
+        return rnds, wgts, *extra
 
     #### Abstract methods
     @abstractmethod
@@ -259,6 +286,22 @@ class MonteCarloFlow(ABC):
         result = self.event()
         return result, pow(result, 2)
 
+    def _digest_random_generation(self, xrand, *args):
+        """All implemented algorithms will take a vector of uniform noise (n_events, n_dim)
+        and make it into a vector of random numbers (n_events, n_dim) with an associated weight.
+
+        It must return at least a tensor (n_events, n_dim) of random numbers
+        and of weights (n_events,) and can return any extra parameters
+        which will pass untouched by _generate_random_array
+        """
+        return xrand, 1.0  # , any extra param
+
+    def _apply_integration_limits(self, rand):
+        """Apply the integration limits (if any)
+        Receives a tensor of random numbers (n_events, n_dim) and returns
+        a transformed array (n_events, n_dim) and the associated jacobian (n_events,)
+        """
+
     def _can_run_vectorial(self, expected_shape=None):
         """Accepting vectorial integrands depends on the algorithm,
         if an algorithm can run on vectorial algorithms it should implement this method and return True"""
@@ -266,7 +309,7 @@ class MonteCarloFlow(ABC):
 
     #### Integration management
     def set_seed(self, seed):
-        """Sets the interation seed"""
+        """Sets the random seed"""
         tf.random.set_seed(seed)
 
     #### Device management methods
@@ -345,7 +388,7 @@ class MonteCarloFlow(ABC):
         """Modifies the attributes of the integration so that it can be compiled inside
         Tensorflow functions (and, therefore, gradients calculated)
         Returns a reference to `run_event`, a method that upon calling it with no arguments
-        will produce results and uncertainties for an intergation iteration of ncalls number of events
+        will produce results and uncertainties for an integration iteration of ncalls number of events
         """
         if self.distribute:
             raise ValueError("Differentiation is not compatible with distribution")
@@ -623,7 +666,7 @@ if you believe this to be a bug please open an issue in https://github.com/N3PDF
                 monte carlo error
 
         Note: it is possible not to pass any histogram variable and still fill
-        some histogram variable at integration time, but then it is the responsability
+        some histogram variable at integration time, but then it is the responsibility
         of the integrand to empty the histograms each iteration and accumulate them.
 
         """
diff --git a/src/vegasflow/plain.py b/src/vegasflow/plain.py
index 1ddd762..d9fd081 100644
--- a/src/vegasflow/plain.py
+++ b/src/vegasflow/plain.py
@@ -21,7 +21,7 @@ class PlainFlow(MonteCarloFlow):
             n_events = ncalls
 
         # Generate all random number for this iteration
-        rnds, _, xjac = self._generate_random_array(n_events)
+        rnds, xjac = self._generate_random_array(n_events)
 
         # Compute the integrand
         tmp = integrand(rnds, weight=xjac) * xjac
diff --git a/src/vegasflow/vflow.py b/src/vegasflow/vflow.py
index a400098..48323c8 100644
--- a/src/vegasflow/vflow.py
+++ b/src/vegasflow/vflow.py
@@ -1,5 +1,5 @@
 """
-    This module contains the VegasFlow class and all its auxuliary functions
+    This module contains the VegasFlow class and all its auxiliary functions
 
     The main interfaces of this class are the class `VegasFlow` and the
     `vegas_wrapper`
@@ -103,10 +103,10 @@ def _generate_random_array(rnds, divisions):
     -------
         x: array (None, n_dim)
             Vegas random output
-        div_index: array (None, n_dim)
-            division index in which each (n_dim) set of random numbers fall
         w: array (None,)
             Weight of each set of (n_dim) random numbers
+        div_index: array (None, n_dim)
+            division index in which each (n_dim) set of random numbers fall
     """
     # Get the boundaries of the random numbers
     #     reg_i = fzero
@@ -121,7 +121,7 @@ def _generate_random_array(rnds, divisions):
     # Compute the random number between the limits
     # commented, for now only from 0 to 1
     #     x = reg_i + rand_x * (reg_f - reg_i)
-    return x, ind_xn, weights
+    return x, weights, ind_xn
 
 
 @tf.function(
@@ -362,12 +362,11 @@ class VegasFlow(MonteCarloFlow):
             new_divisions = refine_grid_per_dimension(arr_res2[j, :], self.divisions[j, :])
             self.divisions[j, :].assign(new_divisions)
 
-    def _generate_random_array(self, n_events):
-        """Uses the internal array to generate ``n_events`` random numbers"""
-        rnds, _, xjac = super()._generate_random_array(n_events)
-        # Pass them through the Vegas digestion
-        x, ind, w = _generate_random_array(rnds, self.divisions)
-        return x, ind, w * xjac
+    def _digest_random_generation(self, rnds):
+        """Generates ``n_events`` random numbers sampled in the
+        adapted Vegas Grid"""
+        x, w, ind = _generate_random_array(rnds, self.divisions)
+        return x, w, ind
 
     def _importance_sampling_array_filling(self, results2, indices):
         """Receives an array of results squared for every event
@@ -408,7 +407,7 @@ class VegasFlow(MonteCarloFlow):
             n_events = ncalls
 
         # Generate all random number for this iteration
-        x, ind, xjac = self._generate_random_array(n_events)
+        x, xjac, ind = self._generate_random_array(n_events)
 
         # Now compute the integrand
         int_result = integrand(x, weight=xjac)
@@ -422,7 +421,7 @@ class VegasFlow(MonteCarloFlow):
         res = tf.reduce_sum(tmp, axis=0)
         res2 = tf.reduce_sum(tmp2, axis=0)
 
-        # If this is a vectorial integrnad, make sure that only the main dimenison
+        # If this is a vectorial integrand, make sure that only the main dimension
         # is used for the grid training
         if self._vectorial:
             tmp2 = tmp2[:, self._main_dimension]
diff --git a/src/vegasflow/vflowplus.py b/src/vegasflow/vflowplus.py
index e3fdb58..57e228f 100644
--- a/src/vegasflow/vflowplus.py
+++ b/src/vegasflow/vflowplus.py
@@ -60,9 +60,9 @@ def generate_samples_in_hypercubes(rnds, n_strat, n_ev, hypercubes, divisions):
         `x` : random numbers collocated in hypercubes
         `w` : weight of each event
         `ind`: division index in which each (n_dim) set of random numbers fall
-        `segm` : segmentantion for later computations
+        `segm` : segmentation for later computations
     """
-    # Use the event-per-hypercube information to fix each random event to a hypercub
+    # Use the event-per-hypercube information to fix each random event to a hypercube
     indices = tf.repeat(tf.range(tf.shape(hypercubes, out_type=DTYPEINT)[0]), n_ev)
     points = float_me(tf.gather(hypercubes, indices))
     n_evs = float_me(tf.gather(n_ev, indices))
@@ -72,11 +72,11 @@ def generate_samples_in_hypercubes(rnds, n_strat, n_ev, hypercubes, divisions):
 
     ind_xn, x, weights = importance_sampling_digest(xn, divisions)
 
-    # Reweight taking into account the number of events per hypercub
+    # Reweighs taking into account the number of events per hypercube
     final_weights = weights / n_evs
 
     segm = indices
-    return x, ind_xn, final_weights, segm
+    return x, final_weights, ind_xn, segm
 
 
 class VegasFlowPlus(VegasFlow):
@@ -135,11 +135,15 @@ class VegasFlowPlus(VegasFlow):
         self.n_ev = tf.fill([1, len(hypercubes)], self.min_neval_hcube)
         self.n_ev = int_me(tf.reshape(self.n_ev, [-1]))
         self._n_events = int(tf.reduce_sum(self.n_ev))
-        self.my_xjac = float_me(1 / len(hypercubes))
+        self._modified_jac = float_me(1 / len(hypercubes))
 
         if self._adaptive:
             logger.warning("Variable number of events requires function signatures all across")
 
+    @property
+    def xjac(self):
+        return self._modified_jac
+
     def make_differentiable(self):
         """Overrides make_differentiable to make sure the runner has a reference to n_ev"""
         runner = super().make_differentiable()
@@ -157,24 +161,31 @@ class VegasFlowPlus(VegasFlow):
         self.n_ev = int_me(new_n_ev)
         self.n_events = int(tf.reduce_sum(self.n_ev))
 
-    def _generate_random_array(self, n_events):
-        """Interface compatible with other algorithms dropping the segmentation in hypercubes"""
-        x, ind, w, _ = self._generate_random_array_plus(n_events, self.n_ev)
-        return x, ind, w
-
-    def _generate_random_array_plus(self, n_events, n_ev):
+    def _digest_random_generation(self, rnds, n_ev):
         """Generate a random array for a given number of events divided in hypercubes"""
-        # Needs to skip parent and go directly to the random array generation of MonteCarloFlow
-        rnds, _, _ = MonteCarloFlow._generate_random_array(self, n_events)
         # Get random numbers from hypercubes
-        x, ind, w, segm = generate_samples_in_hypercubes(
+        x, w, ind, segm = generate_samples_in_hypercubes(
             rnds,
             self._n_strat,
             n_ev,
             self._hypercubes,
             self.divisions,
         )
-        return x, ind, w * self.my_xjac, segm
+        return x, w, ind, segm
+
+    def generate_random_array(self, n_events, *args):
+        """Override the behaviour of ``generate_random_array``
+        to accomodate for the peculiarities of VegasFlowPlus
+        """
+        rnds = []
+        wgts = []
+        for _ in range(n_events // self.n_events + 1):
+            r, w = super().generate_random_array(self.n_events, self.n_ev)
+            rnds.append(r)
+            wgts.append(w)
+        final_r = tf.concat(rnds, axis=0)[:n_events]
+        final_w = tf.concat(wgts, axis=0)[:n_events] * self.n_events / n_events
+        return final_r, final_w
 
     def _run_event(self, integrand, ncalls=None, n_ev=None):
         """Run one step of VegasFlowPlus
@@ -190,12 +201,12 @@ class VegasFlowPlus(VegasFlow):
 
         Returns
         -------
-            `res`: sum of the result of the integrand for all events per segement
+            `res`: sum of the result of the integrand for all events per segment
             `res2`: sum of the result squared of the integrand for all events per segment
             `arr_res2`: result of the integrand squared per dimension and grid bin
         """
         # NOTE: needs to receive both ncalls and n_ev
-        x, ind, xjac, segm = self._generate_random_array_plus(ncalls, n_ev)
+        x, xjac, ind, segm = self._generate_random_array(ncalls, n_ev)
 
         # compute integrand
         tmp = xjac * integrand(x, weight=xjac)


Test patch:
diff --git a/src/vegasflow/tests/test_algs.py b/src/vegasflow/tests/test_algs.py
index caec02b..9d18531 100644
--- a/src/vegasflow/tests/test_algs.py
+++ b/src/vegasflow/tests/test_algs.py
@@ -58,12 +58,12 @@ def instance_and_compile(Integrator, mode=0, integrand_function=example_integran
     return int_instance
 
 
-def check_is_one(result, sigmas=3):
+def check_is_one(result, sigmas=3, target_result=1.0):
     """Wrapper for convenience"""
     res = result[0]
     err = np.mean(result[1] * sigmas)
     # Check that it passes by {sigmas} number of sigmas
-    np.testing.assert_allclose(res, 1.0, atol=err)
+    np.testing.assert_allclose(res, target_result, atol=err)
 
 
 @pytest.mark.parametrize("mode", range(4))
@@ -166,21 +166,38 @@ def test_PlainFlow_change_nevents():
 def helper_rng_tester(sampling_function, n_events):
     """Ensure the random number generated have the correct shape
     Return the random numbers and the jacobian"""
-    rnds, _, px = sampling_function(n_events)
+    rnds, px = sampling_function(n_events)
     np.testing.assert_equal(rnds.shape, (n_events, dim))
     return rnds, px
 
 
-def test_rng_generation(n_events=100):
-    """Test that the random generation genrates the correct type of arrays"""
+def test_rng_generation_plain(n_events=100):
+    """Test the random number generation with plainflow"""
     plain_sampler_instance = instance_and_compile(PlainFlow)
     _, px = helper_rng_tester(plain_sampler_instance.generate_random_array, n_events)
     np.testing.assert_equal(px.numpy(), 1.0 / n_events)
+
+
+def test_rng_generation_vegasflow(n_events=100):
+    """Test the random number generation with vegasflow"""
     vegas_sampler_instance = instance_and_compile(VegasFlow)
+    # Train a bit the grid
     vegas_sampler_instance.run_integration(2)
     _, px = helper_rng_tester(vegas_sampler_instance.generate_random_array, n_events)
     np.testing.assert_equal(px.shape, (n_events,))
-    # Test the wrappers
+
+
+def test_rng_generation_vegasflowplus(n_events=100):
+    """Test the random number generation with vegasflow"""
+    vegas_sampler_instance = instance_and_compile(VegasFlowPlus)
+    # Train a bit the grid
+    #     vegas_sampler_instance.run_integration(2)
+    _, px = helper_rng_tester(vegas_sampler_instance.generate_random_array, n_events)
+    np.testing.assert_equal(px.shape, (n_events,))
+
+
+def test_rng_generation_wrappers(n_events=100):
+    """Test the wrappers for the samplers"""
     p = plain_sampler(example_integrand, dim, n_events, training_steps=2, return_class=True)
     _ = helper_rng_tester(p.generate_random_array, n_events)
     v = vegas_sampler(example_integrand, dim, n_events, training_steps=2)
diff --git a/src/vegasflow/tests/test_misc.py b/src/vegasflow/tests/test_misc.py
index c913ffd..62d50ed 100644
--- a/src/vegasflow/tests/test_misc.py
+++ b/src/vegasflow/tests/test_misc.py
@@ -2,6 +2,7 @@
     Miscellaneous tests that don't really fit anywhere else
 """
 import pytest
+import numpy as np
 
 from vegasflow import VegasFlow, VegasFlowPlus, PlainFlow
 import tensorflow as tf
@@ -19,6 +20,18 @@ def _wrong_integrand(xarr):
     return tf.reduce_sum(xarr)
 
 
+def _simple_integrand(xarr):
+    """Integrand f(x) = x"""
+    return tf.reduce_prod(xarr, axis=1)
+
+
+def _simple_integral(xmin, xmax):
+    """Integated version of simple_ingrand"""
+    xm = np.array(xmin) ** 2 / 2.0
+    xp = np.array(xmax) ** 2 / 2.0
+    return np.prod(xp - xm)
+
+
 def _wrong_vector_integrand(xarr):
     """Vector integrand with the wrong output shape"""
     return tf.transpose(xarr)
@@ -30,7 +43,7 @@ def test_working_vectorial(alg, mode):
     """Check that the algorithms that accept integrating vectorial functions can really do so"""
     inst = instance_and_compile(alg, mode=mode, integrand_function=_vector_integrand)
     result = inst.run_integration(2)
-    check_is_one(result, sigmas=4)
+    check_is_one(result, sigmas=5)
 
 
 @pytest.mark.parametrize("alg", [VegasFlowPlus])
@@ -53,3 +66,31 @@ def test_wrong_shape(wrong_fun):
     """Check that an error is raised by the compilation if the integrand has the wrong shape"""
     with pytest.raises(ValueError):
         _ = instance_and_compile(PlainFlow, integrand_function=wrong_fun)
+
+
+@pytest.mark.parametrize("alg", [PlainFlow, VegasFlow, VegasFlowPlus])
+def test_integration_limits(alg, ncalls=int(1e4)):
+    """Test an integration where the integration limits are modified"""
+    dims = np.random.randint(1, 5)
+    xmin = -1.0 + np.random.rand(dims) * 2.0
+    xmax = 3.0 + np.random.rand(dims)
+    inst = alg(dims, ncalls, xmin=xmin, xmax=xmax)
+    inst.compile(_simple_integrand)
+    result = inst.run_integration(5)
+    expected_result = _simple_integral(xmin, xmax)
+    check_is_one(result, target_result=expected_result)
+
+
+def test_integration_limits_checks():
+    """Test that the errors for wrong limits actually work"""
+    # use hypothesis to check other corner cases
+    with pytest.raises(ValueError):
+        PlainFlow(1, 10, xmin=[10], xmax=[1])
+    with pytest.raises(ValueError):
+        PlainFlow(1, 10, xmin=[10])
+    with pytest.raises(ValueError):
+        PlainFlow(1, 10, xmax=[10])
+    with pytest.raises(ValueError):
+        PlainFlow(2, 10, xmin=[0], xmax=[1])
+    with pytest.raises(ValueError):
+        PlainFlow(2, 10, xmin=[0, 1], xmax=[1])