Source code for pysisyphus.drivers.merge

import argparse
import sys
import warnings

import numpy as np

    from openbabel import pybel

except ModuleNotFoundError:

from pysisyphus.calculators import (
from pysisyphus.calculators.OBabel import OBabel
from pysisyphus.drivers.precon_pos_rot import (
from pysisyphus.Geometry import Geometry
from pysisyphus.helpers import align_coords, geom_loader
from import geom_to_crd_str
from pysisyphus.linalg import get_rot_mat_for_coords
from pysisyphus.optimizers.LBFGS import LBFGS
from pysisyphus.xyzloader import coords_to_trj

[docs] def merge_geoms(geom1, geom2, geom1_del=None, geom2_del=None, make_bonds=None): """Merge geom1 and geom2 while keeping the original coordinates. Supports deleting certain atoms. """ if geom1_del is None: geom1_del = list() if geom2_del is None: geom2_del = list() geom1_del = np.array(geom1_del) geom2_del = np.array(geom2_del) # Update indices of atoms to be deleted in geom2 by the number of atoms in geom1, # as geom1 and geom2 will be merged. # If we want to delete atom 0 and atom 1 in geom2 and geom1 comprises 10 atoms, # then the updated indices in geom2 will be 0 + 10 = 10 and 1 + 10 = 11 in the # merged geometry. atom_num1 = len(geom1.atoms) geom2_del += atom_num1 ndel1 = len(geom1_del) if make_bonds is not None: make_bonds = np.array(make_bonds, dtype=int) geom1_bond_inds, geom2_bond_inds = make_bonds.T geom2_bond_inds += atom_num1 # Correct original bond indices given in make_bonds # # Determe how many atoms to be deleted lie below the atom(s) that are # used to form new bonds. corr1 = (geom1_del < geom1_bond_inds[:, None]).sum(axis=1) geom1_bond_inds -= corr1 # Correct bond indices in geom2 by the number of deleted atoms in geom1. geom2_bond_inds -= ndel1 corr2 = (geom2_del < geom2_bond_inds[:, None]).sum(axis=1) geom2_bond_inds -= corr2 make_bonds_cor = np.stack((geom1_bond_inds, geom2_bond_inds), axis=1) else: make_bonds_cor = None union = geom1 + geom2 union = union.del_atoms(list(geom1_del) + list(geom2_del)) # Set appropriate fragments atom_num1 -= len(geom1_del) frag1 = np.arange(atom_num1) lig_atoms = geom2.atoms frag2 = np.arange(atom_num1, atom_num1 + len(lig_atoms) - len(geom2_del)) union.fragments = {"geom1": frag1.tolist(), "geom2": frag2.tolist()} return union, make_bonds_cor
[docs] def hardsphere_merge(geom1, geom2): union = geom1 + geom2 atom_num = len(geom1.atoms) # Set up lists containing the atom indices for the two fragments frag_lists = [ [i for i, _ in enumerate(geom1.atoms)], [atom_num + i for i, _ in enumerate(geom2.atoms)], ] def get_hs(kappa=1.0): return HardSphere( union, frag_lists, permutations=True, kappa=kappa, ) union.set_calculator(get_hs(1.0)) opt_kwargs = { "max_cycles": 1000, "max_step": 0.5, "rms_force": 0.0005, } opt = SteepestDescent(union, **opt_kwargs) return union
[docs] def prepare_merge(geom1, bond_diff, geom2=None, del1=None, del2=None, dump=False): bond_diff = np.array(bond_diff, dtype=int) if del1 is not None: geom1 = geom1.del_atoms(del1) if geom2: if del2 is not None: geom2 = geom2.del_atoms(del2) union = geom1 + geom2 atom_num = len(geom1.atoms) # Set up lists containing the atom indices for the two fragments frag_lists = [ [i for i, _ in enumerate(geom1.atoms)], [atom_num + i for i, _ in enumerate(geom2.atoms)], ] fragments = {"geom1": frag_lists[0], "geom2": frag_lists[1]} union.fragments = fragments else: union = geom1 frag_lists = [union.fragments["geom1"], union.fragments["geom2"]] backup = list() def keep(comment): backup.append((union.cart_coords.copy(), comment)) keep("Initial union") which_frag = get_which_frag(frag_lists) AR = form_A(frag_lists, which_frag, bond_diff) # Center fragments at their geometric average center_fragments(frag_lists, union) keep("Centered fragments") # Translate centroid of active atoms to origin alphas = get_steps_to_active_atom_mean(frag_lists, frag_lists, AR, union.coords3d) for frag, alpha in zip(frag_lists, alphas): union.coords3d[frag] += alpha keep("Shifted centroids to active atoms") def get_hs(kappa=1.0, **kwargs): return HardSphere( union, frag_lists, permutations=True, kappa=kappa, **kwargs, ) union.set_calculator(get_hs(1.0)) opt_kwargs = { "max_cycles": 1000, "max_step": 0.5, "rms_force": 0.0005, } opt = SteepestDescent(union, **opt_kwargs) keep("Initial hardsphere optimization") # Corresponds to A.3 S_3 initial orientation of molecules in Habershon paper rotate_inplace(frag_lists, union, bond_diff) keep("Rotated after inital hardsphere optimization") sub_frags = [bond_diff[:, 0].tolist(), bond_diff[:, 1].tolist()] keys_calcs = { "hs": get_hs(100.0, radii_offset=3.0), "tt": TransTorque(frag_lists, frag_lists, AR, AR, kappa=2, do_trans=True), "pwhs": PWHardSphere(union, frag_lists, sub_frags=sub_frags, kappa=10.0), } comp = Composite("hs + tt + pwhs", keys_calcs, remove_translation=True) union.set_calculator(comp) max_cycles = 15_000 max_dist = 50 max_step = max_dist / max_cycles opt_kwargs = { "max_cycles": max_cycles, "max_step": max_step, "rms_force": 0.1, # "dump": True, } opt = SteepestDescent(union, **opt_kwargs) keep("Hardsphere + TransTorque optimization") if dump: kept_coords, kept_comments = zip(*backup) align_coords(kept_coords) fn = "merge_dump.trj" coords_to_trj(fn, union.atoms, kept_coords, comments=kept_comments) return union
[docs] def merge_opt(union, bond_diff, ff="mmff94"): """Fragment merging along given bond by forcefield optimization.""" geom1 = union.get_fragments("geom1") freeze = list(range(len(geom1.atoms))) # Create pybel.Molecule/OBMol to set the missing bonds mol = pybel.readstring("xyz", union.as_xyz()) obmol = mol.OBMol for from_, to_ in bond_diff: obmol.AddBond(int(from_ + 1), int(to_ + 1), 1) # Use modified pybel.Molecule calc = OBabel(mol=mol, ff=ff) funion = union.copy() # Only releax second fragment funion.freeze_atoms = freeze funion.set_calculator(calc) opt = LBFGS(funion, max_cycles=1000, max_step=0.5, dump=False, print_every=25) return funion
[docs] def align_on_subset(geom1, union, del1=None): """Align 'union' onto subset of 'geom1'""" # Delete some coordinates (atoms) in geom1, that are not present in union coords3d_1 = geom1.coords3d.copy() atoms1 = geom1.atoms if del1 is not None: atoms1 = [atom for i, atom in enumerate(atoms1) if i not in del1] coords3d_1 = np.delete(coords3d_1, del1, axis=0) atoms1 = tuple(atoms1) num1 = coords3d_1.shape[0] # Restrict length of union to length of coords3d_1 coords3d_2 = union.coords3d.copy() coords3d_2_subset = coords3d_2[:num1] atoms2 = union.atoms assert atoms1 == tuple(atoms2[:num1]) *_, rot_mat = get_rot_mat_for_coords(coords3d_1, coords3d_2_subset) # Align merged system coords3d_2_aligned = (coords3d_2 - coords3d_2.mean(axis=0)[None, :]).dot(rot_mat) # Translate aligned system so that centroids of subsets match centroid_1 = coords3d_1.mean(axis=0) centroid_2 = coords3d_2_aligned[:num1].mean(axis=0) coords3d_2_aligned += centroid_1 - centroid_2 aligned = Geometry(atoms2, coords3d_2_aligned) subset = Geometry(atoms2[:num1], coords3d_2_aligned[:num1]) rest = Geometry(atoms2[num1:], coords3d_2_aligned[num1:]) return aligned, subset, rest
[docs] def merge_with_frozen_geom( frozen_geom, lig_geom, make_bonds, frozen_del, lig_del, ff="mmff94" ): union, make_bonds_cor = merge_geoms( frozen_geom, lig_geom, frozen_del, lig_del, make_bonds ) atoms = union.atoms print("Docking to form bonds:") for i, (from_, to_) in enumerate(make_bonds_cor): from_atom = atoms[from_] to_atom = atoms[to_] print(f"\t{i:02d} {from_atom}{from_}-{to_atom}{to_}") union = prepare_merge(union, make_bonds_cor, dump=True) if HAS_OPENBABEL: opt_union = merge_opt(union, make_bonds_cor, ff=ff) else: warnings.warn(f"Could not import openbabel. Skipping '{ff}' optimization.") opt_union = union # aligned: whole system # subset: frozen_geom - deleted atoms # rest: aligned ligand aligned, subset, rest = align_on_subset(frozen_geom, opt_union, frozen_del) return aligned, subset, rest
[docs] def parse_args(args): parser = argparse.ArgumentParser() parser.add_argument("frozen_fn", help="Filename of the frozen geometry.") parser.add_argument("lig_fn", help="Filename of the ligand to be merged.") parser.add_argument( "--frozen-del", nargs="+", type=int, help="0-based atom indices to be deleted from frozen_fn.", ) parser.add_argument( "--lig-del", nargs="+", type=int, help="0-based atom indices to be deleted from lig_fn.", ) parser.add_argument( "--make-bonds", nargs="+", type=int, help="0-based indices of atom pairs (frozen, lig), forming bonds " "in the merged geometry. lig indices should start at 0.", ) parser.add_argument("--res", default="LIG") parser.add_argument("--resno", type=int) return parser.parse_args(args)
[docs] def run_merge(): args = parse_args(sys.argv[1:]) frozen_fn = args.frozen_fn lig_fn = args.lig_fn protein_pdb = frozen_fn lig_pdb = lig_fn frozen_geom = geom_loader(protein_pdb) lig_geom = geom_loader(lig_pdb) prot_del = args.frozen_del lig_del = args.lig_del make_bonds = args.make_bonds make_bonds = np.array(make_bonds, dtype=int).reshape(-1, 2) aligned, subset, rest = merge_with_frozen_geom( frozen_geom, lig_geom, make_bonds, prot_del, lig_del ) trj_fn = "merged.trj" trj = "\n".join([geom.as_xyz() for geom in (aligned, rest)]) with open(trj_fn, "w") as handle: handle.write(trj) print(f"Dumped geometries to '{trj_fn}'.") res = args.res resno = args.resno if (res is not None) and (resno is not None): crd_str = geom_to_crd_str( rest, res=res, resno=resno, ref_atoms=lig_geom.atoms, del_atoms=lig_del ) crd_fn = "lig_aligned.crd" with open(crd_fn, "w") as handle: handle.write(crd_str) print(f"Dumped ligand coordinates to to '{crd_fn}'.")
if __name__ == "__main__": run_merge()