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))