def test_translate(input_atoms, ndim, as_list, random_seed):
"""
Translate and translate back and check if the coordinates are still
the same.
"""
if ndim > len(input_atoms.shape):
# Cannot run tests if translation vector has more dimensions
# as input coordinates/atoms
return
np.random.seed(random_seed)
vectors = np.random.rand(*struc.coord(input_atoms).shape[-ndim:])
vectors *= 10
neg_vectors = -vectors
if as_list:
vectors = vectors.tolist()
neg_vectors = neg_vectors.tolist()
translated = struc.translate(input_atoms, vectors)
restored = struc.translate(translated, neg_vectors)
assert type(restored) == type(input_atoms)
assert struc.coord(restored).shape == struc.coord(input_atoms).shape
assert np.allclose(
struc.coord(restored), struc.coord(input_atoms), atol=1e-5
)
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_dihedral():
coord1 = struc.coord([-0.5, -1, 0])
coord2 = struc.coord([0, 0, 0])
coord3 = struc.coord([1, 0, 0])
coord4 = struc.coord([0, 0, -1])
assert struc.dihedral(coord1, coord2, coord3, coord4) \
== pytest.approx(0.5*np.pi)
def test_align_vectors(input_atoms, as_list, use_support, random_seed):
"""
Align, swap alignment vectors and align again.
Expect that the original coordinates have been restored.
"""
np.random.seed(random_seed)
source_direction = np.random.rand(3)
target_direction = np.random.rand(3)
if use_support:
source_position = np.random.rand(3)
target_position = np.random.rand(3)
else:
source_position = None
target_position = None
if as_list:
source_direction = source_direction.tolist()
target_direction = target_direction.tolist()
if use_support:
source_position = source_position.tolist()
target_position = target_position.tolist()
transformed = struc.align_vectors(input_atoms, source_direction,
target_direction, source_position,
target_position)
restored = struc.align_vectors(transformed, target_direction,
source_direction, target_position,
source_position)
assert type(restored) == type(input_atoms)
assert struc.coord(restored).shape == struc.coord(input_atoms).shape
assert np.allclose(struc.coord(restored),
struc.coord(input_atoms),
atol=1e-5)
def test_rotate_about_axis_consistency(input_atoms, axis, random_seed):
"""
Compare the outcome of :func:`rotate_about_axis()` with
:func:`rotate()`.
"""
np.random.seed(random_seed)
angle = np.random.rand() * 2 * np.pi
angles = np.zeros(3)
angles[axis] = angle
ref_rotated = struc.rotate(input_atoms, angles)
rot_axis = np.zeros(3)
# Length of axis should be irrelevant
rot_axis[axis] = np.random.rand()
test_rotated = struc.rotate_about_axis(
input_atoms,
rot_axis,
angle,
)
assert type(test_rotated) == type(ref_rotated)
assert struc.coord(test_rotated).shape == struc.coord(ref_rotated).shape
assert np.allclose(struc.coord(test_rotated),
struc.coord(ref_rotated),
atol=1e-5)
def test_rotate_about_axis_360(input_atoms, random_seed, use_support):
"""
Rotate by 360 degrees around an arbitrary axis and expect that the
coordinates have not changed.
"""
np.random.seed(random_seed)
axis = np.random.rand(3)
support = np.random.rand(3) if use_support else None
rotated = struc.rotate_about_axis(input_atoms, axis, 2*np.pi, support)
assert type(rotated) == type(input_atoms)
assert struc.coord(rotated).shape == struc.coord(input_atoms).shape
assert np.allclose(
struc.coord(rotated), struc.coord(input_atoms), atol=1e-5
)
def test_rotate_360(input_atoms, x, y, z, centered):
"""
Rotate by 360 degrees and expect that the coordinates have not
changed.
"""
func = struc.rotate_centered if centered else struc.rotate
rotated = func(input_atoms, [x, y, z])
assert type(rotated) == type(input_atoms)
assert struc.coord(rotated).shape == struc.coord(input_atoms).shape
assert np.allclose(
struc.coord(rotated), struc.coord(input_atoms), atol=1e-5
)
if centered and struc.coord(input_atoms).ndim > 1:
assert np.allclose(
struc.centroid(rotated), struc.centroid(input_atoms), atol=1e-5
)
def test_rotate(input_atoms, as_list, axis, random_seed, centered):
"""
Rotate and rotate back and check if the coordinates are still
the same.
"""
np.random.seed(random_seed)
angles = np.zeros(3)
angles[axis] = np.random.rand() * 2*np.pi
neg_angles = -angles
if as_list:
angles = angles.tolist()
neg_angles = neg_angles.tolist()
func = struc.rotate_centered if centered else struc.rotate
rotated = func(input_atoms, angles)
restored = func(rotated, neg_angles)
assert type(restored) == type(input_atoms)
assert struc.coord(restored).shape == struc.coord(input_atoms).shape
print(np.max(np.abs(struc.coord(restored) - struc.coord(input_atoms))))
assert np.allclose(
struc.coord(restored), struc.coord(input_atoms), atol=1e-5
)
if centered and struc.coord(input_atoms).ndim > 1:
assert np.allclose(
struc.centroid(restored), struc.centroid(input_atoms), atol=1e-5
)
def test_orient_principal_components(input_atoms, as_list, order):
"""
Orient atoms such that the variance in each axis is greatest
where the value for `order` is smallest.
The number of dimensions in input_atoms must be 2.
"""
if struc.coord(input_atoms).ndim != 2:
with pytest.raises(ValueError):
struc.orient_principal_components(input_atoms)
return
if as_list:
order = order.tolist()
result = struc.orient_principal_components(input_atoms, order=order)
neg_variance = -struc.coord(result).var(axis=0)
assert type(result) == type(input_atoms)
assert (neg_variance.argsort() == np.argsort(order)).all()
def test_rotate_about_axis(input_atoms, as_list, use_support, random_seed):
"""
Rotate and rotate back and check if the coordinates are still
the same.
"""
np.random.seed(random_seed)
axis = np.random.rand(3)
angle = np.random.rand()
neg_angle = -angle
support = np.random.rand(3) if use_support else None
if as_list:
axis = axis.tolist()
support = support.tolist() if support is not None else None
rotated = struc.rotate_about_axis(input_atoms, axis, angle, support)
restored = struc.rotate_about_axis(rotated, axis, neg_angle, support)
assert type(restored) == type(input_atoms)
assert struc.coord(restored).shape == struc.coord(input_atoms).shape
assert np.allclose(
struc.coord(restored), struc.coord(input_atoms), atol=1e-5
)
def test_angle():
coord1 = struc.coord([0, 0, 1])
coord2 = struc.coord([0, 0, 0])
coord3 = struc.coord([0, 1, 1])
assert struc.angle(coord1, coord2, coord3) == pytest.approx(0.25 * np.pi)
def test_centroid():
coord = struc.coord([[1, 1, 1], [0, -1, -1], [-1, 0, 0]])
assert struc.centroid(coord).tolist() == [0, 0, 0]
def test_distance():
coord1 = struc.coord([0, 1, 1])
coord2 = struc.coord([0, 2, 2])
assert struc.distance(coord1, coord2) == pytest.approx(np.sqrt(2))
