def test_4():
# using a really big box, we should get the same results with and without
# pbcs
box = np.array([[100, 0, 0], [0, 200, 0], [0, 0, 300]])
box = np.zeros((N_FRAMES, 3, 3)) + box #broadcast it out
a = _displacement_mic(xyz, pairs, box)
b = _displacement(xyz, pairs)
eq(a, b, decimal=3)
def test_4():
# using a really big box, we should get the same results with and without
# pbcs
box = np.array([[100, 0, 0], [0, 200, 0], [0, 0, 300]])
box = np.zeros((N_FRAMES, 3, 3)) + box # broadcast it out
a = _displacement_mic(xyz, pairs, box, False)
b = _displacement(xyz, pairs)
eq(a, b, decimal=3)
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)