Merge pull request #427 from shihchengli/fix_reaction_atom_map

Fix reaction atom mapping
chemprop · Aug 1, 2023 · 69fd4ed · 69fd4ed
2 parents 1f367d5 + 5be03f4
commit 69fd4ed
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 7 deletions.
diff --git a/chemprop/data/scaffold.py b/chemprop/data/scaffold.py
@@ -3,6 +3,7 @@
 from random import Random
 from typing import Dict, List, Set, Tuple, Union
 import warnings
+import copy
 
 from rdkit import Chem
 from rdkit.Chem.Scaffolds import MurckoScaffold
@@ -23,7 +24,9 @@ def generate_scaffold(mol: Union[str, Chem.Mol, Tuple[Chem.Mol, Chem.Mol]], incl
     if isinstance(mol, str):
         mol = make_mol(mol, keep_h = False, add_h = False, keep_atom_map = False)
     if isinstance(mol, tuple):
-        mol = mol[0]
+        mol = copy.deepcopy(mol[0])
+        for atom in mol.GetAtoms():
+            atom.SetAtomMapNum(0)
     scaffold = MurckoScaffold.MurckoScaffoldSmiles(mol = mol, includeChirality = include_chirality)
 
     return scaffold

diff --git a/chemprop/features/featurization.py b/chemprop/features/featurization.py
@@ -135,7 +135,7 @@ def is_keeping_atom_map(is_mol: bool = True) -> bool:
     r"""Returns whether to keep the original atom mapping (not for reactions)"""
     if is_mol:
         return PARAMS.KEEP_ATOM_MAP
-    return False
+    return True
 
 
 def is_reaction(is_mol: bool = True) -> bool:
@@ -337,14 +337,15 @@ def __init__(self, mol: Union[str, Chem.Mol, Tuple[Chem.Mol, Chem.Mol]],
         self.is_reaction = is_reaction(self.is_mol)
         self.is_explicit_h = is_explicit_h(self.is_mol)
         self.is_adding_hs = is_adding_hs(self.is_mol)
+        self.is_keeping_atom_map = is_keeping_atom_map(self.is_mol)
         self.reaction_mode = reaction_mode()
 
         # Convert SMILES to RDKit molecule if necessary
         if type(mol) == str:
             if self.is_reaction:
-                mol = (make_mol(mol.split(">")[0], self.is_explicit_h, self.is_adding_hs, self.is_keeping_atom_map_list), make_mol(mol.split(">")[-1], self.is_explicit_h, self.is_adding_hs, self.is_keeping_atom_map_list)) 
+                mol = (make_mol(mol.split(">")[0], self.is_explicit_h, self.is_adding_hs, self.is_keeping_atom_map), make_mol(mol.split(">")[-1], self.is_explicit_h, self.is_adding_hs, self.is_keeping_atom_map)) 
             else:
-                mol = make_mol(mol, self.is_explicit_h, self.is_adding_hs, self.is_keeping_atom_map_list)
+                mol = make_mol(mol, self.is_explicit_h, self.is_adding_hs, self.is_keeping_atom_map)
 
         self.n_atoms = 0  # number of atoms
         self.n_bonds = 0  # number of bonds

diff --git a/chemprop/rdkit.py b/chemprop/rdkit.py
@@ -12,7 +12,7 @@ def make_mol(s: str, keep_h: bool, add_h: bool, keep_atom_map: bool):
     :return: RDKit molecule.
     """
     params = Chem.SmilesParserParams()
-    params.removeHs = not keep_h if not keep_atom_map else False
+    params.removeHs = not keep_h
     mol = Chem.MolFromSmiles(s, params)
 
     if add_h:

diff --git a/tests/test_integration.py b/tests/test_integration.py
@@ -817,7 +817,7 @@ def test_predict_spectra(self,
         (
                 'chemprop_scaffold_split',
                 'chemprop',
-                2.11470476,
+                2.095871,
                 ['--reaction', '--data_path', os.path.join(TEST_DATA_DIR, 'reaction_regression.csv'),'--split_type', 'scaffold_balanced']
         ),
         (
@@ -1032,7 +1032,7 @@ def test_multimolecule_fingerprint_with_single_input(self,
         (
                 'chemprop_reaction_solvent_diff_mpn_size',
                 'chemprop',
-                2.899513794,
+                2.734318,
                 ['--reaction_solvent', '--number_of_molecules', '2',
                  '--data_path', os.path.join(TEST_DATA_DIR, 'reaction_solvent_regression.csv'), '--hidden_size', '500',
                  '--hidden_size_solvent', '250']