def draw()

in botorch/utils/sampling.py [0:0]

• 27 lines of code
• 4 McCabe index (conditional complexity)

    def draw(self, n: int = 1, seed: Optional[int] = None) -> Tensor:
        r"""Draw samples from the polytope.

        Args:
            n: The number of samples.
            seed: The random seed.

        Returns:
            A `n x d` Tensor of samples from the polytope.
        """
        if self.dim == 1:
            with manual_seed(seed):
                e = torch.rand(n, 1, device=self.new_A.device, dtype=self.new_A.dtype)
            transformed_samples = self.y_min + (self.y_max - self.y_min) * e
        else:
            if seed is None:
                generator = None
            else:
                generator = torch.Generator(device=self.A.device)
                generator.manual_seed(seed)
            index_rvs = torch.multinomial(
                self._p,
                num_samples=n,
                replacement=True,
                generator=generator,
            )
            simplex_rvs = sample_simplex(
                d=self.dim + 1, n=n, seed=seed, device=self.A.device, dtype=self.A.dtype
            )
            transformed_samples = torch.stack(
                [rv @ self._polytopes[idx] for rv, idx in zip(simplex_rvs, index_rvs)]
            )
        init_shift = self.x0.transpose(-1, -2)
        samples = init_shift + transformed_samples @ self.nullC.transpose(-1, -2)
        return samples