Source code for schrodinger.application.desmond.stage.app.fragment_linking.utils

import os
import typing

from schrodinger.application.desmond import measurement
from schrodinger.application.desmond import struc
from schrodinger.application.desmond.constants import FEP_DG_PREFIX
from schrodinger.application.desmond.constants import FEP_RESTRAINT_CORRECTION
from schrodinger.application.desmond.constants import FEP_TRANSFER_FREE_ENERGY
from schrodinger.application.desmond.constants import FRAGMENT_LINKING_JOBS
from schrodinger.application.desmond.constants import FRAGMENT_LINKING_SIGN
from schrodinger.application.desmond.constants import CorrectionTerm

if typing.TYPE_CHECKING:
    from schrodinger.application.scisol.packages.fep import graph


class MissingValueError(Exception):
    pass


def collect_dg_values(g: "graph.Graph", launch_dir: str, jobname: str):
    """
    For each fragment linking edge, collect the dg values from the `*out.mae`
    files and update the edge properties.

    :param g: Graph to update in place.
    :param launch_dir: Path containing the FepLauncher job outputs.
    :param jobname: Name of the main job.

    :raises MissingValueError: If one or more edges were missing values.
    """
    err = ''
    for e in g.edges_iter():
        if not e.is_fragment_linking:
            continue
        e.complex_dg = None
        e.solvent_dg = None
        e.fragment_dgs = []
        from_id = e.short_id[0]
        to_id = e.short_id[1]
        from_title = e[0].struc.title
        to_title = e[1].struc.title

        job_to_dG = {}
        restr_correction = None
        sign = None
        for fragment_linking_job in FRAGMENT_LINKING_JOBS.VAL:
            outpath = f'{launch_dir}/{jobname}_{from_id}_{to_id}_{fragment_linking_job}-out.mae'
            if not os.path.exists(outpath):
                continue
            cts = struc.read_all_ct(outpath)
            for ct in cts:
                if ct.property.get(FEP_DG_PREFIX):
                    dG = measurement.Measurement.from_string(
                        ct.property[FEP_DG_PREFIX])
                    job_to_dG[fragment_linking_job] = dG
                elif ct.property.get(FEP_TRANSFER_FREE_ENERGY):
                    dG = measurement.Measurement.from_string(
                        ct.property[FEP_TRANSFER_FREE_ENERGY])
                    job_to_dG[fragment_linking_job] = dG
                if ct.property.get(FEP_RESTRAINT_CORRECTION):
                    restr_correction = measurement.Measurement(
                        ct.property[FEP_RESTRAINT_CORRECTION], 0.0)
                if ct.property.get(FRAGMENT_LINKING_SIGN):
                    sign = ct.property[FRAGMENT_LINKING_SIGN]

        # Check for missing properties
        missing_jobs = set(FRAGMENT_LINKING_JOBS.VAL) - set(job_to_dG.keys())
        if len(missing_jobs):
            err += f"Could not find fragment linking job(s) for '{from_title}' '{to_title}': {missing_jobs}.\n"
            continue
        if restr_correction is None:
            err += f"Could not find fragment linking restraint correction for '{from_title}' '{to_title}'.\n"
            continue
        if sign is None:
            err += f"Could not find fragment linking sign term for '{from_title}' '{to_title}'.\n"
            continue
        # Assign the results
        print(f'Fragment linking for {e[0].struc.title} {e[1].struc.title}:\n')

        e.complex_dg = job_to_dG[FRAGMENT_LINKING_JOBS.VAL.COMPLEX]
        e.solvent_dg = job_to_dG[FRAGMENT_LINKING_JOBS.VAL.SOLVENT]

        ddg_solvation = job_to_dG[FRAGMENT_LINKING_JOBS.VAL.FRAGMENT_HYDRATION]
        ddg_solvation -= job_to_dG[
            FRAGMENT_LINKING_JOBS.VAL.RESTRAINED_FRAGMENT_HYDRATION]
        ddg_solvation += restr_correction

        e.update_ddg_corrections_map(
            {CorrectionTerm.FRAGMENT_SOLVATION: ddg_solvation * sign})

        print(f'complex: {job_to_dG[FRAGMENT_LINKING_JOBS.VAL.COMPLEX]}')
        print(f'solvent: {job_to_dG[FRAGMENT_LINKING_JOBS.VAL.SOLVENT]}')
        print(
            f'fragment hydration free energy: {job_to_dG[FRAGMENT_LINKING_JOBS.VAL.FRAGMENT_HYDRATION]}'
        )
        print(
            f'restrained fragment hydration free energy: {job_to_dG[FRAGMENT_LINKING_JOBS.VAL.RESTRAINED_FRAGMENT_HYDRATION]}'
        )
        print(f'restraint correction: {restr_correction}')
        print(f'ddG: {e.bennett_ddg}')

    if err:
        raise MissingValueError(err)