diff --git a/heat/core/manipulations.py b/heat/core/manipulations.py index deea81333a..b80faac41c 100644 --- a/heat/core/manipulations.py +++ b/heat/core/manipulations.py @@ -24,6 +24,7 @@ from . import _operations __all__ = [ + "argsort", "balance", "broadcast_arrays", "broadcast_to", @@ -65,6 +66,64 @@ ] +def argsort( + a: DNDarray, axis: int = -1, *args, descending: bool | None = None, **kwargs +) -> DNDarray: + """ + Returns the indices that would sort an array. This is the distributed equivalent of `np.argsort`. + The sorting is not stable which means that equal elements in the result may have a different ordering than in the + original array. + Sorting with `axis==a.split` needs a lot of communication between the processes. + + Parameters + ---------- + a : DNDarray + Input array to be sorted. + axis : int, optional + The dimension to sort along. + Default is the last axis. + *args: Any + Any arguments that are not specified in the function header will be ignored. + descending : bool, optional + If set to `True`, indices are sorted in descending order. + **kwargs: Any + Any keyword arguments that are not specified in the function header will be ignored. + + Raises + ------ + ValueError + If `axis` is not consistent with the available dimensions. + + Examples + -------- + >>> x = ht.array([[4, 1], [2, 3]], split=0) + >>> x.shape + (1, 2) + (1, 2) + >>> y = ht.argsort(x, axis=0) + >>> y + (array([[1, 0], + [0, 1]])) + >>> ht.argsort(x, descending=True) + (array([[0, 1], + [1, 0]])) + """ + for arg in args: + warnings.warn(f"[ht.argsort] Argument: '{arg}' gets ignored.") + + for k in kwargs.keys(): + warnings.warn(f"[ht.argsort] Keyword Argument: '{k}' gets ignored.") + + _, indices = sort( + a=a, + axis=axis, + descending=descending, + out=None, + return_sort_indices=True, + ) + return indices + + def balance(array: DNDarray, copy=False) -> DNDarray: """ Out of place balance function. More information on the meaning of balance can be found in @@ -2543,12 +2602,20 @@ def shape(a: DNDarray) -> Tuple[int, ...]: return a.gshape -def sort(a: DNDarray, axis: int = -1, descending: bool = False, out: Optional[DNDarray] = None): +def sort( + a: DNDarray, + axis: int = -1, + *args, + descending: bool | None = False, + out: Optional[DNDarray] = None, + return_sort_indices: bool = False, + **kwargs, +): """ Sorts the elements of `a` along the given dimension (by default in ascending order) by their value. The sorting is not stable which means that equal elements in the result may have a different ordering than in the original array. - Sorting where `axis==a.split` needs a lot of communication between the processes of MPI. + Sorting with `axis==a.split` needs a lot of communication between the processes of MPI. Returns a tuple `(values, indices)` with the sorted local results and the indices of the elements in the original data Parameters @@ -2558,11 +2625,18 @@ def sort(a: DNDarray, axis: int = -1, descending: bool = False, out: Optional[DN axis : int, optional The dimension to sort along. Default is the last axis. + *args: Any + Any arguments that are not specified in the function header will be ignored. descending : bool, optional If set to `True`, values are sorted in descending order. out : DNDarray, optional A location in which to store the results. If provided, it must have a broadcastable shape. If not provided or set to `None`, a fresh array is allocated. + return_sort_indices: bool, optional + Wether to return the indices by which the array was sorted. + If ``out`` is provided, returns the indices if ``True``, otherwise ``None``. + **kwargs: Any + Any keyword arguments that are not specified in the function header will be ignored. Raises ------ @@ -2591,6 +2665,18 @@ def sort(a: DNDarray, axis: int = -1, descending: bool = False, out: Optional[DN ) stride_tricks.sanitize_axis(a.shape, axis) + for arg in args: + warnings.warn(f"[ht.sort] Argument: '{arg}' gets ignored.") + + for k in kwargs.keys(): + warnings.warn(f"[ht.sort] Keyword Argument: '{k}' gets ignored.") + + if axis is None: + a = flatten(a) + axis = 0 + + descending = descending or False + if not a.is_distributed() or axis != a.split: # sorting is not affected by split -> we can just sort along the axis final_result, final_indices = torch.sort(a.larray, dim=axis, descending=descending) @@ -2801,12 +2887,16 @@ def sort(a: DNDarray, axis: int = -1, descending: bool = False, out: Optional[DN ) if out is not None: out.larray = final_result - return return_indices + if return_sort_indices: + return return_indices + return None else: tensor = factories.array( final_result, dtype=a.dtype, is_split=a.split, device=a.device, comm=a.comm ) - return tensor, return_indices + if return_sort_indices: + return tensor, return_indices + return tensor def split(x: DNDarray, indices_or_sections: Iterable, axis: int = 0) -> List[DNDarray, ...]: diff --git a/heat/core/statistics.py b/heat/core/statistics.py index 3bd1cf9f58..d1afa17bd7 100644 --- a/heat/core/statistics.py +++ b/heat/core/statistics.py @@ -1543,7 +1543,8 @@ def _create_sketch( # create a random sample of indices indices = manipulations.sort( - randint(0, a.shape[axis], sketch_size, device=a.device, dtype=types.int64) + randint(0, a.shape[axis], sketch_size, device=a.device, dtype=types.int64), + return_sort_indices=True, )[0] sketch = a.swapaxes(0, axis) sketch = a[indices, ...].resplit_(None) @@ -1697,7 +1698,7 @@ def _create_sketch( ) # sort data - sorted_x, _ = manipulations.sort(x, axis=axis) + sorted_x, _ = manipulations.sort(x, axis=axis, return_sort_indices=True) del _ sorted_x = sorted_x.astype(output_dtype) diff --git a/heat/utils/data/matrixgallery.py b/heat/utils/data/matrixgallery.py index 47b6c0b19c..779dec9749 100644 --- a/heat/utils/data/matrixgallery.py +++ b/heat/utils/data/matrixgallery.py @@ -197,7 +197,9 @@ def random_known_rank( raise RuntimeError("rank must not exceed matrix dimensions.") singular_values = rand(r, dtype=dtype, comm=comm, device=device) - singular_values = sort(quantile_function(singular_values), descending=True)[0] + singular_values = sort( + quantile_function(singular_values), descending=True, return_sort_indices=True + )[0] return random_known_singularvalues( m, n, singular_values, split=split, device=device, comm=comm, dtype=dtype diff --git a/tests/core/test_manipulations.py b/tests/core/test_manipulations.py index 20e10a2ad8..51d123a9d7 100644 --- a/tests/core/test_manipulations.py +++ b/tests/core/test_manipulations.py @@ -5,7 +5,6 @@ import heat as ht from heat.testing.basic_test import TestCase - class TestManipulations(TestCase): def test_broadcast_arrays(self): a = ht.array([[1], [2]]) @@ -2787,12 +2786,12 @@ def test_sort(self): ) data = ht.array(tensor, split=None) - result, result_indices = ht.sort(data, axis=0, descending=True) + result, result_indices = ht.sort(data, axis=0, descending=True, return_sort_indices=True) expected, exp_indices = torch.sort(tensor, dim=0, descending=True) self.assertTrue(torch.equal(result.larray, expected)) self.assertTrue(torch.equal(result_indices.larray, exp_indices.int())) - result, result_indices = ht.sort(data, axis=1, descending=True) + result, result_indices = ht.sort(data, axis=1, descending=True, return_sort_indices=True) expected, exp_indices = torch.sort(tensor, dim=1, descending=True) self.assertTrue(torch.equal(result.larray, expected)) self.assertTrue(torch.equal(result_indices.larray, exp_indices.int())) @@ -2801,7 +2800,7 @@ def test_sort(self): exp_axis_zero = torch.arange(size, device=self.device.torch_device).reshape(1, size) exp_indices = torch.tensor([[rank] * size], device=self.device.torch_device) - result, result_indices = ht.sort(data, descending=True, axis=0) + result, result_indices = ht.sort(data, descending=True, axis=0, return_sort_indices=True) self.assertTrue(torch.equal(result.larray, exp_axis_zero)) self.assertTrue(torch.equal(result_indices.larray, exp_indices.int())) @@ -2810,12 +2809,12 @@ def test_sort(self): .reshape(1, size) .sort(dim=1, descending=True) ) - result, result_indices = ht.sort(data, descending=True, axis=1) + result, result_indices = ht.sort(data, descending=True, axis=1, return_sort_indices=True) self.assertTrue(torch.equal(result.larray, exp_axis_one)) self.assertTrue(torch.equal(result_indices.larray, exp_indices.int())) - result1 = ht.sort(data, axis=1, descending=True) - result2 = ht.sort(data, descending=True) + result1 = ht.sort(data, axis=1, descending=True, return_sort_indices=True) + result2 = ht.sort(data, descending=True, return_sort_indices=True) self.assertTrue(ht.equal(result1[0], result2[0])) self.assertTrue(ht.equal(result1[1], result2[1])) @@ -2827,7 +2826,7 @@ def test_sort(self): indices_axis_zero = torch.arange( size, dtype=torch.int64, device=self.device.torch_device ).reshape(size, 1) - result, result_indices = ht.sort(data, axis=0, descending=True) + result, result_indices = ht.sort(data, axis=0, descending=True, return_sort_indices=True) self.assertTrue(torch.equal(result.larray, exp_axis_zero)) # comparison value is only true on CPU if result_indices.larray.is_cuda is False: @@ -2838,7 +2837,7 @@ def test_sort(self): .repeat(size) .reshape(size, 1) ) - result, result_indices = ht.sort(data, descending=True, axis=1) + result, result_indices = ht.sort(data, descending=True, axis=1, return_sort_indices=True) self.assertTrue(torch.equal(result.larray, exp_axis_one)) self.assertTrue(torch.equal(result_indices.larray, exp_axis_one.int())) @@ -2860,7 +2859,7 @@ def test_sort(self): indices_axis_zero = torch.tensor( [[0, 2, 2], [3, 0, 0]], dtype=torch.int32, device=self.device.torch_device ) - result, result_indices = ht.sort(data, axis=0) + result, result_indices = ht.sort(data, axis=0, return_sort_indices=True) first = result[0].larray first_indices = result_indices[0].larray if rank == 0: @@ -2879,7 +2878,7 @@ def test_sort(self): indices_axis_one = torch.tensor( [[0, 1, 1]], dtype=torch.int32, device=self.device.torch_device ) - result, result_indices = ht.sort(data, axis=1) + result, result_indices = ht.sort(data, axis=1, return_sort_indices=True) first = result[0].larray[:1] first_indices = result_indices[0].larray[:1] if rank == 0: @@ -2891,7 +2890,7 @@ def test_sort(self): indices_axis_two = torch.tensor( [[0], [1]], dtype=torch.int32, device=self.device.torch_device ) - result, result_indices = ht.sort(data, axis=2) + result, result_indices = ht.sort(data, axis=2, return_sort_indices=True) first = result[0].larray[:, :1] first_indices = result_indices[0].larray[:, :1] if rank == 0: @@ -2899,7 +2898,7 @@ def test_sort(self): self.assertTrue(torch.equal(first_indices, indices_axis_two)) # out = ht.empty_like(data) - indices = ht.sort(data, axis=2, out=out) + indices = ht.sort(data, axis=2, out=out, return_sort_indices=True) self.assertTrue(ht.equal(out, result)) self.assertTrue(ht.equal(indices, result_indices)) @@ -2911,7 +2910,7 @@ def test_sort(self): rank = ht.MPI_WORLD.rank ht.random.seed(1) data = ht.random.randn(100, 1, split=0) - result, _ = ht.sort(data, axis=0) + result, _ = ht.sort(data, axis=0, return_sort_indices=True) counts, _, _ = ht.get_comm().counts_displs_shape(data.gshape, axis=0) for i, c in enumerate(counts): for idx in range(c - 1): diff --git a/tests/core/test_sorting.py b/tests/core/test_sorting.py new file mode 100644 index 0000000000..34806e6502 --- /dev/null +++ b/tests/core/test_sorting.py @@ -0,0 +1,48 @@ +import numpy as np +import torch +import pytest +import os +import heat as ht +from heat.testing.basic_test import TestCase + + +class TestSorting: + def __init__(self, *args, **kwargs): + TestCase.setUpClass() + + @pytest.mark.parametrize("split", [None, 0, 1, 2]) + @pytest.mark.parametrize("descending", [None, True, False]) + @pytest.mark.parametrize("axis", [None, 0, 1, 2]) + def test_sort(self, axis, descending, split): + kwargs = {"axis": axis} + + if descending in [True, False]: + if np.lib.NumpyVersion(np.__version__) < np.lib.NumpyVersion("2.5.0"): + pytest.skip(f"NumPy {np.__version__} does not support the 'descending' keyword.") + else: + kwargs["descending"] = descending + + data = ht.random.rand(2, 3, 4, split=split) + result, idx = ht.sort(data, return_sort_indices=True, **kwargs) + exp = np.sort(data.numpy(), **kwargs) + exp_idx = np.argsort(data.numpy(), **kwargs) + + assert np.allclose(result.numpy(), exp) + assert np.allclose(idx.numpy(), exp_idx) + + @pytest.mark.parametrize("split", [None, 0, 1, 2]) + @pytest.mark.parametrize("descending", [None, True, False]) + @pytest.mark.parametrize("axis", [None, 0, 1, 2]) + def test_argsort_random(self, axis, descending, split): + kwargs = {"axis": axis} + + if descending in [True, False]: + if np.lib.NumpyVersion(np.__version__) < np.lib.NumpyVersion("2.5.0"): + pytest.skip(f"NumPy {np.__version__} does not support the 'descending' keyword.") + else: + kwargs["descending"] = descending + + data = ht.random.rand(2, 3, 4, split=split) + result_indices = ht.argsort(data, **kwargs) + exp_indices = np.argsort(data.numpy(), **kwargs) + assert np.allclose(result_indices.numpy(), exp_indices)