Element wise between a 2-D numpy array and a list-CodePudding

I'm trying to apply a function

def lead(x,n):
    if n>0:
        x = np.roll(x,-n)
        x[-n:]=1
    return x

to each element of Qxx, a 2-D numpy array (121,121), BUT WITH ROLLING the "n" argument from a list [0,1,2,3,4,....121] for example and in a element wise way.

the following code is working but SLOW !

xx = [[lead(qx,n) for n in range(len(qx))] for qx in Qxx]

how can I do it with apply_long_axis or map or...

smtg like :

xx = np.apply_along_axis(lead,1,arr = Qxx,n=range(121))

thanks

CodePudding user response：

This seems to be much faster:

list(map(lambda x: lead(Qxx[x], x), range(121)))

Performance:

The OP's solution:

%%timeit

[[lead(qx,n) for n in range(len(qx))] for qx in Qxx]

178 ms ± 3.32 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

My solution:

%%timeit

list(map(lambda x: lead(Qxx[x], x), range(121)))

1.63 ms ± 60.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

Data:

Qxx = np.array(np.tile(np.arange(121), 121)).reshape((121, 121))