Repository: mdtraj/mdtraj
Pull request #2002: Allow ArcTrajectoryFile to optionally take topology
Author: login: jeremyleung521
Created: 2025-03-14T16:14:26Z | Merged: 2025-03-24T20:52:28Z

Been recently working on `.arc` trajectory files (yay Tinker). While the file itself contains _some_ information on topology, it does not have residue information. This means all atoms are [funneled into a single residue](https://github.com/mdtraj/mdtraj/blob/main/mdtraj/formats/arc.py#L100-L103). ([actual code](https://github.com/mdtraj/mdtraj/blob/main/mdtraj/formats/arc.py#L353-L356))

This PR adds in the functionality to feed in a topology (file name, trajectory objects, or topology objects) to `load_arc()`, allowing users to better read in arc trajectory files for analysis. Most of the code is essentially modeled off of the [NetCDF trajectory loader](https://github.com/mdtraj/mdtraj/blob/main/mdtraj/formats/netcdf.py), which requires a topology. Existing behavior for not feeding in a topology is retained.


Testing script:
```
import mdtraj
a = mdtraj.load('seg.arc', top='seg.pdb')
b = mdtraj.load('seg.arc')

print(len(a.topology._residues))  # prints 2112
print(len(b.topology._residues))  # prints 1

print(a.topology._residues[0])  # prints TYR1
print(b.topology._residues[0])  # prints RES1
```

Remove the `.txt` to run. Wouldn't upload to GitHub otherwise.
[seg.arc.txt](https://github.com/user-attachments/files/19250715/seg.arc.txt)
[seg.pdb.txt](https://github.com/user-attachments/files/19250717/seg.pdb.txt)



Diff:
diff --git a/mdtraj/formats/arc.py b/mdtraj/formats/arc.py
index 776375dc4..295025480 100644
--- a/mdtraj/formats/arc.py
+++ b/mdtraj/formats/arc.py
@@ -37,7 +37,7 @@ class _EOF(IOError):
 
 
 @FormatRegistry.register_loader(".arc")
-def load_arc(filename, stride=None, atom_indices=None, frame=None):
+def load_arc(filename, stride=None, atom_indices=None, frame=None, top=None):
     """Load a TINKER .arc file from disk.
 
     Parameters
@@ -53,6 +53,10 @@ def load_arc(filename, stride=None, atom_indices=None, frame=None):
         Use this option to load only a single frame from a trajectory on disk.
         If frame is None, the default, the entire trajectory will be loaded.
         If supplied, ``stride`` will be ignored.
+    top : {str, Trajectory, Topology}, optional
+        While the ARC format does does contain minimal topology information, it does
+        not include residue information. Pass in either a file path 
+        (e.g. to pdb), a trajectory, or a topology object to supply this information.
 
     Returns
     -------
@@ -69,9 +73,20 @@ def load_arc(filename, stride=None, atom_indices=None, frame=None):
             "filename must be of type path-like for load_arc. " "you supplied %s" % type(filename),
         )
 
+    # Process topology, if given.
+    if top is not None:
+        from mdtraj.core.trajectory import _parse_topology
+        topology = _parse_topology(top)
+    else:
+        topology = None
+
+
     atom_indices = cast_indices(atom_indices)
 
     with ArcTrajectoryFile(filename) as f:
+        # Set topology. If not given, it's None.
+        f.topology = topology
+
         if frame is not None:
             f.seek(frame)
             n_frames = 1
diff --git a/tests/test_arc.py b/tests/test_arc.py
index b7f4cacd9..8804c2646 100644
--- a/tests/test_arc.py
+++ b/tests/test_arc.py
@@ -1,10 +1,10 @@
 ##############################################################################
 # MDTraj: A Python Library for Loading, Saving, and Manipulating
 #         Molecular Dynamics Trajectories.
-# Copyright 2012-2017 Stanford University and the Authors
+# Copyright 2012-2025 Stanford University and the Authors
 #
 # Authors: Lee-Ping Wang
-# Contributors: Robert McGibbon
+# Contributors: Robert McGibbon, Jeremy M. G. Leung
 #
 # MDTraj is free software: you can redistribute it and/or modify
 # it under the terms of the GNU Lesser General Public License as
@@ -25,10 +25,12 @@
 
 import mdtraj as md
 from mdtraj.formats import ArcTrajectoryFile, PDBTrajectoryFile
+from mdtraj.formats.arc import load_arc
 from mdtraj.testing import eq
 
 
 def test_read_0(get_fn):
+    '''Check if .arc positions are read correctly'''
     with ArcTrajectoryFile(get_fn("4waters.arc")) as f:
         xyz, leng, ang = f.read()
     with PDBTrajectoryFile(get_fn("4waters.pdb")) as f:
@@ -42,6 +44,31 @@ def test_read_arctraj(get_fn):
     eq(traj.xyz, owntop.xyz)
 
 
+def test_read_arctraj_top(get_fn):
+    '''Check if anything other than residues differ between loading 
+    a ``.arc`` file with two different topologies'''
+    pdbtop = load_arc(get_fn("4waters.arc"), top=get_fn("4waters.pdb"))
+    arctop = load_arc(get_fn("4waters.arc"))
+ 
+    eq(pdbtop.n_atoms, arctop.n_atoms)
+    eq(pdbtop.xyz, arctop.xyz)
+
+    # Only the residues should differ
+    assert pdbtop.n_residues == 4
+    assert arctop.n_residues == 1
+
+    for ipdb, iarc in zip(pdbtop.top._atoms, arctop.top._atoms):
+        # check if atoms are in the same order with same element identity
+        eq(ipdb.name, iarc.name)
+        eq(ipdb.element, iarc.element)
+ 
+    for ipdb, iarc in zip(pdbtop.top._bonds, arctop.top._bonds):
+        # Check if connectivity is the same, based on index, which is 
+        # implicitly checked in the for loop above
+        eq(ipdb.atom1.index, iarc.atom1.index)
+        eq(ipdb.atom2.index, iarc.atom2.index)
+
+
 def test_read_pbc(get_fn):
     with ArcTrajectoryFile(get_fn("thf200.arc")) as f:
         xyz, leng, ang = f.read()
@@ -49,3 +76,4 @@ def test_read_pbc(get_fn):
     eq(xyz[0, -1], np.array([29.013880, 6.255754, 44.074519]), decimal=3)
     eq(leng, np.array([[45.119376, 45.119376, 45.119376]]), decimal=3)
     eq(ang, np.array([[90.0, 90.0, 90.0]]), decimal=3)
+
diff --git a/tests/test_trajectory.py b/tests/test_trajectory.py
index d9821ca2f..a774bd0a9 100644
--- a/tests/test_trajectory.py
+++ b/tests/test_trajectory.py
@@ -784,7 +784,10 @@ def test_open_and_load(get_fn):
         else:
             md.open(get_fn(file))
 
-        md.load(get_fn(file), top=get_fn("native.pdb"))
+        if file.endswith('.arc'):
+            md.load(get_fn(file), top=get_fn("4waters.pdb"))
+        else:
+            md.load(get_fn(file), top=get_fn("frame0.pdb"))
 
 
 def test_length(get_fn):


Reviews:
user: login: janash
body: I think it's a good addition, but I have a few comments/questions. I'd also like to investigate this more thoroughly myself, but I am at ACS this week! Hopefully can find the time.

I'd like to see a test added to `test_arc.py` to check for topology information in the `top` keyword is provided.
state: CHANGES_REQUESTED
author_association: COLLABORATOR
submitted_at: 2025-03-24T01:02:52Z
user: login: jeremyleung521
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-24T12:35:03Z
user: login: jeremyleung521
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-24T12:51:35Z
user: login: jeremyleung521
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-24T15:28:16Z
user: login: jeremyleung521
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-24T15:34:36Z
user: login: janash
body: LGTM - thanks! I think this makes using with Tinker much better.

[code that checks for top](https://github.com/mdtraj/mdtraj/blob/25b7fe4015944b892ffcb158008053506ec47389/mdtraj/core/trajectory.py#L423)

I'd be curious to know what other formats currently do not accept top because it does not seem too bad to add that capability
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-03-24T17:44:53Z

Review comments:
diff_hunk: @@ -37,13 +37,17 @@ class _EOF(IOError):
 
 
 @FormatRegistry.register_loader(".arc")
-def load_arc(filename, stride=None, atom_indices=None, frame=None):
+def load_arc(filename, top=None, stride=None, atom_indices=None, frame=None):
path: mdtraj/formats/arc.py
user: login: janash
body: Could putting `top` first cause any problems if someone is using one of the following arguments without the keyword in existing scripts? If coming to this through `load`, I assume it's handled elsewhere and won't be an issue
created_at: 2025-03-24T00:47:36Z
author_association: COLLABORATOR
diff_hunk: @@ -784,7 +784,10 @@ def test_open_and_load(get_fn):
         else:
             md.open(get_fn(file))
 
-        md.load(get_fn(file), top=get_fn("native.pdb"))
+        if file.endswith('.arc'):
+            md.load(get_fn(file), top=get_fn("4waters.pdb"))
path: tests/test_trajectory.py
user: login: janash
body: To me it seems that a test that there is actually a functioning topology associated with `.arc` files should be added. This function seems to essentially just test that the load function does not fail.

Could an additional test or condition be added to [test_arc.py](https://github.com/mdtraj/mdtraj/blob/main/tests/test_arc.py), to ensure that the topology is added and stored as expected?

Also, as an aside, this test [`test_read_arc_traj`](https://github.com/mdtraj/mdtraj/blob/25b7fe4015944b892ffcb158008053506ec47389/tests/test_arc.py#L39)  seems to assert that an ARC file loaded without a top and with a top are equivalent. It doesn't seem like this should be correct if a topology is now associated when using `top` (Am I misunderstanding something about this test?)
created_at: 2025-03-24T00:59:15Z
author_association: COLLABORATOR
diff_hunk: @@ -37,13 +37,17 @@ class _EOF(IOError):
 
 
 @FormatRegistry.register_loader(".arc")
-def load_arc(filename, stride=None, atom_indices=None, frame=None):
+def load_arc(filename, top=None, stride=None, atom_indices=None, frame=None):
path: mdtraj/formats/arc.py
user: login: jeremyleung521
body: Most load functions have top as the first keyword (second after filename, except for `.gro`), just trying to unify syntax. There will be potential errors if keywords arguments are provided positionally, but I'm assuming most people would use the keyword during input anyway, or through `mdtraj.load`.
created_at: 2025-03-24T12:35:03Z
author_association: COLLABORATOR
diff_hunk: @@ -784,7 +784,10 @@ def test_open_and_load(get_fn):
         else:
             md.open(get_fn(file))
 
-        md.load(get_fn(file), top=get_fn("native.pdb"))
+        if file.endswith('.arc'):
+            md.load(get_fn(file), top=get_fn("4waters.pdb"))
path: tests/test_trajectory.py
user: login: jeremyleung521
body: I'll add those tests later today. The `test_read_arc_traj` test assumes that the coordinates (xyz) are the same, which should be the same regardless of file.

What `test_open_and_load` was actually testing (before this commit) was if a topology was ignored for `.arc` and whether the other formats (`.mdcrd`, `.lammpstrj`) load with topology with the same number of atoms. I still don't know what the intention of this test is because it's testing different things for different formats within the same test and the name is not descriptive enough? (also why just those 3 formats?) Anyhow, at a minimum it checks if `open` and `load` pass without error with this limited selection of formats.
created_at: 2025-03-24T12:51:35Z
author_association: COLLABORATOR
diff_hunk: @@ -784,7 +784,10 @@ def test_open_and_load(get_fn):
         else:
             md.open(get_fn(file))
 
-        md.load(get_fn(file), top=get_fn("native.pdb"))
+        if file.endswith('.arc'):
+            md.load(get_fn(file), top=get_fn("4waters.pdb"))
path: tests/test_trajectory.py
user: login: jeremyleung521
body: DONE! (see https://github.com/mdtraj/mdtraj/pull/2002/commits/899067d4149aeb922f527a207e9af43a853ad0e4 )
created_at: 2025-03-24T15:28:16Z
author_association: COLLABORATOR
diff_hunk: @@ -37,13 +37,17 @@ class _EOF(IOError):
 
 
 @FormatRegistry.register_loader(".arc")
-def load_arc(filename, stride=None, atom_indices=None, frame=None):
+def load_arc(filename, top=None, stride=None, atom_indices=None, frame=None):
path: mdtraj/formats/arc.py
user: login: jeremyleung521
body: thought more about it. Since `.gro` is also putting it at the end (also optional for `.gro`), it should be fine at the end.
created_at: 2025-03-24T15:34:36Z
author_association: COLLABORATOR

PR comments:
user: login: janash
created_at: 2025-03-14T20:06:25Z
body: Awesome Jeremy! I am so in favor of better Tinker support! I've struggled with analyzing Tinker trajectories.
author_association: COLLABORATOR
user: login: jeremyleung521
created_at: 2025-03-24T17:49:02Z
body: > LGTM - thanks! I think this makes using with Tinker much better.
> 
> [code that checks for top](https://github.com/mdtraj/mdtraj/blob/25b7fe4015944b892ffcb158008053506ec47389/mdtraj/core/trajectory.py#L423)
> 
> I'd be curious to know what other formats currently do not accept top because it does not seem too bad to add that capability

only mol2 and hdf5... 
author_association: COLLABORATOR