def _find_closest_contact(traj, group1, group2, frame=0, periodic=True):
xyz = ensure_type(
traj.xyz[frame],
dtype=np.float32,
ndim=2,
name="xyz",
shape=(None, 3),
warn_on_cast=False,
cast_da_to_np=True,
)
atoms1 = ensure_type(group1,
dtype=np.int32,
ndim=1,
name="group1",
warn_on_cast=False)
atoms2 = ensure_type(group2,
dtype=np.int32,
ndim=1,
name="group2",
warn_on_cast=False)
if periodic and traj._have_unitcell:
box = ensure_type(
traj.unitcell_vectors,
dtype=np.float32,
ndim=3,
name="unitcell_vectors",
shape=(len(traj.xyz), 3, 3),
warn_on_cast=False,
)[frame]
else:
box = None
ans = _geometry._find_closest_contact(xyz, atoms1, atoms2, box)
return np.asarray(ans)
def test_raising_on_wrong_ndim_arrays(self):
test = np.array([1.0], dtype=np.int32)
self.kwargs["ndim"] = 2
self.kwargs.pop("add_newaxis_on_deficient_ndim")
with pytest.raises(ValueError) as err:
_ = ensure_type(test, **self.kwargs)
assert "must be ndim" in str(err.value)
def _compute_displacements_chunk(xyz,
pairs,
box=None,
periodic=True,
opt=True,
orthogonal=False):
"""Compute displacements for a single chunk
Parameters
----------
xyz : ndarray of shape (any, any, 3)
The xyz coordinates of the chunk
pairs : array of shape (any, 2)
The indices for which to compute an angle
box : ndarray of shape (any, 3, 3)
The box vectors of the chunk
periodic : bool
Wether to use the periodc boundary during the calculation.
opt : bool, default=True
Use an optimized native library to calculate distances. MDTraj's
optimized SSE angle calculation implementation is 10-20x faster than
the (itself optimized) numpy implementation.
orthogonal : bool or da.bool
Wether all angles are close to 90 degrees
"""
# Cast orthogonal to a bool, just incase we got a delayed object
orthogonal = bool(orthogonal)
xyz = ensure_type(
xyz,
dtype=np.float32,
ndim=3,
name="xyz",
shape=(None, None, 3),
warn_on_cast=False,
cast_da_to_np=True,
)
if periodic and box is not None:
if opt:
out = np.empty((xyz.shape[0], pairs.shape[0], 3), dtype=np.float32)
_geometry._dist_mic_displacement(xyz, pairs,
box.transpose(0, 2, 1).copy(),
out, orthogonal)
return out
else:
return _displacement_mic(xyz, pairs, box.transpose(0, 2, 1),
orthogonal)
# Either there are no unitcell vectors or they dont want to use them
if opt:
out = np.empty((xyz.shape[0], pairs.shape[0], 3), dtype=np.float32)
_geometry._dist_displacement(xyz, pairs, out)
return out
else:
return _displacement(xyz, pairs)
def _compute_angles_chunk(xyz,
triplets,
box=None,
periodic=True,
opt=True,
orthogonal=False):
"""Compute the angles for a single chunk
Parameters
----------
xyz : ndarray of shape (any, any, 3)
The xyz coordinates of the chunk
triplets : array of shape (any, 3)
The indices for which to compute an angle
box : ndarray of shape (any, 3, 3)
The box vectors of the chunk
periodic : bool
Wether to use the periodc boundary during the calculation.
opt : bool, default=True
Use an optimized native library to calculate distances. MDTraj's
optimized SSE angle calculation implementation is 10-20x faster than
the (itself optimized) numpy implementation.
orthogonal : bool or da.bool
Wether all angles are close to 90 degrees
"""
# Cast dask.bool to a true bool
orthogonal = bool(orthogonal)
xyz = ensure_type(
xyz,
dtype=np.float32,
ndim=3,
name="xyz",
shape=(None, None, 3),
warn_on_cast=False,
cast_da_to_np=True,
)
out = np.empty((xyz.shape[0], triplets.shape[0]), dtype=np.float32)
if opt:
if periodic and box is not None:
_geometry._angle_mic(xyz, triplets,
box.transpose(0, 2, 1).copy(), out,
orthogonal)
else:
_geometry._angle(xyz, triplets, out)
else:
out = _angle(xyz, triplets, periodic, out).compute()
return out
def unitcell_lengths(self, value):
"""Set the lengths that define the shape of the unit cell in each frame
Parameters
----------
value : ndarray, shape=(n_frames, 3)
The distances ``a``, ``b``, and ``c`` that define the shape of the
unit cell in each frame, or None
"""
self._unitcell_lengths = ensure_type(
value,
np.float32,
2,
"unitcell_lengths",
can_be_none=True,
shape=(len(self), 3),
warn_on_cast=False,
add_newaxis_on_deficient_ndim=True,
)
def xyz(self, value):
"Set the cartesian coordinates of each atom in each simulation frame"
if self.top is not None:
# if we have a topology and its not None
shape = (None, self.topology._numAtoms, 3)
else:
shape = (None, None, 3)
value = ensure_type(
value,
np.float32,
3,
"xyz",
shape=shape,
warn_on_cast=True,
add_newaxis_on_deficient_ndim=True,
)
self._xyz = value
self._rmsd_traces = None
def unitcell_angles(self, value):
"""Set the lengths that define the shape of the unit cell in each frame
Parameters
----------
value : ndarray, shape=(n_frames, 3)
The angles ``alpha``, ``beta`` and ``gamma`` that define the
shape of the unit cell in each frame. The angles should be in
degrees.
"""
self._unitcell_angles = ensure_type(
value,
np.float32,
2,
"unitcell_angles",
can_be_none=True,
shape=(len(self), 3),
warn_on_cast=False,
add_newaxis_on_deficient_ndim=True,
)
def compute_angles(traj, angle_indices, periodic=True, opt=True, **kwargs):
""" Daskified version of mdtraj.compute_angles().
This mimics py:method:`mdtraj.compute_angles()` but returns the answer
as a py:class:`dask.array` object
Parameters
----------
traj : :py:class:`dask_traj.Trajectory`
The trajectory to compute the angles for.
angle_indices : array of shape(any, 3)
The indices for which to compute an angle.
periodic : bool
Wether to use the periodc boundary during the calculation.
opt : bool, default=True
Use an optimized native library to calculate distances. MDTraj's
optimized SSE angle calculation implementation is 10-20x faster than
the (itself optimized) numpy implementation.
Returns
-------
angles : dask.array, shape(n_frames, angle_indices)
Dask array with the delayed calculated angle for each item in
angle_indices for each frame.
"""
xyz = traj.xyz
length = len(xyz)
atoms = len(angle_indices)
triplets = ensure_type(
angle_indices,
dtype=np.int32,
ndim=2,
name="angle_indices",
shape=(None, 3),
warn_on_cast=False,
)
if not np.all(np.logical_and(triplets < traj.n_atoms, triplets >= 0)):
raise ValueError("angle_indices must be between 0 and %d" %
traj.n_atoms)
if len(triplets) == 0:
return da.zeros((len(xyz), 0), dtype=np.float32)
if periodic and traj._have_unitcell:
box = ensure_type(
traj.unitcell_vectors,
dtype=np.float32,
ndim=3,
name="unitcell_vectors",
shape=(len(xyz), 3, 3),
warn_on_cast=False,
)
else:
box = None
orthogonal = False
lazy_results = []
current_frame = 0
for frames in xyz.chunks[0]:
next_frame = current_frame + frames
if box is not None:
current_box = box[current_frame:next_frame]
orthogonal = da.allclose(
traj.unitcell_angles[current_frame:next_frame], 90)
else:
current_box = None
chunk_size = (frames, atoms)
lazy_results.append(
wrap_da(_compute_angles_chunk,
chunk_size,
xyz=xyz[current_frame:next_frame],
triplets=triplets,
box=current_box,
orthogonal=orthogonal,
opt=opt,
**kwargs))
current_frame = next_frame
max_result = da.concatenate(lazy_results)
result = max_result[:length]
return result
def test_wrong_shape(self):
test = np.array([1.0], dtype=np.int32)
with pytest.raises(ValueError) as err:
_ = ensure_type(test, shape=(3,), **self.kwargs)
assert "must be shape" in str(err.value)
def test_raising_on_wrong_length(self):
test = np.array([1.0], dtype=np.int32)
with pytest.raises(ValueError) as err:
_ = ensure_type(test, length=3, **self.kwargs)
assert "must be length" in str(err.value)
def test_adding_ndim_to_arrays(self):
test = np.array([1.0], dtype=np.int32)
self.kwargs["ndim"] = 2
b = ensure_type(test, **self.kwargs)
assert all(b == np.array([[1.0]], dtype=np.int32))
def test_type_error(self):
test = 1.0
self.kwargs.pop("add_newaxis_on_deficient_ndim")
with pytest.raises(TypeError) as err:
_ = ensure_type(test, **self.kwargs)
assert "must be numpy array" in str(err.value)
def test_scalar_conversion(self):
test = 1
self.kwargs["warn_on_cast"] = False
b = ensure_type(test, **self.kwargs)
assert all(b == np.array([1], dtype=np.int32))
def test_generators(self):
test = itt.repeat(1, 3) # This is a generator in python3
b = ensure_type(test, **self.kwargs)
assert all(b == np.array([1, 1, 1], dtype=np.int32))
