Reimplemented index reader

This commit is contained in:
sebastiankloth 2022-05-02 17:20:34 +02:00
parent cefdcf7e21
commit 6c54848105
2 changed files with 37 additions and 11 deletions

View File

@ -15,7 +15,7 @@ __version__ = '22.5.dev1'
def open(directory='', topology='*.tpr', trajectory='*.xtc', cached=False,
nojump=False):
nojump=False, index_file=None):
"""
Open a simulation from a directory.
@ -64,7 +64,7 @@ def open(directory='', topology='*.tpr', trajectory='*.xtc', cached=False,
raise FileNotFoundError('Trajectory file could not be identified.')
atom_set, frames = reader.open_with_mdanalysis(
top_file, traj_file, cached=cached
top_file, traj_file, cached=cached, index_file=index_file
)
coords = coordinates.Coordinates(frames, atom_subset=atom_set)
if nojump:

View File

@ -18,12 +18,14 @@ import builtins
import warnings
import subprocess
import re
import itertools
import numpy as np
import MDAnalysis as mdanalysis
from scipy import sparse
from dask import delayed, __version__ as DASK_VERSION
import pandas as pd
import re
class NojumpError(Exception):
@ -32,7 +34,7 @@ class NojumpError(Exception):
class WrongTopologyError(Exception):
pass
def open_with_mdanalysis(topology, trajectory, cached=False):
def open_with_mdanalysis(topology, trajectory, cached=False, index_file=None):
"""Open a the topology and trajectory with mdanalysis."""
uni = mdanalysis.Universe(topology, trajectory, convert_units=False)
if cached is not False:
@ -45,19 +47,43 @@ def open_with_mdanalysis(topology, trajectory, cached=False):
reader = BaseReader(uni.trajectory)
reader.universe = uni
if topology.endswith('.tpr'):
atms = atoms.Atoms(
np.stack((uni.atoms.resids, uni.atoms.resnames, uni.atoms.names), axis=1),
charges=uni.atoms.charges, masses=uni.atoms.masses
).subset()
charges = uni.atoms.charges
masses = uni.atoms.masses
elif topology.endswith('.gro'):
atms = atoms.Atoms(
np.stack((uni.atoms.resids, uni.atoms.resnames, uni.atoms.names), axis=1),
charges=None, masses=None
).subset()
charges = None
masses = None
else:
raise WrongTopologyError('Topology file should end with ".tpr" or ".gro"')
indices = None
if index_file:
indices = load_indices(indexfile)
atms = atoms.Atoms(
np.stack((uni.atoms.resids, uni.atoms.resnames, uni.atoms.names), axis=1),
charges=charges, masses=masses, indices=indices
).subset()
return atms, reader
group_re = re.compile('\[ ([-+\w]+) \]')
def load_indices(indexfile):
indices = {}
index_array = None
with open(indexfile) as idx_file:
for line in idx_file:
m = group_re.search(line)
if m is not None:
group_name = m.group(1)
index_array = indices.get(group_name, [])
indices[group_name] = index_array
else:
elements = line.strip().split('\t')
elements = [x.split(' ') for x in elements]
elements = itertools.chain(*elements) # make a flat iterator
elements = [x for x in elements if x != '']
index_array += [int(x) - 1 for x in elements]
return indices
def is_writeable(fname):
"""Test if a directory is actually writeable, by writing a temporary file."""