fix bug in function latin_hypercube_sampling_standard.

src/evox/operators/sampling/latin_hypercube.py

@@ -1,35 +1,38 @@
 import torch
 
 
-def latin_hypercube_sampling_standard(n: int, d: int, device: torch.device, scramble: bool = True):
+def latin_hypercube_sampling_standard(n: int, d: int, device: torch.device, smooth: bool = True):
     """Generate Latin Hypercube samples in the unit hypercube.
 
     :param n: The number of sample points to generate.
     :param d: The dimensionality of the samples.
     :param device: The device on which to generate the samples.
-    :param scramble: Whether to scramble the order of the samples. Defaults to True.
+    :param smooth: Whether to generate sample in random positions in the cells or not. Defaults to True.
 
     :return: A tensor of shape (n, d), where each row represents a sample point and each column represents a dimension.
     """
-    samples = (torch.arange(0, n, device=device).view(-1, 1) + torch.rand(n, d, device=device)) / n
-    if scramble:
-        samples = samples[torch.randperm(n)]
+    cells = torch.arange(0, n, device=device).view(-1, 1).expand(n, d).contiguous()
+    cells_perms = torch.rand(n, d, device=device).argsort(dim=0)
+    cells = cells.gather(0, cells_perms)
+    if smooth:
+        samples = (cells + torch.rand(n, d, device=device)) / n
+    else:
+        samples = (cells + 0.5) / n
     return samples
 
 
-def latin_hypercube_sampling(n: int, d: int, lb: torch.Tensor, ub: torch.Tensor, scramble: bool = True):
+def latin_hypercube_sampling(n: int, lb: torch.Tensor, ub: torch.Tensor, smooth: bool = True):
     """Generate Latin Hypercube samples in the given hypercube defined by `lb` and `ub`.
 
     :param n: The number of sample points to generate.
-    :param d: The dimensionality of the samples.
-    :param lb: The lower bounds of the hypercube. Must be a 1D tensor with same shape, dtype, and device as `ub`.
-    :param ub: The upper bounds of the hypercube. Must be a 1D tensor with same shape, dtype, and device as `lb`.
-    :param scramble: Whether to scramble the order of the samples. Defaults to True.
+    :param lb: The lower bounds of the hypercube. Must be a 1D tensor of size `d` with same shape, dtype, and device as `ub`.
+    :param ub: The upper bounds of the hypercube. Must be a 1D tensor of size `d` with same shape, dtype, and device as `lb`.
+    :param smooth: Whether to generate sample in random positions in the cells or not. Defaults to True.
 
     :return: A tensor of shape (n, d), where each row represents a sample point and each column represents a dimension whose device is the same as `lb` and `ub`.
     """
     assert lb.device == ub.device and lb.dtype == ub.dtype and lb.ndim == 1 and ub.ndim == 1 and lb.size() == ub.size()
-    samples = latin_hypercube_sampling_standard(n, d, lb.device, scramble)
+    samples = latin_hypercube_sampling_standard(n, lb.size(0), lb.device, smooth)
     lb = lb[None, :]
     ub = ub[None, :]
     return lb + samples * (ub - lb)