graphical-sampling 0.1.0__py3-none-any.whl

graphical_sampling/clustering/soft_balanced_kmeans.py

@@ -0,0 +1,161 @@
+import numpy as np
+from numpy.typing import NDArray
+from sklearn.cluster import KMeans
+
+
+class SoftBalancedKMeans:
+    def __init__(
+        self, k: int, *, initial_centroids: NDArray = None, tolerance: int = 3
+    ) -> None:
+        self.k = k
+        self.tolerance = tolerance
+        self.coords: NDArray = None
+        self.centroids = initial_centroids
+        self.labels: NDArray = None
+        self.fractional_labels: NDArray = None
+        self.clusters_sum: NDArray = None
+        self.rng = np.random.default_rng()
+
+    def _generate_fractional_labels(self, probs: NDArray):
+        fractional_labels = np.zeros((*self.labels.shape, self.k))
+        for i in range(self.labels.shape[0]):
+            fractional_labels[i, self.labels[i]] = probs[i]
+        return fractional_labels
+
+    def _transfer_score(
+        self,
+        data_point: NDArray,
+        current_cluster_indx: float,
+        other_cluster_indx: float,
+    ) -> float:
+        if (
+            self.clusters_sum[current_cluster_indx]
+            - self.clusters_sum[other_cluster_indx]
+            > 10**-self.tolerance
+        ):
+            return (
+                np.linalg.norm(data_point - self.centroids[other_cluster_indx]) ** 2
+                - np.linalg.norm(data_point - self.centroids[current_cluster_indx]) ** 2
+            ) / (
+                self.clusters_sum[current_cluster_indx]
+                - self.clusters_sum[other_cluster_indx]
+                + 1e-9
+            )
+        else:
+            return np.inf
+
+    def _get_transfer_records(self, data: NDArray, top_m: int):
+        costs = []
+
+        for i in range(data.shape[0]):
+            for j in np.nonzero(self.fractional_labels[i])[0]:
+                t_min = np.argmin(
+                    [self._transfer_score(data[i], j, t) for t in range(self.k)]
+                )
+                cost = self._transfer_score(data[i], j, t_min)
+                costs.append((cost, i, j, t_min))
+
+        costs = np.array(costs)
+
+        return costs[np.argsort(costs[:, 0])][:top_m, 1:].astype(int)
+
+    def _transfer(self, data_index: int, from_index: int, to_index: int) -> None:
+        if (
+            self.clusters_sum[from_index] >= 1 - 10**-self.tolerance
+            and self.clusters_sum[to_index] >= 1 - 10**-self.tolerance
+        ) or (
+            self.clusters_sum[from_index] <= 1 + 10**-self.tolerance
+            and self.clusters_sum[to_index] <= 1 + 10**-self.tolerance
+        ):
+            transfer_prob = min(
+                self.fractional_labels[data_index, from_index],
+                (self.clusters_sum[from_index] - self.clusters_sum[to_index]) / 2,
+            )
+        else:
+            transfer_prob = min(
+                self.fractional_labels[data_index, from_index],
+                self.clusters_sum[from_index] - 1,
+                1 - self.clusters_sum[to_index],
+            )
+        self.fractional_labels[data_index, from_index] = (
+            self.fractional_labels[data_index, from_index] - transfer_prob
+        )
+        self.fractional_labels[data_index, to_index] = (
+            self.fractional_labels[data_index, to_index] + transfer_prob
+        )
+
+    def _no_transfer_possible(self, transfer_records: NDArray) -> bool:
+        return transfer_records[0, 0] == np.inf
+
+    def _is_transfer_possible(self, from_cluster: int, to_cluster: int) -> bool:
+        return (
+            self.clusters_sum[from_cluster] - self.clusters_sum[to_cluster]
+            > 10**-self.tolerance
+        )
+
+    def _stop_codition(self, tol) -> bool:
+        return np.all(np.abs(self.clusters_sum - 1) < 10**-tol)
+
+    def _expected_num_transfers(self) -> float:
+        max_diff_sum = np.max(self.clusters_sum - self.clusters_sum[:, None])
+        mean_nonzero_probs = np.mean(
+            self.fractional_labels[np.nonzero(self.fractional_labels)]
+        )
+        return max(int(np.floor(max_diff_sum / (2 * mean_nonzero_probs))), 1)
+
+    def _update_centroids(self, data: NDArray) -> None:
+        self.centroids = np.array(
+            [
+                np.mean(data[np.nonzero(self.fractional_labels[:, i])[0]], axis=0)
+                for i in range(self.k)
+            ]
+        )
+
+    def _numerical_stabilizer(self) -> float:
+        self.fractional_labels = np.round(self.fractional_labels, self.tolerance)
+        self.fractional_labels *= 1 / np.sum(self.fractional_labels, axis=0)
+        self.clusters_sum = np.sum(self.fractional_labels, axis=0)
+
+    def fit(self, data: NDArray, probs: NDArray) -> None:
+        self.coords = data
+
+        kmeans = KMeans(
+            n_clusters=self.k,
+            init=self.centroids if self.centroids is not None else "k-means++",
+            n_init=1 if self.centroids is not None else 10,
+            tol=10**-self.tolerance,
+        )
+        kmeans.fit(self.coords)
+
+        self.centroids = kmeans.cluster_centers_
+        self.labels = kmeans.labels_
+        self.fractional_labels = self._generate_fractional_labels(probs)
+        self.clusters_sum = np.sum(self.fractional_labels, axis=0)
+
+        while not self._stop_codition(self.tolerance):
+            transfer_records = self._get_transfer_records(
+                self.coords, top_m=self._expected_num_transfers()
+            )
+            if self._no_transfer_possible(transfer_records):
+                break
+            for data_index, from_cluster_index, to_cluster_index in transfer_records:
+                if self._is_transfer_possible(from_cluster_index, to_cluster_index):
+                    self._transfer(data_index, from_cluster_index, to_cluster_index)
+                    self.clusters_sum = np.sum(self.fractional_labels, axis=0)
+            self._update_centroids(self.coords)
+
+        self._numerical_stabilizer()
+
+    def get_clusters(self) -> NDArray:
+        clusters = []
+
+        for i in range(self.k):
+            probs = self.fractional_labels[:, i]
+            ids = np.nonzero(probs)[0]
+            units = np.concatenate(
+                [ids.reshape(-1, 1), self.coords[ids], probs[ids].reshape(-1, 1)],
+                axis=1,
+            )
+            clusters.append(units)
+
+        return clusters

