Minimizing the used memory in diffusion simulation using Python

Question

I am recently dealing with a diffusion simulation project and I have come up with the following code:

N = 10000000
num_steps = 100
dim = 3

particles2 = npr.uniform(-1, 1, (N, num_steps, dim))
particles = np.cumsum(particles2, axis=1)

To briefly explain: the first 3 lines are the simulation parameters; the number of particles ($N$), the number of simulation timesteps (num_step) and the number of dimensions (dim) of the simulation. I am simulating diffusion by considering each particle as a random walker, so I used the numpy.random module to generate a random translation vector for each timestep, such that each coordinate is simply a random number between $-1$ and $1$. Now, in the particles2 line, I generate a tensor filled with random numbers between $-1$ and $1$ and then in the next line use cumulative sum to sum just over different timesteps. This is exactly the same as the following (more intuitive) code, but much, much more efficient:

particles = np.zeros((N, num_steps, dim))

for i in range(N):  
    prev_vec = np.zeros((dim))

    for t in range(num_steps):
        trans_vec = npr.uniform(-1, 1, dim)
        particles[i, t, :] = prev_vec
        prev_vec = prev_vec + trans_vec

where prev_vec stands for "previous vector" and trans_vec stands for the random translation vector of each timestep.

I have a problem, however - the memory usage. Since I am storing my entire trajectory for each particle, I (relatively) quickly hit the available memory cap. For the analysis, I only need the last timestep, so I don't really need the whole trajectory. I see how to get around this in the bottom code:

particles = np.zeros((N, dim))

for i in range(N):
    prev_vec = np.zeros((dim))

    for t in range(num_steps):
        trans_vec = npr.uniform(-1, 1, dim)
        prev_vec = prev_vec + trans_vec

    particles[i,:] = prev_vec

but I don't know how I would get around this in the above, a lot more efficient code. I would like to do this since the 2 for-loops are really making the calculation slow.

Answer 1

It seems that you are going from one extreme to the other: you probably want to generate all $N$ particles at once without the for-loop; however, you don't want to generate all the num_steps at once since you only need two last ones.

So, I think you are looking for something like:

import numpy as np
import numpy.random as npr

N = 10000000
dim = 3
num_steps = 100

# First, generating N initial particle positions (or you can initialize them to zeros)
particles = npr.uniform(-1, 1, (N, dim)) 

# start looping over time
for t in range(num_steps):
    # store the previous particles position
    prev_particles = particles
    # overwrite positions with the addition of random translation (no additional storage!)
    particles = np.add(particles,npr.uniform(-1,1,(N, dim)))  
    # notice, I could have saved the translation in the following way (commented out)
    # trans_vec = npr.uniform(-1,1,(N, dim)); particles = np.add(particles,trans_vec)
    # but that would have required additional storage.
    #   
    # do whatever you need with particles and prev_particles that store positions at the two preceding time_steps.

I might have a little bit misunderstood what goes where inside your code, but with that structure, you are doing the for-loop only for the time-dimension, where you need a lot of steps and you want to get the memory savings by getting rid of things you don't need to be storing.

The code above runs on my machine in less than a minute consuming less than 500 MB of storage (which makes sense, I am effectively storing 2 real vectors of the size $10\cdot10^6\times 1$ which would result in around 160 MB of storage at least + whatever overhead Python adds).