def test_adjacency_matrix(cell_size, threshold, periodic):
"""
Compare the construction of an adjacency matrix using a cell list
and using a computationally expensive but simpler distance matrix.
"""
array = strucio.load_structure(join(data_dir, "3o5r.mmtf"))
if periodic:
# Create an orthorhombic box
# with the outer coordinates as bounds
array.box = np.diag(
np.max(array.coord, axis=-2) - np.min(array.coord, axis=-2))
cell_list = struc.CellList(array, cell_size=cell_size, periodic=periodic)
matrix = cell_list.create_adjacency_matrix(threshold)
# Create distance matrix
# Convert to float64 to avoid errorenous warning
# https://github.com/ContinuumIO/anaconda-issues/issues/9129
array.coord = array.coord.astype(np.float64)
length = array.array_length()
distance = struc.index_distance(
array,
np.stack([
np.repeat(np.arange(length), length),
np.tile(np.arange(length), length)
],
axis=-1), periodic)
distance = np.reshape(distance, (length, length))
# Create adjacency matrix from distance matrix
expected_matrix = (distance <= threshold)
# Both ways to create an adjacency matrix
# should give the same result
assert np.array_equal(matrix, expected_matrix)
def test_index_functions():
"""
The `index_xxx()` functions should give the same result as the
corresponding `xxx` functions.
"""
stack = strucio.load_structure(join(data_dir, "1l2y.mmtf"))
array = stack[0]
# Test for atom array, stack and raw coordinates
samples = (array, stack, struc.coord(array), struc.coord(stack))
# Generate random indices
random.seed(42)
indices = random.randint(array.array_length(), size=(100, 4), dtype=int)
for sample in samples:
if isinstance(sample, np.ndarray):
atoms1 = sample[..., indices[:, 0], :]
atoms2 = sample[..., indices[:, 1], :]
atoms3 = sample[..., indices[:, 2], :]
atoms4 = sample[..., indices[:, 3], :]
else:
atoms1 = sample[..., indices[:, 0]]
atoms2 = sample[..., indices[:, 1]]
atoms3 = sample[..., indices[:, 2]]
atoms4 = sample[..., indices[:, 3]]
assert np.allclose(struc.displacement(atoms1, atoms2),
struc.index_displacement(sample, indices[:, :2]),
atol=1e-5)
assert np.allclose(struc.distance(atoms1, atoms2),
struc.index_distance(sample, indices[:, :2]),
atol=1e-5)
assert np.allclose(struc.angle(atoms1, atoms2, atoms3),
struc.index_angle(sample, indices[:, :3]),
atol=1e-5)
assert np.allclose(struc.dihedral(atoms1, atoms2, atoms3, atoms4),
struc.index_dihedral(sample, indices[:, :4]),
atol=1e-5)
def test_index_distance_periodic_triclinic(shift, angles):
"""
The PBC aware computation, should give the same results,
irrespective of which atoms are centered in the box
"""
array = strucio.load_structure(join(data_dir("structure"), "3o5r.mmtf"))
# Use a box based on the boundaries of the structure
# '+1' to add a margin
boundaries = np.max(array.coord, axis=0) - np.min(array.coord, axis=0) + 1
angles = np.deg2rad(angles)
array.box = struc.vectors_from_unitcell(boundaries[0], boundaries[1],
boundaries[2], angles[0],
angles[1], angles[2])
length = array.array_length()
dist_indices = np.stack([
np.repeat(np.arange(length), length),
np.tile(np.arange(length), length)
],
axis=1)
# index_distance() creates a large ndarray
try:
ref_dist = struc.index_distance(array, dist_indices, periodic=True)
except MemoryError:
pytest.skip("Not enough memory")
# Compare with MDTraj
import mdtraj
traj = mdtraj.load(join(data_dir("structure"), "3o5r.pdb"))
# Angstrom to Nanometers
traj.unitcell_vectors = array.box[np.newaxis, :, :] / 10
# Nanometers to Angstrom
mdtraj_dist = mdtraj.compute_distances(traj, dist_indices)[0] * 10
ind = np.where(~np.isclose(ref_dist, mdtraj_dist, atol=1e-5, rtol=1e-3))[0]
assert np.allclose(ref_dist, mdtraj_dist, atol=1e-5, rtol=1e-3)
# Compare with shifted variant
array.coord += shift
array.coord = struc.move_inside_box(array.coord, array.box)
# index_distance() creates a large ndarray
try:
test_dist = struc.index_distance(array, dist_indices, periodic=True)
except MemoryError:
pytest.skip("Not enough memory")
assert np.allclose(test_dist, ref_dist, atol=1e-5)
def test_index_distance_periodic_orthogonal(shift):
"""
The PBC aware computation, should give the same results,
irrespective of which atoms are centered in the box
"""
array = strucio.load_structure(join(data_dir, "3o5r.mmtf"))
# Use a box based on the boundaries of the structure
# '+1' to add a margin
array.box = np.diag(
np.max(array.coord, axis=0) - np.min(array.coord, axis=0) + 1)
length = array.array_length()
dist_indices = np.stack([
np.repeat(np.arange(length), length),
np.tile(np.arange(length), length)
],
axis=1)
ref_dist = struc.index_distance(array, dist_indices, periodic=True)
array.coord += shift
array.coord = struc.move_inside_box(array.coord, array.box)
dist = struc.index_distance(array, dist_indices, periodic=True)
assert np.allclose(dist, ref_dist, atol=1e-5)
def test_adjacency_matrix(cell_size, threshold, periodic, use_selection):
"""
Compare the construction of an adjacency matrix using a cell list
and using a computationally expensive but simpler distance matrix.
"""
array = strucio.load_structure(join(data_dir, "3o5r.mmtf"))
if periodic:
# Create an orthorhombic box
# with the outer coordinates as bounds
array.box = np.diag(
np.max(array.coord, axis=-2) - np.min(array.coord, axis=-2)
)
if use_selection:
np.random.seed(0)
selection = np.random.choice((False, True), array.array_length())
else:
selection = None
cell_list = struc.CellList(
array, cell_size=cell_size, periodic=periodic, selection=selection
)
test_matrix = cell_list.create_adjacency_matrix(threshold)
length = array.array_length()
distance = struc.index_distance(
array,
np.stack(
[
np.repeat(np.arange(length), length),
np.tile(np.arange(length), length)
],
axis=-1
),
periodic
)
distance = np.reshape(distance, (length, length))
# Create adjacency matrix from distance matrix
exp_matrix = (distance <= threshold)
if use_selection:
# Set rows and columns to False for filtered out atoms
exp_matrix[~selection, :] = False
exp_matrix[:, ~selection] = False
# Both ways to create an adjacency matrix
# should give the same result
assert np.array_equal(test_matrix, exp_matrix)
def test_index_distance_non_periodic():
"""
Without PBC the result should be equal to the normal distance
calculation.
"""
array = strucio.load_structure(join(data_dir, "3o5r.mmtf"))
ref_dist = struc.distance(array.coord[np.newaxis, :, :],
array.coord[:, np.newaxis, :]).flatten()
length = array.array_length()
dist = struc.index_distance(array,
indices=np.stack([
np.repeat(np.arange(length), length),
np.tile(np.arange(length), length)
],
axis=1))
assert np.allclose(dist, ref_dist)
def interacting_pairs(structure_path: str,
distance_threshold: float,
atom_name: str = 'CA',
positions: t.Optional[t.Iterable[int]] = None):
"""
Finds residues in structure within distance threshold.
:param structure_path: path to a structure file
:param distance_threshold: min distance between elements (non-inclusive)
:param atom_name: filter atoms to this names (CA, CB, and so on)
:param positions: filter positions to the ones in this list
:return: numpy array with shape (N, 2) where N is a number of interacting pairs
"""
st = io.load_structure(structure_path)
ca = st[(st.atom_name == atom_name) & bst.filter_amino_acids(st)]
if positions is not None:
ca = ca[np.isin(ca.res_id, list(positions))]
pairs = np.array(list(combinations(np.unique(ca.res_id), 2)))
pairs_idx = np.array(list(combinations(np.arange(len(ca)), 2)))
dist = bst.index_distance(ca, pairs_idx)
return pairs[dist < distance_threshold]
