can I pass a function a list index in python?

Question

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.

Welcome to Stack Overflow. I tried to [edit] the question to ask more clearly and directly. However, it's still not quite clear how best to help you. Could we see a corresponding example of the code you would try for the other approaches, and an explanation of what is wrong with the result? — Karl Knechtel, Jun 09 '22 at 18:41
Does https://stackoverflow.com/questions/48710508 answer your question? — Karl Knechtel, Jun 09 '22 at 18:47
Tell us about the function. Does it only accept scalar `x` and `y`. If so, it has to be called once for each `x,y` pair. `python `map` significantly faster than the for loop - and you'll need 2 maps. Your example takes arrays, since it just uses operators, `**` and `+`. — hpaulj, Jun 09 '22 at 18:49
The function I'm applying in my code takes the x and y indices of an array element and does some operations on them to convert it into a scalar (it actually does the x**2+y**2 and multiplies by a few constants) and then takes that number and does some corrections to make it between 0 and 2 pi. — JBowers, Jun 09 '22 at 19:25

hpaulj · Answer 1 · 2022-06-09T19:13:19.237

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.

Thank you for your input. Right now the function I'm applying in my code requires some additional computations:`def phaseconvert(x,y): #position of a pixel to the desired phase shift for that spatial position. phase = (x**2+y**2) floorcorrect = math.floor(phase/twopi) phase -= floorcorrect*twopi return(phase)` could np.fromfunction still be able to do this calculation? — JBowers, Jun 09 '22 at 19:16
@JBowers yes, you would do `np.fromfunction(phaseconvert, (1920, 1080))` I believe. However I'd suggest taking a look at my answer, which doesn't require having to fiddle with broadcasting. — ddejohn, Jun 09 '22 at 19:22
With the `math` function, you can only give it scalars. But most `math` functions have `numpy` equivalents. — hpaulj, Jun 09 '22 at 23:13

ddejohn · Accepted Answer · 2022-06-09T19:31:03.580

2

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 other benefits to using grids.

EDIT:

Using your function

import numpy as np
import matplotlib.pyplot as plt


twopi = 2 * np.pi


def phaseconvert(x, y):
    phase = x**2 + y**2
    floorcorrect = np.floor_divide(phase, twopi)
    return phase - floorcorrect * twopi


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


plt.contourf(x, y, phaseconvert(xx, yy))

edited Jun 09 '22 at 19:31

answered Jun 09 '22 at 18:46

ddejohn

8,775
3
17
30

Not really clear why you do phase - (phase / 2pi) * 2pi, fwiw. – ddejohn Jun 09 '22 at 19:31
thank you for your input. I need to do the phase - floorcorrect*2pi because the end result needs to be between 0 and 2pi (the application to this code is for some optics research where I need a specific phase shift on an incoming light wave, and subtracting a multiple of 2 pi from a phase is equivalent to the same phase) – JBowers Jun 09 '22 at 19:49

can I pass a function a list index in python?

2 Answers2