API reference

Generated from the source. Everything here is exported at the top level: from fastumap import ....

Projection

fastumap.projection

Public projection entry points: umap_project, spectral_project, and the fit / transform pair for placing new points into an existing layout.

UMAPModel dataclass

A fitted layout plus everything transform needs to place new points.

train is kept because placing a new point needs its nearest neighbours in the training data. Those neighbours must be found in the same space the fit used: the PCA-reduced space when pca_dim was set. embedding is the frozen training layout that new points are positioned against.

to_npz

to_npz(file: str | PathLike[str] | BinaryIO) -> None

Save to a versioned .npz. Portable across releases, unlike pickle, which welds the model to the exact dataclass layout. Load it back with :meth:from_npz.

from_npz classmethod

from_npz(file: str | PathLike[str] | BinaryIO) -> UMAPModel

Load a model written by :meth:to_npz. If the file was written by an incompatible release the error names both format versions rather than failing obscurely.

transform

transform(matrix: FloatArray) -> FloatArray

Place new points into this fitted layout without refitting, in the same space (so the same input lands in the same place across scope changes). A method over the module-level :func:transform; identical result. matrix is (n_new, n_features).

umap_project

umap_project(
    matrix: FloatArray,
    dimensions: int = ...,
    *,
    n_neighbors: int = ...,
    min_dist: float = ...,
    spread: float = ...,
    n_epochs: int | None = ...,
    negative_sample_rate: int = ...,
    initial_alpha: float = ...,
    random_state: int = ...,
    metric: str = ...,
    chunk_count: int = ...,
    pca_dim: int | None = ...,
    knn: str = ...,
    precomputed_knn: tuple[IntArray, FloatArray]
    | None = ...,
    init: FloatArray | None = ...,
    y: IntArray | None = ...,
    target_weight: float = ...,
    densmap: bool = ...,
    dens_lambda: float = ...,
    dens_frac: float = ...,
    dens_var_shift: float = ...,
    return_model: Literal[False] = ...,
) -> FloatArray

umap_project(
    matrix: FloatArray,
    dimensions: int = ...,
    *,
    n_neighbors: int = ...,
    min_dist: float = ...,
    spread: float = ...,
    n_epochs: int | None = ...,
    negative_sample_rate: int = ...,
    initial_alpha: float = ...,
    random_state: int = ...,
    metric: str = ...,
    chunk_count: int = ...,
    pca_dim: int | None = ...,
    knn: str = ...,
    precomputed_knn: tuple[IntArray, FloatArray]
    | None = ...,
    init: FloatArray | None = ...,
    y: IntArray | None = ...,
    target_weight: float = ...,
    densmap: bool = ...,
    dens_lambda: float = ...,
    dens_frac: float = ...,
    dens_var_shift: float = ...,
    return_model: Literal[True],
) -> tuple[FloatArray, UMAPModel]

umap_project(
    matrix: FloatArray,
    dimensions: int = 2,
    *,
    n_neighbors: int = 15,
    min_dist: float = 0.1,
    spread: float = 1.0,
    n_epochs: int | None = None,
    negative_sample_rate: int = 5,
    initial_alpha: float = 1.0,
    random_state: int = 42,
    metric: str = "euclidean",
    chunk_count: int = 1,
    pca_dim: int | None = None,
    knn: str = "auto",
    precomputed_knn: tuple[IntArray, FloatArray]
    | None = None,
    init: FloatArray | None = None,
    y: IntArray | None = None,
    target_weight: float = 0.5,
    densmap: bool = False,
    dens_lambda: float = 2.0,
    dens_frac: float = 0.3,
    dens_var_shift: float = 0.1,
    return_model: bool = False,
) -> FloatArray | tuple[FloatArray, UMAPModel]

Project matrix (n_samples, n_features) into dimensions output dims.

Use 2 or 3 for a picture. Use a mid dimension (~10) when the goal is to cluster the result: UMAP's own guidance is to reduce to ~10 dims, not 2, before clustering, so the layout keeps more structure. dimensions must be at least 1 and at most n_features.

