Time-resolved Radial Distribution Functions (TRRDFs)

GPU memory in notebooks and HPC jobs

JAX normally reserves 75% of GPU memory at its first operation and retains its memory pool for reuse. Large, successively smaller allocation failures can be allocator reservation retries; those messages alone do not establish the size of the TRRDF working arrays or prove a memory leak.

Restart the notebook kernel, then run this as the first cell, before importing JAX, SPEADI, or another library that initializes JAX:

import os
os.environ.pop("XLA_PYTHON_CLIENT_ALLOCATOR", None)
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"

import speadi as sp

This enables allocation on demand, but JAX still caches freed buffers. If returning unused allocations to CUDA is required between calculations, also set os.environ["XLA_PYTHON_CLIENT_ALLOCATOR"] = "platform" in that first cell, before the import. Use this only for memory debugging or when returning memory is essential: the platform allocator can be dramatically slower, particularly for tiled calculations. Keep the default pooling allocator for normal analysis. The pop above removes a previous platform setting. Changing either setting after GPU initialization requires a kernel restart. See JAX's GPU memory documentation.

The JAX TRRDF path processes at most 16 frames and 128 atoms from each group per distance block by default (262,144 distances, or 1 MiB for a float32 distance array; intermediates, compilation and allocator overhead are extra). Completed histograms and the full returned g_rt array stay in NumPy host memory. The atom tile loop runs inside compiled JAX code, with one upload of each group's coordinates and one histogram download per frame chunk and group pair. Tiles reuse their device buffers; each frame chunk completes before another is submitted. Final tiles are padded and masked so ragged groups do not require separate kernels for their remainder tiles.

The distance bound excludes group coordinates, whose device storage scales linearly with the selected atoms and frame chunk size. It also does not cap the JAX allocator's reserved pool or host memory: MDTraj still loads one trajectory window, and the host output scales with group pairs, windows and bins.

Use pair_chunk_size and frame_chunk_size to tune the JAX working blocks:

r, g_rt = sp.trrdf(trajectory, g1, g2, top=top,
                   pair_chunk_size=64, frame_chunk_size=8)

Chunk sizes change execution memory and dispatch overhead, not the averaging window, stride or normalization. They are ignored by the NumPy and Numba backends. This blocking currently applies to trrdf, not int_trrdf or vanhove. Deleting NumPy results releases host memory; it does not flush JAX's GPU pool. jax.clear_caches() clears compilation caches and is not a GPU pool release API. Restarting the kernel terminates its GPU context.

Pair-correlation function g(r,t)

Provides functions to calculate the radial distribution function (RDF) between two groups of particles for specified windows along a trajectory g(r,t). Groups can also consist of single particles.

trrdf

trrdf(
    traj,
    g1,
    g2,
    top=None,
    pbc="ortho",
    n_windows=100,
    window_size=200,
    skip=1,
    stride=10,
    r_range=(0.0, 2.0),
    nbins=400,
    *,
    pair_chunk_size=128,
    frame_chunk_size=16
)

Calculate g(r,t) for two groups given in a trajectory. g(r) is calculated for each frame in the trajectory, then averaged over specified windows of time, returning g(r,t) (where t represents the window time along the trajectory).

PARAMETER	DESCRIPTION
traj	String pointing to the location of a trajectory that MDTraj is able to load TYPE: str
g1	List of numpy arrays of atom indices representing the group to calculate G(r,t) for TYPE: list
g2	List of numpy arrays of atom indices representing the group to calculate G(r,t) with TYPE: list
top	Topology object. Needed if trajectory given as a path to lazy-load. TYPE: topology DEFAULT: None
pbc	String representing the periodic boundary conditions of the simulation cell. Currently, only 'ortho' for orthogonal simulation cells is implemented. TYPE: (string, NoneType) DEFAULT: string
n_windows	Number of windows in which to split the trajectory (if a whole trajectory is supplied). TYPE: int DEFAULT: 100
window_size	Number of frames in each window. TYPE: int DEFAULT: 200
skip	Number of frames to skip at the beginning if giving a path as trajectory. TYPE: int DEFAULT: 1
stride	Number of frames in the original trajectory to skip between each calculation. E.g. stride = 10 means calculate distances only every 10th frame. TYPE: int DEFAULT: 10
r_range	Tuple over which r in g(r,t) is defined. TYPE: tuple(float, float) DEFAULT: (0.0, 2.0)
nbins	Number of bins (points in r to consider) in g(r,t) TYPE: int DEFAULT: 400
pair_chunk_size	Maximum atoms from each group per JAX distance block (default 128). Smaller values reduce device memory at the cost of more dispatches. Ignored by the NumPy and Numba backends. TYPE: int DEFAULT: 128
frame_chunk_size	Maximum frames per JAX distance block (default 16). Does not change the averaging window or stride. Ignored by other backends. TYPE: int DEFAULT: 16

