add atom_indices(): allows filtering for atoms using the mdtraj selec…

…tor language
axsk · Aug 21, 2024 · 4346c0d · 4346c0d
1 parent 93e97c2
commit 4346c0d
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/src/molutils.jl b/src/molutils.jl
@@ -138,4 +138,10 @@ function save_trajectory(filename, coords::AbstractMatrix; top::String)
     xyz = reshape(coords, 3, :, size(coords, 2))
     traj.xyz = PyReverseDims(xyz)
     traj.save(filename)
+end
+
+function atom_indices(filename::String, selector::String)
+    mdtraj = pyimport_conda("mdtraj", "mdtraj", "conda-forge")
+    traj = mdtraj.load(filename, stride=-1)
+    inds = traj.top.select(selector) .+ 1
 end