Script to reimage Anton2 trajectories

anton2_reimage.py

#!/usr/bin/env python

"""
Python script to reimage/unwrap an MD simulation trajectory
ran on the Anton2 supercomputer.

Expects a topology file in .dms format and the trajectory
folder '.dtr'.

Requires:
  - mdtraj (http://mdtraj.org/1.9.3/)
  - openmm (http://openmm.org/)

These libraries can be installed with:
  - conda install mdtraj
  - conda -c omnia install openmm

Joao Rodrigues @ 2019
"""

import argparse
import os

import mdtraj as md
import simtk.openmm.app as app

##

ap = argparse.ArgumentParser()
ap.add_argument('topology', type=str, 
                help='Topology file in .dms format')
ap.add_argument('trajectory', type=str,
                help='Trajectory folder in .dtr format')
ap.add_argument('--output', type=str, default='reimaged.dcd',
                help="Reimaged trajectory name and format. Default 'reimaged.dcd'")
args = ap.parse_args()

if not os.path.isfile(args.topology):
    print(f'Topology file not found: {args.topology}')
    sys.exit(1)

if not os.path.isdir(args.trajectory):
    print(f'Trajectory folder not found: {args.trajectory}')
    sys.exit(1)

# Parse topology and trajectory   
dms = app.DesmondDMSFile(args.topology)
print(f'Parsed topology file: {args.topology}')

t = md.load(args.trajectory, top=md.Topology.from_openmm(dms.topology))
print(f'Parsed trajectory: {args.trajectory}')
print(f'  no. atoms: {t.n_atoms}')
print(f'  no. frames: {t.n_frames}')

# Reimage using as anchor largest molecule
print('Reimaging trajectory')
mols = t.topology.find_molecules()
t_reimaged = t.image_molecules(anchor_molecules=mols[:1])

# # Write topology in PDB format.
# with open('reimaged.pdb', 'w') as handle:
#     app.PDBFile.writeFile(dms.topology, t_reimaged.openmm_positions(0), handle, keepIds=True)

# Write trajectory in DCD format.
print('Saving reimaged trajectory')
t_reimaged.save(args.output)