dockformerpp/utils/relax.py

import sys
import os.path
import time
from random import shuffle

import numpy as np
import pdbfixer
import openmm as mm
import openmm.app as mm_app
import openmm.unit as mm_unit
from openmm import CustomExternalForce
from openmm.app import Modeller
from openmmforcefields.generators import SystemGenerator
from openff.toolkit import Molecule
from openff.toolkit.utils.exceptions import UndefinedStereochemistryError, RadicalsNotSupportedError
import mdtraj
from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit.Geometry import Point3D
import Bio.PDB
from Bio.SVDSuperimposer import SVDSuperimposer


# -- Relax protein and ligand. Code adapted from:
# https://github.com/patrickbryant1/Umol/blob/f7cd2b4de09b4e7cc1b68606791dd1cc81deeebc/src/relax/openmm_relax.py
def fix_pdb(pdb_path, hydrogen_added_pdb_path):
    """Add hydrogens to the PDB file
    """
    fixer = pdbfixer.PDBFixer(pdb_path)
    fixer.findMissingResidues()
    fixer.findNonstandardResidues()
    fixer.replaceNonstandardResidues()
    fixer.findMissingAtoms()
    fixer.addMissingAtoms()
    fixer.addMissingHydrogens(7.0)
    mm_app.PDBFile.writeFile(fixer.topology, fixer.positions, open(hydrogen_added_pdb_path, 'w'))
    return fixer.topology, fixer.positions


def minimize_energy(topology, system, positions, output_pdb_path):
    '''Function that minimizes energy, given topology, OpenMM system, and positions '''
    # Use a Brownian Integrator
    integrator = mm.BrownianIntegrator(
        100 * mm.unit.kelvin,
        100. / mm.unit.picoseconds,
        2.0 * mm.unit.femtoseconds
    )
    simulation = mm.app.Simulation(topology, system, integrator)

    # Initialize the DCDReporter
    reportInterval = 100  # Adjust this value as needed
    reporter = mdtraj.reporters.DCDReporter('positions.dcd', reportInterval)

    # Add the reporter to the simulation
    simulation.reporters.append(reporter)

    simulation.context.setPositions(positions)

    simulation.minimizeEnergy(1, 1000)
    # Save positions
    minpositions = simulation.context.getState(getPositions=True).getPositions()
    mm_app.PDBFile.writeFile(topology, minpositions, open(output_pdb_path, "w"))

    reporter.close()

    return topology, minpositions


def add_restraints(system, topology, positions, restraint_type):
    # Code adapted from https://gist.github.com/peastman/ad8cda653242d731d75e18c836b2a3a5
    restraint = CustomExternalForce('k*periodicdistance(x, y, z, x0, y0, z0)^2')
    system.addForce(restraint)
    restraint.addGlobalParameter('k', 100.0*mm_unit.kilojoules_per_mole/mm_unit.nanometer**2)
    restraint.addPerParticleParameter('x0')
    restraint.addPerParticleParameter('y0')
    restraint.addPerParticleParameter('z0')

    for atom in topology.atoms():
        if restraint_type == 'protein':
            if 'x' not in atom.name:
                restraint.addParticle(atom.index, positions[atom.index])
        elif restraint_type == 'CA+ligand':
            if ('x' in atom.name) or (atom.name == "CA"):
                restraint.addParticle(atom.index, positions[atom.index])

    return system


