Source code for schrodinger.protein.tasks.optimize_alignment

from schrodinger.models import parameters
from schrodinger.protein import sequence
from schrodinger.structutils import analyze
from schrodinger.tasks import tasks


class OptimizeAlignmentTask(tasks.BlockingFunctionTask):
    """
    Task to optimize the alignment between a target and template sequence by
    moving gaps out of secondary structure elements. See move_gaps_from_sse.py
    for additional documentation.
    """

    class Input(parameters.CompoundParam):
        template_seq: sequence.ProteinSequence = None
        target_seq: sequence.ProteinSequence = None
        ligand_asl: str
        residue_asl: str
        radius: float = 10.0

    class Output(parameters.CompoundParam):
        template_seq: sequence.ProteinSequence = None
        target_seq: sequence.ProteinSequence = None

    @staticmethod
    def _makeSeqStr(seq, require_structure=False):
        if require_structure:

            def accept(res):
                if res.is_res:
                    # skip structureless residues and X
                    return res.hasStructure() and res.short_code != 'X'
                return True

            elems = (ele for ele in seq if accept(ele))
        else:
            elems = seq

        return "".join(ele.short_code if ele.is_res else "." for ele in elems)

    @tasks.preprocessor
    def _checkInputs(self):
        if self.input.template_seq is None or self.input.target_seq is None:
            return False, "Both template and target seq must be set"
        if not self.input.template_seq.hasStructure():
            return False, "Template seq must have structure"
        if self.input.ligand_asl and self.input.residue_asl:
            return False, "Only one of ligand_asl and residue_asl can be set"

    def mainFunction(self):
        # Import inline so the module can be imported when prime is missing
        from schrodinger.application.prime.packages import move_gaps_from_sse
        full_template = self.input.template_seq.getStructure()
        chain_id = self.input.template_seq.structure_chain
        template = full_template.chain[chain_id].extractStructure()
        seq1 = self._makeSeqStr(self.input.target_seq)
        seq2 = self._makeSeqStr(self.input.template_seq, require_structure=True)
        in_aln = [seq1, seq2]
        atoms = None
        centroids = None
        if self.input.ligand_asl:
            ligand_asl = self.input.ligand_asl
            ligand_iatoms = analyze.evaluate_asl(template, ligand_asl)
            centroids = [template.atom[i].xyz for i in ligand_iatoms]
            asl = f'fillres within {self.input.radius} ({ligand_asl})'
            atoms = analyze.evaluate_asl(template, asl)
        elif self.input.residue_asl:
            atoms = analyze.evaluate_asl(template, self.input.residue_asl)
        out_aln = move_gaps_from_sse.move_gaps_from_sse(template,
                                                        in_aln,
                                                        atoms=atoms,
                                                        centroids=centroids)
        target_seq, template_seq = out_aln
        self.output.target_seq = sequence.ProteinSequence(target_seq)
        self.output.template_seq = sequence.ProteinSequence(template_seq)