Source code for schrodinger.application.steps.transformers

from rdkit import Chem

from schrodinger.infra import mmerr
from schrodinger.models.parameters import CompoundParam
from schrodinger.structutils.build import desalt_structure
from schrodinger.structutils.build import neutralize_structure
from schrodinger.stepper import ReduceStep
from schrodinger.utils import license

from . import utils
from .basesteps import MolMapStep
from .dataclasses import ScoredSmiles
from .dataclasses import ScoredSmilesIOMixin

STANDARDIZER_MMERR_LEVEL = mmerr.MMERR_OFF


class Standardizer(MolMapStep):
    """
    Creates a coordinate-less molecule that is neutralized and desalted.

    This involves a conversion to a `schrodinger.structure.Structure`.
    """

    def getLicenseRequirements(self):
        return {license.LIGPREP_MAIN: 1}

    def mapFunction(self, mol):
        st = utils.mol_to_structure(mol, self, generate_coordinates=False)
        if st is not None:
            with mmerr.Level(STANDARDIZER_MMERR_LEVEL):
                st = desalt_structure(neutralize_structure(st))
            mol = utils.structure_to_mol(st, self, mol)
            if mol is not None:
                yield mol


class ScoredSmilesStandardizer(ScoredSmilesIOMixin, ReduceStep):
    """
    Standardizes the SMILES.

    If more than one input has the same standardized SMILES `keep_lowest` in the
    settings determines whether the lowest value is retained or the highest.
    """

    class Settings(CompoundParam):
        keep_lowest: bool = True

    def getLicenseRequirements(self):
        return {license.LIGPREP_MAIN: 1}

    def reduceFunction(self, inputs):
        standardizer = Standardizer()
        scored_smiles = dict()
        for input in inputs:
            if mol := Chem.MolFromSmiles(input.smiles):
                try:
                    mol = next(standardizer.mapFunction(mol))
                except StopIteration:
                    continue
                smiles = Chem.MolToSmiles(mol)
                score = scored_smiles.get(smiles)
                if score is None:
                    score = input.score
                elif self.settings.keep_lowest:
                    score = min(input.score, score)
                else:
                    score = max(input.score, score)
                scored_smiles[smiles] = score

        for smiles, score in scored_smiles.items():
            yield ScoredSmiles(smiles=smiles, score=score)


class StereoChemistryRemover(MolMapStep):
    """
    Removes the stereo chemistry from a molecule.
    """

    def mapFunction(self, mol):
        Chem.RemoveStereochemistry(mol)
        yield mol