python - applying a function to rows of a ndarray && converting itertool object to numpy array

Question

I am trying to create permutations of size 4 from a group of real numbers. After that, I'd like to know the position of the first element in a permutation after I sort it. Here is what I have tried so far. What's the best way to do this?

import numpy as np
from itertools import chain, permutations

N_PLAYERS = 4
N_STATES = 60
np.random.seed(0)
state_space = np.linspace(0.0, 1.0, num=N_STATES, retstep=True)[0].tolist()
perms = permutations(state_space, N_PLAYERS)
perms_arr = np.fromiter(chain(*perms),dtype=np.float16)
def loc(row):
  return np.where(np.argsort(row) == 0)[0].tolist()[0]
locs = np.apply_along_axis(loc, 0, perms)

Answer 1

In [153]: N_PLAYERS = 4
     ...: N_STATES = 60
     ...: np.random.seed(0)
     ...: state_space = np.linspace(0.0, 1.0, num=N_STATES, retstep=True)[0].tolist()
     ...: perms = itertools.permutations(state_space, N_PLAYERS)
In [154]: alist = list(perms)
In [155]: len(alist)
Out[155]: 11703240

Simply making a list from the permuations produces a list of lists, with all sublists of length N_PLAYERS.

Making an array from that with chain flattens it:

In [156]: perms = itertools.permutations(state_space, N_PLAYERS)
In [158]: perms_arr = np.fromiter(itertools.chain(*perms),dtype=np.float16)
In [159]: perms_arr.shape
Out[159]: (46812960,)
In [160]: alist[0]

Which could be reshaped to (11703240,4).

Using apply on that 1d array doesn't work (or make sense):

In [170]: perms_arr.shape
Out[170]: (46812960,)
In [171]: locs = np.apply_along_axis(loc, 0, perms_arr)
In [172]: locs.shape
Out[172]: ()

Reshape to 4 columns:

In [173]: locs = np.apply_along_axis(loc, 0, perms_arr.reshape(-1,4))
In [174]: locs.shape
Out[174]: (4,)
In [175]: locs
Out[175]: array([     0, 195054, 578037, 769366])

This applies loc to each column, returning one value for each. But loc has a row variable. Is that supposed to be significant?

I could switch the axis; this takes much longer, and al

In [176]: locs = np.apply_along_axis(loc, 1, perms_arr.reshape(-1,4))
In [177]: locs.shape
Out[177]: (11703240,)

list comprehension

This iteration does the same thing as your apply_along_axis, and I expect is faster (though I haven't timed it - it's too slow).

In [188]: locs1 = np.array([loc(row) for row in perms_arr.reshape(-1,4)])
In [189]: np.allclose(locs, locs1)
Out[189]: True

whole array sort

But argsort takes an axis, so I can sort all rows at once (instead of iterating):

In [185]: np.nonzero(np.argsort(perms_arr.reshape(-1,4), axis=1)==0)
Out[185]: 
(array([       0,        1,        2, ..., 11703237, 11703238, 11703239]),
 array([0, 0, 0, ..., 3, 3, 3]))
In [186]: np.allclose(_[1],locs)
Out[186]: True

Or going the other direction: - cf with Out[175]

In [187]: np.nonzero(np.argsort(perms_arr.reshape(-1,4), axis=0)==0)
Out[187]: (array([     0, 195054, 578037, 769366]), array([0, 1, 2, 3]))