graphical_sampling/criteria/__init__.py

@@ -0,0 +1,4 @@
+from .var_nht import VarNHT
+
+
+__all__ = ["VarNHT"]

graphical_sampling/criteria/criteria.py

@@ -0,0 +1,15 @@
+from dataclasses import dataclass
+from abc import abstractmethod
+
+from numpy.typing import NDArray
+
+from ..design import Design
+
+
+@dataclass
+class Criteria:
+    auxiliary_variable: NDArray
+    inclusion_probability: NDArray
+
+    @abstractmethod
+    def __call__(self, design: Design) -> float: ...

graphical_sampling/criteria/var_nht.py

@@ -0,0 +1,26 @@
+import numpy as np
+
+from ..design import Design
+from .criteria import Criteria
+
+
+class VarNHT(Criteria):
+    def __call__(self, design: Design) -> float:
+        nht_estimator = np.array(
+            [
+                np.sum(
+                    self.auxiliary_variable[list(sample.ids)]
+                    / self.inclusion_probability[list(sample.ids)]
+                )
+                for sample in design
+            ]
+        )
+
+        samples_probabilities = np.array([sample.probability for sample in design])
+
+        variance_nht = (
+            np.sum((nht_estimator**2) * samples_probabilities)
+            - (np.sum(self.auxiliary_variable)) ** 2
+        )
+
+        return variance_nht

