Source code for schrodinger.application.phase.shape_screen_reporter.tasks
import os
from typing import Set
from typing import List
from typing import Tuple
from typing import Optional
from hit_analysis_gui_dir import label_property_utils
from schrodinger import structure
from schrodinger.models import parameters
from schrodinger.tasks import jobtasks
from schrodinger.tasks import tasks
from schrodinger.utils import fileutils
from schrodinger.infra import phase
from schrodinger.application.phase.shape_screen_reporter import prop_utils
from schrodinger.application.phase.shape_screen_reporter import task_utils
ORDER_DB_CHECK = 0
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
[docs]class CreateScreeningDBTask(_AbstractShapeScreenReporterTaskMixin,
jobtasks.CmdJobTask):
"""
Task to create a virtual screening database. Adds a label property to the
newly created DB in a postprocessor.
"""
MODE = "create"
@tasks.preprocessor(tasks.AFTER_TASKDIR)
def _checkInput(self):
if self.input.label_prop and (self.input.label_prop
not in self.input.props):
return False, ("Expected the specified label property to be one of "
"the included properties")
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(ORDER_DB_CHECK)
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}"
@tasks.postprocessor(ORDER_DB_CHECK + 1)
def _writeLabelPropertyToDB(self):
"""
Write the specified label property to the output DB file.
"""
label_property_utils.set_label_property(self.input.label_prop,
self.getDBFile())
[docs] def getDBFile(self) -> str:
"""
@overrides: _AbstractShapeScreenReporterTaskMixin
"""
return self.getTaskFilename(f"{self.name}.vsdb")
[docs]class FilterDBTask(tasks.ComboSubprocessTask):
"""
Task to run filtering on a virtual screening database.
"""
@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}"
[docs] def mainFunction(self):
"""
Create the shape screening reporter and filter based on the property
and screening filters. Write the filtered structs to the hits file.
"""
reporter = phase.PhpShapeScreenReporter(self.input.screening_db_file)
prop_filters = []
for prop_tuple in self.input.property_filters:
prop_filters.append(','.join(str(item) for item in prop_tuple))
filtered_rows = None
if prop_filters:
property_filter = prop_utils.read_property_filters(prop_filters)
filtered_rows = task_utils.run_property_filter(
reporter, property_filter)
if self.input.sel_features:
hypo = task_utils.create_user_hypo(reporter,
list(self.input.sel_features))
min_sites = len(self.input.sel_features)
if prop_filters:
results = reporter.screenSitesInPlace(hypo, min_sites,
filtered_rows)
else:
results = reporter.screenSitesInPlace(hypo, min_sites)
filtered_rows = [result.row_number for result in results]
if filtered_rows is None:
raise tasks.TaskFailure('No valid filters are provided.')
structs = task_utils.get_structures(reporter, filtered_rows)
with structure.StructureWriter(
self.getTaskFilename(f'{self.name}.maegz')) as writer:
writer.extend(structs)
[docs] def getDBFile(self) -> str:
"""
@overrides: _AbstractShapeScreenReporterTaskMixin
"""
return self.getTaskFilename(self.input.screening_db_file)
[docs] 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')