[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,69 @@
 ]
 
 
+def argsort(
+    a: DNDarray,
+    axis: int = -1,
+    kind: str | None = None,
+    order: str | list[str] | None = None,
+    *,
+    stable: bool | None = None,
+    descending: bool | None = None,
+) -> 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 where `axis==a.split` needs a lot of communication between the processes of MPI.
+
+    Parameters
+    ----------
+    a : DNDarray
+        Input array to be sorted.
+    axis : int, optional
+        The dimension to sort along.
+        Default is the last axis.
+    kind : str, optional
+        Sorting algorithm. Gets ignored.
+    order : str, list[str], optional
+        Used in numpy for array with fields, which are not possible in HeAT. Gets ignored.
+    stable : bool, optional
+        Sort stability, currenty not supported.
+    descending : bool, optional
+        If set to `True`, indices are sorted in descending order.
+
+    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]]))
+    """
+    _, indices = sort(
+        a=a,
+        axis=axis,
+        kind=kind,
+        order=order,
+        stable=stable,
+        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
@@ -2549,7 +2613,17 @@ 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,
+    kind: str | None = None,
+    order: str | list[str] | None = None,
+    *,
+    stable: bool | None = None,
+    descending: bool | None = False,
+    out: Optional[DNDarray] = None,
+    return_sort_indices: bool | None = None,
+):
     """
     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
@@ -2564,11 +2638,20 @@ 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.
+    kind : str, optional
+        Sorting algorithm. Gets ignored.
+    order : str, list[str], optional
+        Used in numpy for array with fields, which are not possible in HeAT. Gets ignored.
+    stable : bool, optional
+        Sort stability, currenty not supported.
     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``.
 
     Raises
     ------
@@ -2597,6 +2680,13 @@ def sort(a: DNDarray, axis: int = -1, descending: bool = False, out: Optional[DN
     )
     stride_tricks.sanitize_axis(a.shape, axis)
 
+    if axis is None:
+        a = flatten(a)
+        axis = 0
+
+    descending = descending or False
+    return_sort_indices = return_sort_indices 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)
@@ -2807,12 +2897,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