debug atom mapping from rdkit mol to smiles

ps/mol_bpe.py

@@ -235,7 +235,7 @@ def tokenize(self, mol):
                 if i != j:
                     ad_mat[i][j] = ad_mat[j][i] = 1
         group_idxs = [x[1] for x in res]
-        return Molecule(smiles, group_idxs, self.kekulize)
+        return Molecule(rdkit_mol, group_idxs, self.kekulize)
 
     def idx_to_subgraph(self, idx):
         return self.idx2subgraph[idx]

ps/molecule.py

@@ -1,10 +1,12 @@
 #!/usr/bin/python
 # -*- coding:utf-8 -*-
 from copy import copy, deepcopy
+from typing import Union
 
 import networkx as nx
 from rdkit import Chem
 from rdkit.Chem.Draw import rdMolDraw2D
+from rdkit.Chem.rdchem import Mol as RDKitMol
 import numpy as np
 
 from utils.chem_utils import smi2mol, mol2smi
@@ -65,12 +67,16 @@ def __str__(self):
 class Molecule(nx.Graph):
     '''molecule represented in subgraph-level'''
 
-    def __init__(self, smiles: str=None, groups: list=None, kekulize: bool=False):
+    def __init__(self, mol: Union[str, RDKitMol]=None, groups: list=None, kekulize: bool=False):
         super().__init__()
-        if smiles is None:
+        if mol is None:
             return
+
+        if isinstance(mol, str):
+            smiles, rdkit_mol = mol, smi2mol(mol, kekulize)
+        else:
+            smiles, rdkit_mol = mol2smi(mol), mol
         self.graph['smiles'] = smiles
-        rdkit_mol = smi2mol(smiles, kekulize)
         # processing atoms
         aid2pos = {}
         for pos, group in enumerate(groups):

ps/utils/chem_utils.py

@@ -11,8 +11,8 @@ def smi2mol(smiles: str, kekulize=False, sanitize=True):
     return mol
 
 
-def mol2smi(mol):
-    return Chem.MolToSmiles(mol)
+def mol2smi(mol, canonical=True):
+    return Chem.MolToSmiles(mol, canonical=canonical)
 
 
 def get_submol(mol, atom_indices, kekulize=False):