from collections.abc import Callable, Sequence
from functools import partial
import numpy as np
from tadasets import sphere, torus, swiss_roll
[docs]
def point_clouds_from_constructors(
*,
constructors: Sequence[ Callable[[int, int], np.ndarray] ],
seed: int = 1_000,
nb_point_clouds: int = 40,
nb_points: int = 125,
nb_noise: int = 125,
nb_subsamples: int = 30,
nb_coarsened: int = 30
) -> np.ndarray:
clean_point_clouds = np.stack( [
np.stack(
[ constructor( nb_points, seed + i * nb_point_clouds + j ) for
j in range( nb_point_clouds )
]
) for
i, constructor in enumerate ( constructors )
] )
rng = np.random.default_rng( seed + 10_000 )
point_clouds = np.concatenate( (
clean_point_clouds,
rng.random( ( len(constructors),
nb_point_clouds,
nb_noise,
3) ) * 6 - 3
), axis = 2 )
return np.take_along_axis(
axis = 3,
arr = point_clouds[:, :, np.newaxis, :, :],
indices = rng.integers(
low = 0,
high = nb_points + nb_noise,
size = ( len(constructors),
nb_point_clouds,
nb_coarsened,
nb_subsamples,
1 )
)
)
[docs]
def my_sphere(nb_points: int, seed: int) -> np.ndarray:
return sphere(
r = 2,
n = nb_points,
seed = seed
)
[docs]
def my_torus(nb_points: int, seed: int) -> np.ndarray:
return torus(
c = 2,
a = 1,
n = nb_points,
seed = seed
)
[docs]
def my_swiss_roll(nb_points: int, seed: int) -> np.ndarray:
return (
swiss_roll(
r = 20,
n = nb_points,
seed = seed
) - np.array( [0.0, 0.0, 7.0] )
) / 7.0
create_point_clouds = partial(
point_clouds_from_constructors,
constructors = [ my_sphere,
my_torus,
my_swiss_roll
]
)