is_alkyl fg prop

Author: aditya0by0

chebai_graph/preprocessing/properties/__init__.py

@@ -24,6 +24,9 @@
     AtomNodeLevel,
     AtomFunctionalGroup,
     AtomRingSize,
+    IsHydrogenBondDonorFG,
+    IsHydrogenBondAcceptorFG,
+    IsFGAlkyl,
     BondLevel,
     AugAtomType,
     AugNumAtomBonds,
@@ -59,6 +62,9 @@
     "AtomNodeLevel",
     "AtomFunctionalGroup",
     "AtomRingSize",
+    "IsHydrogenBondDonorFG",
+    "IsHydrogenBondAcceptorFG",
+    "IsFGAlkyl",
     "BondLevel",
     "AugAtomType",
     "AugNumAtomBonds",

chebai_graph/preprocessing/properties/augmented_properties.py

@@ -145,6 +145,14 @@ def get_atom_value(self, atom: Chem.rdchem.Atom | Dict):
         return fg in self._hydrogen_bond_acceptor
 
 
+class IsFGAlkyl(AugmentedAtomProperty):
+    def __init__(self, encoder: Optional[PropertyEncoder] = None):
+        super().__init__(encoder or BoolEncoder(self))
+
+    def get_atom_value(self, atom: Chem.rdchem.Atom | Dict):
+        return int(self._check_modify_atom_prop_value(atom, "is_alkyl"))
+
+
 class AugNodeValueDefaulter(AugmentedAtomProperty, FrozenPropertyAlias, ABC):
     def get_atom_value(self, atom: Chem.rdchem.Atom | Dict):
         if isinstance(atom, Chem.rdchem.Atom):

chebai_graph/preprocessing/reader/augmented_reader.py

@@ -1,3 +1,4 @@
+import re
 import sys
 from abc import ABC, abstractmethod
 from typing import Dict, List, Optional, Tuple
@@ -379,9 +380,10 @@ def _construct_fg_to_atom_structure(
         fg_to_atoms_map = {}
 
         molecule_atoms_set = set()
-        for _, fg_group in structure.items():
+        for fg_smiles, fg_group in structure.items():
             fg_to_atoms_map[self._num_of_nodes] = fg_group
             is_ring_fg = fg_group["is_ring_fg"]
+            is_alkyl = 0
 
             connected_atoms = []
             # Build edge index for fg to atom nodes connections
@@ -406,7 +408,7 @@ def _construct_fg_to_atom_structure(
             if is_ring_fg:
                 self._set_ring_fg_prop(connected_atoms, fg_nodes)
             else:
-                self._set_fg_prop(connected_atoms, fg_nodes)
+                self._set_fg_prop(connected_atoms, fg_nodes, fg_smiles)
 
             self._num_of_nodes += 1
 
@@ -419,6 +421,7 @@ def _set_ring_fg_prop(self, connected_atoms, fg_nodes):
             k.NODE_LEVEL: k.FG_NODE_LEVEL,
             "FG": f"RING_{ring_size}",
             "RING": ring_size,
+            "is_alkyl": "0",
         }
         # In this case, all atoms of Ring/Fused Ring are assigned the ring size as functional group
         for atom in connected_atoms:
@@ -429,8 +432,9 @@ def _set_ring_fg_prop(self, connected_atoms, fg_nodes):
             # An atom belonging to multiple rings in fused Ring has size "5-6", indicating size of each ring
             max_ring_size = max(list(map(int, ring_prop.split("-"))))
             atom.SetProp("FG", f"RING_{max_ring_size}")
+            atom.SetProp("is_alkyl", "0")
 
-    def _set_fg_prop(self, connected_atoms, fg_nodes):
+    def _set_fg_prop(self, connected_atoms, fg_nodes, fg_smiles):
         fg_set = {atom.GetProp("FG") for atom in connected_atoms}
         if not fg_set:
             raise ValueError(
@@ -458,17 +462,23 @@ def _set_fg_prop(self, connected_atoms, fg_nodes):
                 "All Connected atoms must belong to one functional group or None"
             )
 
-        # Select any one connected atom to get FG type and ring size
-        representative_atom = next(
-            (atom for atom in connected_atoms if atom.GetProp("FG")), None
-        )
+        check = re.sub(r"[CH\-\(\)\[\]/\\@]", "", fg_smiles)
+        is_alkyl = "1" if len(check) == 0 else "0"
+
+        representative_atom = None
+        for atom in connected_atoms:
+            if atom.GetProp("FG"):
+                representative_atom = atom
+            atom.SetProp("is_alkyl", is_alkyl)
+
         if representative_atom is None:
             raise AssertionError("Expected at least one atom with a functional group.")
 
         fg_nodes[self._num_of_nodes] = {
             k.NODE_LEVEL: k.FG_NODE_LEVEL,
             "FG": representative_atom.GetProp("FG"),
             "RING": 0,
+            "is_alkyl": is_alkyl,
         }
 
     def _construct_fg_level_structure(