def test_atoms_set():
a1 = Atom(1)
a2 = Atom(2)
a3 = Atom(3)
a = Atoms(a1, 3)
assert len(a) == 3
assert len(a.atom) == 1
a[1] = a2
assert len(a) == 3
assert len(a.atom) == 2
# Add the atom, but do not specify any
# atoms to have the species
a[[]] = a3
assert len(a) == 3
assert len(a.atom) == 3
for atom in a:
assert atom in [a1, a2]
assert atom != a3
found = 0
for atom, i_s in a.iter(True):
if atom == a1:
assert len(i_s) == 2
found += 1
elif atom == a2:
assert len(i_s) == 1
found += 1
elif atom == a3:
assert len(i_s) == 0
found += 1
assert found == 3
def test_iter1():
# Add new atoms to the set
atom = Atoms(['C', 'C'])
for a in atom.iter():
assert a == Atom[6]
for a, idx in atom.iter(True):
assert a == Atom[6]
assert len(idx) == 2
atom = Atoms(['C', 'Au', 'C', 'Au'])
for i, aidx in enumerate(atom.iter(True)):
a, idx = aidx
if i == 0:
assert a == Atom[6]
assert (idx == [0, 2]).all()
elif i == 1:
assert a == Atom['Au']
assert (idx == [1, 3]).all()
assert len(idx) == 2
