Compute RMS distances between frames in an mdtraj Trajectory object

compute_distance_matrix.py

import numpy as np
import mdtraj as md
import tqdm


def compute_distance_matrix(trajectory: md.Trajectory) -> np.ndarray:
    """Compute distance matrix.

    """
    topology = trajectory.topology
    CA_indices = [atom.index for atom in topology.atoms_by_name('CA')]
    num_frames = len(trajectory)
    distance_matrix = np.zeros((num_frames, num_frames), dtype=float)

    for i in tqdm.tqdm(range(num_frames)):
        rmsds = md.rmsd(trajectory, trajectory, frame=i,
                        atom_indices=CA_indices)

        distance_matrix[i, :] = rmsds

    return distance_matrix