def generate_gaussians_spiral(n_samples=100,
n_gaussians=7,
n_gaussians_per_loop=4,
dim=2,
radius_start=1,
radius_end=0.2,
std_gaussians_start=0.3,
std_gaussians_end=0.1,
noise=1e-3):
"""Creates `dim`-dimensional `n_gaussians` on a spiral.
:param n_samples: number of data points in the generated dataset
:type n_samples: int
:param n_gaussians: number of total gaussians distributions placed on the spirals
:type n_gaussians: int
:param n_gaussians_per_loop: number of gaussians distributions per loop of the spiral
:type n_gaussians_per_loop: int
:param dim: dimension of the dataset. The distributions are placed on the hyperplane (x1, x2, 0, 0..) if dim > 2
:type dim: int
:param radius_start: starting radius of the spiral
:type radius_start: int
:param radius_end: end radius of the spiral
:type radius_end: int
:param std_gaussians_start: standard deviation of the gaussians at the start of the spiral. Linear interpolation (start, end, num_gaussians)
:type std_gaussians_start: int
:param std_gaussians_end: standard deviation of the gaussians at the end of the spiral
:type std_gaussians_end: int
:param noise: standard deviation of noise magnitude added to each data point
:type noise: float
"""
X = torch.zeros(n_samples * n_gaussians, dim)
y = torch.zeros(n_samples * n_gaussians).long()
angle = torch.zeros(1)
radiuses = torch.linspace(radius_start, radius_end, n_gaussians)
std_devs = torch.linspace(std_gaussians_start, std_gaussians_end,
n_gaussians)
if dim > 2:
loc = torch.cat([
radiuses[0] * torch.cos(angle), radiuses[0] * torch.sin(angle),
torch.zeros(dim - 2)
])
else:
loc = torch.cat(
[radiuses[0] * torch.cos(angle), radiuses[0] * torch.sin(angle)])
dist = Normal(loc, scale=std_devs[0])
for i in range(n_gaussians):
angle += 2 * math.pi / n_gaussians_per_loop
if dim > 2:
dist.loc = torch.Tensor([
radiuses[i] * torch.cos(angle),
torch.sin(angle), radiuses[i] * torch.zeros(dim - 2)
])
else:
dist.loc = torch.Tensor([
radiuses[i] * torch.cos(angle), radiuses[i] * torch.sin(angle)
])
dist.scale = std_devs[i]
X[i * n_samples:(i + 1) * n_samples] = dist.sample(
sample_shape=(n_samples, )) + torch.randn(dim) * noise
y[i * n_samples:(i + 1) * n_samples] = i
return X, y
def std_normal(shape):
N = Normal(torch.zeros(shape), torch.ones(shape))
if torch.cuda.is_available():
N.loc = N.loc.cuda()
N.scale = N.scale.cuda()
return N
def get_normal_dist(mean, std):
normal = Normal(mean, std)
normal.loc = normal.loc.to(device)
normal.scale = normal.scale.to(device)
return normal
