def testSaveRead(self):
# Setup
atoms = Atoms('4Ge', positions=np.array([[0, 1*i, -2*i] for i in range(4)]))
atoms2 = atoms.copy()
atoms2.setTag('hello', [1,2])
atoms2.setRegion(
'second',
picker=[1,2],
cell=np.array([[2.,0,0], [0, 1, 0], [0,0,3]]),
pbc=[True, False, True],
voltage=-3.0,
)
w_traj = Trajectory(atoms=atoms2, mode='w', filename='test.traj')
w_traj.write()
w_traj.close()
r_traj = Trajectory(mode='r', filename='test.traj')
atoms3 = r_traj[-1]
self.assertEqual(atoms3.tagStrings(), ['hello', 'second'])
self.assertIsInstance(atoms3.tag('hello'), IndexTag)
self.assertArraysEqual(atoms3.tag('hello').picker(), [1,2])
self.assertIsInstance(atoms3.tag('second'), Region)
self.assertArraysEqual(atoms3.tag('second').cell(), [[2.0, 0,0], [0, 1, 0], [0, 0 , 3]])
self.assertEqual(atoms3.tag('second').pbc(), (True, False, True))
self.assertEqual(atoms3.tag('second').voltage(), -3.0)
def testTagAdd(self):
atoms = Atoms('4Ge', positions=np.array([[0, 1*i, -2*i] for i in range(4)]))
atoms2 = atoms.copy()
atoms2.setTag('hello', [1,2])
self.assertArraysEqual(atoms2.tag('hello').picker(), [1,2])
atoms = atoms + atoms2
self.assertArraysEqual(atoms.tag('hello').picker(), [5,6])
def testTagCut(self):
n = 6
atoms = Atoms('%dGe'%n, positions=np.array([[0, 1*i, -2*i] for i in range(n)]))
atoms.setTag('hello', [2, 3, 4])
atoms2 = atoms[[1, 2, 3]]
atoms3 = atoms[[0]]
atoms4 = atoms[[4, 5]]
atoms5 = atoms[-3:]
atoms6 = atoms[:3]
atoms7 = atoms[:]
atoms8 = atoms[[]]
self.assertArraysEqual(atoms.tag('hello').picker(), [2, 3, 4])
self.assertArraysEqual(atoms7.tag('hello').picker(), [2, 3, 4])
self.assertArraysEqual(atoms8.tag('hello').picker(), [])
self.assertArraysEqual(atoms2.tag('hello').picker(), [1, 2])
self.assertArraysEqual(atoms3.tag('hello').picker(), [])
self.assertArraysEqual(atoms4.tag('hello').picker(), [0])
self.assertArraysEqual(atoms5.tag('hello').picker(), [0, 1])
self.assertArraysEqual(atoms6.tag('hello').picker(), [2])