graphical_sampling/design.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+from typing import Iterator, Collection, Optional
+
+import numpy as np
+from matplotlib import pyplot as plt
+from .structs import MaxHeap, Sample
+
+
+class Design:
+    def __init__(
+        self,
+        inclusions: Optional[Collection[float]] = None,
+        rng: np.random.Generator = np.random.default_rng(),
+    ):
+        self.heap = MaxHeap[Sample](rng=rng)
+        self.rng = rng
+        self.changes = 0
+        if inclusions is not None:
+            self.push_initial_design(inclusions)
+
+    def push_initial_design(self, inclusions: Collection[float]):
+        events: list[tuple[float, str, int]] = []
+        level: float = 0
+        for i, p in enumerate(inclusions):
+            next_level = level + p
+            if next_level < 1 - 1e-9:
+                events.append((level, "start", i))
+                events.append((next_level, "end", i))
+                level = next_level
+            elif next_level > 1 + 1e-9:
+                events.append((level, "start", i))
+                events.append((1, "end", i))
+                events.append((0, "start", i))
+                events.append((next_level - 1, "end", i))
+                level = next_level - 1
+            else:
+                events.append((level, "start", i))
+                events.append((1, "end", i))
+                level = 0
+
+        events.sort()
+        active = set()
+        last_point: float = 0
+
+        for point, event_type, bar_index in events:
+            if event_type == "start":
+                active.add(bar_index)
+            elif event_type == "end":
+                if last_point != point:
+                    self.push(Sample(round(point - last_point, 9), frozenset(active)))
+                active.remove(bar_index)
+
+            last_point = point
+
+    def copy(self) -> Design:
+        new_design = Design(
+            rng=self.rng,
+        )
+        new_design.heap = self.heap.copy()
+        new_design.changes = self.changes
+        return new_design
+
+    def pull(self, random: bool = False) -> Sample:
+        if random:
+            return self.heap.randompop()
+        return self.heap.pop()
+
+    def push(self, *args: Sample) -> None:
+        for r in args:
+            if not r.almost_zero():
+                self.heap.push(r)
+
+    def merge_identical(self):
+        dic = {}
+        for r in self.heap:
+            dic.setdefault(r.ids, 0)
+            dic[r.ids] += r.probability
+        self.heap = MaxHeap[Sample](
+            initial_heap=[Sample(length, ids) for ids, length in dic.items()],
+            rng=self.rng,
+        )
+
+    def switch(
+        self,
+        r1: Sample,
+        r2: Sample,
+        coefficient: float = 0.5,
+    ) -> tuple[Sample, Sample, Sample, Sample]:
+        length = coefficient * min(r1.probability, r2.probability)
+        n1 = self.rng.choice(list(r1.ids - r2.ids))
+        n2 = self.rng.choice(list(r2.ids - r1.ids))
+        return (
+            Sample(length, r1.ids - {n1} | {n2}),
+            Sample(r1.probability - length, r1.ids),
+            Sample(length, r2.ids - {n2} | {n1}),
+            Sample(r2.probability - length, r2.ids),
+        )
+
+    def iterate(
+        self, random_pull: bool = False, switch_coefficient: float = 0.5
+    ) -> None:
+        r1 = self.pull(random_pull)
+        r2 = self.pull(random_pull)
+        if r1.ids == r2.ids:
+            self.push(Sample(r1.probability + r2.probability, r1.ids))
+        else:
+            self.push(*self.switch(r1, r2, switch_coefficient))
+        self.changes += 1
+
+    def show(self) -> None:
+        initial_level: float = 0
+        for r in self.heap:
+            for i in r.ids:
+                plt.plot([i, i], [initial_level, initial_level + r.probability])
+            initial_level += r.probability
+        plt.show()
+
+    def __iter__(self) -> Iterator[Sample]:
+        return iter(self.heap)
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Design):
+            return NotImplemented
+        return self.heap == other.heap
+
+    def __hash__(self) -> int:
+        return hash(self.heap)

