Home > front end >  How do I convert numpy mgrid function as a function?
How do I convert numpy mgrid function as a function?

Time:03-23

Here is the way how numpy.mgrid is used.

grid = np.mgrid[x1:y1:100j , x2:y2:100j, ..., xn:yn:100j]

However, I find this structure very irritating. Therefore, I would like to create function f which works as follows:

f([(x1,y1,100),...,(xn,yn,100)]) = np.mgrid[x1:y1:100j , x2:y2:100j, ..., xn:yn:100j]

How can I create f?

(Here is the source code for np.mgrid)

CodePudding user response:

Just loop over each item passed to f and make a slice out of it with slice, and to get 100j from 100, multiply 100 by 1j:

def f(items):
    slices = [slice(i[0], i[1], 1j * i[2]) for i in items]
    return np.mgrid[slices]

Output:

>>> np.all( f([(1,2,5), (2,3,5)]) == np.mgrid[1:2:5j, 2:3:5j] )
True

You could make calling the function even simpler by using *items instead of items:

def f(*items):
    slices = [slice(i[0], i[1], 1j * i[2]) for i in items]
    return np.mgrid[slices]

Output:

>>> np.all( f([1,2,5], [2,3,5]) == np.mgrid[1:2:5j, 2:3:5j] )
True
  • Related