RETURNS	DESCRIPTION
r	bin centers of g(r,t) TYPE: array
g_rt	Averaged function values of g(r,t) for each time from t=0 considered. Returned in NumPy host memory, including when using JAX. TYPE: array

Examples:

First, import both MDTraj and SPEADI together.

>>> import mdtraj as md
>>> import speadi as sp

Then, point to a particle simulation topology and trajectory (e.g. a Molecular Dynamics Simulation using Gromacs).

>>> topology = './topology.gro'
>>> trajectory = './trajectory.xtc'

Next, load the topology file using MDTraj and start defining reference and target groups.

>>> top = md.load_topology(topology)
>>> na = top.select('name NA')
>>> cl = top.select('name CL')
>>> protein_by_atom = [top.select(f'index {ix}') for
>>>                    ix in top.select('protein and not type H')]

Finally, run the Time-Resolved Radial Distribution Function (TRRDF) by calling trrdf().

>>> r, g_rt = sp.trrdf(trajectory, protein_by_atom, [na, cl], top=top,
>>>                     n_windows=1000, window_size=500, skip=0,
>>>                     pbc='general', stride=1, nbins=400)

The outputs are

• the centre points of the radial bins r

• the \(g(r,t)\) function with shape \(N\)(reference groups)\(\times N\)(target groups)\(\times N\)(windows)\(\times N\)(radial bins)

Source code in speadi/time_resolved_rdf/trrdf.py

def trrdf(traj, g1, g2, top=None, pbc='ortho', n_windows=100, window_size=200, skip=1, stride=10, r_range=(0.0, 2.0),
          nbins=400, *, pair_chunk_size=128, frame_chunk_size=16):
    """
    Calculate g(r,t) for two groups given in a trajectory. g(r) is
    calculated for each frame in the trajectory, then averaged over
    specified windows of time, returning g(r,t) (where t represents
    the window time along the trajectory).

    Parameters
    ----------
    traj : str
    	String pointing to the location of a trajectory that MDTraj is
    	able to load
    g1 : list
        List of numpy arrays of atom indices representing the group to
    	calculate G(r,t) for
    g2 : list
        List of numpy arrays of atom indices representing the group to
    	calculate G(r,t) with
    top : mdtraj.topology
        Topology object. Needed if trajectory given as a path to lazy-load.
    pbc : {string, NoneType}
        String representing the periodic boundary conditions of the
    	simulation cell. Currently, only 'ortho' for orthogonal simulation cells
    	is implemented.
    n_windows : int
        Number of windows in which to split the trajectory (if a whole
    	trajectory is supplied).
    window_size : int
        Number of frames in each window.
    skip : int
        Number of frames to skip at the beginning if giving a path as
    	trajectory.
    stride : int
        Number of frames in the original trajectory to skip between each
        calculation. E.g. stride = 10 means calculate distances only every
    	10th frame.
    r_range : tuple(float, float)
        Tuple over which r in g(r,t) is defined.
    nbins : int
        Number of bins (points in r to consider) in g(r,t)
    pair_chunk_size : int
        Maximum atoms from each group per JAX distance block (default 128).
        Smaller values reduce device memory at the cost of more dispatches.
        Ignored by the NumPy and Numba backends.
    frame_chunk_size : int
        Maximum frames per JAX distance block (default 16). Does not change
        the averaging window or stride. Ignored by other backends.

    Returns
    -------
    r : np.array
        bin centers of g(r,t)
    g_rt : np.array
        Averaged function values of g(r,t) for each time from t=0 considered.
        Returned in NumPy host memory, including when using JAX.

    Examples
    --------
    First, import both `MDTraj` and `SPEADI` together.
    >>> import mdtraj as md
    >>> import speadi as sp

    Then, point to a particle simulation topology and trajectory (e.g. a Molecular Dynamics Simulation using `Gromacs`).
    >>> topology = './topology.gro'
    >>> trajectory = './trajectory.xtc'

    Next, load the topology file using `MDTraj` and start defining reference and target groups.
    >>> top = md.load_topology(topology)
    >>> na = top.select('name NA')
    >>> cl = top.select('name CL')
    >>> protein_by_atom = [top.select(f'index {ix}') for
    >>>                    ix in top.select('protein and not type H')]

    Finally, run the Time-Resolved Radial Distribution Function (TRRDF) by calling `trrdf()`.
    >>> r, g_rt = sp.trrdf(trajectory, protein_by_atom, [na, cl], top=top,
    >>>                     n_windows=1000, window_size=500, skip=0,
    >>>                     pbc='general', stride=1, nbins=400)

    The outputs are

    - the centre points of the radial bins `r`

    - the $g(r,t)$ function with shape $N$(reference groups)$\\times N$(target groups)$\\times N$(windows)$\\times N$(radial bins)

    """
    for name, value in (("pair_chunk_size", pair_chunk_size), ("frame_chunk_size", frame_chunk_size)):
        if isinstance(value, bool) or not isinstance(value, Integral) or value < 1:
            raise ValueError(f"{name} must be a positive integer")
    g_rt, g1, g2 = _construct_results_array(g1, g2, n_windows, nbins)

    r, g_rt = _calculate_according_to_inputs(g1, g2, g_rt, n_windows, nbins, pbc, r_range, skip, stride,
                                             top, traj, window_size, pair_chunk_size, frame_chunk_size)

    return r, g_rt

