new chebi data class for new readers

Author: aditya0by0

chebai_graph/preprocessing/datasets/__init__.py

@@ -1,5 +1,9 @@
 from .chebi import (
     ChEBI50_Atom_WGNOnly_GraphProp,
+    ChEBI50_GN_WithAllNodes_FG_WithAtoms_FGE,
+    ChEBI50_GN_WithAllNodes_FG_WithAtoms_NoFGE,
+    ChEBI50_GN_WithAtoms_FG_WithAtoms_FGE,
+    ChEBI50_GN_WithAtoms_FG_WithAtoms_NoFGE,
     ChEBI50_NFGE_NGN_GraphProp,
     ChEBI50_NFGE_WGN_GraphProp,
     ChEBI50_StaticGNI,
@@ -23,4 +27,8 @@
     "ChEBI50_WFGE_WGN_GraphProp",
     "ChEBI50_StaticGNI",
     "ChEBI50_WFGE_WGN_AsPerNodeType",
+    "ChEBI50_GN_WithAllNodes_FG_WithAtoms_FGE",
+    "ChEBI50_GN_WithAllNodes_FG_WithAtoms_NoFGE",
+    "ChEBI50_GN_WithAtoms_FG_WithAtoms_FGE",
+    "ChEBI50_GN_WithAtoms_FG_WithAtoms_NoFGE",
 ]

chebai_graph/preprocessing/datasets/chebi.py

@@ -30,6 +30,10 @@
     AtomFGReader_WithFGEdges_WithGraphNode,
     AtomReader_WithGraphNodeOnly,
     AtomsFGReader_NoFGEdges_NoGraphNode,
+    GN_WithAllNodes_FG_WithAtoms_FGE,
+    GN_WithAllNodes_FG_WithAtoms_NoFGE,
+    GN_WithAtoms_FG_WithAtoms_FGE,
+    GN_WithAtoms_FG_WithAtoms_NoFGE,
     GraphPropertyReader,
     GraphReader,
     RandomNodeInitializationReader,
@@ -595,6 +599,50 @@ class ChEBI50_WFGE_WGN_GraphProp(AugGraphPropMixIn_WithGraphNode, ChEBIOver50):
     READER = AtomFGReader_WithFGEdges_WithGraphNode
 
 
+class ChEBI50_GN_WithAllNodes_FG_WithAtoms_FGE(
+    AugGraphPropMixIn_WithGraphNode, ChEBIOver50
+):
+    """
+    ChEBIOver50 with FG nodes (connected to their respective atom nodes) with functional group
+    edges, and adds a graph-level node connected to all nodes (fg + atoms).
+    """
+
+    READER = GN_WithAllNodes_FG_WithAtoms_FGE
+
+
+class ChEBI50_GN_WithAllNodes_FG_WithAtoms_NoFGE(
+    AugGraphPropMixIn_WithGraphNode, ChEBIOver50
+):
+    """
+    ChEBIOver50 with FG nodes (connected to their respective atom nodes) without functional group
+    edges, and adds a graph-level node connected to all nodes (fg + atoms).
+    """
+
+    READER = GN_WithAllNodes_FG_WithAtoms_NoFGE
+
+
+class ChEBI50_GN_WithAtoms_FG_WithAtoms_FGE(
+    AugGraphPropMixIn_WithGraphNode, ChEBIOver50
+):
+    """
+    ChEBIOver50 with FG nodes (connected to their respective atom nodes) with functional group
+    edges, and adds a graph-level node connected to all atom nodes.
+    """
+
+    READER = GN_WithAtoms_FG_WithAtoms_FGE
+
+
+class ChEBI50_GN_WithAtoms_FG_WithAtoms_NoFGE(
+    AugGraphPropMixIn_WithGraphNode, ChEBIOver50
+):
+    """
+    ChEBIOver50 with FG nodes (connected to their respective atom nodes) without functional group
+    edges, and adds a graph-level node connected to all atom nodes.
+    """
+
+    READER = GN_WithAtoms_FG_WithAtoms_NoFGE
+
+
 class ChEBI50_NFGE_WGN_GraphProp(AugGraphPropMixIn_WithGraphNode, ChEBIOver50):
     """ChEBIOver50 with FG nodes but without FG edges, with graph node."""