Source code for pysisyphus.intcoords.setup_fast

import itertools as it
import logging

import numpy as np
from sklearn.neighbors import KDTree

from pysisyphus.elem_data import COVALENT_RADII as CR
from pysisyphus.helpers_pure import log, timed
from pysisyphus.intcoords.valid import bend_valid, dihedral_valid

logger = logging.getLogger("internal_coords")


[docs] def get_max_bond_dists(atoms, bond_factor, covalent_radii=None): if covalent_radii is None: cr = CR else: cr = {atom: covrad for atom, covrad in zip(atoms, covalent_radii)} unique_atoms = set(atoms) atom_pairs = it.combinations_with_replacement(unique_atoms, 2) max_bond_dists = { frozenset((from_, to_)): bond_factor * (cr[from_] + cr[to_]) for from_, to_ in atom_pairs } return max_bond_dists
[docs] def find_bonds( atoms, coords3d, covalent_radii=None, bond_factor=BOND_FACTOR, min_dist=0.1 ): atoms = [atom.lower() for atom in atoms] c3d = coords3d.reshape(-1, 3) if covalent_radii is None: covalent_radii = [CR[atom] for atom in atoms] cr = np.array(covalent_radii) max_bond_dists = get_max_bond_dists(atoms, bond_factor, covalent_radii=cr) radii = bond_factor * (cr.copy() + max(cr)) kdt = KDTree(c3d) res, dists = kdt.query_radius(c3d, radii, return_distance=True) bonds_ = list() for i, (atom, bonds, dists_) in enumerate(zip(atoms, res, dists)): fsets = [frozenset((atom, atoms[a])) for a in bonds] ref_dists = np.array([max_bond_dists[fset] for fset in fsets]) keep = np.logical_and(dists_ <= ref_dists, min_dist <= dists_) for to_ in bonds[keep]: bonds_.append(frozenset((i, to_))) bonds = np.array([tuple((from_, to_)) for from_, to_ in set(bonds_)]) return bonds
[docs] def find_bonds_for_geom(geom, bond_factor=BOND_FACTOR): return find_bonds(geom.atoms, geom.coords3d, geom.covalent_radii)
[docs] def get_bond_vec_getter( atoms, covalent_radii, bonds_for_inds, no_bonds_with=None, bond_factor=BOND_FACTOR ): max_bond_dists = get_max_bond_dists( atoms, bond_factor, covalent_radii=covalent_radii ) max_bond_dists_for_inds = [ np.array([max_bond_dists[frozenset((atoms[ind], atom_))] for atom_ in atoms]) for ind in bonds_for_inds ] if no_bonds_with is None: # List of empty lists no_bonds_with = [[] * len(bonds_for_inds)] # Set it to a negative value, so the calculated distance, which is always positive # can't be smaller than this value. for mbd, nbw in zip(max_bond_dists_for_inds, no_bonds_with): mbd[nbw] = -1 all_inds = np.arange(len(atoms)) def get_bond_vecs(coords, return_bonded_inds=False): coords3d = coords.reshape(-1, 3) all_bond_vecs = list() all_bonded_inds = list() for ind, max_dists in zip(bonds_for_inds, max_bond_dists_for_inds): distance_vecs = coords3d - coords3d[ind] distances = np.linalg.norm(distance_vecs, axis=1) # Set 0.0 distance of atom with itself to a high value to not form # and ind-ind bond. distances[ind] = 10_000 bond_mask = distances <= max_dists all_bond_vecs.append(distance_vecs[bond_mask]) all_bonded_inds.append(all_inds[bond_mask]) if return_bonded_inds: return all_bond_vecs, all_bonded_inds else: return all_bond_vecs return get_bond_vecs
[docs] def get_bend_candidates(bonds): """Also yields duplicates [a, b, c] and [c, b, a].""" bond_dict = {} bonds = [tuple(bond) for bond in bonds] # Construct a dictionary holding neighbours for a given atom. for from_, to_ in bonds: bond_dict.setdefault(from_, list()).append(to_) bond_dict.setdefault(to_, list()).append(from_) for bond in bonds: from_, to_ = bond # Look up neighbours of from_ and to_ in the dictionary from_neighs = set(bond_dict[from_]) - set((to_,)) to_neighs = set(bond_dict[to_]) - set((from_,)) bend_candidates = [(neigh,) + bond for neigh in from_neighs] + [ bond + (neigh,) for neigh in to_neighs ] yield from bend_candidates
[docs] def find_bends(coords3d, bonds, min_deg, max_deg, logger=None): bend_set = set() for indices in get_bend_candidates(bonds): if ( (not bend_valid(coords3d, indices, min_deg, max_deg)) or (indices in bend_set) or (indices[::-1] in bend_set) ): continue bend_set.add(indices) return [list(bend) for bend in bend_set]
[docs] def find_dihedrals(coords3d, bonds, bends, max_deg, logger=None): bond_dict = {} bonds = [tuple(bond) for bond in bonds] for from_, to_ in bonds: bond_dict.setdefault(from_, list()).append(to_) bond_dict.setdefault(to_, list()).append(from_) dihedral_set = set() for bend in bends: bend = tuple(bend) from_, central, to_ = bend from_neighs = set(bond_dict[from_]) - set((to_, central)) to_neighs = set(bond_dict[to_]) - set((from_, central)) dihedral_candidates = [(neigh,) + bend for neigh in from_neighs] + [ bend + (neigh,) for neigh in to_neighs ] for indices in dihedral_candidates: if not dihedral_valid(coords3d, indices, deg_thresh=max_deg): continue dihedral_set.add(indices) return [list(dihedral) for dihedral in dihedral_set]
[docs] def find_bonds_bends(geom, bond_factor=BOND_FACTOR, min_deg=15, max_deg=175): log(logger, "Starting detection of bonds and bends.") bonds = find_bonds_for_geom(geom, bond_factor=bond_factor) log(logger, f"Found {len(bonds)} bonds.") bends = find_bends(geom.coords3d, bonds, min_deg=min_deg, max_deg=max_deg) log(logger, f"Found {len(bends)} bends.") return bonds, bends
@timed(logger) def find_bonds_bends_dihedrals(geom, bond_factor=BOND_FACTOR, min_deg=15, max_deg=175): log(logger, f"Detecting bonds, bends and dihedrals for {len(geom.atoms)} atoms.") bonds, bends = find_bonds_bends( geom, bond_factor=bond_factor, min_deg=min_deg, max_deg=max_deg ) proper_dihedrals = find_dihedrals(geom.coords3d, bonds, bends, max_deg) log(logger, f"Found {len(proper_dihedrals)} proper dihedrals.") return bonds, bends, proper_dihedrals