Synthetic data¶

datasets ¶

Synthetic datasets with realistic structure.

Uniform Gaussian noise is a bad benchmark for nearest-neighbour search. Past a few dozen dimensions every point sits at roughly the same distance from every other, so recall stops telling you anything about the index. Each generator here puts back structure that real single-cell data has, and each stresses a different part of an index.

These are the same generators, with the same seeds, behind the benchmark tables in the Rust crate's docs/. A Python benchmark and a cargo run --example gridsearch_hnsw run see identical points, so the numbers are comparable.

>>> from ann_search import datasets
>>> X, labels = datasets.make_clustered(n_samples=50_000, dim=32,
...                                     n_clusters=25, seed=42)
>>> Q = datasets.subsample_queries(X, n_samples=5_000, seed=42)

Every generator returns (X, labels), so ground-truth cluster labels come free: useful for recall, and for scoring whatever clustering you run downstream.

Output is always float32, which is what the benchmark tables use. Cast with X.astype(np.float64) if you want to exercise the f64 paths.

The tuning constants behind each generator (bridge fractions, spectral decay, rogue-dimension counts) are fixed rather than exposed. They're what the published tables were produced with, and changing one quietly invalidates the comparison.

make_clustered ¶

make_clustered(
    n_samples: int = 150000,
    dim: int = 32,
    n_clusters: int = 25,
    *,
    seed: int = 42,
) -> tuple[ndarray, ndarray]

Separated Gaussian clusters joined by inter-cluster bridges.

The baseline. Well-separated blobs with a fifth of the points placed on bridges between neighbouring clusters, so the boundaries aren't trivially clean.

Parameters:

Name	Type	Description	Default
`n_samples`	`int`	Rows to generate.	`150000`
`dim`	`int`	Features per row.	`32`
`n_clusters`	`int`	Distinct clusters.	`25`
`seed`	`int`	Fixes the whole draw.	`42`

Returns:

Type	Description
`ndarray`	`(X, labels)`, shapes `(n_samples, dim)` float32 and
`ndarray`	`(n_samples,)` int64.

make_correlated ¶

make_correlated(
    n_samples: int = 150000,
    dim: int = 32,
    n_clusters: int = 25,
    *,
    seed: int = 42,
    cor_strength: float | None = None,
) -> tuple[ndarray, ndarray]

Clusters with local anisotropy plus a shared off-axis subspace.

The interesting case for quantisation. That shared subspace is exactly what OPQ's learned rotation exploits and what PQ's axis-aligned split can't see, so the two pull apart here in a way they don't on make_clustered.

Parameters:

Name	Type	Description	Default
`n_samples`	`int`	Rows to generate.	`150000`
`dim`	`int`	Features per row.	`32`
`n_clusters`	`int`	Distinct clusters.	`25`
`seed`	`int`	Fixes the whole draw.	`42`
`cor_strength`	`float \| None`	Share of structured variance routed to the global off-axis subspace, 0.0 to 1.0. `None` uses the value behind the published tables.	`None`

Returns:

Type	Description
`ndarray`	`(X, labels)`, shapes `(n_samples, dim)` float32 and
`ndarray`	`(n_samples,)` int64.

Raises:

Type	Description
`ValueError`	If `cor_strength` is outside 0.0 to 1.0.

make_low_rank ¶

make_low_rank(
    n_samples: int = 150000,
    dim: int = 32,
    n_clusters: int = 25,
    *,
    intrinsic_dim: int = 16,
    seed: int = 42,
) -> tuple[ndarray, ndarray]

Data on a low-dimensional manifold inside a high-dimensional space.

The manifold hypothesis made concrete: cell types sit on an intrinsic_dim-dimensional surface embedded in dim ambient dimensions, with differentiation trajectories running between them.

Parameters:

Name	Type	Description	Default
`n_samples`	`int`	Rows to generate.	`150000`
`dim`	`int`	Ambient features per row.	`32`
`n_clusters`	`int`	Distinct cell types.	`25`
`intrinsic_dim`	`int`	True dimensionality of the manifold. Must not exceed `dim`.	`16`
`seed`	`int`	Fixes the whole draw.	`42`

Returns:

Type	Description
`ndarray`	`(X, labels)`, shapes `(n_samples, dim)` float32 and
`ndarray`	`(n_samples,)` int64.

Raises:

Type	Description
`ValueError`	If `intrinsic_dim` exceeds `dim`.

make_cell_embeddings ¶

make_cell_embeddings(
    n_samples: int = 150000,
    dim: int = 32,
    n_clusters: int = 25,
    *,
    seed: int = 42,
) -> tuple[ndarray, ndarray]

Foundation-model cell embeddings, Geneformer or scGPT flavoured.

The nastiest of the four. Heavy-tailed spectrum, a handful of high-variance rogue dimensions, and a shared mean offset that crams everything into an anisotropy cone. Quantised indices get painful here.

Parameters:

Name	Type	Description	Default
`n_samples`	`int`	Rows to generate.	`150000`
`dim`	`int`	Embedding width.	`32`
`n_clusters`	`int`	Distinct cell types.	`25`
`seed`	`int`	Fixes the whole draw.	`42`

Returns:

Type	Description
`ndarray`	`(X, labels)`, shapes `(n_samples, dim)` float32 and
`ndarray`	`(n_samples,)` int64.

subsample_queries ¶

subsample_queries(
    x: ndarray, n_samples: int, *, seed: int = 42
) -> ndarray

Draw a query set from a dataset, with light Gaussian noise added.

Querying an index with rows it was built from flatters it: every query has an exact hit at distance zero. This perturbs the draw so queries sit near the data instead of on it.

Parameters:

Name	Type	Description	Default
`x`	`ndarray`	Dataset to draw from. Cast to float32 if it isn't already.	required
`n_samples`	`int`	Rows to draw. Capped at the number available.	required
`seed`	`int`	Fixes both the draw and the noise.	`42`

Returns:

Type	Description
`ndarray`	`(min(n_samples, len(x)), dim)` float32.

Raises:

Type	Description
`ValueError`	If `x` is not 2-D.