umap

package module
v0.0.0-...-79bd843 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Mar 1, 2026 License: BSD-3-Clause Imports: 10 Imported by: 0

README

umap-go

umap-go is a pure-Go implementation of Uniform Manifold Approximation and Projection (UMAP).

This project was built with a strict mandate: achieve exact mathematical parity with the reference Python umap-learn library (pinned to v0.5.11) and its underlying pynndescent neighbor search. It is designed to be highly accurate, reproducible, and dependency-free (pure Go, no CGO).

Features

  • Pure Go: No reliance on CGO, Cython, or external C libraries.
  • Mathematical Parity: Built to replicate the exact float math, graph constructions, and algorithmic stages of Python's UMAP, including exact PRNG sequence replication for reproducibility.
  • Complete Fit & Transform: Supports training a model (Fit / FitTransform) and mapping new, out-of-sample data into an existing embedding space (Transform).
  • Custom NN-Descent: Includes a fully replicated PyNNDescent implementation, including RP-Trees, graph-informed Hub Trees, and constrained graph search for exact nearest neighbor approximations.
  • Metrics: Supports all standard UMAP distance metrics (Euclidean, Cosine, etc.).

Installation

go get github.com/nozzle/umap-go

Capability Matrix

Capability Status Notes
Fit(X) / FitTransform(X, nil) Supported Unsupervised training and embedding generation.
FitTransform(X, y) Supported Supervised mode (y []float64) via target-graph intersection.
Transform(XNew) Supported Out-of-sample mapping after Fit/FitTransform.
InverseTransform(XEmbedded) Deferred Method currently returns umap: InverseTransform not yet implemented; deferred while a pure-Go replacement for Python's Delaunay/QHull path is unresolved.

Reproducibility (seeded RandSource)

Use a fixed seed for deterministic runs:

seed := uint64(42)
opts := umap.DefaultOptions()
opts.RandSource = umaprand.NewProduction(&seed)
model := umap.New(opts)

For repeatable independent runs, create a fresh RandSource from the same seed per run (do not reuse a previously advanced source).

Usage

package main

import (
	"github.com/nozzle/umap-go"
	umaprand "github.com/nozzle/umap-go/rand"
)

func seededOptions(seed uint64) umap.Options {
	opts := umap.DefaultOptions()
	opts.RandSource = umaprand.NewProduction(&seed)
	return opts
}
1) Unsupervised Fit and FitTransform
var X [][]float64

modelA := umap.New(seededOptions(42))
if err := modelA.Fit(X); err != nil {
	panic(err)
}
embeddingA := modelA.Embedding()

modelB := umap.New(seededOptions(42))
embeddingB, err := modelB.FitTransform(X, nil)
if err != nil {
	panic(err)
}
_ = embeddingA
_ = embeddingB
2) Supervised FitTransform
var X [][]float64
var y []float64 // class labels or continuous targets

model := umap.New(seededOptions(42))
embedding, err := model.FitTransform(X, y)
if err != nil {
	panic(err)
}
_ = embedding
3) Out-of-sample Transform
var XTrain, XNew [][]float64

model := umap.New(seededOptions(42))
_, err := model.FitTransform(XTrain, nil)
if err != nil {
	panic(err)
}

newEmbedding, err := model.Transform(XNew)
if err != nil {
	panic(err)
}
_ = newEmbedding

Parallel execution

The library can use multiple CPU cores for parallel-capable stages.

  • Options.NWorkers: number of workers (default: runtime.GOMAXPROCS(0))
  • Options.ParallelMode:
    • "auto" (default): use deterministic parallel paths where available
    • "serial": force single-thread behavior
    • "parallel": favor throughput for parallel-capable paths

Example:

opts := umap.DefaultOptions()
opts.NWorkers = runtime.GOMAXPROCS(0) // auto CPU count
opts.ParallelMode = "auto"

Python vs Go FitTransform benchmark

