Source code for schrodinger.application.steps.dataclasses

"""
Data classes to be used as input/output objects for steps and their associated
serializers which may be used when batching those steps.
"""

import warnings

from rdkit import Chem

from schrodinger import structure
from schrodinger.models import parameters
from schrodinger import stepper


class StepsError(RuntimeError):
    pass


class MolToSmilesSerializer(stepper.Serializer):
    """
    A serializer using RDKit's canonical SMILES of the Mol as the serialized
    form.
    """
    DataType = Chem.Mol

    def deserialize(self, fname):
        with open(fname, 'r') as infile:
            for line in infile:
                line = line.strip()
                inp = self.fromString(line)
                if inp is None:
                    warning_msg = f'ERROR: FAILED TO PARSE SMILES - "{line.strip()}"'
                    warnings.warn(RuntimeWarning(warning_msg))
                else:
                    yield inp

    def fromString(self, smiles_str):
        return Chem.MolFromSmiles(smiles_str)

    def toString(self, output_mol):
        return Chem.MolToSmiles(output_mol)


class MolInMixin:
    Input = Chem.Mol
    InputSerializer = MolToSmilesSerializer


class MolOutMixin:
    Output = Chem.Mol
    OutputSerializer = MolToSmilesSerializer


class MolMolMixin(MolInMixin, MolOutMixin):
    pass


class ScoredMol(parameters.CompoundParam):
    mol: Chem.Mol
    score: float = None

    def __str__(self):
        return f'{Chem.MolToSmiles(self.mol)} {self.score}'


class ScoredMolSerializer(stepper.Serializer):
    """
    A serializer using RDKit's canonical SMILES of the Mol and the score value
    string separated by a ',' as the serialized form.
    """
    DataType = ScoredMol
    SEPARATOR = ','

    def toString(self, scored_mol):
        smiles = Chem.MolToSmiles(scored_mol.mol)
        return smiles + self.SEPARATOR + str(scored_mol.score)

    def fromString(self, scored_mol_str):
        smiles, score = scored_mol_str.split(self.SEPARATOR)
        mol = Chem.MolFromSmiles(smiles)
        score = float(score)
        return ScoredMol(mol=mol, score=score)


class ScoredMolInMixin:
    Input = ScoredMol
    InputSerializer = ScoredMolSerializer


class ScoredMolOutMixin:
    Output = ScoredMol
    OutputSerializer = ScoredMolSerializer


class ScorerMixin:
    Input = Chem.Mol
    InputSerializer = MolToSmilesSerializer
    Output = ScoredMol
    OutputSerializer = ScoredMolSerializer


class ScoredSmiles(parameters.CompoundParam):
    smiles: str
    score: float = None

    def __str__(self):
        return f'{self.smiles} {self.score}'


class ScoredSmilesSerializer(stepper.Serializer):
    """
    A serializer using the smiles and score value strings separated by a ',' as
    the serialized form
    """
    DataType = ScoredSmiles
    SEPARATOR = ','

    def fromString(self, input_str):
        smiles, score = input_str.split(self.SEPARATOR)
        return ScoredSmiles(smiles=smiles, score=float(score))

    def toString(self, scored_smiles):
        return f'{scored_smiles.smiles}{self.SEPARATOR}{scored_smiles.score}'


class ScoredSmilesInMixin:
    Input = ScoredSmiles
    InputSerializer = ScoredSmilesSerializer


class ScoredSmilesOutMixin:
    Output = ScoredSmiles
    OutputSerializer = ScoredSmilesSerializer


class MaeInMixin:
    # the default serializer for structure objects will be used
    Input = structure.Structure


class MaeOutMixin:
    # the default serializer for structure objects will be used
    Output = structure.Structure


class MaeMaeMixin(MaeInMixin, MaeOutMixin):
    pass


class FileInMixin:
    Input = stepper.StepperFile


class ScoredSmilesIOMixin:
    Input = ScoredSmiles
    InputSerializer = ScoredSmilesSerializer
    Output = ScoredSmiles
    OutputSerializer = ScoredSmilesSerializer