update readme

Author: houpc

MCintegration/base.py

@@ -9,10 +9,8 @@
 from MCintegration.utils import get_device
 
 # Constants for numerical stability
-# Small but safe non-zero value
-MINVAL = 10 ** (sys.float_info.min_10_exp + 50)
-MAXVAL = 10 ** (sys.float_info.max_10_exp - 50)  # Large but safe value
 EPSILON = 1e-16  # Small value to ensure numerical stability
+# EPSILON = sys.float_info.epsilon * 1e4  # Small value to ensure numerical stability
 
 
 class BaseDistribution(nn.Module):
@@ -98,10 +96,8 @@ def sample(self, batch_size=1, **kwargs):
             tuple: (uniform samples, log_det_jacobian=0)
         """
         # torch.manual_seed(0) # test seed
-        u = torch.rand((batch_size, self.dim),
-                       device=self.device, dtype=self.dtype)
-        log_detJ = torch.zeros(
-            batch_size, device=self.device, dtype=self.dtype)
+        u = torch.rand((batch_size, self.dim), device=self.device, dtype=self.dtype)
+        log_detJ = torch.zeros(batch_size, device=self.device, dtype=self.dtype)
         return u, log_detJ
 
 
@@ -133,16 +129,14 @@ def __init__(self, A, b, device=None, dtype=torch.float32):
         elif isinstance(A, torch.Tensor):
             self.A = A.to(dtype=self.dtype, device=self.device)
         else:
-            raise ValueError(
-                "'A' must be a list, numpy array, or torch tensor.")
+            raise ValueError("'A' must be a list, numpy array, or torch tensor.")
 
         if isinstance(b, (list, np.ndarray)):
             self.b = torch.tensor(b, dtype=self.dtype, device=self.device)
         elif isinstance(b, torch.Tensor):
             self.b = b.to(dtype=self.dtype, device=self.device)
         else:
-            raise ValueError(
-                "'b' must be a list, numpy array, or torch tensor.")
+            raise ValueError("'b' must be a list, numpy array, or torch tensor.")
 
         # Pre-compute determinant of Jacobian for efficiency
         self._detJ = torch.prod(self.A)

MCintegration/maps.py

@@ -9,7 +9,8 @@
 from MCintegration.utils import get_device
 import sys
 
-TINY = 10 ** (sys.float_info.min_10_exp + 50)  # Small but safe non-zero value
+# TINY = 10 ** (sys.float_info.min_10_exp + 50)  # Small but safe non-zero value
+TINY = 1e-45
 
 
 class Configuration:
@@ -38,14 +39,10 @@ def __init__(self, batch_size, dim, f_dim, device=None, dtype=torch.float32):
         self.f_dim = f_dim
         self.batch_size = batch_size
         # Initialize tensors for storing samples and results
-        self.u = torch.empty(
-            (batch_size, dim), dtype=dtype, device=self.device)
-        self.x = torch.empty(
-            (batch_size, dim), dtype=dtype, device=self.device)
-        self.fx = torch.empty((batch_size, f_dim),
-                              dtype=dtype, device=self.device)
-        self.weight = torch.empty(
-            (batch_size,), dtype=dtype, device=self.device)
+        self.u = torch.empty((batch_size, dim), dtype=dtype, device=self.device)
+        self.x = torch.empty((batch_size, dim), dtype=dtype, device=self.device)
+        self.fx = torch.empty((batch_size, f_dim), dtype=dtype, device=self.device)
+        self.weight = torch.empty((batch_size,), dtype=dtype, device=self.device)
         self.detJ = torch.empty((batch_size,), dtype=dtype, device=self.device)
 
 
@@ -202,8 +199,7 @@ def __init__(self, dim, ninc=1000, device=None, dtype=torch.float32):
                 (self.dim,), ninc, dtype=torch.int32, device=self.device
             )
         elif isinstance(ninc, (list, np.ndarray)):
-            self.ninc = torch.tensor(
-                ninc, dtype=torch.int32, device=self.device)
+            self.ninc = torch.tensor(ninc, dtype=torch.int32, device=self.device)
         elif isinstance(ninc, torch.Tensor):
             self.ninc = ninc.to(dtype=torch.int32, device=self.device)
         else:
@@ -223,8 +219,9 @@ def __init__(self, dim, ninc=1000, device=None, dtype=torch.float32):
         self.sum_f = torch.zeros(
             self.dim, self.max_ninc, dtype=self.dtype, device=self.device
         )
-        self.n_f = torch.zeros(
-            self.dim, self.max_ninc, dtype=self.dtype, device=self.device
+        self.n_f = (
+            torch.zeros(self.dim, self.max_ninc, dtype=self.dtype, device=self.device)
+            + TINY
         )
         self.avg_f = torch.ones(
             (self.dim, self.max_ninc), dtype=self.dtype, device=self.device
@@ -308,19 +305,16 @@ def adapt(self, alpha=0.5):
         """
         # Aggregate training data across distributed processes if applicable
         if torch.distributed.is_initialized():