This repository includes a benchmark comparison for FitTransform:

  • Python runner: testdata/benchmark_fit_transform.py (umap-learn==0.5.11)
  • Go benchmark: BenchmarkFitTransformCompare
  • Parallel-mode benchmark: BenchmarkFitTransformParallelModes
  • Worker-count benchmark: BenchmarkFitTransformWorkerCounts
  • Python worker-count runner: testdata/benchmark_fit_transform_worker_counts.py

Run from repository root:

go test -run '^$' -bench BenchmarkFitTransformCompare -benchmem .
go test -run '^$' -bench BenchmarkFitTransformParallelModes -benchmem .
go test -run '^$' -bench BenchmarkFitTransformWorkerCounts -benchmem .
python3 testdata/benchmark_worker_counts_readme.py --update-readme

When Python tooling is available (uv or python3 with testdata dependencies), the Go benchmark reports extra metrics:

  • py_ns/op: Python mean nanoseconds per operation
  • go_ns/op: Go mean nanoseconds per operation
  • py/go: Python-to-Go time ratio

If Python dependencies are unavailable, the benchmark still runs and reports Go-only metrics.

Worker count scaling (FitTransform)

The section below is generated by testdata/benchmark_worker_counts_readme.py from BenchmarkFitTransformWorkerCounts and Python umap-learn worker-count runs.

For Python worker scaling, random_state is intentionally unset because umap-learn forces n_jobs=1 when random_state is set.

Environment: darwin/arm64 on Apple M3 Max; Go command: go test -run ^$ -bench ^BenchmarkFitTransformWorkerCounts$ -benchmem -count=1 .; Python command: uv run --directory testdata python benchmark_fit_transform_worker_counts.py.

Workers Go ns/op Go speedup vs 1 Python ns/op Python speedup vs 1 Python/Go Go B/op Go allocs/op
1 176,903,696 1.00x 77,777,983 1.00x 0.44x 12,967,797 10,613
2 112,461,212 1.57x 56,082,200 1.39x 0.50x 20,300,187 42,002
4 83,422,583 2.12x 47,916,275 1.62x 0.57x 21,419,971 43,003
8 71,069,789 2.49x 45,046,549 1.73x 0.63x 23,658,834 45,001
auto 73,158,011 2.42x 46,056,166 1.69x 0.63x 26,943,766 47,999 
xychart-beta
    title "FitTransform ns/op by worker count (Go vs Python; lower is better)"
    x-axis ["1", "2", "4", "8", "auto"]
    y-axis "ns/op" 0 --> 194594065
    bar [176903696, 112461212, 83422583, 71069789, 73158011]
    bar [77777983, 56082200, 47916275, 45046549, 46056166]

Mermaid bar order: first series is Go, second series is Python (umap-learn).

License

This project is licensed under the BSD 3-Clause License. See LICENSE.

Credits / Upstream Attribution

This library is a pure-Go port built for mathematical parity with:

umap-go is an independent implementation and is not affiliated with or endorsed by the upstream projects.

Documentation

Index

Constants

This section is empty.

Variables

This section is empty.

Functions

func ClipValue 

func ClipValue(x, clipLimit float64) float64

ClipValue clamps a value to [-clipLimit, clipLimit]. Used in SGD gradient updates. Matches UMAP's clip().

func ComputeMembershipStrengths 

func ComputeMembershipStrengths(knnIndices [][]int, knnDists [][]float64,
	sigmas, rhos []float64, nSamples int) *sparse.COO

ComputeMembershipStrengths computes the sparse fuzzy set (COO matrix) from kNN indices, distances, sigmas, and rhos.

For each point i and neighbor j: 

w_ij = exp(-(d_ij - rho_i) / sigma_i)

Matches umap_.py compute_membership_strengths().

func ComputeMembershipStrengthsBipartite 

func ComputeMembershipStrengthsBipartite(knnIndices [][]int, knnDists [][]float64,
	sigmas, rhos []float64, nQueries, nTrain int) *sparse.COO

ComputeMembershipStrengthsBipartite is like ComputeMembershipStrengths but for out-of-sample mapping. It creates a bipartite graph of shape (nQueries x nTrain).

func EpochsOfNextNegativeSample 

func EpochsOfNextNegativeSample(epochsPerSample []float64, negativeSampleRate float64) []float64

