Edit on GitHub

starplast.gpu

GPU acceleration, offered where it helps and absent everywhere else.

Three backends, in the order they are preferred, because they cover different parts of the work:

cuml    UMAP, t-SNE, HDBSCAN, k-means and DBSCAN themselves
cupy    array work: scaling, ranking, pairwise distances
torch   the same array work, and far more likely to be installed already

None of them is a dependency. available() reports what is actually importable, and every function here falls back to the NumPy path it mirrors, so a machine with no GPU runs the same code it always did and a machine with one does not have to be configured. The switch in Preferences turns the whole thing off even where a backend exists, because "it was slow" and "it was wrong" are different complaints and the way to tell them apart is to rerun on the CPU.

A GPU result that differs from the CPU result is a bug, not a speed-up -- with two exceptions, both stated rather than discovered later:

  • The device path is single precision. A 3090 does float64 at a thirty-second of its float32 rate, so distances computed here carry a relative error of about 1e-4 against SciPy's float64. That is invisible in what they are used for -- the continuity proxy is a Spearman correlation of distances, and a rank does not move for a part in ten thousand -- and the test checks the score itself, not only the matrix.
  • cuml's UMAP is not umap-learn. A map built on the GPU is a different map of the same data, not the same map faster. It is announced in the log when it happens, because a walk whose rows came from both would be comparing the two libraries rather than the settings.

Ranking is exact either way: it is a permutation, and a permutation has no precision.

ENV_GPU = 'STARPLAST_GPU'
MIN_ELEMENTS = 120000
def available() -> dict:

What is importable right now: {"cuml": bool, "cupy": bool, "torch": bool, "device": str}.

def enabled() -> bool:

Whether GPU work is both wanted and possible.

Possible is whether anything imports. Wanted is the switch in Preferences -- and where nobody has touched that switch, wanted defaults to whether a backend is there at all: installing two gigabytes of CUDA wheels is a deliberate act, and making the user then find a setting to get what they installed is a way of wasting their afternoon.

An explicit choice always wins, in both directions, and survives a restart. Read on every call rather than cached: someone who installs cuml while the program is open should get it.

@contextlib.contextmanager
def pinned():

Hold the selected CPU/GPU implementation fixed for one complete search.

The Preferences switch remains usable, but a change affects the next search. The yielded backend snapshot is the exact record to put in that search's manifest.

def backend() -> dict:

Which library will do each job, with its version -- the record a run has to carry.

A map built by cuml's UMAP is a different map of the same data, so two rows of one table built by two implementations would be a comparison of the libraries wearing the look of a comparison of settings. Every saved run and every results row records this.

def backend_id() -> str:

The backend as one short string, for a table column: "cuml 26.8.0" or "umap-learn 0.5.9".

def describe() -> str:

One line for the interface and for a saved recipe: what will do the work.

def worth_it(a) -> bool:

Whether an array is large enough that moving it to the device pays for itself.

def pairwise_distances(X):

Euclidean distances, on the device when that is worth doing.

The walk's own cost centre: trustworthiness and the continuity proxy are computed from the full distance matrix of the subsample, which is O(n^2) in both time and memory, and it is recomputed for every configuration.

def umap_class():

cuml's UMAP when it is there and wanted, else None so the caller uses umap-learn.

Named rather than wrapped: cuml's UMAP takes the same arguments this project passes, and a wrapper that silently mapped them would be one more place for the two paths to disagree.

def hdbscan_class():

cuml's HDBSCAN when it is there and wanted, else None.

def kmeans_class():

cuml's k-means when it is there and wanted, else None.

def dbscan_class():

cuml's DBSCAN when it is there and wanted, else None.

def tsne_class():

cuml's t-SNE when it is there and wanted, else None.