Deterministic given random_state: the spectral start vector is pinned and the negative sampler is seeded, so the same input and seed give bit-identical output.

metric is "euclidean" or "cosine". chunk_count > 1 enables partial in-epoch feedback (higher local overlap, slower). pca_dim pre-reduces the input to that many principal components before the kNN, to speed up high-dimensional inputs.

densmap=True turns on densMAP: it keeps dense regions dense and sparse regions sparse, instead of UMAP's usual flattening of local density. dens_lambda (default 2.0) is the strength; dens_frac and dens_var_shift tune it as in :func:fit. It runs on the numpy path (the native kernel does plain UMAP), so it is slower. See the README.

return_model=True returns (embedding, model) instead of just the coordinates, so a fit can be pinned and reused: call model.transform(new_points) (or the module-level :func:transform) to place later points into the same layout. That is what keeps a scatter stable across a scope change: the same points land in the same places instead of the axes being reshuffled by a refit. The model also serialises with :meth:UMAPModel.to_npz.

spectral_project

spectral_project(
    matrix: FloatArray,
    dimensions: int = 2,
    *,
    n_neighbors: int = 15,
    random_state: int = 42,
    metric: str = "euclidean",
    pca_dim: int | None = None,
    knn: str = "auto",
) -> FloatArray

Just the spectral initialisation: kNN graph, fuzzy set, spectral embedding.

Useful on its own as a fast, deterministic layout, and as the seed the full optimisation refines. pca_dim pre-reduces the input before the kNN.

fit

fit(
    matrix: FloatArray,
    dimensions: int = 2,
    *,
    n_neighbors: int = 15,
    min_dist: float = 0.1,
    spread: float = 1.0,
    n_epochs: int | None = None,
    negative_sample_rate: int = 5,
    initial_alpha: float = 1.0,
    random_state: int = 42,
    metric: str = "euclidean",
    chunk_count: int = 1,
    pca_dim: int | None = None,
    knn: str = "auto",
    precomputed_knn: tuple[IntArray, FloatArray]
    | None = None,
    init: FloatArray | None = None,
    y: IntArray | None = None,
    target_weight: float = 0.5,
    densmap: bool = False,
    dens_lambda: float = 2.0,
    dens_frac: float = 0.3,
    dens_var_shift: float = 0.1,
) -> UMAPModel

Fit a UMAP layout and return a :class:UMAPModel (use .embedding for the coords).

Deterministic given random_state. See :func:umap_project for the parameters.

pca_dim (default None = off) pre-reduces the input to that many principal components before the kNN. This speeds up high-dimensional inputs. transform projects new points through the same basis.

precomputed_knn skips the built-in neighbour search. Pass (indices, distances), each (n, k) with column 0 the point itself, e.g. a faiss query you already ran against your own index. Then n_neighbors, metric and knn no longer affect the graph, and it cannot be combined with pca_dim. transform still uses matrix's own kNN, so pass the kNN of matrix (in matrix's space) if you also want transform.

init overrides the spectral start with a given (n, dimensions) embedding. Pass the previous view's coordinates so a refreshed layout stays comparable, instead of rotating or reflecting. init sets where the optimiser starts, not where it ends: the learning rate still decays from initial_alpha over all n_epochs, so at the default epoch count a warm start drifts almost as far as a cold one. For a refreshed view that should stay close to the previous one, also reduce n_epochs (a lighter touch) or lower initial_alpha (gentler corrections across the full schedule). Place any new rows yourself.

initial_alpha (default 1.0) is the starting SGD learning rate. Lower it with init to make small corrections rather than a full re-layout.

y (categorical class labels, one per point, -1 for unlabelled) turns on supervised projection: same-label points attract, weighted by target_weight in [0, 1] (default 0.5; 1.0 severs every inter-class edge). Default y=None is unchanged unsupervised UMAP.

transform

transform(
    model: UMAPModel,
    matrix: FloatArray,
    *,
    n_epochs: int | None = ...,
    return_distances: Literal[False] = ...,
) -> FloatArray

transform(
    model: UMAPModel,
    matrix: FloatArray,
    *,
    n_epochs: int | None = ...,
    return_distances: Literal[True],
) -> tuple[FloatArray, FloatArray]