EpochsOfNextNegativeSample precomputes the epoch at which each edge should next receive a negative sample.

func EpochsOfNextSample 

func EpochsOfNextSample(epochsPerSample []float64) []float64

EpochsOfNextSample precomputes the epoch at which each edge should first be sampled. Used to drive the SGD loop.

func FindABParams 

func FindABParams(spread, minDist float64) (float64, float64)

FindABParams computes the a, b parameters for the UMAP membership function via curve fitting (Levenberg-Marquardt style).

The target curve is: 

y = 1.0  if x <= min_dist
y = exp(-(x - min_dist) / (spread - min_dist))  otherwise

And we fit: y = 1 / (1 + a * x^(2b))

Returns (a, b).

func MakeEpochsPerSample 

func MakeEpochsPerSample(weights []float64, nEpochs int) []float64

MakeEpochsPerSample computes the number of epochs between samples for each edge weight.

Given weights and n_epochs: 

n_samples = n_epochs * (weight / max_weight)
epochs_per_sample = n_epochs / n_samples

Which simplifies to: epochs_per_sample = max_weight / weight

Weights that are zero or negative get -1 (never sampled).

Matches umap_.py make_epochs_per_sample() exactly.

func OptimizeLayoutEuclidean 

func OptimizeLayoutEuclidean(
	headEmbedding [][]float64,
	tailEmbedding [][]float64,
	head []int,
	tail []int,
	epochsPerSample []float64,
	rngStates []nn.TauRandState,
	cfg OptimizeLayoutConfig,
) [][]float64

OptimizeLayoutEuclidean performs SGD optimization of the embedding layout.

Parameters: 

headEmbedding: initial embedding positions (modified in place)
tailEmbedding: reference embedding (same as headEmbedding for fit, different for transform)
head: source indices of edges
tail: target indices of edges
epochsPerSample: how many epochs between samples for each edge
rngState: per-sample Tausworthe PRNG states [nSamples][3]
cfg: optimization configuration

Returns the optimized embedding (same slice as headEmbedding).

Matches umap/layouts.py optimize_layout_euclidean().

func OptimizeLayoutGeneric 

func OptimizeLayoutGeneric(
	headEmbedding [][]float64,
	tailEmbedding [][]float64,
	head []int,
	tail []int,
	epochsPerSample []float64,
	rngStates []nn.TauRandState,
	cfg OptimizeLayoutConfig,
	outputMetricGrad func(x, y []float64) (float64, []float64),
) [][]float64

OptimizeLayoutGeneric performs SGD with a custom output distance metric. Used for inverse_transform or non-Euclidean output spaces.

TODO: Implement this for inverse_transform support.

func Rdist 

func Rdist(x, y []float64) float64

Rdist computes the squared Euclidean distance between two vectors. Used in the SGD inner loop. Matches UMAP's rdist().

func SelectNEpochs 

func SelectNEpochs(nSamples int) int

SelectNEpochs chooses the number of optimization epochs based on dataset size if the user hasn't specified one. Matches umap_.py default n_epochs selection.

func SpectralLayout 

func SpectralLayout(graph *sparse.CSR, nComponents int) [][]float64

SpectralLayout computes the spectral initialization embedding.

Algorithm: 1. Compute the normalized graph Laplacian: L = D^{-1/2} (D - A) D^{-1/2} 2. Find the smallest nComponents+1 eigenvectors 3. Discard the trivial eigenvector (constant, eigenvalue ~0) 4. Return the next nComponents eigenvectors, scaled by sqrt(eigenvalue)

If the graph has multiple connected components, each component is handled independently and results are combined.

graph: the symmetrized fuzzy simplicial set (CSR) nComponents: number of embedding dimensions

Types

type FuzzySimplicialSetConfig 

type FuzzySimplicialSetConfig struct {
	NWorkers     int
	ParallelMode string
}

FuzzySimplicialSetConfig holds internal execution controls.

type FuzzySimplicialSetResult 

type FuzzySimplicialSetResult struct {
	Graph       *sparse.CSR     // the symmetrized fuzzy simplicial set graph
	Sigmas      []float64       // per-point sigma values
	Rhos        []float64       // per-point rho values
	SearchIndex *nn.SearchIndex // the nearest neighbor search index
}

