import argparse
import sys
import warnings
import numpy as np
try:
from openbabel import pybel
HAS_OPENBABEL = True
except ModuleNotFoundError:
HAS_OPENBABEL = False
from pysisyphus.calculators import (
Composite,
HardSphere,
PWHardSphere,
TransTorque,
)
from pysisyphus.calculators.OBabel import OBabel
from pysisyphus.drivers.precon_pos_rot import (
center_fragments,
form_A,
get_steps_to_active_atom_mean,
get_which_frag,
rotate_inplace,
SteepestDescent,
)
from pysisyphus.Geometry import Geometry
from pysisyphus.helpers import align_coords, geom_loader
from pysisyphus.io 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)
opt.run()
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)
opt.run()
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)
opt.run()
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)
opt.run()
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()