can I pass a function a list index in python?-CodePudding

I want to create a 1920x1080 Numpy array where the value at each index is dependent on the index itself (for example Array[x][y] = f(x,y)).

I've been doing this so far with for loops, for example:

import numpy as np
xlength = 1920
ylength = 1080
def function(x,y):
    return(x**2 y**2)
indexarray = np.zeros((ylength,xlength),dtype = float)
for y in range(ylength):
    for x in range(xlength):
        indexarray[y][x]=function(x,y)

However, this is simply too slow for my purposes.

Is there a faster way to do this? I've tried using map() and numpy.apply_over_axes() to apply the function over a pre-built array where each array element is just a list of its x and y coordinates (like Array[x][y]=(x,y)), but I couldn't figure out how to make them work properly.

CodePudding user response：

You talk as though function requires scalar inputs:

In [111]: xlength = 5
     ...: ylength = 6
     ...: def function(x,y):
     ...:     return(x**2 y**2)
     ...: indexarray = np.zeros((ylength,xlength),dtype = float)
     ...: for y in range(ylength):
     ...:     for x in range(xlength):
     ...:         indexarray[y,x]=function(x,y)
     ...:         

In [112]: indexarray.shape
Out[112]: (6, 5)

In [113]: indexarray
Out[113]: 
array([[ 0.,  1.,  4.,  9., 16.],
       [ 1.,  2.,  5., 10., 17.],
       [ 4.,  5.,  8., 13., 20.],
       [ 9., 10., 13., 18., 25.],
       [16., 17., 20., 25., 32.],
       [25., 26., 29., 34., 41.]])

But it works just fine with arrays, for example with (6,1) and (5,) shape:

In [114]: np.arange(6)[:,None]**2   np.arange(5)**2
Out[114]: 
array([[ 0,  1,  4,  9, 16],
       [ 1,  2,  5, 10, 17],
       [ 4,  5,  8, 13, 20],
       [ 9, 10, 13, 18, 25],
       [16, 17, 20, 25, 32],
       [25, 26, 29, 34, 41]], dtype=int32)

Same as:

 function(np.arange(6)[:,None],np.arange(5))

There are lots of ways of generating the arrays for such a function. The other answer suggests np.meshgrid.

fromfunction takes a function and shape, and uses np.indices to create the arrays. But it still does just one call to the function.

In [117]: np.fromfunction(function,(6,5))
Out[117]: 
array([[ 0.,  1.,  4.,  9., 16.],
       [ 1.,  2.,  5., 10., 17.],
       [ 4.,  5.,  8., 13., 20.],
       [ 9., 10., 13., 18., 25.],
       [16., 17., 20., 25., 32.],
       [25., 26., 29., 34., 41.]])

np.indices((6,5)), np.meshgrid(np.arange(6),np.arange(5), indexing='ij') produce the same 2 arrays. Also np.mgrid.

If the function really only accepts scalars, then you are stuck with some sort of iteration, one that calls it each time for each x,y pair.

Pretending function only works with scalars, we can use:

In [126]: f = np.vectorize(function, otypes=['float']) 
In [127]: f(np.arange(6)[:,None], np.arange(5))

This still calls function for each x,y pair. For small examples, your iteration is faster, though for very large cases, this vectorize does better. But don't use either if you can write the function to work with arrays directly.

CodePudding user response：

For applying functions on grids, use np.meshgrid:

import numpy as np
import matplotlib.pyplot as plt


def f(x, y):
    return x**2   y**2


x = np.linspace(1, 1920, 1920)
y = np.linspace(1, 1080, 1080)
xx, yy = np.meshgrid(x, y, sparse=True)


plt.contourf(x, y, f(xx, yy))
plt.show()

The same can be achieved through broadcasting, but there are