Reimplemented index reader
This commit is contained in:
parent
cefdcf7e21
commit
6c54848105
@ -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:
|
||||
|
@ -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."""
|
||||
|
Loading…
Reference in New Issue
Block a user