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: