def interpret_sel(s, sel):
if sel is None:
return AtomSelection.all(s)
elif is_string(sel):
return AtomSelection.from_element(s, sel)
elif hasattr(sel, '__call__'):
return sel(s)
def test_arrays(self):
a = Atoms('HCHC',
positions=[[i] * 3 for i in range(4)],
cell=[4] * 3,
pbc=[True] * 3)
s = AtomSelection.from_element(a, 'C')
s.set_array('testarr', [1, 2])
self.assertTrue(all(s.subset(a).get_array('testarr') == [1, 2]))
# Test that arrays are reordered
a.set_array('testarr', np.array([1, 2, 3, 4]))
s = AtomSelection(a, [2, 0])
a2 = s.subset(a)
self.assertTrue((a2.get_array('testarr') == np.array([3, 1])).all())
# Cell indices test!
s = AtomSelection(a, [0, 3])
s.set_array('cell_indices', [[0, 0, 0], [-1, 0, 0]])
a2 = s.subset(a, True)
self.assertTrue(np.allclose(a2.get_positions()[-1], [-1, 3, 3]))
def test_selectors(self):
# Multiple tests for various methods
a = Atoms('HCHC',
positions=[[i] * 3 for i in range(4)],
cell=[4] * 3,
pbc=[True] * 3)
# Element test
s1 = AtomSelection.from_element(a, 'C')
self.assertTrue(set(s1.indices) == set([1, 3]))
# Box test
s1 = AtomSelection.from_box(a, [1.5] * 3, [4.5] * 3, periodic=True)
s2 = AtomSelection.from_box(a, [1.5] * 3, [4.5] * 3, periodic=False)
s3 = AtomSelection.from_box(a, [0.375] * 3, [1.125] * 3,
periodic=True,
scaled=True)
self.assertTrue(set(s1.indices) == set([0, 2, 3]))
self.assertTrue(set(s2.indices) == set([2, 3]))
self.assertTrue(set(s3.indices) == set([0, 2, 3]))
# Sphere test
s1 = AtomSelection.from_sphere(a, [0.5] * 3, 3, periodic=True)
s2 = AtomSelection.from_sphere(a, [0.5] * 3, 3, periodic=False)
self.assertTrue(set(s1.indices) == set([0, 1, 2, 3]))
self.assertTrue(set(s2.indices) == set([0, 1, 2]))
def test_transformprops(self):
from ase.quaternions import Quaternion
from soprano.selection import AtomSelection
from soprano.properties.transform import (Translate, Rotate, Mirror)
a = Atoms('CH', positions=[[0, 0, 0], [0.5, 0, 0]])
sel = AtomSelection.from_element(a, 'C')
transl = Translate(selection=sel, vector=[0.5, 0, 0])
rot = Rotate(selection=sel, center=[0.25, 0.0, 0.25],
quaternion=Quaternion([np.cos(np.pi/4.0),
0,
np.sin(np.pi/4.0),
0]))
mirr = Mirror(selection=sel, plane=[1, 0, 0, -0.25])
aT = transl(a)
aR = rot(a)
aM = mirr(a)
self.assertAlmostEqual(np.linalg.norm(aT.get_positions()[0]),
np.linalg.norm(aT.get_positions()[1]))
self.assertAlmostEqual(np.linalg.norm(aR.get_positions()[0]),
np.linalg.norm(aR.get_positions()[1]))
self.assertAlmostEqual(np.linalg.norm(aM.get_positions()[0]),
np.linalg.norm(aM.get_positions()[1]))
def test_arrays(self):
a = Atoms('HCHC', positions=[[i]*3 for i in range(4)],
cell=[4]*3, pbc=[True]*3)
s = AtomSelection.from_element(a, 'C')
s.set_array('testarr', [1, 2])
self.assertTrue(all(s.subset(a).get_array('testarr') == [1, 2]))
Besides allowing to manipulate information about multiple structures, Soprano provides tools to edit them as well.
This is accomplished by combining selection of atoms and transformation operations that change their positions.
As an example we will use again the ammonia molecule.
Selections can be carried with multiple criteria. The basic ones are selection by element, selection of all atoms
in a box, and selection of all atoms in a sphere.
"""
from soprano.selection import AtomSelection
nh3coords = np.array([[2.5, 2.5, 2.5], [3.4373, 2.5, 2.1193],
[2.0314, 3.3117, 2.1193], [2.0314, 1.6883, 2.1193]])
nh3l = Atoms('NHHH', nh3coords, cell=[5, 5,
5]) # The cell is just an empty box
# Now instead of switching the coordinates by hand let's do this with selections.
nh3Hsel = AtomSelection.from_element(nh3l, 'H') # All H atoms in nh3l
# Selections can be manipulated in interesting ways. To begin with, we can create an Atoms object containing
# only the selected atoms
h3 = nh3Hsel.subset(nh3l)
print "---- Selected atoms contained in nh3Hsel ----\n"
print h3.get_chemical_symbols(), "\n\n"
# Also, selections can be summed, subtracted, or multiplied (representing intersection)
sel1 = AtomSelection(nh3l, [1]) # A custom generated selection
sel2 = AtomSelection(nh3l, [0, 2]) # A custom generated selection
print "---- Indices of selected atoms for various combinations ----\n"
print "sel1:\t", sel1.indices