Source code for schrodinger.application.steps.basesteps

from schrodinger.tasks import stepper
from schrodinger.utils import log

from .dataclasses import FileInMixin
from .dataclasses import MaeMaeMixin
from .dataclasses import MaeOutMixin
from .dataclasses import MolInMixin
from .dataclasses import MolMolMixin
from .dataclasses import MolOutMixin
from .utils import to_string


class LoggerMixin:
    """
    A mixin for stepper._BaseStep to allow debugging information about every
    input and output that is processed
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.logger = log.get_output_logger(self.getStepId())

    def reduceFunction(self, inputs):
        for inp in inputs:
            self.logger.debug(f'{self.getStepId()} <- {to_string(inp)}')
            for output in self.mapFunction(inp):
                self.logger.debug(f'{self.getStepId()} -> {to_string(output)}')
                yield output


class MolReduceStep(MolMolMixin, stepper.ReduceStep):
    pass


class MolMapStep(LoggerMixin, MolReduceStep):
    pass


class MaeReduceStep(MaeMaeMixin, stepper.ReduceStep):
    pass


class MaeMapStep(LoggerMixin, MaeReduceStep):
    pass


class OptionalStepsMixin:
    """
    A mixin that provides optional steps for a `Workflow` class.

    Whether the step is optional is determined by the `OPTIONAL_SETTINGS_MAP`, a
    dictionary with step classes as the key and the bool settings attributes as
    the value.

    E.g.
    class ExampleOptionalStepChain(OptionalStepMixin, MolMolWorkflow):
        STEPS = (filters.UniqueSmilesFilter, filters.SmartsFilter,
                 filters.PropertyFilter)

        OPTIONAL_SETTINGS_MAP = {filters.SmartsFilter: 'smarts_filter'}

        class Settings(parameters.CompoundParam):
            smarts_filter: bool = False
    """
    OPTIONAL_SETTINGS_MAP = {}

    def buildChain(self):
        for step in self.STEPS:
            use_step = True
            if step in self.OPTIONAL_SETTINGS_MAP:
                use_step = self.settings.getSubParam(
                    self.OPTIONAL_SETTINGS_MAP[step])
            if use_step:
                self.addStep(step())


class Workflow(stepper.Chain):
    """
    A chain of steps.

    To define which steps should be present, populate the `STEPS` variable.
    """
    STEPS = ()

    def buildChain(self):
        for step in self.STEPS:
            self.addStep(step())

    def _getStepsSettingAttributeValues(self, attr):
        """
        :return: recursively obtain all non None values of `attr` in all steps.
        :rtype: Set[object]
        """
        values = set()
        for step in self:
            if isinstance(step, stepper.Chain):
                values = values.union(
                    step._getStepsSettingAttributeValues(attr))
            elif getattr(step.settings, attr, None) is not None:
                values.add(getattr(step.settings, attr))
        return values


class MolMolWorkflow(MolMolMixin, Workflow):
    pass


class MolMaeWorkflow(MolInMixin, MaeOutMixin, Workflow):
    pass


class MaeMaeWorkflow(MaeMaeMixin, Workflow):
    pass


class FileMolWorkflow(FileInMixin, MolOutMixin, Workflow):
    pass