keep input_reaction for unsolvable cases

TieuLongPhan · Apr 15, 2024 · fc198de · fc198de
1 parent e11610c
commit fc198de
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/synrbl/main.py b/synrbl/main.py
@@ -59,7 +59,7 @@ def __run_pipeline(self, reactions, stats=None):
 
         logger.info("Run rule-based method.")
         self.rb_method.run(reactions, stats=stats)
-        self.rb_validator.check(reactions)
+        self.rb_validator.check(reactions, override_unsolved=True)
 
         self.mcs.find(reactions)
 
@@ -70,7 +70,7 @@ def __run_pipeline(self, reactions, stats=None):
             "Run rule-based method again to fix remaining non-carbon imbalance."
         )
         self.rb_method.run(reactions)
-        self.mcs_validator.check(reactions)
+        self.mcs_validator.check(reactions, override_unsolved=True)
 
         self.conf_predictor.predict(
             reactions, stats=stats, threshold=self.confidence_threshold

diff --git a/synrbl/postprocess.py b/synrbl/postprocess.py
@@ -23,7 +23,7 @@ def __init__(
         self.carbon_balance_col = carbon_balance_col
         self.n_jobs = n_jobs
 
-    def check(self, reactions):
+    def check(self, reactions, override_unsolved=False):
         update_reactants_and_products(reactions, self.reaction_col)
         decompose = RSMIDecomposer(
             smiles=None,  # type: ignore
@@ -61,5 +61,6 @@ def check(self, reactions):
             ):
                 reaction[self.solved_col] = True
                 reaction[self.solved_method_col] = self.method
-
+            if override_unsolved and not reaction[self.solved_col]:
+                reaction[self.reaction_col] = reaction["input_reaction"]
         return reactions