transform(
    model: UMAPModel,
    matrix: FloatArray,
    *,
    n_epochs: int | None = None,
    return_distances: bool = False,
) -> FloatArray | tuple[FloatArray, FloatArray]

Place new points into model's existing layout without refitting it.

New points are attracted to their nearest neighbours in the training set and repelled from random training points. The training layout stays fixed, so the picture is stable across calls. Deterministic given the model and input. Returns (n_new, dimensions).

With return_distances=True it also returns, per new point, the distance to its nearest training neighbour. This is a per-point measure of fit. Watch it against the training set's own mean nearest-neighbour distance. Points landing 2-3x further out are extrapolations. A rising batch mean is the signal to refit (see the README).

Clustering

fastumap.clustering

Clustering, implemented ourselves. numpy + scipy plus a built-in native k-means kernel. No third-party clustering library.

fastumap reimplements UMAP from scratch. Clustering is the same. cluster() does NOT compute a 2-D/3-D layout first. That layout is an iterative SGD and takes minutes at a million points. Clustering does not need it.

Four methods, all ours:

• method="kmeans" — k-means (k-means++ seeding, Lloyd iterations, blocked assignment so the n-by-k distance matrix is never materialised). Fastest. You pick n_clusters. Above ~50k points it routes through the built-in native kernel (fastumap._accel): ~10 s at 1M x 128. Below that, the deterministic numpy path here.
• method="spectral" — spectral clustering. Build the fuzzy kNN graph. Take the leading normalised-adjacency eigenvectors (fastumap's own spectral machinery). k-means those. Finds non-convex clusters k-means alone misses. Needs fastumap[ann] only above ~16k points, where the kNN switches to the approximate backend.
• method="dbscan" — density clustering. Finds clusters of any shape, labels outliers -1, and picks the number of clusters itself. You give it eps (or let it auto-pick).
• method="hdbscan" — the hierarchical version: no eps at all, so clusters of different densities come out together. You give it min_cluster_size. The usual choice after a UMAP reduction when you don't know k.

Memory: the native kernel runs in float32, but the numpy entry points upcast the input to float64 (np.ascontiguousarray(..., dtype=np.float64)) for a deterministic-to-the-byte result and for the inertia. So on any path the float64 work array is the memory driver at large n: a 1,000,000 x 1024 float32 input becomes an ~8 GB float64 array. The assignment step is blocked, so the n-by-k distance matrix on top of that stays bounded by the block size, not by n. At 1M the native kernel meets the ~10 s budget; without it (an unbuilt checkout) the numpy fallback is minutes, since a numpy inner loop is not the threaded float32 GEMM that budget assumes.

ClusterResult dataclass

What cluster(..., return_centroids=True) returns.

centroids, inertia and dispersion are in the space the clustering ran in: the input space, or unit-normalised rows for metric="cosine". inertia is the k-means objective, the sum of squared distances from each point to its assigned centroid. Lower is tighter. It is the number to compare across n_clusters when picking k.

dispersion[c] is the mean distance of cluster c's points to centroid c, and sizes[c] is how many points it has (both length n_clusters, indexed by label). Where inertia is one number for the whole fit, dispersion describes each cluster on its own: a diffuse cluster and a tight one of the same size look identical by size and total inertia, but differ in dispersion, which is what tells a caller which clusters hold real structure worth splitting. Computed in the fit's own metric, so a cosine fit is not silently measured in euclidean.

cluster

cluster(
    embeddings: FloatArray,
    *,
    method: str = ...,
    n_clusters: int | None = ...,
    n_neighbors: int | None = ...,
    metric: str = ...,
    n_iter: int = ...,
    n_init: int | None = ...,
    random_state: int = ...,
    eps: float | None = ...,
    min_samples: int | None = ...,
    min_cluster_size: int = ...,
    return_centroids: Literal[False] = ...,
) -> IntArray

cluster(
    embeddings: FloatArray,
    *,
    method: str = ...,
    n_clusters: int | None = ...,
    n_neighbors: int | None = ...,
    metric: str = ...,
    n_iter: int = ...,
    n_init: int | None = ...,
    random_state: int = ...,
    eps: float | None = ...,
    min_samples: int | None = ...,
    min_cluster_size: int = ...,
    return_centroids: Literal[True],
) -> ClusterResult

cluster(
    embeddings: FloatArray,
    *,
    method: str = "kmeans",
    n_clusters: int | None = None,
    n_neighbors: int | None = None,
    metric: str = "euclidean",
    n_iter: int = 50,
    n_init: int | None = None,
    random_state: int = 42,
    eps: float | None = None,
    min_samples: int | None = None,
    min_cluster_size: int = 5,
    return_centroids: bool = False,
) -> IntArray | ClusterResult

Cluster embeddings (n, d) into integer labels (n,).

Four methods:

• "kmeans" and "spectral" take n_clusters (you pick k).
• "dbscan" picks the number of clusters itself and labels outliers -1. It takes eps (None = auto) and min_samples instead of n_clusters. Good when you don't know k or want noise handling; see :func:dbscan.
• "hdbscan" does the same but with no eps at all, so clusters of different densities are found together. It takes min_cluster_size (and min_samples, which defaults to min_cluster_size here, not to dbscan's 5); see :func:hdbscan.

metric is "euclidean" or "cosine" (cosine unit-normalises the rows). Deterministic given random_state. Does NOT compute a layout. For a picture, call umap_project.

n_init ("kmeans" only) restarts the fit from that many seeds and keeps the lowest inertia. The default single run is fine on the overlapping data real embeddings produce; raise it (n_init=5) if your data is well-separated, where mini-batch can otherwise settle into a worse optimum.

With return_centroids=True the return is a :class:ClusterResult (labels, centroids, inertia) instead of just labels ("kmeans" / "spectral" only).

kmeans

kmeans(
    x: FloatArray,
    n_clusters: int,
    *,
    n_iter: int = 50,
    random_state: int = 42,
    batch_size: int | None = None,
    n_init: int | None = None,
) -> tuple[IntArray, FloatArray]

Lloyd's k-means with k-means++ init. Returns (labels, centroids).

Uses the built-in native kernel (fastumap._accel) above _ACCEL_MIN points when present (~10s at 1M x 128 vs minutes for numpy). Otherwise the pure-numpy path here, which is deterministic to the byte. Both converge to the same objective within the usual tolerance.

batch_size picks between full Lloyd and mini-batch k-means (Sculley 2010). Full Lloyd touches all n rows every iteration, which is fine with cores to spare and slow on a 0.5-2 vCPU container. Mini-batch samples batch_size rows per iteration and moves only the centroids they land on, so per-iteration work drops from O(n*k*d) to O(batch*k*d) at a small cost in inertia.

• None (default) — mini-batch above _MINIBATCH_MIN rows, full Lloyd below. The switch is on n only, so the same input and seed give the same labels on any machine.
• 0 — force full Lloyd at any size.
• a positive int — force mini-batch with that batch size.

Mini-batch needs the native kernel; without it the numpy path runs full Lloyd regardless.

n_init restarts the fit from that many seeds and keeps the lowest-inertia result. None (default) is a single run. Mini-batch can settle into a bad optimum on well-separated clusters (4.8x worse inertia than full Lloyd on adversarial blobs); it is fine on the overlapping data real embeddings produce (~1.01x), so the default stays a single fast run. If your data is well-separated, pass n_init=5 to recover full-Lloyd parity. Runtime scales with n_init.

dbscan

dbscan(
    embeddings: FloatArray,
    *,
    eps: float | None = None,
    min_samples: int = 5,
    metric: str = "euclidean",
    n_neighbors: int = 30,
) -> IntArray

Density clustering (DBSCAN). Finds clusters of any shape and labels outliers -1, with no n_clusters to choose.

A point is a core point if at least min_samples points (itself included) lie within eps of it. Core points within eps of each other join one cluster. A non-core point within eps of a core joins as a border. Everything else is noise (-1).

eps=None picks eps as the median distance to each point's min_samples-th neighbour (the usual k-distance heuristic). Neighbours come from fastumap's kNN (blocked exact, or the faiss approx backend above ~16k points), so this scales past a naive O(n^2) DBSCAN. Only the n_neighbors nearest are considered, so a very dense cluster may need a larger n_neighbors (which must exceed min_samples). Deterministic given the kNN.

hdbscan

hdbscan(
    embeddings: FloatArray,
    *,
    min_cluster_size: int = 5,
    min_samples: int | None = None,
    metric: str = "euclidean",
    n_neighbors: int = 30,
) -> IntArray

Hierarchical density clustering (HDBSCAN). Like :func:dbscan but with no eps to pick, so clusters of different densities are all found at once. Outliers are -1.

DBSCAN cuts the density hierarchy at one global eps: pick it loose and two nearby clusters merge, pick it tight and the sparse cluster becomes noise. HDBSCAN builds the whole hierarchy and then keeps the branches that persist longest over it, so each cluster is effectively cut at its own density.

Four steps, all over fastumap's kNN: the core distance of a point (its min_samples-th neighbour distance) inflates every edge into a mutual reachability distance max(core_i, core_j, d_ij); an MST over those edges gives the single-linkage hierarchy; branches smaller than min_cluster_size are condensed away as noise; what is left is scored by stability (excess of mass) and the best non-overlapping set wins.

min_cluster_size is the real knob: the smallest group you are willing to call a cluster. min_samples (default: min_cluster_size) sets how conservative the density estimate is — raise it to push more points into noise.

One consequence to know before you call it: the root of the hierarchy is never a cluster, so data with only one dense group in it comes back all -1. HDBSCAN answers "which groups stand out from each other", not "is there a group here". This matches scikit-learn's allow_single_cluster=False default.

This is approximate HDBSCAN: the MST is built over the n_neighbors kNN edges, not the full mutual-reachability graph (the same tradeoff RAPIDS makes), so a cluster whose points are all further apart than their n_neighbors-th neighbour can split. The kNN graph may also be disconnected; the pieces are joined at infinite distance, which is the honest answer for points that share no neighbourhood. Deterministic given the kNN.

choose_k

choose_k(
    embeddings: FloatArray,
    *,
    k_min: int = ...,
    k_max: int = ...,
    metric: str = ...,
    n_iter: int = ...,
    random_state: int = ...,
    max_samples: int = ...,
    n_jobs: int = ...,
    return_evidence: Literal[False] = ...,
) -> int

choose_k(
    embeddings: FloatArray,
    *,
    k_min: int = ...,
    k_max: int = ...,
    metric: str = ...,
    n_iter: int = ...,
    random_state: int = ...,
    max_samples: int = ...,
    n_jobs: int = ...,
    return_evidence: Literal[True],
) -> ChooseKResult

choose_k(
    embeddings: FloatArray,
    *,
    k_min: int = 2,
    k_max: int = 10,
    metric: str = "euclidean",
    n_iter: int = 50,
    random_state: int = 42,
    max_samples: int = 20000,
    n_jobs: int = 1,
    return_evidence: bool = False,
) -> int | ChooseKResult

Pick n_clusters automatically by the inertia elbow, the selection criterion here. It is the "kneedle" knee of the inertia-vs-k curve: the k past which adding clusters stops buying much tightness. This is separate from metric, which is the distance (euclidean or cosine).

Runs k-means for each k in [k_min, k_max] and returns the elbow. On more than max_samples rows it searches on a deterministic random subsample to bound the cost (the returned k still applies to the full data).

n_jobs runs the independent per-k fits in parallel. The fits are pinned to one BLAS thread each for the duration, so n_jobs of them run concurrently without oversubscribing; the result is identical to n_jobs=1. return_evidence=True returns a :class:ChooseKResult (the chosen k plus the per-candidate inertias and the sample size) instead of the bare int, so a caller can publish the evidence behind the choice.

It is a heuristic. For a clear cluster structure it lands on the true k give or take one. For gradual structure, treat it as a starting point.

assign_clusters

assign_clusters(
    embeddings: FloatArray,
    centroids: FloatArray,
    *,
    metric: str = ...,
    max_distance: float | None = ...,
    return_distances: Literal[False] = ...,
) -> IntArray

assign_clusters(
    embeddings: FloatArray,
    centroids: FloatArray,
    *,
    metric: str = ...,
    max_distance: float | None = ...,
    return_distances: Literal[True],
) -> tuple[IntArray, FloatArray]

assign_clusters(
    embeddings: FloatArray,
    centroids: FloatArray,
    *,
    metric: str = "euclidean",
    max_distance: float | None = None,
    return_distances: bool = False,
) -> IntArray | tuple[IntArray, FloatArray]

Assign each row of embeddings to its nearest of centroids. This is the predict half of a fit/predict split. New points get labels without re-clustering.

centroids come from a prior cluster(..., return_centroids=True) or kmeans run. metric must match the one they were fitted with ("cosine" unit-normalises the rows, as fitting did). Uses the same blocked nearest-centroid step as cluster, so peak memory is bounded by the block size, no matter how many rows you pass.

By default a point is always given its nearest centroid, however far that is. To stop far points being absorbed into whatever cluster is least far:

• max_distance labels any point whose nearest centroid is beyond it -1 (noise), the same outlier label the density methods use. The distance is in the space the fit ran in (unit-normalised rows for "cosine"), so the threshold is comparable to the inertia per point.
• return_distances=True also returns the distance to the assigned centroid for every row (the raw distance, even for -1 rows), so a caller can pick a threshold from the data.

Metrics

fastumap.metrics

Grading metrics: how much structure a projection preserves.

Every quality claim in this project is a number these functions produce. They are runtime-importable (numpy + scipy only) so a caller can grade a layout without pulling in umap-learn or scikit-learn.

neighbor_overlap

neighbor_overlap(
    original: FloatArray, projected: FloatArray, k: int
) -> float

Mean share of each point's k input-space neighbours kept after projection.

Chance level is about k/(n-1) — always compare against a random-layout control (see :func:random_layout).

global_distance_correlation

global_distance_correlation(
    original: FloatArray,
    projected: FloatArray,
    *,
    sample: int = 2000,
    random_state: int = 0,
) -> float

Spearman correlation of all pairwise distances, before vs after projection.

For n above sample a deterministic random subset is used so the condensed distance vector stays bounded (its length is quadratic in n).

random_layout

random_layout(
    n: int, dimensions: int, *, random_state: int = 0
) -> FloatArray

A Gaussian random layout — the control every overlap number is read against.

Environment

fastumap.layout

Initialisation and optimisation of the low-dimensional embedding.

Two steps:

• spectral_layout gives the deterministic starting positions. These are the leading eigenvectors of the normalised fuzzy graph, guarded against disconnected components.
• optimize_layout runs the attract/repel SGD. UMAP's cross-entropy objective reduces to that SGD.

accelerator_active

accelerator_active() -> bool

Whether the native fastumap._accel SGD kernel is present and will be used.

Analogous to torch.cuda.is_available(). The kernel is compiled into the wheel, so this is normally True. It is False only when the extension is not present, e.g. an unbuilt source checkout. Then the pure-numpy path runs instead.

The kernel does umap-learn's true in-place walk. So it produces a different (higher-quality) layout than the numpy fallback for the same seed.

Use this to record which path produced a stored projection. The default path uses the kernel whenever present; chunk_count > 1 forces numpy.

fastumap.neighbors

k-nearest-neighbour graph by blocked brute force.

Exact kNN, computed one row-block at a time so the n-by-n distance matrix is never materialised. At 5000 points that matrix would be 100MB and its argpartition index another 200MB. Blocking bounds peak memory to block_size-by-n. Brute force is the right call below ~10k points and needs no NN-descent dependency. Above that, use the approximate backend (roadmap #15).

Two metrics: euclidean (default) and cosine. Cosine is what text/CLS embeddings usually want, since euclidean on unnormalised encoder output is rarely the structure you care about. knn_between finds neighbours of one set inside another. That is what transform needs to place new points against a fitted training set.

ann_available

ann_available() -> bool

Whether the approximate-kNN backend (pip install fastumap[ann] → faiss) is present.