fixed size_threshold check

ReactionMechanismGenerator · May 14, 2024 · f30cce0 · f30cce0
1 parent ee6971e
commit f30cce0
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 20 deletions.
diff --git a/rmgpy/molecule/fragment.py b/rmgpy/molecule/fragment.py
@@ -945,8 +945,7 @@ def sliceitup_arom(self, molecule, size_threshold=5):
                     # mol_set contains new set of fragments
                     mol_set = Chem.GetMolFrags(new_mol, asMols=True)
                     # check all fragments' size
-                    frag_sizes = [sum(1 for atom in mol.GetAtoms() if atom.GetAtomicNum() == 6) for mol in mol_set]
-                    if all(x >= size_threshold for x in frag_sizes):
+                    if all(sum(1 for atom in mol.GetAtoms() if atom.GetAtomicNum() == 6) >= size_threshold for mol in mol_set):
                         if len(mol_set) == 2:
                             frag1 = Chem.MolToSmiles(mol_set[0])
                             frag2 = Chem.MolToSmiles(mol_set[1])
@@ -1090,27 +1089,47 @@ def sliceitup_aliph(self, molecule, size_threshold=5):
                     # mol_set contains new set of fragments
                     mol_set = Chem.GetMolFrags(new_mol, asMols=True)
                     # check all fragments' size
-                    if all(
-                        sum(1 for atom in mol.GetAtoms() if atom.GetAtomicNum() == 6)
-                        >= size_threshold
-                        for mol in mol_set
-                    ):
-                        # replace * at cutting position with cutting label
-                        for ind, rdmol in enumerate(mol_set):
-                            frag = Chem.MolToSmiles(rdmol)
-                            if len(mol_set) > 2:  # means it cut into 3 fragments
+                    if all(sum(1 for atom in mol.GetAtoms() if atom.GetAtomicNum() == 6) >= size_threshold for mol in mol_set):
+                        if len(mol_set) == 2:
+                            frag1 = Chem.MolToSmiles(mol_set[0])
+                            frag2 = Chem.MolToSmiles(mol_set[1])
+
+                            frag1_R = frag1.count("Na")
+                            frag1_L = frag1.count("K")
+                            frag2_R = frag2.count("Na")
+                            frag2_L = frag2.count("K")
+
+                            if frag1_R > frag2_R and frag1_L <= frag2_L:
+                                frag1_smi = frag1.replace("*", "L")
+                                frag2_smi = frag2.replace("*", "R")
+                            elif frag1_L > frag2_L and frag1_R <= frag2_R:
+                                frag1_smi = frag1.replace("*", "R")
+                                frag2_smi = frag2.replace("*", "L")
+                            elif frag2_L > frag1_L and frag2_R <= frag1_R:
+                                frag1_smi = frag1.replace("*", "R")
+                                frag2_smi = frag2.replace("*", "L")
+                            elif frag2_R > frag1_R and frag2_L <= frag1_L:
+                                frag1_smi = frag1.replace("*", "R")
+                                frag2_smi = frag2.replace("*", "L")
+                            elif randint(0,1)==1:
+                                frag1_smi = frag1.replace("*", "L")
+                                frag2_smi = frag2.replace("*", "R")
+                            else:
+                                frag1_smi = frag1.replace("*", "R")
+                                frag2_smi = frag2.replace("*", "L")
+
+                            frag_list = [frag1_smi, frag2_smi]
+
+                        elif len(mol_set) > 2: # means it cut into 3 fragments
+                            frag_list = []
+                            for ind, rdmol in enumerate(mol_set):
+                                frag = Chem.MolToSmiles(rdmol)
                                 if frag.count("*") > 1:
-                                    # replace both with R
-                                    frag_smi = frag.replace("*", "R")
-                                else:
-                                    frag_smi = frag.replace("*", "L")
-                            else:  # means it only cut once, generate 2 fragments
-                                if ind == 0:
                                     frag_smi = frag.replace("*", "R")
                                 else:
                                     frag_smi = frag.replace("*", "L")
-                            frag_list.append(frag_smi)
-                        break
+                                frag_list.append(frag_smi)                          
+                            break
                     else:
                         # turn to next matched_atom_map
                         continue

diff --git a/rmgpy/rmg/model.py b/rmgpy/rmg/model.py
@@ -59,7 +59,7 @@
 from rmgpy.rmg.decay import decay_species
 from rmgpy.rmg.reactors import PhaseSystem, Phase, Interface, Reactor
 from rmgpy.molecule.fragment import Fragment
-
+from rmgpy.molecule.molecule import Molecule
 ################################################################################