_pipeline.py

'''High-level functionality for modelling module to build pipelines. Added in 
the __init__.py file. To be used directly by passing a ModellingHandle instance
as argument.
'''

# internal
from promod3 import loop
from promod3 import sidechain
from _modelling import *
from _closegaps import *
# external
import ost
from ost import mol,conop
from ost.mol import mm
import os

def BuildSidechains(mhandle):
    '''Build sidechains for model.

    This is esentially a wrapper for :func:`promod3.sidechain.Reconstruct`.

    :param mhandle: Modelling handle on which to apply change.
    :type mhandle:  :class:`ModellingHandle`
    '''
    ost.LogInfo("Rebuilding sidechains.")
    sidechain.Reconstruct(mhandle.model, keep_sidechains=True)

def MinimizeModelEnergy(mhandle, max_iterations=3, max_iter_sd=30, 
                        max_iter_lbfgs=20):
    '''Minimize energy of final model using molecular mechanics.

    Uses :mod:`ost.mol.mm` to perform energy minimization.
    It will iteratively (at most *max_iterations* times):
    
    - run up to *max_iter_sd* minimization iter. of a steepest descend method
    - run up to *max_iter_lbfgs* minimization iter. of a Limited-memory 
      Broyden-Fletcher-Goldfarb-Shanno method
    - abort if no stereochemical problems found

    The idea is that we don't want to minimize "too much". So, we iteratively
    minimize until there are no stereochemical problems and not more.

    To speed things up, this can run on multiple CPU threads by setting the
    env. variable ``PM3_OPENMM_CPU_THREADS`` to the number of desired threads.
    If the variable is not set, 1 thread will be used by default.

    :param mhandle: Modelling handle on which to apply change.
    :type mhandle:  :class:`ModellingHandle`

    :param max_iterations: Max. number of iterations for SD+LBFGS
    :type max_iterations:  :class:`int`

    :param max_iter_sd: Max. number of iterations within SD method
    :type max_iter_sd:  :class:`int`

    :param max_iter_lbfgs: Max. number of iterations within LBFGS method
    :type max_iter_lbfgs:  :class:`int`
    '''
    ost.LogInfo("Minimize energy.")
    # ignore LogInfo in stereochemical problems if output up to info done
    ignore_stereo_log = (ost.GetVerbosityLevel() == 3)
    # setup mm simulation
    settings = mm.Settings()
    settings.integrator = mm.LangevinIntegrator(310, 1, 0.002)
    settings.init_temperature = 0
    settings.forcefield = mm.LoadCHARMMForcefield()
    settings.nonbonded_method = mm.NonbondedMethod.CutoffNonPeriodic
    settings.keep_ff_specific_naming = False
    # switch this to "mm.Platform.Reference" for debugging
    settings.platform = mm.Platform.CPU
    if os.getenv('PM3_OPENMM_CPU_THREADS') is None:
        settings.cpu_properties["CpuThreads"] = "1"
    else:
        settings.cpu_properties["CpuThreads"] = os.getenv('PM3_OPENMM_CPU_THREADS')
    sim = mm.Simulation(mhandle.model,settings)
    # settings to check for stereochemical problems
    clashing_distances = mol.alg.DefaultClashingDistances()
    bond_stereo_chemical_param = mol.alg.DefaultBondStereoChemicalParams()
    angle_stereo_chemical_param = mol.alg.DefaultAngleStereoChemicalParams()

    for i in range(max_iterations):
        # update atoms
        ost.LogInfo("Perform energy minimization (iteration %d)" % (i+1))
        sim.ApplySD(tolerance = 1.0, max_iterations = max_iter_sd)
        sim.ApplyLBFGS(tolerance = 1.0, max_iterations = max_iter_lbfgs)
        sim.UpdatePositions()

        # check for stereochemical problems
        if ignore_stereo_log:
            ost.PushVerbosityLevel(2)
        temp_ent = sim.GetEntity()
        temp_ent = temp_ent.Select("aname!=OXT")
        temp_ent_clash_filtered = mol.alg.FilterClashes(\
                                                temp_ent, clashing_distances)[0]
        # note: 10,10 parameters below are hard coded bond-/angle-tolerances
        temp_ent_stereo_checked = mol.alg.CheckStereoChemistry(\
                                                temp_ent_clash_filtered, 
                                                bond_stereo_chemical_param,
                                                angle_stereo_chemical_param, 
                                                10, 10)[0]
        if ignore_stereo_log:
            ost.PopVerbosityLevel()
        # checks above would remove bad atoms
        if len(temp_ent_stereo_checked.Select("ele!=H").atoms) \
           == len(temp_ent.Select("ele!=H").atoms):
            break
        else:
            ost.LogInfo("Stereo-chemical problems found, need more energy"
                        + " minimization (iteration %d)" % (i+1))

    # update model
    simulation_ent = sim.GetEntity()
    mhandle.model = mol.CreateEntityFromView(\
                            simulation_ent.Select("peptide=true and ele!=H"),
                            True)

def BuildFromRawModel(mhandle):
    '''Build a model starting with a raw model (see :func:`BuildRawModel`).

    This function implements a recommended pipeline to generate complete models
    from a raw model. The steps are shown in detail in the code example
    :ref:`above <modelling_steps_example>`. If you wish to use your own
    pipeline, you can use that code as a starting point for your own custom
    modelling pipeline. For reproducibility, we recommend that you keep copies
    of custom pipelines. 

    If the function fails to close all gaps, it will produce a warning and
    return an incomplete model.

    :param mhandle: The prepared template coordinates loaded with the input
                    alignment.
    :type mhandle:  :class:`~promod3.modelling.ModellingHandle`

    :return: Delivers the model as an |ost_s| entity.
    :rtype: :class:`Entity <ost.mol.EntityHandle>`
    '''
    ost.LogInfo("Starting modelling based on a raw model.")

    # a bit of setup
    fragment_db = loop.LoadFragDB()
    structure_db = loop.LoadStructureDB()
    torsion_sampler = loop.LoadTorsionSamplerCoil()
    merge_distance = 4

    scorer = SetupBackboneScorer(mhandle)

    # close small deletions and remove terminal gaps
    CloseSmallDeletions(mhandle, scorer)
    RemoveTerminalGaps(mhandle)

    # iteratively merge gaps of distance i and fill loops by database
    for distance in range(merge_distance):
        MergeGapsByDistance(mhandle, distance)
        FillLoopsByDatabase(mhandle, scorer, fragment_db, structure_db,
                            torsion_sampler, min_loops_required=-1,
                            max_res_extension=6)
    FillLoopsByDatabase(mhandle, scorer, fragment_db, structure_db,
                        torsion_sampler, min_loops_required=-1)
    FillLoopsByDatabase(mhandle, scorer, fragment_db, structure_db,
                        torsion_sampler)

    # close remaining gaps by Monte Carlo
    FillLoopsByMonteCarlo(mhandle, scorer, torsion_sampler)
    CloseLargeDeletions(mhandle, scorer, structure_db)

    # check if we succeeded
    if len(mhandle.gaps) > 0:
        ost.LogWarning("Failed to close %d gap(s). Returning incomplete model!" % \
                            len(mhandle.gaps))

    # build sidechains
    BuildSidechains(mhandle)

    # minimize energy of final model using molecular mechanics
    MinimizeModelEnergy(mhandle)

    # done
    return mhandle.model

# these methods will be exported into module
__all__ = ('BuildFromRawModel', 'BuildSidechains', 'MinimizeModelEnergy',)