[Breaking Change] Implement argsort and make sort numpy API compliant

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, ...]:

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)
 

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

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,15 +2890,15 @@ 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:
             self.assertTrue(torch.equal(first, exp_axis_two))
             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):