App Domain Pathway Prediction (#47)

Co-authored-by: Tim Lorsbach <tim@lorsba.ch>
Reviewed-on: enviPath/enviPy#47

epdb/logic.py

@@ -895,9 +895,10 @@ class SearchManager(object):
 
 class SNode(object):
 
-    def __init__(self, smiles: str, depth: int):
+    def __init__(self, smiles: str, depth: int, app_domain_assessment: dict = None):
         self.smiles = smiles
         self.depth = depth
+        self.app_domain_assessment = app_domain_assessment
 
     def __hash__(self):
         return hash(self.smiles)
@@ -1040,7 +1041,7 @@ class SPathway(object):
     def depth(self):
         return max([v.depth for v in self.smiles_to_node.values()])
 
-    def _get_nodes_for_depth(self, depth: int):
+    def _get_nodes_for_depth(self, depth: int) -> List[SNode]:
         if depth == 0:
             return self.root_nodes
 
@@ -1051,7 +1052,7 @@ class SPathway(object):
 
         return sorted(res, key=lambda x: x.smiles)
 
-    def _get_edges_for_depth(self, depth: int):
+    def _get_edges_for_depth(self, depth: int) -> List[SEdge]:
         res = []
         for e in self.edges:
             for n in e.educts:
@@ -1076,15 +1077,44 @@ class SPathway(object):
         new_tp = False
         if substrates:
             for sub in substrates:
+
+                if sub.app_domain_assessment is None:
+                    if self.prediction_setting.model:
+                        if self.prediction_setting.model.app_domain:
+                            app_domain_assessment = self.prediction_setting.model.app_domain.assess(sub.smiles)[0]
+
+                            if self.persist is not None:
+                                n = self.snode_persist_lookup[sub]
+
+                                assert n.id is not None, "Node has no id! Should have been saved already... aborting!"
+                                node_data = n.simple_json()
+                                node_data['image'] = f"{n.url}?image=svg"
+                                app_domain_assessment['assessment']['node'] = node_data
+
+                                n.kv['app_domain_assessment'] = app_domain_assessment
+                                n.save()
+
+                            sub.app_domain_assessment = app_domain_assessment
+
+
                 candidates = self.prediction_setting.expand(self, sub)
+                # candidates is a List of PredictionResult. The length of the List is equal to the number of rules
                 for cand_set in candidates:
                     if cand_set:
                         new_tp = True
+                        # cand_set is a PredictionResult object that can consist of multiple candidate reactions
                         for cand in cand_set:
                             cand_nodes = []
+                            # candidate reactions can have multiple fragments
                             for c in cand:
                                 if c not in self.smiles_to_node:
-                                    self.smiles_to_node[c] = SNode(c, sub.depth + 1)
+                                    # For new nodes do an AppDomain Assessment if an AppDomain is attached
+                                    app_domain_assessment = None
+                                    if self.prediction_setting.model:
+                                        if self.prediction_setting.model.app_domain:
+                                            app_domain_assessment = self.prediction_setting.model.app_domain.assess(c)[0]
+
+                                    self.smiles_to_node[c] = SNode(c, sub.depth + 1, app_domain_assessment)
 
                                 node = self.smiles_to_node[c]
                                 cand_nodes.append(node)
@@ -1097,18 +1127,30 @@ class SPathway(object):
         if len(substrates) == 0 or from_node is not None:
             self.done = True
 
-        # Check if we need to write back data to database
+        # Check if we need to write back data to the database
         if new_tp and self.persist:
             self._sync_to_pathway()
-            # call save to update internal modified field
+            # call save to update the internal modified field
             self.persist.save()
 
-    def _sync_to_pathway(self):
+    def _sync_to_pathway(self) -> None:
         logger.info("Updating Pathway with SPathway")
 
         for snode in self.smiles_to_node.values():
             if snode not in self.snode_persist_lookup:
                 n = Node.create(self.persist, snode.smiles, snode.depth)
+
+                if snode.app_domain_assessment is not None:
+                    app_domain_assessment = snode.app_domain_assessment
+
+                    assert n.id is not None, "Node has no id! Should have been saved already... aborting!"
+                    node_data = n.simple_json()
+                    node_data['image'] = f"{n.url}?image=svg"
+                    app_domain_assessment['assessment']['node'] = node_data
+
+                    n.kv['app_domain_assessment'] = app_domain_assessment
+                    n.save()
+
                 self.snode_persist_lookup[snode] = n
 
         for sedge in self.edges:
@@ -1130,7 +1172,6 @@ class SPathway(object):
                 self.sedge_persist_lookup[sedge] = e
 
         logger.info("Update done!")
-        pass
 
     def to_json(self):
         nodes = []