Doing computations on a very large numpy array: streaming the calculation vs out-of-core memory

Question

I am trying to perform a calculation in numpy that depends on several parameters, and involved the creation of many intermediate arrays. These intermediate steps involve integrals over more parameters. Because of this, part of the way through my calculation I attempt to create an array that I realize is around 30GB in size at the lower end. My system can't handle this.

As a general example of the problem I am having, consider this Python code:

param1s = np.linspace(min1, max1, 100)
param2s = np.linspace(min2, max2, 100)
param3s = np.linspace(min3, max3, 100)

def main(param1, param2, param3):
    dummy_param_as = np.linspace(0, 10, 100)
    dummy_param_bs = np.linspace(-1, 1, 100)

    huge_array = do_big_calculation(params1, params2, params3,
                                    dummy_param_as, dummy_param_bs)

    integral_over_dummy_b = np.trapz(huge_array, x=dummy_param_bs axis=-1)

    return np.trapz(integral_over_dummy_b, x=dummy_param_as, axis=-1)

Specifically, do_big_calculation broadcasts all of the parameters together along separate axes, creating an array of shape (100, 100, 100, 100, 100). However, this array does not need to exist very long, just until it can be integrated over its last two axes.

Broadly there seems to be one highly recommended solution for this kind of situation: use something like h5py or dask to write the data to storage, and perform the calculation by loading data in blocks from the stored file. This type of solution seems to generally be referred to as an out-of-core method.

However, I have found another avenue that seems to be less discussed: stream the results out and iteratively construct the result, without ever touching storage. One such solution I found is npstreams.

For my toy example, the way this would work is to build the trapezoid rule sum of huge_array piece-by-piece in memory so that huge_array doesn't get loaded in memory at all. npstreams looks like one way of implementing this.

I am curious, given I am calculating my data on the fly rather than loading in stored data, what are the upsides and downsides of each of these approaches? Is one likely to be faster? Are there scaling issues with one or both?

Answer 1

A lot of this will depend on the details of your do_big_calculation function.

In general you want to avoid pushing data to disk for performance reasons. Disk I/O speed is significantly slower than memory speed.

There are some strategies that might help avoid creating that huge matrix in the first place.

If the output value at each location depends only on the parameters, i.e you could write

def do_big_calculation(a,b,c,d,e):
   result = np.zeros((len(a),len(b),len(c),len(d),len(e)):
   for ia in range(a):
       for ib in range(b):
           for ic in range(c):
               for id in range(d):
                   for ie in range(e):
                       result[ia,ib,ic,id]=do_big_calculation_core(
                             a[ia],b[ib],c[ic],d[id],e[ie]
                       )
   return result

Then there's no advantage to creating the big matrix, and not even any need to stream, you can just do the trapezoidal integration inside the loops.

def do_big_calculation_and_integrate(a,b,c,d,e):
   result = np.zeros((len(a),len(b),len(c)):
   for ia in range(a):
       for ib in range(b):
           for ic in range(c):
               temp1 = np.zeros(len(d))
               for id in range(d):
                   temp2 = np.zeros(len(e))
                   for ie in range(e):
                       temp2[ie]=do_big_calculation_core(
                             a[ia],b[ib],c[ic],d[id],e[ie]
                       )
                   temp1[id] = np.trapz(temp2,...)
               result[ia,ib,ic] = np.trapz(temp1)
   return result

Of course, this kind of explicit looping kills performance in numpy, so you'll need to be careful - In such cases reordering the loops and combining can help (at the cost of slightly increasing the memory. Also I've found that the numba package can optimise these loops reasonably well (especially if you can write do_big_calculation_core as a numba function).

Another option is to produce slices or blocks, and accumulate those into your final result. You might need to do this if there is significant work shared between parameter values, or numpy+numba is choking on your explicit loops.