cleaned up and commented

Author: Nora Khalil

rmgpy/data/kinetics/family.py

@@ -2994,8 +2994,7 @@ def get_extension_edge(self, parent, template_rxn_map, obj, T, iter_max=np.inf,
                     #parent
                     if typr != 'intNewBondExt' and typr != 'extNewBondExt':  # these dimensions should be regularized
                         if typr == 'atomExt':
-                            pass 
-                            #grp.atoms[indcr[0]].reg_dim_atm = list(reg_val)
+                            pass #no longer passing regularization info to the parent here. Doing this instead in `extend_node`
                         elif typr == 'elExt':
                             grp.atoms[indcr[0]].reg_dim_u = list(reg_val)
                         elif typr == 'ringExt':
@@ -3091,10 +3090,27 @@ def get_extension_edge(self, parent, template_rxn_map, obj, T, iter_max=np.inf,
         return out, gave_up_split
 
     def get_compliment_reg_dim(self, parent, template_rxn_map, new_ext, comp_ext):
+        """
+        Function takes in a parent node (`parent`), an extension node (`new_ext`) and its compliment (`comp_ext`). 
+        Reactions of the parent node are split to extension and compliment.
+        Iterating over all the reactions that fit the complimentary node, the atomtypes of each labeled atom in each reaction are saved to a dictionary `atom_labeling_in_comp_rxns`, 
+        where the key is the integer of the atom label (i.e. 5 in '*5') and the value is a set of all the atomtypes in all the complimentary reactions with that atom label.  
+
+        Additionally, when iterating over all the reactions that fit the complimentary node, the atomtypes of each unlabeled atom in each reaction are saved to a list `unlabeled_atoms_in_comp_rxns`. 
+        """
+        
+        
+        assert comp_ext is not None, "This extension does not include a complimentary node. Cannot get regularization dimensions of complimentary node."
+        
+        #divide parent reactions into the extension node and its compliment
         rxns_from_parent = template_rxn_map[parent.label]
         new_ext_rxns, comp_ext_rxns, _ = self._split_reactions(rxns_from_parent, new_ext)
+        
+        #for saving data
         atom_labeling_in_comp_rxns = dict()
         unlabeled_atoms_in_comp_rxns = []
+
+        #iterate through each complimentary rxn
         for rxn_c in comp_ext_rxns: 
             for reactant in rxn_c.reactants: 
                 for mol in reactant.molecule:
@@ -3105,13 +3121,13 @@ def get_compliment_reg_dim(self, parent, template_rxn_map, new_ext, comp_ext):
                             if unlabeled_atmtype not in unlabeled_atoms_in_comp_rxns:
                                 unlabeled_atoms_in_comp_rxns.append(unlabeled_atmtype)
                         else: 
+                            #this is a labeled atom
                             atm_label = int(atm.label.replace('*',''))
                             if atm_label not in atom_labeling_in_comp_rxns.keys():
                                 atom_labeling_in_comp_rxns[atm_label] = [ATOMTYPES[atm.symbol]]
                             else: 
                                 existing_atomtypes = atom_labeling_in_comp_rxns[atm_label]
                                 existing_atomtypes.append(ATOMTYPES[atm.symbol])
-                    #print(f'count of missing * is {count}')
         atom_labeling_in_comp_rxns_set = {k: set(v) for k, v in atom_labeling_in_comp_rxns.items()}
 
         return atom_labeling_in_comp_rxns_set, unlabeled_atoms_in_comp_rxns
@@ -3191,45 +3207,35 @@ def extend_node(self, parent, template_rxn_map, obj=None, T=1000.0, iter_max=np.
 
         extname = ext[2]
 
-        print(extname, ext[3])
+
         if ext[3] == 'atomExt':
             ext[0].atoms[ext[4][0]].reg_dim_atm = [ext[0].atoms[ext[4][0]].atomtype, ext[0].atoms[ext[4][0]].atomtype] #passing regularization information to the selected extension node
 
-            #handling regularization in complement below
+            #handling regularization in complement below:
             atom_labeling_in_comp_rxns_set, unlabeled_atoms_in_comp_rxns = self.get_compliment_reg_dim(parent, template_rxn_map, ext[0], ext[1])
-            #print(ext[0].atoms[ext[4][0]], ext[0].atoms[ext[4][0]].label, ext[1].atoms[ext[4][0]], ext[1].atoms[ext[4][0]].label)
 
             #regularize the atom in which the extension was performed on
-            if ext[1] is not None: 
-                if ext[1].atoms[ext[4][0]].label=='':
-                    #extension was performed on an unlabeled atom
-                    limited_atomtypes_comp = set(ext[1].atoms[ext[4][0]].atomtype).intersection(set(unlabeled_atoms_in_comp_rxns))
-                    #print(ext[1].atoms[ext[4][0]].atomtype, list(limited_atomtypes_comp))
-                    ext[1].atoms[ext[4][0]].reg_dim_atm = [ext[1].atoms[ext[4][0]].atomtype, list(limited_atomtypes_comp)]
-                else: 
-                    adjusted_index = int(ext[1].atoms[ext[4][0]].label.replace('*','')) #i.e. ext[4]= (3,), ext[4][0] = 3, ext[0].atoms[3]=<GroupAtom [*5 'N', 'C']>, ext[0].atoms[3].label = '*5'
-                    ext[1].atoms[ext[4][0]].reg_dim_atm = [ext[1].atoms[ext[4][0]].atomtype, list(atom_labeling_in_comp_rxns_set[adjusted_index])]
-            
-        #make sure the rest of the atoms in the extension take on the same regularization dimensions as the parent. 
+            if ext[1].atoms[ext[4][0]].label=='':
+                #extension was performed on an unlabeled atom, so pass in regularization dimensions that are at least limited to the atomtypes of all the unlabeled atoms
+                limited_atomtypes_comp = set(ext[1].atoms[ext[4][0]].atomtype).intersection(set(unlabeled_atoms_in_comp_rxns))
+                ext[1].atoms[ext[4][0]].reg_dim_atm = [ext[1].atoms[ext[4][0]].atomtype, list(limited_atomtypes_comp)]
+            else: 
+                #extension was performed on a labeled atom. For each labeled atom, we know all the atomtypes in the training reactions. Let's limit regularization dimensions to these known atomtypes
+                adjusted_index = int(ext[1].atoms[ext[4][0]].label.replace('*','')) #i.e. ext[4]= (3,), ext[4][0] = 3, ext[0].atoms[3]=<GroupAtom [*5 'N', 'C']>, ext[0].atoms[3].label = '*5'
+                ext[1].atoms[ext[4][0]].reg_dim_atm = [ext[1].atoms[ext[4][0]].atomtype, list(atom_labeling_in_comp_rxns_set[adjusted_index])]
+        
+        #make sure the rest of the atoms in the extension take on the same regularization dimensions as the parent. Ensures subgraph isomorphism. 
         for i, parent_atm in enumerate(parent.item.atoms): 
             if i == ext[4][0]: 
-                print('extension atom')
-                continue #this is the atom that the extension is focused on, handled above
+                continue #this is the atom that the extension is focused on, handled above if the extension was an 'atomExt' extension type
             elif parent_atm.reg_dim_atm[1]==[]:
-                print('parent atm reg_dim is empty')
                 continue #only take on regularization dimensions of parent if there is some
             else: 
                 ext[0].atoms[i].reg_dim_atm[1] = parent_atm.reg_dim_atm[1] #passing regularization info from parent to the extension
-                if ext[1] is not None: 
+                if ext[1] is not None: #check if there's a complimentary node
                     ext[1].atoms[i].reg_dim_atm[1] = parent_atm.reg_dim_atm[1] #passing regularization info from parent to the complimentary extension
 
-                
-                #print(ext[1].atoms[i].atomtype,'   ', ext[1].atoms[i].reg_dim_atm[1])
-                
 
-
-            # print(ext[1].atoms[ext[4][0]].atomtype, )
-            # ext[1].atoms[ext[4][0]].reg_dim_atm = [ext[1].atoms[ext[4][0]].atomtype, ext[1].atoms[ext[4][0]].atomtype] #must also pass regularization information to the compliment
         if ext[3] == 'elExt':
             ext[0].atoms[ext[4][0]].reg_dim_u = [ext[0].atoms[ext[4][0]].radical_electrons,
                                                  ext[0].atoms[ext[4][0]].radical_electrons]
@@ -3316,8 +3322,6 @@ def extend_node(self, parent, template_rxn_map, obj=None, T=1000.0, iter_max=np.
         if complement:
             template_rxn_map[parent.label] = []
             template_rxn_map[cextname] = comp_entries
-            if cextname=="Root_N-4R!H->O":
-                print(f'end of extend_node: {self.groups.entries["Root_N-4R!H->O"].item.atoms[3].reg_dim_atm}')
         else:
             template_rxn_map[parent.label] = comp_entries
 
@@ -3389,8 +3393,6 @@ def rxnkey(rxn):
                 logging.error("built tree with {} nodes".format(len(list(self.groups.entries))))
 
             self.auto_generated = True
-        print(f'end of generate_tree: {self.groups.entries["Root_N-4R!H->O"].item.atoms[3].reg_dim_atm}')
-
 
     def get_rxn_batches(self, rxns, T=1000.0, max_batch_size=800, outlier_fraction=0.02, stratum_num=8):
         """
@@ -3558,8 +3560,6 @@ def make_tree_nodes(self, template_rxn_map=None, obj=None, T=1000.0, nprocs=0, d
                         continue
                     boo2 = self.extend_node(entry, template_rxn_map, obj, T, iter_max=extension_iter_max, iter_item_cap=extension_iter_item_cap)
                     if boo2:  # extended node so restart while loop
-                        # if "Root_N-4R!H->O" in template_rxn_map.keys(): 
-                        #     print(f'at boo2: {self.groups.entries["Root_N-4R!H->O"].item.atoms[3].reg_dim_atm}')
                         break
                     else:  # no extensions could be generated since all reactions were identical
                         mult_completed_nodes.append(entry)
@@ -3591,8 +3591,6 @@ def make_tree_nodes(self, template_rxn_map=None, obj=None, T=1000.0, nprocs=0, d
                 entry.parent = self.groups.entries[pname]
                 entry.parent.children.append(entry)
 
-        print(f'end of make_tree_nodes: {self.groups.entries["Root_N-4R!H->O"].item.atoms[3].reg_dim_atm}')
-
         return
 
     def _absorb_process(self, p, conn, name):