sakura.utils.distributions.gaussian_mixture

sakura.utils.distributions.gaussian_mixture(batch_size, n_dim, n_labels=10, x_var=0.5, y_var=0.1, label_indices=None, shift_magnitude=1.4, angular_offset=0.0)

Generates samples from a Gaussian mixture distribution arranged in 2D subspace(s).

The distribution consists of Gaussian components arranged in circle(s) within multiple 2D subspace(s). For even dimensions, each pair of dimensions forms a 2D subspace with rotated Gaussian components. For odd dimensions, the last dimension contains independent Gaussian noise.

Parameters:

• batch_size (int) – Number of samples to generate

• n_dim (Literal[2], optional) – Dimension of output samples

• n_labels (int, optional) – Number of Gaussian components arranged around the circle, defaults to 10

• x_var (float, optional) – Variance along the x-axis for each component, defaults to 0.5

• y_var (float, optional) – Variance along the y-axis for each component, defaults to 0.1

• label_indices (list[int], optional) – List of label indices (0 <= index < n_labels) for each batch sample, randomly assigned if none provided

• shift_magnitude (float) – Distance from the origin (center) to the center of each Gaussian component, defaults to 1.4

• angular_offset (float) – Angle offset coefficient between different 2D subspaces, defaults to 0.0

Returns:

Tensor of shape (batch_size, n_dim) with uniform samples.

Return type:

torch.FloatTensor