FuzzySimplicialSetResult holds the output of FuzzySimplicialSet.

func FuzzySimplicialSet 

func FuzzySimplicialSet(
	data [][]float64,
	nNeighbors int,
	rng umaprand.Source,
	metric string,
	metricKwds map[string]any,
	localConnectivity float64,
	setOpMixRatio float64,
) *FuzzySimplicialSetResult

FuzzySimplicialSet constructs the fuzzy simplicial set from raw data. This is the main graph construction pipeline:

1. Compute kNN (brute-force or NN-Descent) 2. Smooth kNN distances to get sigma, rho 3. Compute membership strengths → sparse matrix 4. Symmetrize via fuzzy set union: W + W^T - W*W^T 5. Reset local connectivity

Matches umap_.py fuzzy_simplicial_set().

func FuzzySimplicialSetWithConfig 

func FuzzySimplicialSetWithConfig(
	data [][]float64,
	nNeighbors int,
	rng umaprand.Source,
	metric string,
	metricKwds map[string]any,
	localConnectivity float64,
	setOpMixRatio float64,
	cfg FuzzySimplicialSetConfig,
) *FuzzySimplicialSetResult

FuzzySimplicialSetWithConfig constructs the fuzzy simplicial set with execution config.

type OptimizeLayoutConfig 

type OptimizeLayoutConfig struct {
	A                  float64 // UMAP a parameter
	B                  float64 // UMAP b parameter
	Gamma              float64 // repulsive force weight (default: 1.0)
	InitialAlpha       float64 // initial learning rate (default: 1.0)
	NegativeSampleRate float64 // negative samples per positive (default: 5.0)
	NEpochs            int     // number of optimization epochs
	MoveOther          bool    // update both head and tail embeddings
	NWorkers           int
	ParallelMode       string
}

OptimizeLayoutConfig holds the configuration for layout optimization.

type Options 

type Options struct {
	// NNeighbors is the number of nearest neighbors to use for graph construction.
	// Larger values capture more global structure at the cost of local detail.
	// Default: 15.
	NNeighbors int

	// NComponents is the dimensionality of the output embedding.
	// Default: 2.
	NComponents int

	// MinDist controls how tightly points are packed together.
	// Smaller values produce more clustered embeddings.
	// Default: 0.1.
	MinDist float64

	// Spread determines the scale of the embedding.
	// Together with MinDist, controls the membership function.
	// Default: 1.0.
	Spread float64

	// Metric is the distance metric to use on the input data.
	// Default: "euclidean".
	Metric string

	// MetricKwds are additional parameters for parameterized metrics
	// (e.g., "p" for Minkowski, "sigma" for StandardisedEuclidean).
	MetricKwds map[string]any

	// NEpochs is the number of SGD optimization epochs.
	// If 0, automatically chosen based on dataset size.
	NEpochs int

	// InitMethod controls the initial embedding.
	// "spectral" (default), "random", or a custom [][]float64.
	InitMethod string

	// CustomInit provides a pre-computed initial embedding.
	// Only used if InitMethod == "custom".
	CustomInit [][]float64

	// LocalConnectivity is the number of nearest neighbors that should
	// be assumed to be connected at a local level.
	// Default: 1.0.
	LocalConnectivity float64

	// SetOpMixRatio controls the blend between fuzzy union and intersection
	// for the symmetrization of the kNN graph.
	// 1.0 = pure union (default), 0.0 = pure intersection.
	SetOpMixRatio float64

	// DisconnectionDistance removes edges with distances greater than or equal to this value.
	// Default: +inf (no disconnection).
	DisconnectionDistance float64

	// NegativeSampleRate controls the number of negative samples per
	// positive sample in SGD optimization.
	// Default: 5.
	NegativeSampleRate float64

	// RepulsionStrength controls the weight of the repulsive force.
	// Default: 1.0.
	RepulsionStrength float64

	// LearningRate is the initial SGD learning rate.
	// Default: 1.0.
	LearningRate float64

	// RandSource provides the random number generator.
	// Default: Production (wrapping math/rand/v2 with seed 42).
	RandSource umaprand.Source

	// TargetNNeighbors is the number of nearest neighbors for the target
	// (y) space in supervised mode.
	// Default: NNeighbors.
	TargetNNeighbors int

	// TargetMetric is the distance metric for the target space.
	// Default: "categorical" for discrete labels, "euclidean" for continuous.
	TargetMetric string

	// TargetWeight controls the balance between data topology and target
	// topology in supervised mode. 0.0 = data only, 1.0 = target only.
	// Default: 0.5.
	TargetWeight float64

	// Verbose controls whether progress information is printed.
	Verbose bool

	// NWorkers controls the number of worker goroutines used by parallel-capable stages.
	// If 0, runtime.GOMAXPROCS(0) is used.
	// Default: runtime.GOMAXPROCS(0).
	NWorkers int

	// ParallelMode controls how parallel-capable stages execute.
	// "auto" (default), "serial", or "parallel".
	ParallelMode string
}

