def topology(self, top):
"""Set the topology in the file
Parameters
----------
top : mdtraj.Topology
A topology object
"""
_check_mode(self.mode, ('w',))
if self._needs_initialization:
self._initialize_headers(top.n_atoms)
self._needs_initialization = False
top, bonds = top.to_dataframe()
data = {
"AtomID": top.index.values + 1,
"AtomNames": top.name.values,
"ResidueNames": top.resName.values,
"ChainID": top.chainID.values,
"ResidueID": top.resSeq.values + 1,
}
for key, val in iteritems(data):
node = self._get_node(where='/', name=key)[:] = val[:]
node[:] = val[:]
def _initialize_headers(self, n_atoms):
_check_mode(self.mode, ('w'))
self._create_carray(where='/',
name='AtomID',
atom=self.tables.Int64Atom(),
shape=(n_atoms, ))
self._create_carray(where='/',
name='AtomNames',
atom=self.tables.StringAtom(itemsize=4),
shape=(n_atoms, ))
self._create_carray(where='/',
name='ResidueNames',
atom=self.tables.StringAtom(itemsize=4),
shape=(n_atoms, ))
self._create_carray(where='/',
name='ChainID',
atom=self.tables.StringAtom(itemsize=1),
shape=(n_atoms, ))
self._create_carray(where='/',
name='ResidueID',
atom=self.tables.Int64Atom(),
shape=(n_atoms, ))
self._create_earray(where='/',
name='XYZList',
atom=self.tables.Int16Atom(),
shape=(0, n_atoms, 3))
def topology(self, top):
"""Set the topology in the file
Parameters
----------
top : mdtraj.Topology
A topology object
"""
_check_mode(self.mode, ('w', ))
if self._needs_initialization:
self._initialize_headers(top.n_atoms)
self._needs_initialization = False
top, bonds = top.to_dataframe()
data = {
"AtomID": top.index.values + 1,
"AtomNames": top.name.values,
"ResidueNames": top.resName.values,
"ChainID": top.chainID.values,
"ResidueID": top.resSeq.values + 1,
}
for key, val in iteritems(data):
node = self._get_node(where='/', name=key)[:] = val[:]
node[:] = val[:]
def read_as_traj(self, iteration=None, segment=None, atom_indices=None):
_check_mode(self.mode, ('r', ))
pnode = self._get_node(where='/', name='pointer')
iter_labels = pnode[:, 0]
seg_labels = pnode[:, 1]
if iteration is None and segment is None:
frame_indices = slice(None)
elif isinstance(iteration, (np.integer, int)) and isinstance(
segment, (np.integer, int)):
frame_torf = np.logical_and(iter_labels == iteration,
seg_labels == segment)
frame_indices = np.arange(len(iter_labels))[frame_torf]
else:
raise ValueError(
"iteration and segment must be integers and provided at the same time"
)
if len(frame_indices) == 0:
raise ValueError(
f"no frame was selected: iteration={iteration}, segment={segment}, atom_indices={atom_indices}"
)
iter_labels = iter_labels[frame_indices]
seg_labels = seg_labels[frame_indices]
topology = self.topology
if atom_indices is not None:
topology = topology.subset(atom_indices)
data = self.read(frame_indices=frame_indices,
atom_indices=atom_indices)
if len(data) == 0:
return Trajectory(xyz=np.zeros((0, topology.n_atoms, 3)),
topology=topology)
in_units_of(data.coordinates,
self.distance_unit,
Trajectory._distance_unit,
inplace=True)
in_units_of(data.cell_lengths,
self.distance_unit,
Trajectory._distance_unit,
inplace=True)
return WESTTrajectory(
data.coordinates,
topology=topology,
time=data.time,
unitcell_lengths=data.cell_lengths,
unitcell_angles=data.cell_angles,
iter_labels=iter_labels,
seg_labels=seg_labels,
pcoords=None,
)
def read(self, n_frames=None, stride=None, atom_indices=None):
"""Read one or more frames of data from the file
Parameters
----------
n_frames : {int, None}
The number of frames to read. If not supplied, all of the
remaining frames will be read.
stride : {int, None}
By default all of the frames will be read, but you can pass this
flag to read a subset of of the data by grabbing only every
`stride`-th frame from disk.
atom_indices : {int, None}
By default all of the atom will be read, but you can pass this
flag to read only a subsets of the atoms for the `coordinates` and
`velocities` fields. Note that you will have to carefully manage
the indices and the offsets, since the `i`-th atom in the topology
will not necessarily correspond to the `i`-th atom in your subset.
Returns
-------
xyz : np.ndarray, shape=(n_frames, n_atoms, 3), dtype=np.float32
The cartesian coordinates, in nanometers
"""
_check_mode(self.mode, ('r'))
if n_frames is None:
n_frames = np.inf
if stride is not None:
stride = int(stride)
if atom_indices is None:
atom_slice = slice(None)
else:
atom_slice = ensure_type(atom_indices,
dtype=np.int,
ndim=1,
name='atom_indices',
warn_on_cast=False)
total_n_frames = len(self._handle.root.XYZList)
frame_slice = slice(self._frame_index,
min(self._frame_index + n_frames, total_n_frames),
stride)
if frame_slice.stop - frame_slice.start == 0:
return np.array([], dtype=np.float32)
xyz = self._handle.root.XYZList.__getitem__((frame_slice, atom_slice))
if xyz.dtype == np.int16 or xyz.dtype == np.int32:
xyz = _convert_from_lossy_integers(xyz)
self._frame_index += (frame_slice.stop - frame_slice.start)
return xyz
def read_as_traj(self, n_frames=None, stride=None, atom_indices=None):
"""Read a trajectory from the LH5 file
Parameters
----------
n_frames : {int, None}
The number of frames to read. If not supplied, all of the
remaining frames will be read.
stride : {int, None}
By default all of the frames will be read, but you can pass this
flag to read a subset of of the data by grabbing only every
`stride`-th frame from disk.
atom_indices : {int, None}
By default all of the atom will be read, but you can pass this
flag to read only a subsets of the atoms for the `coordinates` and
`velocities` fields. Note that you will have to carefully manage
the indices and the offsets, since the `i`-th atom in the topology
will not necessarily correspond to the `i`-th atom in your subset.
Returns
-------
trajectory : Trajectory
A trajectory object containing the loaded portion of the file.
"""
_check_mode(self.mode, ('r', ))
from mdtraj.core.trajectory import Trajectory
topology = self.topology
if atom_indices is not None:
topology = topology.subset(atom_indices)
initial = int(self._frame_index)
xyz = self.read(n_frames=n_frames,
stride=stride,
atom_indices=atom_indices)
if len(xyz) == 0:
return Trajectory(xyz=np.zeros((0, topology.n_atoms, 3)),
topology=topology)
in_units_of(xyz,
self.distance_unit,
Trajectory._distance_unit,
inplace=True)
if stride is None:
stride = 1
time = (stride * np.arange(len(xyz))) + initial
return Trajectory(xyz=xyz, topology=topology, time=time)
def read(self, n_frames=None, stride=None, atom_indices=None):
"""Read one or more frames of data from the file
Parameters
----------
n_frames : {int, None}
The number of frames to read. If not supplied, all of the
remaining frames will be read.
stride : {int, None}
By default all of the frames will be read, but you can pass this
flag to read a subset of of the data by grabbing only every
`stride`-th frame from disk.
atom_indices : {int, None}
By default all of the atom will be read, but you can pass this
flag to read only a subsets of the atoms for the `coordinates` and
`velocities` fields. Note that you will have to carefully manage
the indices and the offsets, since the `i`-th atom in the topology
will not necessarily correspond to the `i`-th atom in your subset.
Returns
-------
xyz : np.ndarray, shape=(n_frames, n_atoms, 3), dtype=np.float32
The cartesian coordinates, in nanometers
"""
_check_mode(self.mode, ('r'))
if n_frames is None:
n_frames = np.inf
if stride is not None:
stride = int(stride)
if atom_indices is None:
atom_slice = slice(None)
else:
atom_slice = ensure_type(atom_indices, dtype=np.int, ndim=1,
name='atom_indices', warn_on_cast=False)
total_n_frames = len(self._handle.root.XYZList)
frame_slice = slice(self._frame_index, min(
self._frame_index + n_frames, total_n_frames), stride)
if frame_slice.stop - frame_slice.start == 0:
return np.array([], dtype=np.float32)
xyz = self._handle.root.XYZList.__getitem__((frame_slice, atom_slice))
if xyz.dtype == np.int16 or xyz.dtype == np.int32:
xyz = _convert_from_lossy_integers(xyz)
self._frame_index += (frame_slice.stop - frame_slice.start)
return xyz
def _initialize_headers(self, n_atoms):
_check_mode(self.mode, ('w'))
self._create_carray(
where='/', name='AtomID', atom=self.tables.Int64Atom(), shape=(n_atoms,))
self._create_carray(
where='/', name='AtomNames', atom=self.tables.StringAtom(itemsize=4),
shape=(n_atoms,))
self._create_carray(
where='/', name='ResidueNames', atom=self.tables.StringAtom(itemsize=4),
shape=(n_atoms,))
self._create_carray(
where='/', name='ChainID', atom=self.tables.StringAtom(itemsize=1),
shape=(n_atoms,))
self._create_carray(
where='/', name='ResidueID', atom=self.tables.Int64Atom(), shape=(n_atoms,))
self._create_earray(
where='/', name='XYZList', atom=self.tables.Int16Atom(),
shape=(0, n_atoms, 3))
def write(self, coordinates):
"""Write one or more frames of data to the file
Parameters
----------
coordinates : np.ndarray, dtype=np.float32, shape=(n_frames, n_atoms, 3)
The cartesian coordinates of the atoms in every frame, in nanometers.
"""
_check_mode(self.mode, ('w'))
coordinates = ensure_type(coordinates, dtype=np.float32, ndim=3,
name='coordinates', shape=(None, None, 3), can_be_none=False,
warn_on_cast=False, add_newaxis_on_deficient_ndim=True)
if self._needs_initialization:
self._initialize_headers(coordinates.shape[1])
self._needs_initialization = False
coordinates = _convert_to_lossy_integers(coordinates)
self._get_node(where='/', name='XYZList').append(coordinates)
def read_as_traj(self, n_frames=None, stride=None, atom_indices=None):
"""Read a trajectory from the LH5 file
Parameters
----------
n_frames : {int, None}
The number of frames to read. If not supplied, all of the
remaining frames will be read.
stride : {int, None}
By default all of the frames will be read, but you can pass this
flag to read a subset of of the data by grabbing only every
`stride`-th frame from disk.
atom_indices : {int, None}
By default all of the atom will be read, but you can pass this
flag to read only a subsets of the atoms for the `coordinates` and
`velocities` fields. Note that you will have to carefully manage
the indices and the offsets, since the `i`-th atom in the topology
will not necessarily correspond to the `i`-th atom in your subset.
Returns
-------
trajectory : Trajectory
A trajectory object containing the loaded portion of the file.
"""
_check_mode(self.mode, ('r',))
from mdtraj.core.trajectory import Trajectory
topology = self.topology
if atom_indices is not None:
topology = topology.subset(atom_indices)
initial = int(self._frame_index)
xyz = self.read(n_frames=n_frames, stride=stride, atom_indices=atom_indices)
if len(xyz) == 0:
return Trajectory(xyz=np.zeros((0, topology.n_atoms, 3)), topology=topology)
in_units_of(xyz, self.distance_unit, Trajectory._distance_unit, inplace=True)
if stride is None:
stride = 1
time = (stride*np.arange(len(xyz))) + initial
return Trajectory(xyz=xyz, topology=topology, time=time)
def seek(self, offset, whence=0):
"""Move to a new file position
Parameters
----------
offset : int
A number of frames.
whence : {0, 1, 2}
0: offset from start of file, offset should be >=0.
1: move relative to the current position, positive or negative
2: move relative to the end of file, offset should be <= 0.
Seeking beyond the end of a file is not supported
"""
_check_mode(self.mode, ('r', ))
if whence == 0 and offset >= 0:
self._frame_index = offset
elif whence == 1:
self._frame_index = self._frame_index + offset
elif whence == 2 and offset <= 0:
self._frame_index = len(self._handle.root.XYZList) + offset
else:
raise IOError('Invalid argument')
def seek(self, offset, whence=0):
"""Move to a new file position
Parameters
----------
offset : int
A number of frames.
whence : {0, 1, 2}
0: offset from start of file, offset should be >=0.
1: move relative to the current position, positive or negative
2: move relative to the end of file, offset should be <= 0.
Seeking beyond the end of a file is not supported
"""
_check_mode(self.mode, ('r',))
if whence == 0 and offset >= 0:
self._frame_index = offset
elif whence == 1:
self._frame_index = self._frame_index + offset
elif whence == 2 and offset <= 0:
self._frame_index = len(self._handle.root.XYZList) + offset
else:
raise IOError('Invalid argument')
def write(self, coordinates):
"""Write one or more frames of data to the file
Parameters
----------
coordinates : np.ndarray, dtype=np.float32, shape=(n_frames, n_atoms, 3)
The cartesian coordinates of the atoms in every frame, in nanometers.
"""
_check_mode(self.mode, ('w'))
coordinates = ensure_type(coordinates,
dtype=np.float32,
ndim=3,
name='coordinates',
shape=(None, None, 3),
can_be_none=False,
warn_on_cast=False,
add_newaxis_on_deficient_ndim=True)
if self._needs_initialization:
self._initialize_headers(coordinates.shape[1])
self._needs_initialization = False
coordinates = _convert_to_lossy_integers(coordinates)
self._get_node(where='/', name='XYZList').append(coordinates)
def read(self, frame_indices=None, atom_indices=None):
_check_mode(self.mode, ('r', ))
if frame_indices is None:
frame_slice = slice(None)
self._frame_index += frame_slice.stop - frame_slice.start
else:
frame_slice = ensure_type(frame_indices,
dtype=np.int,
ndim=1,
name='frame_indices',
warn_on_cast=False)
if not np.all(
frame_slice < self._handle.root.coordinates.shape[0]):
raise ValueError(
'As a zero-based index, the entries in '
'frame_slice must all be less than the number of frames '
'in the trajectory, %d' %
self._handle.root.coordinates.shape[0])
if not np.all(frame_slice >= 0):
raise ValueError(
'The entries in frame_indices must be greater '
'than or equal to zero')
self._frame_index += frame_slice[-1] - frame_slice[0]
if atom_indices is None:
# get all of the atoms
atom_slice = slice(None)
else:
atom_slice = ensure_type(atom_indices,
dtype=np.int,
ndim=1,
name='atom_indices',
warn_on_cast=False)
if not np.all(atom_slice < self._handle.root.coordinates.shape[1]):
raise ValueError(
'As a zero-based index, the entries in '
'atom_indices must all be less than the number of atoms '
'in the trajectory, %d' %
self._handle.root.coordinates.shape[1])
if not np.all(atom_slice >= 0):
raise ValueError('The entries in atom_indices must be greater '
'than or equal to zero')
def get_item(node, key):
if not isinstance(key, tuple):
return node.__getitem__(key)
n_list_like = 0
new_keys = []
for item in key:
if not isinstance(item, slice):
try:
d = np.diff(item)
if len(d) == 0:
item = item[0]
elif np.all(d == d[0]):
item = slice(item[0], item[-1] + d[0], d[0])
else:
n_list_like += 1
except Exception:
n_list_like += 1
new_keys.append(item)
new_keys = tuple(new_keys)
if n_list_like <= 1:
return node.__getitem__(new_keys)
data = node
for i, item in enumerate(new_keys):
dkey = [slice(None)] * len(key)
dkey[i] = item
dkey = tuple(dkey)
data = data.__getitem__(dkey)
return data
def get_field(name, slice, out_units, can_be_none=True):
try:
node = self._get_node(where='/', name=name)
data = get_item(node, slice)
in_units = node.attrs.units
if not isinstance(in_units, string_types):
in_units = in_units.decode()
data = in_units_of(data, in_units, out_units)
return data
except self.tables.NoSuchNodeError:
if can_be_none:
return None
raise
frames = Frames(
coordinates=get_field('coordinates',
(frame_slice, atom_slice, slice(None)),
out_units='nanometers',
can_be_none=False),
time=get_field('time', frame_slice, out_units='picoseconds'),
cell_lengths=get_field('cell_lengths', (frame_slice, slice(None)),
out_units='nanometers'),
cell_angles=get_field('cell_angles', (frame_slice, slice(None)),
out_units='degrees'),
velocities=get_field('velocities',
(frame_slice, atom_slice, slice(None)),
out_units='nanometers/picosecond'),
kineticEnergy=get_field('kineticEnergy',
frame_slice,
out_units='kilojoules_per_mole'),
potentialEnergy=get_field('potentialEnergy',
frame_slice,
out_units='kilojoules_per_mole'),
temperature=get_field('temperature',
frame_slice,
out_units='kelvin'),
alchemicalLambda=get_field('lambda',
frame_slice,
out_units='dimensionless'),
)
return frames