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
– Karl Knechtel, Commented Jun 9, 2022 at 18:41
Does stackoverflow.com/questions/48710508 answer your question? — Karl Knechtel
– Karl Knechtel, Commented Jun 9, 2022 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
– hpaulj, Commented Jun 9, 2022 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 x2+y2 and multiplies by a few constants) and then takes that number and does some corrections to make it between 0 and 2 pi. — JBowers
– JBowers, Commented Jun 9, 2022 at 19:25

hpaulj · Accepted Answer · 2022-06-09 19:13:19Z

3

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.

edited Jun 9, 2022 at 19:13

answered Jun 9, 2022 at 18:59

hpaulj

233k14 gold badges260 silver badges392 bronze badges

Sign up to request clarification or add additional context in comments.

3 Comments

JBowers Over a year ago

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?

ddejohn Over a year ago

@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.

hpaulj Over a year ago

With the math function, you can only give it scalars. But most math functions have numpy equivalents.

ddejohn · Accepted Answer · 2022-06-09 19:31:03Z

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 9, 2022 at 19:31

answered Jun 9, 2022 at 18:46

ddejohn

9,0043 gold badges21 silver badges31 bronze badges

2 Comments

ddejohn Over a year ago

Not really clear why you do phase - (phase / 2pi) * 2pi, fwiw.

JBowers Over a year ago

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)

Collectives™ on Stack Overflow

can I pass a function a list index in python?

2 Answers 2

3 Comments

2 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

3 Comments

2 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related