def atom_indices(zma, symb, match=True):
""" Obtain the indices of a atoms of a particular type in the geometry.
:param zma: Z-Matrix
:type zma: automol Z-Matrix data structure
:param match: grab idxs that match given atom type
:param symb: atomic symbol
:type symb: str
:param match: obtain indices of symbols that match the type?
:type match: bool
"""
return vmat.atom_indices(zma, symb, match=match)
def shift_up(zma, idxs):
""" shift up from the dummy idxs
"""
dummy_idxs = sorted(atom_indices(zma, 'X', match=True))
shift_idxs = []
for idx in idxs:
new_idx = idx
for dummy_idx in dummy_idxs:
if idx >= dummy_idx:
new_idx += 1
shift_idxs.append(new_idx)
return tuple(shift_idxs)
def shift_down(zma, vals):
"""
Build a remdummy list that tells how to shift the groups
"""
dummy_idxs = sorted(atom_indices(zma, 'X', match=True))
if dummy_idxs:
remdummy = [0 for _ in range(count(zma))]
for dummy in dummy_idxs:
for idx, _ in enumerate(remdummy):
if dummy < idx:
remdummy[idx] += 1
vals1 = tuple(val + 1 for val in vals)
vals2 = tuple(val - remdummy[val - 1] for val in vals1)
shift_vals = tuple(val - 1 for val in vals2)
else:
shift_vals = vals
return shift_vals
def test__atom_indices():
""" test vmat.atom_indices
"""
angle_idxs = vmat.atom_indices(C5H8O_VMA, 'C', match=False)
assert angle_idxs == (2, 3, 4, 7, 8, 10, 11, 12, 13, 14, 15)