def create_joined_relaxed(protein_pdb_path: str, ligand_sdf_path: str, hydorgen_added_protein_pdb_path: str,
                          relaxed_joined_path: str):
    restraint_type = 'CA+ligand'

    start_time = time.time()
    print('Reading ligand')
    try:
        ligand_mol = Molecule.from_file(ligand_sdf_path)
    # Check for undefined stereochemistry, allow undefined stereochemistry to be loaded
    except UndefinedStereochemistryError:
        print('Undefined Stereochemistry Error found! Trying with undefined stereo flag True')
        ligand_mol = Molecule.from_file(ligand_sdf_path, allow_undefined_stereo=True)
    # Check for radicals -- break out of script if radical is encountered
    except RadicalsNotSupportedError:
        print('OpenFF does not currently support radicals -- use unrelaxed structure')
        sys.exit()
    # Assigning partial charges first because the default method (am1bcc) does not work
    ligand_mol.assign_partial_charges(partial_charge_method='gasteiger')

    # Read protein PDB and add hydrogens
    protein_topology, protein_positions = fix_pdb(protein_pdb_path, hydorgen_added_protein_pdb_path)
    print('Added all atoms...')

    modeller = Modeller(protein_topology, protein_positions)
    print('System has %d atoms' % modeller.topology.getNumAtoms())

    print('Adding ligand...')
    lig_top = ligand_mol.to_topology()
    modeller.add(lig_top.to_openmm(), lig_top.get_positions().to_openmm())
    print('System has %d atoms' % modeller.topology.getNumAtoms())

    print('Preparing system')
    # Initialize a SystemGenerator using the GAFF for the ligand and implicit water.
    # forcefield_kwargs = {'constraints': mm_app.HBonds, 'rigidWater': True, 'removeCMMotion': False,
    # 'hydrogenMass': 4*mm_unit.amu }
    system_generator = SystemGenerator(
        forcefields=['amber14-all.xml', 'implicit/gbn2.xml'],
        small_molecule_forcefield='gaff-2.11',
        molecules=[ligand_mol],
        # forcefield_kwargs=forcefield_kwargs
    )

    system = system_generator.create_system(modeller.topology, molecules=ligand_mol)

    print('Adding restraints on protein CAs and ligand atoms')

    system = add_restraints(system, modeller.topology, modeller.positions, restraint_type=restraint_type)

    minimize_energy(modeller.topology, system, modeller.positions, relaxed_joined_path)

    print(f'Time taken for relax calculation is {time.time() - start_time:.1f} seconds')


# -- Fix ligand changed structure. Code adapted from:
# https://github.com/patrickbryant1/Umol/blob/f7cd2b4de09b4e7cc1b68606791dd1cc81deeebc/src/relax/align_ligand_conformer.py
def generate_best_conformer(pred_coords, ligand_smiles, max_confs=100):
    """Generate conformers and compare the coords with the predicted atom positions

    Generating with constraints doesn't seem to work.
    cids = Chem.rdDistGeom.EmbedMultipleConfs(m,max_confs,ps)
    if len([x for x in m.GetConformers()])<1:
        print('Could not generate conformer with constraints')
    """
    # Generate conformers
    m = Chem.AddHs(Chem.MolFromSmiles(ligand_smiles))
    # Embed in 3D to get distance matrix
    AllChem.EmbedMolecule(m, maxAttempts=500)
    bounds = AllChem.Get3DDistanceMatrix(m)
    # Get pred distance matrix
    pred_dmat = np.sqrt(1e-10 + np.sum((pred_coords[:, None] - pred_coords[None, :]) ** 2 ,axis=-1))
    # Go through the atom types and add the constraints if not H
    # The order here will be the same as for the pred ligand as the smiles are identical
    ai, mi = 0, 0
    bounds_mapping = {}
    for atom in m.GetAtoms():
        if atom.GetSymbol() != 'H':
            bounds_mapping[ai] = mi
            ai += 1
        mi += 1

    # Assign available pred bound atoms
    bounds_keys = [*bounds_mapping.keys()]
    for i in range(len(bounds_keys)):
        key_i = bounds_keys[i]
        for j in range(i+1, len(bounds_keys)):
            key_j = bounds_keys[j]
            try:
                bounds[bounds_mapping[key_i], bounds_mapping[key_j]] = pred_dmat[i, j]
                bounds[bounds_mapping[key_j], bounds_mapping[key_i]] = pred_dmat[j, i]
            except:
                continue
    # Now generate conformers using the bounds
    ps = Chem.rdDistGeom.ETKDGv3()
    ps.randomSeed = 0xf00d
    ps.SetBoundsMat(bounds)
    cids = Chem.rdDistGeom.EmbedMultipleConfs(m, max_confs)
    # Get all conformer dmats
    nonH_inds = [*bounds_mapping.values()]
    conf_errs = []
    for conf in m.GetConformers():
        pos = conf.GetPositions()
        nonH_pos = pos[nonH_inds]
        conf_dmat = np.sqrt(1e-10 + np.sum((nonH_pos[:,None]-nonH_pos[None,:])**2,axis=-1))
        err = np.mean(np.sqrt(1e-10 + (conf_dmat-pred_dmat)**2))
        conf_errs.append(err)

    # Get the best
    best_conf_id = np.argmin(conf_errs)
    best_conf_err = conf_errs[best_conf_id]
    best_conf = [x for x in m.GetConformers()][best_conf_id]
    best_conf_pos = best_conf.GetPositions()

    return best_conf, best_conf_pos, best_conf_err, [atom.GetSymbol() for atom in m.GetAtoms()], nonH_inds, m, best_conf_id


