diff --git a/chebai_graph/preprocessing/datasets/__init__.py b/chebai_graph/preprocessing/datasets/__init__.py index 9ee9856..c1e0a23 100644 --- a/chebai_graph/preprocessing/datasets/__init__.py +++ b/chebai_graph/preprocessing/datasets/__init__.py @@ -1,4 +1,5 @@ from .chebi import ( + ChEBI25GraphProperties, ChEBI50_Atom_WGNOnly_GraphProp, ChEBI50_GN_WithAllNodes_FG_WithAtoms_FGE, ChEBI50_GN_WithAllNodes_FG_WithAtoms_NoFGE, @@ -9,6 +10,7 @@ ChEBI50_StaticGNI, ChEBI50_WFGE_NGN_GraphProp, ChEBI50_WFGE_WGN_AsPerNodeType, + ChEBI50_WFGE_WGN_ForAllPropsExceptNodeLevel, ChEBI50_WFGE_WGN_ForAtomAndFGLevelOnly, ChEBI50_WFGE_WGN_ForAtomLevelAndGraphNodeOnly, ChEBI50_WFGE_WGN_ForAtomLevelOnly, @@ -18,7 +20,6 @@ ChEBI50_WFGE_WGN_ForNodeLevelPropOnlyForAllNodes, ChEBI50_WFGE_WGN_GraphProp, ChEBI50GraphData, - ChEBI25GraphProperties, ChEBI50GraphProperties, ChEBI100GraphProperties, ) @@ -60,6 +61,7 @@ "ChEBI50_WFGE_WGN_ForAtomLevelAndGraphNodeOnly", "ChEBI50_WFGE_WGN_ForFGLevelAndGraphNodeOnly", "ChEBI50_WFGE_WGN_ForNodeLevelPropOnlyForAllNodes", + "ChEBI50_WFGE_WGN_ForAllPropsExceptNodeLevel", "BACE_WFGE_WGN_AsPerNodeType", "BBBP_WFGE_WGN_AsPerNodeType", "ClinTox_WFGE_WGN_AsPerNodeType", diff --git a/chebai_graph/preprocessing/datasets/augmentation_base.py b/chebai_graph/preprocessing/datasets/augmentation_base.py index 9bf6cdc..20ce564 100644 --- a/chebai_graph/preprocessing/datasets/augmentation_base.py +++ b/chebai_graph/preprocessing/datasets/augmentation_base.py @@ -177,3 +177,10 @@ def _fill_node_tensor_with_molecule_type_property( is_graph_node: torch.Tensor, ) -> torch.Tensor: return node_tensor + + +class GraphPropAllPropsExceptNodeLevel(GraphPropAsPerNodeType): + def _fill_node_tensor_with_all_node_type_property( + self, node_tensor: torch.Tensor, property_values: torch.Tensor, offset: int + ) -> torch.Tensor: + return node_tensor diff --git a/chebai_graph/preprocessing/datasets/chebi.py b/chebai_graph/preprocessing/datasets/chebi.py index 1962c31..c4ac4ed 100644 --- a/chebai_graph/preprocessing/datasets/chebi.py +++ b/chebai_graph/preprocessing/datasets/chebi.py @@ -25,6 +25,7 @@ from .augmentation_base import ( AugGraphPropMixIn_NoGraphNode, AugGraphPropMixIn_WithGraphNode, + GraphPropAllPropsExceptNodeLevel, GraphPropForAtomAndFGLevelOnly, GraphPropForAtomLevelAndGraphNodeOnly, GraphPropForAtomLevelOnly, @@ -204,6 +205,12 @@ class ChEBI50_WFGE_WGN_ForNodeLevelPropOnlyForAllNodes( READER = AtomFGReader_WithFGEdges_WithGraphNode +class ChEBI50_WFGE_WGN_ForAllPropsExceptNodeLevel( + GraphPropAllPropsExceptNodeLevel, ChEBIOver50 +): + READER = AtomFGReader_WithFGEdges_WithGraphNode + + # ---------- Final Augmentation: Different Thresholds ------------------------------ class ChEBI100_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, ChEBIOver100): READER = AtomFGReader_WithFGEdges_WithGraphNode diff --git a/configs/data/augmented/properties/chebi50_all_props_except_node_lvl.yml b/configs/data/augmented/properties/chebi50_all_props_except_node_lvl.yml new file mode 100644 index 0000000..a03da69 --- /dev/null +++ b/configs/data/augmented/properties/chebi50_all_props_except_node_lvl.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_WFGE_WGN_ForAllPropsExceptNodeLevel +init_args: + properties: + # All Node type properties + - chebai_graph.preprocessing.properties.AtomNodeLevel + # Atom Node type properties + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + # FG Node type properties + - chebai_graph.preprocessing.properties.AtomFunctionalGroup + - chebai_graph.preprocessing.properties.IsHydrogenBondDonorFG + - chebai_graph.preprocessing.properties.IsHydrogenBondAcceptorFG + - chebai_graph.preprocessing.properties.IsFGAlkyl + # Graph Node type properties + - chebai_graph.preprocessing.properties.AugRDKit2DNormalized + # Bond properties + - chebai_graph.preprocessing.properties.BondLevel + - chebai_graph.preprocessing.properties.AugBondAromaticity + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondType