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 # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import logging
from towhee import register from towhee import register
from towhee.operator import Operator from towhee.operator import Operator
from rdkit import DataStructs, Chem from rdkit import DataStructs, Chem
log = logging.getLogger()
@register(output_schema=['fingerprint']) @register(output_schema=['fingerprint'])
class Rdkit(Operator): class Rdkit(Operator):
""" """
@ -35,18 +38,22 @@ class Rdkit(Operator):
def __call__(self, smiles: str): def __call__(self, smiles: str):
mols = Chem.MolFromSmiles(smiles) 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) hex_fp = DataStructs.BitVectToFPSText(fp)
fingerprint = bytes.fromhex(hex_fp) fingerprint = bytes.fromhex(hex_fp)

Loading…
Cancel
Save