diff --git a/rdkit.py b/rdkit.py index 0d90cde..7646974 100644 --- a/rdkit.py +++ b/rdkit.py @@ -12,11 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +import logging from towhee import register from towhee.operator import Operator from rdkit import DataStructs, Chem +log = logging.getLogger() + @register(output_schema=['fingerprint']) class Rdkit(Operator): """ @@ -35,18 +38,22 @@ class Rdkit(Operator): def __call__(self, smiles: str): mols = Chem.MolFromSmiles(smiles) - if self.algorithm == 'daylight': - fp = Chem.RDKFingerprint(mols, fpSize=self.size) - elif self.algorithm == 'morgan': - from rdkit.Chem import AllChem - AllChem.GetMorganFingerprint - fp = AllChem.GetMorganFingerprintAsBitVect(mols, 2, self.size) - elif self.algorithm == 'ap': - from rdkit.Chem.AtomPairs import Pairs - fp = Pairs.GetAtomPairFingerprintAsBitVect(mols) - elif self.algorithm == 'maccs': - from rdkit.Chem import MACCSkeys - fp = MACCSkeys.GenMACCSKeys(mols) + try: + if self.algorithm == 'daylight': + fp = Chem.RDKFingerprint(mols, fpSize=self.size) + elif self.algorithm == 'morgan': + from rdkit.Chem import AllChem + AllChem.GetMorganFingerprint + fp = AllChem.GetMorganFingerprintAsBitVect(mols, 2, self.size) + elif self.algorithm == 'ap': + from rdkit.Chem.AtomPairs import Pairs + fp = Pairs.GetAtomPairFingerprintAsBitVect(mols) + elif self.algorithm == 'maccs': + from rdkit.Chem import MACCSkeys + fp = MACCSkeys.GenMACCSKeys(mols) + except Exception as e: + log.error(f'{e}, cannot generate fingerprint of {smiles}.') + raise KeyError(e) hex_fp = DataStructs.BitVectToFPSText(fp) fingerprint = bytes.fromhex(hex_fp)