How do I convert numpy mgrid function as a function?

Question

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)

Answer 1

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

Answer 2

mgrid is an instance of a cute class that lets us use indexing notation. Under the covers it uses np.linspace (or np.arange) to generate the ranges.

In [29]: x1,y1 = 0,1; x2,y2 = 1,3
In [30]: np.mgrid[x1:y1:3j, x2:y2:4j]
Out[30]: 
array([[[0.        , 0.        , 0.        , 0.        ],
        [0.5       , 0.5       , 0.5       , 0.5       ],
        [1.        , 1.        , 1.        , 1.        ]],

       [[1.        , 1.66666667, 2.33333333, 3.        ],
        [1.        , 1.66666667, 2.33333333, 3.        ],
        [1.        , 1.66666667, 2.33333333, 3.        ]]])

meshgrid is the function equivalent. I suspect it was the original function, and mgrid (and ogrid) was secondary version:

In [31]: np.meshgrid(np.linspace(x1,y1,3), np.linspace(x2,y2,4), indexing='ij')
Out[31]: 
[array([[0. , 0. , 0. , 0. ],
        [0.5, 0.5, 0.5, 0.5],
        [1. , 1. , 1. , 1. ]]),
 array([[1.        , 1.66666667, 2.33333333, 3.        ],
        [1.        , 1.66666667, 2.33333333, 3.        ],
        [1.        , 1.66666667, 2.33333333, 3.        ]])]

mgrid creates a n-d array; meshgrid returns a list of arrays. Otherwise they are equivalent. np.array(Out[31]) creates the array.

sparse versions

ogrid produces a "sparse" pair of arrays that, with broadcasting, functions the same way:

In [37]: np.ogrid[x1:y1:3j, x2:y2:4j]
Out[37]: 
[array([[0. ],
        [0.5],
        [1. ]]),
 array([[1.        , 1.66666667, 2.33333333, 3.        ]])]

meshgrid has an equivalent sparse mode:

In [38]: np.meshgrid(np.linspace(x1,y1,3), np.linspace(x2,y2,4), indexing='ij',
    ...: sparse=True)
Out[38]: 
[array([[0. ],
        [0.5],
        [1. ]]),
 array([[1.        , 1.66666667, 2.33333333, 3.        ]])]

We can create the same pair of arrays with:

In [39]: np.ix_(np.linspace(x1,y1,3), np.linspace(x2,y2,4))
Out[39]: 
(array([[0. ],
        [0.5],
        [1. ]]),
 array([[1.        , 1.66666667, 2.33333333, 3.        ]]))

or even:

In [40]: (np.linspace(x1,y1,3)[:,None], np.linspace(x2,y2,4)[None,:])
Out[40]: 
(array([[0. ],
        [0.5],
        [1. ]]),
 array([[1.        , 1.66666667, 2.33333333, 3.        ]]))