def _extend_chain_to_include_anchoring_atoms(gra, keys, zma_keys): """ extend chain to include three atoms already specified in v-matrix :param gra: the graph :param keys: keys in the chain; the first atom should already be specified :param zma_keys: keys currently in the v-matrix """ ngb_keys_dct = atoms_sorted_neighbor_atom_keys( gra, symbs_first=('X', 'C',), symbs_last=('H',), ords_last=(0.1,)) symb_dct = atom_symbols(gra) key3 = keys[0] assert key3 in zma_keys key2 = next(k for k in ngb_keys_dct[key3] if k in zma_keys) if symb_dct[key2] == 'X': key1 = next(k for k in ngb_keys_dct[key3][1:] if k in zma_keys) else: key1 = next(k for k in ngb_keys_dct[key2] if k in zma_keys and k != key3) keys = (key1, key2,) + tuple(keys) return keys
def start_at(gra, key): """ start a v-matrix at a specific atom Returns the started vmatrix, along with keys to atoms whose neighbors are missing from it """ symb_dct = atom_symbols(gra) ngb_keys_dct = atoms_sorted_neighbor_atom_keys(gra, symbs_first=( 'X', 'C', ), symbs_last=('H', ), ords_last=(0.1, )) ngb_keys = ngb_keys_dct[key] if not ngb_keys: zma_keys = [] elif len(ngb_keys) == 1: # Need special handling for atoms with only one neighbor if symb_dct[key] in ('H', 'X'): key2 = ngb_keys[0] zma_keys = (key2, ) + ngb_keys_dct[key2] else: key2 = ngb_keys[0] ngb_keys = tuple(k for k in ngb_keys_dct[key2] if k != key) zma_keys = (key, key2) + ngb_keys else: zma_keys = (key, ) + ngb_keys_dct[key] vma = () for row, key_ in enumerate(zma_keys): idx1 = idx2 = idx3 = None if row > 0: key1 = next(k for k in ngb_keys_dct[key_] if k in zma_keys[:row]) idx1 = zma_keys.index(key1) if row > 1: key2 = next(k for k in ngb_keys_dct[key1] if k in zma_keys[:row] and k != key_) idx2 = zma_keys.index(key2) if row > 2: key3 = next(k for k in zma_keys[:row] if k not in (key_, key1, key2)) idx3 = zma_keys.index(key3) sym = symb_dct[key_] key_row = [idx1, idx2, idx3] vma = automol.vmat.add_atom(vma, sym, key_row) return vma, zma_keys
def complete_branch(gra, key, vma, zma_keys, branch_keys=None): """ continue constructing a v-matrix along a chain All neighboring atoms along the chain will be included Exactly one atom in the chain must already be in the v-matrix :param gra: the graph for which the v-matrix will be constructed :param keys: the keys for atoms along the chain, which must be contiguous; the first atom must already appear in the v-matrix :param vma: a partial v-matrix from which to continue :param zma_keys: row keys for the partial v-matrix, identifying the atom specified by each row of `vma` in order :param branch_keys: optionally, restrict the v-matrix to these keys and their neighbors; if `None`, the entire branch will be included """ branch_keys = atom_keys(gra) if branch_keys is None else branch_keys keys = _extend_chain_to_include_anchoring_atoms(gra, [key], zma_keys) zma_keys = list(zma_keys) symb_dct = atom_symbols(gra) ngb_keys_dct = atoms_sorted_neighbor_atom_keys( gra, symbs_first=('X', 'C',), symbs_last=('H',), ords_last=(0.1,), prioritize_keys=branch_keys) # If this z-matrix is being continued from a partial z-matrix, the leading # atom for a torsion may have already be defined. To handle this case, I # make a dictionary of these leading atoms and use them below where needed. lead_key_dct = {} for idx, key_row in enumerate(automol.vmat.key_matrix(vma)): axis = key_row[:2] if None not in axis: axis = tuple(map(zma_keys.__getitem__, axis)) if axis not in lead_key_dct: lead_key_dct[axis] = zma_keys[idx] def _continue(key1, key2, key3, vma, zma_keys): k3ns = list(ngb_keys_dct[key3]) for k3n in set(k3ns) & set(zma_keys): k3ns.remove(k3n) if k3ns: key4 = k3ns.pop(0) lead_key = None if (key3, key2) in lead_key_dct: lead_key = lead_key_dct[(key3, key2)] # Add the leading atom to the v-matrix symb = symb_dct[key4] dkey = key1 if lead_key is None else lead_key key_row = list(map(zma_keys.index, (key3, key2, dkey))) vma = automol.vmat.add_atom(vma, symb, key_row) assert key4 not in zma_keys, ("Atom {:d} already in v-matrix." .format(key4)) zma_keys.append(key4) dkey = key4 if lead_key is None else lead_key # Add the neighbors of atom 3 (if any) to the v-matrix, decoupled # from atom 1 for properly decopuled torsions for k3n in k3ns: sym = symb_dct[k3n] if symb_dct[dkey] == 'X': key_row = list(map(zma_keys.index, (key3, dkey, key2))) else: key_row = list(map(zma_keys.index, (key3, key2, dkey))) vma = automol.vmat.add_atom(vma, sym, key_row) assert k3n not in zma_keys, ("Atom {:d} already in v-matrix." .format(k3n)) zma_keys.append(k3n) # Recursion if key4 in branch_keys: vma, zma_keys = _continue(key2, key3, key4, vma, zma_keys) if symb_dct[key4] == 'X': key2 = key4 for k3n in k3ns: if k3n in branch_keys: vma, zma_keys = _continue(key2, key3, k3n, vma, zma_keys) return vma, zma_keys key1, key2, key3 = keys[:3] vma, zma_keys = _continue(key1, key2, key3, vma, zma_keys) return vma, zma_keys
def complete_branch(gra, key, vma, zma_keys, branch_keys=None): """ continue constructing a v-matrix along a chain All neighboring atoms along the chain will be included Exactly one atom in the chain must already be in the v-matrix :param gra: the graph for which the v-matrix will be constructed :param vma: a partial v-matrix from which to continue :param zma_keys: row keys for the partial v-matrix, identifying the atom specified by each row of `vma` in order :param branch_keys: optionally, restrict the v-matrix to these keys and their neighbors; if `None`, the entire branch will be included """ branch_keys = atom_keys(gra) if branch_keys is None else branch_keys keys = _extend_chain_to_include_anchoring_atoms(gra, [key], zma_keys) zma_keys = list(zma_keys) symb_dct = atom_symbols(gra) ngb_keys_dct = atoms_sorted_neighbor_atom_keys(gra, symbs_first=( 'X', 'C', ), symbs_last=('H', ), ords_last=(0.1, )) def _continue(key1, key2, key3, vma, zma_keys): k3ns = list(ngb_keys_dct[key3]) for k3n in set(k3ns) & set(zma_keys): k3ns.remove(k3n) if k3ns: key4 = k3ns.pop(0) # Add the leading atom to the v-matrix symb = symb_dct[key4] key_row = list(map(zma_keys.index, (key3, key2, key1))) vma = automol.vmat.add_atom(vma, symb, key_row) assert key4 not in zma_keys, f"Atom {key4:d} already in v-matrix" zma_keys.append(key4) # Add the neighbors of atom 3 (if any) to the v-matrix, decoupled # from atom 1 for properly decopuled torsions for k3n in k3ns: sym = symb_dct[k3n] if symb_dct[key4] == 'X': key_row = list(map(zma_keys.index, (key3, key4, key2))) else: key_row = list(map(zma_keys.index, (key3, key2, key4))) vma = automol.vmat.add_atom(vma, sym, key_row) assert k3n not in zma_keys, f"Atom {k3n:d} already in v-matrix" zma_keys.append(k3n) # Recursion if key4 in branch_keys: vma, zma_keys = _continue(key2, key3, key4, vma, zma_keys) if symb_dct[key4] == 'X': key2 = key4 for k3n in k3ns: if k3n in branch_keys: vma, zma_keys = _continue(key2, key3, k3n, vma, zma_keys) return vma, zma_keys key1, key2, key3 = keys[0], keys[1], keys[2] vma, zma_keys = _continue(key1, key2, key3, vma, zma_keys) return vma, zma_keys