Source code for schrodinger.application.phase.shape_screen_reporter.tasks

import os
from typing import List
from typing import Optional

from schrodinger.models import parameters
from schrodinger.tasks import jobtasks
from schrodinger.tasks import tasks
from schrodinger.utils import fileutils


class _AbstractShapeScreenReporterTaskMixin(parameters.CompoundParamMixin):
    """
    Abstract task to run shape_screen_reporter.

    Subclasses must mix this class into a Task subclass and define `MODE`. MODE
    must be one of the modes of shape_screen_reporter. Subclasses must also
    define getDBFile in order to specify the path to the database file used in
    makeCmd.

    Subclasses may override _getCmdOptions to add additional command options.

    The .vsdb file will be added to the command as `self.name` + .vsdb
    """

    CMD = "shape_screen_reporter"
    MODE = NotImplemented

    def makeCmd(self):
        cmd = [self.CMD, self.MODE]
        if not jobtasks.is_jobtask(self):
            cmd.append('-NOJOBID')
        cmd.extend(self._getCmdOptions())
        cmd.append(self.getDBFile())
        return cmd

    def _getCmdOptions(self):
        return []

    def getDBFile(self) -> str:
        """
        Return the path to the screening database file.
        """
        raise NotImplementedError


class FilterDBTask(_AbstractShapeScreenReporterTaskMixin,
                   tasks.SubprocessCmdTask):
    """
    Task to run filtering on a virtual screening database.
    """

    MODE = "filter"

    class Input(parameters.CompoundParam):
        screening_db_file: tasks.TaskFile
        property_filters_file: tasks.TaskFile

    @tasks.preprocessor(tasks.AFTER_TASKDIR)
    def _checkDB(self):
        db_file = self.input.screening_db_file
        if not os.path.isfile(db_file):
            return False, f"Did not find expected database file {db_file}"

    @tasks.preprocessor(tasks.AFTER_TASKDIR)
    def _checkFilters(self):
        filters_file = self.input.property_filters_file
        if not os.path.isfile(filters_file):
            msg = f"Did not find expected property filters file {filters_file}"
            return False, msg

    def _getCmdOptions(self):
        opts = ['-hits', self.getHitsFile()]
        opts.extend(['-props', self.input.property_filters_file])
        return opts

    def getDBFile(self) -> str:
        """
        @overrides: _AbstractShapeScreenReporterTaskMixin
        """
        return self.getTaskFilename(self.input.screening_db_file)

    def getHitsFile(self) -> Optional[str]:
        """
        Return the path to the output hits file or None if the taskdir has not
        yet been created.
        """
        if not os.path.exists(self.getTaskDir()):
            return None
        return self.getTaskFilename(f'{self.name}.maegz')


class CreateScreeningDBTask(_AbstractShapeScreenReporterTaskMixin,
                            jobtasks.CmdJobTask):
    """
    Task to create a virtual screening database.
    """

    MODE = "create"

    class Input(parameters.CompoundParam):
        hits_file: tasks.TaskFile
        query_file: tasks.TaskFile = None
        props: List = []

    @tasks.preprocessor(tasks.AFTER_TASKDIR)
    def _checkInput(self):
        if not self.input.hits_file:
            return False, "Hits file must be specified"
        if not os.path.isfile(self.input.hits_file):
            return False, f"Specified hits file {self.input.hits_file} does not exist"
        is_st_file = fileutils.get_structure_file_format(
            self.input.hits_file) is not None
        if not is_st_file:
            return False, ("Expected the specified hit file to be a structure"
                           "file")

        if self.input.query_file is None:
            return
        if not os.path.isfile(self.input.query_file):
            return False, (f"Specified query file {self.input.query_file} does"
                           "not exist")
        is_st_file = fileutils.get_structure_file_format(
            self.input.query_file) is not None
        if not is_st_file:
            return False, ("Expected the specified query file to be a structure"
                           "file")

    def _getCmdOptions(self):
        opts = ["-hits", self.input.hits_file]
        if self.input.query_file:
            opts.extend(['-query', self.input.query_file])
        if self.input.props:
            opts.extend(['-props', ','.join(self.input.props)])
        return opts

    @tasks.postprocessor()
    def _checkDB(self):
        db_file = self.getDBFile()
        # Check that DB file exists in the (auto-populated) output files
        if db_file not in self.output.output_files:
            return False, "Did not find expected database file {db_file}"

    def getDBFile(self) -> str:
        """
        @overrides: _AbstractShapeScreenReporterTaskMixin
        """
        return self.getTaskFilename(f"{self.name}.vsdb")