graphical_sampling/measure/__init__.py

@@ -0,0 +1,4 @@
+from .density import Density
+
+
+__all__ = ["Density"]

graphical_sampling/measure/density.py

@@ -0,0 +1,94 @@
+from ..clustering import AggregateBalancedKMeans
+
+import numpy as np
+from numpy._typing import NDArray
+from scipy.optimize import linear_sum_assignment
+from sklearn.neighbors import KernelDensity
+
+
+class Density:
+    def __init__(self, coordinates: NDArray, probabilities: NDArray, k: int):
+        # self.coords = (coordinates - coordinates.min(axis=0))/np.ptp(coordinates, axis=0)[0]
+        self.coords = coordinates
+        self.probs = probabilities
+        self.kde = self._kde(coordinates)
+        self.labels, self.centroids, self.clusters = self._generate_labels_centroids(k)
+
+    def _kde(self, coords: NDArray) -> KernelDensity:
+        kde = KernelDensity(kernel="tophat", bandwidth="scott")
+        kde.fit(coords)
+        return kde
+
+    def scale(self, scores):
+        scaled_scores = np.zeros_like(scores)
+        limit = np.sin(np.pi/8)/np.sin(np.pi/4)
+        for i, score in enumerate(scores):
+            if score >= 0:
+                scaled_scores[i] = min(score, limit)/limit
+            else:
+                scaled_scores[i] = max(score, -limit)/limit
+        return scaled_scores
+
+    def _density(self, shifted_coords: NDArray) -> float:
+        shifted_kde = self._kde(shifted_coords)
+        density = np.exp(self.kde.score_samples(self.coords))
+        shifted_density = np.exp(shifted_kde.score_samples(shifted_coords))
+        spread = np.mean(self.scale((density-shifted_density)/np.sqrt(density**2+shifted_density**2)))
+        var = np.mean(1 - (density+shifted_density)/(np.sqrt(2)*np.sqrt(density**2+shifted_density**2)))
+        scale_for_var = 1-np.cos(np.pi/8)
+        var_scaled = min(var, scale_for_var)/scale_for_var
+        measure = [
+            spread,
+            var_scaled,
+            spread + (np.sign(spread) - spread) * var_scaled,
+            # spread + (np.sign(spread) - spread) * var_scaled**2
+        ]
+        return density, shifted_density, measure
+
+    def _generate_labels_centroids(self, k):
+        agg = AggregateBalancedKMeans(k=k, tolerance=5)
+        agg.fit(self.coords, self.probs.reshape(-1, 1), np.array([1]))
+        labels = np.argmax(agg.membership, axis=1)
+        centroids = np.array(
+            [
+                np.mean(self.coords[labels == i], axis=0)
+                for i in range(k)
+            ]
+        )
+        return labels, centroids, agg.get_clusters()
+
+    def _assign_samples_to_centroids(self, samples, centroids):
+        cost_matrix = np.linalg.norm(samples[:, :, np.newaxis] - centroids, axis=3).transpose(0, 2, 1)
+        return np.array([samples[i][linear_sum_assignment(cost_matrix[i])[1]] for i in range(samples.shape[0])])
+
+    def _generate_shifted_coords(self, shifts: NDArray, labels: NDArray) -> NDArray:
+        shifted_coords = self.coords.copy()
+        for j, shift in enumerate(shifts):
+            shifted_coords[labels == j] += shift
+        return shifted_coords
+
+    def _score_sample(
+        self, sample: NDArray, labels: NDArray, centroids: NDArray
+    ) -> float:
+        shifts = sample - centroids
+        shifted_coords = self._generate_shifted_coords(shifts, labels)
+        return self._density(shifted_coords)
+
+    def score(self, samples: NDArray) -> NDArray:
+        scores = []
+        densities = []
+        samples_assigned = self._assign_samples_to_centroids(
+                self.coords[samples], self.centroids
+            )
+        for sample in samples_assigned:
+            density, shifted_density, score = self._score_sample(sample, self.labels, self.centroids)
+            densities.append([density, shifted_density])
+            scores.append(score)
+        return self.zippify(scores), densities
+
+    def zippify(self, scores):
+        zipped_scores = [[] for _ in range(len(scores[0]))]
+        for score in scores:
+            for i, den in enumerate(score):
+                zipped_scores[i].append(den)
+        return zipped_scores

