How to modify the sorting results of a numpy array

user_14125436 0
Viewed 161

I want to sort an array based on the values of two columns and a condition on the third. This is my array:

my_array = np.array([[1., 2., 5.1],
                     [1., 1., 5.],
                     [2., 2., 2.],
                     [2., 1., 2.],
                     [2., 2., 5.],
                     [3., 2., 2.5],
                     [3., 1., 2.5],
                     [2., 1., 5.]])

I must sort it based on the first and second column and also a condition based on the third column. I tried this method:

my_sorted_array = my_array[np.lexsort((my_array[:, 1], my_array[:, 0]))]

but it does not regard my third column. It gives me:

result = np.array([[1. , 1. , 5. ],
                   [1. , 2. , 5.1 ],
                   [2., 1. , 2. ],
                   [2., 1. , 5. ],
                   [2. , 2. , 2. ],
                   [2. , 2. , 5. ],
                   [3. , 1. , 2.5 ],
                   [3. , 2. , 2.5 ]])

I want to have the following output:

my_sorted_array = np.array([[1., 1., 5.],
                            [1., 2., 5.1],
                            [2., 1., 5.],
                            [2., 2., 5.],
                            [2., 1., 2.],
                            [2., 2., 2.],
                            [3., 1., 2.5],
                            [3., 2., 2.5]])

I also tried to set coefficients to the third column using this method:

sort_func = my_array[:, 0] * c1 + my_array[:, 1] * c2 + my_array[:, 2] * c3 # c1, c2 and c3 are coefficient
sort_index = np.argsort(sort_func)

This method also is time consuming because I should tune coefficinet for each new data set. Is it possible to put an if_condition in the sorting? How can I rearrange the result into my_sorted_array? Data of two first columns are always gentle and regular (they x and y of a regular grid).

To make it more visual, I uploaded a figure here. The figure shows the trend I want to use to sort my data.

enter image description here

numpy sorting arrays python
1 Answers

Based on our discussion, you want to sort first on some categorical function of z, then x, then y. So in general, something very similar to your original code, rearranged ever so slightly for clarity:

x, y, z = my_array.T
index = np.lexsort((y, x, f(z)))
my_sorted_array = my_array[index, :]

For the simplest case, you can make f(x) return a boolean mask. This would work perfectly for your toy example, since it splits the data into two categories. Keep in mind that points labeled False will come before those labeled True, and apply your threshold accordingly:

def f(z):
    return z < z.mean()

But why stop at two categories? You can use np.digitize to split the data into an arbitrary number of labels:

breaks = [2.25, 4.9]

def f(z):
    return np.digitize(z, np.sort(breaks)[::-1])

Even more specifically, if your data has a pair of approximately equal clumps, you can isolate the gunk in between with something like this:

def f(z):
    mask0 = z > z.mean()
    z1 = z[mask0]
    z2 = z[~mask0]
    breaks = [z1.mean() - 3 * z1.std(), z2.mean() + 3 * z2.std()]
    return np.digitize(z, breaks)
user_29887300
Related