sort numpy array with custom predicate

Question

I'd like to sort my numpy array of shape [n,4], along first dimension (size:n) using a custom predicate operating on the 2nd dimension vector (size:4). The C++ version of what I'd like to do is below, it's quite simple really. I've seen how to do this with python lists, but I can't find the syntax to do it with numpy arrays. Is this possible? The documentation on np.sort, np.argsort, np.lexsort doesn't mention custom predicates.

// c++ version
vector< float[4] > v = init_v(); 
float[4] p = init_p();
std::sort(v.begin(), v.end(), [&p](const auto& lhs, const auto& rhs) {
   return myfn(p, lhs) > myfn(p, rhs); });

EDIT: below is the python code I would like to use for the sorting. I.e. for each 'row' (n:4) of my array, I'd calculate the square of the euclidean 3D distance (i.e. only the first 3 columns) to a fixed point.

# these both operate on numpy vectors of shape [4] (i.e. a single row of my data matrix)
def dist_sq(a,b):
    d = a[:3]-b[:3]
    return np.dot(d*d)

def sort_pred(lhs, rhs, p):
    return dist_sq(lhs, p) > dist_sq(rhs, p)

Answer 1

In numpy you would apply the (vectorized) order defining function to the array, then use np.argsort to sort by the result.

This is less space efficient than the C++ version, but that is how you usually achieve performance with numpy.

import numpy as np    

def myfn(x):
    return np.sin(x[:, 1])  # example: sort by the sine of the second column

a = np.random.randn(10, 4)

predicate = myfn(a)  # not sure if predicate is the best name for this variable
order = np.argsort(predicate)

a_sorted = a[order]