Integral of g(r,t): n(r,t)

Provides functions to calculate the integral of theradial distribution function (RDF) between two groups of particles for specified windows along a trajectory n(r,t). Groups can also consist of single particles.

int_trrdf

int_trrdf(
    traj,
    g1,
    g2,
    top=None,
    pbc="ortho",
    n_windows=100,
    window_size=200,
    skip=1,
    stride=10,
    r_range=(0.0, 2.0),
    nbins=400,
)

Calculate n(r,t) for two groups given in a trajectory. n(r) is calculated for each frame in the trajectory, then averaged over specified windows of time, returning n(r,t) (where t represents the window time along the trajectory).

PARAMETER	DESCRIPTION
traj	String pointing to the location of a trajectory that MDTraj is able to load TYPE: str
g1	List of numpy arrays of atom indices representing the group to calculate G(r,t) for TYPE: list
g2	List of numpy arrays of atom indices representing the group to calculate G(r,t) with TYPE: list
top	Topology object. Needed if trajectory given as a path to lazy-load. TYPE: topology DEFAULT: None
pbc	String representing the periodic boundary conditions of the simulation cell. Currently, only 'ortho' for orthogonal simulation cells is implemented. TYPE: (string, NoneType) DEFAULT: string
n_windows	Number of windows in which to split the trajectory (if a whole trajectory is supplied). TYPE: int DEFAULT: 100
window_size	Number of frames in each window. TYPE: int DEFAULT: 200
skip	Number of frames to skip at the beginning if giving a path as trajectory. TYPE: int DEFAULT: 1
stride	Number of frames in the original trajectory to skip between each calculation. E.g. stride = 10 means calculate distances only every 10th frame. TYPE: int DEFAULT: 10
r_range	Tuple over which r in n(r,t) is defined. TYPE: tuple(float, float) DEFAULT: (0.0, 2.0)
nbins	Number of bins (points in r to consider) in n(r,t) TYPE: int DEFAULT: 400

RETURNS	DESCRIPTION
r	bin centers of n(r,t) TYPE: ndarray
n_rt	averaged function values of n(r,t) for each time from t=0 considered TYPE: ndarray

Examples:

First, import both MDTraj and SPEADI together.

>>> import mdtraj as md
>>> import speadi as sp

Then, point to a particle simulation topology and trajectory (e.g. a Molecular Dynamics Simulation using Gromacs).

