Fixed edge case in vegas

esa · Aug 20, 2021 · 6be0131 · 6be0131
1 parent 27d8fe1
commit 6be0131
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/torchquad/integration/vegas.py b/torchquad/integration/vegas.py
@@ -179,7 +179,8 @@ def _warmup_grid(self, warmup_N_it=5, N_samples=1000):
             jf = 0  # jacobians * function
             jf2 = 0
 
-            yrnd = torch.rand(size=[N_samples, self._dim])
+            # Multiplying by 0.99999999 as the edge case of y=1 leads to an error
+            yrnd = torch.rand(size=[N_samples, self._dim]) * 0.999999
             x = self.map.get_X(yrnd)
             f_eval = self._eval(x).squeeze()
             jac = self.map.get_Jac(yrnd)

diff --git a/torchquad/integration/vegas_stratification.py b/torchquad/integration/vegas_stratification.py
@@ -133,9 +133,9 @@ def get_Y(self, nevals):
         indices = self._get_indices(indices)
 
         # Get random numbers (we get a few more just to vectorize properly)
-        # This might increase the memorz requirements slightly but is probably
-        # worth it
-        random_uni = torch.rand(size=[len(nevals), nevals.max(), self.dim])
+        # This might increase the memory requirements slightly but is probably
+        # worth it.
+        random_uni = torch.rand(size=[len(nevals), nevals.max(), self.dim]) * 0.999999
 
         # Sum the random numbers onto the index locations and scale with dy
         # Note that the resulting tensor is still slightly too large