Source code for schrodinger.protein.tasks.binding_site_align

import csv
from collections import namedtuple
from typing import List

from schrodinger import structure
from schrodinger.models import parameters
from schrodinger.protein import residue
from schrodinger.tasks import jobtasks
from schrodinger.tasks import tasks
from schrodinger.utils.csv_unicode import reader_open

Result = namedtuple("Result", ["other_name", "rmsd"])

OTHER_NAME_COLUMN = "Mobile"
RMSD_COLUMN = "RMSD_after"


class BindingSiteAlignmentTask(jobtasks.CmdJobTask):
    """
    Manages binding site alignment job for input structures.
    """

    class Input(parameters.CompoundParam):
        sts: List[structure.Structure] = []
        binding_site_cutoff: float = 5.0
        mapping_dist: float = 5.0
        previously_aligned: bool = False
        align_sel_res_only: bool = False
        res_list: List[residue.Residue] = []

    class Output(jobtasks.CmdJobTask.Output):
        sts: List[structure.Structure] = []
        rmsd_rows: List[Result] = []

    DEFAULT_TASKDIR_SETTING = tasks.TEMP_TASKDIR

    @tasks.preprocessor(order=tasks.AFTER_TASKDIR)
    def writeInputFiles(self):
        self._in_file = self.getTaskFilename('input.mae')
        self._out_file = self.getTaskFilename(
            'align-binding-sites-align-final.maegz')
        structure.write_cts(self.input.sts, self._in_file)

    def makeCmd(self):
        """
        Returns command list based on input params
        :return: Command for align_binding_sites backend
        :rtype: list(str)
        """
        cmd = ['align_binding_sites.py',
               '-jobname', 'align-binding-sites',
               '-cutoff', str(self.input.binding_site_cutoff),
               '-dist', str(self.input.mapping_dist)]  # yapf: disable
        if self.input.previously_aligned:
            cmd.extend(['-prealigned'])
        if self.input.align_sel_res_only and self.input.res_list:
            cmd.extend(['-res', self._makeResiduesArgs(self.input.res_list)])
        cmd.extend([self._in_file])
        return cmd

    def _makeResiduesArgs(self, res_list):
        """
        Returns a str of selected residues to align in
        a binding site alignment. Residues correspond to residues in the
        reference sequence.

        :param res_list: List of residues to align
        :type res_list: list(protein.residue.Residue)

        :return: Residues in format chain:resnum separated by commas.
            If the residue does not have a chain, the format is _:resnum
        :rtype: str
        """
        res_str = []
        for res in res_list:
            seq_chain = res.sequence.structure_chain
            if not seq_chain or seq_chain == ' ':
                res_str.append(f'_:{res.resnum}')
            else:
                res_str.append(f'{seq_chain}:{res.resnum}')
        return ",".join(res_str)

    @tasks.postprocessor()
    def readOutputFiles(self):
        with structure.StructureReader(self._out_file) as st_reader:
            self.output.sts = list(st_reader)

        rmsd_file = self.getTaskFilename("align-binding-sites.csv")
        rmsd_rows = []
        with reader_open(rmsd_file) as fh:
            for row in csv.DictReader(fh):
                rmsd = float(row[RMSD_COLUMN])
                result = Result(row[OTHER_NAME_COLUMN], rmsd)
                rmsd_rows.append(result)
        self.output.rmsd_rows = rmsd_rows