Coefficient of Variation and NumPy

Question

I'd like to create a function with two arguments (a, axis=0) that computes the coefficient of variation of each column or row (2-dimensional array) and returns the index of the column or row with the maximum coefficient of variation.

I understand that .argmaxreturns the indices of the maximum values along an axis, but I'm unsure of how to proceed after that.

I'd like for the code to pass the following tests:

print(np.asscalar(arg_cvmax(b)) is 2)
print(np.asscalar(arg_cvmax(b,1)) is 0)
print(np.asscalar(arg_cvmax(b,axis=1)) is 0)

Answer 1

About putonspectacles's answer.

Since the variation calculate the ratio of the biased standard deviation to the mean, it make sense to change the np.var in original lambda to np.std.

The non-stats import version should be like this:

A = np.random.randn(10, 10)
cv =  lambda x: np.std(x) / np.mean(x)
var = np.apply_along_axis(cv, axis=0, arr=A)
idmax = np.argmax(var)