Add argsort1d function

fluids/numerics/__init__.py

@@ -29,7 +29,7 @@
 import sys
 from fluids.numerics.arrays import (solve as py_solve, inv, dot, norm2, inner_product, eye,
                      array_as_tridiagonals, tridiagonals_as_array, transpose,
-                     solve_tridiagonal, subset_matrix)
+                     solve_tridiagonal, subset_matrix, argsort1d)
 
 from fluids.numerics.special import (py_hypot, py_cacos, py_catan, py_catanh, 
                                      trunc_exp, trunc_log, cbrt, factorial, comb)
@@ -123,6 +123,7 @@
            'sort_nelder_mead_points_numba', 'sort_nelder_mead_points_python', 
            'bounds_clip_naive', 'nelder_mead', 'cbrt',
            'polyint_stable', 'polyint_over_x_stable',
+           'argsort1d',
            ]
 
 from fluids.numerics import doubledouble

fluids/numerics/arrays.py

@@ -44,7 +44,8 @@
 #        np = None
 
 __all__ = ['dot', 'inv', 'det', 'solve', 'norm2', 'inner_product', 'transpose',
-           'eye', 'array_as_tridiagonals', 'solve_tridiagonal', 'subset_matrix']
+           'eye', 'array_as_tridiagonals', 'solve_tridiagonal', 'subset_matrix',
+           'argsort1d']
 primitive_containers = frozenset([list, tuple])
 
 def transpose(x):
@@ -611,3 +612,34 @@ def subset_matrix(whole, subset):
 #        for j in subset:
 #            r.append(whole_i[j])
     return new
+
+
+
+## argsort implementation
+
+def argsort1d(arr):
+    """
+    Returns the indices that would sort a 1D list.
+
+    Parameters
+    ----------
+    arr : list
+        Input array [-]
+
+    Returns
+    -------
+    indices : list[int]
+        List of indices that sort the input array [-]
+
+    Notes
+    -----
+    This function uses the built-in sorted function with a custom key to get the indices.
+    Note this does not match numpy's sorting for nan and inf values.
+
+    Examples
+    --------
+    >>> arr = [3, 1, 2]
+    >>> argsort1d(arr)
+    [1, 2, 0]
+    """
+    return [i[0] for i in sorted(enumerate(arr), key=lambda x: x[1])]

tests/test_numerics.py

@@ -88,6 +88,7 @@
     trunc_log_numpy,
     zeros,
     is_increasing,
+    argsort1d,
 )
 from fluids.numerics import numpy as np
 
@@ -1552,3 +1553,33 @@ def test_py_lambertw():
     # Test the boundary at -1/e
     minus_one_over_e = -1/exp(1)
     res = abs(py_lambertw(minus_one_over_e, k=-1).real + 1)
+
+
+def test_argsort1d():
+
+    def check_argsort1d(input_list, expected, error_message):
+        numpy_argsort1d = lambda x: list(np.argsort(x))
+        assert argsort1d(input_list) == expected, error_message
+        assert argsort1d(input_list) == numpy_argsort1d(input_list), error_message
+
+
+    check_argsort1d([3, 1, 2], [1, 2, 0], "Failed on simple test case")
+    check_argsort1d([-1, -3, -2], [1, 2, 0], "Failed with negative numbers")
+    check_argsort1d([], [], "Failed on empty list")
+    check_argsort1d([42], [0], "Failed with single element list")
+    check_argsort1d([99, 21, 31, 80, 70], [1, 2, 4, 3, 0], "Mismatch with expected output")
+    check_argsort1d([2, 3, 1, 5, 4], [2, 0, 1, 4, 3], "Mismatch with expected output")
+
+    check_argsort1d([3.5, 1, 2.2], [1, 2, 0], "Failed with mixed floats and ints")
+    check_argsort1d([0.1, 0.2, 0.3], [0, 1, 2], "Failed with floats")
+
+    check_argsort1d([True, False, True], [1, 0, 2], "Failed with boolean values")
+
+    check_argsort1d(['apple', 'banana', 'cherry'], [0, 1, 2], "Failed with strings")
+
+    check_argsort1d([2, 3, 2, 3, 3], [0, 2, 1, 3, 4], "Failed with duplicate numbers")
+
+    check_argsort1d([-3, -1, 0, 1, 3], [0, 1, 2, 3, 4], "Failed with negative and positive numbers")
+
+    # infinities and nan behavior does not match
+    # check_argsort1d([-np.inf, np.inf, np.nan, 0, -1], [0, 4, 3, 2, 1], "Failed with infinities and NaN")