332 lines
10 KiB
Python
Executable File
332 lines
10 KiB
Python
Executable File
"""
|
|
Module that provides different readers for trajectory files.
|
|
|
|
It also provides a common interface layer between the file IO packages,
|
|
namely mdanalysis, and mdevaluate.
|
|
"""
|
|
from collections import namedtuple
|
|
import os
|
|
from os import path
|
|
from array import array
|
|
from zipfile import BadZipFile
|
|
import builtins
|
|
import re
|
|
import itertools
|
|
|
|
import numpy as np
|
|
import numpy.typing as npt
|
|
import MDAnalysis
|
|
from scipy import sparse
|
|
|
|
from .checksum import checksum
|
|
from .logging_util import logger
|
|
from . import atoms
|
|
from .coordinates import Coordinates
|
|
|
|
from unittest.mock import MagicMock
|
|
|
|
CSR_ATTRS = ("data", "indices", "indptr")
|
|
NOJUMP_MAGIC = 2016
|
|
Group_RE = re.compile(r"\[ ([-+\w]+) \]")
|
|
|
|
|
|
class NojumpError(Exception):
|
|
pass
|
|
|
|
|
|
class WrongTopologyError(Exception):
|
|
pass
|
|
|
|
|
|
class BaseReader:
|
|
"""Base class for trajectory readers."""
|
|
|
|
@property
|
|
def filename(self):
|
|
return self.rd.filename
|
|
|
|
@property
|
|
def nojump_matrices(self):
|
|
if self._nojump_matrices is None:
|
|
raise NojumpError("Nojump Data not available: {}".format(self.filename))
|
|
return self._nojump_matrices
|
|
|
|
@nojump_matrices.setter
|
|
def nojump_matrices(self, mats):
|
|
self._nojump_matrices = mats
|
|
|
|
def __init__(self, rd):
|
|
self.rd = rd
|
|
self._nojump_matrices = None
|
|
if path.exists(nojump_load_filename(self)):
|
|
load_nojump_matrices(self)
|
|
|
|
def __getitem__(self, item):
|
|
return self.rd[item]
|
|
|
|
def __len__(self):
|
|
return len(self.rd)
|
|
|
|
def __checksum__(self):
|
|
cache = array("L", self.rd._xdr.offsets.tobytes())
|
|
return checksum(self.filename, str(cache))
|
|
|
|
|
|
def open_with_mdanalysis(
|
|
topology: str,
|
|
trajectory: str,
|
|
index_file: str = None,
|
|
charges: npt.ArrayLike = None,
|
|
masses: npt.ArrayLike = None,
|
|
) -> (atoms.Atoms, BaseReader):
|
|
"""Open the topology and trajectory with mdanalysis."""
|
|
uni = MDAnalysis.Universe(topology, trajectory, convert_units=False)
|
|
reader = BaseReader(uni.trajectory)
|
|
reader.universe = uni
|
|
if topology.endswith(".tpr"):
|
|
charges = uni.atoms.charges
|
|
masses = uni.atoms.masses
|
|
elif topology.endswith(".gro"):
|
|
charges = charges
|
|
masses = masses
|
|
else:
|
|
raise WrongTopologyError('Topology file should end with ".tpr" or ".gro"')
|
|
indices = None
|
|
if index_file:
|
|
indices = load_indices(index_file)
|
|
|
|
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
|
|
|
|
|
|
def load_indices(index_file: str):
|
|
indices = {}
|
|
index_array = None
|
|
with open(index_file) 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: str):
|
|
"""Test if a directory is actually writeable, by writing a temporary file."""
|
|
fdir = os.path.dirname(fname)
|
|
ftmp = os.path.join(fdir, str(np.random.randint(999999999)))
|
|
while os.path.exists(ftmp):
|
|
ftmp = os.path.join(fdir, str(np.random.randint(999999999)))
|
|
|
|
if os.access(fdir, os.W_OK):
|
|
try:
|
|
with builtins.open(ftmp, "w"):
|
|
pass
|
|
os.remove(ftmp)
|
|
return True
|
|
except PermissionError:
|
|
pass
|
|
return False
|
|
|
|
|
|
def nojump_load_filename(reader: BaseReader):
|
|
directory, fname = path.split(reader.filename)
|
|
full_path = path.join(directory, ".{}.nojump.npz".format(fname))
|
|
if not is_writeable(directory):
|
|
user_data_dir = os.path.join("/data/", os.environ["HOME"].split("/")[-1])
|
|
full_path_fallback = os.path.join(
|
|
os.path.join(user_data_dir, ".mdevaluate/nojump"),
|
|
directory.lstrip("/"),
|
|
".{}.nojump.npz".format(fname),
|
|
)
|
|
if os.path.exists(full_path_fallback):
|
|
return full_path_fallback
|
|
if os.path.exists(full_path) or is_writeable(directory):
|
|
return full_path
|
|
else:
|
|
user_data_dir = os.path.join("/data/", os.environ["HOME"].split("/")[-1])
|
|
full_path = os.path.join(
|
|
os.path.join(user_data_dir, ".mdevaluate/nojump"),
|
|
directory.lstrip("/"),
|
|
".{}.nojump.npz".format(fname),
|
|
)
|
|
return full_path
|
|
|
|
|
|
def nojump_save_filename(reader: BaseReader):
|
|
directory, fname = path.split(reader.filename)
|
|
full_path = path.join(directory, ".{}.nojump.npz".format(fname))
|
|
if is_writeable(directory):
|
|
return full_path
|
|
else:
|
|
user_data_dir = os.path.join("/data/", os.environ["HOME"].split("/")[-1])
|
|
full_path_fallback = os.path.join(
|
|
os.path.join(user_data_dir, ".mdevaluate/nojump"),
|
|
directory.lstrip("/"),
|
|
".{}.nojump.npz".format(fname),
|
|
)
|
|
logger.info(
|
|
"Saving nojump to {}, since original location is not writeable.".format(
|
|
full_path_fallback
|
|
)
|
|
)
|
|
os.makedirs(os.path.dirname(full_path_fallback), exist_ok=True)
|
|
return full_path_fallback
|
|
|
|
|
|
def parse_jumps(trajectory: Coordinates, whole: bool=True, fractional_inverted: bool=True):
|
|
if whole:
|
|
prev = trajectory[0].whole
|
|
else:
|
|
prev = trajectory[0]
|
|
box = prev.box
|
|
if fractional_inverted:
|
|
s_prev = prev @ np.linalg.inv(box)
|
|
|
|
SparseData = namedtuple("SparseData", ["data", "row", "col"])
|
|
jump_data = (
|
|
SparseData(data=array("b"), row=array("l"), col=array("l")),
|
|
SparseData(data=array("b"), row=array("l"), col=array("l")),
|
|
SparseData(data=array("b"), row=array("l"), col=array("l")),
|
|
)
|
|
|
|
for i, curr in enumerate(trajectory):
|
|
if i % 500 == 0:
|
|
logger.debug("Parse jumps Step: %d", i)
|
|
if not fractional_inverted:
|
|
r3 = np.subtract(curr, prev)
|
|
delta_z = np.array(np.rint(np.divide(r3[:, 2], box[2][2])), dtype=np.int8)
|
|
r2 = np.subtract(
|
|
r3,
|
|
(np.rint(np.divide(r3[:, 2], box[2][2])))[:, np.newaxis]
|
|
* box[2][np.newaxis, :],
|
|
)
|
|
delta_y = np.array(np.rint(np.divide(r2[:, 1], box[1][1])), dtype=np.int8)
|
|
r1 = np.subtract(
|
|
r2,
|
|
(np.rint(np.divide(r2[:, 1], box[1][1])))[:, np.newaxis]
|
|
* box[1][np.newaxis, :],
|
|
)
|
|
delta_x = np.array(np.rint(np.divide(r1[:, 0], box[0][0])), dtype=np.int8)
|
|
delta = np.array([delta_x, delta_y, delta_z]).T
|
|
prev = curr
|
|
box = prev.box
|
|
else:
|
|
s_curr = curr @ np.linalg.inv(curr.box)
|
|
|
|
ds = s_curr - s_prev
|
|
delta = np.array(np.rint(ds), dtype=np.int8)
|
|
|
|
s_prev = s_curr
|
|
|
|
|
|
for d in range(3):
|
|
(col,) = np.where(delta[:, d] != 0)
|
|
jump_data[d].col.extend(col)
|
|
jump_data[d].row.extend([i] * len(col))
|
|
jump_data[d].data.extend(delta[col, d])
|
|
|
|
return jump_data
|
|
|
|
|
|
def generate_nojump_matrices(trajectory: Coordinates):
|
|
"""
|
|
Create the matrices with pbc jumps for a trajectory.
|
|
"""
|
|
logger.info("generate Nojump matrices for: {}".format(trajectory))
|
|
|
|
jump_data = parse_jumps(trajectory)
|
|
N = len(trajectory)
|
|
M = len(trajectory[0])
|
|
|
|
trajectory.frames.nojump_matrices = tuple(
|
|
sparse.csr_matrix((np.array(m.data), (m.row, m.col)), shape=(N, M))
|
|
for m in jump_data
|
|
)
|
|
save_nojump_matrices(trajectory.frames)
|
|
|
|
|
|
def _ensure_xdr(reader: BaseReader):
|
|
"""Patch missing _xdr attribute for non-XDR readers (e.g. LAMMPS DumpReader)
|
|
with a stable mock so checksums are consistent across runs."""
|
|
if not hasattr(reader.rd, '_xdr'):
|
|
mock_xdr = MagicMock()
|
|
mock_xdr.offsets = np.arange(len(reader))
|
|
print(f"Adding mock _xdr attribute for to reader of length {len(reader)}.")
|
|
reader.rd._xdr = mock_xdr
|
|
|
|
|
|
def save_nojump_matrices(reader: BaseReader, matrices: npt.ArrayLike = None):
|
|
_ensure_xdr(reader)
|
|
if matrices is None:
|
|
matrices = reader.nojump_matrices
|
|
data = {"checksum": checksum(NOJUMP_MAGIC, checksum(reader))}
|
|
for d, mat in enumerate(matrices):
|
|
data["shape"] = mat.shape
|
|
for attr in CSR_ATTRS:
|
|
data["{}_{}".format(attr, d)] = getattr(mat, attr)
|
|
|
|
np.savez(nojump_save_filename(reader), **data)
|
|
|
|
|
|
def load_nojump_matrices(reader: BaseReader):
|
|
_ensure_xdr(reader)
|
|
zipname = nojump_load_filename(reader)
|
|
try:
|
|
data = np.load(zipname, allow_pickle=True)
|
|
except (AttributeError, BadZipFile, OSError):
|
|
# npz-files can be corrupted, probably a bug for big arrays saved with
|
|
# savez_compressed?
|
|
logger.info("Removing zip-File: %s", zipname)
|
|
os.remove(nojump_load_filename(reader))
|
|
return
|
|
try:
|
|
if data["checksum"] == checksum(NOJUMP_MAGIC, checksum(reader)):
|
|
reader.nojump_matrices = tuple(
|
|
sparse.csr_matrix(
|
|
tuple(data["{}_{}".format(attr, d)] for attr in CSR_ATTRS),
|
|
shape=data["shape"],
|
|
)
|
|
for d in range(3)
|
|
)
|
|
logger.info(
|
|
"Loaded Nojump matrices: {}".format(nojump_load_filename(reader))
|
|
)
|
|
else:
|
|
logger.info("Invalid Nojump Data: {}".format(nojump_load_filename(reader)))
|
|
except KeyError:
|
|
logger.info("Removing zip-File: %s", zipname)
|
|
os.remove(nojump_load_filename(reader))
|
|
return
|
|
|
|
|
|
def correct_nojump_matrices_for_whole(trajectory: Coordinates):
|
|
reader = trajectory.frames
|
|
frame = trajectory[0]
|
|
box = frame.box.diagonal()
|
|
cor = ((frame - frame.whole) / box).round().astype(np.int8)
|
|
for d in range(3):
|
|
reader.nojump_matrices[d][0] = cor[:, d]
|
|
save_nojump_matrices(reader)
|
|
|
|
|
|
def energy_reader(file: str):
|
|
"""Reads a gromacs energy file with mdanalysis and returns an auxiliary file.
|
|
Args:
|
|
file: Filename of the energy file
|
|
"""
|
|
return MDAnalysis.auxiliary.EDR.EDRReader(file)
|