>>> topology = './topology.gro'
>>> trajectory = './trajectory.xtc'

Next, load the topology file using MDTraj and start defining reference and target groups.

>>> top = md.load_topology(topology)
>>> na = top.select('name NA')
>>> cl = top.select('name CL')
>>> protein_by_atom = [top.select(f'index {ix}') for
>>>                    ix in top.select('protein and not type H')]

Finally, run the Time-Resolved Radial Distribution Function (TRRDF) by calling trrdf().

>>> r, n_rt = sp.int_trrdf(trajectory, protein_by_atom, [na, cl], top=top,
>>>                         n_windows=1000, window_size=500, skip=0,
>>>                         pbc='general', stride=1, nbins=400)

The outputs are

• the centre points of the radial bins r

• the \(n(r,t)\) function with shape \(N\)(reference groups)\(\times N\)(target groups)\(\times N\)(windows)\(\times N\)(radial bins)

Source code in speadi/int_trrdf/int_trrdf.py

def int_trrdf(traj, g1, g2, top=None, pbc='ortho', n_windows=100, window_size=200, skip=1, stride=10, r_range=(0.0, 2.0),
          nbins=400):
    """
    Calculate n(r,t) for two groups given in a trajectory. n(r) is
    calculated for each frame in the trajectory, then averaged over
    specified windows of time, returning n(r,t) (where t represents
    the window time along the trajectory).

    Parameters
    ----------
    traj : str
    	String pointing to the location of a trajectory that MDTraj is
    	able to load
    g1 : list
        List of numpy arrays of atom indices representing the group to
    	calculate G(r,t) for
    g2 : list
        List of numpy arrays of atom indices representing the group to
    	calculate G(r,t) with
    top : mdtraj.topology
        Topology object. Needed if trajectory given as a path to lazy-load.
    pbc : {string, NoneType}
        String representing the periodic boundary conditions of the
    	simulation cell. Currently, only 'ortho' for orthogonal simulation cells
    	is implemented.
    n_windows : int
        Number of windows in which to split the trajectory (if a whole
    	trajectory is supplied).
    window_size : int
        Number of frames in each window.
    skip : int
        Number of frames to skip at the beginning if giving a path as
    	trajectory.
    stride : int
        Number of frames in the original trajectory to skip between each
        calculation. E.g. stride = 10 means calculate distances only every
    	10th frame.
    r_range : tuple(float, float)
        Tuple over which r in n(r,t) is defined.
    nbins : int
        Number of bins (points in r to consider) in n(r,t)

    Returns
    -------
    r : np.ndarray
        bin centers of n(r,t)
    n_rt : np.ndarray
        averaged function values of n(r,t) for each time from t=0 considered

    Examples
    --------
    First, import both `MDTraj` and `SPEADI` together.
    >>> import mdtraj as md
    >>> import speadi as sp

    Then, point to a particle simulation topology and trajectory (e.g. a Molecular Dynamics Simulation using `Gromacs`).
    >>> topology = './topology.gro'
    >>> trajectory = './trajectory.xtc'

    Next, load the topology file using `MDTraj` and start defining reference and target groups.
    >>> top = md.load_topology(topology)
    >>> na = top.select('name NA')
    >>> cl = top.select('name CL')
    >>> protein_by_atom = [top.select(f'index {ix}') for
    >>>                    ix in top.select('protein and not type H')]

    Finally, run the Time-Resolved Radial Distribution Function (TRRDF) by calling `trrdf()`.
    >>> r, n_rt = sp.int_trrdf(trajectory, protein_by_atom, [na, cl], top=top,
    >>>                         n_windows=1000, window_size=500, skip=0,
    >>>                         pbc='general', stride=1, nbins=400)

    The outputs are

    - the centre points of the radial bins `r`

    - the $n(r,t)$ function with shape $N$(reference groups)$\\times N$(target groups)$\\times N$(windows)$\\times N$(radial bins)

    """
    n_rt, g1, g2 = _construct_results_array(g1, g2, n_windows, nbins)

    r, n_rt = _calculate_according_to_inputs(g1, g2, n_rt, n_windows, nbins, pbc, r_range, skip, stride,
                                             top, traj, window_size)

    return r, n_rt