fix sample

Author: houpc

MCintegration/integrators.py

@@ -125,13 +125,11 @@ def __init__(
                 self.device = device
 
         if isinstance(bounds, (list, np.ndarray)):
-            self.bounds = torch.tensor(
-                bounds, dtype=self.dtype, device=self.device)
+            self.bounds = torch.tensor(bounds, dtype=self.dtype, device=self.device)
         elif isinstance(bounds, torch.Tensor):
             self.bounds = bounds.to(dtype=self.dtype, device=self.device)
         else:
-            raise TypeError(
-                "bounds must be a list, NumPy array, or torch.Tensor.")
+            raise TypeError("bounds must be a list, NumPy array, or torch.Tensor.")
 
         assert self.bounds.shape[1] == 2, "bounds must be a 2D array"
 
@@ -177,13 +175,18 @@ def __call__(self, **kwargs):
     def sample(self, config, **kwargs):
         """
         Generate samples for integration.
-        This should be implemented by subclasses.
 
         Args:
             config (Configuration): The configuration object to store samples
             **kwargs: Additional parameters
         """
-        raise NotImplementedError("Subclasses must implement this method.")
+        config.u, config.detJ = self.q0.sample_with_detJ(config.batch_size)
+        if not self.maps:
+            config.x[:] = config.u
+        else:
+            config.x[:], detj = self.maps.forward_with_detJ(config.u)
+            config.detJ *= detj
+        self.f(config.x, config.fx)
 
     def statistics(self, means, vars, neval=None):
         """
@@ -212,11 +215,9 @@ def statistics(self, means, vars, neval=None):
                     gathered_vars = [
                         torch.zeros_like(vars) for _ in range(self.world_size)
                     ]
-            dist.gather(means, gathered_means if self.rank ==
-                        0 else None, dst=0)
+            dist.gather(means, gathered_means if self.rank == 0 else None, dst=0)
             if weighted:
-                dist.gather(vars, gathered_vars if self.rank ==
-                            0 else None, dst=0)
+                dist.gather(vars, gathered_vars if self.rank == 0 else None, dst=0)
 
             if self.rank == 0:
                 results = np.array([RAvg() for _ in range(f_dim)])
@@ -237,7 +238,7 @@ def statistics(self, means, vars, neval=None):
                             nblock_total, dtype=self.dtype, device=self.device
                         )
                         for igpu in range(self.world_size):
-                            _means[igpu * nblock: (igpu + 1) * nblock] = (
+                            _means[igpu * nblock : (igpu + 1) * nblock] = (
                                 gathered_means[igpu][:, i]
                             )
                         results[i].update(
@@ -341,8 +342,7 @@ def __call__(self, neval, nblock=32, verbose=-1, **kwargs):
         integ_values = torch.zeros(
             (self.batch_size, self.f_dim), dtype=self.dtype, device=self.device
         )
-        means = torch.zeros((nblock, self.f_dim),
-                            dtype=self.dtype, device=self.device)
+        means = torch.zeros((nblock, self.f_dim), dtype=self.dtype, device=self.device)
         vars = torch.zeros_like(means)
 
         for iblock in range(nblock):
@@ -354,8 +354,7 @@ def __call__(self, neval, nblock=32, verbose=-1, **kwargs):
             vars[iblock, :] = integ_values.var(dim=0) / self.batch_size
             integ_values.zero_()
 
-        results = self.statistics(
-            means, vars, epoch_perblock * self.batch_size)
+        results = self.statistics(means, vars, epoch_perblock * self.batch_size)
 
         if self.rank == 0:
             if self.f_dim == 1:
@@ -382,8 +381,7 @@ def random_walk(dim, device, dtype, u, **kwargs):
     """
     step_size = kwargs.get("step_size", 0.2)
     step_sizes = torch.ones(dim, device=device) * step_size
-    step = torch.empty(dim, device=device,
-                       dtype=dtype).uniform_(-1, 1) * step_sizes
+    step = torch.empty(dim, device=device, dtype=dtype).uniform_(-1, 1) * step_sizes
     new_u = (u + step) % 1.0
     return new_u
 
@@ -497,8 +495,7 @@ def sample(self, config, nsteps=1, mix_rate=0.5, **kwargs):
             acceptance_probs = new_weight / config.weight * new_detJ / config.detJ
 
             accept = (
-                torch.rand(self.batch_size, dtype=self.dtype,
-                           device=self.device)
+                torch.rand(self.batch_size, dtype=self.dtype, device=self.device)
                 <= acceptance_probs
             )
 
@@ -560,8 +557,7 @@ def __call__(
         config.x, detj = self.maps.forward_with_detJ(config.u)
         config.detJ *= detj
         config.weight = (
-            mix_rate / config.detJ + (1 - mix_rate) *
-            self.f(config.x, config.fx).abs_()
+            mix_rate / config.detJ + (1 - mix_rate) * self.f(config.x, config.fx).abs_()
         )
         config.weight.masked_fill_(config.weight < EPSILON, EPSILON)
 
@@ -571,14 +567,11 @@ def __call__(
         values = torch.zeros(
             (self.batch_size, self.f_dim), dtype=self.dtype, device=self.device
         )
-        refvalues = torch.zeros(
-            self.batch_size, dtype=self.dtype, device=self.device)
+        refvalues = torch.zeros(self.batch_size, dtype=self.dtype, device=self.device)
 
-        means = torch.zeros((nblock, self.f_dim),
-                            dtype=self.dtype, device=self.device)
+        means = torch.zeros((nblock, self.f_dim), dtype=self.dtype, device=self.device)
         vars = torch.zeros_like(means)
-        means_ref = torch.zeros(
-            (nblock, 1), dtype=self.dtype, device=self.device)
+        means_ref = torch.zeros((nblock, 1), dtype=self.dtype, device=self.device)
         vars_ref = torch.zeros_like(means_ref)
 
         for iblock in range(nblock):
@@ -588,17 +581,15 @@ def __call__(
 
                 config.fx.div_(config.weight.unsqueeze(1))
                 values += config.fx / n_meas_perblock
-                refvalues += 1 / \
-                    (config.detJ * config.weight) / n_meas_perblock
+                refvalues += 1 / (config.detJ * config.weight) / n_meas_perblock
             means[iblock, :] = values.mean(dim=0)
             vars[iblock, :] = values.var(dim=0) / self.batch_size
             means_ref[iblock, 0] = refvalues.mean()
             vars_ref[iblock, 0] = refvalues.var() / self.batch_size
             values.zero_()
             refvalues.zero_()
 
-        results_unnorm = self.statistics(
-            means, vars, nsteps_perblock * self.batch_size)
+        results_unnorm = self.statistics(means, vars, nsteps_perblock * self.batch_size)
         results_ref = self.statistics(
             means_ref, vars_ref, nsteps_perblock * self.batch_size
         )