-            torch.distributed.all_reduce(
-                self.sum_f, op=torch.distributed.ReduceOp.SUM)
-            torch.distributed.all_reduce(
-                self.n_f, op=torch.distributed.ReduceOp.SUM)
+            torch.distributed.all_reduce(self.sum_f, op=torch.distributed.ReduceOp.SUM)
+            torch.distributed.all_reduce(self.n_f, op=torch.distributed.ReduceOp.SUM)
 
         # Initialize a new grid tensor
         new_grid = torch.empty(
             (self.dim, self.max_ninc + 1), dtype=self.dtype, device=self.device
         )
 
         if alpha > 0:
-            tmp_f = torch.empty(
-                self.max_ninc, dtype=self.dtype, device=self.device)
+            tmp_f = torch.empty(self.max_ninc, dtype=self.dtype, device=self.device)
 
         # avg_f = torch.ones(self.inc.shape[1], dtype=self.dtype, device=self.device)
         # print(self.ninc.shape, self.dim)
@@ -338,14 +332,12 @@ def adapt(self, alpha=0.5):
 
                 if alpha > 0:
                     # Smooth avg_f
-                    tmp_f[0] = (7.0 * avg_f[0] + avg_f[1]
-                                ).abs() / 8.0  # Shape: ()
+                    tmp_f[0] = (7.0 * avg_f[0] + avg_f[1]).abs() / 8.0  # Shape: ()
                     tmp_f[ninc - 1] = (
                         7.0 * avg_f[ninc - 1] + avg_f[ninc - 2]
                     ).abs() / 8.0  # Shape: ()
-                    tmp_f[1: ninc - 1] = (
-                        6.0 * avg_f[1: ninc - 1] +
-                        avg_f[: ninc - 2] + avg_f[2:ninc]
+                    tmp_f[1 : ninc - 1] = (
+                        6.0 * avg_f[1 : ninc - 1] + avg_f[: ninc - 2] + avg_f[2:ninc]
                     ).abs() / 8.0
 
                     # Normalize tmp_f to ensure the sum is 1
@@ -393,8 +385,7 @@ def adapt(self, alpha=0.5):
                 ) / avg_f_relevant
 
                 # Calculate the new grid points using vectorized operations
-                new_grid[d, 1:ninc] = grid_left + \
-                    fractional_positions * inc_relevant
+                new_grid[d, 1:ninc] = grid_left + fractional_positions * inc_relevant
             else:
                 # If alpha == 0 or no training data, retain the existing grid
                 new_grid[d, :] = self.grid[d, :]
@@ -407,8 +398,7 @@ def adapt(self, alpha=0.5):
         self.inc.zero_()  # Reset increments to zero
         for d in range(self.dim):
             self.inc[d, : self.ninc[d]] = (
-                self.grid[d, 1: self.ninc[d] + 1] -
-                self.grid[d, : self.ninc[d]]
+                self.grid[d, 1 : self.ninc[d] + 1] - self.grid[d, : self.ninc[d]]
             )
 
         # Clear accumulated training data for the next adaptation cycle
@@ -432,8 +422,7 @@ def make_uniform(self):
                 device=self.device,
             )
             self.inc[d, : self.ninc[d]] = (
-                self.grid[d, 1: self.ninc[d] + 1] -
-                self.grid[d, : self.ninc[d]]
+                self.grid[d, 1 : self.ninc[d] + 1] - self.grid[d, : self.ninc[d]]
             )
         self.clear()
 
@@ -459,8 +448,7 @@ def forward(self, u):
 
         batch_size = u.size(0)
         # Clamp iu to [0, ninc-1] to handle out-of-bounds indices
-        min_tensor = torch.zeros(
-            (1, self.dim), dtype=iu.dtype, device=self.device)
+        min_tensor = torch.zeros((1, self.dim), dtype=iu.dtype, device=self.device)
         # Shape: (1, dim)
         max_tensor = (self.ninc - 1).unsqueeze(0).to(iu.dtype)
         iu_clamped = torch.clamp(iu, min=min_tensor, max=max_tensor)
@@ -471,8 +459,7 @@ def forward(self, u):
         grid_gather = torch.gather(grid_expanded, 2, iu_clamped.unsqueeze(2)).squeeze(
             2
         )  # Shape: (batch_size, dim)
