mlmm-toolkit 0.2.2.dev0__py3-none-any.whl

pysisyphus/drivers/merge.py

@@ -0,0 +1,368 @@
+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
+
+
+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
+
+
+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
+
+
+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.xyz"
+        coords_to_trj(fn, union.atoms, kept_coords, comments=kept_comments)
+
+    return union
+
+
+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
+
+
+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
+
+
+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
+
+
+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)
+
+
+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.xyz"
+    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()

pysisyphus/drivers/merge_mol2.py

@@ -0,0 +1,322 @@
+import argparse
+import copy
+import sys
+from typing import Optional
+import warnings
+
+import numpy as np
+
+from pysisyphus.constants import BOHR2ANG
+from pysisyphus.drivers.merge import merge_with_frozen_geom
+from pysisyphus.helpers import geom_loader
+from pysisyphus.io.mol2 import parse_mol2, dict_to_mol2_string
+
+
+def delete_atoms_bonds_inplace(
+    as_dict: dict, inds: list[int], atom_offset: int = 0, bond_offset: int = 0
+) -> dict[int, int]:
+    """Update atom_ids and bonds in mol2-dict inplace.
+
+    Parameter
+    ---------
+    as_dict
+        mol2-dict as returned from pysisyphus.io.mol2.parse_mol2.
+    inds
+        List of positive integer atom_ids to be deleted.
+    atom_offset
+        Integer >= 0. atom_ids will be shifted by this number.
+    bond_offset
+        Integer >= 0. bond_ids will be shifted by this number.
+
+    Returns
+    -------
+    atom_map
+        Dictionary w/ origin atom_ids as keys and updated atom_ids as values.
+    """
+    to_del = set(inds)
+
+    # Atoms
+    axs = as_dict["atoms_xyzs"]
+    axs_mod = list()
+    atoms_del = 0
+    atom_map = {}
+    for ax in axs:
+        if ax["atom_id"] in to_del:
+            warnings.warn("Charges were not updated after deleting atoms!")
+            print(f"Deleted atom {ax['atom_name']} with atom_id {ax['atom_id']}")
+            atoms_del += 1
+            continue
+        atom_id = ax["atom_id"]
+        mod_atom_id = atom_id - atoms_del + atom_offset
+        # Store updated atom_id in dict with original atom_id as key,
+        # so we can later acces them to update the bond atom_ids.
+        atom_map[atom_id] = mod_atom_id
+        ax["atom_id"] = mod_atom_id
+        axs_mod.append(ax)
+    as_dict["atoms_xyzs"] = axs_mod
+
+    # Bonds
+    bonds_mod = list()
+    bonds_del = 0
+    for bond in as_dict["bond"]:
+        bond_inds = set((bond["origin_atom_id"], bond["target_atom_id"]))
+        if bond_inds & to_del:
+            print(f"Deleted bond with bond_id {bond['bond_id']}")
+            bonds_del += 1
+            continue
+        bond["bond_id"] -= bonds_del
+        bond["bond_id"] += bond_offset
+        bond["origin_atom_id"] = atom_map[bond["origin_atom_id"]]
+        bond["target_atom_id"] = atom_map[bond["target_atom_id"]]
+        bonds_mod.append(bond)
+    as_dict["bond"] = bonds_mod
+    return atom_map
+
+
+def update_coords_inplace(as_dict: dict, coords3d: np.ndarray):
+    """Update xyz coordinates in mol2 dict inplace."""
+    axs = as_dict["atoms_xyzs"]
+    assert len(axs) == len(coords3d)
+
+    for ax, nc3d in zip(axs, coords3d):
+        ax["xyz"] = (nc3d * BOHR2ANG).tolist()
+
+
+def merge_mol2_dicts(as_dict1: dict, as_dict2: dict) -> dict:
+    """Merge two mol2 dict. Atoms, bonds and associated numbers will be updated.
+
+    Parameters
+    ----------
+    as_dict1
+        Mol2 dict.
+    as_dict2
+        Mo2 dict.
+
+    Returns
+    -------
+    merged
+        Merged mol2 dict. Most entries will be copied from as_dict1 but atoms,
+        bonds and associated counts will be updated with entries from as_dict2.
+    """
+    # Atoms
+    atoms_xyzs = as_dict1["atoms_xyzs"] + as_dict2["atoms_xyzs"]
+
+    # Bonds
+    bond = as_dict1["bond"] + as_dict2["bond"]
+
+    merged = copy.deepcopy(as_dict1)
+    merged.update(
+        {
+            "atoms_xyzs": atoms_xyzs,
+            "num_atoms": len(atoms_xyzs),
+            "bond": bond,
+            "num_bonds": len(bond),
+            "mol_name": "merged",
+        }
+    )
+    return merged
+
+
+def get_substs(as_dict: dict) -> tuple[set[tuple[int, str]], int]:
+    keys = set()
+    for ax in as_dict["atoms_xyzs"]:
+        key = (ax["subst_id"], ax["subst_name"])
+        keys.add(key)
+    subst_ids, _ = zip(*keys)
+    max_subst_id = max(subst_ids)
+    return keys, max_subst_id
+
+
+def update_subst_ids_inplace(as_dict: dict, offset: int):
+    for ax in as_dict["atoms_xyzs"]:
+        ax["subst_id"] += offset
+
+
+def merge_mol2(
+    fn1: str,
+    fn2: str,
+    del1: Optional[list[int]] = None,
+    del2: Optional[list[int]] = None,
+    bonds: Optional[list[int]] = None,
+    new_coords: Optional[np.ndarray] = None,
+) -> dict:
+    if del1 is None:
+        del1 = list()
+    if del2 is None:
+        del2 = list()
+    if bonds is None:
+        bonds = list()
+    bonds = np.array(bonds, dtype=int).reshape(-1, 2)
+
+    dict1 = parse_mol2(fn1).as_dict()
+    # dict1 = m1.as_dict()
+    dict2 = parse_mol2(fn2).as_dict()
+    # dict2 = m2.as_dict()
+
+    # Delete atoms and assoicated bonds
+    print(f"Deleting atoms/bonds from '{fn1}'")
+    atom_map1 = delete_atoms_bonds_inplace(dict1, del1)
+    print()
+    natoms1 = len(dict1["atoms_xyzs"])
+    nbonds1 = len(dict1["bond"])
+
+    print(f"Deleting atoms/bonds from '{fn2}'")
+    atom_map2 = delete_atoms_bonds_inplace(
+        dict2, del2, atom_offset=natoms1, bond_offset=nbonds1
+    )
+
+    # Determine highest subst_id in dict1, to shift subst_ids in dict2
+    _, max_subst_id1 = get_substs(dict1)
+    print(f"Highest subst_id in '{fn1}' is {max_subst_id1}.")
+    update_subst_ids_inplace(dict2, max_subst_id1)
+
+    if new_coords is not None:
+        coords1 = new_coords[:natoms1]
+        coords2 = new_coords[natoms1:]
+        assert len(coords1) + len(coords2) == len(new_coords)
+        update_coords_inplace(dict1, coords1)
+        update_coords_inplace(dict2, coords2)
+        print("Updated coordinates")
+
+    # Add new bond(s) to dict2
+    bond2 = dict2["bond"]
+    nbonds2 = len(bond2)
+    for from_, to_ in bonds:
+        # Use updated atom_ids for the bond target/origin
+        origin = atom_map1[from_]
+        target = atom_map2[to_]
+        nbonds2 += 1
+        bond_type = 1
+        new_bond = {
+            "bond_id": nbonds1 + nbonds2,
+            "bond_type": bond_type,
+            "origin_atom_id": origin,
+            "target_atom_id": target,
+        }
+        print(f"Added new bond '{new_bond}'")
+        warnings.warn(f"Set bond_type to '{bond_type}'.")
+        bond2.append(new_bond)
+
+    merged = merge_mol2_dicts(dict1, dict2)
+    return merged
+
+
+def merge_mol2_geoms(
+    fn1: str,
+    fn2: str,
+    bonds: list[list[int]],
+    del1: Optional[list[int]] = None,
+    del2: Optional[list[int]] = None,
+    out_fn: Optional[str] = None,
+) -> str:
+    """Merge geometries in mol2 files w/ atom deletion and bond formation.
+
+    Parameters
+    ----------
+    fn1
+        Filename of first mol2 file.
+    fn2
+        Filename of second mol2 file.
+    bonds
+        List of list of two integers. Each integer pair comprises an atom
+        of fn1 and an atom of fn2. Original indices/atom_ids must be used,
+        regardless of any atom deletion.
+    del1
+        Optional list of integers of atoms to be deleted. As for 'bonds',
+        atom_ids as appearing in the mol2 file must used.
+    del2
+        Same as 'del1' but deletes atoms in 'fn2'.
+    out_fn
+        Optional str. If given the resulting mol2-file will be written to
+        this filename.
+
+    Returns
+    -------
+    rendered
+        Merged mol2 w/ deleted atoms.
+    """
+
+    if del1 is None:
+        del1 = []
+    if del2 is None:
+        del2 = []
+
+    del1 = list(sorted(del1))
+    del2 = list(sorted(del2))
+    # Make 0-based indices from 1-based mol2 indices
+    del10 = [d - 1 for d in del1]
+    del20 = [d - 1 for d in del2]
+    bonds0 = np.array(bonds, dtype=int).reshape(-1, 2) - 1
+
+    geom1 = geom_loader(fn1)
+    geom2 = geom_loader(fn2)
+    # Get new coordinates by merging both fragments.
+    # Function takes 0-based indices
+    new_geom, *_ = merge_with_frozen_geom(geom1, geom2, bonds0, del10, del20)
+    new_coords3d = new_geom.coords3d
+
+    # Merge mol2 files w/ updated coordinates
+    # Takes 1-based indices, bad?!
+    merged = merge_mol2(fn1, fn2, del1, del2, bonds, new_coords3d)
+
+    # Render mol2 string from merged data
+    rendered = dict_to_mol2_string(merged)
+
+    # Dump merged mol2 dict to file
+    if out_fn is not None:
+        with open(out_fn, "w") as handle:
+            handle.write(rendered)
+        print(f"Dumped merged geometry to '{out_fn}'.")
+    return rendered
+
+
+def parse_args(args):
+    parser = argparse.ArgumentParser(
+        description="Merge two mol2 files w/ atom deletion & bond formation."
+    )
+
+    parser.add_argument("fn1", help="Name of frozen mol2 file.")
+    parser.add_argument("fn2", help="Name of mobile mol2 file.")
+    parser.add_argument(
+        "--bonds",
+        nargs="+",
+        type=int,
+        help="1-based atom id pairs (id1, id2), between which bonds are formed."
+        "The original atom ids must be used, regardless of any atom deletions.",
+    )
+    parser.add_argument(
+        "--del1",
+        nargs="+",
+        type=int,
+        help="1-based atom ids to be deleted from fn1.",
+    )
+    parser.add_argument(
+        "--del2",
+        nargs="+",
+        type=int,
+        help="1-based atom ids to be deleted from fn2.",
+    )
+    parser.add_argument(
+        "--out", default="merged.mol2", help="Name of the final mol2 file."
+    )
+
+    return parser.parse_args(args)
+
+
+def run() -> str:
+    """CLI frontend for merge_mol2_geoms()."""
+
+    args = parse_args(sys.argv[1:])
+
+    fn1 = args.fn1
+    fn2 = args.fn2
+    bonds = args.bonds
+    del1 = args.del1
+    del2 = args.del2
+    out = args.out
+
+    return merge_mol2_geoms(fn1, fn2, bonds, del1, del2, out)
+
+
+if __name__ == "__main__":
+    run()