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.