-        inc_gather = torch.gather(
-            inc_expanded, 2, iu_clamped.unsqueeze(2)).squeeze(2)
+        inc_gather = torch.gather(inc_expanded, 2, iu_clamped.unsqueeze(2)).squeeze(2)
 
         x = grid_gather + inc_gather * du_ninc
         log_detJ = (inc_gather * self.ninc).log_().sum(dim=1)
@@ -484,17 +471,15 @@ def forward(self, u):
             # For each sample and dimension, set x to grid[d, ninc[d]]
             # and log_detJ += log(inc[d, ninc[d]-1] * ninc[d])
             boundary_grid = (
-                self.grid[torch.arange(
-                    self.dim, device=self.device), self.ninc]
+                self.grid[torch.arange(self.dim, device=self.device), self.ninc]
                 .unsqueeze(0)
                 .expand(batch_size, -1)
             )
             # x = torch.where(out_of_bounds, boundary_grid, x)
             x[out_of_bounds] = boundary_grid[out_of_bounds]
 
             boundary_inc = (
-                self.inc[torch.arange(
-                    self.dim, device=self.device), self.ninc - 1]
+                self.inc[torch.arange(self.dim, device=self.device), self.ninc - 1]
                 .unsqueeze(0)
                 .expand(batch_size, -1)
             )
@@ -522,8 +507,7 @@ def inverse(self, x):
 
         # Initialize output tensors
         u = torch.empty_like(x)
-        log_detJ = torch.zeros(
-            batch_size, device=self.device, dtype=self.dtype)
+        log_detJ = torch.zeros(batch_size, device=self.device, dtype=self.dtype)
 
         # Loop over each dimension to perform inverse mapping
         for d in range(dim):
@@ -537,8 +521,7 @@ def inverse(self, x):
             # Perform searchsorted to find indices where x should be inserted to maintain order
             # torch.searchsorted returns indices in [0, max_ninc +1]
             iu = (
-                torch.searchsorted(
-                    grid_d, x[:, d].contiguous(), right=True) - 1
+                torch.searchsorted(grid_d, x[:, d].contiguous(), right=True) - 1
             )  # Shape: (batch_size,)
 
             # Clamp indices to [0, ninc_d - 1] to ensure they are within valid range
@@ -551,8 +534,7 @@ def inverse(self, x):
             inc_gather = inc_d[iu_clamped]  # Shape: (batch_size,)
 
             # Compute du: fractional part within the increment
-            du = (x[:, d] - grid_gather) / \
-                (inc_gather + TINY)  # Shape: (batch_size,)
+            du = (x[:, d] - grid_gather) / (inc_gather + TINY)  # Shape: (batch_size,)
 
             # Compute u for dimension d
             u[:, d] = (du + iu_clamped) / ninc_d  # Shape: (batch_size,)

MCintegration/utils.py

@@ -11,8 +11,8 @@
 
 # Constants for numerical stability
 # Small but safe non-zero value
-MINVAL = 10 ** (sys.float_info.min_10_exp + 50)
-MAXVAL = 10 ** (sys.float_info.max_10_exp - 50)  # Large but safe value
+# MINVAL = 10 ** (sys.float_info.min_10_exp + 50)
+MINVAL = 1e-45
 _VECTOR_TYPES = [np.ndarray, list]
 
 
@@ -83,8 +83,7 @@ def add(self, res):
             self._wlist.append(1 / (res.var if res.var > MINVAL else MINVAL))
             var = 1.0 / np.sum(self._wlist)
             sdev = np.sqrt(var)
-            mean = np.sum(
-                [w * m for w, m in zip(self._wlist, self._mlist)]) * var
+            mean = np.sum([w * m for w, m in zip(self._wlist, self._mlist)]) * var
             super(RAvg, self).__init__(*gvar.gvar(mean, sdev).internaldata)
         else:
             # Simple average
@@ -93,8 +92,7 @@ def add(self, res):
             self._count += 1
             mean = self._sum / self._count
             var = self._varsum / self._count**2
-            super(RAvg, self).__init__(
-                *gvar.gvar(mean, np.sqrt(var)).internaldata)
+            super(RAvg, self).__init__(*gvar.gvar(mean, np.sqrt(var)).internaldata)
 
     def extend(self, ravg):
         """

README.md

@@ -178,9 +178,8 @@ def singular_func(x, f):
     Integrand with singularity at x=0.
     The integral ∫₀¹ log(x)/√x dx = -4 (analytical result)
     """
-    f[:, 0] = torch.where(x[:, 0] < 1e-14, 
-                      0.0,
-                      torch.log(x[:, 0]) / torch.sqrt(x[:, 0]))
+    x_safe = torch.clamp(x[:, 0], min=1e-32)
+    f[:, 0] = torch.log(x_safe) / torch.sqrt(x_safe)
     return f[:, 0]
 
 # Integration parameters