# Minimizing the used memory in diffusion simulation using Python

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.

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!)

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).