Source code for create_point_clouds

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 ] )