logo
Browse Source

add logger

Signed-off-by: shiyu22 <shiyu.chen@zilliz.com>
main
shiyu22 2 years ago
parent
commit
0a5f781cfe
  1. 31
      rdkit.py

31
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)

Loading…
Cancel
Save