Juan Nunez-Iglesias
2017-03-09 02:48:45 UTC
I was a bit surprised to discover that both meshgrid nor mgrid return fully instantiated arrays, when simple broadcasting (ie with stride=0 for other axes) is functionally identical and happens much, much faster.
I wrote my own function to do this:
def broadcast_mgrid(arrays):
  shape = tuple(map(len, arrays))
  ndim = len(shape)
  result = []
  for i, arr in enumerate(arrays, start=1):
    reshaped = np.broadcast_to(arr[[...] + [np.newaxis] * (ndim - i)],
                  shape)
    result.append(reshaped)
  return result
For even a modest-sized 512 x 512 grid, this version is close to 100x faster:
In [154]: %timeit th.broadcast_mgrid((np.arange(512), np.arange(512)))
10000 loops, best of 3: 25.9 µs per loop
In [156]: %timeit np.meshgrid(np.arange(512), np.arange(512))
100 loops, best of 3: 2.02 ms per loop
In [157]: %timeit np.mgrid[:512, :512]
100 loops, best of 3: 4.84 ms per loop
Is there a conscious design decision as to why this isnât what meshgrid/mgrid do already? Or would a PR be welcome to do this?
Thanks,
Juan.
I wrote my own function to do this:
def broadcast_mgrid(arrays):
  shape = tuple(map(len, arrays))
  ndim = len(shape)
  result = []
  for i, arr in enumerate(arrays, start=1):
    reshaped = np.broadcast_to(arr[[...] + [np.newaxis] * (ndim - i)],
                  shape)
    result.append(reshaped)
  return result
For even a modest-sized 512 x 512 grid, this version is close to 100x faster:
In [154]: %timeit th.broadcast_mgrid((np.arange(512), np.arange(512)))
10000 loops, best of 3: 25.9 µs per loop
In [156]: %timeit np.meshgrid(np.arange(512), np.arange(512))
100 loops, best of 3: 2.02 ms per loop
In [157]: %timeit np.mgrid[:512, :512]
100 loops, best of 3: 4.84 ms per loop
Is there a conscious design decision as to why this isnât what meshgrid/mgrid do already? Or would a PR be welcome to do this?
Thanks,
Juan.