diff --git a/arc/species/zmat.py b/arc/species/zmat.py index aba7c95a5b..a92b366525 100644 --- a/arc/species/zmat.py +++ b/arc/species/zmat.py @@ -2092,3 +2092,44 @@ def map_index_to_int(index: Union[int, str]) -> int: if isinstance(index, str) and all(char.isdigit() for char in index[1:]): return int(index[1:]) raise TypeError(f'Expected either an int or a string on the format "X15", got {index}') + + +def add_atom_to_zmat(zmat: dict, + element: str, + r_atoms: list, + a_atoms: list, + d_atoms: list, + r: float, + a: float, + d: float, + ) -> dict: + """ + Add a new atom to an existing ZMAT. + + Args: + zmat (dict): The zmat. + element (str): The element of the new atom. + r_atoms (list): The R atom index descriptors. + a_atoms (list): The A atom index descriptors. + d_atoms (list): The D atom index descriptors. + r (float): The R value. + a (float): The A value. + d (float): The D value. + + Returns: + dict: The updated zmat. + """ + zmat = zmat.copy() + symbols = list(zmat['symbols']) + coords = list(zmat['coords']) + vars_ = zmat['vars'] + symbols.append(element) + r_str = f'R_{r_atoms[0]}_{r_atoms[1]}' + a_str = f'A_{a_atoms[0]}_{a_atoms[1]}_{a_atoms[2]}' + d_str = f'D_{d_atoms[0]}_{d_atoms[1]}_{d_atoms[2]}_{d_atoms[3]}' + coords.append((r_str, a_str, d_str)) + vars_[r_str] = r + vars_[a_str] = a + vars_[d_str] = d + zmat = {'symbols': tuple(symbols), 'coords': tuple(coords), 'vars': vars_} + return zmat