Efficiently Managing RAM Usage When Iterating Over Trajectories

[xiki-tempula]

Hi MDAnalysis developers,

First, I want to acknowledge MDAnalysis for its impressive ability to handle large trajectory files while keeping RAM usage under control. I'm seeking advice on optimizing RAM usage in a specific use case.

Use Case

I have a list of trajectories and need to:

1. Iterate through all frames in these trajectories.
2. Compute a collective variable (CV) for each frame (based only on the current frame’s data).
3. Perform calculations on the collected CV values.
4. Extract specific frames from the trajectories based on the results.

Ideally, the best-case scenario would involve loading and processing one frame at a time, ensuring constant and minimal RAM usage throughout.

Issue

To test this, I created an example with a system of 511,244 atoms and 10 trajectories, each containing 200 frames (1 GB each). The script iterates through the trajectories, computes a CV, and extracts a randomly selected frame (for simplicity).

Upon checking the RAM usage, it appears that a new frame is indeed loaded only when needed. However, the frame remains in RAM after it has been analysed, which is not ideal. Are there ways to unload the frame from RAM after it has been processed? Similarly, is there a RAM-efficient method to access a specific frame in the trajectory without having to load all preceding frames?

Test Script

Here’s the example code:

import MDAnalysis as mda
from memory_profiler import profile

@profile
def main():
    u = mda.Universe("amber.prm7", [f"amber_{i}.nc" for i in range(1, 11)], topology_format='PARM7')
    CV_list = []
    for ts in u.trajectory:
        # Dummy way of computing CV
        CV_list.append(ts.positions[0])
    # index = cluster(CV_list)
    index = 100
    u.trajectory[index]
    u.atoms.write('test.pdb')

if __name__ == "__main__":
    main()

Here is the RAM usage of each command

Line #    Mem usage    Increment  Occurrences   Line Contents
=============================================================
     3    139.2 MiB    139.2 MiB           1   @profile
     4                                         def main():
     5    465.9 MiB    326.7 MiB          13       u = mda.Universe("amber.prm7", [f"amber_{i}.nc" for i in range(1, 11)], topology_format='PARM7')
     6    465.9 MiB      0.0 MiB           1       CV_list = []
     7  12108.1 MiB  11642.1 MiB        2001       for ts in u.trajectory:
     8  12108.1 MiB      0.0 MiB        2000           CV_list.append(ts.positions[0])
     9                                             # index = cluster(CV_list)
    10  12108.1 MiB      0.0 MiB           1       index = 100
    11  12108.1 MiB      0.0 MiB           1       u.trajectory[index]
    12  12491.6 MiB    383.6 MiB           1       u.atoms.write('test.pdb')

How I Tested It

I used the memory profiler from Conda (conda install -c conda-forge memory_profiler):

mprof run --multiprocess --interval 0.001 python run.py
mprof plot

Question

What’s the best way to ensure that MDAnalysis processes only one frame at a time without loading the entire trajectory into RAM?

Thank you for your guidance!