Source code for schrodinger.application.steps.ligprep

import os

from schrodinger.infra import mmerr
from schrodinger.models import parameters
from schrodinger import stepper
from schrodinger.utils import license
from schrodinger.utils import log

from . import dataclasses
from . import utils

# for unit test environment
try:
    from schrodinger.application.macromodel.packages.ligprep3.ligprepworkflow import \
        LigPrepWorkflow
except ImportError:
    LigPrepWorkflow = None

LIGPREP_MMERR_LEVEL = mmerr.MMERR_OFF
LIGPREP_LOG_LEVEL = log.CRITICAL

# We can only instantiate a single LigPrepWorkflow instances b/c we are not
# allowed to instantiate Epik more than once.
mm_lig_prep_workflow = None
mm_lig_prep_workflow_args = None


class LigPrepStep(stepper.MapStep):
    """
    A step that implements the basic lig prep functionality of `Structure`
    objects.

    The settings is the list of command line arguments that are usually passed
    to the ligprep script, excluding the input and output arguments.

    Within one process more than one LigPrepStep is allowed to be created only
    if the settings are the same, to avoid problems if both use '-epik'.

    There are issues with generating tautomers even if the first generator is
    closed (LIGPREP-1941).
    """

    def getLicenseRequirements(self):
        return {license.LIGPREP_MAIN: 1}

    class Settings(parameters.CompoundParam):
        arg_string: str = '-bff 16 -epik -s 32'
        ligprep_filter_file: stepper.StepperFile = None

    def validateSettings(self):
        return utils.validate_file(
            self, 'ligprep_filter_file')  # optional StepperFile

    def _makeLigPrepWorkflowArgs(self):
        """
        :return: the arguments for the LigPrepWorkflow initialization
        :rtype: List[str]
        :raises AutoDesignerError: if the ligprep_filter file was defined but
            not found.
        """
        # the initializer needs to have input and output files, but for my
        # use case they are not used.
        args = self.settings.arg_string.split() + ['-isd', 'foo', '-osd', 'bar']
        flt_file = self.settings.ligprep_filter_file
        if flt_file:
            if not os.path.exists(flt_file):
                msg = (f'{self._step_id}: ligprep_filter_file "{flt_file}" not'
                       f' found.')
                raise utils.StepsError(msg)
            args += ['-f', flt_file]
        return args

    def _getLigPrepWorkflow(self):
        global mm_lig_prep_workflow, mm_lig_prep_workflow_args
        args = self._makeLigPrepWorkflowArgs()
        if mm_lig_prep_workflow is None:
            lpwf_logger = log.get_output_logger("LigandPrepper")
            lpwf_logger.setLevel(LIGPREP_LOG_LEVEL)
            mm_lig_prep_workflow_args = args
            mm_lig_prep_workflow = LigPrepWorkflow(lpwf_logger, args)
        elif args != mm_lig_prep_workflow_args:
            msg = (f'{self._step_id}: a LigandPrepper with different arguments'
                   f' was already instantiated.')
            raise utils.StepsError(msg)
        return mm_lig_prep_workflow

    def mapFunction(self, struc):
        lpwf = self._getLigPrepWorkflow()
        with mmerr.Level(LIGPREP_MMERR_LEVEL):
            outcomes, _, _ = lpwf.run(struc)
        yield from outcomes

    def cleanUp(self):
        global mm_lig_prep_workflow, mm_lig_prep_workflow_args
        super().cleanUp()
        if mm_lig_prep_workflow:
            mm_lig_prep_workflow.close()
            mm_lig_prep_workflow = None
            mm_lig_prep_workflow_args = None


class LigandPrepper(dataclasses.MolMolMixin, LigPrepStep):
    """
    See also `LigPrepMixin`
    """

    def mapFunction(self, mol):
        st = utils.mol_to_structure(mol, self)
        if st is not None:
            for ligprepped_st in super().mapFunction(st):
                mol = utils.structure_to_mol(ligprepped_st, self, mol)
                if mol is not None:
                    yield mol


class MaeLigandPrepper(dataclasses.MaeMaeMixin, LigPrepStep):
    """
    See also `LigPrepStep`
    """
    # the parent basically implements the `Structure` input/output
    pass