From 41ed4184bce75a59381e661a956648b2e43a033c Mon Sep 17 00:00:00 2001 From: Ao Zhang Date: Mon, 26 Oct 2026 18:19:00 +0000 Subject: [PATCH] Compute pairwise_distance with matrix multiplications Write both metrics as matrix products of the values, with missing entries set to zero, and the mask of observed entries. For example, the squared Euclidean distance over the entries observed in both rows is sum(x_i^2) + sum(x_j^2) - 2 * sum(x_i * x_j), where each sum is a matrix product. Build them from Dask array operations instead of Numba gufuncs, CUDA kernels and hand-written blockwise and reduction steps, with each metric as a function in sgkit.distance.metrics, looked up by name in pairwise_distance, so that new metrics can be added there. This is several times faster on the CPU, and no longer passes any functions for Dask to tokenize (see issue 1267). The GPU path now uses CuPy instead of Numba's CUDA support: arrays with CuPy chunks are computed on the GPU, and device="gpu" moves the chunks to the GPU and the result back to NumPy. This fixes wrong results for int8 input on the GPU, where the partial sums were stored in an int8 array. The GPU workflow now installs CuPy for CUDA 12, with the CUDA libraries it uses from PyPI, instead of the CUDA toolkit from conda, which needs NVIDIA driver 525 or later on the runner, and checks that CuPy can use the GPU before running the tests, which are skipped otherwise. The distances are computed in float64, and returned in the floating point type of the input, or float64 for other types, so integer input now reports the float64 dtype of the result before it is computed. Subtracting the sums of squares and products loses small distances in float32: with float32 dosages, two rows 0.11 apart could come out 0 apart, and copies of a row could have a negative correlation distance. In float64 these sums are exact for small integers, such as allele counts. On GPUs with slow float64, such as the T4 that the GPU workflow runs on, float32 input now takes as long as float64 input. Infinite values are now missing, like negative and nan values, since an infinite value times zero is nan in the matrix products, which made every distance of its row nan. The correlation distance is clipped to [0, 2], and its denominator is sqrt(var_i * var_j), which unlike sqrt(var_i) * sqrt(var_j) keeps copies of a row of integers exactly zero apart. The diagonal is now exactly zero, or nan where the correlation is undefined. This needs Dask 2024.12 or later: the blockwise fusion in earlier releases computed wrong correlations from these expressions. --- .github/workflows/build-gpu.yml | 13 +- docs/changelog.rst | 16 +- docs/contributing.rst | 7 +- sgkit/distance/api.py | 147 ++++------- sgkit/distance/metrics.py | 447 +++----------------------------- sgkit/tests/test_distance.py | 144 +++++++++- 6 files changed, 244 insertions(+), 530 deletions(-) diff --git a/.github/workflows/build-gpu.yml b/.github/workflows/build-gpu.yml index b5744eba1..42efe71e5 100644 --- a/.github/workflows/build-gpu.yml +++ b/.github/workflows/build-gpu.yml @@ -45,15 +45,16 @@ jobs: conda info conda list - - name: Installing cudatoolkit and dependencies + - name: Install dependencies run: | - conda install -c nvidia cudatoolkit - # Avoid building SciPy from source in CI. - uv pip install --only-binary scipy --python "$CONDA_PREFIX/bin/python" -e . --group test + # Avoid building SciPy from source in CI, and install the CUDA + # libraries that CuPy uses with its ctk extra. + uv pip install --only-binary scipy --python "$CONDA_PREFIX/bin/python" -e . --group test "cupy-cuda12x[ctk]" - - name: Numba Information + - name: CuPy information + # Fail here if CuPy can't use the GPU, since the GPU tests would be skipped. run: | - numba -s + python -c "import cupy; cupy.show_config(); print(cupy.ones((2, 2)) @ cupy.ones((2, 2)))" - name: Run GPU tagged tests run: | diff --git a/docs/changelog.rst b/docs/changelog.rst index b69f4860c..7f7a29eb0 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -35,6 +35,10 @@ Breaking changes Use bio2zarr for conversion to Zarr and vcztools for export. See :ref:`plink_conversion_migration`, including companion-file metadata differences. (:pr:`1324`, :pr:`1362`, :issue:`1318`) +- :func:`pairwise_distance` with ``device="gpu"`` now requires CuPy instead of + Numba's CUDA support. +- :func:`pairwise_distance` treats infinite values as missing, like negative + and nan values. Deprecations ~~~~~~~~~~~~ @@ -53,8 +57,12 @@ Deprecations `vcztools `_ packages instead. (:user:`tomwhite`, :pr:`1324`, :issue:`1318`) -.. Improvements -.. ~~~~~~~~~~~~ +Improvements +~~~~~~~~~~~~ + +- :func:`pairwise_distance` is computed with matrix multiplications in float64 + instead of Numba functions, which is several times faster on the CPU, and + arrays with CuPy chunks are computed on the GPU. (:issue:`1267`) Bug fixes ~~~~~~~~~ @@ -70,6 +78,10 @@ Bug fixes (:user:`timothymillar`, :pr:`1335`, :issue:`1333`) - :func:`Tajimas_D` now works with newer versions of Dask, and no longer computes allele counts when it is called. +- :func:`pairwise_distance` no longer gives wrong results for int8 input with + ``device="gpu"``, the dtype it reports for integer and float16 input now + matches the computed result, and correlation distances are no longer + slightly negative. .. Documentation .. ~~~~~~~~~~~~~ diff --git a/docs/contributing.rst b/docs/contributing.rst index e7aba6275..33c9a9e0b 100644 --- a/docs/contributing.rst +++ b/docs/contributing.rst @@ -99,9 +99,10 @@ the dependencies needed to test compatibility with bio2zarr. Use ``uv run`` to r commands in this environment without activating it. -If you have a Nvidia GPU you will need to make sure that it is configured properly, -as in you have cudatoolkit installed, the instructions for the same can be found on -`nvidia docs. `_ +The tests marked ``gpu`` run on an NVIDIA GPU with CuPy, and are skipped if CuPy +or a GPU is not available. To run them, install the CuPy package that matches +your CUDA version, as described in the +`CuPy installation guide `_. Install the pre-commit hook, which is used to enforce coding standards:: diff --git a/sgkit/distance/api.py b/sgkit/distance/api.py index fb0518584..b95ba0edf 100644 --- a/sgkit/distance/api.py +++ b/sgkit/distance/api.py @@ -1,9 +1,11 @@ import typing from typing import Literal +import dask import dask.array as da import numpy as np +from sgkit.distance.metrics import METRICS from sgkit.typing import ArrayLike MetricTypes = Literal["euclidean", "correlation"] @@ -46,8 +48,17 @@ def pairwise_distance( The (i, j) position in the resulting array (matrix) denotes the distance between vi and vj vectors. - Negative and nan values are considered as missing values. They are ignored - for all distance metric calculations. + Negative, nan and infinite values are considered as missing values. They + are ignored for all distance metric calculations. + + The distances are computed with matrix multiplications in float64, and + returned in the floating point type of ``x``, or in float64 if ``x`` has + another type. They are computed from sums of squares and products, which + are exact for small integers such as allele counts. For other values, a + distance much smaller than the values it is computed from, such as between + two nearly identical rows, can be inaccurate. Arrays with CuPy chunks are + computed on the GPU, where float64 is much slower than float32 on many + GPUs. Parameters ---------- @@ -60,13 +71,14 @@ def pairwise_distance( 'euclidean' or 'correlation'. split_every Determines the depth of the recursive aggregation in the reduction - step. This argument is directly passed to the call to``dask.reduction`` - function in the reduce step of the map reduce. + step over the chunks of the second dimension of ``x``. Omit to let dask heuristically decide a good default. A default can also be set globally with the split_every key in dask.config. device - The architecture to run the calculation on, either of "cpu" or "gpu" + The architecture to run the calculation on, either of "cpu" or "gpu". + With "gpu", the chunks of ``x`` are moved to the GPU with CuPy, and + the result is moved back to NumPy, so this requires CuPy. Returns ------- @@ -97,106 +109,45 @@ def pairwise_distance( >>> x = np.array([[6, 4, 1,], [4, 5, 2], [9, 7, 3]]) >>> pairwise_distance(x, metric='correlation').compute() - array([[-4.44089210e-16, 2.62956526e-01, 2.82353505e-03], - [ 2.62956526e-01, 0.00000000e+00, 2.14285714e-01], - [ 2.82353505e-03, 2.14285714e-01, 0.00000000e+00]]) + array([[0. , 0.26295653, 0.00282354], + [0.26295653, 0. , 0.21428571], + [0.00282354, 0.21428571, 0. ]]) """ - from sgkit.distance import metrics # defer numba compilation - valid_devices = DeviceTypes.__args__ # type: ignore[attr-defined] if device not in valid_devices: raise ValueError( f"Invalid Device, expected one of {valid_devices}, got: {device}" ) - try: - map_func_name = f"{metric}_map_{device}" - reduce_func_name = f"{metric}_reduce_{device}" - map_func = getattr(metrics, map_func_name) - reduce_func = getattr(metrics, reduce_func_name) - n_map_param = metrics.N_MAP_PARAM[metric] - except AttributeError: - raise NotImplementedError( - f"Given metric: '{metric}' is not implemented for '{device}'." - ) + if metric not in METRICS: + raise NotImplementedError(f"Given metric: '{metric}' is not implemented.") x = da.asarray(x) if x.ndim != 2: raise ValueError(f"2-dimensional array expected, got '{x.ndim}'") - - # setting this variable outside of _pairwise to avoid it's recreation - # in every iteration, which eventually leads to increase in dask - # graph serialisation/deserialisation time significantly - metric_param = np.empty(n_map_param, dtype=x.dtype) - - def _pairwise_cpu(f: ArrayLike, g: ArrayLike) -> ArrayLike: - result: ArrayLike = map_func(f[:, None, :], g, metric_param) - # Adding a new axis to help combine chunks along this axis in the - # reduction step (see the _aggregate and _combine functions below). - return result[..., np.newaxis] - - def _pairwise_gpu(f: ArrayLike, g: ArrayLike) -> ArrayLike: # pragma: no cover - result = map_func(f, g) - return result[..., np.newaxis] - - pairwise_func = _pairwise_cpu - if device == "gpu": - pairwise_func = _pairwise_gpu # pragma: no cover - - # concatenate in blockwise leads to high memory footprints, so instead - # we perform blockwise without contraction followed by reduction. - # More about this issue: https://github.com/dask/dask/issues/6874 - out = da.blockwise( - pairwise_func, - "ijk", - x, - "ik", - x, - "jk", - dtype=x.dtype, - concatenate=False, - ) - - def _aggregate(x_chunk: ArrayLike, **_: typing.Any) -> ArrayLike: - """Last function to be executed when resolving the dask graph, - producing the final output. It is always invoked, even when the reduced - Array counts a single chunk along the reduced axes. - - Parameters - ---------- - x_chunk - [array-like, shape: (M, M, N, 1)] - An array like two dimensional matrix. The dimension is as follows: - M is the chunk size on axis=0 for `x` (input for `pairwise_distance` - function). - N: is the number of chunks along axis=1 - """ - x_chunk = x_chunk.reshape(x_chunk.shape[:-2] + (-1, n_map_param)) - result: ArrayLike = reduce_func(x_chunk) - return result - - def _chunk(x_chunk: ArrayLike, **_: typing.Any) -> ArrayLike: - return x_chunk - - def _combine(x_chunk: ArrayLike, **_: typing.Any) -> ArrayLike: - """Function used for intermediate recursive aggregation (see - split_every argument to ``da.reduction below``). If the - reduction can be performed in less than 3 steps, it will - not be invoked at all.""" - # reduce chunks by summing along the -2 axis - x_chunk_reshaped = x_chunk.reshape(x_chunk.shape[:-2] + (-1, n_map_param)) - return x_chunk_reshaped.sum(axis=-2)[..., np.newaxis] - - r = da.reduction( - out, - chunk=_chunk, - combine=_combine, - aggregate=_aggregate, - axis=-1, - dtype=x.dtype, - meta=np.ndarray((0, 0), dtype=x.dtype), - split_every=split_every, - name="pairwise", - ) - - t = da.triu(r) - return t + t.T + if device == "gpu": # pragma: no cover + import cupy + + x = x.map_blocks(cupy.asarray) + + dtype = x.dtype if x.dtype.kind == "f" else np.dtype(np.float64) + # false for negative, nan and infinite values + observed = (x >= 0) & da.isfinite(x) + # zero the missing values so they drop out of the matrix products in the + # metrics, which are computed in at least float64: small distances are + # lost when the sums of squares and products are subtracted in float32 + compute_dtype = np.promote_types(dtype, np.float64) + x = da.where(observed, x, 0).astype(compute_dtype) + observed = observed.astype(compute_dtype) + # split_every controls the sums over the column chunks inside da.matmul + config = {} if split_every is None else {"split_every": split_every} + with dask.config.set(config): + d = METRICS[metric](x, observed).astype(dtype) + + # make the result exactly symmetric, with zeros on the diagonal (or nan + # where the distance of a vector to itself is undefined) + u = da.triu(d, 1) + diagonal = da.where(da.isnan(da.diagonal(d)), np.nan, 0).astype(dtype) + d = u + u.T + da.diag(diagonal) + if device == "gpu": # pragma: no cover + d = d.map_blocks(cupy.asnumpy) + return d diff --git a/sgkit/distance/metrics.py b/sgkit/distance/metrics.py index 167c49e73..e0cda0444 100644 --- a/sgkit/distance/metrics.py +++ b/sgkit/distance/metrics.py @@ -1,419 +1,44 @@ """ This module implements various distance metrics. To implement a new distance -metric, two methods needs to be written, one of them suffixed by 'map' and other -suffixed by 'reduce'. An entry for the same should be added in the N_MAP_PARAM -dictionary below. +metric, write it with Dask array operations, such as matrix products, so that it +runs on NumPy chunks on the CPU and on CuPy chunks on the GPU. Each metric takes +the values, with the missing values set to zero, and the mask of the observed +values, both in float64. """ -import math -from typing import Any - +import dask.array as da import numpy as np -from numba import cuda, types -from sgkit.accelerate import numba_guvectorize -from sgkit.typing import ArrayLike -# The number of parameters for the map step of the respective distance metric -N_MAP_PARAM = { - "correlation": 6, - "euclidean": 1, +def euclidean(x: da.Array, observed: da.Array) -> da.Array: + """Euclidean distance between all pairs of rows of ``x``.""" + # Sum of squared differences over the entries observed in both vectors, + # expanded as sum(x_i^2) + sum(x_j^2) - 2 * sum(x_i * x_j). + a = (x * x) @ observed.T + d2 = a + a.T - 2 * (x @ x.T) + # rounding can make d2 slightly negative when two vectors are nearly equal + return da.sqrt(da.maximum(d2, 0)) + + +def correlation(x: da.Array, observed: da.Array) -> da.Array: + """One minus the Pearson correlation between all pairs of rows of ``x``, or + nan where the correlation is undefined.""" + n = observed @ observed.T + sx = x @ observed.T # sum of x_i, so sum of x_j is sx.T + sxx = (x * x) @ observed.T # sum of x_i^2, so sum of x_j^2 is sxx.T + sxy = x @ x.T + num = n * sxy - sx * sx.T + var = da.maximum(n * sxx - sx * sx, 0) + # sqrt(v * v) is exactly v, unlike sqrt(v) * sqrt(v), so copies of a row of + # integers are exactly zero apart + den = da.sqrt(var * var.T) + defined = den > 0 + # rounding can put the distance slightly outside [0, 2] + d = da.clip(1 - num / da.where(defined, den, 1), 0, 2) + return da.where(defined, d, np.nan) + + +METRICS = { + "euclidean": euclidean, + "correlation": correlation, } - - -@numba_guvectorize( # type: ignore - [ - "void(float32[:], float32[:], float32[:], float32[:])", - "void(float64[:], float64[:], float64[:], float64[:])", - "void(int8[:], int8[:], int8[:], float64[:])", - ], - "(n),(n),(p)->(p)", -) -def euclidean_map_cpu( - x: ArrayLike, y: ArrayLike, _: ArrayLike, out: ArrayLike -) -> None: # pragma: no cover - """Euclidean distance "map" function for partial vector pairs. - - Parameters - ---------- - x - An array chunk, a partial vector - y - Another array chunk, a partial vector - _ - A dummy variable to map the size of output - out - The output array, which has the squared sum of partial vector pairs. - - Returns - ------- - An ndarray, which contains the output of the calculation of the application - of euclidean distance on the given pair of chunks, without the aggregation. - """ - square_sum = 0.0 - m = x.shape[0] - # Ignore missing values - for i in range(m): - if x[i] >= 0 and y[i] >= 0: - square_sum += (x[i] - y[i]) ** 2 - out[:] = square_sum - - -def euclidean_reduce_cpu(v: ArrayLike) -> ArrayLike: # pragma: no cover - """Corresponding "reduce" function for euclidean distance. - - Parameters - ---------- - v - The euclidean array on which map step of euclidean distance has been - applied. - - Returns - ------- - An ndarray, which contains square root of the sum of the squared sums obtained from - the map step of `euclidean_map`. - """ - out: ArrayLike = np.sqrt(np.einsum("ijkl -> ij", v)) - return out - - -@numba_guvectorize( # type: ignore - [ - "void(float32[:], float32[:], float32[:], float32[:])", - "void(float64[:], float64[:], float64[:], float64[:])", - "void(int8[:], int8[:], int8[:], float64[:])", - ], - "(n),(n),(p)->(p)", -) -def correlation_map_cpu( - x: ArrayLike, y: ArrayLike, _: ArrayLike, out: ArrayLike -) -> None: # pragma: no cover - """Pearson correlation "map" function for partial vector pairs. - - Parameters - ---------- - x - An array chunk, a partial vector - y - Another array chunk, a partial vector - _ - A dummy variable to map the size of output - out - The output array, which has the output of pearson correlation. - - Returns - ------- - An ndarray, which contains the output of the calculation of the application - of pearson correlation on the given pair of chunks, without the aggregation. - """ - - m = x.shape[0] - valid_indices = np.zeros(m, dtype=np.float64) - - for i in range(m): - if x[i] >= 0 and y[i] >= 0: - valid_indices[i] = 1 - - valid_shape = valid_indices.sum() - _x = np.zeros(int(valid_shape), dtype=x.dtype) - _y = np.zeros(int(valid_shape), dtype=y.dtype) - - # Ignore missing values - valid_idx = 0 - for i in range(valid_indices.shape[0]): - if valid_indices[i] > 0: - _x[valid_idx] = x[i] - _y[valid_idx] = y[i] - valid_idx += 1 - - out[:] = np.array( - [ - np.sum(_x), - np.sum(_y), - np.sum(_x * _x), - np.sum(_y * _y), - np.sum(_x * _y), - len(_x), - ] - ) - - -@numba_guvectorize( # type: ignore - [ - "void(float32[:, :], float32[:])", - "void(float64[:, :], float64[:])", - ], - "(p, m)->()", -) -def correlation_reduce_cpu(v: ArrayLike, out: ArrayLike) -> None: # pragma: no cover - """Corresponding "reduce" function for pearson correlation - Parameters - ---------- - v - The correlation array on which pearson corrections has been - applied on chunks - out - An ndarray, which is a symmetric matrix of pearson correlation - - Returns - ------- - An ndarray, which contains the result of the calculation of the application - of euclidean distance on all the chunks. - """ - v = v.sum(axis=0) - n = v[5] - num = n * v[4] - v[0] * v[1] - denom1 = np.sqrt(n * v[2] - v[0] ** 2) - denom2 = np.sqrt(n * v[3] - v[1] ** 2) - denom = denom1 * denom2 - value = np.nan - if denom > 0: - value = 1 - (num / denom) - out[0] = value - - -def call_metric_kernel( - f: ArrayLike, g: ArrayLike, metric: str, metric_kernel: Any -) -> ArrayLike: # pragma: no cover. - # Numba's 0.54.0 version is required, which is not released yet - # We install numba from numba conda channel: conda install -c numba/label/dev numba - # Relevant issue https://github.com/numba/numba/issues/6824 - f = np.ascontiguousarray(f) - g = np.ascontiguousarray(g) - - # move input data to the device - d_a = cuda.to_device(f) - d_b = cuda.to_device(g) - # create output data on the device - out = np.zeros((f.shape[0], g.shape[0], N_MAP_PARAM[metric]), dtype=f.dtype) - d_out = cuda.to_device(out) - - # https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#features-and-technical-specifications__technical-specifications-per-compute-capability - # These apply to compute capability 2.0 and higher and all GPUs NVIDIA has - # shipped in the past 10+ years have compute capability > 3.0. - # One way to get the compute capability programmatically is via: - # from numba import cuda - # cuda.get_current_device().compute_capability - - # In future when we have an average GPU with ability to have - # more number of threads per block, we can increase this to that value - # or parameterise this from the pairwise function or get the maximum - # possible value for a given compute capability. - - threads_per_block = (32, 32) - # Minimum grid size is 128 (12*12 = 144) - # Ref: https://github.com/numba/numba/blob/008077553b558bd183668ecd581d4d0bc54bd32c/numba/cuda/dispatcher.py#L534 - default_blocks_per_grid = (12, 12) - - blocks_per_grid = ( - max(default_blocks_per_grid[0], math.ceil(out.shape[0] / threads_per_block[0])), - max(default_blocks_per_grid[1], math.ceil(out.shape[1] / threads_per_block[1])), - ) - - metric_kernel[blocks_per_grid, threads_per_block](d_a, d_b, d_out) - # copy the output array back to the host system - d_out_host = d_out.copy_to_host() - return d_out_host - - -@cuda.jit(device=True) # type: ignore -def _correlation( - x: ArrayLike, y: ArrayLike, out: ArrayLike -) -> None: # pragma: no cover. - # Note: assigning variable and only saving the final value in the - # array made this significantly faster. - - # aggressively making all variables explicitly typed - # makes it more performant by a factor of ~2-3x - v0 = types.float32(0) - v1 = types.float32(0) - v2 = types.float32(0) - v3 = types.float32(0) - v4 = types.float32(0) - v5 = types.float32(0) - - m = types.uint32(x.shape[types.uint32(0)]) - i = types.uint32(0) - - zero = types.uint32(0) - - while i < m: - if x[i] >= zero and y[i] >= zero: - v0 += x[i] - v1 += y[i] - v2 += x[i] * x[i] - v3 += y[i] * y[i] - v4 += x[i] * y[i] - v5 += 1 - i = types.uint32(i + types.uint32(1)) - - out[0] = v0 - out[1] = v1 - out[2] = v2 - out[3] = v3 - out[4] = v4 - out[5] = v5 - - -@cuda.jit # type: ignore -def correlation_map_kernel( - x: ArrayLike, y: ArrayLike, out: ArrayLike -) -> None: # pragma: no cover. - i1 = types.uint32(cuda.grid(2)[types.uint32(0)]) - i2 = types.uint32(cuda.grid(2)[types.uint32(1)]) - - out_shape_0 = types.uint32(out.shape[types.uint32(0)]) - out_shape_1 = types.uint32(out.shape[types.uint32(1)]) - - if i1 >= out_shape_0 or i2 >= out_shape_1: - # Quit if (x, y) is outside of valid output array boundary - return - - _correlation(x[i1], y[i2], out[i1][i2]) - - -def correlation_map_gpu(x: ArrayLike, y: ArrayLike) -> ArrayLike: # pragma: no cover. - """Pearson correlation "map" function for partial vector pairs on GPU - - Parameters - ---------- - x - [array-like, shape: (m, n)] - y - [array-like, shape: (p, n)] - - Returns - ------- - An ndarray, which contains the output of the calculation of the application - of pearson correlation on the given pair of chunks, without the aggregation. - """ - - return call_metric_kernel(x, y, "correlation", correlation_map_kernel) - - -def correlation_reduce_gpu(v: ArrayLike) -> None: # pragma: no cover. - """GPU implementation of the corresponding "reduce" function for pearson - correlation. - - Parameters - ---------- - v - [array-like, shape: (1, n)] - The correlation array on which map step pearson corrections has been - applied. - - Returns - ------- - An ndarray, which contains the result of the calculation of the reduction step - of correlation metric. - """ - return correlation_reduce_cpu(v) # type: ignore - - -@cuda.jit(device=True) # type: ignore -def _euclidean_distance_map( - a: ArrayLike, b: ArrayLike, out: ArrayLike -) -> None: # pragma: no cover. - """Helper function for the map step of euclidean distance which runs on - the device (GPU) itself. - - Parameters - ---------- - a - [array-like, shape: (1, n)] - b - [array-like, shape: (1, n)] - out - [array-like, shape: (1)] - The output array for returning the result. - - Returns - ------- - An ndarray, which contains the squared sum of the corresponding elements - of the given pair of vectors. - """ - square_sum = types.float32(0) - - zero = types.uint32(0) - a_shape_0 = types.uint32(a.shape[types.uint32(0)]) - i = types.uint32(0) - - while i < a_shape_0: - if a[i] >= zero and b[i] >= zero: - square_sum += (a[i] - b[i]) ** types.uint32(2) - i = types.uint32(i + types.uint32(1)) - out[0] = square_sum - - -@cuda.jit # type: ignore -def euclidean_map_kernel( - x: ArrayLike, y: ArrayLike, out: ArrayLike -) -> None: # pragma: no cover. - """Euclidean map CUDA kernel. - - Parameters - ---------- - x - [array-like, shape: (m, n)] - y - [array-like, shape: (p, n)] - out - [array-like, shape: (m, p, 1)] - The zeros array of shape (m, p, 1) for returning the result. - - Returns - ------- - An ndarray, which contains the output of the calculation of the application - of euclidean distance on all pairs of vectors from x and y arrays. - """ - # Aggresive typecasting of all the variables is done to improve performance. - - # Unique index of the thread in the whole grid. - i1 = types.uint32(cuda.grid(2)[types.uint32(0)]) - i2 = types.uint32(cuda.grid(2)[types.uint32(1)]) - - out_shape_0 = types.uint32(out.shape[types.uint32(0)]) - out_shape_1 = types.uint32(out.shape[types.uint32(1)]) - - if i1 >= out_shape_0 or i2 >= out_shape_1: - # Quit if (x, y) is outside of valid output array boundary - # This is required because we may spin up more threads than we need. - return - _euclidean_distance_map(x[i1], y[i2], out[i1][i2]) - - -def euclidean_map_gpu(x: ArrayLike, y: ArrayLike) -> ArrayLike: # pragma: no cover. - """GPU implementation of Euclidean distance "map" function for partial - vector pairs. This runs on GPU by using the euclidean_map_kernel cuda kernel. - - Parameters - ---------- - x - [array-like, shape: (m, n)] - y - [array-like, shape: (p, n)] - - Returns - ------- - An ndarray, which contains the output of the calculation of the application - of euclidean distance on the given pair of chunks, without the aggregation. - """ - return call_metric_kernel(x, y, "euclidean", euclidean_map_kernel) - - -def euclidean_reduce_gpu(v: ArrayLike) -> ArrayLike: # pragma: no cover. - """GPU Implementation of the Corresponding "reduce" function for euclidean - distance. - - Parameters - ---------- - v - The euclidean array on which map step of euclidean distance has been - applied. - - Returns - ------- - An ndarray, which contains square root of the sum of the squared sums obtained from - the map step of `euclidean_map`. - """ - return euclidean_reduce_cpu(v) diff --git a/sgkit/tests/test_distance.py b/sgkit/tests/test_distance.py index 67b99475e..312fe55a9 100644 --- a/sgkit/tests/test_distance.py +++ b/sgkit/tests/test_distance.py @@ -3,23 +3,24 @@ import dask.array as da import numpy as np import pytest -from numba import cuda from scipy.spatial.distance import correlation, euclidean, pdist, squareform from sgkit.distance.api import DeviceTypes, MetricTypes, pairwise_distance from sgkit.typing import ArrayLike -def detect_cuda_driver() -> bool: +def detect_gpu() -> bool: try: - return len(cuda.list_devices()) > 0 - except cuda.CudaSupportError: + import cupy + + return cupy.cuda.runtime.getDeviceCount() > 0 + except (ImportError, RuntimeError): return False def skip_gpu_tests_if_no_gpu(device: DeviceTypes) -> None: - if device == "gpu" and not detect_cuda_driver(): - pytest.skip("Cuda driver not found") + if device == "gpu" and not detect_gpu(): + pytest.skip("CuPy or a GPU not found") def get_vectors( @@ -143,9 +144,6 @@ def test_distance_ndarray() -> None: pytest.param("correlation", correlation, "i8", "gpu", marks=pytest.mark.gpu), ], ) -@pytest.mark.filterwarnings( - "ignore:invalid value encountered in correlation_reduce_cpu:RuntimeWarning" -) def test_missing_values( metric: MetricTypes, metric_func: typing.Callable[[ArrayLike, ArrayLike], np.float64], @@ -169,15 +167,140 @@ def test_missing_values( np.testing.assert_almost_equal(distance_matrix, expected_matrix) +@pytest.mark.parametrize( + "metric, metric_func", + [("euclidean", euclidean), ("correlation", correlation)], +) +@pytest.mark.parametrize("device", ["cpu", pytest.param("gpu", marks=pytest.mark.gpu)]) +def test_genotype_calls( + metric: MetricTypes, + metric_func: typing.Callable[[ArrayLike, ArrayLike], np.float64], + device: DeviceTypes, +) -> None: + skip_gpu_tests_if_no_gpu(device) + # int8 calls with -1 for missing values, as in call_genotype + rs = np.random.RandomState(0) + x = rs.randint(0, 3, size=(30, 300)).astype(np.int8) + x[rs.rand(*x.shape) < 0.1] = -1 + x[0] = -1 # no observed values + x[1] = 1 # constant + x[2, 2:] = -1 # two observed values + + distance_matrix = pairwise_distance( + da.from_array(x, chunks=(10, 150)), metric=metric, device=device + ) + expected_matrix = create_distance_matrix(x, metric_func) + np.testing.assert_almost_equal(distance_matrix, expected_matrix) + + +@pytest.mark.parametrize("metric", ["euclidean", "correlation"]) +@pytest.mark.parametrize("device", ["cpu", pytest.param("gpu", marks=pytest.mark.gpu)]) +def test_identical_genotype_calls(metric: MetricTypes, device: DeviceTypes) -> None: + skip_gpu_tests_if_no_gpu(device) + # the sums of squares and products of integers are exact, so copies of a + # row, in the same chunk or in another one, are exactly zero apart + rs = np.random.RandomState(0) + x = rs.randint(0, 3, size=(4, 100_000)).astype(np.int8) + x[rs.rand(*x.shape) < 0.05] = -1 + x[1] = x[0] + x[3] = x[0] + + distance_matrix = pairwise_distance( + da.from_array(x, chunks=(2, 25_000)), metric=metric, device=device + ).compute() + assert distance_matrix[0, 1] == distance_matrix[0, 3] == 0 + assert distance_matrix[1, 3] == 0 + + +@pytest.mark.parametrize( + "metric, metric_func, atol", + [("euclidean", euclidean, 1e-3), ("correlation", correlation, 1e-6)], +) +@pytest.mark.parametrize("device", ["cpu", pytest.param("gpu", marks=pytest.mark.gpu)]) +def test_nearly_identical_float32_rows( + metric: MetricTypes, + metric_func: typing.Callable[[ArrayLike, ArrayLike], np.float64], + atol: float, + device: DeviceTypes, +) -> None: + skip_gpu_tests_if_no_gpu(device) + # float32 dosages, where the second row is a copy of the first and the third + # differs from it in five entries. When the sums of squares and products + # were subtracted in float32, these distances were lost. + rs = np.random.RandomState(0) + x = rs.uniform(0, 2, size=(4, 100_000)).round(3).astype(np.float32) + x[rs.rand(*x.shape) < 0.05] = np.nan + x[1] = x[0] + x[2] = x[0] + x[2, :5] += 0.05 + + distance_matrix = pairwise_distance( + da.from_array(x, chunks=(2, 25_000)), metric=metric, device=device + ).compute() + expected_matrix = create_distance_matrix(x.astype(np.float64), metric_func) + np.testing.assert_allclose(distance_matrix, expected_matrix, rtol=1e-6, atol=atol) + assert distance_matrix.min() >= 0 + + +@pytest.mark.parametrize("metric", ["euclidean", "correlation"]) +@pytest.mark.parametrize("device", ["cpu", pytest.param("gpu", marks=pytest.mark.gpu)]) +def test_float16(metric: MetricTypes, device: DeviceTypes) -> None: + skip_gpu_tests_if_no_gpu(device) + # the sums of squares and products are larger than the float16 maximum + x = np.array([[0, 2] * 10_000, [2, 0] * 10_000, [0, 2] * 10_000], dtype="f2") + + distance_matrix = pairwise_distance(x, metric=metric, device=device) + assert distance_matrix.dtype == np.float16 + expected_matrix = squareform(pdist(x.astype(np.float64), metric=metric)) + np.testing.assert_allclose(distance_matrix, expected_matrix, rtol=1e-3) + + +@pytest.mark.parametrize( + "metric, expected_matrix", + [ + ("euclidean", np.sqrt([[0, 2, 2], [2, 0, 0], [2, 0, 0]])), + ("correlation", [[0, 2, 2], [2, 0, 0], [2, 0, 0]]), + ], +) +@pytest.mark.parametrize("device", ["cpu", pytest.param("gpu", marks=pytest.mark.gpu)]) +def test_infinite_values( + metric: MetricTypes, expected_matrix: ArrayLike, device: DeviceTypes +) -> None: + skip_gpu_tests_if_no_gpu(device) + # infinite values are missing, like negative and nan values, so the first + # column is only observed in the last row. In the matrix products, the + # infinite value made every distance of the first row nan. + x = np.array([[np.inf, 0, 1], [-1, 1, 0], [2, 1, 0]]) + + distance_matrix = pairwise_distance(x, metric=metric, device=device) + np.testing.assert_allclose(distance_matrix, expected_matrix) + + +@pytest.mark.gpu +@pytest.mark.parametrize("metric", ["euclidean", "correlation"]) +def test_cupy_backed_array(metric: MetricTypes) -> None: + cupy = pytest.importorskip("cupy") + skip_gpu_tests_if_no_gpu("gpu") + x = get_vectors() + distance_matrix = pairwise_distance(x.map_blocks(cupy.asarray), metric=metric) + assert isinstance(distance_matrix._meta, cupy.ndarray) + expected_matrix = squareform(pdist(x, metric=metric)) + np.testing.assert_almost_equal( + cupy.asnumpy(distance_matrix.compute()), expected_matrix + ) + + @pytest.mark.parametrize( "metric, dtype, expected, device", [ + ("euclidean", "i1", "float64", "cpu"), ("euclidean", "i8", "float64", "cpu"), ("euclidean", "f4", "float32", "cpu"), ("euclidean", "f8", "float64", "cpu"), pytest.param("euclidean", "i8", "float64", "gpu", marks=pytest.mark.gpu), pytest.param("euclidean", "f4", "float32", "gpu", marks=pytest.mark.gpu), pytest.param("euclidean", "f8", "float64", "gpu", marks=pytest.mark.gpu), + ("correlation", "i1", "float64", "cpu"), ("correlation", "i8", "float64", "cpu"), ("correlation", "f4", "float32", "cpu"), ("correlation", "f8", "float64", "cpu"), @@ -191,8 +314,9 @@ def test_data_types( ) -> None: skip_gpu_tests_if_no_gpu(device) x = get_vectors(dtype=dtype) - distance_matrix = pairwise_distance(x, metric=metric, device=device).compute() + distance_matrix = pairwise_distance(x, metric=metric, device=device) assert distance_matrix.dtype.name == expected + assert distance_matrix.compute().dtype.name == expected def test_undefined_metric() -> None: