Numpy indexing of large meshgrid

user_15552943 0
Viewed 378

I would like to have a way to index a d-dimensional numpy meshgrid:

  • without actually storing the full dense meshgrid in memory
  • which supports all of the indexing types supported by the full dense meshgrid

As an example:

x = np.random.randn(Nx)
y = np.random.randn(Ny)
z = np.random.randn(Nz)

x_g, y_g, z_g = np.meshgrid(x, y, z, indexing='ij')
X = np.concatenate([x_g[..., None], y_g[..., None], z_g[..., None]], axis=-1)

Once I have X, I can use all of numpy's indexing methods, for example:

X[10:20,1:9:3,:]

X[(
    [0, 1, 3, 5],
    [1, 1, 3, 3],
    [2, 7, 8, 9]
    )]

I have cases in which the product Nx * Ny * Nz is just too large to fit in memory (e.g. Nx = Ny = Nz = 4000), yet I'd like to be able to take out bits of my cube using indexing.

Is there a way I can achieve this without re-coding numpy's indexing logic? The answer in 3d would look like:

def meshgrid_index(x, y, z, index):
    # index is the argument normally passed to X.__getitem__
    # should support all types of indexing
    ...

Thanks!

large-data indexing numpy python
1 Answers

Unfortunately, Numpy mostly works with array views. Indeed, a view is just a reference to a raw memory buffer, a datatype, a dimension (d), strides (d relative integers) and a shape (d positive integers). Strides are great to generate views of repeated elements, to reverse the order of items, but they are not powerful enough to generate the mesh-grid you want (I think this is possible to generate a subset but not the full grid since strides need somehow to change at some point).

Thus, you need either to create the mesh array (possibly just a part) or to do that yourself without the Numpy mesh feature. A Numba code with loops is often quite efficient for that and quite readable. What you want is likely a (generalized) lazy ND view which is very rarely supported by languages and libraries (because it is complex to implement and/or often not efficient).

Initial answer:

When there is not enough memory, the usual way to solve this problem is to work chunk-by-chunk. The idea is to iterate over all possible chunks using basic Python loops and slice the x, y and z arrays to produce a chunk of meshgrid (eg. np.meshgrid(x[xChunkId:xChunkId+xChunkSize], y[yChunkId:yChunkId+yChunkSize], z[zChunkId:zChunkId+zChunkSize], indexing='ij')). This approach often speed up computation of huge array that can fit in RAM because of CPU cache (assuming the chunks are small enough to fit in the cache). Loops do not introduce a significant overhead as long as chunks are big enough.

Note that generating such a huge (temporary) array may not even be mandatory. Indeed, tools like Numba or Cython can help you to iterate over the mesh grid without generating temporary array like this. This takes far less memory and is much faster in practice (due to CPU cache), not to mention these tool can be used to make your code run in parallel (assuming your algorithm can be parallelized). Note that these tools can also strongly reduce the cost of loops for small chunks (and using small chunks is often good for performance, again because of CPU cache).

When this is not possible because the full array needs to be computed at once, then the only solution is to store/load the huge arrays to storage devices. A fast and convenient way to do that is to use memory-mapped arrays. Note that this will be much slower though (especially with HDD).

Note that it would be very inefficient to store this specific array on storage device since it contains a lot of repeated values.

Reducing the size of arrays is critical for the sake of performance. Indeed, your array will probably take 1.4 TiB of memory by default and Numpy tends to often create many temporary arrays. Thus, the final storage capacity required can quickly exceed 10 TiB. HDD are very slow for such a use and high-performance SSD/persistent-memory are very expensive. Working on smaller types like 32-bits floats or integers (using the dtype parameter and the astype method) can help a lot. This is not free though this is is done at the expense of a lower accuracy and possible overflows.

user_129395570
Related