def align_coords_transform(pred_pos, conf_pos, nonH_inds):
    """Align the predicted and conformer positions
    """
    sup = SVDSuperimposer()

    sup.set(pred_pos, conf_pos[nonH_inds])  # (reference_coords, coords)
    sup.run()
    rot, tran = sup.get_rotran()

    # Rotate coords from new chain to its new relative position/orientation
    tr_coords = np.dot(conf_pos, rot) + tran

    return tr_coords


def write_sdf(mol, conf, aligned_conf_pos, best_conf_id, outname):
    for i in range(mol.GetNumAtoms()):
        x, y, z = aligned_conf_pos[i]
        conf.SetAtomPosition(i, Point3D(x, y, z))

    writer = Chem.SDWriter(outname)
    writer.write(mol, confId=int(best_conf_id))


# Main function
def relax_complex(protein_pdb_path: str, ligand_sdf_path: str, relaxed_protein_path: str, relaxed_ligand_path: str):
    hydorgen_added_protein_pdb_path = protein_pdb_path + "_hydrogen_added.pdb"
    relaxed_joined_path = protein_pdb_path + "_joined_relaxed.pdb"

    create_joined_relaxed(protein_pdb_path, ligand_sdf_path, hydorgen_added_protein_pdb_path, relaxed_joined_path)

    parser = Bio.PDB.PDBParser(QUIET=True)
    joined_structure = next(iter(parser.get_structure('', relaxed_joined_path)))

    # save the relaxed protein
    io = Bio.PDB.PDBIO()
    io.set_structure(joined_structure["A"])
    io.save(relaxed_protein_path)

    relaxed_ligand_coords = np.array([atom.get_coord() for atom in joined_structure["B"].get_atoms()
                                      if atom.get_id()[0] != "H"])
    original_ligand = Chem.SDMolSupplier(ligand_sdf_path)[0]
    ligand_smiles = Chem.MolToSmiles(original_ligand)

    best_conf, best_conf_pos, best_conf_err, atoms, nonH_inds, mol, best_conf_id = generate_best_conformer(
        relaxed_ligand_coords, ligand_smiles, max_confs=100
    )

    aligned_conf_pos = align_coords_transform(relaxed_ligand_coords, best_conf_pos, nonH_inds)

    write_sdf(mol, best_conf, aligned_conf_pos, best_conf_id, relaxed_ligand_path)


def relax_folder(folder_path: str):
    all_jobnames = []
    filenames = os.listdir(folder_path)
    shuffle(filenames)
    for filename in filenames:
        if filename.endswith("_predicted_protein.pdb"):
            jobname = filename.split("_predicted_protein.pdb")[0]
            ligand_path = os.path.join(folder_path, jobname + "_predicted_ligand_0.sdf")
            if not os.path.exists(ligand_path):
                continue
            all_jobnames.append(jobname)

    success = 0
    for jobname in all_jobnames:
        protein_pdb_path = os.path.join(folder_path, jobname + "_predicted_protein.pdb")
        ligand_sdf_path = os.path.join(folder_path, jobname + "_predicted_ligand_0.sdf")
        relaxed_protein_path = os.path.join(folder_path, jobname + "_protein_relaxed.pdb")
        relaxed_ligand_path = os.path.join(folder_path, jobname + "_ligand_relaxed.sdf")
        if os.path.exists(relaxed_protein_path) and os.path.exists(relaxed_ligand_path):
            print("Already has relaxed", jobname)
            success += 1
            continue
        print("Relaxing", jobname)
        try:
            relax_complex(protein_pdb_path, ligand_sdf_path, relaxed_protein_path, relaxed_ligand_path)
            success += 1
        except Exception as e:
            print("Failed to relax", jobname, e)
    print(f"Relaxed {success}/{len(all_jobnames)}")


if __name__ == "__main__":
    relax_folder(os.path.abspath(sys.argv[1]))