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!

[orbeckst]

• Perhaps do some manual garbage collection, gc.collect(), to force Python to release memory?
• The NCDFReader uses scipy.io.netcdf_file and exposes the mmap kwarg — try Universe(..., mmap=True), even though that should be the default. False loads the whole trajectory into memory.
• Does this happen if you only load a single trajectory fragment (a single nc file)? (To check that this is not the ChainReader keeping references that it shouldn't)
• Does this happen with different format (eg XTC or TRR)? (To check that it's not the ncdf reader)

[orbeckst]

Similarly, is there a RAM-efficient method to access a specific frame in the trajectory without having to load all preceding frames?

With a true random access file format such as AMBER netcdf, it should never be required to load the whole trajectory into memory. The readers seek to the frame position in the file and just load the data from the file offset.

For some formats that do not natively support random access (XTC/TRR, XYZ, ...) we scan the trajectory once as fast as we can and build an index of frames/offsets and then store it as a hidden file next to the trajectory file. Thus loading such a trajectory for the first time can take a while. Once the offset file is created, the next loading of the trajectory will be fast because the offsets are read from the hidden file.

[erikfransson]

I also ran into this with netcdf trajectories, after some debugging I could resolve it afaict but Im not fully sure about all the inner workings of MDAnalysis.

mmap=True doesnt read the trajectory into memory, but it doesnt solve the problem. The problem is the pages are never released: nothing unmaps them as you iterate, so after a full pass RSS equals the file size.

It's not a Python-object leak, which is why gc.collect() doesn't help when I tested.

Dropping each frame's pages once it's been copied into the Timestep fixes it:

import mmap

mm = u.trajectory.trjfile._mm
for ts in u.trajectory:
    ...
    mm.madvise(mmap.MADV_DONTNEED)

From my test, reading 10 000 frames of a 4.2 GiB .nc peaked RSS 4.26 GiB, which changed to 0.08 GiB (and was not any slower).

---

For the record Im using Python3.10 on ubuntu with MDAnalysis 2.9.0 / scipy 1.13.1 / numpy 2.2.6.

I could try a PR if that is of interest to fix this, although as said im not too familiar with MDAnalysis inner workings.

[orbeckst]

Interesting! I'll have to learn more about mmap. However, if you want to put up a PR then we can discuss more and see if it doesn't break anything (eg on windows).