Options configures the UMAP algorithm.

func DefaultOptions 

func DefaultOptions() Options

DefaultOptions returns Options with all defaults set.

type SmoothKNNResult 

type SmoothKNNResult struct {
	Sigmas []float64 // per-point bandwidth parameters
	Rhos   []float64 // per-point nearest neighbor distances
}

SmoothKNNResult holds the output of SmoothKNNDist.

func SmoothKNNDist 

func SmoothKNNDist(knnDists [][]float64, k float64, localConnectivity float64) *SmoothKNNResult

SmoothKNNDist computes the smooth nearest-neighbor distance parameters (sigma and rho) for each point using binary search.

For each point, we find sigma such that: 

sum(exp(-(d_i - rho) / sigma)) = log2(k)

where d_i are the distances to k nearest neighbors and rho is the distance to the nearest neighbor (local connectivity adjustment).

Matches umap_.py smooth_knn_dist().

type UMAP 

type UMAP struct {
	// contains filtered or unexported fields
}

UMAP is the main UMAP dimensionality reduction model.

func New 

func New(opts Options) *UMAP

New creates a new UMAP model with the given options.

func (*UMAP) A 

func (u *UMAP) A() float64

A returns the fitted 'a' parameter.

func (*UMAP) B 

func (u *UMAP) B() float64

B returns the fitted 'b' parameter.

func (*UMAP) Embedding 

func (u *UMAP) Embedding() [][]float64

Embedding returns the fitted embedding. Nil if not fitted.

func (*UMAP) Fit 

func (u *UMAP) Fit(X [][]float64) error

Fit fits the UMAP model to the data X without returning the embedding. X is a matrix of shape (n_samples, n_features) stored as [][]float64.

func (*UMAP) FitTransform 

func (u *UMAP) FitTransform(X [][]float64, y []float64) ([][]float64, error)

FitTransform fits the UMAP model and returns the embedding. X: input data (n_samples x n_features) y: optional target labels for supervised mode (nil for unsupervised)

func (*UMAP) Graph 

func (u *UMAP) Graph() *sparse.CSR

Graph returns the fitted fuzzy simplicial set graph. Nil if not fitted.

func (*UMAP) InverseTransform 

func (u *UMAP) InverseTransform(XEmbedded [][]float64) ([][]float64, error)

InverseTransform maps points from the embedding space back to data space.

TODO: Implement inverse transform.

func (*UMAP) Transform 

func (u *UMAP) Transform(XNew [][]float64) ([][]float64, error)

Transform projects new data into the existing embedding space. The model must be fitted first.

Source Files

View all Source files

Directories

Path Synopsis
Package distance provides distance metric implementations for UMAP.
 Package distance provides distance metric implementations for UMAP.
Package nn implements nearest neighbor search for UMAP.
 Package nn implements nearest neighbor search for UMAP.
Package rand provides a random number source abstraction for UMAP.
 Package rand provides a random number source abstraction for UMAP.
Package sparse provides COO and CSR sparse matrix types for UMAP.
 Package sparse provides COO and CSR sparse matrix types for UMAP.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.