graphical_sampling/random/__init__.py

@@ -0,0 +1,4 @@
+from .generator import rng
+
+
+__all__ = ["rng"]

graphical_sampling/random/generator.py

@@ -0,0 +1,251 @@
+import numpy as np
+from numpy.typing import NDArray
+
+
+class Generator:
+    def __init__(self, seed: int = None) -> None:
+        self.rng = np.random.default_rng(seed)
+
+    def grid_coordinates(self, size: int | tuple[int, ...] = None) -> NDArray:
+        """
+        Generate grid coordinates.
+
+        Args:
+            size (int|Tuple[int, ...]]): Specifies the grid dimensions.
+                - If None, generates a single grid of size `(1,)`.
+                - If an integer N, generates a single grid of size `(N,)`.
+                - If a tuple of `(N, D)`, generates a single grid with `N` points per `D` dimensions.
+                  The size of the grids is `(N**D, D)`.
+                - If a tuple of `(B, N, D)`, generates a batch of `B` grids of size `(N**D, D)`.
+
+        Returns:
+            NDArray: An array containing the coordinates of the grid(s), normalized to [0, 1].
+        """
+        batch, N, dim = self._check_size(size)
+
+        linspace = np.linspace(0, 1, N)
+        grid = np.meshgrid(*[linspace] * dim)
+        base_coordinates = np.stack([indices.ravel() for indices in grid], axis=-1)
+        coordinates = np.repeat(base_coordinates[np.newaxis, :, :], batch, axis=0)
+
+        return coordinates.squeeze()
+
+    def random_coordinates(self, size: int | tuple[int, ...] = None) -> NDArray:
+        """
+        Generate random coordinates uniformly from [0, 1].
+
+        Args:
+            size (int|Tuple[int, ...]]): Specifies the grid dimensions.
+                - If None, generates a single coordinate of size `(1,)`.
+                - If an integer N, generates one dimensional coordinates for N points.
+                - If a tuple of `(N, D)`, generates a D-dimensional coordinates for N points.
+                - If a tuple of `(B, N, D)`, generates a batch of `B` coordinates of size `(N, D)`.
+
+        Returns:
+            NDArray: An array containing random coordinates in [0, 1].
+        """
+        batch, N, dim = self._check_size(size)
+        coordinates = self.rng.random((batch, N, dim))
+        return coordinates.squeeze()
+
+    def uniform_coordinates(
+        self, low: float = 0.0, high: float = 1.0, size: int | tuple[int, ...] = None
+    ) -> NDArray:
+        """
+        Generate random coordinates uniformly from [low, high].
+
+        Args:
+            low (float): Lower boundary of the output interval. All values generated will be greater than or equal to low.
+                The default value is 0.
+            high (float): Upper boundary of the output interval. All values generated will be less than high.
+                high - low must be non-negative. The default value is 1.0.
+            size (int|Tuple[int, ...]]): Specifies the grid dimensions.
+                - If None, generates a single coordinate of size `(1,)`.
+                - If an integer N, generates one dimensional coordinates for N points.
+                - If a tuple of `(N, D)`, generates a D-dimensional coordinates for N points.
+                - If a tuple of `(B, N, D)`, generates a batch of `B` coordinates of size `(N, D)`.
+
+        Returns:
+            NDArray: An array containing random coordinates in [0, 1].
+        """
+        batch, N, dim = self._check_size(size)
+        coordinates = self.rng.uniform(low, high, (batch, N, dim)).squeeze()
+        return coordinates.squeeze()
+
+    def normal_1D_coordinates(
+        self, mean: float = 0.0, std: float = 1.0, size: int | tuple[int, ...] = None
+    ) -> NDArray:
+        """
+        Generate random 1D coordinates sampled from a normal distribution.
+
+        Args:
+            mean (float): Mean of the normal distribution (default 0).
+            std (float): Standard deviation of the normal distribution (default 1).
+            size (int|Tuple[int, ...]]): Specifies the coordinates dimensions.
+                - If None, generates a single coordinate of size `(1,)`.
+                - If an integer N, generates one dimensional coordinates for N points.
+                - If a tuple of `(B, N)`, generates a batch of `B` coordinates of size `(N,)`.
+                - If a tuple of `(B, N, D)`, D has to be 1.
+
+        Returns:
+            NDArray: An array containing coordinates sampled from a normal distribution.
+        """
+        batch, N, dim = self._check_size(size, fixed_dim=1)
+        coordinates = self.rng.normal(mean, std, (batch, N)).squeeze()
+        return coordinates.squeeze()
+
+    def normal_mD_coordinates(
+        self, mean: NDArray, cov: NDArray, size: int | tuple[int, ...] = None
+    ) -> NDArray:
+        """
+        Generate random m-dimensional coordinates sampled from a multivariate normal distribution.
+
+        Args:
+            mean (ArrayLike): Mean of the m-dimensional distribution.
+            cov (ArrayLike): Covariance matrix of the distribution.
+            size (int|Tuple[int, ...]]): Specifies the coordinates dimensions.
+                - If None, generates a single coordinate of size `(1, m)`.
+                - If an integer N, generates m-dimensional coordinates for N points. the shape of coordinates is `(N, m)`.
+                - If a tuple of `(B, N)`, generates a batch of `B` coordinates of size `(N, m)`.
+                - If a tuple of `(B, N, D)`, D has to be m.
+
+        Returns:
+            NDArray: An array containing coordinates sampled from a multivariate normal distribution.
+        """
+        batch, N, dim = self._check_size(size, fixed_dim=mean.shape[0])
+        coordinates = self.rng.multivariate_normal(mean, cov, (batch, N)).squeeze()
+        return coordinates.squeeze()
+
+    def cluster_coordinates(
+        self,
+        n_clusters: int,
+        cluster_std: float | NDArray,
+        size: int | tuple[int, ...] = None,
+    ) -> NDArray:
+        """
+        Generate clustered coordinates by placing points around cluster centers.
+
+        Args:
+            n_clusters (int): Number of clusters to generate.
+            cluster_std (float|NDArray): Standard deviation of points around each cluster center.
+            size (int|Tuple[int, ...]]): Specifies the grid dimensions.
+                - If None, generates a single coordinate of size `(1,)`.
+                - If an integer N, generates one dimensional coordinates for N points.
+                - If a tuple of `(N, D)`, generates a D-dimensional coordinates for N points.
+                - If a tuple of `(B, N, D)`, generates a batch of `B` coordinates of size `(N, D)`.
+
+        Returns:
+            NDArray: An array containing clustered coordinates.
+        """
+        if not isinstance(n_clusters, int):
+            raise ValueError("n_clusters must be an integer.")
+        if not isinstance(cluster_std, (float, np.ndarray)):
+            raise ValueError("cluster_std must be a float or an array.")
+
+        batch, N, dim = self._check_size(size)
+
+        cluster_centers = self.rng.random((batch, n_clusters, dim))
+        cluster_assignments = self.rng.integers(0, n_clusters, size=(batch, N))
+        coordinates = np.zeros((batch, N, dim))
+
+        for b in range(batch):
+            for n in range(N):
+                cluster_id = cluster_assignments[b, n]
+                coordinates[b, n, :] = self.rng.normal(
+                    cluster_centers[b, cluster_id],
+                    cluster_std
+                    if isinstance(cluster_std, float)
+                    else cluster_std[cluster_id],
+                    dim,
+                )
+
+        return coordinates.squeeze()
+
+    def equal_probabilities(
+        self, n: int, size: int | tuple[int, ...] = None
+    ) -> NDArray:
+        """
+        Generate equal probabilities that sum up to n.
+
+        Args:
+            n (int): The total sum of the probabilities.
+            size (int): The size of probabilities.
+                - If None, generates a single probability of size `(1,)`.
+                - If an integer N, generates N probabilities.
+                - If a tuple of `(B, N)`, generates a batch of `B` probabilities of size `(N,)`.
+                - If a tuple of `(B, N, D)`, D has to be 1.
+
+        Returns:
+            NDArray: An array of size `size` with equal probabilities summing to `n`.
+        """
+        if not isinstance(n, int) or n <= 0:
+            raise ValueError("n must be an positive integer.")
+
+        batch, N, dim = self._check_size(size, fixed_dim=1)
+        probabilities = np.full((batch, N, dim), n / N)
+        return probabilities.squeeze()
+
+    def unequal_probabilities(
+        self, n: int, size: int | tuple[int, ...] = None
+    ) -> NDArray:
+        """
+        Generate equal probabilities that sum up to n.
+
+        Args:
+            n (int): The total sum of the probabilities.
+            size (int): The size of probabilities.
+                - If None, generates a single probability of size `(1,)`.
+                - If an integer N, generates N probabilities.
+                - If a tuple of `(B, N)`, generates a batch of `B` probabilities of size `(N,)`.
+                - If a tuple of `(B, N, D)`, D has to be 1.
+
+        Returns:
+            NDArray: An array of size `size` with equal probabilities summing to `n`.
+        """
+        if not isinstance(n, int) or n <= 0:
+            raise ValueError("n must be an positive integer.")
+
+        batch, N, dim = self._check_size(size, fixed_dim=1)
+        probabilities = self.rng.random((batch, N))
+        probabilities *= n / probabilities.sum(axis=1)
+        return probabilities.squeeze()
+
+    def _check_size(
+        self, size: int | tuple[int, ...], fixed_dim: int = None
+    ) -> tuple[int, int, int]:
+        """
+        Validate and parse the size input.
+
+        Args:
+            size (int|Tuple[int, ...]]): Specifies the dimensions.
+
+        Returns:
+            Tuple[int, int, int]: A tuple `(B, N, D)` where:
+                - `B` is the batch size.
+                - `N` is the number of points per in the batch (or per dimension in some cases).
+                - `D` is the dimensionality of points.
+        """
+        if isinstance(size, int):
+            return (1, size, fixed_dim or 1)
+        if size is None or len(size) == 0:
+            return (1, 1, fixed_dim or 1)
+        if not isinstance(size, tuple) or len(size) > 3:
+            raise ValueError("Size must be a tuple of at most 3 positive integers.")
+        if not all(isinstance(x, int) and x > 0 for x in size):
+            raise ValueError("All elements of size must be positive integers.")
+        if len(size) == 3 and fixed_dim is not None and size[2] != fixed_dim:
+            raise ValueError(
+                "The dimension is fixed. you cannot set a different dimension."
+            )
+
+        match len(size):
+            case 1:
+                return (1, *size, fixed_dim or 1)
+            case 2:
+                return (1, *size) if fixed_dim is None else (*size, fixed_dim)
+            case 3:
+                return size
+
+
+def rng(seed=None):
+    return Generator(seed=seed)