From 9abbe9fe76c245ff9352902e47b408a719c0ac4b Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Sat, 8 Aug 2026 11:23:01 +0200 Subject: [PATCH 01/58] chebi aug variant with all props except node level --- .../preprocessing/datasets/__init__.py | 4 +++- .../datasets/augmentation_base.py | 7 ++++++ chebai_graph/preprocessing/datasets/chebi.py | 7 ++++++ .../chebi50_all_props_except_node_lvl.yml | 24 +++++++++++++++++++ 4 files changed, 41 insertions(+), 1 deletion(-) create mode 100644 configs/data/augmented/properties/chebi50_all_props_except_node_lvl.yml 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 From 05462dacf410bd216416aff568507cb100576a57 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Mon, 10 Aug 2026 20:42:36 +0200 Subject: [PATCH 02/58] allow adding new tokens through augmented class for atom specific props --- chebai_graph/preprocessing/properties/base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/chebai_graph/preprocessing/properties/base.py b/chebai_graph/preprocessing/properties/base.py index f148df8..1f2e0aa 100644 --- a/chebai_graph/preprocessing/properties/base.py +++ b/chebai_graph/preprocessing/properties/base.py @@ -209,8 +209,8 @@ def on_finish(self) -> None: hasattr(self.encoder, "cache") and len(self.encoder.cache) > self.encoder.index_length_start ): - raise ValueError( - f"{self.__class__.__name__} attempted to add new tokens " + print( + f"{self.__class__.__name__} added new tokens " f"to a frozen encoder at {self.encoder.index_path}" ) super().on_finish() From d26b9b9e04179b0ce851c75f74c47ec6e221d7f0 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Mon, 10 Aug 2026 20:42:52 +0200 Subject: [PATCH 03/58] BACE dataset tokens --- .../bin/bace/AtomCharge/indices_one_hot.txt | 3 ++ .../AtomFunctionalGroup/indices_one_hot.txt | 44 +++++++++++++++++++ .../AtomHybridization/indices_one_hot.txt | 3 ++ .../bace/AtomNodeLevel/indices_one_hot.txt | 3 ++ .../bin/bace/AtomNumHs/indices_one_hot.txt | 4 ++ .../bin/bace/AtomType/indices_one_hot.txt | 8 ++++ .../bin/bace/BondLevel/indices_one_hot.txt | 4 ++ .../bin/bace/BondType/indices_one_hot.txt | 4 ++ .../bin/bace/NumAtomBonds/indices_one_hot.txt | 4 ++ 9 files changed, 77 insertions(+) create mode 100644 chebai_graph/preprocessing/bin/bace/AtomCharge/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bace/AtomFunctionalGroup/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bace/AtomHybridization/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bace/AtomNodeLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bace/AtomNumHs/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bace/AtomType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bace/BondLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bace/BondType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bace/NumAtomBonds/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/bace/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..d7e52d1 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/AtomCharge/indices_one_hot.txt @@ -0,0 +1,3 @@ +0 +1 +-1 diff --git a/chebai_graph/preprocessing/bin/bace/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..b621c2a --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,44 @@ +RING_6 +difluoromethyl +RING_5 +ether +alkyne +alkyl +primary_amine +fluoro +amide +graph_fg +NO_FG +alkene +tertiary_carbon +alkene_carbon +nitrile +trifluoromethyl +ketone +chloro +quaternary_carbon +sulfonyl +hydroxyl +bromo +secondary_amine +ester +aldoxime +sulfide +iodo +tertiary_amine +RING_4 +secondary_ketimine +ketoxime +RING_3 +carboxylate +acetal +sulfinyl +RING_7 +RING_16 +RING_15 +nitro +primary_ketimine +sulfonate_ester +RING_8 +RING_14 +RING_13 diff --git a/chebai_graph/preprocessing/bin/bace/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..7c99da3 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,3 @@ +SP2 +SP3 +SP diff --git a/chebai_graph/preprocessing/bin/bace/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/bace/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..c4b216a --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,4 @@ +0 +1 +3 +2 diff --git a/chebai_graph/preprocessing/bin/bace/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..6da4980 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/AtomType/indices_one_hot.txt @@ -0,0 +1,8 @@ +6 +7 +8 +9 +17 +16 +35 +53 diff --git a/chebai_graph/preprocessing/bin/bace/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/bace/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/BondType/indices_one_hot.txt new file mode 100644 index 0000000..f1a2ce5 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/BondType/indices_one_hot.txt @@ -0,0 +1,4 @@ +SINGLE +AROMATIC +TRIPLE +DOUBLE diff --git a/chebai_graph/preprocessing/bin/bace/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..cf286d6 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,4 @@ +3 +2 +4 +1 From 1ce48efaba747c6e84938c5a78ade85800ebae46 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Tue, 11 Aug 2026 23:59:57 +0200 Subject: [PATCH 04/58] check len of base_df after merge, to avoid memory blow due to duplicated entries leading to cross merge --- chebai_graph/preprocessing/datasets/base.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/chebai_graph/preprocessing/datasets/base.py b/chebai_graph/preprocessing/datasets/base.py index 4e8e6c3..84d42be 100644 --- a/chebai_graph/preprocessing/datasets/base.py +++ b/chebai_graph/preprocessing/datasets/base.py @@ -462,6 +462,7 @@ def load_processed_data( base_data = super().load_processed_data(kind, filename) base_df = pd.DataFrame(base_data) base_df["ident"] = base_df["ident"].astype(str) + len_base_df_before_merge = len(base_df) props_categories = { "AllNodeTypeProperties": [], "FGNodeTypeProperties": [], @@ -524,11 +525,21 @@ def load_processed_data( ) property_df = pd.DataFrame(property_data) + assert len(property_df) == len(property_df["ident"].unique()), ( + f"Duplicate entries found in property {property.name} data. " + f"Unique entries {len(property_df['ident'].unique())}. " + f"Total entries {len(property_df)}. " + ) property_df["ident"] = property_df["ident"].astype(str) property_df.rename( columns={property.name: f"{property.name}"}, inplace=True ) + base_df = base_df.merge(property_df, on="ident", how="left") + assert len(base_df) == len_base_df_before_merge, ( + f"Length mismatch after merging property {property.name}. " + f"Expected {len_base_df_before_merge}, got {len(base_df)}" + ) base_df["features"] = base_df.apply( lambda row: self._merge_props_into_base( From f590aebb32fcc2a7549680f41896fcc32588ca2e Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Wed, 12 Aug 2026 13:56:52 +0200 Subject: [PATCH 05/58] BBBP tokens --- .../bin/bbbp/AtomCharge/indices_one_hot.txt | 4 + .../AtomFunctionalGroup/indices_one_hot.txt | 80 +++++++++++++++++++ .../AtomHybridization/indices_one_hot.txt | 4 + .../bbbp/AtomNodeLevel/indices_one_hot.txt | 3 + .../bin/bbbp/AtomNumHs/indices_one_hot.txt | 4 + .../bin/bbbp/AtomType/indices_one_hot.txt | 13 +++ .../bin/bbbp/BondLevel/indices_one_hot.txt | 4 + .../bin/bbbp/BondType/indices_one_hot.txt | 4 + .../bin/bbbp/NumAtomBonds/indices_one_hot.txt | 5 ++ 9 files changed, 121 insertions(+) create mode 100644 chebai_graph/preprocessing/bin/bbbp/AtomCharge/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bbbp/AtomFunctionalGroup/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bbbp/AtomHybridization/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bbbp/AtomNodeLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bbbp/AtomNumHs/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bbbp/AtomType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bbbp/BondLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bbbp/BondType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bbbp/NumAtomBonds/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..7b6a00a --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/AtomCharge/indices_one_hot.txt @@ -0,0 +1,4 @@ +0 +1 +-1 +2 diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..d2ee211 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,80 @@ +alkyl +ester +chloro +RING_6 +tertiary_amine +graph_fg +tertiary_carbon +primary_amine +dichloromethyl +sulfonyl +ketone +secondary_amine +hydroxyl +ether +quaternary_carbon +amide +NO_FG +carboxyl +nitro +alkene +alkene_carbon +secondary_ketimine +secondary_aldimine +carboxylate +iodo +aldoxime +bromo +alkyne +fluoro +carbamate +trifluoromethyl +sulfide +trichloromethyl +hemiacetal +primary_ketimine +azide +4_ammonium_ion +sulfonate_ester +ketoxime +nitroso +sulfhydryl +difluoromethyl +fluorochloromethyl +chlorobromomethyl +phosphono +carbodithio +disulfide +difluorochloromethyl +carbonate_ester +phosphate +RING_5 +thiolester +RING_7 +nitrile +phosphoryl +aldehyde +sulfinyl +thioketone +acetal +RING_4 +sulfonic_acid +imide +RING_14 +RING_24 +amidine +RING_3 +RING_9 +RING_19 +RING_16 +RING_8 +isonitrile +RING_23 +RING_12 +RING_21 +RING_17 +RING_30 +RING_33 +RING_36 +azo +RING_15 diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..0e4a9aa --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,4 @@ +SP3 +SP2 +S +SP diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..35f5032 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,4 @@ +2 +0 +1 +3 diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..9500570 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/AtomType/indices_one_hot.txt @@ -0,0 +1,13 @@ +6 +17 +8 +7 +16 +11 +53 +35 +9 +1 +15 +5 +20 diff --git a/chebai_graph/preprocessing/bin/bbbp/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/bbbp/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/BondType/indices_one_hot.txt new file mode 100644 index 0000000..a134437 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/BondType/indices_one_hot.txt @@ -0,0 +1,4 @@ +DOUBLE +SINGLE +AROMATIC +TRIPLE diff --git a/chebai_graph/preprocessing/bin/bbbp/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..5d07eee --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,5 @@ +2 +4 +1 +3 +0 From f4cb1ace533ddbbebcb2a40d046c59bb4317ef37 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Wed, 12 Aug 2026 14:27:20 +0200 Subject: [PATCH 06/58] hiv tokens --- .../bin/hiv/AtomCharge/indices_one_hot.txt | 9 ++ .../AtomFunctionalGroup/indices_one_hot.txt | 122 ++++++++++++++++++ .../hiv/AtomHybridization/indices_one_hot.txt | 7 + .../bin/hiv/AtomNodeLevel/indices_one_hot.txt | 3 + .../bin/hiv/AtomNumHs/indices_one_hot.txt | 5 + .../bin/hiv/AtomType/indices_one_hot.txt | 54 ++++++++ .../bin/hiv/BondLevel/indices_one_hot.txt | 0 .../bin/hiv/BondType/indices_one_hot.txt | 0 .../bin/hiv/NumAtomBonds/indices_one_hot.txt | 0 9 files changed, 200 insertions(+) create mode 100644 chebai_graph/preprocessing/bin/hiv/AtomCharge/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/hiv/AtomFunctionalGroup/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/hiv/AtomHybridization/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/hiv/AtomNodeLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/hiv/AtomNumHs/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/hiv/AtomType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/hiv/BondLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/hiv/BondType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/hiv/NumAtomBonds/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/hiv/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..22aacd7 --- /dev/null +++ b/chebai_graph/preprocessing/bin/hiv/AtomCharge/indices_one_hot.txt @@ -0,0 +1,9 @@ +0 +1 +-1 +3 +-4 +-3 +-2 +2 +-5 diff --git a/chebai_graph/preprocessing/bin/hiv/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..e5b33d4 --- /dev/null +++ b/chebai_graph/preprocessing/bin/hiv/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,122 @@ +alkyl +RING_6 +chloro +NO_FG +secondary_amine +graph_fg +carboxyl +primary_amine +sulfonic_acid +nitro +ketone +hydroxyl +sulfonyl +ester +nitroso +sulfhydryl +tertiary_carbon +tertiary_amine +dichloromethyl +nitrile +primary_ketimine +alkyne +alkene +ether +iodo +sulfide +secondary_aldimine +phosphoryl +thioketone +alkene_carbon +secondary_ketimine +nitrate +4_ammonium_ion +ketoxime +amide +amidine +trichloromethyl +fluoro +quaternary_carbon +hydrazone +disulfide +thiocyanate +carbamate +trifluoromethyl +phosphate +sulfonate_ester +sulfinyl +azo +bromo +imide +carbodithioic_acid +carbodithio +sulfino +silyl_ether +aldehyde +phosphono +azide +aldoxime +phosphino +trimethylsilyl +thionoester +thiolester +isonitrile +acetal +difluorochloromethyl +dibromomethyl +ketal +diiodomethyl +borono +isothiocyanate +difluoromethyl +tribromomethyl +carbonate_ester +hemiacetal +phosphodiester +carboxylate +orthoester +isocyanate +primary_aldimine +hemiketal +borinate +RING_5 +cyanate +RING_4 +RING_24 +RING_3 +thial +RING_7 +RING_12 +RING_9 +RING_16 +hydroperoxy +RING_15 +RING_18 +RING_8 +RING_35 +RING_10 +RING_19 +RING_17 +RING_29 +RING_23 +RING_27 +RING_11 +RING_14 +RING_20 +RING_33 +RING_13 +RING_31 +RING_21 +RING_28 +RING_22 +RING_30 +RING_34 +RING_54 +RING_38 +RING_36 +RING_26 +RING_32 +carboxylic_anhydride +RING_25 +carbothioic_S-acid +RING_40 diff --git a/chebai_graph/preprocessing/bin/hiv/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..cdc6985 --- /dev/null +++ b/chebai_graph/preprocessing/bin/hiv/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,7 @@ +SP3 +SP2 +SP +S +SP3D2 +SP3D +UNSPECIFIED diff --git a/chebai_graph/preprocessing/bin/hiv/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/hiv/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/hiv/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..df4a3e9 --- /dev/null +++ b/chebai_graph/preprocessing/bin/hiv/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,5 @@ +3 +1 +0 +2 +4 diff --git a/chebai_graph/preprocessing/bin/hiv/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..95c9cba --- /dev/null +++ b/chebai_graph/preprocessing/bin/hiv/AtomType/indices_one_hot.txt @@ -0,0 +1,54 @@ +6 +17 +8 +7 +15 +16 +33 +53 +9 +35 +14 +11 +34 +79 +50 +80 +32 +51 +82 +5 +19 +3 +83 +29 +67 +44 +65 +74 +27 +55 +25 +26 +75 +42 +78 +47 +24 +46 +89 +20 +31 +12 +30 +28 +77 +1 +40 +81 +45 +22 +52 +64 +92 +23 diff --git a/chebai_graph/preprocessing/bin/hiv/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..e69de29 diff --git a/chebai_graph/preprocessing/bin/hiv/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/BondType/indices_one_hot.txt new file mode 100644 index 0000000..e69de29 diff --git a/chebai_graph/preprocessing/bin/hiv/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..e69de29 From d2b439dfd8a3fa0eb38e94d5bcb839b2e4f8566e Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Wed, 12 Aug 2026 15:02:52 +0200 Subject: [PATCH 07/58] clintox tokens --- .../clin_tox/AtomCharge/indices_one_hot.txt | 6 ++ .../AtomFunctionalGroup/indices_one_hot.txt | 84 +++++++++++++++++++ .../AtomHybridization/indices_one_hot.txt | 7 ++ .../AtomNodeLevel/indices_one_hot.txt | 3 + .../clin_tox/AtomNumHs/indices_one_hot.txt | 4 + .../bin/clin_tox/AtomType/indices_one_hot.txt | 30 +++++++ .../clin_tox/BondLevel/indices_one_hot.txt | 4 + .../bin/clin_tox/BondType/indices_one_hot.txt | 5 ++ .../clin_tox/NumAtomBonds/indices_one_hot.txt | 6 ++ 9 files changed, 149 insertions(+) create mode 100644 chebai_graph/preprocessing/bin/clin_tox/AtomCharge/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/clin_tox/AtomFunctionalGroup/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/clin_tox/AtomHybridization/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/clin_tox/AtomNodeLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/clin_tox/AtomNumHs/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/clin_tox/AtomType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/clin_tox/BondLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/clin_tox/BondType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/clin_tox/NumAtomBonds/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..374de68 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/AtomCharge/indices_one_hot.txt @@ -0,0 +1,6 @@ +0 +-1 +1 +3 +2 +-2 diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..36bc343 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,84 @@ +alkyl +chloro +RING_6 +ether +secondary_amine +fluoro +graph_fg +tertiary_carbon +hydroxyl +ketone +sulfonyl +carboxylate +RING_5 +NO_FG +amide +sulfide +RING_4 +quaternary_carbon +alkyne +alkene +primary_amine +carboxyl +trifluoromethyl +ester +alkene_carbon +RING_20 +sulfhydryl +nitrile +RING_3 +RING_29 +RING_7 +4_ammonium_ion +acetal +RING_9 +iodo +RING_38 +sulfinyl +difluoromethyl +imide +RING_16 +tertiary_amine +nitro +RING_8 +secondary_ketimine +nitrate +sulfonic_acid +carbamate +secondary_aldimine +bromo +thioketone +RING_14 +azo +RING_33 +trichloromethyl +RING_15 +ketoxime +phosphodiester +RING_24 +RING_22 +aldehyde +RING_18 +RING_21 +nitroso +dichloromethyl +carbonate_ester +silyl_ether +thiolester +borono +aldoxime +sulfonate_ester +RING_23 +RING_19 +RING_25 +carbodithio +disulfide +RING_36 +RING_31 +fluorochloromethyl +RING_12 +phosphate +azide +chlorobromomethyl +bromodifluoromethyl +phosphino diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..a96b9b3 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,7 @@ +SP3 +SP2 +SP +S +SP3D2 +UNSPECIFIED +SP3D diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..1127527 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,4 @@ +3 +0 +2 +1 diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..2c8a324 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/AtomType/indices_one_hot.txt @@ -0,0 +1,30 @@ +6 +17 +7 +8 +9 +16 +15 +53 +80 +1 +13 +35 +79 +0 +27 +22 +78 +33 +34 +30 +14 +5 +29 +24 +81 +43 +83 +25 +20 +26 diff --git a/chebai_graph/preprocessing/bin/clin_tox/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/clin_tox/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/BondType/indices_one_hot.txt new file mode 100644 index 0000000..bdc90d7 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/BondType/indices_one_hot.txt @@ -0,0 +1,5 @@ +SINGLE +AROMATIC +DOUBLE +TRIPLE +DATIVE diff --git a/chebai_graph/preprocessing/bin/clin_tox/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..3c2db51 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,6 @@ +1 +3 +2 +4 +0 +6 From 3a8b8ed3370ee628d3ef1ce0717135a9a95f6260 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Wed, 12 Aug 2026 15:03:55 +0200 Subject: [PATCH 08/58] sider tokens --- .../bin/sider/AtomCharge/indices_one_hot.txt | 6 ++ .../AtomFunctionalGroup/indices_one_hot.txt | 93 +++++++++++++++++++ .../AtomHybridization/indices_one_hot.txt | 6 ++ .../sider/AtomNodeLevel/indices_one_hot.txt | 3 + .../bin/sider/AtomNumHs/indices_one_hot.txt | 5 + .../bin/sider/AtomType/indices_one_hot.txt | 40 ++++++++ .../bin/sider/BondLevel/indices_one_hot.txt | 4 + .../bin/sider/BondType/indices_one_hot.txt | 4 + .../sider/NumAtomBonds/indices_one_hot.txt | 6 ++ 9 files changed, 167 insertions(+) create mode 100644 chebai_graph/preprocessing/bin/sider/AtomCharge/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/sider/AtomFunctionalGroup/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/sider/AtomHybridization/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/sider/AtomNodeLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/sider/AtomNumHs/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/sider/AtomType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/sider/BondLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/sider/BondType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/sider/NumAtomBonds/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/sider/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..6f89a5d --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/AtomCharge/indices_one_hot.txt @@ -0,0 +1,6 @@ +0 +1 +-1 +3 +2 +-2 diff --git a/chebai_graph/preprocessing/bin/sider/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..6fbc2c2 --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,93 @@ +alkyl +RING_6 +carboxyl +alkene +fluoro +sulfinyl +RING_5 +graph_fg +tertiary_carbon +secondary_amine +ketone +primary_amine +primary_ketimine +NO_FG +sulfide +hydroxyl +sulfonic_acid +bromo +4_ammonium_ion +chloro +quaternary_carbon +ether +carboxylate +tertiary_amine +amide +RING_3 +trifluoromethyl +ester +nitro +sulfonyl +iodo +secondary_ketimine +RING_7 +aldehyde +phosphono +alkyne +nitrile +carbamate +alkene_carbon +carbonate_ester +RING_4 +RING_18 +RING_24 +RING_14 +imide +sulfonate_ester +RING_81 +RING_20 +RING_8 +nitroso +RING_29 +RING_16 +RING_12 +RING_21 +sulfhydryl +thiolester +difluoromethyl +thioketone +RING_25 +RING_78 +RING_68 +RING_77 +nitrate +azo +RING_36 +phosphate +ketoxime +secondary_aldimine +dichloromethyl +carboxylic_anhydride +RING_33 +borono +RING_15 +diiodomethyl +RING_23 +carbodithio +disulfide +acetal +RING_40 +phosphodiester +fluorochloromethyl +RING_22 +RING_10 +RING_9 +azide +phosphino +RING_31 +RING_28 +RING_53 +RING_56 +RING_41 +RING_43 +RING_35 diff --git a/chebai_graph/preprocessing/bin/sider/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..e9958eb --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,6 @@ +SP3 +SP2 +S +SP3D +SP +SP3D2 diff --git a/chebai_graph/preprocessing/bin/sider/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/sider/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..df4a3e9 --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,5 @@ +3 +1 +0 +2 +4 diff --git a/chebai_graph/preprocessing/bin/sider/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..384eb37 --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/AtomType/indices_one_hot.txt @@ -0,0 +1,40 @@ +6 +8 +9 +16 +7 +35 +17 +11 +64 +15 +24 +53 +47 +49 +1 +25 +29 +20 +62 +3 +19 +78 +39 +43 +31 +26 +5 +12 +38 +79 +33 +30 +56 +98 +34 +32 +27 +88 +81 +57 diff --git a/chebai_graph/preprocessing/bin/sider/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/sider/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/BondType/indices_one_hot.txt new file mode 100644 index 0000000..f14a241 --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/BondType/indices_one_hot.txt @@ -0,0 +1,4 @@ +SINGLE +DOUBLE +AROMATIC +TRIPLE diff --git a/chebai_graph/preprocessing/bin/sider/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..dfb5c6d --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,6 @@ +1 +2 +3 +4 +0 +6 From 13234d125ac7674e72131531ab3c37118e1be4cb Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Wed, 12 Aug 2026 17:31:41 +0200 Subject: [PATCH 09/58] HIV tokens --- .../bin/hiv/BondLevel/indices_one_hot.txt | 4 ++++ .../bin/hiv/BondType/indices_one_hot.txt | 5 +++++ .../bin/hiv/NumAtomBonds/indices_one_hot.txt | 11 +++++++++++ 3 files changed, 20 insertions(+) diff --git a/chebai_graph/preprocessing/bin/hiv/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/BondLevel/indices_one_hot.txt index e69de29..6b48d66 100644 --- a/chebai_graph/preprocessing/bin/hiv/BondLevel/indices_one_hot.txt +++ b/chebai_graph/preprocessing/bin/hiv/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/hiv/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/BondType/indices_one_hot.txt index e69de29..58c0a5c 100644 --- a/chebai_graph/preprocessing/bin/hiv/BondType/indices_one_hot.txt +++ b/chebai_graph/preprocessing/bin/hiv/BondType/indices_one_hot.txt @@ -0,0 +1,5 @@ +SINGLE +DOUBLE +AROMATIC +TRIPLE +DATIVE diff --git a/chebai_graph/preprocessing/bin/hiv/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/NumAtomBonds/indices_one_hot.txt index e69de29..88ddb48 100644 --- a/chebai_graph/preprocessing/bin/hiv/NumAtomBonds/indices_one_hot.txt +++ b/chebai_graph/preprocessing/bin/hiv/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,11 @@ +1 +2 +3 +4 +0 +5 +6 +9 +7 +8 +10 From 185b78d7f5afe966257497820e33888b64d3355a Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Thu, 13 Aug 2026 12:18:36 +0200 Subject: [PATCH 10/58] muv tokens --- .../bin/muv/AtomCharge/indices_one_hot.txt | 3 + .../AtomFunctionalGroup/indices_one_hot.txt | 68 +++++++++++++++++++ .../muv/AtomHybridization/indices_one_hot.txt | 3 + .../bin/muv/AtomNodeLevel/indices_one_hot.txt | 3 + .../bin/muv/AtomNumHs/indices_one_hot.txt | 4 ++ .../bin/muv/AtomType/indices_one_hot.txt | 7 ++ .../bin/muv/BondLevel/indices_one_hot.txt | 4 ++ .../bin/muv/BondType/indices_one_hot.txt | 4 ++ .../bin/muv/NumAtomBonds/indices_one_hot.txt | 4 ++ 9 files changed, 100 insertions(+) create mode 100644 chebai_graph/preprocessing/bin/muv/AtomCharge/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/muv/AtomFunctionalGroup/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/muv/AtomHybridization/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/muv/AtomNodeLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/muv/AtomNumHs/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/muv/AtomType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/muv/BondLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/muv/BondType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/muv/NumAtomBonds/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/muv/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..d7e52d1 --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/AtomCharge/indices_one_hot.txt @@ -0,0 +1,3 @@ +0 +1 +-1 diff --git a/chebai_graph/preprocessing/bin/muv/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..c55dc92 --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,68 @@ +primary_amine +RING_6 +ketone +carboxyl +alkyl +tertiary_carbon +secondary_amine +graph_fg +hydroxyl +alkene +thioketone +bromo +fluoro +ester +ether +tertiary_amine +NO_FG +chloro +sulfonyl +difluoromethyl +quaternary_carbon +nitro +amide +alkyne +trifluoromethyl +aldoxime +alkene_carbon +nitrile +sulfide +secondary_aldimine +ketoxime +carbodithio +secondary_ketimine +thionoester +thiolester +sulfinyl +carbamate +thiocyanate +aldehyde +sulfonate_ester +amidine +trichloromethyl +sulfhydryl +sulfonic_acid +RING_5 +nitroso +acetal +carbonate_ester +imide +azo +primary_ketimine +dichloromethyl +RING_7 +RING_3 +RING_8 +hydrazone +RING_4 +RING_15 +RING_12 +azide +RING_13 +RING_18 +difluorochloromethyl +RING_14 +RING_9 +RING_11 +disulfide +4_ammonium_ion diff --git a/chebai_graph/preprocessing/bin/muv/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..7c99da3 --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,3 @@ +SP2 +SP3 +SP diff --git a/chebai_graph/preprocessing/bin/muv/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/muv/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..52db5ae --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,4 @@ +2 +1 +0 +3 diff --git a/chebai_graph/preprocessing/bin/muv/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..c36ddc7 --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/AtomType/indices_one_hot.txt @@ -0,0 +1,7 @@ +7 +6 +8 +35 +9 +16 +17 diff --git a/chebai_graph/preprocessing/bin/muv/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/muv/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/BondType/indices_one_hot.txt new file mode 100644 index 0000000..f14a241 --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/BondType/indices_one_hot.txt @@ -0,0 +1,4 @@ +SINGLE +DOUBLE +AROMATIC +TRIPLE diff --git a/chebai_graph/preprocessing/bin/muv/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..94ebaf9 --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,4 @@ +1 +2 +3 +4 From b6bc71d12d42bb20ab1e84028c02c525bbd53b81 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Fri, 14 Aug 2026 23:11:50 +0200 Subject: [PATCH 11/58] pcba tokens --- .../bin/pcba/AtomCharge/indices_one_hot.txt | 8 ++ .../AtomFunctionalGroup/indices_one_hot.txt | 127 ++++++++++++++++++ .../AtomHybridization/indices_one_hot.txt | 7 + .../pcba/AtomNodeLevel/indices_one_hot.txt | 3 + .../bin/pcba/AtomNumHs/indices_one_hot.txt | 6 + .../bin/pcba/AtomType/indices_one_hot.txt | 44 ++++++ .../bin/pcba/BondLevel/indices_one_hot.txt | 4 + .../bin/pcba/BondType/indices_one_hot.txt | 4 + .../bin/pcba/NumAtomBonds/indices_one_hot.txt | 6 + 9 files changed, 209 insertions(+) create mode 100644 chebai_graph/preprocessing/bin/pcba/AtomCharge/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/pcba/AtomFunctionalGroup/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/pcba/AtomHybridization/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/pcba/AtomNodeLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/pcba/AtomNumHs/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/pcba/AtomType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/pcba/BondLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/pcba/BondType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/pcba/NumAtomBonds/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/pcba/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..7897486 --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/AtomCharge/indices_one_hot.txt @@ -0,0 +1,8 @@ +0 +1 +-1 +3 +2 +4 +5 +6 diff --git a/chebai_graph/preprocessing/bin/pcba/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..89e00f3 --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,127 @@ +alkyl +RING_6 +sulfide +RING_5 +ketone +trifluoromethyl +secondary_amine +graph_fg +NO_FG +amide +sulfonyl +ether +quaternary_carbon +tertiary_carbon +imide +carboxyl +nitrile +ester +primary_amine +alkyne +fluoro +tertiary_amine +nitro +RING_7 +chloro +hydroxyl +RING_3 +alkene +secondary_aldimine +bromo +thioketone +sulfinyl +RING_4 +secondary_ketimine +RING_12 +primary_ketimine +alkene_carbon +carbamate +aldehyde +phosphoryl +trimethylsilyl +carboxylate +ketoxime +RING_23 +difluoromethyl +sulfonate_ester +RING_14 +iodo +thiocyanate +RING_8 +disulfide +carbonate_ester +RING_19 +sulfhydryl +nitroso +acetal +aldoxime +borono +RING_16 +phosphono +RING_9 +RING_15 +hydrazone +RING_11 +4_ammonium_ion +thiolester +trichloromethyl +carbodithio +azo +azide +sulfonic_acid +dichloromethyl +phosphodiester +borinate +RING_13 +RING_18 +thial +RING_21 +RING_10 +isocyanate +isonitrile +RING_22 +RING_36 +RING_20 +phosphate +amidine +dibromomethyl +RING_17 +peroxy +RING_40 +silyl_ether +thionoester +nitrate +haloformyl +hemiacetal +RING_25 +isothiocyanate +difluorochloromethyl +RING_24 +nitrosooxy +RING_32 +carbodithioic_acid +hydroperoxy +sulfino +RING_26 +tribromomethyl +RING_30 +RING_29 +diiodomethyl +phosphino +chlorobromomethyl +RING_33 +RING_38 +RING_34 +bromodifluoromethyl +carboxylic_anhydride +bromodichloromethyl +RING_27 +dichlorosilane +RING_42 +ketal +RING_28 +RING_56 +RING_31 +RING_35 +fluorochloromethyl +RING_54 diff --git a/chebai_graph/preprocessing/bin/pcba/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..fb3232b --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,7 @@ +SP3 +SP2 +SP +SP3D +S +UNSPECIFIED +SP3D2 diff --git a/chebai_graph/preprocessing/bin/pcba/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/pcba/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..cad2399 --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,6 @@ +3 +1 +0 +2 +6 +4 diff --git a/chebai_graph/preprocessing/bin/pcba/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..166ffe2 --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/AtomType/indices_one_hot.txt @@ -0,0 +1,44 @@ +6 +16 +7 +8 +9 +17 +35 +15 +14 +53 +5 +29 +79 +38 +50 +30 +33 +34 +48 +51 +1 +26 +82 +27 +78 +80 +3 +28 +22 +56 +44 +11 +13 +46 +25 +83 +47 +60 +24 +64 +32 +49 +74 +31 diff --git a/chebai_graph/preprocessing/bin/pcba/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/pcba/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/BondType/indices_one_hot.txt new file mode 100644 index 0000000..c819d38 --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/BondType/indices_one_hot.txt @@ -0,0 +1,4 @@ +SINGLE +AROMATIC +DOUBLE +TRIPLE diff --git a/chebai_graph/preprocessing/bin/pcba/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..2efbda4 --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,6 @@ +1 +2 +3 +4 +0 +5 From 695fe571dcf3620867b451732d115e508742da0b Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sat, 15 Aug 2026 11:52:27 +0200 Subject: [PATCH 12/58] tox21 tokens --- .../bin/tox21/AtomCharge/indices_one_hot.txt | 6 + .../AtomFunctionalGroup/indices_one_hot.txt | 115 +++++++++++++++++ .../AtomHybridization/indices_one_hot.txt | 7 ++ .../tox21/AtomNodeLevel/indices_one_hot.txt | 3 + .../bin/tox21/AtomNumHs/indices_one_hot.txt | 6 + .../bin/tox21/AtomType/indices_one_hot.txt | 50 ++++++++ .../bin/tox21/BondLevel/indices_one_hot.txt | 4 + .../bin/tox21/BondType/indices_one_hot.txt | 4 + .../tox21/NumAtomBonds/indices_one_hot.txt | 7 ++ .../toxcast/AtomCharge/indices_one_hot.txt | 6 + .../AtomFunctionalGroup/indices_one_hot.txt | 117 ++++++++++++++++++ .../AtomHybridization/indices_one_hot.txt | 7 ++ .../toxcast/AtomNodeLevel/indices_one_hot.txt | 3 + .../bin/toxcast/AtomNumHs/indices_one_hot.txt | 5 + .../bin/toxcast/AtomType/indices_one_hot.txt | 53 ++++++++ .../bin/toxcast/BondLevel/indices_one_hot.txt | 4 + .../bin/toxcast/BondType/indices_one_hot.txt | 4 + .../toxcast/NumAtomBonds/indices_one_hot.txt | 7 ++ 18 files changed, 408 insertions(+) create mode 100644 chebai_graph/preprocessing/bin/tox21/AtomCharge/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/tox21/AtomFunctionalGroup/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/tox21/AtomHybridization/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/tox21/AtomNodeLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/tox21/AtomNumHs/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/tox21/AtomType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/tox21/BondLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/tox21/BondType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/tox21/NumAtomBonds/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/toxcast/AtomCharge/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/toxcast/AtomFunctionalGroup/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/toxcast/AtomHybridization/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/toxcast/AtomNodeLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/toxcast/AtomNumHs/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/toxcast/AtomType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/toxcast/BondLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/toxcast/BondType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/toxcast/NumAtomBonds/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/tox21/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..a8bd8c3 --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/AtomCharge/indices_one_hot.txt @@ -0,0 +1,6 @@ +0 +1 +-1 +2 +3 +-2 diff --git a/chebai_graph/preprocessing/bin/tox21/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..5aa0016 --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,115 @@ +alkyl +hemiacetal +ether +graph_fg +RING_18 +NO_FG +RING_6 +ester +tertiary_carbon +secondary_amine +ketone +carboxyl +RING_5 +amide +primary_amine +RING_4 +carboxylate +alkene +nitro +imide +4_ammonium_ion +bromo +chloro +hydroxyl +sulfonyl +fluoro +quaternary_carbon +alkene_carbon +RING_3 +alkyne +secondary_ketimine +nitrile +isocyanate +azo +tertiary_amine +trifluoromethyl +aldehyde +thiolester +sulfide +nitroso +phosphate +phosphoryl +RING_7 +iodo +acetal +carbodithio +carbonate_ester +RING_21 +difluoromethyl +disulfide +RING_25 +nitrate +carbamate +secondary_aldimine +sulfonic_acid +haloformyl +amidine +sulfhydryl +phosphono +dibromomethyl +RING_23 +trichloromethyl +sulfonate_ester +RING_40 +RING_24 +RING_15 +primary_ketimine +RING_29 +thioketone +aldoxime +RING_16 +nitrosooxy +sulfino +RING_36 +peroxy +RING_14 +ketoxime +thiocyanate +RING_8 +RING_17 +sulfinyl +dichloromethyl +primary_aldimine +tribromomethyl +RING_19 +RING_9 +silyl_ether +phosphodiester +RING_33 +azide +RING_27 +RING_26 +RING_20 +phosphino +RING_12 +trimethylsilyl +carboxylic_anhydride +orthoester +RING_10 +thionoester +isothiocyanate +borono +chlorobromomethyl +dichlorosilane +hydroperoxy +RING_35 +RING_30 +RING_28 +fluorochloromethyl +RING_13 +bromodichloromethyl +ketal +difluorochloromethyl +bromodifluoromethyl +RING_11 diff --git a/chebai_graph/preprocessing/bin/tox21/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..cdc6985 --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,7 @@ +SP3 +SP2 +SP +S +SP3D2 +SP3D +UNSPECIFIED diff --git a/chebai_graph/preprocessing/bin/tox21/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/tox21/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..5a694e8 --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,6 @@ +3 +2 +1 +0 +4 +6 diff --git a/chebai_graph/preprocessing/bin/tox21/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..64e481b --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/AtomType/indices_one_hot.txt @@ -0,0 +1,50 @@ +6 +8 +7 +16 +15 +35 +17 +9 +56 +30 +26 +80 +53 +28 +5 +51 +14 +47 +11 +78 +40 +79 +4 +29 +20 +34 +81 +50 +82 +23 +32 +38 +48 +46 +60 +27 +83 +64 +25 +19 +33 +22 +3 +12 +42 +66 +24 +1 +49 +70 diff --git a/chebai_graph/preprocessing/bin/tox21/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/tox21/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/BondType/indices_one_hot.txt new file mode 100644 index 0000000..c819d38 --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/BondType/indices_one_hot.txt @@ -0,0 +1,4 @@ +SINGLE +AROMATIC +DOUBLE +TRIPLE diff --git a/chebai_graph/preprocessing/bin/tox21/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..21b19a8 --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,7 @@ +1 +2 +3 +4 +0 +5 +6 diff --git a/chebai_graph/preprocessing/bin/toxcast/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..c857ab5 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/AtomCharge/indices_one_hot.txt @@ -0,0 +1,6 @@ +0 +-1 +1 +2 +3 +-2 diff --git a/chebai_graph/preprocessing/bin/toxcast/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..500ebf5 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,117 @@ +ester +RING_6 +sulfonyl +alkyl +NO_FG +4_ammonium_ion +graph_fg +sulfide +primary_amine +carboxyl +chloro +quaternary_carbon +hydroxyl +secondary_amine +fluoro +ketone +tertiary_carbon +ether +alkene +nitro +tertiary_amine +aldehyde +RING_5 +alkyne +nitrile +RING_7 +alkene_carbon +carbonate_ester +trifluoromethyl +imide +sulfonic_acid +primary_ketimine +disulfide +sulfonate_ester +amide +nitroso +iodo +RING_3 +carboxylate +RING_8 +RING_20 +acetal +bromo +carbamate +RING_12 +ketoxime +RING_18 +silyl_ether +RING_4 +dichloromethyl +sulfinyl +aldoxime +thioketone +secondary_ketimine +isothiocyanate +isocyanate +RING_14 +sulfhydryl +phosphono +carbodithioic_acid +secondary_aldimine +haloformyl +phosphoryl +azo +nitrate +RING_35 +thiocyanate +hemiacetal +thiolester +difluorochloromethyl +dibromomethyl +difluoromethyl +phosphino +RING_29 +chlorobromomethyl +RING_9 +tribromomethyl +trichloromethyl +RING_28 +RING_16 +nitrosooxy +bromodichloromethyl +carbodithio +phosphodiester +RING_23 +phosphate +carboxylic_anhydride +RING_13 +RING_24 +RING_30 +RING_11 +RING_17 +azide +RING_40 +primary_aldimine +RING_15 +peroxy +RING_33 +trimethylsilyl +sulfino +ketal +RING_10 +thionoester +borono +amidine +orthoester +RING_27 +RING_26 +RING_38 +hydroperoxy +RING_21 +RING_19 +dichlorosilane +diiodomethyl +fluorochloromethyl +bromodifluoromethyl +RING_36 diff --git a/chebai_graph/preprocessing/bin/toxcast/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..5d039f5 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,7 @@ +SP2 +SP3 +SP +S +UNSPECIFIED +SP3D2 +SP3D diff --git a/chebai_graph/preprocessing/bin/toxcast/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/toxcast/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..ed1daa5 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,5 @@ +0 +1 +3 +2 +4 diff --git a/chebai_graph/preprocessing/bin/toxcast/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..9c59d72 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/AtomType/indices_one_hot.txt @@ -0,0 +1,53 @@ +8 +6 +7 +16 +17 +9 +50 +19 +22 +14 +11 +53 +35 +33 +49 +78 +15 +51 +12 +3 +83 +28 +26 +13 +20 +47 +29 +5 +1 +30 +38 +80 +27 +34 +79 +56 +81 +24 +23 +21 +64 +70 +60 +46 +32 +4 +82 +66 +25 +48 +63 +40 +42 diff --git a/chebai_graph/preprocessing/bin/toxcast/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/toxcast/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/BondType/indices_one_hot.txt new file mode 100644 index 0000000..f14a241 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/BondType/indices_one_hot.txt @@ -0,0 +1,4 @@ +SINGLE +DOUBLE +AROMATIC +TRIPLE diff --git a/chebai_graph/preprocessing/bin/toxcast/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..a61fff5 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,7 @@ +2 +1 +3 +4 +0 +6 +5 From 8f51ea7fdbc09033f2631699c752918b41a91cbf Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Sat, 15 Aug 2026 15:39:38 +0200 Subject: [PATCH 13/58] batch script for augmentation ablation --- .../aug-ablation/aug-ablation-script.sh | 208 ++++++++++++++++++ 1 file changed, 208 insertions(+) create mode 100644 configs/data/augmented/aug-ablation/aug-ablation-script.sh diff --git a/configs/data/augmented/aug-ablation/aug-ablation-script.sh b/configs/data/augmented/aug-ablation/aug-ablation-script.sh new file mode 100644 index 0000000..41f8506 --- /dev/null +++ b/configs/data/augmented/aug-ablation/aug-ablation-script.sh @@ -0,0 +1,208 @@ +#!/bin/bash + +#SBATCH --time=48:00:00 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=32 +#SBATCH --threads-per-core=1 +#SBATCH --mem=256000 + +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --constraint="A100|H100.80gb" + +# ============================================================ +# Job array +# +# 9 data configurations × 3 seeds = 27 experiments +# Maximum 3 experiments running simultaneously +# ============================================================ + +#SBATCH --array=0-26%9 + +#SBATCH --job-name=aug-ablation + +# Separate output/error file for every array task +#SBATCH --output=aug-ablation_%A_%a.out +#SBATCH --error=aug-ablation_%A_%a.err + + +# ============================================================ +# Shell settings +# ============================================================ + +set -x +set -euo pipefail + +# ============================================================ +# Seeds +# ============================================================ + +SEEDS=(0 42 12345) + + +# ============================================================ +# Data configurations +# +# Add/remove configs here. +# +# IMPORTANT: +# The number of array tasks must be: +# +# number of data configs × number of seeds +# +# ============================================================ + +DATA_CONFIGS=( + "../python-chebai-graph/configs/data/augmented/aug-ablation/FGN.yml" + "../python-chebai-graph/configs/data/augmented/aug-ablation/FGN+E.yml" + "../python-chebai-graph/configs/data/augmented/aug-ablation/FGN+E+WGN.yml" + "../python-chebai-graph/configs/data/augmented/aug-ablation/FGN+WGN.yml" + "../python-chebai-graph/configs/data/augmented/aug-ablation/gn_wall_fgwa_nfge.yml" + "../python-chebai-graph/configs/data/augmented/aug-ablation/gn_wall_fgwa_wfge.yml" + "../python-chebai-graph/configs/data/augmented/aug-ablation/gnwa_fgwa_nfge.yml" + "../python-chebai-graph/configs/data/augmented/aug-ablation/gnwa_fgwa_wfge.yml" + "../python-chebai-graph/configs/data/augmented/aug-ablation/WGN.yml" +) + + +# ============================================================ +# Determine data config and seed from array task ID +# ============================================================ + +# Array task ordering: +# All data configurations are run for seed 0 first, +# followed by all data configurations for seed 42, +# and finally all data configurations for seed 12345. +# This is intentionally because the data is generated in first run for each data configuration. +# Then the following runs for the same data configuration can re-use the same generated data. + +NUM_SEEDS=${#SEEDS[@]} +NUM_DATA_CONFIGS=${#DATA_CONFIGS[@]} + +SEED_INDEX=$((SLURM_ARRAY_TASK_ID / NUM_DATA_CONFIGS)) +DATA_INDEX=$((SLURM_ARRAY_TASK_ID % NUM_DATA_CONFIGS)) + +DATA_CONFIG=${DATA_CONFIGS[$DATA_INDEX]} +SEED=${SEEDS[$SEED_INDEX]} + + +# ============================================================ +# Create experiment name +# ============================================================ + +DATA_NAME=$(basename "$DATA_CONFIG" .yml) + +RUN_NAME="${DATA_NAME}_s${SEED}" + + +# ============================================================ +# Print job information +# ============================================================ + +echo "============================================================" +echo "Job ID: $SLURM_JOB_ID" +echo "Array Job ID: $SLURM_ARRAY_JOB_ID" +echo "Array Task ID: $SLURM_ARRAY_TASK_ID" +echo "Data index: $DATA_INDEX" +echo "Seed index: $SEED_INDEX" +echo "Data config: $DATA_CONFIG" +echo "Data name: $DATA_NAME" +echo "Seed: $SEED" +echo "Run name: $RUN_NAME" +echo "Node: $(hostname)" +echo "Date: $(date)" +echo "============================================================" + +echo "GPU information:" +nvidia-smi + +echo "============================================================" + + +# ============================================================ +# Temporary directory +# ============================================================ + +export TMPDIR=/home/staff/a/akhedekar/atmp_dir/ + +export SSL_CERT_FILE=$(python -m certifi) + + +# ============================================================ +# Activate Python environment +# ============================================================ + +source /home/staff/a/akhedekar/python-chebai-graph/.venv/bin/activate + + +# ============================================================ +# Set working directory +# ============================================================ + +CHEBAI_DIR="/home/staff/a/akhedekar/python-chebai" + +cd "$CHEBAI_DIR" + + +# ============================================================ +# Check selected configuration +# ============================================================ + +if [[ ! -f "$DATA_CONFIG" ]]; then + echo "ERROR: Data configuration does not exist:" + echo "$DATA_CONFIG" + exit 1 +fi + + +# ============================================================ +# Set seed +# ============================================================ + +echo "Using seed: $SEED" + +export SEED="$SEED" + + +# ============================================================ +# Run training +# ============================================================ + +python -m chebai fit \ + --trainer=configs/training/default_trainer.yml \ + --trainer.logger=configs/training/wandb_logger.yml \ + --model=../python-chebai-graph/configs/model/baselines/gat.yml \ + --model.train_metrics=configs/metrics/micro-macro-f1.yml \ + --model.test_metrics=configs/metrics/micro-macro-f1.yml \ + --model.val_metrics=configs/metrics/micro-macro-f1.yml \ + --data="$DATA_CONFIG" \ + --data.init_args.batch_size=64 \ + --trainer.accumulate_grad_batches=1 \ + --data.init_args.num_workers=10 \ + --model.pass_loss_kwargs=false \ + --data.init_args.chebi_version=252 \ + --trainer.min_epochs=200 \ + --trainer.max_epochs=200 \ + --model.criterion=configs/loss/bce_unweighted.yml \ + --trainer.logger.init_args.name="$RUN_NAME" \ + --model.init_args.optimizer_kwargs.lr=0.002 \ + --data.init_args.splits_file_path=data/chebi_v252/ChEBI50/processed/splits.csv \ + --trainer.logger.init_args.tags='["augmented_paper","aug-ablation"]' \ + --seed_everything="$SEED" + + +# ============================================================ +# Finished +# ============================================================ + +echo "============================================================" +echo "Training finished successfully" +echo "============================================================" +echo "Job ID: $SLURM_JOB_ID" +echo "Array Task ID: $SLURM_ARRAY_TASK_ID" +echo "Data config: $DATA_CONFIG" +echo "Seed: $SEED" +echo "Run name: $RUN_NAME" +echo "Node: $(hostname)" +echo "Date: $(date)" +echo "============================================================" From a94c7fd748b02addda2980eab6e8abb3a146f7ca Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Sat, 15 Aug 2026 17:40:59 +0200 Subject: [PATCH 14/58] fix ssl error --- .../aug-ablation/aug-ablation-script.sh | 23 +------------------ 1 file changed, 1 insertion(+), 22 deletions(-) diff --git a/configs/data/augmented/aug-ablation/aug-ablation-script.sh b/configs/data/augmented/aug-ablation/aug-ablation-script.sh index 41f8506..c21b85b 100644 --- a/configs/data/augmented/aug-ablation/aug-ablation-script.sh +++ b/configs/data/augmented/aug-ablation/aug-ablation-script.sh @@ -39,7 +39,6 @@ set -euo pipefail SEEDS=(0 42 12345) - # ============================================================ # Data configurations # @@ -64,7 +63,6 @@ DATA_CONFIGS=( "../python-chebai-graph/configs/data/augmented/aug-ablation/WGN.yml" ) - # ============================================================ # Determine data config and seed from array task ID # ============================================================ @@ -85,16 +83,13 @@ DATA_INDEX=$((SLURM_ARRAY_TASK_ID % NUM_DATA_CONFIGS)) DATA_CONFIG=${DATA_CONFIGS[$DATA_INDEX]} SEED=${SEEDS[$SEED_INDEX]} - # ============================================================ # Create experiment name # ============================================================ DATA_NAME=$(basename "$DATA_CONFIG" .yml) - RUN_NAME="${DATA_NAME}_s${SEED}" - # ============================================================ # Print job information # ============================================================ @@ -118,31 +113,25 @@ nvidia-smi echo "============================================================" - # ============================================================ # Temporary directory # ============================================================ export TMPDIR=/home/staff/a/akhedekar/atmp_dir/ -export SSL_CERT_FILE=$(python -m certifi) - - # ============================================================ # Activate Python environment # ============================================================ source /home/staff/a/akhedekar/python-chebai-graph/.venv/bin/activate - # ============================================================ # Set working directory # ============================================================ CHEBAI_DIR="/home/staff/a/akhedekar/python-chebai" - cd "$CHEBAI_DIR" - +export SSL_CERT_FILE=$(python -m certifi) # ============================================================ # Check selected configuration @@ -154,16 +143,6 @@ if [[ ! -f "$DATA_CONFIG" ]]; then exit 1 fi - -# ============================================================ -# Set seed -# ============================================================ - -echo "Using seed: $SEED" - -export SEED="$SEED" - - # ============================================================ # Run training # ============================================================ From 1a3e3c5e5b766336004834ca38d5b4748bacdf72 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Tue, 18 Aug 2026 15:19:51 +0200 Subject: [PATCH 15/58] add baseline to aug-ablation script too --- configs/data/augmented/aug-ablation/aug-ablation-script.sh | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/configs/data/augmented/aug-ablation/aug-ablation-script.sh b/configs/data/augmented/aug-ablation/aug-ablation-script.sh index c21b85b..d19dd95 100644 --- a/configs/data/augmented/aug-ablation/aug-ablation-script.sh +++ b/configs/data/augmented/aug-ablation/aug-ablation-script.sh @@ -13,11 +13,11 @@ # ============================================================ # Job array # -# 9 data configurations × 3 seeds = 27 experiments -# Maximum 3 experiments running simultaneously +# 10 data configurations × 3 seeds = 30 experiments +# Maximum 10 experiments running simultaneously # ============================================================ -#SBATCH --array=0-26%9 +#SBATCH --array=0-29%10 #SBATCH --job-name=aug-ablation @@ -61,6 +61,7 @@ DATA_CONFIGS=( "../python-chebai-graph/configs/data/augmented/aug-ablation/gnwa_fgwa_nfge.yml" "../python-chebai-graph/configs/data/augmented/aug-ablation/gnwa_fgwa_wfge.yml" "../python-chebai-graph/configs/data/augmented/aug-ablation/WGN.yml" + "../python-chebai-graph/configs/data/chebi50_baseline.yml" ) # ============================================================ From a5a2c2e57ef1390b4802b1fb533f7ff544a54e44 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Wed, 19 Aug 2026 22:33:31 +0200 Subject: [PATCH 16/58] aug props for gnwa_wfga_wfge --- .../preprocessing/datasets/__init__.py | 18 ++++++++ chebai_graph/preprocessing/datasets/chebi.py | 46 +++++++++++++++++++ .../chebi50_all_props_except_node_lvl.yml | 0 .../chebi50_atom_fg_prop_only.yml | 0 .../chebi50_atom_graph_node_prop_only.yml | 0 .../chebi50_atom_prop_only.yml | 0 .../chebi50_fg_graph_node_prop_only.yml | 0 .../{ => FGN+E+WGN}/chebi50_fg_prop_only.yml | 0 .../chebi50_graph_node_prop_only.yml | 0 .../chebi50_node_level_prop_only.yml | 0 .../gnwa_fgwa_wfge/chebi50_all_props.yml | 24 ++++++++++ .../chebi50_all_props_except_node_lvl.yml | 24 ++++++++++ .../chebi50_atom_fg_prop_only.yml | 24 ++++++++++ .../chebi50_atom_graph_node_prop_only.yml | 24 ++++++++++ .../gnwa_fgwa_wfge/chebi50_atom_prop_only.yml | 24 ++++++++++ .../chebi50_fg_graph_node_prop_only.yml | 24 ++++++++++ .../gnwa_fgwa_wfge/chebi50_fg_prop_only.yml | 24 ++++++++++ .../chebi50_graph_node_prop_only.yml | 24 ++++++++++ .../chebi50_node_level_prop_only.yml | 24 ++++++++++ 19 files changed, 280 insertions(+) rename configs/data/augmented/properties/{ => FGN+E+WGN}/chebi50_all_props_except_node_lvl.yml (100%) rename configs/data/augmented/properties/{ => FGN+E+WGN}/chebi50_atom_fg_prop_only.yml (100%) rename configs/data/augmented/properties/{ => FGN+E+WGN}/chebi50_atom_graph_node_prop_only.yml (100%) rename configs/data/augmented/properties/{ => FGN+E+WGN}/chebi50_atom_prop_only.yml (100%) rename configs/data/augmented/properties/{ => FGN+E+WGN}/chebi50_fg_graph_node_prop_only.yml (100%) rename configs/data/augmented/properties/{ => FGN+E+WGN}/chebi50_fg_prop_only.yml (100%) rename configs/data/augmented/properties/{ => FGN+E+WGN}/chebi50_graph_node_prop_only.yml (100%) rename configs/data/augmented/properties/{ => FGN+E+WGN}/chebi50_node_level_prop_only.yml (100%) create mode 100644 configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_all_props.yml create mode 100644 configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_all_props_except_node_lvl.yml create mode 100644 configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_atom_fg_prop_only.yml create mode 100644 configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_atom_graph_node_prop_only.yml create mode 100644 configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_atom_prop_only.yml create mode 100644 configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_fg_graph_node_prop_only.yml create mode 100644 configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_fg_prop_only.yml create mode 100644 configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_graph_node_prop_only.yml create mode 100644 configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_node_level_prop_only.yml diff --git a/chebai_graph/preprocessing/datasets/__init__.py b/chebai_graph/preprocessing/datasets/__init__.py index c1e0a23..b0b1df5 100644 --- a/chebai_graph/preprocessing/datasets/__init__.py +++ b/chebai_graph/preprocessing/datasets/__init__.py @@ -5,6 +5,15 @@ ChEBI50_GN_WithAllNodes_FG_WithAtoms_NoFGE, ChEBI50_GN_WithAtoms_FG_WithAtoms_FGE, ChEBI50_GN_WithAtoms_FG_WithAtoms_NoFGE, + ChEBI50_GNWA_FGWA_WFGE_AsPerNodeType, + ChEBI50_GNWA_FGWA_WFGE_ForAllPropsExceptNodeLevel, + ChEBI50_GNWA_FGWA_WFGE_ForAtomAndFGLevelOnly, + ChEBI50_GNWA_FGWA_WFGE_ForAtomLevelAndGraphNodeOnly, + ChEBI50_GNWA_FGWA_WFGE_ForAtomLevelOnly, + ChEBI50_GNWA_FGWA_WFGE_ForFGLevelAndGraphNodeOnly, + ChEBI50_GNWA_FGWA_WFGE_ForFGLevelOnly, + ChEBI50_GNWA_FGWA_WFGE_ForGraphNodeOnly, + ChEBI50_GNWA_FGWA_WFGE_ForNodeLevelPropOnlyForAllNodes, ChEBI50_NFGE_NGN_GraphProp, ChEBI50_NFGE_WGN_GraphProp, ChEBI50_StaticGNI, @@ -71,4 +80,13 @@ "Tox21_WFGE_WGN_AsPerNodeType", "ToxCast_WFGE_WGN_AsPerNodeType", "PCBA_WFGE_WGN_AsPerNodeType", + "ChEBI50_GNWA_FGWA_WFGE_AsPerNodeType", + "ChEBI50_GNWA_FGWA_WFGE_ForAllPropsExceptNodeLevel", + "ChEBI50_GNWA_FGWA_WFGE_ForAtomAndFGLevelOnly", + "ChEBI50_GNWA_FGWA_WFGE_ForAtomLevelAndGraphNodeOnly", + "ChEBI50_GNWA_FGWA_WFGE_ForAtomLevelOnly", + "ChEBI50_GNWA_FGWA_WFGE_ForFGLevelAndGraphNodeOnly", + "ChEBI50_GNWA_FGWA_WFGE_ForFGLevelOnly", + "ChEBI50_GNWA_FGWA_WFGE_ForGraphNodeOnly", + "ChEBI50_GNWA_FGWA_WFGE_ForNodeLevelPropOnlyForAllNodes", ] diff --git a/chebai_graph/preprocessing/datasets/chebi.py b/chebai_graph/preprocessing/datasets/chebi.py index c4ac4ed..7a68dc6 100644 --- a/chebai_graph/preprocessing/datasets/chebi.py +++ b/chebai_graph/preprocessing/datasets/chebi.py @@ -169,48 +169,94 @@ class ChEBI50_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, ChEBIOver50): READER = AtomFGReader_WithFGEdges_WithGraphNode +class ChEBI50_GNWA_FGWA_WFGE_AsPerNodeType(GraphPropAsPerNodeType, ChEBIOver50): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class ChEBI50_WFGE_WGN_ForAtomLevelOnly(GraphPropForAtomLevelOnly, ChEBIOver50): READER = AtomFGReader_WithFGEdges_WithGraphNode +class ChEBI50_GNWA_FGWA_WFGE_ForAtomLevelOnly(GraphPropForAtomLevelOnly, ChEBIOver50): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class ChEBI50_WFGE_WGN_ForFGLevelOnly(GraphPropForFGLevelOnly, ChEBIOver50): READER = AtomFGReader_WithFGEdges_WithGraphNode +class ChEBI50_GNWA_FGWA_WFGE_ForFGLevelOnly(GraphPropForFGLevelOnly, ChEBIOver50): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class ChEBI50_WFGE_WGN_ForGraphNodeOnly(GraphPropForGraphNodeOnly, ChEBIOver50): READER = AtomFGReader_WithFGEdges_WithGraphNode +class ChEBI50_GNWA_FGWA_WFGE_ForGraphNodeOnly(GraphPropForGraphNodeOnly, ChEBIOver50): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class ChEBI50_WFGE_WGN_ForAtomAndFGLevelOnly( GraphPropForAtomAndFGLevelOnly, ChEBIOver50 ): READER = AtomFGReader_WithFGEdges_WithGraphNode +class ChEBI50_GNWA_FGWA_WFGE_ForAtomAndFGLevelOnly( + GraphPropForAtomAndFGLevelOnly, ChEBIOver50 +): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class ChEBI50_WFGE_WGN_ForAtomLevelAndGraphNodeOnly( GraphPropForAtomLevelAndGraphNodeOnly, ChEBIOver50 ): READER = AtomFGReader_WithFGEdges_WithGraphNode +class ChEBI50_GNWA_FGWA_WFGE_ForAtomLevelAndGraphNodeOnly( + GraphPropForAtomLevelAndGraphNodeOnly, ChEBIOver50 +): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class ChEBI50_WFGE_WGN_ForFGLevelAndGraphNodeOnly( GraphPropForFGLevelAndGraphNodeOnly, ChEBIOver50 ): READER = AtomFGReader_WithFGEdges_WithGraphNode +class ChEBI50_GNWA_FGWA_WFGE_ForFGLevelAndGraphNodeOnly( + GraphPropForFGLevelAndGraphNodeOnly, ChEBIOver50 +): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class ChEBI50_WFGE_WGN_ForNodeLevelPropOnlyForAllNodes( GraphPropNodeLevelPropOnlyForAllNodes, ChEBIOver50 ): READER = AtomFGReader_WithFGEdges_WithGraphNode +class ChEBI50_GNWA_FGWA_WFGE_ForNodeLevelPropOnlyForAllNodes( + GraphPropNodeLevelPropOnlyForAllNodes, ChEBIOver50 +): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class ChEBI50_WFGE_WGN_ForAllPropsExceptNodeLevel( GraphPropAllPropsExceptNodeLevel, ChEBIOver50 ): READER = AtomFGReader_WithFGEdges_WithGraphNode +class ChEBI50_GNWA_FGWA_WFGE_ForAllPropsExceptNodeLevel( + GraphPropAllPropsExceptNodeLevel, ChEBIOver50 +): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + # ---------- 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/FGN+E+WGN/chebi50_all_props_except_node_lvl.yml similarity index 100% rename from configs/data/augmented/properties/chebi50_all_props_except_node_lvl.yml rename to configs/data/augmented/properties/FGN+E+WGN/chebi50_all_props_except_node_lvl.yml diff --git a/configs/data/augmented/properties/chebi50_atom_fg_prop_only.yml b/configs/data/augmented/properties/FGN+E+WGN/chebi50_atom_fg_prop_only.yml similarity index 100% rename from configs/data/augmented/properties/chebi50_atom_fg_prop_only.yml rename to configs/data/augmented/properties/FGN+E+WGN/chebi50_atom_fg_prop_only.yml diff --git a/configs/data/augmented/properties/chebi50_atom_graph_node_prop_only.yml b/configs/data/augmented/properties/FGN+E+WGN/chebi50_atom_graph_node_prop_only.yml similarity index 100% rename from configs/data/augmented/properties/chebi50_atom_graph_node_prop_only.yml rename to configs/data/augmented/properties/FGN+E+WGN/chebi50_atom_graph_node_prop_only.yml diff --git a/configs/data/augmented/properties/chebi50_atom_prop_only.yml b/configs/data/augmented/properties/FGN+E+WGN/chebi50_atom_prop_only.yml similarity index 100% rename from configs/data/augmented/properties/chebi50_atom_prop_only.yml rename to configs/data/augmented/properties/FGN+E+WGN/chebi50_atom_prop_only.yml diff --git a/configs/data/augmented/properties/chebi50_fg_graph_node_prop_only.yml b/configs/data/augmented/properties/FGN+E+WGN/chebi50_fg_graph_node_prop_only.yml similarity index 100% rename from configs/data/augmented/properties/chebi50_fg_graph_node_prop_only.yml rename to configs/data/augmented/properties/FGN+E+WGN/chebi50_fg_graph_node_prop_only.yml diff --git a/configs/data/augmented/properties/chebi50_fg_prop_only.yml b/configs/data/augmented/properties/FGN+E+WGN/chebi50_fg_prop_only.yml similarity index 100% rename from configs/data/augmented/properties/chebi50_fg_prop_only.yml rename to configs/data/augmented/properties/FGN+E+WGN/chebi50_fg_prop_only.yml diff --git a/configs/data/augmented/properties/chebi50_graph_node_prop_only.yml b/configs/data/augmented/properties/FGN+E+WGN/chebi50_graph_node_prop_only.yml similarity index 100% rename from configs/data/augmented/properties/chebi50_graph_node_prop_only.yml rename to configs/data/augmented/properties/FGN+E+WGN/chebi50_graph_node_prop_only.yml diff --git a/configs/data/augmented/properties/chebi50_node_level_prop_only.yml b/configs/data/augmented/properties/FGN+E+WGN/chebi50_node_level_prop_only.yml similarity index 100% rename from configs/data/augmented/properties/chebi50_node_level_prop_only.yml rename to configs/data/augmented/properties/FGN+E+WGN/chebi50_node_level_prop_only.yml diff --git a/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_all_props.yml b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_all_props.yml new file mode 100644 index 0000000..3ff3232 --- /dev/null +++ b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_all_props.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GNWA_FGWA_WFGE_AsPerNodeType +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 diff --git a/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_all_props_except_node_lvl.yml b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_all_props_except_node_lvl.yml new file mode 100644 index 0000000..fd33150 --- /dev/null +++ b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_all_props_except_node_lvl.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GNWA_FGWA_WFGE_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 diff --git a/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_atom_fg_prop_only.yml b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_atom_fg_prop_only.yml new file mode 100644 index 0000000..d9053ae --- /dev/null +++ b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_atom_fg_prop_only.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GNWA_FGWA_WFGE_ForAtomAndFGLevelOnly +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 diff --git a/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_atom_graph_node_prop_only.yml b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_atom_graph_node_prop_only.yml new file mode 100644 index 0000000..7ab3ecb --- /dev/null +++ b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_atom_graph_node_prop_only.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GNWA_FGWA_WFGE_ForAtomLevelAndGraphNodeOnly +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 diff --git a/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_atom_prop_only.yml b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_atom_prop_only.yml new file mode 100644 index 0000000..50d18ba --- /dev/null +++ b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_atom_prop_only.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GNWA_FGWA_WFGE_ForAtomLevelOnly +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 diff --git a/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_fg_graph_node_prop_only.yml b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_fg_graph_node_prop_only.yml new file mode 100644 index 0000000..c56edff --- /dev/null +++ b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_fg_graph_node_prop_only.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GNWA_FGWA_WFGE_ForFGLevelAndGraphNodeOnly +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 diff --git a/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_fg_prop_only.yml b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_fg_prop_only.yml new file mode 100644 index 0000000..9b42371 --- /dev/null +++ b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_fg_prop_only.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GNWA_FGWA_WFGE_ForFGLevelOnly +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 diff --git a/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_graph_node_prop_only.yml b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_graph_node_prop_only.yml new file mode 100644 index 0000000..f1921ff --- /dev/null +++ b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_graph_node_prop_only.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GNWA_FGWA_WFGE_ForGraphNodeOnly +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 diff --git a/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_node_level_prop_only.yml b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_node_level_prop_only.yml new file mode 100644 index 0000000..db61465 --- /dev/null +++ b/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_node_level_prop_only.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GNWA_FGWA_WFGE_ForNodeLevelPropOnlyForAllNodes +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 From 6edc8b98acd8c05b5197fd81ce9aa0ec65c754f6 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Thu, 20 Aug 2026 15:31:47 +0200 Subject: [PATCH 17/58] valid mask label for graph collate --- chebai_graph/models/architectures/base.py | 10 ++++- chebai_graph/preprocessing/collate.py | 17 ++++++-- tests/unit/test_collate.py | 48 +++++++++++++++++++++++ 3 files changed, 71 insertions(+), 4 deletions(-) create mode 100644 tests/unit/test_collate.py diff --git a/chebai_graph/models/architectures/base.py b/chebai_graph/models/architectures/base.py index 514e287..11ac3af 100644 --- a/chebai_graph/models/architectures/base.py +++ b/chebai_graph/models/architectures/base.py @@ -26,7 +26,15 @@ def _get_prediction_and_labels( Returns: tuple[torch.Tensor, torch.Tensor]: Tuple of (predictions, labels). """ - return torch.sigmoid(output), labels.int() + valid_label_mask = data["loss_kwargs"]["valid_label_mask"] + predictions = torch.sigmoid(output) + labels = labels.int() + + if valid_label_mask is not None: + predictions = predictions[valid_label_mask] + labels = labels[valid_label_mask] + + return predictions, labels def _process_labels_in_batch(self, batch: XYData) -> torch.Tensor | None: """ diff --git a/chebai_graph/preprocessing/collate.py b/chebai_graph/preprocessing/collate.py index 53bdde5..b2bc4a0 100644 --- a/chebai_graph/preprocessing/collate.py +++ b/chebai_graph/preprocessing/collate.py @@ -15,11 +15,16 @@ def __call__(self, data): # Unpack labels and optional identifiers y, idents = zip(*((d["labels"], d.get("ident")) for d in data)) - # Replace labels with `y` inside graph features and collect them + valid_label_mask = self._get_valid_label_mask(y) + + # Keep labels in XYGraphData.y instead of graph attributes. Mixed labelled + # and unlabelled batches otherwise make PyG collate a partial `y` key. merged_data = [] for row in data: - row["features"].y = row["labels"] - merged_data.append(row["features"]) + features = row["features"] + if "y" in features: + del features.y + merged_data.append(features) # Add empty edge_attr for graphs with no edges to prevent PyG errors for mdata in merged_data: @@ -31,6 +36,11 @@ def __call__(self, data): for attr in merged_data[0].keys(): for data in merged_data: for store in data.stores: + if attr not in store: + continue + if store[attr] is None: + del store[attr] + continue # Im not sure why the following conversion is needed, but it solves this error: # packages/torch_geometric/data/collate.py", line 177, in _collate # value = torch.cat(values, dim=cat_dim or 0, out=out) @@ -69,6 +79,7 @@ def __call__(self, data): x[0].x = x[0].x.to(dtype=torch.int64) # x is a Tuple[BaseData, Mapping, Mapping] + loss_kwargs["valid_label_mask"] = valid_label_mask return XYGraphData( x, y, diff --git a/tests/unit/test_collate.py b/tests/unit/test_collate.py new file mode 100644 index 0000000..d55ad9a --- /dev/null +++ b/tests/unit/test_collate.py @@ -0,0 +1,48 @@ +import unittest + +import torch +from torch_geometric.data import Data as GeomData + +from chebai_graph.preprocessing.collate import GraphCollator + + +class TestGraphCollator(unittest.TestCase): + @staticmethod + def _graph(): + return GeomData( + x=torch.tensor([[6]]), + edge_index=torch.empty((2, 0), dtype=torch.long), + ) + + def test_mixed_missing_labels_do_not_enter_graph_collate(self): + batch = GraphCollator()( + [ + { + "features": self._graph(), + "labels": [1, None, 0, 0, 1], + "ident": "a", + }, + {"features": self._graph(), "labels": None, "ident": "b"}, + { + "features": self._graph(), + "labels": [0, 1, None, 0, 1], + "ident": "c", + }, + ] + ) + + expected_labels = torch.tensor( + [ + [True, False, False, False, True], + [False, True, False, False, True], + ] + ) + self.assertTrue(torch.equal(batch.y, expected_labels)) + self.assertEqual( + batch.additional_fields["loss_kwargs"]["non_null_labels"], [0, 2] + ) + self.assertNotIn("y", batch.x[0]) + + +if __name__ == "__main__": + unittest.main() From 905ed524ae187ffe0df46b398b53631a65f95acc Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Fri, 21 Aug 2026 00:11:02 +0200 Subject: [PATCH 18/58] set invalid label to -1 to so that ignore_index=-1 works --- chebai_graph/models/architectures/base.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/chebai_graph/models/architectures/base.py b/chebai_graph/models/architectures/base.py index 11ac3af..c86baa8 100644 --- a/chebai_graph/models/architectures/base.py +++ b/chebai_graph/models/architectures/base.py @@ -31,8 +31,13 @@ def _get_prediction_and_labels( labels = labels.int() if valid_label_mask is not None: - predictions = predictions[valid_label_mask] - labels = labels[valid_label_mask] + labels[~valid_label_mask] = -1 # Mark invalid labels as -1 + # https://lightning.ai/docs/torchmetrics/stable/classification/auroc#multilabelauroc + # -1 as we torchmetrics ignores -1 labels in multilabel metrics + # metric = MultilabelAUROC( + # num_labels=labels.shape[1], + # ignore_index=-1, + # ) return predictions, labels From ec958e10176a38083daa97b58e66a3d444bcc0f6 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Mon, 24 Aug 2026 18:46:39 +0200 Subject: [PATCH 19/58] aug ablation study for tox21, toxcast --- .gitignore | 3 + .../datasets/molecule_net_classification.py | 154 +++++++++++++- .../aug-ablation/aug-ablation-local.sh | 123 ++++++++++++ .../aug-ablation/aug-ablation-script.sh | 188 ------------------ .../aug-ablation/{ => chebi}/FGN+E+WGN.yml | 0 .../aug-ablation/{ => chebi}/FGN+E.yml | 0 .../aug-ablation/{ => chebi}/FGN+WGN.yml | 0 .../aug-ablation/{ => chebi}/FGN.yml | 0 .../aug-ablation/{ => chebi}/WGN.yml | 0 .../{ => chebi}/gn_wall_fgwa_nfge.yml | 0 .../{ => chebi}/gn_wall_fgwa_wfge.yml | 0 .../{ => chebi}/gnwa_fgwa_nfge.yml | 0 .../{ => chebi}/gnwa_fgwa_wfge.yml | 0 .../aug-ablation/tox21/FGN+E+WGN.yml | 12 ++ .../augmented/aug-ablation/tox21/FGN+E.yml | 12 ++ .../augmented/aug-ablation/tox21/FGN+WGN.yml | 12 ++ .../data/augmented/aug-ablation/tox21/FGN.yml | 12 ++ .../data/augmented/aug-ablation/tox21/WGN.yml | 12 ++ .../augmented/aug-ablation/tox21/baseline.yml | 12 ++ .../aug-ablation/tox21/gn_wall_fgwa_nfge.yml | 12 ++ .../aug-ablation/tox21/gn_wall_fgwa_wfge.yml | 12 ++ .../aug-ablation/tox21/gnwa_fgwa_nfge.yml | 12 ++ .../aug-ablation/tox21/gnwa_fgwa_wfge.yml | 12 ++ .../aug-ablation/toxcast/FGN+E+WGN.yml | 12 ++ .../augmented/aug-ablation/toxcast/FGN+E.yml | 12 ++ .../aug-ablation/toxcast/FGN+WGN.yml | 12 ++ .../augmented/aug-ablation/toxcast/FGN.yml | 12 ++ .../augmented/aug-ablation/toxcast/WGN.yml | 12 ++ .../aug-ablation/toxcast/baseline.yml | 12 ++ .../toxcast/gn_wall_fgwa_nfge.yml | 12 ++ .../toxcast/gn_wall_fgwa_wfge.yml | 12 ++ .../aug-ablation/toxcast/gnwa_fgwa_nfge.yml | 12 ++ .../aug-ablation/toxcast/gnwa_fgwa_wfge.yml | 12 ++ 33 files changed, 519 insertions(+), 189 deletions(-) create mode 100644 configs/data/augmented/aug-ablation/aug-ablation-local.sh delete mode 100644 configs/data/augmented/aug-ablation/aug-ablation-script.sh rename configs/data/augmented/aug-ablation/{ => chebi}/FGN+E+WGN.yml (100%) rename configs/data/augmented/aug-ablation/{ => chebi}/FGN+E.yml (100%) rename configs/data/augmented/aug-ablation/{ => chebi}/FGN+WGN.yml (100%) rename configs/data/augmented/aug-ablation/{ => chebi}/FGN.yml (100%) rename configs/data/augmented/aug-ablation/{ => chebi}/WGN.yml (100%) rename configs/data/augmented/aug-ablation/{ => chebi}/gn_wall_fgwa_nfge.yml (100%) rename configs/data/augmented/aug-ablation/{ => chebi}/gn_wall_fgwa_wfge.yml (100%) rename configs/data/augmented/aug-ablation/{ => chebi}/gnwa_fgwa_nfge.yml (100%) rename configs/data/augmented/aug-ablation/{ => chebi}/gnwa_fgwa_wfge.yml (100%) create mode 100644 configs/data/augmented/aug-ablation/tox21/FGN+E+WGN.yml create mode 100644 configs/data/augmented/aug-ablation/tox21/FGN+E.yml create mode 100644 configs/data/augmented/aug-ablation/tox21/FGN+WGN.yml create mode 100644 configs/data/augmented/aug-ablation/tox21/FGN.yml create mode 100644 configs/data/augmented/aug-ablation/tox21/WGN.yml create mode 100644 configs/data/augmented/aug-ablation/tox21/baseline.yml create mode 100644 configs/data/augmented/aug-ablation/tox21/gn_wall_fgwa_nfge.yml create mode 100644 configs/data/augmented/aug-ablation/tox21/gn_wall_fgwa_wfge.yml create mode 100644 configs/data/augmented/aug-ablation/tox21/gnwa_fgwa_nfge.yml create mode 100644 configs/data/augmented/aug-ablation/tox21/gnwa_fgwa_wfge.yml create mode 100644 configs/data/augmented/aug-ablation/toxcast/FGN+E+WGN.yml create mode 100644 configs/data/augmented/aug-ablation/toxcast/FGN+E.yml create mode 100644 configs/data/augmented/aug-ablation/toxcast/FGN+WGN.yml create mode 100644 configs/data/augmented/aug-ablation/toxcast/FGN.yml create mode 100644 configs/data/augmented/aug-ablation/toxcast/WGN.yml create mode 100644 configs/data/augmented/aug-ablation/toxcast/baseline.yml create mode 100644 configs/data/augmented/aug-ablation/toxcast/gn_wall_fgwa_nfge.yml create mode 100644 configs/data/augmented/aug-ablation/toxcast/gn_wall_fgwa_wfge.yml create mode 100644 configs/data/augmented/aug-ablation/toxcast/gnwa_fgwa_nfge.yml create mode 100644 configs/data/augmented/aug-ablation/toxcast/gnwa_fgwa_wfge.yml diff --git a/.gitignore b/.gitignore index d01eecf..f57d30f 100644 --- a/.gitignore +++ b/.gitignore @@ -174,3 +174,6 @@ electra_pretrained.ckpt *.out *.sh *.ckpt + +!configs/data/augmented/aug-ablation/aug-ablation-local.sh +!configs/data/augmented/aug-ablation/aug-ablation-script.sh \ No newline at end of file diff --git a/chebai_graph/preprocessing/datasets/molecule_net_classification.py b/chebai_graph/preprocessing/datasets/molecule_net_classification.py index 9afaf71..e8d1aa4 100644 --- a/chebai_graph/preprocessing/datasets/molecule_net_classification.py +++ b/chebai_graph/preprocessing/datasets/molecule_net_classification.py @@ -13,8 +13,21 @@ from chebai_graph.preprocessing.datasets.base import ( GraphPropAsPerNodeType, ) -from chebai_graph.preprocessing.reader.augmented_reader import ( +from chebai_graph.preprocessing.reader import ( + AtomFGReader_NoFGEdges_WithGraphNode, + AtomFGReader_WithFGEdges_NoGraphNode, 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, +) + +from .augmentation_base import ( + AugGraphPropMixIn_NoGraphNode, + AugGraphPropMixIn_WithGraphNode, ) @@ -54,6 +67,145 @@ class ToxCast_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, ToxCast): READER = AtomFGReader_WithFGEdges_WithGraphNode +# ---- Augmentation: Variants with graph Node connected to FG nodes only ------------- +class Tox21_WFGE_WGN_GraphProp(AugGraphPropMixIn_WithGraphNode, Tox21): + """Tox21 with with FG nodes and FG edges and graph node.""" + + READER = AtomFGReader_WithFGEdges_WithGraphNode + + +class ToxCast_WFGE_WGN_GraphProp(AugGraphPropMixIn_WithGraphNode, ToxCast): + """ToxCast with with FG nodes and FG edges and graph node.""" + + READER = AtomFGReader_WithFGEdges_WithGraphNode + + +class Tox21_NFGE_WGN_GraphProp(AugGraphPropMixIn_WithGraphNode, Tox21): + """Tox21 with FG nodes but without FG edges, with graph node.""" + + READER = AtomFGReader_NoFGEdges_WithGraphNode + + +class ToxCast_NFGE_WGN_GraphProp(AugGraphPropMixIn_WithGraphNode, ToxCast): + """ToxCast with FG nodes but without FG edges, with graph node.""" + + READER = AtomFGReader_NoFGEdges_WithGraphNode + + +class Tox21_WFGE_NGN_GraphProp(AugGraphPropMixIn_NoGraphNode, Tox21): + """Tox21 with FG nodes and FG edges, no graph node.""" + + READER = AtomFGReader_WithFGEdges_NoGraphNode + + +class ToxCast_WFGE_NGN_GraphProp(AugGraphPropMixIn_NoGraphNode, ToxCast): + """ToxCast with FG nodes and FG edges, no graph node.""" + + READER = AtomFGReader_WithFGEdges_NoGraphNode + + +class Tox21_NFGE_NGN_GraphProp(AugGraphPropMixIn_NoGraphNode, Tox21): + """Tox21 with FG nodes but without FG edges or graph node.""" + + READER = AtomsFGReader_NoFGEdges_NoGraphNode + + +class ToxCast_NFGE_NGN_GraphProp(AugGraphPropMixIn_NoGraphNode, ToxCast): + """ToxCast with FG nodes but without FG edges or graph node.""" + + READER = AtomsFGReader_NoFGEdges_NoGraphNode + + +class Tox21_Atom_WGNOnly_GraphProp(AugGraphPropMixIn_WithGraphNode, Tox21): + """Tox21 with atom-level nodes and graph node only.""" + + READER = AtomReader_WithGraphNodeOnly + + +class ToxCast_Atom_WGNOnly_GraphProp(AugGraphPropMixIn_WithGraphNode, ToxCast): + """ToxCast with atom-level nodes and graph node only.""" + + READER = AtomReader_WithGraphNodeOnly + + +# ------- Augmentation: Variants with graph Node connected to all others nodes (FG and atoms) -------------- +class Tox21_GN_WithAllNodes_FG_WithAtoms_FGE(AugGraphPropMixIn_WithGraphNode, Tox21): + """ + Tox21 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 ToxCast_GN_WithAllNodes_FG_WithAtoms_FGE( + AugGraphPropMixIn_WithGraphNode, ToxCast +): + """ + ToxCast 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 Tox21_GN_WithAllNodes_FG_WithAtoms_NoFGE(AugGraphPropMixIn_WithGraphNode, Tox21): + """ + Tox21 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 ToxCast_GN_WithAllNodes_FG_WithAtoms_NoFGE( + AugGraphPropMixIn_WithGraphNode, ToxCast +): + """ + ToxCast 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 + + +# ------- Augmentation: Variants with graph node connected to atom nodes ONLY ----------- +class Tox21_GN_WithAtoms_FG_WithAtoms_FGE(AugGraphPropMixIn_WithGraphNode, Tox21): + """ + Tox21 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 ToxCast_GN_WithAtoms_FG_WithAtoms_FGE(AugGraphPropMixIn_WithGraphNode, ToxCast): + """ + ToxCast 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 Tox21_GN_WithAtoms_FG_WithAtoms_NoFGE(AugGraphPropMixIn_WithGraphNode, Tox21): + """ + Tox21 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 ToxCast_GN_WithAtoms_FG_WithAtoms_NoFGE(AugGraphPropMixIn_WithGraphNode, ToxCast): + """ + ToxCast 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 + + if __name__ == "__main__": dataset = BACE_WFGE_WGN_AsPerNodeType() dataset.prepare_data() diff --git a/configs/data/augmented/aug-ablation/aug-ablation-local.sh b/configs/data/augmented/aug-ablation/aug-ablation-local.sh new file mode 100644 index 0000000..a40c3e4 --- /dev/null +++ b/configs/data/augmented/aug-ablation/aug-ablation-local.sh @@ -0,0 +1,123 @@ +#!/bin/bash + +# ============================================================ +# Local single-GPU adaptation of ../aug-ablation-script.sh +# +# The original is a Slurm job-array script (10 data configs x 3 +# seeds = 30 experiments, 10 concurrent). This version runs the +# SAME experiments serially on a single local GPU, without Slurm. +# +# Changes vs. the Slurm version: +# - Removed all #SBATCH scheduling directives. +# - Fixed hardcoded /home/staff/a/akhedekar paths -> local user. +# - Replaced SLURM_ARRAY_TASK_ID array fan-out with a serial +# loop over [seeds] x [data configs]. +# +# Requires (verified on this machine): +# - local venv: /home/aditya/python-chebai-graph/.venv +# - chebai CLI: /home/aditya/python-chebai (python -m chebai) +# - chebai_graph src: /home/aditya/python-chebai-graph. +# - wandb credentials in ~/.netrc (or run with a local logger). +# ============================================================ + +set -x +set -euo pipefail + +# ============================================================ +# Local paths +# ============================================================ +CHEBAI_DIR="/home/aditya/python-chebai" +CHEBAI_GRAPH_DIR="/home/aditya/python-chebai-graph" + + +# ============================================================ +# Seeds +# ============================================================ +SEEDS=(0 42 12345) + +# ============================================================ +# Data configurations (all resolved to absolute local paths) +# ============================================================ +DATA_CONFIGS=( + "$CHEBAI_GRAPH_DIR/configs/data/augmented/aug-ablation/FGN.yml" + "$CHEBAI_GRAPH_DIR/configs/data/augmented/aug-ablation/FGN+E.yml" + "$CHEBAI_GRAPH_DIR/configs/data/augmented/aug-ablation/FGN+E+WGN.yml" + "$CHEBAI_GRAPH_DIR/configs/data/augmented/aug-ablation/FGN+WGN.yml" + "$CHEBAI_GRAPH_DIR/configs/data/augmented/aug-ablation/gn_wall_fgwa_nfge.yml" + "$CHEBAI_GRAPH_DIR/configs/data/augmented/aug-ablation/gn_wall_fgwa_wfge.yml" + "$CHEBAI_GRAPH_DIR/configs/data/augmented/aug-ablation/gnwa_fgwa_nfge.yml" + "$CHEBAI_GRAPH_DIR/configs/data/augmented/aug-ablation/gnwa_fgwa_wfge.yml" + "$CHEBAI_GRAPH_DIR/configs/data/augmented/aug-ablation/WGN.yml" + "$CHEBAI_GRAPH_DIR/configs/data/chebi50_baseline.yml" +) + +# ============================================================ +# Activate Python environment +# ============================================================ +source "$CHEBAI_GRAPH_DIR/.venv/bin/activate" + +# ============================================================ +# Set working directory (chebai CLI looks up configs/... from cwd) +# ============================================================ +cd "$CHEBAI_DIR" +export SSL_CERT_FILE=$(python -m certifi) + +echo "============================================================" +echo "Node: $(hostname)" +echo "Date: $(date)" +echo "TMPDIR: $TMPDIR" +echo "GPU:" +nvidia-smi +echo "============================================================" + +# ============================================================ +# Serial loop over all seeds and data configurations +# ============================================================ +for SEED in "${SEEDS[@]}"; do + for DATA_CONFIG in "${DATA_CONFIGS[@]}"; do + + if [[ ! -f "$DATA_CONFIG" ]]; then + echo "ERROR: Data configuration does not exist: $DATA_CONFIG" + exit 1 + fi + + DATA_NAME=$(basename "$DATA_CONFIG" .yml) + RUN_NAME="${DATA_NAME}_s${SEED}" + + echo "============================================================" + echo "Starting run: $RUN_NAME" + echo "Seed: $SEED" + echo "Config: $DATA_CONFIG" + echo "============================================================" + + python -m chebai fit \ + --trainer=configs/training/default_trainer.yml \ + --trainer.logger=configs/training/wandb_logger.yml \ + --model="$CHEBAI_GRAPH_DIR/configs/model/baselines/gat.yml" \ + --model.train_metrics=configs/metrics/micro-macro-f1.yml \ + --model.test_metrics=configs/metrics/micro-macro-f1.yml \ + --model.val_metrics=configs/metrics/micro-macro-f1.yml \ + --data="$DATA_CONFIG" \ + --data.init_args.batch_size=64 \ + --trainer.accumulate_grad_batches=1 \ + --data.init_args.num_workers=10 \ + --model.pass_loss_kwargs=false \ + --data.init_args.chebi_version=252 \ + --trainer.min_epochs=200 \ + --trainer.max_epochs=200 \ + --model.criterion=configs/loss/bce_unweighted.yml \ + --trainer.logger.init_args.name="$RUN_NAME" \ + --model.init_args.optimizer_kwargs.lr=0.002 \ + --data.init_args.splits_file_path=data/chebi_v252/ChEBI50/processed/splits.csv \ + --trainer.logger.init_args.tags='["augmented_paper","aug-ablation"]' \ + --seed_everything="$SEED" + + echo "============================================================" + echo "Finished run: $RUN_NAME" + echo "============================================================" + done +done + +echo "============================================================" +echo "All experiments finished successfully" +echo "============================================================" \ No newline at end of file diff --git a/configs/data/augmented/aug-ablation/aug-ablation-script.sh b/configs/data/augmented/aug-ablation/aug-ablation-script.sh deleted file mode 100644 index d19dd95..0000000 --- a/configs/data/augmented/aug-ablation/aug-ablation-script.sh +++ /dev/null @@ -1,188 +0,0 @@ -#!/bin/bash - -#SBATCH --time=48:00:00 -#SBATCH --ntasks=1 -#SBATCH --cpus-per-task=32 -#SBATCH --threads-per-core=1 -#SBATCH --mem=256000 - -#SBATCH --partition=gpu -#SBATCH --gres=gpu:1 -#SBATCH --constraint="A100|H100.80gb" - -# ============================================================ -# Job array -# -# 10 data configurations × 3 seeds = 30 experiments -# Maximum 10 experiments running simultaneously -# ============================================================ - -#SBATCH --array=0-29%10 - -#SBATCH --job-name=aug-ablation - -# Separate output/error file for every array task -#SBATCH --output=aug-ablation_%A_%a.out -#SBATCH --error=aug-ablation_%A_%a.err - - -# ============================================================ -# Shell settings -# ============================================================ - -set -x -set -euo pipefail - -# ============================================================ -# Seeds -# ============================================================ - -SEEDS=(0 42 12345) - -# ============================================================ -# Data configurations -# -# Add/remove configs here. -# -# IMPORTANT: -# The number of array tasks must be: -# -# number of data configs × number of seeds -# -# ============================================================ - -DATA_CONFIGS=( - "../python-chebai-graph/configs/data/augmented/aug-ablation/FGN.yml" - "../python-chebai-graph/configs/data/augmented/aug-ablation/FGN+E.yml" - "../python-chebai-graph/configs/data/augmented/aug-ablation/FGN+E+WGN.yml" - "../python-chebai-graph/configs/data/augmented/aug-ablation/FGN+WGN.yml" - "../python-chebai-graph/configs/data/augmented/aug-ablation/gn_wall_fgwa_nfge.yml" - "../python-chebai-graph/configs/data/augmented/aug-ablation/gn_wall_fgwa_wfge.yml" - "../python-chebai-graph/configs/data/augmented/aug-ablation/gnwa_fgwa_nfge.yml" - "../python-chebai-graph/configs/data/augmented/aug-ablation/gnwa_fgwa_wfge.yml" - "../python-chebai-graph/configs/data/augmented/aug-ablation/WGN.yml" - "../python-chebai-graph/configs/data/chebi50_baseline.yml" -) - -# ============================================================ -# Determine data config and seed from array task ID -# ============================================================ - -# Array task ordering: -# All data configurations are run for seed 0 first, -# followed by all data configurations for seed 42, -# and finally all data configurations for seed 12345. -# This is intentionally because the data is generated in first run for each data configuration. -# Then the following runs for the same data configuration can re-use the same generated data. - -NUM_SEEDS=${#SEEDS[@]} -NUM_DATA_CONFIGS=${#DATA_CONFIGS[@]} - -SEED_INDEX=$((SLURM_ARRAY_TASK_ID / NUM_DATA_CONFIGS)) -DATA_INDEX=$((SLURM_ARRAY_TASK_ID % NUM_DATA_CONFIGS)) - -DATA_CONFIG=${DATA_CONFIGS[$DATA_INDEX]} -SEED=${SEEDS[$SEED_INDEX]} - -# ============================================================ -# Create experiment name -# ============================================================ - -DATA_NAME=$(basename "$DATA_CONFIG" .yml) -RUN_NAME="${DATA_NAME}_s${SEED}" - -# ============================================================ -# Print job information -# ============================================================ - -echo "============================================================" -echo "Job ID: $SLURM_JOB_ID" -echo "Array Job ID: $SLURM_ARRAY_JOB_ID" -echo "Array Task ID: $SLURM_ARRAY_TASK_ID" -echo "Data index: $DATA_INDEX" -echo "Seed index: $SEED_INDEX" -echo "Data config: $DATA_CONFIG" -echo "Data name: $DATA_NAME" -echo "Seed: $SEED" -echo "Run name: $RUN_NAME" -echo "Node: $(hostname)" -echo "Date: $(date)" -echo "============================================================" - -echo "GPU information:" -nvidia-smi - -echo "============================================================" - -# ============================================================ -# Temporary directory -# ============================================================ - -export TMPDIR=/home/staff/a/akhedekar/atmp_dir/ - -# ============================================================ -# Activate Python environment -# ============================================================ - -source /home/staff/a/akhedekar/python-chebai-graph/.venv/bin/activate - -# ============================================================ -# Set working directory -# ============================================================ - -CHEBAI_DIR="/home/staff/a/akhedekar/python-chebai" -cd "$CHEBAI_DIR" -export SSL_CERT_FILE=$(python -m certifi) - -# ============================================================ -# Check selected configuration -# ============================================================ - -if [[ ! -f "$DATA_CONFIG" ]]; then - echo "ERROR: Data configuration does not exist:" - echo "$DATA_CONFIG" - exit 1 -fi - -# ============================================================ -# Run training -# ============================================================ - -python -m chebai fit \ - --trainer=configs/training/default_trainer.yml \ - --trainer.logger=configs/training/wandb_logger.yml \ - --model=../python-chebai-graph/configs/model/baselines/gat.yml \ - --model.train_metrics=configs/metrics/micro-macro-f1.yml \ - --model.test_metrics=configs/metrics/micro-macro-f1.yml \ - --model.val_metrics=configs/metrics/micro-macro-f1.yml \ - --data="$DATA_CONFIG" \ - --data.init_args.batch_size=64 \ - --trainer.accumulate_grad_batches=1 \ - --data.init_args.num_workers=10 \ - --model.pass_loss_kwargs=false \ - --data.init_args.chebi_version=252 \ - --trainer.min_epochs=200 \ - --trainer.max_epochs=200 \ - --model.criterion=configs/loss/bce_unweighted.yml \ - --trainer.logger.init_args.name="$RUN_NAME" \ - --model.init_args.optimizer_kwargs.lr=0.002 \ - --data.init_args.splits_file_path=data/chebi_v252/ChEBI50/processed/splits.csv \ - --trainer.logger.init_args.tags='["augmented_paper","aug-ablation"]' \ - --seed_everything="$SEED" - - -# ============================================================ -# Finished -# ============================================================ - -echo "============================================================" -echo "Training finished successfully" -echo "============================================================" -echo "Job ID: $SLURM_JOB_ID" -echo "Array Task ID: $SLURM_ARRAY_TASK_ID" -echo "Data config: $DATA_CONFIG" -echo "Seed: $SEED" -echo "Run name: $RUN_NAME" -echo "Node: $(hostname)" -echo "Date: $(date)" -echo "============================================================" diff --git a/configs/data/augmented/aug-ablation/FGN+E+WGN.yml b/configs/data/augmented/aug-ablation/chebi/FGN+E+WGN.yml similarity index 100% rename from configs/data/augmented/aug-ablation/FGN+E+WGN.yml rename to configs/data/augmented/aug-ablation/chebi/FGN+E+WGN.yml diff --git a/configs/data/augmented/aug-ablation/FGN+E.yml b/configs/data/augmented/aug-ablation/chebi/FGN+E.yml similarity index 100% rename from configs/data/augmented/aug-ablation/FGN+E.yml rename to configs/data/augmented/aug-ablation/chebi/FGN+E.yml diff --git a/configs/data/augmented/aug-ablation/FGN+WGN.yml b/configs/data/augmented/aug-ablation/chebi/FGN+WGN.yml similarity index 100% rename from configs/data/augmented/aug-ablation/FGN+WGN.yml rename to configs/data/augmented/aug-ablation/chebi/FGN+WGN.yml diff --git a/configs/data/augmented/aug-ablation/FGN.yml b/configs/data/augmented/aug-ablation/chebi/FGN.yml similarity index 100% rename from configs/data/augmented/aug-ablation/FGN.yml rename to configs/data/augmented/aug-ablation/chebi/FGN.yml diff --git a/configs/data/augmented/aug-ablation/WGN.yml b/configs/data/augmented/aug-ablation/chebi/WGN.yml similarity index 100% rename from configs/data/augmented/aug-ablation/WGN.yml rename to configs/data/augmented/aug-ablation/chebi/WGN.yml diff --git a/configs/data/augmented/aug-ablation/gn_wall_fgwa_nfge.yml b/configs/data/augmented/aug-ablation/chebi/gn_wall_fgwa_nfge.yml similarity index 100% rename from configs/data/augmented/aug-ablation/gn_wall_fgwa_nfge.yml rename to configs/data/augmented/aug-ablation/chebi/gn_wall_fgwa_nfge.yml diff --git a/configs/data/augmented/aug-ablation/gn_wall_fgwa_wfge.yml b/configs/data/augmented/aug-ablation/chebi/gn_wall_fgwa_wfge.yml similarity index 100% rename from configs/data/augmented/aug-ablation/gn_wall_fgwa_wfge.yml rename to configs/data/augmented/aug-ablation/chebi/gn_wall_fgwa_wfge.yml diff --git a/configs/data/augmented/aug-ablation/gnwa_fgwa_nfge.yml b/configs/data/augmented/aug-ablation/chebi/gnwa_fgwa_nfge.yml similarity index 100% rename from configs/data/augmented/aug-ablation/gnwa_fgwa_nfge.yml rename to configs/data/augmented/aug-ablation/chebi/gnwa_fgwa_nfge.yml diff --git a/configs/data/augmented/aug-ablation/gnwa_fgwa_wfge.yml b/configs/data/augmented/aug-ablation/chebi/gnwa_fgwa_wfge.yml similarity index 100% rename from configs/data/augmented/aug-ablation/gnwa_fgwa_wfge.yml rename to configs/data/augmented/aug-ablation/chebi/gnwa_fgwa_wfge.yml diff --git a/configs/data/augmented/aug-ablation/tox21/FGN+E+WGN.yml b/configs/data/augmented/aug-ablation/tox21/FGN+E+WGN.yml new file mode 100644 index 0000000..e8520f3 --- /dev/null +++ b/configs/data/augmented/aug-ablation/tox21/FGN+E+WGN.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_WFGE_WGN_GraphProp +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/tox21/FGN+E.yml b/configs/data/augmented/aug-ablation/tox21/FGN+E.yml new file mode 100644 index 0000000..c50baac --- /dev/null +++ b/configs/data/augmented/aug-ablation/tox21/FGN+E.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_WFGE_NGN_GraphProp +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/tox21/FGN+WGN.yml b/configs/data/augmented/aug-ablation/tox21/FGN+WGN.yml new file mode 100644 index 0000000..cae7d3a --- /dev/null +++ b/configs/data/augmented/aug-ablation/tox21/FGN+WGN.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_NFGE_WGN_GraphProp +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/tox21/FGN.yml b/configs/data/augmented/aug-ablation/tox21/FGN.yml new file mode 100644 index 0000000..44b3705 --- /dev/null +++ b/configs/data/augmented/aug-ablation/tox21/FGN.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_NFGE_NGN_GraphProp +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/tox21/WGN.yml b/configs/data/augmented/aug-ablation/tox21/WGN.yml new file mode 100644 index 0000000..7b5f253 --- /dev/null +++ b/configs/data/augmented/aug-ablation/tox21/WGN.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_Atom_WGNOnly_GraphProp +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/tox21/baseline.yml b/configs/data/augmented/aug-ablation/tox21/baseline.yml new file mode 100644 index 0000000..ce9dfa0 --- /dev/null +++ b/configs/data/augmented/aug-ablation/tox21/baseline.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50GraphProperties +init_args: + properties: + - chebai_graph.preprocessing.properties.AtomType + - chebai_graph.preprocessing.properties.NumAtomBonds + - chebai_graph.preprocessing.properties.AtomCharge + - chebai_graph.preprocessing.properties.AtomAromaticity + - chebai_graph.preprocessing.properties.AtomHybridization + - chebai_graph.preprocessing.properties.AtomNumHs + - chebai_graph.preprocessing.properties.BondType + - chebai_graph.preprocessing.properties.BondInRing + - chebai_graph.preprocessing.properties.BondAromaticity diff --git a/configs/data/augmented/aug-ablation/tox21/gn_wall_fgwa_nfge.yml b/configs/data/augmented/aug-ablation/tox21/gn_wall_fgwa_nfge.yml new file mode 100644 index 0000000..c75cbbf --- /dev/null +++ b/configs/data/augmented/aug-ablation/tox21/gn_wall_fgwa_nfge.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GN_WithAllNodes_FG_WithAtoms_NoFGE +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/tox21/gn_wall_fgwa_wfge.yml b/configs/data/augmented/aug-ablation/tox21/gn_wall_fgwa_wfge.yml new file mode 100644 index 0000000..4eda133 --- /dev/null +++ b/configs/data/augmented/aug-ablation/tox21/gn_wall_fgwa_wfge.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GN_WithAllNodes_FG_WithAtoms_FGE +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/tox21/gnwa_fgwa_nfge.yml b/configs/data/augmented/aug-ablation/tox21/gnwa_fgwa_nfge.yml new file mode 100644 index 0000000..5d6ff23 --- /dev/null +++ b/configs/data/augmented/aug-ablation/tox21/gnwa_fgwa_nfge.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GN_WithAtoms_FG_WithAtoms_NoFGE +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/tox21/gnwa_fgwa_wfge.yml b/configs/data/augmented/aug-ablation/tox21/gnwa_fgwa_wfge.yml new file mode 100644 index 0000000..648ecfc --- /dev/null +++ b/configs/data/augmented/aug-ablation/tox21/gnwa_fgwa_wfge.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GN_WithAtoms_FG_WithAtoms_FGE +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/toxcast/FGN+E+WGN.yml b/configs/data/augmented/aug-ablation/toxcast/FGN+E+WGN.yml new file mode 100644 index 0000000..e8520f3 --- /dev/null +++ b/configs/data/augmented/aug-ablation/toxcast/FGN+E+WGN.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_WFGE_WGN_GraphProp +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/toxcast/FGN+E.yml b/configs/data/augmented/aug-ablation/toxcast/FGN+E.yml new file mode 100644 index 0000000..c50baac --- /dev/null +++ b/configs/data/augmented/aug-ablation/toxcast/FGN+E.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_WFGE_NGN_GraphProp +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/toxcast/FGN+WGN.yml b/configs/data/augmented/aug-ablation/toxcast/FGN+WGN.yml new file mode 100644 index 0000000..cae7d3a --- /dev/null +++ b/configs/data/augmented/aug-ablation/toxcast/FGN+WGN.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_NFGE_WGN_GraphProp +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/toxcast/FGN.yml b/configs/data/augmented/aug-ablation/toxcast/FGN.yml new file mode 100644 index 0000000..44b3705 --- /dev/null +++ b/configs/data/augmented/aug-ablation/toxcast/FGN.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_NFGE_NGN_GraphProp +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/toxcast/WGN.yml b/configs/data/augmented/aug-ablation/toxcast/WGN.yml new file mode 100644 index 0000000..7b5f253 --- /dev/null +++ b/configs/data/augmented/aug-ablation/toxcast/WGN.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_Atom_WGNOnly_GraphProp +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/toxcast/baseline.yml b/configs/data/augmented/aug-ablation/toxcast/baseline.yml new file mode 100644 index 0000000..ce9dfa0 --- /dev/null +++ b/configs/data/augmented/aug-ablation/toxcast/baseline.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50GraphProperties +init_args: + properties: + - chebai_graph.preprocessing.properties.AtomType + - chebai_graph.preprocessing.properties.NumAtomBonds + - chebai_graph.preprocessing.properties.AtomCharge + - chebai_graph.preprocessing.properties.AtomAromaticity + - chebai_graph.preprocessing.properties.AtomHybridization + - chebai_graph.preprocessing.properties.AtomNumHs + - chebai_graph.preprocessing.properties.BondType + - chebai_graph.preprocessing.properties.BondInRing + - chebai_graph.preprocessing.properties.BondAromaticity diff --git a/configs/data/augmented/aug-ablation/toxcast/gn_wall_fgwa_nfge.yml b/configs/data/augmented/aug-ablation/toxcast/gn_wall_fgwa_nfge.yml new file mode 100644 index 0000000..c75cbbf --- /dev/null +++ b/configs/data/augmented/aug-ablation/toxcast/gn_wall_fgwa_nfge.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GN_WithAllNodes_FG_WithAtoms_NoFGE +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/toxcast/gn_wall_fgwa_wfge.yml b/configs/data/augmented/aug-ablation/toxcast/gn_wall_fgwa_wfge.yml new file mode 100644 index 0000000..4eda133 --- /dev/null +++ b/configs/data/augmented/aug-ablation/toxcast/gn_wall_fgwa_wfge.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GN_WithAllNodes_FG_WithAtoms_FGE +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/toxcast/gnwa_fgwa_nfge.yml b/configs/data/augmented/aug-ablation/toxcast/gnwa_fgwa_nfge.yml new file mode 100644 index 0000000..5d6ff23 --- /dev/null +++ b/configs/data/augmented/aug-ablation/toxcast/gnwa_fgwa_nfge.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GN_WithAtoms_FG_WithAtoms_NoFGE +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity diff --git a/configs/data/augmented/aug-ablation/toxcast/gnwa_fgwa_wfge.yml b/configs/data/augmented/aug-ablation/toxcast/gnwa_fgwa_wfge.yml new file mode 100644 index 0000000..648ecfc --- /dev/null +++ b/configs/data/augmented/aug-ablation/toxcast/gnwa_fgwa_wfge.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GN_WithAtoms_FG_WithAtoms_FGE +init_args: + properties: + - chebai_graph.preprocessing.properties.AugAtomType + - chebai_graph.preprocessing.properties.AugNumAtomBonds + - chebai_graph.preprocessing.properties.AugAtomCharge + - chebai_graph.preprocessing.properties.AugAtomAromaticity + - chebai_graph.preprocessing.properties.AugAtomHybridization + - chebai_graph.preprocessing.properties.AugAtomNumHs + - chebai_graph.preprocessing.properties.AugBondType + - chebai_graph.preprocessing.properties.AugBondInRing + - chebai_graph.preprocessing.properties.AugBondAromaticity From 0611c493312a9c5a94901b9870fcced100f115c1 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Mon, 24 Aug 2026 20:06:07 +0200 Subject: [PATCH 20/58] add tox ablation classes to configs --- .../preprocessing/datasets/__init__.py | 62 ++++++++++++++++--- .../datasets/molecule_net_classification.py | 14 +++++ .../aug-ablation/tox21/FGN+E+WGN.yml | 2 +- .../augmented/aug-ablation/tox21/FGN+E.yml | 2 +- .../augmented/aug-ablation/tox21/FGN+WGN.yml | 2 +- .../data/augmented/aug-ablation/tox21/FGN.yml | 2 +- .../data/augmented/aug-ablation/tox21/WGN.yml | 2 +- .../augmented/aug-ablation/tox21/baseline.yml | 2 +- .../aug-ablation/tox21/gn_wall_fgwa_nfge.yml | 2 +- .../aug-ablation/tox21/gn_wall_fgwa_wfge.yml | 2 +- .../aug-ablation/tox21/gnwa_fgwa_nfge.yml | 2 +- .../aug-ablation/tox21/gnwa_fgwa_wfge.yml | 2 +- .../aug-ablation/toxcast/FGN+E+WGN.yml | 2 +- .../augmented/aug-ablation/toxcast/FGN+E.yml | 2 +- .../aug-ablation/toxcast/FGN+WGN.yml | 2 +- .../augmented/aug-ablation/toxcast/FGN.yml | 2 +- .../augmented/aug-ablation/toxcast/WGN.yml | 2 +- .../aug-ablation/toxcast/baseline.yml | 2 +- .../toxcast/gn_wall_fgwa_nfge.yml | 2 +- .../toxcast/gn_wall_fgwa_wfge.yml | 2 +- .../aug-ablation/toxcast/gnwa_fgwa_nfge.yml | 2 +- .../aug-ablation/toxcast/gnwa_fgwa_wfge.yml | 2 +- 22 files changed, 88 insertions(+), 28 deletions(-) diff --git a/chebai_graph/preprocessing/datasets/__init__.py b/chebai_graph/preprocessing/datasets/__init__.py index b0b1df5..ad2aca9 100644 --- a/chebai_graph/preprocessing/datasets/__init__.py +++ b/chebai_graph/preprocessing/datasets/__init__.py @@ -40,29 +40,45 @@ MUV_WFGE_WGN_AsPerNodeType, PCBA_WFGE_WGN_AsPerNodeType, SIDER_WFGE_WGN_AsPerNodeType, + Tox21_Atom_WGNOnly_GraphProp, + Tox21_GN_WithAllNodes_FG_WithAtoms_FGE, + Tox21_GN_WithAllNodes_FG_WithAtoms_NoFGE, + Tox21_GN_WithAtoms_FG_WithAtoms_FGE, + Tox21_GN_WithAtoms_FG_WithAtoms_NoFGE, + Tox21_NFGE_NGN_GraphProp, + Tox21_NFGE_WGN_GraphProp, + Tox21_WFGE_NGN_GraphProp, Tox21_WFGE_WGN_AsPerNodeType, + Tox21_WFGE_WGN_GraphProp, + Tox21GraphProperties, + ToxCast_Atom_WGNOnly_GraphProp, + ToxCast_GN_WithAllNodes_FG_WithAtoms_FGE, + ToxCast_GN_WithAllNodes_FG_WithAtoms_NoFGE, + ToxCast_GN_WithAtoms_FG_WithAtoms_FGE, + ToxCast_GN_WithAtoms_FG_WithAtoms_NoFGE, + ToxCast_NFGE_NGN_GraphProp, + ToxCast_NFGE_WGN_GraphProp, + ToxCast_WFGE_NGN_GraphProp, ToxCast_WFGE_WGN_AsPerNodeType, + ToxCast_WFGE_WGN_GraphProp, + ToxCastGraphProperties, ) from .pubchem import PubChemGraphProperties __all__ = [ - "ChEBI50GraphFGAugmentorReader", - "ChEBI25GraphProperties", - "ChEBI50GraphProperties", - "ChEBI100GraphProperties", - "ChEBI50GraphData", - "PubChemGraphProperties", - "ChEBI50_Atom_WGNOnly_GraphProp", + # ---------- CHEBI Augmentation ablation classes -------- + "ChEBI50GraphProperties", # -- baseline class for chebi "ChEBI50_NFGE_NGN_GraphProp", "ChEBI50_NFGE_WGN_GraphProp", "ChEBI50_WFGE_NGN_GraphProp", "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", + "ChEBI50_Atom_WGNOnly_GraphProp", + # ----- CHEBI FGN+E+FGE Properties ablation classes ----- "ChEBI50_WFGE_WGN_ForAtomLevelOnly", "ChEBI50_WFGE_WGN_ForFGLevelOnly", "ChEBI50_WFGE_WGN_ForGraphNodeOnly", @@ -71,6 +87,7 @@ "ChEBI50_WFGE_WGN_ForFGLevelAndGraphNodeOnly", "ChEBI50_WFGE_WGN_ForNodeLevelPropOnlyForAllNodes", "ChEBI50_WFGE_WGN_ForAllPropsExceptNodeLevel", + # ----- MoleculeNet Dataset FGN+E+FGE Final classes ----- "BACE_WFGE_WGN_AsPerNodeType", "BBBP_WFGE_WGN_AsPerNodeType", "ClinTox_WFGE_WGN_AsPerNodeType", @@ -80,6 +97,7 @@ "Tox21_WFGE_WGN_AsPerNodeType", "ToxCast_WFGE_WGN_AsPerNodeType", "PCBA_WFGE_WGN_AsPerNodeType", + # ----- CHEBI GNWA_FGWA_WFGE Properties ablation classes ----- "ChEBI50_GNWA_FGWA_WFGE_AsPerNodeType", "ChEBI50_GNWA_FGWA_WFGE_ForAllPropsExceptNodeLevel", "ChEBI50_GNWA_FGWA_WFGE_ForAtomAndFGLevelOnly", @@ -89,4 +107,32 @@ "ChEBI50_GNWA_FGWA_WFGE_ForFGLevelOnly", "ChEBI50_GNWA_FGWA_WFGE_ForGraphNodeOnly", "ChEBI50_GNWA_FGWA_WFGE_ForNodeLevelPropOnlyForAllNodes", + # ----- Tox21 Augmentation Ablation classes ----- + "Tox21_Atom_WGNOnly_GraphProp", + "Tox21_GN_WithAllNodes_FG_WithAtoms_FGE", + "Tox21_GN_WithAllNodes_FG_WithAtoms_NoFGE", + "Tox21_GN_WithAtoms_FG_WithAtoms_FGE", + "Tox21_GN_WithAtoms_FG_WithAtoms_NoFGE", + "Tox21_NFGE_NGN_GraphProp", + "Tox21_NFGE_WGN_GraphProp", + "Tox21_WFGE_NGN_GraphProp", + "Tox21_WFGE_WGN_GraphProp", + "Tox21GraphProperties", # -- baseline class for Tox21 + # ----- ToxCast Augmentation Ablation classes ----- + "ToxCast_Atom_WGNOnly_GraphProp", + "ToxCast_GN_WithAllNodes_FG_WithAtoms_FGE", + "ToxCast_GN_WithAllNodes_FG_WithAtoms_NoFGE", + "ToxCast_GN_WithAtoms_FG_WithAtoms_FGE", + "ToxCast_GN_WithAtoms_FG_WithAtoms_NoFGE", + "ToxCast_NFGE_NGN_GraphProp", + "ToxCast_NFGE_WGN_GraphProp", + "ToxCast_WFGE_NGN_GraphProp", + "ToxCast_WFGE_WGN_GraphProp", + "ToxCastGraphProperties", # -- baseline class for ToxCast + # ------ Others ----- + "ChEBI50_StaticGNI", + "ChEBI25GraphProperties", + "ChEBI100GraphProperties", + "ChEBI50GraphData", + "PubChemGraphProperties", ] diff --git a/chebai_graph/preprocessing/datasets/molecule_net_classification.py b/chebai_graph/preprocessing/datasets/molecule_net_classification.py index e8d1aa4..d015d5f 100644 --- a/chebai_graph/preprocessing/datasets/molecule_net_classification.py +++ b/chebai_graph/preprocessing/datasets/molecule_net_classification.py @@ -12,6 +12,7 @@ from chebai_graph.preprocessing.datasets.base import ( GraphPropAsPerNodeType, + GraphPropertiesMixIn, ) from chebai_graph.preprocessing.reader import ( AtomFGReader_NoFGEdges_WithGraphNode, @@ -206,6 +207,19 @@ class ToxCast_GN_WithAtoms_FG_WithAtoms_NoFGE(AugGraphPropMixIn_WithGraphNode, T READER = GN_WithAtoms_FG_WithAtoms_NoFGE +# ---------------------------Baselines classes for Tox21 and ToxCast datasets +class Tox21GraphProperties(GraphPropertiesMixIn, Tox21): + """Tox21 dataset with molecular property encodings.""" + + pass + + +class ToxCastGraphProperties(GraphPropertiesMixIn, ToxCast): + """ToxCast dataset with molecular property encodings.""" + + pass + + if __name__ == "__main__": dataset = BACE_WFGE_WGN_AsPerNodeType() dataset.prepare_data() diff --git a/configs/data/augmented/aug-ablation/tox21/FGN+E+WGN.yml b/configs/data/augmented/aug-ablation/tox21/FGN+E+WGN.yml index e8520f3..9b1f918 100644 --- a/configs/data/augmented/aug-ablation/tox21/FGN+E+WGN.yml +++ b/configs/data/augmented/aug-ablation/tox21/FGN+E+WGN.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.ChEBI50_WFGE_WGN_GraphProp +class_path: chebai_graph.preprocessing.datasets.Tox21_WFGE_WGN_GraphProp init_args: properties: - chebai_graph.preprocessing.properties.AugAtomType diff --git a/configs/data/augmented/aug-ablation/tox21/FGN+E.yml b/configs/data/augmented/aug-ablation/tox21/FGN+E.yml index c50baac..193eec0 100644 --- a/configs/data/augmented/aug-ablation/tox21/FGN+E.yml +++ b/configs/data/augmented/aug-ablation/tox21/FGN+E.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.ChEBI50_WFGE_NGN_GraphProp +class_path: chebai_graph.preprocessing.datasets.Tox21_WFGE_NGN_GraphProp init_args: properties: - chebai_graph.preprocessing.properties.AugAtomType diff --git a/configs/data/augmented/aug-ablation/tox21/FGN+WGN.yml b/configs/data/augmented/aug-ablation/tox21/FGN+WGN.yml index cae7d3a..3515dfb 100644 --- a/configs/data/augmented/aug-ablation/tox21/FGN+WGN.yml +++ b/configs/data/augmented/aug-ablation/tox21/FGN+WGN.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.ChEBI50_NFGE_WGN_GraphProp +class_path: chebai_graph.preprocessing.datasets.Tox21_NFGE_WGN_GraphProp init_args: properties: - chebai_graph.preprocessing.properties.AugAtomType diff --git a/configs/data/augmented/aug-ablation/tox21/FGN.yml b/configs/data/augmented/aug-ablation/tox21/FGN.yml index 44b3705..50feae1 100644 --- a/configs/data/augmented/aug-ablation/tox21/FGN.yml +++ b/configs/data/augmented/aug-ablation/tox21/FGN.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.ChEBI50_NFGE_NGN_GraphProp +class_path: chebai_graph.preprocessing.datasets.Tox21_NFGE_NGN_GraphProp init_args: properties: - chebai_graph.preprocessing.properties.AugAtomType diff --git a/configs/data/augmented/aug-ablation/tox21/WGN.yml b/configs/data/augmented/aug-ablation/tox21/WGN.yml index 7b5f253..bdb347f 100644 --- a/configs/data/augmented/aug-ablation/tox21/WGN.yml +++ b/configs/data/augmented/aug-ablation/tox21/WGN.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.ChEBI50_Atom_WGNOnly_GraphProp +class_path: chebai_graph.preprocessing.datasets.Tox21_Atom_WGNOnly_GraphProp init_args: properties: - chebai_graph.preprocessing.properties.AugAtomType diff --git a/configs/data/augmented/aug-ablation/tox21/baseline.yml b/configs/data/augmented/aug-ablation/tox21/baseline.yml index ce9dfa0..c0aeeda 100644 --- a/configs/data/augmented/aug-ablation/tox21/baseline.yml +++ b/configs/data/augmented/aug-ablation/tox21/baseline.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.ChEBI50GraphProperties +class_path: chebai_graph.preprocessing.datasets.Tox21GraphProperties init_args: properties: - chebai_graph.preprocessing.properties.AtomType diff --git a/configs/data/augmented/aug-ablation/tox21/gn_wall_fgwa_nfge.yml b/configs/data/augmented/aug-ablation/tox21/gn_wall_fgwa_nfge.yml index c75cbbf..88ddd67 100644 --- a/configs/data/augmented/aug-ablation/tox21/gn_wall_fgwa_nfge.yml +++ b/configs/data/augmented/aug-ablation/tox21/gn_wall_fgwa_nfge.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.ChEBI50_GN_WithAllNodes_FG_WithAtoms_NoFGE +class_path: chebai_graph.preprocessing.datasets.Tox21_GN_WithAllNodes_FG_WithAtoms_NoFGE init_args: properties: - chebai_graph.preprocessing.properties.AugAtomType diff --git a/configs/data/augmented/aug-ablation/tox21/gn_wall_fgwa_wfge.yml b/configs/data/augmented/aug-ablation/tox21/gn_wall_fgwa_wfge.yml index 4eda133..a227629 100644 --- a/configs/data/augmented/aug-ablation/tox21/gn_wall_fgwa_wfge.yml +++ b/configs/data/augmented/aug-ablation/tox21/gn_wall_fgwa_wfge.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.ChEBI50_GN_WithAllNodes_FG_WithAtoms_FGE +class_path: chebai_graph.preprocessing.datasets.Tox21_GN_WithAllNodes_FG_WithAtoms_FGE init_args: properties: - chebai_graph.preprocessing.properties.AugAtomType diff --git a/configs/data/augmented/aug-ablation/tox21/gnwa_fgwa_nfge.yml b/configs/data/augmented/aug-ablation/tox21/gnwa_fgwa_nfge.yml index 5d6ff23..5102fdf 100644 --- a/configs/data/augmented/aug-ablation/tox21/gnwa_fgwa_nfge.yml +++ b/configs/data/augmented/aug-ablation/tox21/gnwa_fgwa_nfge.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.ChEBI50_GN_WithAtoms_FG_WithAtoms_NoFGE +class_path: chebai_graph.preprocessing.datasets.Tox21_GN_WithAtoms_FG_WithAtoms_NoFGE init_args: properties: - chebai_graph.preprocessing.properties.AugAtomType diff --git a/configs/data/augmented/aug-ablation/tox21/gnwa_fgwa_wfge.yml b/configs/data/augmented/aug-ablation/tox21/gnwa_fgwa_wfge.yml index 648ecfc..d79a0ea 100644 --- a/configs/data/augmented/aug-ablation/tox21/gnwa_fgwa_wfge.yml +++ b/configs/data/augmented/aug-ablation/tox21/gnwa_fgwa_wfge.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.ChEBI50_GN_WithAtoms_FG_WithAtoms_FGE +class_path: chebai_graph.preprocessing.datasets.Tox21_GN_WithAtoms_FG_WithAtoms_FGE init_args: properties: - chebai_graph.preprocessing.properties.AugAtomType diff --git a/configs/data/augmented/aug-ablation/toxcast/FGN+E+WGN.yml b/configs/data/augmented/aug-ablation/toxcast/FGN+E+WGN.yml index e8520f3..662da11 100644 --- a/configs/data/augmented/aug-ablation/toxcast/FGN+E+WGN.yml +++ b/configs/data/augmented/aug-ablation/toxcast/FGN+E+WGN.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.ChEBI50_WFGE_WGN_GraphProp +class_path: chebai_graph.preprocessing.datasets.ToxCast_WFGE_WGN_GraphProp init_args: properties: - chebai_graph.preprocessing.properties.AugAtomType diff --git a/configs/data/augmented/aug-ablation/toxcast/FGN+E.yml b/configs/data/augmented/aug-ablation/toxcast/FGN+E.yml index c50baac..1c5b321 100644 --- a/configs/data/augmented/aug-ablation/toxcast/FGN+E.yml +++ b/configs/data/augmented/aug-ablation/toxcast/FGN+E.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.ChEBI50_WFGE_NGN_GraphProp +class_path: chebai_graph.preprocessing.datasets.ToxCast_WFGE_NGN_GraphProp init_args: properties: - chebai_graph.preprocessing.properties.AugAtomType diff --git a/configs/data/augmented/aug-ablation/toxcast/FGN+WGN.yml b/configs/data/augmented/aug-ablation/toxcast/FGN+WGN.yml index cae7d3a..875a57d 100644 --- a/configs/data/augmented/aug-ablation/toxcast/FGN+WGN.yml +++ b/configs/data/augmented/aug-ablation/toxcast/FGN+WGN.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.ChEBI50_NFGE_WGN_GraphProp +class_path: chebai_graph.preprocessing.datasets.ToxCast_NFGE_WGN_GraphProp init_args: properties: - chebai_graph.preprocessing.properties.AugAtomType diff --git a/configs/data/augmented/aug-ablation/toxcast/FGN.yml b/configs/data/augmented/aug-ablation/toxcast/FGN.yml index 44b3705..d046445 100644 --- a/configs/data/augmented/aug-ablation/toxcast/FGN.yml +++ b/configs/data/augmented/aug-ablation/toxcast/FGN.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.ChEBI50_NFGE_NGN_GraphProp +class_path: chebai_graph.preprocessing.datasets.ToxCast_NFGE_NGN_GraphProp init_args: properties: - chebai_graph.preprocessing.properties.AugAtomType diff --git a/configs/data/augmented/aug-ablation/toxcast/WGN.yml b/configs/data/augmented/aug-ablation/toxcast/WGN.yml index 7b5f253..1464be0 100644 --- a/configs/data/augmented/aug-ablation/toxcast/WGN.yml +++ b/configs/data/augmented/aug-ablation/toxcast/WGN.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.ChEBI50_Atom_WGNOnly_GraphProp +class_path: chebai_graph.preprocessing.datasets.ToxCast_Atom_WGNOnly_GraphProp init_args: properties: - chebai_graph.preprocessing.properties.AugAtomType diff --git a/configs/data/augmented/aug-ablation/toxcast/baseline.yml b/configs/data/augmented/aug-ablation/toxcast/baseline.yml index ce9dfa0..5a7e18c 100644 --- a/configs/data/augmented/aug-ablation/toxcast/baseline.yml +++ b/configs/data/augmented/aug-ablation/toxcast/baseline.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.ChEBI50GraphProperties +class_path: chebai_graph.preprocessing.datasets.ToxCastGraphProperties init_args: properties: - chebai_graph.preprocessing.properties.AtomType diff --git a/configs/data/augmented/aug-ablation/toxcast/gn_wall_fgwa_nfge.yml b/configs/data/augmented/aug-ablation/toxcast/gn_wall_fgwa_nfge.yml index c75cbbf..c38ced5 100644 --- a/configs/data/augmented/aug-ablation/toxcast/gn_wall_fgwa_nfge.yml +++ b/configs/data/augmented/aug-ablation/toxcast/gn_wall_fgwa_nfge.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.ChEBI50_GN_WithAllNodes_FG_WithAtoms_NoFGE +class_path: chebai_graph.preprocessing.datasets.ToxCast_GN_WithAllNodes_FG_WithAtoms_NoFGE init_args: properties: - chebai_graph.preprocessing.properties.AugAtomType diff --git a/configs/data/augmented/aug-ablation/toxcast/gn_wall_fgwa_wfge.yml b/configs/data/augmented/aug-ablation/toxcast/gn_wall_fgwa_wfge.yml index 4eda133..8d28ad9 100644 --- a/configs/data/augmented/aug-ablation/toxcast/gn_wall_fgwa_wfge.yml +++ b/configs/data/augmented/aug-ablation/toxcast/gn_wall_fgwa_wfge.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.ChEBI50_GN_WithAllNodes_FG_WithAtoms_FGE +class_path: chebai_graph.preprocessing.datasets.ToxCast_GN_WithAllNodes_FG_WithAtoms_FGE init_args: properties: - chebai_graph.preprocessing.properties.AugAtomType diff --git a/configs/data/augmented/aug-ablation/toxcast/gnwa_fgwa_nfge.yml b/configs/data/augmented/aug-ablation/toxcast/gnwa_fgwa_nfge.yml index 5d6ff23..bae368f 100644 --- a/configs/data/augmented/aug-ablation/toxcast/gnwa_fgwa_nfge.yml +++ b/configs/data/augmented/aug-ablation/toxcast/gnwa_fgwa_nfge.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.ChEBI50_GN_WithAtoms_FG_WithAtoms_NoFGE +class_path: chebai_graph.preprocessing.datasets.ToxCast_GN_WithAtoms_FG_WithAtoms_NoFGE init_args: properties: - chebai_graph.preprocessing.properties.AugAtomType diff --git a/configs/data/augmented/aug-ablation/toxcast/gnwa_fgwa_wfge.yml b/configs/data/augmented/aug-ablation/toxcast/gnwa_fgwa_wfge.yml index 648ecfc..d754e25 100644 --- a/configs/data/augmented/aug-ablation/toxcast/gnwa_fgwa_wfge.yml +++ b/configs/data/augmented/aug-ablation/toxcast/gnwa_fgwa_wfge.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.ChEBI50_GN_WithAtoms_FG_WithAtoms_FGE +class_path: chebai_graph.preprocessing.datasets.ToxCast_GN_WithAtoms_FG_WithAtoms_FGE init_args: properties: - chebai_graph.preprocessing.properties.AugAtomType From a43cfc84284ed9c17729e4a742256c0b9a57e778 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sat, 5 Sep 2026 13:12:08 +0200 Subject: [PATCH 21/58] add aug ablation scripts --- .gitignore | 3 +- .../aug-ablation/aug-ablation-chebi.sh | 188 ++++++++++++++++++ ...blation-local.sh => aug-ablation-tox21.sh} | 38 ++-- .../aug-ablation/aug-ablation-toxcast.sh | 125 ++++++++++++ 4 files changed, 334 insertions(+), 20 deletions(-) create mode 100644 configs/data/augmented/aug-ablation/aug-ablation-chebi.sh rename configs/data/augmented/aug-ablation/{aug-ablation-local.sh => aug-ablation-tox21.sh} (78%) create mode 100644 configs/data/augmented/aug-ablation/aug-ablation-toxcast.sh diff --git a/.gitignore b/.gitignore index f57d30f..95daf78 100644 --- a/.gitignore +++ b/.gitignore @@ -175,5 +175,4 @@ electra_pretrained.ckpt *.sh *.ckpt -!configs/data/augmented/aug-ablation/aug-ablation-local.sh -!configs/data/augmented/aug-ablation/aug-ablation-script.sh \ No newline at end of file +!configs/data/augmented/aug-ablation/*.sh diff --git a/configs/data/augmented/aug-ablation/aug-ablation-chebi.sh b/configs/data/augmented/aug-ablation/aug-ablation-chebi.sh new file mode 100644 index 0000000..d19dd95 --- /dev/null +++ b/configs/data/augmented/aug-ablation/aug-ablation-chebi.sh @@ -0,0 +1,188 @@ +#!/bin/bash + +#SBATCH --time=48:00:00 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=32 +#SBATCH --threads-per-core=1 +#SBATCH --mem=256000 + +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --constraint="A100|H100.80gb" + +# ============================================================ +# Job array +# +# 10 data configurations × 3 seeds = 30 experiments +# Maximum 10 experiments running simultaneously +# ============================================================ + +#SBATCH --array=0-29%10 + +#SBATCH --job-name=aug-ablation + +# Separate output/error file for every array task +#SBATCH --output=aug-ablation_%A_%a.out +#SBATCH --error=aug-ablation_%A_%a.err + + +# ============================================================ +# Shell settings +# ============================================================ + +set -x +set -euo pipefail + +# ============================================================ +# Seeds +# ============================================================ + +SEEDS=(0 42 12345) + +# ============================================================ +# Data configurations +# +# Add/remove configs here. +# +# IMPORTANT: +# The number of array tasks must be: +# +# number of data configs × number of seeds +# +# ============================================================ + +DATA_CONFIGS=( + "../python-chebai-graph/configs/data/augmented/aug-ablation/FGN.yml" + "../python-chebai-graph/configs/data/augmented/aug-ablation/FGN+E.yml" + "../python-chebai-graph/configs/data/augmented/aug-ablation/FGN+E+WGN.yml" + "../python-chebai-graph/configs/data/augmented/aug-ablation/FGN+WGN.yml" + "../python-chebai-graph/configs/data/augmented/aug-ablation/gn_wall_fgwa_nfge.yml" + "../python-chebai-graph/configs/data/augmented/aug-ablation/gn_wall_fgwa_wfge.yml" + "../python-chebai-graph/configs/data/augmented/aug-ablation/gnwa_fgwa_nfge.yml" + "../python-chebai-graph/configs/data/augmented/aug-ablation/gnwa_fgwa_wfge.yml" + "../python-chebai-graph/configs/data/augmented/aug-ablation/WGN.yml" + "../python-chebai-graph/configs/data/chebi50_baseline.yml" +) + +# ============================================================ +# Determine data config and seed from array task ID +# ============================================================ + +# Array task ordering: +# All data configurations are run for seed 0 first, +# followed by all data configurations for seed 42, +# and finally all data configurations for seed 12345. +# This is intentionally because the data is generated in first run for each data configuration. +# Then the following runs for the same data configuration can re-use the same generated data. + +NUM_SEEDS=${#SEEDS[@]} +NUM_DATA_CONFIGS=${#DATA_CONFIGS[@]} + +SEED_INDEX=$((SLURM_ARRAY_TASK_ID / NUM_DATA_CONFIGS)) +DATA_INDEX=$((SLURM_ARRAY_TASK_ID % NUM_DATA_CONFIGS)) + +DATA_CONFIG=${DATA_CONFIGS[$DATA_INDEX]} +SEED=${SEEDS[$SEED_INDEX]} + +# ============================================================ +# Create experiment name +# ============================================================ + +DATA_NAME=$(basename "$DATA_CONFIG" .yml) +RUN_NAME="${DATA_NAME}_s${SEED}" + +# ============================================================ +# Print job information +# ============================================================ + +echo "============================================================" +echo "Job ID: $SLURM_JOB_ID" +echo "Array Job ID: $SLURM_ARRAY_JOB_ID" +echo "Array Task ID: $SLURM_ARRAY_TASK_ID" +echo "Data index: $DATA_INDEX" +echo "Seed index: $SEED_INDEX" +echo "Data config: $DATA_CONFIG" +echo "Data name: $DATA_NAME" +echo "Seed: $SEED" +echo "Run name: $RUN_NAME" +echo "Node: $(hostname)" +echo "Date: $(date)" +echo "============================================================" + +echo "GPU information:" +nvidia-smi + +echo "============================================================" + +# ============================================================ +# Temporary directory +# ============================================================ + +export TMPDIR=/home/staff/a/akhedekar/atmp_dir/ + +# ============================================================ +# Activate Python environment +# ============================================================ + +source /home/staff/a/akhedekar/python-chebai-graph/.venv/bin/activate + +# ============================================================ +# Set working directory +# ============================================================ + +CHEBAI_DIR="/home/staff/a/akhedekar/python-chebai" +cd "$CHEBAI_DIR" +export SSL_CERT_FILE=$(python -m certifi) + +# ============================================================ +# Check selected configuration +# ============================================================ + +if [[ ! -f "$DATA_CONFIG" ]]; then + echo "ERROR: Data configuration does not exist:" + echo "$DATA_CONFIG" + exit 1 +fi + +# ============================================================ +# Run training +# ============================================================ + +python -m chebai fit \ + --trainer=configs/training/default_trainer.yml \ + --trainer.logger=configs/training/wandb_logger.yml \ + --model=../python-chebai-graph/configs/model/baselines/gat.yml \ + --model.train_metrics=configs/metrics/micro-macro-f1.yml \ + --model.test_metrics=configs/metrics/micro-macro-f1.yml \ + --model.val_metrics=configs/metrics/micro-macro-f1.yml \ + --data="$DATA_CONFIG" \ + --data.init_args.batch_size=64 \ + --trainer.accumulate_grad_batches=1 \ + --data.init_args.num_workers=10 \ + --model.pass_loss_kwargs=false \ + --data.init_args.chebi_version=252 \ + --trainer.min_epochs=200 \ + --trainer.max_epochs=200 \ + --model.criterion=configs/loss/bce_unweighted.yml \ + --trainer.logger.init_args.name="$RUN_NAME" \ + --model.init_args.optimizer_kwargs.lr=0.002 \ + --data.init_args.splits_file_path=data/chebi_v252/ChEBI50/processed/splits.csv \ + --trainer.logger.init_args.tags='["augmented_paper","aug-ablation"]' \ + --seed_everything="$SEED" + + +# ============================================================ +# Finished +# ============================================================ + +echo "============================================================" +echo "Training finished successfully" +echo "============================================================" +echo "Job ID: $SLURM_JOB_ID" +echo "Array Task ID: $SLURM_ARRAY_TASK_ID" +echo "Data config: $DATA_CONFIG" +echo "Seed: $SEED" +echo "Run name: $RUN_NAME" +echo "Node: $(hostname)" +echo "Date: $(date)" +echo "============================================================" diff --git a/configs/data/augmented/aug-ablation/aug-ablation-local.sh b/configs/data/augmented/aug-ablation/aug-ablation-tox21.sh similarity index 78% rename from configs/data/augmented/aug-ablation/aug-ablation-local.sh rename to configs/data/augmented/aug-ablation/aug-ablation-tox21.sh index a40c3e4..756514c 100644 --- a/configs/data/augmented/aug-ablation/aug-ablation-local.sh +++ b/configs/data/augmented/aug-ablation/aug-ablation-tox21.sh @@ -38,17 +38,20 @@ SEEDS=(0 42 12345) # ============================================================ # Data configurations (all resolved to absolute local paths) # ============================================================ +DATA_CONFIG_DIR=( + "$CHEBAI_GRAPH_DIR/configs/data/augmented/aug-ablation/tox21" +) DATA_CONFIGS=( - "$CHEBAI_GRAPH_DIR/configs/data/augmented/aug-ablation/FGN.yml" - "$CHEBAI_GRAPH_DIR/configs/data/augmented/aug-ablation/FGN+E.yml" - "$CHEBAI_GRAPH_DIR/configs/data/augmented/aug-ablation/FGN+E+WGN.yml" - "$CHEBAI_GRAPH_DIR/configs/data/augmented/aug-ablation/FGN+WGN.yml" - "$CHEBAI_GRAPH_DIR/configs/data/augmented/aug-ablation/gn_wall_fgwa_nfge.yml" - "$CHEBAI_GRAPH_DIR/configs/data/augmented/aug-ablation/gn_wall_fgwa_wfge.yml" - "$CHEBAI_GRAPH_DIR/configs/data/augmented/aug-ablation/gnwa_fgwa_nfge.yml" - "$CHEBAI_GRAPH_DIR/configs/data/augmented/aug-ablation/gnwa_fgwa_wfge.yml" - "$CHEBAI_GRAPH_DIR/configs/data/augmented/aug-ablation/WGN.yml" - "$CHEBAI_GRAPH_DIR/configs/data/chebi50_baseline.yml" + "$DATA_CONFIG_DIR/FGN.yml" + "$DATA_CONFIG_DIR/FGN+E.yml" + "$DATA_CONFIG_DIR/FGN+E+WGN.yml" + "$DATA_CONFIG_DIR/FGN+WGN.yml" + "$DATA_CONFIG_DIR/gn_wall_fgwa_nfge.yml" + "$DATA_CONFIG_DIR/gn_wall_fgwa_wfge.yml" + "$DATA_CONFIG_DIR/gnwa_fgwa_nfge.yml" + "$DATA_CONFIG_DIR/gnwa_fgwa_wfge.yml" + "$DATA_CONFIG_DIR/WGN.yml" + "$DATA_CONFIG_DIR/baseline.yml" ) # ============================================================ @@ -94,22 +97,21 @@ for SEED in "${SEEDS[@]}"; do --trainer=configs/training/default_trainer.yml \ --trainer.logger=configs/training/wandb_logger.yml \ --model="$CHEBAI_GRAPH_DIR/configs/model/baselines/gat.yml" \ - --model.train_metrics=configs/metrics/micro-macro-f1.yml \ - --model.test_metrics=configs/metrics/micro-macro-f1.yml \ - --model.val_metrics=configs/metrics/micro-macro-f1.yml \ + --model.train_metrics=configs/metrics/micro-macro-f1-roc-auc.yml \ + --model.test_metrics=configs/metrics/micro-macro-f1-roc-auc.yml \ + --model.val_metrics=configs/metrics/micro-macro-f1-roc-auc.yml \ --data="$DATA_CONFIG" \ --data.init_args.batch_size=64 \ - --trainer.accumulate_grad_batches=1 \ + --trainer.accumulate_grad_batches=2 \ --data.init_args.num_workers=10 \ --model.pass_loss_kwargs=false \ - --data.init_args.chebi_version=252 \ --trainer.min_epochs=200 \ --trainer.max_epochs=200 \ --model.criterion=configs/loss/bce_unweighted.yml \ --trainer.logger.init_args.name="$RUN_NAME" \ - --model.init_args.optimizer_kwargs.lr=0.002 \ - --data.init_args.splits_file_path=data/chebi_v252/ChEBI50/processed/splits.csv \ - --trainer.logger.init_args.tags='["augmented_paper","aug-ablation"]' \ + --data.init_args.splits_file_path=data/Tox21:MNClassification/Tox21/processed/splits.csv \ + --trainer.logger.init_args.tags='["augmented_paper","aug-ablation", "tox21"]' \ + --model.config.edge_dim=10 \ --seed_everything="$SEED" echo "============================================================" diff --git a/configs/data/augmented/aug-ablation/aug-ablation-toxcast.sh b/configs/data/augmented/aug-ablation/aug-ablation-toxcast.sh new file mode 100644 index 0000000..e355623 --- /dev/null +++ b/configs/data/augmented/aug-ablation/aug-ablation-toxcast.sh @@ -0,0 +1,125 @@ +#!/bin/bash + +# ============================================================ +# Local single-GPU adaptation of ../aug-ablation-script.sh +# +# The original is a Slurm job-array script (10 data configs x 3 +# seeds = 30 experiments, 10 concurrent). This version runs the +# SAME experiments serially on a single local GPU, without Slurm. +# +# Changes vs. the Slurm version: +# - Removed all #SBATCH scheduling directives. +# - Fixed hardcoded /home/staff/a/akhedekar paths -> local user. +# - Replaced SLURM_ARRAY_TASK_ID array fan-out with a serial +# loop over [seeds] x [data configs]. +# +# Requires (verified on this machine): +# - local venv: /home/aditya/python-chebai-graph/.venv +# - chebai CLI: /home/aditya/python-chebai (python -m chebai) +# - chebai_graph src: /home/aditya/python-chebai-graph. +# - wandb credentials in ~/.netrc (or run with a local logger). +# ============================================================ + +set -x +set -euo pipefail + +# ============================================================ +# Local paths +# ============================================================ +CHEBAI_DIR="/home/aditya/python-chebai" +CHEBAI_GRAPH_DIR="/home/aditya/python-chebai-graph" + + +# ============================================================ +# Seeds +# ============================================================ +SEEDS=(0 42 12345) + +# ============================================================ +# Data configurations (all resolved to absolute local paths) +# ============================================================ +DATA_CONFIG_DIR=( + "$CHEBAI_GRAPH_DIR/configs/data/augmented/aug-ablation/toxcast" +) +DATA_CONFIGS=( + "$DATA_CONFIG_DIR/FGN.yml" + "$DATA_CONFIG_DIR/FGN+E.yml" + "$DATA_CONFIG_DIR/FGN+E+WGN.yml" + "$DATA_CONFIG_DIR/FGN+WGN.yml" + "$DATA_CONFIG_DIR/gn_wall_fgwa_nfge.yml" + "$DATA_CONFIG_DIR/gn_wall_fgwa_wfge.yml" + "$DATA_CONFIG_DIR/gnwa_fgwa_nfge.yml" + "$DATA_CONFIG_DIR/gnwa_fgwa_wfge.yml" + "$DATA_CONFIG_DIR/WGN.yml" + "$DATA_CONFIG_DIR/baseline.yml" +) + +# ============================================================ +# Activate Python environment +# ============================================================ +source "$CHEBAI_GRAPH_DIR/.venv/bin/activate" + +# ============================================================ +# Set working directory (chebai CLI looks up configs/... from cwd) +# ============================================================ +cd "$CHEBAI_DIR" +export SSL_CERT_FILE=$(python -m certifi) + +echo "============================================================" +echo "Node: $(hostname)" +echo "Date: $(date)" +echo "TMPDIR: $TMPDIR" +echo "GPU:" +nvidia-smi +echo "============================================================" + +# ============================================================ +# Serial loop over all seeds and data configurations +# ============================================================ +for SEED in "${SEEDS[@]}"; do + for DATA_CONFIG in "${DATA_CONFIGS[@]}"; do + + if [[ ! -f "$DATA_CONFIG" ]]; then + echo "ERROR: Data configuration does not exist: $DATA_CONFIG" + exit 1 + fi + + DATA_NAME=$(basename "$DATA_CONFIG" .yml) + RUN_NAME="${DATA_NAME}_s${SEED}" + + echo "============================================================" + echo "Starting run: $RUN_NAME" + echo "Seed: $SEED" + echo "Config: $DATA_CONFIG" + echo "============================================================" + + python -m chebai fit \ + --trainer=configs/training/default_trainer.yml \ + --trainer.logger=configs/training/wandb_logger.yml \ + --model="$CHEBAI_GRAPH_DIR/configs/model/baselines/gat.yml" \ + --model.train_metrics=configs/metrics/micro-macro-f1-roc-auc.yml \ + --model.test_metrics=configs/metrics/micro-macro-f1-roc-auc.yml \ + --model.val_metrics=configs/metrics/micro-macro-f1-roc-auc.yml \ + --data="$DATA_CONFIG" \ + --data.init_args.batch_size=64 \ + --trainer.accumulate_grad_batches=2 \ + --data.init_args.num_workers=10 \ + --model.pass_loss_kwargs=false \ + --trainer.min_epochs=200 \ + --trainer.max_epochs=200 \ + --model.criterion=configs/loss/bce_unweighted.yml \ + --trainer.logger.init_args.name="$RUN_NAME" \ + --data.init_args.splits_file_path=data/ToxCast:MNClassification/ToxCast/processed/splits.csv \ + --trainer.logger.init_args.tags='["augmented_paper","aug-ablation", "toxcast"]' \ + --model.config.edge_dim=10 \ + --seed_everything="$SEED" + + echo "============================================================" + echo "Finished run: $RUN_NAME" + echo "============================================================" + done +done + +echo "============================================================" +echo "All experiments finished successfully" +echo "============================================================" \ No newline at end of file From e1e1d77ba0b9fe0c6a72eb0ac9e91a0fbdea138a Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sun, 6 Sep 2026 17:24:43 +0200 Subject: [PATCH 22/58] rectify script --- .../data/augmented/aug-ablation/aug-ablation-tox21.sh | 10 ++++++---- .../augmented/aug-ablation/aug-ablation-toxcast.sh | 10 ++++++---- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/configs/data/augmented/aug-ablation/aug-ablation-tox21.sh b/configs/data/augmented/aug-ablation/aug-ablation-tox21.sh index 756514c..8a66149 100644 --- a/configs/data/augmented/aug-ablation/aug-ablation-tox21.sh +++ b/configs/data/augmented/aug-ablation/aug-ablation-tox21.sh @@ -2,6 +2,7 @@ # ============================================================ # Local single-GPU adaptation of ../aug-ablation-script.sh +# Run this script using: `bash configs/data/augmented/aug-ablation/aug-ablation-tox21.sh` # # The original is a Slurm job-array script (10 data configs x 3 # seeds = 30 experiments, 10 concurrent). This version runs the @@ -68,7 +69,6 @@ export SSL_CERT_FILE=$(python -m certifi) echo "============================================================" echo "Node: $(hostname)" echo "Date: $(date)" -echo "TMPDIR: $TMPDIR" echo "GPU:" nvidia-smi echo "============================================================" @@ -96,6 +96,7 @@ for SEED in "${SEEDS[@]}"; do python -m chebai fit \ --trainer=configs/training/default_trainer.yml \ --trainer.logger=configs/training/wandb_logger.yml \ + --trainer.callbacks=configs/training/roc-auc_callbacks.yml \ --model="$CHEBAI_GRAPH_DIR/configs/model/baselines/gat.yml" \ --model.train_metrics=configs/metrics/micro-macro-f1-roc-auc.yml \ --model.test_metrics=configs/metrics/micro-macro-f1-roc-auc.yml \ @@ -105,13 +106,14 @@ for SEED in "${SEEDS[@]}"; do --trainer.accumulate_grad_batches=2 \ --data.init_args.num_workers=10 \ --model.pass_loss_kwargs=false \ - --trainer.min_epochs=200 \ - --trainer.max_epochs=200 \ + --trainer.min_epochs=100 \ + --trainer.max_epochs=100 \ --model.criterion=configs/loss/bce_unweighted.yml \ --trainer.logger.init_args.name="$RUN_NAME" \ --data.init_args.splits_file_path=data/Tox21:MNClassification/Tox21/processed/splits.csv \ --trainer.logger.init_args.tags='["augmented_paper","aug-ablation", "tox21"]' \ - --model.config.edge_dim=10 \ + --model.config.edge_dim=6 \ + --model.config.in_channels=77 \ --seed_everything="$SEED" echo "============================================================" diff --git a/configs/data/augmented/aug-ablation/aug-ablation-toxcast.sh b/configs/data/augmented/aug-ablation/aug-ablation-toxcast.sh index e355623..5ab542f 100644 --- a/configs/data/augmented/aug-ablation/aug-ablation-toxcast.sh +++ b/configs/data/augmented/aug-ablation/aug-ablation-toxcast.sh @@ -2,6 +2,7 @@ # ============================================================ # Local single-GPU adaptation of ../aug-ablation-script.sh +# # Run this script using: `bash configs/data/augmented/aug-ablation/aug-ablation-toxcast.sh` # # The original is a Slurm job-array script (10 data configs x 3 # seeds = 30 experiments, 10 concurrent). This version runs the @@ -68,7 +69,6 @@ export SSL_CERT_FILE=$(python -m certifi) echo "============================================================" echo "Node: $(hostname)" echo "Date: $(date)" -echo "TMPDIR: $TMPDIR" echo "GPU:" nvidia-smi echo "============================================================" @@ -96,6 +96,7 @@ for SEED in "${SEEDS[@]}"; do python -m chebai fit \ --trainer=configs/training/default_trainer.yml \ --trainer.logger=configs/training/wandb_logger.yml \ + --trainer.callbacks=configs/training/roc-auc_callbacks.yml \ --model="$CHEBAI_GRAPH_DIR/configs/model/baselines/gat.yml" \ --model.train_metrics=configs/metrics/micro-macro-f1-roc-auc.yml \ --model.test_metrics=configs/metrics/micro-macro-f1-roc-auc.yml \ @@ -105,13 +106,14 @@ for SEED in "${SEEDS[@]}"; do --trainer.accumulate_grad_batches=2 \ --data.init_args.num_workers=10 \ --model.pass_loss_kwargs=false \ - --trainer.min_epochs=200 \ - --trainer.max_epochs=200 \ + --trainer.min_epochs=100 \ + --trainer.max_epochs=100 \ --model.criterion=configs/loss/bce_unweighted.yml \ --trainer.logger.init_args.name="$RUN_NAME" \ --data.init_args.splits_file_path=data/ToxCast:MNClassification/ToxCast/processed/splits.csv \ --trainer.logger.init_args.tags='["augmented_paper","aug-ablation", "toxcast"]' \ - --model.config.edge_dim=10 \ + --model.config.edge_dim=6 \ + --model.config.in_channels=79 \ --seed_everything="$SEED" echo "============================================================" From 705000ef840a47dd3fd395082cbf1fe977d77b82 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sun, 6 Sep 2026 20:58:18 +0200 Subject: [PATCH 23/58] update chebi aug ablation script --- .../aug-ablation/aug-ablation-chebi.sh | 23 +++++++++++-------- .../augmented/aug-ablation/chebi/baseline.yml | 12 ++++++++++ 2 files changed, 25 insertions(+), 10 deletions(-) create mode 100644 configs/data/augmented/aug-ablation/chebi/baseline.yml diff --git a/configs/data/augmented/aug-ablation/aug-ablation-chebi.sh b/configs/data/augmented/aug-ablation/aug-ablation-chebi.sh index d19dd95..eb82031 100644 --- a/configs/data/augmented/aug-ablation/aug-ablation-chebi.sh +++ b/configs/data/augmented/aug-ablation/aug-ablation-chebi.sh @@ -51,17 +51,20 @@ SEEDS=(0 42 12345) # # ============================================================ +DATA_CONFIG_DIR=( + "$CHEBAI_GRAPH_DIR/configs/data/augmented/aug-ablation/tox21" +) DATA_CONFIGS=( - "../python-chebai-graph/configs/data/augmented/aug-ablation/FGN.yml" - "../python-chebai-graph/configs/data/augmented/aug-ablation/FGN+E.yml" - "../python-chebai-graph/configs/data/augmented/aug-ablation/FGN+E+WGN.yml" - "../python-chebai-graph/configs/data/augmented/aug-ablation/FGN+WGN.yml" - "../python-chebai-graph/configs/data/augmented/aug-ablation/gn_wall_fgwa_nfge.yml" - "../python-chebai-graph/configs/data/augmented/aug-ablation/gn_wall_fgwa_wfge.yml" - "../python-chebai-graph/configs/data/augmented/aug-ablation/gnwa_fgwa_nfge.yml" - "../python-chebai-graph/configs/data/augmented/aug-ablation/gnwa_fgwa_wfge.yml" - "../python-chebai-graph/configs/data/augmented/aug-ablation/WGN.yml" - "../python-chebai-graph/configs/data/chebi50_baseline.yml" + "$DATA_CONFIG_DIR/FGN.yml" + "$DATA_CONFIG_DIR/FGN+E.yml" + "$DATA_CONFIG_DIR/FGN+E+WGN.yml" + "$DATA_CONFIG_DIR/FGN+WGN.yml" + "$DATA_CONFIG_DIR/gn_wall_fgwa_nfge.yml" + "$DATA_CONFIG_DIR/gn_wall_fgwa_wfge.yml" + "$DATA_CONFIG_DIR/gnwa_fgwa_nfge.yml" + "$DATA_CONFIG_DIR/gnwa_fgwa_wfge.yml" + "$DATA_CONFIG_DIR/WGN.yml" + "$DATA_CONFIG_DIR/baseline.yml" ) # ============================================================ diff --git a/configs/data/augmented/aug-ablation/chebi/baseline.yml b/configs/data/augmented/aug-ablation/chebi/baseline.yml new file mode 100644 index 0000000..ce9dfa0 --- /dev/null +++ b/configs/data/augmented/aug-ablation/chebi/baseline.yml @@ -0,0 +1,12 @@ +class_path: chebai_graph.preprocessing.datasets.ChEBI50GraphProperties +init_args: + properties: + - chebai_graph.preprocessing.properties.AtomType + - chebai_graph.preprocessing.properties.NumAtomBonds + - chebai_graph.preprocessing.properties.AtomCharge + - chebai_graph.preprocessing.properties.AtomAromaticity + - chebai_graph.preprocessing.properties.AtomHybridization + - chebai_graph.preprocessing.properties.AtomNumHs + - chebai_graph.preprocessing.properties.BondType + - chebai_graph.preprocessing.properties.BondInRing + - chebai_graph.preprocessing.properties.BondAromaticity From 730bd70a344b7c042321c8252ca7a92b47ca9331 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Wed, 9 Sep 2026 20:24:32 +0200 Subject: [PATCH 24/58] final aug chang to gnwa_fgwa_wfge --- configs/data/augmented/final/chebi50_final_augmented.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/configs/data/augmented/final/chebi50_final_augmented.yml b/configs/data/augmented/final/chebi50_final_augmented.yml index 576cf75..3ff3232 100644 --- a/configs/data/augmented/final/chebi50_final_augmented.yml +++ b/configs/data/augmented/final/chebi50_final_augmented.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.ChEBI50_WFGE_WGN_AsPerNodeType +class_path: chebai_graph.preprocessing.datasets.ChEBI50_GNWA_FGWA_WFGE_AsPerNodeType init_args: properties: # All Node type properties From 5ca1a57a24234a58c35e6fb4b3a0e79199cdbfa2 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Thu, 10 Sep 2026 16:58:49 +0200 Subject: [PATCH 25/58] update molenet final config to gnwa_fgwa_wfge config --- .../preprocessing/datasets/__init__.py | 19 ++++++++++ .../datasets/molecule_net_classification.py | 36 +++++++++++++++++++ .../augmented/final/BACE_final_augmented.yml | 2 +- .../augmented/final/BBBP_final_augmented.yml | 2 +- .../final/ClinTox_final_augmented.yml | 2 +- .../augmented/final/HIV_final_augmented.yml | 2 +- .../augmented/final/MUV_final_augmented.yml | 2 +- .../augmented/final/PCBA_final_augmented.yml | 2 +- .../augmented/final/SIDER_final_augmented.yml | 2 +- .../augmented/final/Tox21_final_augmented.yml | 2 +- .../final/ToxCast_final_augmented.yml | 2 +- 11 files changed, 64 insertions(+), 9 deletions(-) diff --git a/chebai_graph/preprocessing/datasets/__init__.py b/chebai_graph/preprocessing/datasets/__init__.py index ad2aca9..24e9648 100644 --- a/chebai_graph/preprocessing/datasets/__init__.py +++ b/chebai_graph/preprocessing/datasets/__init__.py @@ -33,18 +33,26 @@ ChEBI100GraphProperties, ) from .molecule_net_classification import ( + BACE_GNWA_FGWA_WFGE_AsPerNodeType, BACE_WFGE_WGN_AsPerNodeType, + BBBP_GNWA_FGWA_WFGE_AsPerNodeType, BBBP_WFGE_WGN_AsPerNodeType, + ClinTox_GNWA_FGWA_WFGE_AsPerNodeType, ClinTox_WFGE_WGN_AsPerNodeType, + HIV_GNWA_FGWA_WFGE_AsPerNodeType, HIV_WFGE_WGN_AsPerNodeType, + MUV_GNWA_FGWA_WFGE_AsPerNodeType, MUV_WFGE_WGN_AsPerNodeType, + PCBA_GNWA_FGWA_WFGE_AsPerNodeType, PCBA_WFGE_WGN_AsPerNodeType, + SIDER_GNWA_FGWA_WFGE_AsPerNodeType, SIDER_WFGE_WGN_AsPerNodeType, Tox21_Atom_WGNOnly_GraphProp, Tox21_GN_WithAllNodes_FG_WithAtoms_FGE, Tox21_GN_WithAllNodes_FG_WithAtoms_NoFGE, Tox21_GN_WithAtoms_FG_WithAtoms_FGE, Tox21_GN_WithAtoms_FG_WithAtoms_NoFGE, + Tox21_GNWA_FGWA_WFGE_AsPerNodeType, Tox21_NFGE_NGN_GraphProp, Tox21_NFGE_WGN_GraphProp, Tox21_WFGE_NGN_GraphProp, @@ -56,6 +64,7 @@ ToxCast_GN_WithAllNodes_FG_WithAtoms_NoFGE, ToxCast_GN_WithAtoms_FG_WithAtoms_FGE, ToxCast_GN_WithAtoms_FG_WithAtoms_NoFGE, + ToxCast_GNWA_FGWA_WFGE_AsPerNodeType, ToxCast_NFGE_NGN_GraphProp, ToxCast_NFGE_WGN_GraphProp, ToxCast_WFGE_NGN_GraphProp, @@ -97,6 +106,16 @@ "Tox21_WFGE_WGN_AsPerNodeType", "ToxCast_WFGE_WGN_AsPerNodeType", "PCBA_WFGE_WGN_AsPerNodeType", + # ----- MoleculeNet Dataset GNWA+FGWA+WFGE Final classes ----- + "Tox21_GNWA_FGWA_WFGE_AsPerNodeType", + "ToxCast_GNWA_FGWA_WFGE_AsPerNodeType", + "BACE_GNWA_FGWA_WFGE_AsPerNodeType", + "BBBP_GNWA_FGWA_WFGE_AsPerNodeType", + "ClinTox_GNWA_FGWA_WFGE_AsPerNodeType", + "HIV_GNWA_FGWA_WFGE_AsPerNodeType", + "MUV_GNWA_FGWA_WFGE_AsPerNodeType", + "PCBA_GNWA_FGWA_WFGE_AsPerNodeType", + "SIDER_GNWA_FGWA_WFGE_AsPerNodeType", # ----- CHEBI GNWA_FGWA_WFGE Properties ablation classes ----- "ChEBI50_GNWA_FGWA_WFGE_AsPerNodeType", "ChEBI50_GNWA_FGWA_WFGE_ForAllPropsExceptNodeLevel", diff --git a/chebai_graph/preprocessing/datasets/molecule_net_classification.py b/chebai_graph/preprocessing/datasets/molecule_net_classification.py index d015d5f..2132965 100644 --- a/chebai_graph/preprocessing/datasets/molecule_net_classification.py +++ b/chebai_graph/preprocessing/datasets/molecule_net_classification.py @@ -36,38 +36,74 @@ class PCBA_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, PCBA): READER = AtomFGReader_WithFGEdges_WithGraphNode +class PCBA_GNWA_FGWA_WFGE_AsPerNodeType(GraphPropAsPerNodeType, PCBA): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class BACE_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, BACE): READER = AtomFGReader_WithFGEdges_WithGraphNode +class BACE_GNWA_FGWA_WFGE_AsPerNodeType(GraphPropAsPerNodeType, BACE): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class BBBP_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, BBBP): READER = AtomFGReader_WithFGEdges_WithGraphNode +class BBBP_GNWA_FGWA_WFGE_AsPerNodeType(GraphPropAsPerNodeType, BBBP): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class ClinTox_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, ClinTox): READER = AtomFGReader_WithFGEdges_WithGraphNode +class ClinTox_GNWA_FGWA_WFGE_AsPerNodeType(GraphPropAsPerNodeType, ClinTox): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class HIV_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, HIV): READER = AtomFGReader_WithFGEdges_WithGraphNode +class HIV_GNWA_FGWA_WFGE_AsPerNodeType(GraphPropAsPerNodeType, HIV): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class SIDER_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, SIDER): READER = AtomFGReader_WithFGEdges_WithGraphNode +class SIDER_GNWA_FGWA_WFGE_AsPerNodeType(GraphPropAsPerNodeType, SIDER): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class MUV_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, MUV): READER = AtomFGReader_WithFGEdges_WithGraphNode +class MUV_GNWA_FGWA_WFGE_AsPerNodeType(GraphPropAsPerNodeType, MUV): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class Tox21_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, Tox21): READER = AtomFGReader_WithFGEdges_WithGraphNode +class Tox21_GNWA_FGWA_WFGE_AsPerNodeType(GraphPropAsPerNodeType, Tox21): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + class ToxCast_WFGE_WGN_AsPerNodeType(GraphPropAsPerNodeType, ToxCast): READER = AtomFGReader_WithFGEdges_WithGraphNode +class ToxCast_GNWA_FGWA_WFGE_AsPerNodeType(GraphPropAsPerNodeType, ToxCast): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + # ---- Augmentation: Variants with graph Node connected to FG nodes only ------------- class Tox21_WFGE_WGN_GraphProp(AugGraphPropMixIn_WithGraphNode, Tox21): """Tox21 with with FG nodes and FG edges and graph node.""" diff --git a/configs/data/augmented/final/BACE_final_augmented.yml b/configs/data/augmented/final/BACE_final_augmented.yml index 7ea67f0..fc95a55 100644 --- a/configs/data/augmented/final/BACE_final_augmented.yml +++ b/configs/data/augmented/final/BACE_final_augmented.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.BACE_WFGE_WGN_AsPerNodeType +class_path: chebai_graph.preprocessing.datasets.BACE_GNWA_FGWA_WFGE_AsPerNodeType init_args: properties: # All Node type properties diff --git a/configs/data/augmented/final/BBBP_final_augmented.yml b/configs/data/augmented/final/BBBP_final_augmented.yml index 9a7ee90..fdef5ca 100644 --- a/configs/data/augmented/final/BBBP_final_augmented.yml +++ b/configs/data/augmented/final/BBBP_final_augmented.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.BBBP_WFGE_WGN_AsPerNodeType +class_path: chebai_graph.preprocessing.datasets.BBBP_GNWA_FGWA_WFGE_AsPerNodeType init_args: properties: # All Node type properties diff --git a/configs/data/augmented/final/ClinTox_final_augmented.yml b/configs/data/augmented/final/ClinTox_final_augmented.yml index a4023b3..774103a 100644 --- a/configs/data/augmented/final/ClinTox_final_augmented.yml +++ b/configs/data/augmented/final/ClinTox_final_augmented.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.ClinTox_WFGE_WGN_AsPerNodeType +class_path: chebai_graph.preprocessing.datasets.ClinTox_GNWA_FGWA_WFGE_AsPerNodeType init_args: properties: # All Node type properties diff --git a/configs/data/augmented/final/HIV_final_augmented.yml b/configs/data/augmented/final/HIV_final_augmented.yml index 596fc19..126a919 100644 --- a/configs/data/augmented/final/HIV_final_augmented.yml +++ b/configs/data/augmented/final/HIV_final_augmented.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.HIV_WFGE_WGN_AsPerNodeType +class_path: chebai_graph.preprocessing.datasets.HIV_GNWA_FGWA_WFGE_AsPerNodeType init_args: properties: # All Node type properties diff --git a/configs/data/augmented/final/MUV_final_augmented.yml b/configs/data/augmented/final/MUV_final_augmented.yml index 56a3e1e..a3ade9c 100644 --- a/configs/data/augmented/final/MUV_final_augmented.yml +++ b/configs/data/augmented/final/MUV_final_augmented.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.MUV_WFGE_WGN_AsPerNodeType +class_path: chebai_graph.preprocessing.datasets.MUV_GNWA_FGWA_WFGE_AsPerNodeType init_args: properties: # All Node type properties diff --git a/configs/data/augmented/final/PCBA_final_augmented.yml b/configs/data/augmented/final/PCBA_final_augmented.yml index 40688f4..dbacc76 100644 --- a/configs/data/augmented/final/PCBA_final_augmented.yml +++ b/configs/data/augmented/final/PCBA_final_augmented.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.PCBA_WFGE_WGN_AsPerNodeType +class_path: chebai_graph.preprocessing.datasets.PCBA_GNWA_FGWA_WFGE_AsPerNodeType init_args: properties: # All Node type properties diff --git a/configs/data/augmented/final/SIDER_final_augmented.yml b/configs/data/augmented/final/SIDER_final_augmented.yml index 5eef128..bcb071c 100644 --- a/configs/data/augmented/final/SIDER_final_augmented.yml +++ b/configs/data/augmented/final/SIDER_final_augmented.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.SIDER_WFGE_WGN_AsPerNodeType +class_path: chebai_graph.preprocessing.datasets.SIDER_GNWA_FGWA_WFGE_AsPerNodeType init_args: properties: # All Node type properties diff --git a/configs/data/augmented/final/Tox21_final_augmented.yml b/configs/data/augmented/final/Tox21_final_augmented.yml index 1c88112..cebd0dc 100644 --- a/configs/data/augmented/final/Tox21_final_augmented.yml +++ b/configs/data/augmented/final/Tox21_final_augmented.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.Tox21_WFGE_WGN_AsPerNodeType +class_path: chebai_graph.preprocessing.datasets.Tox21_GNWA_FGWA_WFGE_AsPerNodeType init_args: properties: # All Node type properties diff --git a/configs/data/augmented/final/ToxCast_final_augmented.yml b/configs/data/augmented/final/ToxCast_final_augmented.yml index 4831ae5..88f1a17 100644 --- a/configs/data/augmented/final/ToxCast_final_augmented.yml +++ b/configs/data/augmented/final/ToxCast_final_augmented.yml @@ -1,4 +1,4 @@ -class_path: chebai_graph.preprocessing.datasets.ToxCast_WFGE_WGN_AsPerNodeType +class_path: chebai_graph.preprocessing.datasets.ToxCast_GNWA_FGWA_WFGE_AsPerNodeType init_args: properties: # All Node type properties From 1b69b71fa9f3100802397258fdee0f6e81f76c2d Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Thu, 10 Sep 2026 16:59:14 +0200 Subject: [PATCH 26/58] enable seperate learning rate for gnn and final classification layer --- chebai_graph/models/architectures/base.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/chebai_graph/models/architectures/base.py b/chebai_graph/models/architectures/base.py index c86baa8..64dca7c 100644 --- a/chebai_graph/models/architectures/base.py +++ b/chebai_graph/models/architectures/base.py @@ -196,3 +196,25 @@ def forward(self, batch: dict) -> torch.Tensor: if self.use_batch_norm: a = self.batch_norm(a) return self.lin_sequential(a) + + def configure_optimizers(self, **kwargs) -> torch.optim.Optimizer: + optimizer_kwargs = dict(self.optimizer_kwargs) + + default_lr = optimizer_kwargs.pop("lr", None) + gnn_lr = optimizer_kwargs.pop("lr_gnn", default_lr) + linear_lr = optimizer_kwargs.pop("lr_linear", default_lr) + + if gnn_lr is None or linear_lr is None: + raise ValueError("Set lr or both lr_gnn and lr_linear") + + linear_params = list(self.lin_sequential.parameters()) + if self.use_batch_norm: + linear_params.extend(self.batch_norm.parameters()) + + return torch.optim.Adamax( + [ + {"params": self.gnn.parameters(), "lr": gnn_lr}, + {"params": linear_params, "lr": linear_lr}, + ], + **optimizer_kwargs, + ) From 2d6c0a9ffb54de83d066e2196bc1e9a8a0139d80 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Wed, 7 Oct 2026 12:54:05 +0200 Subject: [PATCH 27/58] allow pretrained weights with mismatchning edges --- chebai_graph/models/architectures/base.py | 62 +++++++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/chebai_graph/models/architectures/base.py b/chebai_graph/models/architectures/base.py index 64dca7c..9b5fd81 100644 --- a/chebai_graph/models/architectures/base.py +++ b/chebai_graph/models/architectures/base.py @@ -91,6 +91,7 @@ def __init__( config: dict, n_linear_layers: int, use_batch_norm: bool = False, + pretrained_checkpoint: str | None = None, **kwargs, ): """ @@ -101,6 +102,8 @@ def __init__( """ super().__init__(**kwargs) self.gnn = self._get_gnn(config) + if pretrained_checkpoint is not None: + self._load_pretrained_gnn(pretrained_checkpoint) gnn_out_dim = int(config["out_channels"]) self.activation = torch.nn.ELU self.lin_input_dim = self._get_lin_seq_input_dim( @@ -118,6 +121,65 @@ def __init__( out_dim=self.out_dim, ) + def _load_pretrained_gnn(self, checkpoint_path: str) -> None: + """Load encoder weights while leaving the task-specific head untouched.""" + checkpoint = torch.load(checkpoint_path, weights_only=False) + state_dict = checkpoint.get("state_dict", checkpoint) + + if not isinstance(state_dict, dict): + raise ValueError( + f"Checkpoint {checkpoint_path!r} does not contain a state dictionary." + ) + + gnn_state = self.gnn.state_dict() + gnn_keys = set(gnn_state) + matching_keys = set(state_dict) & gnn_keys + if not matching_keys: + for prefix in ("gnn.", "model.gnn."): + prefixed_keys = { + key.removeprefix(prefix): value + for key, value in state_dict.items() + if key.startswith(prefix) + } + if set(prefixed_keys) & gnn_keys: + state_dict = prefixed_keys + matching_keys = set(state_dict) & gnn_keys + break + + if not matching_keys: + raise ValueError( + f"Checkpoint {checkpoint_path!r} contains no weights matching the GNN." + ) + + pretrained_gnn_state = { + key: value for key, value in state_dict.items() if key in gnn_keys + } + skipped_edge_weights = [] + for key in list(pretrained_gnn_state): + if ( + key.endswith(".lin_edge.weight") + and pretrained_gnn_state[key].shape != gnn_state[key].shape + ): + skipped_edge_weights.append(key) + del pretrained_gnn_state[key] + + load_result = self.gnn.load_state_dict(pretrained_gnn_state, strict=False) + allowed_missing = set(skipped_edge_weights) + unexpected_keys = set(load_result.unexpected_keys) + missing_keys = set(load_result.missing_keys) - allowed_missing + if unexpected_keys or missing_keys: + raise RuntimeError( + "Pretrained GNN checkpoint is incompatible: " + f"missing={sorted(missing_keys)}, " + f"unexpected={sorted(unexpected_keys)}" + ) + + if skipped_edge_weights: + print( + "Skipped incompatible pretrained edge weights: " + + ", ".join(sorted(skipped_edge_weights)) + ) + @abstractmethod def _get_gnn(self, config: dict) -> torch.nn.Module: """ From ed9075910ee57a2bccc24961cff37907fe928e75 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Fri, 9 Oct 2026 12:11:58 +0200 Subject: [PATCH 28/58] update aug ablation script for tox --- .../data/augmented/aug-ablation/aug-ablation-tox21.sh | 10 ++++++---- .../augmented/aug-ablation/aug-ablation-toxcast.sh | 10 ++++++---- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/configs/data/augmented/aug-ablation/aug-ablation-tox21.sh b/configs/data/augmented/aug-ablation/aug-ablation-tox21.sh index 8a66149..bf2c8c2 100644 --- a/configs/data/augmented/aug-ablation/aug-ablation-tox21.sh +++ b/configs/data/augmented/aug-ablation/aug-ablation-tox21.sh @@ -96,14 +96,14 @@ for SEED in "${SEEDS[@]}"; do python -m chebai fit \ --trainer=configs/training/default_trainer.yml \ --trainer.logger=configs/training/wandb_logger.yml \ + --trainer.logger.init_args.tags='["augmented_paper","aug-ablation", "tox21"]' \ --trainer.callbacks=configs/training/roc-auc_callbacks.yml \ --model="$CHEBAI_GRAPH_DIR/configs/model/baselines/gat.yml" \ --model.train_metrics=configs/metrics/micro-macro-f1-roc-auc.yml \ --model.test_metrics=configs/metrics/micro-macro-f1-roc-auc.yml \ --model.val_metrics=configs/metrics/micro-macro-f1-roc-auc.yml \ --data="$DATA_CONFIG" \ - --data.init_args.batch_size=64 \ - --trainer.accumulate_grad_batches=2 \ + --data.init_args.batch_size=32 \ --data.init_args.num_workers=10 \ --model.pass_loss_kwargs=false \ --trainer.min_epochs=100 \ @@ -111,9 +111,11 @@ for SEED in "${SEEDS[@]}"; do --model.criterion=configs/loss/bce_unweighted.yml \ --trainer.logger.init_args.name="$RUN_NAME" \ --data.init_args.splits_file_path=data/Tox21:MNClassification/Tox21/processed/splits.csv \ - --trainer.logger.init_args.tags='["augmented_paper","aug-ablation", "tox21"]' \ --model.config.edge_dim=6 \ --model.config.in_channels=77 \ + --model.config.dropout=0.7 \ + --model.optimizer_kwargs.lr_gnn=0.0005 \ + --model.optimizer_kwargs.lr_linear=0.001 \ --seed_everything="$SEED" echo "============================================================" @@ -124,4 +126,4 @@ done echo "============================================================" echo "All experiments finished successfully" -echo "============================================================" \ No newline at end of file +echo "============================================================" diff --git a/configs/data/augmented/aug-ablation/aug-ablation-toxcast.sh b/configs/data/augmented/aug-ablation/aug-ablation-toxcast.sh index 5ab542f..a4de723 100644 --- a/configs/data/augmented/aug-ablation/aug-ablation-toxcast.sh +++ b/configs/data/augmented/aug-ablation/aug-ablation-toxcast.sh @@ -96,14 +96,14 @@ for SEED in "${SEEDS[@]}"; do python -m chebai fit \ --trainer=configs/training/default_trainer.yml \ --trainer.logger=configs/training/wandb_logger.yml \ + --trainer.logger.init_args.tags='["augmented_paper","aug-ablation", "toxcast"]' \ --trainer.callbacks=configs/training/roc-auc_callbacks.yml \ --model="$CHEBAI_GRAPH_DIR/configs/model/baselines/gat.yml" \ --model.train_metrics=configs/metrics/micro-macro-f1-roc-auc.yml \ --model.test_metrics=configs/metrics/micro-macro-f1-roc-auc.yml \ --model.val_metrics=configs/metrics/micro-macro-f1-roc-auc.yml \ --data="$DATA_CONFIG" \ - --data.init_args.batch_size=64 \ - --trainer.accumulate_grad_batches=2 \ + --data.init_args.batch_size=32 \ --data.init_args.num_workers=10 \ --model.pass_loss_kwargs=false \ --trainer.min_epochs=100 \ @@ -111,9 +111,11 @@ for SEED in "${SEEDS[@]}"; do --model.criterion=configs/loss/bce_unweighted.yml \ --trainer.logger.init_args.name="$RUN_NAME" \ --data.init_args.splits_file_path=data/ToxCast:MNClassification/ToxCast/processed/splits.csv \ - --trainer.logger.init_args.tags='["augmented_paper","aug-ablation", "toxcast"]' \ --model.config.edge_dim=6 \ --model.config.in_channels=79 \ + --model.config.dropout=0.7 \ + --model.optimizer_kwargs.lr_gnn=0.001 \ + --model.optimizer_kwargs.lr_linear=0.001 \ --seed_everything="$SEED" echo "============================================================" @@ -124,4 +126,4 @@ done echo "============================================================" echo "All experiments finished successfully" -echo "============================================================" \ No newline at end of file +echo "============================================================" From 5ea28d249eadf0c625663c61b2f51ee5df810887 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sun, 11 Oct 2026 09:31:45 +0200 Subject: [PATCH 29/58] himol msp pretraining --- .gitignore | 1 + chebai_graph/loss/msp.py | 278 ++++++++++++++++++++ chebai_graph/models/__init__.py | 2 + chebai_graph/models/pretraining/__init__.py | 3 + chebai_graph/models/pretraining/msp.py | 75 ++++++ chebai_graph/preprocessing/collate.py | 6 + 6 files changed, 365 insertions(+) create mode 100644 chebai_graph/loss/msp.py create mode 100644 chebai_graph/models/pretraining/__init__.py create mode 100644 chebai_graph/models/pretraining/msp.py diff --git a/.gitignore b/.gitignore index 95daf78..2a04373 100644 --- a/.gitignore +++ b/.gitignore @@ -176,3 +176,4 @@ electra_pretrained.ckpt *.ckpt !configs/data/augmented/aug-ablation/*.sh +!.vscode/launch.json diff --git a/chebai_graph/loss/msp.py b/chebai_graph/loss/msp.py new file mode 100644 index 0000000..20055b8 --- /dev/null +++ b/chebai_graph/loss/msp.py @@ -0,0 +1,278 @@ +""" +Multi-level Self-supervised Pre-training from +Zang, X., Zhao, X. & Tang, B. Hierarchical Molecular Graph Self-Supervised Learning for property prediction. Commun Chem 6, 34 (2023). https://doi.org/10.1038/s42004-023-00825-5 + +The files containes modified code from: + - https://github.com/ZangXuan/HiMol/blob/main/pretrain.py +""" + +import torch +from sklearn.metrics import average_precision_score, roc_auc_score +from torch_geometric.data import Data + + +class MSPPretrainingLoss(torch.nn.Module): + """HiMol-style molecular structure pretraining loss. + + The model output must be ``(graph_data, node_embeddings)``. The graph data + must contain ``is_atom_node`` and a PyG ``batch`` vector. If present, + ``is_graph_node`` identifies the molecule-level representation. + """ + + def __init__( + self, + hidden_size: int, + atom_num_classes: int = 118, + bond_num_classes: int = 4, + atom_target_start: int = 0, + atom_target_width: int = 1, + bond_target_start: int = 0, + bond_target_width: int = 1, + dropout: float = 0.2, + ) -> None: + super().__init__() + self.atom_target_start = atom_target_start + self.atom_target_width = atom_target_width + self.bond_target_start = bond_target_start + self.bond_target_width = bond_target_width + + # Bond existence projection + self.bond_if_proj = torch.nn.Sequential( + torch.nn.Linear(hidden_size, hidden_size), + torch.nn.ReLU(), + torch.nn.Linear(hidden_size, hidden_size), + ) + # Bond existence classifier + self.bond_if_s = torch.nn.Sequential( + torch.nn.Linear(2 * hidden_size, hidden_size), + torch.nn.ReLU(), + torch.nn.Linear(hidden_size, 1), + ) + + # Bond type classifier + self.bond_type_s = torch.nn.Sequential( + torch.nn.Linear(2 * hidden_size, hidden_size), + torch.nn.ReLU(), + torch.nn.Linear(hidden_size, bond_num_classes), + ) + # Atom type classifier + self.atom_type_s = torch.nn.Sequential( + torch.nn.Linear(hidden_size, hidden_size), + torch.nn.ReLU(), + torch.nn.Linear(hidden_size, atom_num_classes), + ) + + count_hidden = max(hidden_size // 4, 1) + # Atom count regressor + self.atom_num_s = torch.nn.Sequential( + torch.nn.Linear(hidden_size, count_hidden), + torch.nn.Softplus(), + torch.nn.Linear(count_hidden, 1), + ) + # Bond count regressor + self.bond_num_s = torch.nn.Sequential( + torch.nn.Linear(hidden_size, count_hidden), + torch.nn.Softplus(), + torch.nn.Linear(count_hidden, 1), + ) + self.feat_drop = ( + torch.nn.Dropout(dropout) if dropout > 0 else torch.nn.Identity() + ) + self.bond_if_loss = torch.nn.BCEWithLogitsLoss() + self.bond_type_loss = torch.nn.CrossEntropyLoss() + self.atom_type_loss = torch.nn.CrossEntropyLoss() + self.count_loss = torch.nn.SmoothL1Loss() + + self.loss_weight = torch.nn.Parameter(torch.rand(5), requires_grad=True) + + @staticmethod + def _bond_count(edge_index: torch.Tensor) -> int: + if edge_index.numel() == 0: + return 0 + return int(edge_index.t().sort(dim=1).values.unique(dim=0).shape[0]) + + def _super( + self, data: Data, embeddings: torch.Tensor, graphs: list[torch.Tensor] + ) -> torch.Tensor: + if not hasattr(data, "is_graph_node") or data.is_graph_node is None: + raise ValueError( + "MSP pretraining requires a boolean is_graph_node attribute in the graph data" + ) + + graph_mask = data.is_graph_node.bool() + graph_representations = [] + for graph_index in range(len(graphs)): + representations = embeddings[graph_mask & (data.batch == graph_index)] + if representations.numel() == 0: + raise ValueError( + f"Graph {graph_index} has no node marked is_graph_node" + ) + if len(representations) > 1: + raise ValueError( + f"Graph {graph_index} has multiple nodes marked is_graph_node" + ) + graph_representations.append(representations[0]) + return torch.stack(graph_representations) + + @staticmethod + def _molecule_metadata(custom_metadata, graph_index: int) -> dict: + metadata = custom_metadata[graph_index] + if isinstance(metadata, list): + metadata = next( + ( + item + for item in metadata + if isinstance(item, dict) and "nodes" in item + ), + None, + ) + if ( + not isinstance(metadata, dict) + or "nodes" not in metadata + or "edges" not in metadata + ): + raise ValueError( + f"Missing molecular metadata for graph {graph_index}; expected nodes and edges" + ) + return metadata + + @staticmethod + def _metadata_targets(metadata: dict, device: torch.device): + atom_vocab = list(range(1, 119)) # Atomic numbers from 1 to 118 + bond_vocab = ["SINGLE", "DOUBLE", "AROMATIC", "TRIPLE", "DATIVE", "UNSPECIFIED"] + atom_mol = metadata["nodes"]["atom_nodes"] + atom_targets = torch.tensor( + [atom_vocab.index(atom.GetAtomicNum()) for atom in atom_mol.GetAtoms()], + dtype=torch.long, + device=device, + ) + within_atoms = metadata["edges"]["within_atoms_lvl"] + bond_indices = [] + bond_targets = [] + for bond in within_atoms.GetBonds(): + begin = bond.GetBeginAtomIdx() + end = bond.GetEndAtomIdx() + bond_class = bond_vocab.index(str(bond.GetBondType())) + bond_indices.extend(((begin, end), (end, begin))) + bond_targets.extend((bond_class, bond_class)) + edge_index = ( + torch.tensor(bond_indices, dtype=torch.long, device=device).t() + if bond_indices + else torch.empty((2, 0), dtype=torch.long, device=device) + ) + return ( + atom_targets, + torch.tensor(bond_targets, dtype=torch.long, device=device), + edge_index, + ) + + def forward( + self, output, target, features, custom_metadata, **loss_kwargs + ) -> torch.Tensor: + node_embeddings = output + graph_data = features[0] + if not isinstance(graph_data, Data): + raise TypeError("features[0] must be a torch_geometric.data.Data object") + if graph_data.batch is None or not hasattr(graph_data, "is_atom_node"): + raise ValueError("MSP pretraining requires batch and is_atom_node") + is_atom_node = graph_data.is_atom_node.bool() + batch_index = graph_data.batch.long() + graph_nodes = [ + torch.where(batch_index == graph_index)[0] + for graph_index in range(int(batch_index.max().item()) + 1) + ] + graph_node_embeddings = self._super(graph_data, node_embeddings, graph_nodes) + + atom_counts = torch.stack( + [ + (is_atom_node & (batch_index == graph_index)).sum() + for graph_index in range(len(graph_nodes)) + ] + ).to(node_embeddings.dtype) + atom_count_pred = self.atom_num_s(graph_node_embeddings).squeeze(-1) + atom_count_loss = self.count_loss(atom_count_pred, atom_counts) + + bond_if_loss = node_embeddings.new_zeros(()) + bond_type_loss = node_embeddings.new_zeros(()) + atom_type_loss = node_embeddings.new_zeros(()) + bond_counts = [] + atom_correct = bond_correct = 0 + atom_total = bond_total = 0 + auc_values, ap_values = [], [] + for graph_index, nodes in enumerate(graph_nodes): + metadata = self._molecule_metadata(custom_metadata, graph_index) + mol = metadata["nodes"]["atom_nodes"] + atom_targets, bond_targets, local_edges = self._metadata_targets( + metadata, node_embeddings.device + ) + # Bond existence prediction + atoms = nodes[is_atom_node[nodes]] + if atoms.numel() != atom_targets.numel(): + raise ValueError( + f"Graph {graph_index} has {atoms.numel()} atom nodes but metadata has " + f"{atom_targets.numel()} atoms" + ) + atom_embeddings = self.feat_drop(self.bond_if_proj(node_embeddings[atoms])) + bond_counts.append(mol.GetNumBonds()) + n_atoms = atoms.numel() + pairs = torch.cat( + [ + atom_embeddings[:, None, :].expand(-1, n_atoms, -1), + atom_embeddings[None, :, :].expand(n_atoms, -1, -1), + ], + dim=-1, + ) + bond_pred = self.bond_if_s(pairs).squeeze(-1) + bond_target = torch.zeros_like(bond_pred) + if local_edges.numel(): + bond_target[local_edges[0], local_edges[1]] = 1 + bond_if_loss = bond_if_loss + self.bond_if_loss(bond_pred, bond_target) + if loss_kwargs.get("compute_metrics", True) and local_edges.numel(): + target_cpu = bond_target.flatten().detach().cpu() + pred_cpu = bond_pred.flatten().detach().sigmoid().cpu() + if target_cpu.unique().numel() > 1: + auc_values.append(roc_auc_score(target_cpu, pred_cpu)) + ap_values.append(average_precision_score(target_cpu, pred_cpu)) + + atom_pred = self.atom_type_s(node_embeddings[atoms]) + atom_type_loss = atom_type_loss + self.atom_type_loss( + atom_pred, atom_targets + ) + atom_correct += int((atom_pred.argmax(1) == atom_targets).sum()) + atom_total += atom_targets.numel() + if local_edges.numel(): + edge_rep = torch.cat( + [atom_embeddings[local_edges[0]], atom_embeddings[local_edges[1]]], + dim=1, + ) + bond_pred_type = self.bond_type_s(edge_rep) + bond_type_loss = bond_type_loss + self.bond_type_loss( + bond_pred_type, bond_targets + ) + bond_correct += int((bond_pred_type.argmax(1) == bond_targets).sum()) + bond_total += bond_targets.numel() + + bond_counts = torch.tensor( + bond_counts, + dtype=graph_node_embeddings.dtype, + device=graph_node_embeddings.device, + ) + bond_count_pred = self.bond_num_s(graph_node_embeddings).squeeze(-1) + bond_count_loss = self.count_loss(bond_count_pred, bond_counts) + divisor = max(len(graph_nodes), 1) + + losses = torch.stack( + [ + bond_if_loss / divisor, + bond_type_loss / divisor, + atom_type_loss / divisor, + atom_count_loss, + bond_count_loss, + ] + ) + if self.loss_weight.numel() != losses.numel(): + raise RuntimeError( + f"Expected {losses.numel()} loss weights, got {self.loss_weight.numel()}" + ) + loss = (losses * torch.softmax(self.loss_weight, dim=0)).sum() + return loss diff --git a/chebai_graph/models/__init__.py b/chebai_graph/models/__init__.py index 515a513..27a4b64 100644 --- a/chebai_graph/models/__init__.py +++ b/chebai_graph/models/__init__.py @@ -10,6 +10,7 @@ ResGatedAMGPoolGraphPred, ) from .dynamic_gni import ResGatedDynamicGNIGraphPred +from .pretraining import GATMSPretraining __all__ = [ "ResGatedGraphPred", @@ -22,4 +23,5 @@ "GINEGraphPred", "GINEAAPoolGraphPred", "GINEAMGPoolGraphPred", + "GATMSPretraining", ] diff --git a/chebai_graph/models/pretraining/__init__.py b/chebai_graph/models/pretraining/__init__.py new file mode 100644 index 0000000..5e7e638 --- /dev/null +++ b/chebai_graph/models/pretraining/__init__.py @@ -0,0 +1,3 @@ +from .msp import GATMSPretraining + +__all__ = ["GATMSPretraining"] diff --git a/chebai_graph/models/pretraining/msp.py b/chebai_graph/models/pretraining/msp.py new file mode 100644 index 0000000..1f252d6 --- /dev/null +++ b/chebai_graph/models/pretraining/msp.py @@ -0,0 +1,75 @@ +from typing import Any + +import torch +from chebai.preprocessing.structures import XYData +from torch_geometric.data import Data as GraphData + +from chebai_graph.models.architectures.base import GraphBaseNet +from chebai_graph.models.architectures.gat import GATGraphConvNetBase + + +class GATMSPretraining(GraphBaseNet): + def __init__( + self, + config: dict, + **kwargs, + ): + super().__init__(**kwargs) + self.gnn = GATGraphConvNetBase(config=config) + + def _process_batch(self, batch: XYData, batch_idx: int) -> dict[str, Any]: + """ + Processes the batch data. + + Args: + batch (XYData): The input batch of data. + batch_idx (int): The index of the current batch. + + Returns: + Dict[str, Any]: Processed batch data. + """ + return dict( + features=batch.x, + labels=self._process_labels_in_batch(batch), + model_kwargs=batch.additional_fields["model_kwargs"], + loss_kwargs=batch.additional_fields["loss_kwargs"], + idents=batch.additional_fields["idents"], + custom_metadata=batch.additional_fields["custom_metadata"], + ) + + def forward(self, batch: dict) -> torch.Tensor: + """ + Forward pass through GNN, pooling and linear layers. + + Args: + batch (dict): Input batch with graph features. + + Returns: + torch.Tensor: Predicted output. + """ + graph_data = batch["features"][0] + graph_data.to(self.device) + assert isinstance(graph_data, GraphData) + return self.gnn(batch) + + def _process_for_loss( + self, + model_output: torch.Tensor, + labels: torch.Tensor, + data: dict[str, Any], + ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: + """ + Processes the data for loss computation. + + Args: + model_output (torch.Tensor): The model output. + labels (torch.Tensor): The true labels. + data (dict[str, Any]): The processed batch data. + + Returns: + Tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: Model output, labels, and data. + """ + data["valid_label_mask"] = data.get("loss_kwargs", dict()).get( + "valid_label_mask", None + ) + return model_output, labels, data diff --git a/chebai_graph/preprocessing/collate.py b/chebai_graph/preprocessing/collate.py index b2bc4a0..51b1498 100644 --- a/chebai_graph/preprocessing/collate.py +++ b/chebai_graph/preprocessing/collate.py @@ -26,11 +26,16 @@ def __call__(self, data): del features.y merged_data.append(features) + custom_metadata = [] # Add empty edge_attr for graphs with no edges to prevent PyG errors for mdata in merged_data: + graph_metadata = [] for store in mdata.stores: if "edge_attr" not in store: store["edge_attr"] = torch.tensor([]) + if "custom_metadata" in store: + graph_metadata.append(store.pop("custom_metadata", None)) + custom_metadata.append(graph_metadata) # Ensure all attributes are float tensors to prevent torch.cat dtype issues for attr in merged_data[0].keys(): @@ -86,4 +91,5 @@ def __call__(self, data): idents=idents, model_kwargs={}, loss_kwargs=loss_kwargs, + custom_metadata=custom_metadata, ) From 0c3367ee1d968eb9e6eeba3019bdba5146c7eebd Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sun, 11 Oct 2026 09:35:06 +0200 Subject: [PATCH 30/58] zinc pretraining dataset --- .../preprocessing/datasets/__init__.py | 2 ++ chebai_graph/preprocessing/datasets/base.py | 8 ++++--- .../preprocessing/datasets/pretrain_zinc15.py | 18 ++++++++++++++ configs/data/augmented/pretraining/zinc15.yml | 24 +++++++++++++++++++ configs/model/pretraining/gat_msp.yml | 19 +++++++++++++++ 5 files changed, 68 insertions(+), 3 deletions(-) create mode 100644 chebai_graph/preprocessing/datasets/pretrain_zinc15.py create mode 100644 configs/data/augmented/pretraining/zinc15.yml create mode 100644 configs/model/pretraining/gat_msp.yml diff --git a/chebai_graph/preprocessing/datasets/__init__.py b/chebai_graph/preprocessing/datasets/__init__.py index 24e9648..d0b7b72 100644 --- a/chebai_graph/preprocessing/datasets/__init__.py +++ b/chebai_graph/preprocessing/datasets/__init__.py @@ -72,6 +72,7 @@ ToxCast_WFGE_WGN_GraphProp, ToxCastGraphProperties, ) +from .pretrain_zinc15 import Zinc15Pretrain_GNWA_FGWA_WFGE_AsPerNodeType from .pubchem import PubChemGraphProperties __all__ = [ @@ -154,4 +155,5 @@ "ChEBI100GraphProperties", "ChEBI50GraphData", "PubChemGraphProperties", + "Zinc15Pretrain_GNWA_FGWA_WFGE_AsPerNodeType", ] diff --git a/chebai_graph/preprocessing/datasets/base.py b/chebai_graph/preprocessing/datasets/base.py index 84d42be..d94a333 100644 --- a/chebai_graph/preprocessing/datasets/base.py +++ b/chebai_graph/preprocessing/datasets/base.py @@ -571,9 +571,10 @@ def _merge_props_into_base( if geom_data is None: return None if isinstance(geom_data, tuple): - geom_data = geom_data[ - 0 - ] # ignore additional returned data from _read_data (e.g. augmented molecule dict) + # include custom metadata returned data from _read_data (e.g. augmented molecule dict) + custom_metadata = geom_data[1] + geom_data = geom_data[0] + geom_data["custom_metadata"] = custom_metadata assert isinstance(geom_data, GeomData) is_atom_node = geom_data.is_atom_node @@ -643,6 +644,7 @@ def _merge_props_into_base( is_atom_node=is_atom_node, is_fg_node=is_fg_node, is_graph_node=is_graph_node, + custom_metadata=geom_data.custom_metadata, ) def _build_node_property_tensor( diff --git a/chebai_graph/preprocessing/datasets/pretrain_zinc15.py b/chebai_graph/preprocessing/datasets/pretrain_zinc15.py new file mode 100644 index 0000000..3082899 --- /dev/null +++ b/chebai_graph/preprocessing/datasets/pretrain_zinc15.py @@ -0,0 +1,18 @@ +from chebai.preprocessing.datasets.pretrain_zinc15 import Zinc15Pretrain + +from chebai_graph.preprocessing.datasets.base import GraphPropAsPerNodeType +from chebai_graph.preprocessing.reader import GN_WithAtoms_FG_WithAtoms_FGE + + +class Zinc15Pretrain_GNWA_FGWA_WFGE_AsPerNodeType( + GraphPropAsPerNodeType, Zinc15Pretrain +): + READER = GN_WithAtoms_FG_WithAtoms_FGE + + +if __name__ == "__main__": + dataset = Zinc15Pretrain_GNWA_FGWA_WFGE_AsPerNodeType( + root="data/Zinc15", name="Zinc15" + ) + dataset.prepare_data() + dataset.setup() diff --git a/configs/data/augmented/pretraining/zinc15.yml b/configs/data/augmented/pretraining/zinc15.yml new file mode 100644 index 0000000..dd8ab65 --- /dev/null +++ b/configs/data/augmented/pretraining/zinc15.yml @@ -0,0 +1,24 @@ +class_path: chebai_graph.preprocessing.datasets.Zinc15Pretrain_GNWA_FGWA_WFGE_AsPerNodeType +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 diff --git a/configs/model/pretraining/gat_msp.yml b/configs/model/pretraining/gat_msp.yml new file mode 100644 index 0000000..4f4d617 --- /dev/null +++ b/configs/model/pretraining/gat_msp.yml @@ -0,0 +1,19 @@ +class_path: chebai_graph.models.GATMSPretraining +init_args: + criterion: + class_path: chebai_graph.loss.msp.MSPPretrainingLoss + init_args: + hidden_size: 512 + atom_num_classes: 11 + bond_num_classes: 4 + optimizer_kwargs: + lr: 1e-3 + config: + in_channels: 203 # number of node/atom properties + hidden_channels: 256 + out_channels: 512 + num_layers: 4 + edge_dim: 10 # number of bond properties + heads: 8 # the number of heads should be divisible by output channels (hidden channels if output channel not given) + v2: True # This uses `torch_geometric.nn.conv.GATv2Conv` convolution layers, set False to use `GATConv` + dropout: 0.5 From 40076b9ce70257470804328f6ba7bc494ed27244 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sun, 11 Oct 2026 09:38:16 +0200 Subject: [PATCH 31/58] Revert "Separate Token files based on data type (#47)" This reverts commit cf723aa94e0573f1096f97375a3eef032f30e8d5. --- .../AtomCharge/indices_one_hot.txt | 0 .../AtomFunctionalGroup/indices_one_hot.txt | 0 .../AtomHybridization/indices_one_hot.txt | 0 .../AtomNodeLevel/indices_one_hot.txt | 0 .../{chebi => }/AtomNumHs/indices_one_hot.txt | 0 .../{chebi => }/AtomType/indices_one_hot.txt | 0 .../{chebi => }/BondLevel/indices_one_hot.txt | 0 .../{chebi => }/BondType/indices_one_hot.txt | 0 .../NumAtomBonds/indices_one_hot.txt | 0 chebai_graph/preprocessing/datasets/base.py | 2 +- chebai_graph/preprocessing/datasets/utils.py | 10 +-- .../properties/augmented_properties.py | 40 ++++------- chebai_graph/preprocessing/properties/base.py | 14 ++-- .../preprocessing/properties/properties.py | 72 +++++++------------ .../preprocessing/property_encoder.py | 16 ++--- 15 files changed, 51 insertions(+), 103 deletions(-) rename chebai_graph/preprocessing/bin/{chebi => }/AtomCharge/indices_one_hot.txt (100%) rename chebai_graph/preprocessing/bin/{chebi => }/AtomFunctionalGroup/indices_one_hot.txt (100%) rename chebai_graph/preprocessing/bin/{chebi => }/AtomHybridization/indices_one_hot.txt (100%) rename chebai_graph/preprocessing/bin/{chebi => }/AtomNodeLevel/indices_one_hot.txt (100%) rename chebai_graph/preprocessing/bin/{chebi => }/AtomNumHs/indices_one_hot.txt (100%) rename chebai_graph/preprocessing/bin/{chebi => }/AtomType/indices_one_hot.txt (100%) rename chebai_graph/preprocessing/bin/{chebi => }/BondLevel/indices_one_hot.txt (100%) rename chebai_graph/preprocessing/bin/{chebi => }/BondType/indices_one_hot.txt (100%) rename chebai_graph/preprocessing/bin/{chebi => }/NumAtomBonds/indices_one_hot.txt (100%) diff --git a/chebai_graph/preprocessing/bin/chebi/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/AtomCharge/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/chebi/AtomCharge/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/AtomCharge/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/chebi/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/AtomFunctionalGroup/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/chebi/AtomFunctionalGroup/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/AtomFunctionalGroup/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/chebi/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/AtomHybridization/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/chebi/AtomHybridization/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/AtomHybridization/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/chebi/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/AtomNodeLevel/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/chebi/AtomNodeLevel/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/AtomNodeLevel/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/chebi/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/AtomNumHs/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/chebi/AtomNumHs/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/AtomNumHs/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/chebi/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/AtomType/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/chebi/AtomType/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/AtomType/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/chebi/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/BondLevel/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/chebi/BondLevel/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/BondLevel/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/chebi/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/BondType/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/chebi/BondType/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/BondType/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/chebi/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/NumAtomBonds/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/chebi/NumAtomBonds/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/NumAtomBonds/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/datasets/base.py b/chebai_graph/preprocessing/datasets/base.py index d94a333..7513f76 100644 --- a/chebai_graph/preprocessing/datasets/base.py +++ b/chebai_graph/preprocessing/datasets/base.py @@ -50,7 +50,7 @@ def __init__( super().__init__(**kwargs) # atom_properties and bond_properties are given as lists containing class_paths if properties is not None: - properties = [resolve_property(prop, self.data_type) for prop in properties] + properties = [resolve_property(prop) for prop in properties] properties = self._sort_properties(properties) else: properties = [] diff --git a/chebai_graph/preprocessing/datasets/utils.py b/chebai_graph/preprocessing/datasets/utils.py index 79f2fb3..3ec6515 100644 --- a/chebai_graph/preprocessing/datasets/utils.py +++ b/chebai_graph/preprocessing/datasets/utils.py @@ -4,9 +4,7 @@ from chebai_graph.preprocessing.properties import MolecularProperty -def resolve_property( - property: str | MolecularProperty, data_type: str -) -> MolecularProperty: +def resolve_property(property: str | MolecularProperty) -> MolecularProperty: """ Resolves a molecular property specification (either as a class instance or class path string) into a MolecularProperty instance. @@ -21,8 +19,6 @@ def resolve_property( property (str | MolecularProperty): The property to resolve. Can be a class instance, a fully qualified class name (e.g. "module.ClassName"), or a class name assumed to be in `chebai_graph.preprocessing.properties`. - data_type (str): The data type associated with the property. This used to determine or set - tokens file path for the property if applicable. Returns: MolecularProperty: An instance of the resolved MolecularProperty. @@ -41,7 +37,7 @@ def resolve_property( module_name = property[:last_dot] class_name = property[last_dot + 1 :] module = importlib.import_module(module_name) - return getattr(module, class_name)(data_type=data_type) + return getattr(module, class_name)() except ValueError: # if only a class name is given, assume the module is chebai_graph.processing.properties - return getattr(graph_properties, property)(data_type=data_type) + return getattr(graph_properties, property)() diff --git a/chebai_graph/preprocessing/properties/augmented_properties.py b/chebai_graph/preprocessing/properties/augmented_properties.py index 6f112ea..f5f7b1d 100644 --- a/chebai_graph/preprocessing/properties/augmented_properties.py +++ b/chebai_graph/preprocessing/properties/augmented_properties.py @@ -23,17 +23,14 @@ class AtomNodeLevel(AllNodeTypeProperty): - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: + def __init__(self, encoder: PropertyEncoder | None = None): """ Initialize AtomNodeLevel with an optional encoder. Args: encoder (PropertyEncoder | None): Property encoder to use. Defaults to OneHotEncoder. """ - data_type = kwargs.get("data_type") - super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs - ) + super().__init__(encoder or OneHotEncoder(self)) def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> str | int | bool: """ @@ -49,17 +46,14 @@ def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> str | int | bool: class AtomFunctionalGroup(FGNodeTypeProperty): - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: + def __init__(self, encoder: PropertyEncoder | None = None): """ Initialize AtomFunctionalGroup with an optional encoder. Args: encoder (PropertyEncoder | None): Property encoder to use. Defaults to OneHotEncoder. """ - data_type = kwargs.get("data_type") - super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs - ) + super().__init__(encoder or OneHotEncoder(self)) def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> str | int | bool: """ @@ -75,17 +69,14 @@ def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> str | int | bool: class AtomRingSize(FGNodeTypeProperty): - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: + def __init__(self, encoder: PropertyEncoder | None = None): """ Initialize AtomRingSize with an optional encoder. Args: encoder (PropertyEncoder | None): Property encoder to use. Defaults to OneHotEncoder. """ - data_type = kwargs.get("data_type") - super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs - ) + super().__init__(encoder or OneHotEncoder(self)) def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> int: """ @@ -122,14 +113,14 @@ def _check_modify_atom_prop_value( class IsHydrogenBondDonorFG(FGNodeTypeProperty): - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: + def __init__(self, encoder: PropertyEncoder | None = None): """ Initialize IsHydrogenBondDonorFG with an optional encoder. Args: encoder (PropertyEncoder | None): Property encoder to use. Defaults to BoolEncoder. """ - super().__init__(encoder=encoder or BoolEncoder(self), **kwargs) + super().__init__(encoder or BoolEncoder(self)) # fmt: off # https://github.com/thaonguyen217/farm_molecular_representation/blob/main/src/(6)gen_FG_KG.py#L26-L31 self._hydrogen_bond_donor: set[str] = { @@ -155,14 +146,14 @@ def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> bool: class IsHydrogenBondAcceptorFG(FGNodeTypeProperty): - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: + def __init__(self, encoder: PropertyEncoder | None = None): """ Initialize IsHydrogenBondAcceptorFG with an optional encoder. Args: encoder (PropertyEncoder | None): Property encoder to use. Defaults to BoolEncoder. """ - super().__init__(encoder=encoder or BoolEncoder(self), **kwargs) + super().__init__(encoder or BoolEncoder(self)) # fmt: off # https://github.com/thaonguyen217/farm_molecular_representation/blob/main/src/(6)gen_FG_KG.py#L33-L39 self._hydrogen_bond_acceptor: set[str] = { @@ -189,13 +180,13 @@ def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> bool: class IsFGAlkyl(FGNodeTypeProperty): - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: + def __init__(self, encoder: PropertyEncoder | None = None): """ Args: encoder (PropertyEncoder | None): Optional encoder to use for this property. Defaults to BoolEncoder if not provided. """ - super().__init__(encoder=encoder or BoolEncoder(self), **kwargs) + super().__init__(encoder or BoolEncoder(self)) def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> int: """ @@ -332,15 +323,12 @@ class AugAtomAromaticity(AugNodeValueDefaulter, pr.AtomAromaticity): class BondLevel(AugmentedBondProperty): - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: + def __init__(self, encoder: PropertyEncoder | None = None): """ Args: encoder (PropertyEncoder | None): Optional encoder to use. Defaults to OneHotEncoder. """ - data_type = kwargs.get("data_type") - super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs - ) + super().__init__(encoder or OneHotEncoder(self)) def get_bond_value(self, bond: Chem.rdchem.Bond | dict) -> str: """ diff --git a/chebai_graph/preprocessing/properties/base.py b/chebai_graph/preprocessing/properties/base.py index 1f2e0aa..ced740d 100644 --- a/chebai_graph/preprocessing/properties/base.py +++ b/chebai_graph/preprocessing/properties/base.py @@ -31,16 +31,10 @@ class MolecularProperty(ABC): Defaults to IndexEncoder if not provided. """ - def __init__( - self, - data_type: str, - encoder: PropertyEncoder | None = None, - ) -> None: - assert data_type is not None, "data_type must be provided for MolecularProperty" + def __init__(self, encoder: PropertyEncoder | None = None) -> None: if encoder is None: - encoder = IndexEncoder(self, data_type=data_type) + encoder = IndexEncoder(self) self.encoder: PropertyEncoder = encoder - self._data_type = data_type @property def name(self) -> str: @@ -180,8 +174,8 @@ class AugAtomType(FrozenPropertyAlias, AtomType): ... ValueError: If new tokens are added to the frozen encoder during processing. """ - def __init__(self, encoder: PropertyEncoder, **kwargs) -> None: - super().__init__(encoder=encoder, **kwargs) + def __init__(self, encoder: PropertyEncoder | None = None) -> None: + super().__init__(encoder) # Lock the encoder's cache to prevent adding new tokens if hasattr(self.encoder, "cache") and isinstance(self.encoder.cache, dict): self.encoder.cache = MappingProxyType(self.encoder.cache) diff --git a/chebai_graph/preprocessing/properties/properties.py b/chebai_graph/preprocessing/properties/properties.py index 87112da..2154f9c 100644 --- a/chebai_graph/preprocessing/properties/properties.py +++ b/chebai_graph/preprocessing/properties/properties.py @@ -19,11 +19,8 @@ class AtomType(AtomProperty): Uses a one-hot encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") - super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs - ) + def __init__(self, encoder: PropertyEncoder | None = None) -> None: + super().__init__(encoder or OneHotEncoder(self)) def get_atom_value(self, atom: Chem.rdchem.Atom) -> int: """ @@ -45,11 +42,8 @@ class NumAtomBonds(AtomProperty): Uses a one-hot encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") - super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs - ) + def __init__(self, encoder: PropertyEncoder | None = None) -> None: + super().__init__(encoder or OneHotEncoder(self)) def get_atom_value(self, atom: Chem.rdchem.Atom) -> int: """ @@ -71,11 +65,8 @@ class AtomCharge(AtomProperty): Uses a one-hot encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") - super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs - ) + def __init__(self, encoder: PropertyEncoder | None = None) -> None: + super().__init__(encoder or OneHotEncoder(self)) def get_atom_value(self, atom: Chem.rdchem.Atom) -> int: """ @@ -97,11 +88,8 @@ class AtomChirality(AtomProperty): Uses a one-hot encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") - super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs - ) + def __init__(self, encoder: PropertyEncoder | None = None) -> None: + super().__init__(encoder or OneHotEncoder(self)) def get_atom_value(self, atom: Chem.rdchem.Atom) -> Chem.rdchem.ChiralType: """ @@ -123,11 +111,8 @@ class AtomHybridization(AtomProperty): Uses a one-hot encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") - super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs - ) + def __init__(self, encoder: PropertyEncoder | None = None) -> None: + super().__init__(encoder or OneHotEncoder(self)) def get_atom_value(self, atom: Chem.rdchem.Atom) -> Chem.rdchem.HybridizationType: """ @@ -149,11 +134,8 @@ class AtomNumHs(AtomProperty): Uses a one-hot encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") - super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs - ) + def __init__(self, encoder: PropertyEncoder | None = None) -> None: + super().__init__(encoder or OneHotEncoder(self)) def get_atom_value(self, atom: Chem.rdchem.Atom) -> int: """ @@ -175,8 +157,8 @@ class AtomAromaticity(AtomProperty): Uses a boolean encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - super().__init__(encoder=encoder or BoolEncoder(self), **kwargs) + def __init__(self, encoder: PropertyEncoder | None = None) -> None: + super().__init__(encoder or BoolEncoder(self)) def get_atom_value(self, atom: Chem.rdchem.Atom) -> bool: """ @@ -198,8 +180,8 @@ class BondAromaticity(BondProperty): Uses a boolean encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - super().__init__(encoder=encoder or BoolEncoder(self), **kwargs) + def __init__(self, encoder: PropertyEncoder | None = None) -> None: + super().__init__(encoder or BoolEncoder(self)) def get_bond_value(self, bond: Chem.rdchem.Bond) -> bool: """ @@ -221,11 +203,8 @@ class BondType(BondProperty): Uses a one-hot encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") - super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs - ) + def __init__(self, encoder: PropertyEncoder | None = None) -> None: + super().__init__(encoder or OneHotEncoder(self)) def get_bond_value(self, bond: Chem.rdchem.Bond) -> Chem.rdchem.BondType: """ @@ -247,8 +226,8 @@ class BondInRing(BondProperty): Uses a boolean encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - super().__init__(encoder=encoder or BoolEncoder(self), **kwargs) + def __init__(self, encoder: PropertyEncoder | None = None) -> None: + super().__init__(encoder or BoolEncoder(self)) def get_bond_value(self, bond: Chem.rdchem.Bond) -> bool: """ @@ -270,11 +249,8 @@ class MoleculeNumRings(MoleculeProperty): Uses a one-hot encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") - super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs - ) + def __init__(self, encoder: PropertyEncoder | None = None) -> None: + super().__init__(encoder or OneHotEncoder(self)) def get_property_value(self, mol: Chem.rdchem.Mol) -> list[int]: """ @@ -296,8 +272,8 @@ class RDKit2DNormalized(MoleculeProperty): Uses an identity encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - super().__init__(encoder=encoder or AsIsEncoder(self), **kwargs) + def __init__(self, encoder: PropertyEncoder | None = None) -> None: + super().__init__(encoder or AsIsEncoder(self)) self.generator_normalized = rdNormalizedDescriptors.RDKit2DNormalized() # Create a dummy molecule (e.g., methane) to extract the length of descriptor vector dummy_mol = Chem.MolFromSmiles("C") diff --git a/chebai_graph/preprocessing/property_encoder.py b/chebai_graph/preprocessing/property_encoder.py index 5daeb4c..38cb279 100644 --- a/chebai_graph/preprocessing/property_encoder.py +++ b/chebai_graph/preprocessing/property_encoder.py @@ -99,15 +99,12 @@ class IndexEncoder(PropertyEncoder): **kwargs: Additional keyword arguments. """ - def __init__( - self, property, data_type: str, indices_dir: str | None = None, **kwargs - ) -> None: + def __init__(self, property, indices_dir: str | None = None, **kwargs) -> None: super().__init__(property, **kwargs) if indices_dir is None: indices_dir = os.path.dirname(inspect.getfile(self.__class__)) self.dirname = indices_dir # load already existing cache - self._data_type = data_type with open(self.index_path, "r") as pk: self.cache: dict[str, int] = { token.strip(): idx for idx, token in enumerate(pk) @@ -129,15 +126,12 @@ def index_path(self) -> str: Returns: Path to index file. """ - assert self._data_type is not None, "data_type must be set for IndexEncoder" index_path = os.path.join( - self.dirname, - "bin", - self._data_type, - self.property.name, - f"indices_{self.name}.txt", + self.dirname, "bin", self.property.name, f"indices_{self.name}.txt" + ) + os.makedirs( + os.path.join(self.dirname, "bin", self.property.name), exist_ok=True ) - os.makedirs(os.path.dirname(index_path), exist_ok=True) if not os.path.exists(index_path): with open(index_path, "x"): pass From 543bb0e348cfe060bb4028aeee4f6886a7b34b98 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sun, 11 Oct 2026 10:03:10 +0200 Subject: [PATCH 32/58] delete data type based tokens --- .../bin/bace/AtomCharge/indices_one_hot.txt | 3 - .../AtomFunctionalGroup/indices_one_hot.txt | 44 ------ .../AtomHybridization/indices_one_hot.txt | 3 - .../bace/AtomNodeLevel/indices_one_hot.txt | 3 - .../bin/bace/AtomNumHs/indices_one_hot.txt | 4 - .../bin/bace/AtomType/indices_one_hot.txt | 8 -- .../bin/bace/BondLevel/indices_one_hot.txt | 4 - .../bin/bace/BondType/indices_one_hot.txt | 4 - .../bin/bace/NumAtomBonds/indices_one_hot.txt | 4 - .../bin/bbbp/AtomCharge/indices_one_hot.txt | 4 - .../AtomFunctionalGroup/indices_one_hot.txt | 80 ----------- .../AtomHybridization/indices_one_hot.txt | 4 - .../bbbp/AtomNodeLevel/indices_one_hot.txt | 3 - .../bin/bbbp/AtomNumHs/indices_one_hot.txt | 4 - .../bin/bbbp/AtomType/indices_one_hot.txt | 13 -- .../bin/bbbp/BondLevel/indices_one_hot.txt | 4 - .../bin/bbbp/BondType/indices_one_hot.txt | 4 - .../bin/bbbp/NumAtomBonds/indices_one_hot.txt | 5 - .../clin_tox/AtomCharge/indices_one_hot.txt | 6 - .../AtomFunctionalGroup/indices_one_hot.txt | 84 ------------ .../AtomHybridization/indices_one_hot.txt | 7 - .../AtomNodeLevel/indices_one_hot.txt | 3 - .../clin_tox/AtomNumHs/indices_one_hot.txt | 4 - .../bin/clin_tox/AtomType/indices_one_hot.txt | 30 ----- .../clin_tox/BondLevel/indices_one_hot.txt | 4 - .../bin/clin_tox/BondType/indices_one_hot.txt | 5 - .../clin_tox/NumAtomBonds/indices_one_hot.txt | 6 - .../bin/hiv/AtomCharge/indices_one_hot.txt | 9 -- .../AtomFunctionalGroup/indices_one_hot.txt | 122 ----------------- .../hiv/AtomHybridization/indices_one_hot.txt | 7 - .../bin/hiv/AtomNodeLevel/indices_one_hot.txt | 3 - .../bin/hiv/AtomNumHs/indices_one_hot.txt | 5 - .../bin/hiv/AtomType/indices_one_hot.txt | 54 -------- .../bin/hiv/BondLevel/indices_one_hot.txt | 4 - .../bin/hiv/BondType/indices_one_hot.txt | 5 - .../bin/hiv/NumAtomBonds/indices_one_hot.txt | 11 -- .../bin/muv/AtomCharge/indices_one_hot.txt | 3 - .../AtomFunctionalGroup/indices_one_hot.txt | 68 ---------- .../muv/AtomHybridization/indices_one_hot.txt | 3 - .../bin/muv/AtomNodeLevel/indices_one_hot.txt | 3 - .../bin/muv/AtomNumHs/indices_one_hot.txt | 4 - .../bin/muv/AtomType/indices_one_hot.txt | 7 - .../bin/muv/BondLevel/indices_one_hot.txt | 4 - .../bin/muv/BondType/indices_one_hot.txt | 4 - .../bin/muv/NumAtomBonds/indices_one_hot.txt | 4 - .../bin/pcba/AtomCharge/indices_one_hot.txt | 8 -- .../AtomFunctionalGroup/indices_one_hot.txt | 127 ------------------ .../AtomHybridization/indices_one_hot.txt | 7 - .../pcba/AtomNodeLevel/indices_one_hot.txt | 3 - .../bin/pcba/AtomNumHs/indices_one_hot.txt | 6 - .../bin/pcba/AtomType/indices_one_hot.txt | 44 ------ .../bin/pcba/BondLevel/indices_one_hot.txt | 4 - .../bin/pcba/BondType/indices_one_hot.txt | 4 - .../bin/pcba/NumAtomBonds/indices_one_hot.txt | 6 - .../bin/sider/AtomCharge/indices_one_hot.txt | 6 - .../AtomFunctionalGroup/indices_one_hot.txt | 93 ------------- .../AtomHybridization/indices_one_hot.txt | 6 - .../sider/AtomNodeLevel/indices_one_hot.txt | 3 - .../bin/sider/AtomNumHs/indices_one_hot.txt | 5 - .../bin/sider/AtomType/indices_one_hot.txt | 40 ------ .../bin/sider/BondLevel/indices_one_hot.txt | 4 - .../bin/sider/BondType/indices_one_hot.txt | 4 - .../sider/NumAtomBonds/indices_one_hot.txt | 6 - .../bin/tox21/AtomCharge/indices_one_hot.txt | 6 - .../AtomFunctionalGroup/indices_one_hot.txt | 115 ---------------- .../AtomHybridization/indices_one_hot.txt | 7 - .../tox21/AtomNodeLevel/indices_one_hot.txt | 3 - .../bin/tox21/AtomNumHs/indices_one_hot.txt | 6 - .../bin/tox21/AtomType/indices_one_hot.txt | 50 ------- .../bin/tox21/BondLevel/indices_one_hot.txt | 4 - .../bin/tox21/BondType/indices_one_hot.txt | 4 - .../tox21/NumAtomBonds/indices_one_hot.txt | 7 - .../toxcast/AtomCharge/indices_one_hot.txt | 6 - .../AtomFunctionalGroup/indices_one_hot.txt | 117 ---------------- .../AtomHybridization/indices_one_hot.txt | 7 - .../toxcast/AtomNodeLevel/indices_one_hot.txt | 3 - .../bin/toxcast/AtomNumHs/indices_one_hot.txt | 5 - .../bin/toxcast/AtomType/indices_one_hot.txt | 53 -------- .../bin/toxcast/BondLevel/indices_one_hot.txt | 4 - .../bin/toxcast/BondType/indices_one_hot.txt | 4 - .../toxcast/NumAtomBonds/indices_one_hot.txt | 7 - 81 files changed, 1451 deletions(-) delete mode 100644 chebai_graph/preprocessing/bin/bace/AtomCharge/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/bace/AtomFunctionalGroup/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/bace/AtomHybridization/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/bace/AtomNodeLevel/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/bace/AtomNumHs/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/bace/AtomType/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/bace/BondLevel/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/bace/BondType/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/bace/NumAtomBonds/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/bbbp/AtomCharge/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/bbbp/AtomFunctionalGroup/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/bbbp/AtomHybridization/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/bbbp/AtomNodeLevel/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/bbbp/AtomNumHs/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/bbbp/AtomType/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/bbbp/BondLevel/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/bbbp/BondType/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/bbbp/NumAtomBonds/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/clin_tox/AtomCharge/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/clin_tox/AtomFunctionalGroup/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/clin_tox/AtomHybridization/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/clin_tox/AtomNodeLevel/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/clin_tox/AtomNumHs/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/clin_tox/AtomType/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/clin_tox/BondLevel/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/clin_tox/BondType/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/clin_tox/NumAtomBonds/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/hiv/AtomCharge/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/hiv/AtomFunctionalGroup/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/hiv/AtomHybridization/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/hiv/AtomNodeLevel/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/hiv/AtomNumHs/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/hiv/AtomType/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/hiv/BondLevel/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/hiv/BondType/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/hiv/NumAtomBonds/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/muv/AtomCharge/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/muv/AtomFunctionalGroup/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/muv/AtomHybridization/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/muv/AtomNodeLevel/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/muv/AtomNumHs/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/muv/AtomType/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/muv/BondLevel/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/muv/BondType/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/muv/NumAtomBonds/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/pcba/AtomCharge/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/pcba/AtomFunctionalGroup/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/pcba/AtomHybridization/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/pcba/AtomNodeLevel/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/pcba/AtomNumHs/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/pcba/AtomType/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/pcba/BondLevel/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/pcba/BondType/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/pcba/NumAtomBonds/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/sider/AtomCharge/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/sider/AtomFunctionalGroup/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/sider/AtomHybridization/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/sider/AtomNodeLevel/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/sider/AtomNumHs/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/sider/AtomType/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/sider/BondLevel/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/sider/BondType/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/sider/NumAtomBonds/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/tox21/AtomCharge/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/tox21/AtomFunctionalGroup/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/tox21/AtomHybridization/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/tox21/AtomNodeLevel/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/tox21/AtomNumHs/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/tox21/AtomType/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/tox21/BondLevel/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/tox21/BondType/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/tox21/NumAtomBonds/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/toxcast/AtomCharge/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/toxcast/AtomFunctionalGroup/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/toxcast/AtomHybridization/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/toxcast/AtomNodeLevel/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/toxcast/AtomNumHs/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/toxcast/AtomType/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/toxcast/BondLevel/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/toxcast/BondType/indices_one_hot.txt delete mode 100644 chebai_graph/preprocessing/bin/toxcast/NumAtomBonds/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/bace/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomCharge/indices_one_hot.txt deleted file mode 100644 index d7e52d1..0000000 --- a/chebai_graph/preprocessing/bin/bace/AtomCharge/indices_one_hot.txt +++ /dev/null @@ -1,3 +0,0 @@ -0 -1 --1 diff --git a/chebai_graph/preprocessing/bin/bace/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomFunctionalGroup/indices_one_hot.txt deleted file mode 100644 index b621c2a..0000000 --- a/chebai_graph/preprocessing/bin/bace/AtomFunctionalGroup/indices_one_hot.txt +++ /dev/null @@ -1,44 +0,0 @@ -RING_6 -difluoromethyl -RING_5 -ether -alkyne -alkyl -primary_amine -fluoro -amide -graph_fg -NO_FG -alkene -tertiary_carbon -alkene_carbon -nitrile -trifluoromethyl -ketone -chloro -quaternary_carbon -sulfonyl -hydroxyl -bromo -secondary_amine -ester -aldoxime -sulfide -iodo -tertiary_amine -RING_4 -secondary_ketimine -ketoxime -RING_3 -carboxylate -acetal -sulfinyl -RING_7 -RING_16 -RING_15 -nitro -primary_ketimine -sulfonate_ester -RING_8 -RING_14 -RING_13 diff --git a/chebai_graph/preprocessing/bin/bace/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomHybridization/indices_one_hot.txt deleted file mode 100644 index 7c99da3..0000000 --- a/chebai_graph/preprocessing/bin/bace/AtomHybridization/indices_one_hot.txt +++ /dev/null @@ -1,3 +0,0 @@ -SP2 -SP3 -SP diff --git a/chebai_graph/preprocessing/bin/bace/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomNodeLevel/indices_one_hot.txt deleted file mode 100644 index 2d65776..0000000 --- a/chebai_graph/preprocessing/bin/bace/AtomNodeLevel/indices_one_hot.txt +++ /dev/null @@ -1,3 +0,0 @@ -atom_node_lvl -fg_node_lvl -graph_node_level diff --git a/chebai_graph/preprocessing/bin/bace/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomNumHs/indices_one_hot.txt deleted file mode 100644 index c4b216a..0000000 --- a/chebai_graph/preprocessing/bin/bace/AtomNumHs/indices_one_hot.txt +++ /dev/null @@ -1,4 +0,0 @@ -0 -1 -3 -2 diff --git a/chebai_graph/preprocessing/bin/bace/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomType/indices_one_hot.txt deleted file mode 100644 index 6da4980..0000000 --- a/chebai_graph/preprocessing/bin/bace/AtomType/indices_one_hot.txt +++ /dev/null @@ -1,8 +0,0 @@ -6 -7 -8 -9 -17 -16 -35 -53 diff --git a/chebai_graph/preprocessing/bin/bace/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/BondLevel/indices_one_hot.txt deleted file mode 100644 index 6b48d66..0000000 --- a/chebai_graph/preprocessing/bin/bace/BondLevel/indices_one_hot.txt +++ /dev/null @@ -1,4 +0,0 @@ -within_atoms_lvl -atom_fg_lvl -within_fg_lvl -to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/bace/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/BondType/indices_one_hot.txt deleted file mode 100644 index f1a2ce5..0000000 --- a/chebai_graph/preprocessing/bin/bace/BondType/indices_one_hot.txt +++ /dev/null @@ -1,4 +0,0 @@ -SINGLE -AROMATIC -TRIPLE -DOUBLE diff --git a/chebai_graph/preprocessing/bin/bace/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/NumAtomBonds/indices_one_hot.txt deleted file mode 100644 index cf286d6..0000000 --- a/chebai_graph/preprocessing/bin/bace/NumAtomBonds/indices_one_hot.txt +++ /dev/null @@ -1,4 +0,0 @@ -3 -2 -4 -1 diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomCharge/indices_one_hot.txt deleted file mode 100644 index 7b6a00a..0000000 --- a/chebai_graph/preprocessing/bin/bbbp/AtomCharge/indices_one_hot.txt +++ /dev/null @@ -1,4 +0,0 @@ -0 -1 --1 -2 diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomFunctionalGroup/indices_one_hot.txt deleted file mode 100644 index d2ee211..0000000 --- a/chebai_graph/preprocessing/bin/bbbp/AtomFunctionalGroup/indices_one_hot.txt +++ /dev/null @@ -1,80 +0,0 @@ -alkyl -ester -chloro -RING_6 -tertiary_amine -graph_fg -tertiary_carbon -primary_amine -dichloromethyl -sulfonyl -ketone -secondary_amine -hydroxyl -ether -quaternary_carbon -amide -NO_FG -carboxyl -nitro -alkene -alkene_carbon -secondary_ketimine -secondary_aldimine -carboxylate -iodo -aldoxime -bromo -alkyne -fluoro -carbamate -trifluoromethyl -sulfide -trichloromethyl -hemiacetal -primary_ketimine -azide -4_ammonium_ion -sulfonate_ester -ketoxime -nitroso -sulfhydryl -difluoromethyl -fluorochloromethyl -chlorobromomethyl -phosphono -carbodithio -disulfide -difluorochloromethyl -carbonate_ester -phosphate -RING_5 -thiolester -RING_7 -nitrile -phosphoryl -aldehyde -sulfinyl -thioketone -acetal -RING_4 -sulfonic_acid -imide -RING_14 -RING_24 -amidine -RING_3 -RING_9 -RING_19 -RING_16 -RING_8 -isonitrile -RING_23 -RING_12 -RING_21 -RING_17 -RING_30 -RING_33 -RING_36 -azo -RING_15 diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomHybridization/indices_one_hot.txt deleted file mode 100644 index 0e4a9aa..0000000 --- a/chebai_graph/preprocessing/bin/bbbp/AtomHybridization/indices_one_hot.txt +++ /dev/null @@ -1,4 +0,0 @@ -SP3 -SP2 -S -SP diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomNodeLevel/indices_one_hot.txt deleted file mode 100644 index 2d65776..0000000 --- a/chebai_graph/preprocessing/bin/bbbp/AtomNodeLevel/indices_one_hot.txt +++ /dev/null @@ -1,3 +0,0 @@ -atom_node_lvl -fg_node_lvl -graph_node_level diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomNumHs/indices_one_hot.txt deleted file mode 100644 index 35f5032..0000000 --- a/chebai_graph/preprocessing/bin/bbbp/AtomNumHs/indices_one_hot.txt +++ /dev/null @@ -1,4 +0,0 @@ -2 -0 -1 -3 diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomType/indices_one_hot.txt deleted file mode 100644 index 9500570..0000000 --- a/chebai_graph/preprocessing/bin/bbbp/AtomType/indices_one_hot.txt +++ /dev/null @@ -1,13 +0,0 @@ -6 -17 -8 -7 -16 -11 -53 -35 -9 -1 -15 -5 -20 diff --git a/chebai_graph/preprocessing/bin/bbbp/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/BondLevel/indices_one_hot.txt deleted file mode 100644 index 6b48d66..0000000 --- a/chebai_graph/preprocessing/bin/bbbp/BondLevel/indices_one_hot.txt +++ /dev/null @@ -1,4 +0,0 @@ -within_atoms_lvl -atom_fg_lvl -within_fg_lvl -to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/bbbp/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/BondType/indices_one_hot.txt deleted file mode 100644 index a134437..0000000 --- a/chebai_graph/preprocessing/bin/bbbp/BondType/indices_one_hot.txt +++ /dev/null @@ -1,4 +0,0 @@ -DOUBLE -SINGLE -AROMATIC -TRIPLE diff --git a/chebai_graph/preprocessing/bin/bbbp/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/NumAtomBonds/indices_one_hot.txt deleted file mode 100644 index 5d07eee..0000000 --- a/chebai_graph/preprocessing/bin/bbbp/NumAtomBonds/indices_one_hot.txt +++ /dev/null @@ -1,5 +0,0 @@ -2 -4 -1 -3 -0 diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomCharge/indices_one_hot.txt deleted file mode 100644 index 374de68..0000000 --- a/chebai_graph/preprocessing/bin/clin_tox/AtomCharge/indices_one_hot.txt +++ /dev/null @@ -1,6 +0,0 @@ -0 --1 -1 -3 -2 --2 diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomFunctionalGroup/indices_one_hot.txt deleted file mode 100644 index 36bc343..0000000 --- a/chebai_graph/preprocessing/bin/clin_tox/AtomFunctionalGroup/indices_one_hot.txt +++ /dev/null @@ -1,84 +0,0 @@ -alkyl -chloro -RING_6 -ether -secondary_amine -fluoro -graph_fg -tertiary_carbon -hydroxyl -ketone -sulfonyl -carboxylate -RING_5 -NO_FG -amide -sulfide -RING_4 -quaternary_carbon -alkyne -alkene -primary_amine -carboxyl -trifluoromethyl -ester -alkene_carbon -RING_20 -sulfhydryl -nitrile -RING_3 -RING_29 -RING_7 -4_ammonium_ion -acetal -RING_9 -iodo -RING_38 -sulfinyl -difluoromethyl -imide -RING_16 -tertiary_amine -nitro -RING_8 -secondary_ketimine -nitrate -sulfonic_acid -carbamate -secondary_aldimine -bromo -thioketone -RING_14 -azo -RING_33 -trichloromethyl -RING_15 -ketoxime -phosphodiester -RING_24 -RING_22 -aldehyde -RING_18 -RING_21 -nitroso -dichloromethyl -carbonate_ester -silyl_ether -thiolester -borono -aldoxime -sulfonate_ester -RING_23 -RING_19 -RING_25 -carbodithio -disulfide -RING_36 -RING_31 -fluorochloromethyl -RING_12 -phosphate -azide -chlorobromomethyl -bromodifluoromethyl -phosphino diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomHybridization/indices_one_hot.txt deleted file mode 100644 index a96b9b3..0000000 --- a/chebai_graph/preprocessing/bin/clin_tox/AtomHybridization/indices_one_hot.txt +++ /dev/null @@ -1,7 +0,0 @@ -SP3 -SP2 -SP -S -SP3D2 -UNSPECIFIED -SP3D diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomNodeLevel/indices_one_hot.txt deleted file mode 100644 index 2d65776..0000000 --- a/chebai_graph/preprocessing/bin/clin_tox/AtomNodeLevel/indices_one_hot.txt +++ /dev/null @@ -1,3 +0,0 @@ -atom_node_lvl -fg_node_lvl -graph_node_level diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomNumHs/indices_one_hot.txt deleted file mode 100644 index 1127527..0000000 --- a/chebai_graph/preprocessing/bin/clin_tox/AtomNumHs/indices_one_hot.txt +++ /dev/null @@ -1,4 +0,0 @@ -3 -0 -2 -1 diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomType/indices_one_hot.txt deleted file mode 100644 index 2c8a324..0000000 --- a/chebai_graph/preprocessing/bin/clin_tox/AtomType/indices_one_hot.txt +++ /dev/null @@ -1,30 +0,0 @@ -6 -17 -7 -8 -9 -16 -15 -53 -80 -1 -13 -35 -79 -0 -27 -22 -78 -33 -34 -30 -14 -5 -29 -24 -81 -43 -83 -25 -20 -26 diff --git a/chebai_graph/preprocessing/bin/clin_tox/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/BondLevel/indices_one_hot.txt deleted file mode 100644 index 6b48d66..0000000 --- a/chebai_graph/preprocessing/bin/clin_tox/BondLevel/indices_one_hot.txt +++ /dev/null @@ -1,4 +0,0 @@ -within_atoms_lvl -atom_fg_lvl -within_fg_lvl -to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/clin_tox/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/BondType/indices_one_hot.txt deleted file mode 100644 index bdc90d7..0000000 --- a/chebai_graph/preprocessing/bin/clin_tox/BondType/indices_one_hot.txt +++ /dev/null @@ -1,5 +0,0 @@ -SINGLE -AROMATIC -DOUBLE -TRIPLE -DATIVE diff --git a/chebai_graph/preprocessing/bin/clin_tox/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/NumAtomBonds/indices_one_hot.txt deleted file mode 100644 index 3c2db51..0000000 --- a/chebai_graph/preprocessing/bin/clin_tox/NumAtomBonds/indices_one_hot.txt +++ /dev/null @@ -1,6 +0,0 @@ -1 -3 -2 -4 -0 -6 diff --git a/chebai_graph/preprocessing/bin/hiv/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/AtomCharge/indices_one_hot.txt deleted file mode 100644 index 22aacd7..0000000 --- a/chebai_graph/preprocessing/bin/hiv/AtomCharge/indices_one_hot.txt +++ /dev/null @@ -1,9 +0,0 @@ -0 -1 --1 -3 --4 --3 --2 -2 --5 diff --git a/chebai_graph/preprocessing/bin/hiv/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/AtomFunctionalGroup/indices_one_hot.txt deleted file mode 100644 index e5b33d4..0000000 --- a/chebai_graph/preprocessing/bin/hiv/AtomFunctionalGroup/indices_one_hot.txt +++ /dev/null @@ -1,122 +0,0 @@ -alkyl -RING_6 -chloro -NO_FG -secondary_amine -graph_fg -carboxyl -primary_amine -sulfonic_acid -nitro -ketone -hydroxyl -sulfonyl -ester -nitroso -sulfhydryl -tertiary_carbon -tertiary_amine -dichloromethyl -nitrile -primary_ketimine -alkyne -alkene -ether -iodo -sulfide -secondary_aldimine -phosphoryl -thioketone -alkene_carbon -secondary_ketimine -nitrate -4_ammonium_ion -ketoxime -amide -amidine -trichloromethyl -fluoro -quaternary_carbon -hydrazone -disulfide -thiocyanate -carbamate -trifluoromethyl -phosphate -sulfonate_ester -sulfinyl -azo -bromo -imide -carbodithioic_acid -carbodithio -sulfino -silyl_ether -aldehyde -phosphono -azide -aldoxime -phosphino -trimethylsilyl -thionoester -thiolester -isonitrile -acetal -difluorochloromethyl -dibromomethyl -ketal -diiodomethyl -borono -isothiocyanate -difluoromethyl -tribromomethyl -carbonate_ester -hemiacetal -phosphodiester -carboxylate -orthoester -isocyanate -primary_aldimine -hemiketal -borinate -RING_5 -cyanate -RING_4 -RING_24 -RING_3 -thial -RING_7 -RING_12 -RING_9 -RING_16 -hydroperoxy -RING_15 -RING_18 -RING_8 -RING_35 -RING_10 -RING_19 -RING_17 -RING_29 -RING_23 -RING_27 -RING_11 -RING_14 -RING_20 -RING_33 -RING_13 -RING_31 -RING_21 -RING_28 -RING_22 -RING_30 -RING_34 -RING_54 -RING_38 -RING_36 -RING_26 -RING_32 -carboxylic_anhydride -RING_25 -carbothioic_S-acid -RING_40 diff --git a/chebai_graph/preprocessing/bin/hiv/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/AtomHybridization/indices_one_hot.txt deleted file mode 100644 index cdc6985..0000000 --- a/chebai_graph/preprocessing/bin/hiv/AtomHybridization/indices_one_hot.txt +++ /dev/null @@ -1,7 +0,0 @@ -SP3 -SP2 -SP -S -SP3D2 -SP3D -UNSPECIFIED diff --git a/chebai_graph/preprocessing/bin/hiv/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/AtomNodeLevel/indices_one_hot.txt deleted file mode 100644 index 2d65776..0000000 --- a/chebai_graph/preprocessing/bin/hiv/AtomNodeLevel/indices_one_hot.txt +++ /dev/null @@ -1,3 +0,0 @@ -atom_node_lvl -fg_node_lvl -graph_node_level diff --git a/chebai_graph/preprocessing/bin/hiv/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/AtomNumHs/indices_one_hot.txt deleted file mode 100644 index df4a3e9..0000000 --- a/chebai_graph/preprocessing/bin/hiv/AtomNumHs/indices_one_hot.txt +++ /dev/null @@ -1,5 +0,0 @@ -3 -1 -0 -2 -4 diff --git a/chebai_graph/preprocessing/bin/hiv/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/AtomType/indices_one_hot.txt deleted file mode 100644 index 95c9cba..0000000 --- a/chebai_graph/preprocessing/bin/hiv/AtomType/indices_one_hot.txt +++ /dev/null @@ -1,54 +0,0 @@ -6 -17 -8 -7 -15 -16 -33 -53 -9 -35 -14 -11 -34 -79 -50 -80 -32 -51 -82 -5 -19 -3 -83 -29 -67 -44 -65 -74 -27 -55 -25 -26 -75 -42 -78 -47 -24 -46 -89 -20 -31 -12 -30 -28 -77 -1 -40 -81 -45 -22 -52 -64 -92 -23 diff --git a/chebai_graph/preprocessing/bin/hiv/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/BondLevel/indices_one_hot.txt deleted file mode 100644 index 6b48d66..0000000 --- a/chebai_graph/preprocessing/bin/hiv/BondLevel/indices_one_hot.txt +++ /dev/null @@ -1,4 +0,0 @@ -within_atoms_lvl -atom_fg_lvl -within_fg_lvl -to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/hiv/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/BondType/indices_one_hot.txt deleted file mode 100644 index 58c0a5c..0000000 --- a/chebai_graph/preprocessing/bin/hiv/BondType/indices_one_hot.txt +++ /dev/null @@ -1,5 +0,0 @@ -SINGLE -DOUBLE -AROMATIC -TRIPLE -DATIVE diff --git a/chebai_graph/preprocessing/bin/hiv/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/NumAtomBonds/indices_one_hot.txt deleted file mode 100644 index 88ddb48..0000000 --- a/chebai_graph/preprocessing/bin/hiv/NumAtomBonds/indices_one_hot.txt +++ /dev/null @@ -1,11 +0,0 @@ -1 -2 -3 -4 -0 -5 -6 -9 -7 -8 -10 diff --git a/chebai_graph/preprocessing/bin/muv/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/AtomCharge/indices_one_hot.txt deleted file mode 100644 index d7e52d1..0000000 --- a/chebai_graph/preprocessing/bin/muv/AtomCharge/indices_one_hot.txt +++ /dev/null @@ -1,3 +0,0 @@ -0 -1 --1 diff --git a/chebai_graph/preprocessing/bin/muv/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/AtomFunctionalGroup/indices_one_hot.txt deleted file mode 100644 index c55dc92..0000000 --- a/chebai_graph/preprocessing/bin/muv/AtomFunctionalGroup/indices_one_hot.txt +++ /dev/null @@ -1,68 +0,0 @@ -primary_amine -RING_6 -ketone -carboxyl -alkyl -tertiary_carbon -secondary_amine -graph_fg -hydroxyl -alkene -thioketone -bromo -fluoro -ester -ether -tertiary_amine -NO_FG -chloro -sulfonyl -difluoromethyl -quaternary_carbon -nitro -amide -alkyne -trifluoromethyl -aldoxime -alkene_carbon -nitrile -sulfide -secondary_aldimine -ketoxime -carbodithio -secondary_ketimine -thionoester -thiolester -sulfinyl -carbamate -thiocyanate -aldehyde -sulfonate_ester -amidine -trichloromethyl -sulfhydryl -sulfonic_acid -RING_5 -nitroso -acetal -carbonate_ester -imide -azo -primary_ketimine -dichloromethyl -RING_7 -RING_3 -RING_8 -hydrazone -RING_4 -RING_15 -RING_12 -azide -RING_13 -RING_18 -difluorochloromethyl -RING_14 -RING_9 -RING_11 -disulfide -4_ammonium_ion diff --git a/chebai_graph/preprocessing/bin/muv/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/AtomHybridization/indices_one_hot.txt deleted file mode 100644 index 7c99da3..0000000 --- a/chebai_graph/preprocessing/bin/muv/AtomHybridization/indices_one_hot.txt +++ /dev/null @@ -1,3 +0,0 @@ -SP2 -SP3 -SP diff --git a/chebai_graph/preprocessing/bin/muv/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/AtomNodeLevel/indices_one_hot.txt deleted file mode 100644 index 2d65776..0000000 --- a/chebai_graph/preprocessing/bin/muv/AtomNodeLevel/indices_one_hot.txt +++ /dev/null @@ -1,3 +0,0 @@ -atom_node_lvl -fg_node_lvl -graph_node_level diff --git a/chebai_graph/preprocessing/bin/muv/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/AtomNumHs/indices_one_hot.txt deleted file mode 100644 index 52db5ae..0000000 --- a/chebai_graph/preprocessing/bin/muv/AtomNumHs/indices_one_hot.txt +++ /dev/null @@ -1,4 +0,0 @@ -2 -1 -0 -3 diff --git a/chebai_graph/preprocessing/bin/muv/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/AtomType/indices_one_hot.txt deleted file mode 100644 index c36ddc7..0000000 --- a/chebai_graph/preprocessing/bin/muv/AtomType/indices_one_hot.txt +++ /dev/null @@ -1,7 +0,0 @@ -7 -6 -8 -35 -9 -16 -17 diff --git a/chebai_graph/preprocessing/bin/muv/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/BondLevel/indices_one_hot.txt deleted file mode 100644 index 6b48d66..0000000 --- a/chebai_graph/preprocessing/bin/muv/BondLevel/indices_one_hot.txt +++ /dev/null @@ -1,4 +0,0 @@ -within_atoms_lvl -atom_fg_lvl -within_fg_lvl -to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/muv/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/BondType/indices_one_hot.txt deleted file mode 100644 index f14a241..0000000 --- a/chebai_graph/preprocessing/bin/muv/BondType/indices_one_hot.txt +++ /dev/null @@ -1,4 +0,0 @@ -SINGLE -DOUBLE -AROMATIC -TRIPLE diff --git a/chebai_graph/preprocessing/bin/muv/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/NumAtomBonds/indices_one_hot.txt deleted file mode 100644 index 94ebaf9..0000000 --- a/chebai_graph/preprocessing/bin/muv/NumAtomBonds/indices_one_hot.txt +++ /dev/null @@ -1,4 +0,0 @@ -1 -2 -3 -4 diff --git a/chebai_graph/preprocessing/bin/pcba/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/AtomCharge/indices_one_hot.txt deleted file mode 100644 index 7897486..0000000 --- a/chebai_graph/preprocessing/bin/pcba/AtomCharge/indices_one_hot.txt +++ /dev/null @@ -1,8 +0,0 @@ -0 -1 --1 -3 -2 -4 -5 -6 diff --git a/chebai_graph/preprocessing/bin/pcba/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/AtomFunctionalGroup/indices_one_hot.txt deleted file mode 100644 index 89e00f3..0000000 --- a/chebai_graph/preprocessing/bin/pcba/AtomFunctionalGroup/indices_one_hot.txt +++ /dev/null @@ -1,127 +0,0 @@ -alkyl -RING_6 -sulfide -RING_5 -ketone -trifluoromethyl -secondary_amine -graph_fg -NO_FG -amide -sulfonyl -ether -quaternary_carbon -tertiary_carbon -imide -carboxyl -nitrile -ester -primary_amine -alkyne -fluoro -tertiary_amine -nitro -RING_7 -chloro -hydroxyl -RING_3 -alkene -secondary_aldimine -bromo -thioketone -sulfinyl -RING_4 -secondary_ketimine -RING_12 -primary_ketimine -alkene_carbon -carbamate -aldehyde -phosphoryl -trimethylsilyl -carboxylate -ketoxime -RING_23 -difluoromethyl -sulfonate_ester -RING_14 -iodo -thiocyanate -RING_8 -disulfide -carbonate_ester -RING_19 -sulfhydryl -nitroso -acetal -aldoxime -borono -RING_16 -phosphono -RING_9 -RING_15 -hydrazone -RING_11 -4_ammonium_ion -thiolester -trichloromethyl -carbodithio -azo -azide -sulfonic_acid -dichloromethyl -phosphodiester -borinate -RING_13 -RING_18 -thial -RING_21 -RING_10 -isocyanate -isonitrile -RING_22 -RING_36 -RING_20 -phosphate -amidine -dibromomethyl -RING_17 -peroxy -RING_40 -silyl_ether -thionoester -nitrate -haloformyl -hemiacetal -RING_25 -isothiocyanate -difluorochloromethyl -RING_24 -nitrosooxy -RING_32 -carbodithioic_acid -hydroperoxy -sulfino -RING_26 -tribromomethyl -RING_30 -RING_29 -diiodomethyl -phosphino -chlorobromomethyl -RING_33 -RING_38 -RING_34 -bromodifluoromethyl -carboxylic_anhydride -bromodichloromethyl -RING_27 -dichlorosilane -RING_42 -ketal -RING_28 -RING_56 -RING_31 -RING_35 -fluorochloromethyl -RING_54 diff --git a/chebai_graph/preprocessing/bin/pcba/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/AtomHybridization/indices_one_hot.txt deleted file mode 100644 index fb3232b..0000000 --- a/chebai_graph/preprocessing/bin/pcba/AtomHybridization/indices_one_hot.txt +++ /dev/null @@ -1,7 +0,0 @@ -SP3 -SP2 -SP -SP3D -S -UNSPECIFIED -SP3D2 diff --git a/chebai_graph/preprocessing/bin/pcba/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/AtomNodeLevel/indices_one_hot.txt deleted file mode 100644 index 2d65776..0000000 --- a/chebai_graph/preprocessing/bin/pcba/AtomNodeLevel/indices_one_hot.txt +++ /dev/null @@ -1,3 +0,0 @@ -atom_node_lvl -fg_node_lvl -graph_node_level diff --git a/chebai_graph/preprocessing/bin/pcba/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/AtomNumHs/indices_one_hot.txt deleted file mode 100644 index cad2399..0000000 --- a/chebai_graph/preprocessing/bin/pcba/AtomNumHs/indices_one_hot.txt +++ /dev/null @@ -1,6 +0,0 @@ -3 -1 -0 -2 -6 -4 diff --git a/chebai_graph/preprocessing/bin/pcba/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/AtomType/indices_one_hot.txt deleted file mode 100644 index 166ffe2..0000000 --- a/chebai_graph/preprocessing/bin/pcba/AtomType/indices_one_hot.txt +++ /dev/null @@ -1,44 +0,0 @@ -6 -16 -7 -8 -9 -17 -35 -15 -14 -53 -5 -29 -79 -38 -50 -30 -33 -34 -48 -51 -1 -26 -82 -27 -78 -80 -3 -28 -22 -56 -44 -11 -13 -46 -25 -83 -47 -60 -24 -64 -32 -49 -74 -31 diff --git a/chebai_graph/preprocessing/bin/pcba/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/BondLevel/indices_one_hot.txt deleted file mode 100644 index 6b48d66..0000000 --- a/chebai_graph/preprocessing/bin/pcba/BondLevel/indices_one_hot.txt +++ /dev/null @@ -1,4 +0,0 @@ -within_atoms_lvl -atom_fg_lvl -within_fg_lvl -to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/pcba/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/BondType/indices_one_hot.txt deleted file mode 100644 index c819d38..0000000 --- a/chebai_graph/preprocessing/bin/pcba/BondType/indices_one_hot.txt +++ /dev/null @@ -1,4 +0,0 @@ -SINGLE -AROMATIC -DOUBLE -TRIPLE diff --git a/chebai_graph/preprocessing/bin/pcba/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/NumAtomBonds/indices_one_hot.txt deleted file mode 100644 index 2efbda4..0000000 --- a/chebai_graph/preprocessing/bin/pcba/NumAtomBonds/indices_one_hot.txt +++ /dev/null @@ -1,6 +0,0 @@ -1 -2 -3 -4 -0 -5 diff --git a/chebai_graph/preprocessing/bin/sider/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/AtomCharge/indices_one_hot.txt deleted file mode 100644 index 6f89a5d..0000000 --- a/chebai_graph/preprocessing/bin/sider/AtomCharge/indices_one_hot.txt +++ /dev/null @@ -1,6 +0,0 @@ -0 -1 --1 -3 -2 --2 diff --git a/chebai_graph/preprocessing/bin/sider/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/AtomFunctionalGroup/indices_one_hot.txt deleted file mode 100644 index 6fbc2c2..0000000 --- a/chebai_graph/preprocessing/bin/sider/AtomFunctionalGroup/indices_one_hot.txt +++ /dev/null @@ -1,93 +0,0 @@ -alkyl -RING_6 -carboxyl -alkene -fluoro -sulfinyl -RING_5 -graph_fg -tertiary_carbon -secondary_amine -ketone -primary_amine -primary_ketimine -NO_FG -sulfide -hydroxyl -sulfonic_acid -bromo -4_ammonium_ion -chloro -quaternary_carbon -ether -carboxylate -tertiary_amine -amide -RING_3 -trifluoromethyl -ester -nitro -sulfonyl -iodo -secondary_ketimine -RING_7 -aldehyde -phosphono -alkyne -nitrile -carbamate -alkene_carbon -carbonate_ester -RING_4 -RING_18 -RING_24 -RING_14 -imide -sulfonate_ester -RING_81 -RING_20 -RING_8 -nitroso -RING_29 -RING_16 -RING_12 -RING_21 -sulfhydryl -thiolester -difluoromethyl -thioketone -RING_25 -RING_78 -RING_68 -RING_77 -nitrate -azo -RING_36 -phosphate -ketoxime -secondary_aldimine -dichloromethyl -carboxylic_anhydride -RING_33 -borono -RING_15 -diiodomethyl -RING_23 -carbodithio -disulfide -acetal -RING_40 -phosphodiester -fluorochloromethyl -RING_22 -RING_10 -RING_9 -azide -phosphino -RING_31 -RING_28 -RING_53 -RING_56 -RING_41 -RING_43 -RING_35 diff --git a/chebai_graph/preprocessing/bin/sider/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/AtomHybridization/indices_one_hot.txt deleted file mode 100644 index e9958eb..0000000 --- a/chebai_graph/preprocessing/bin/sider/AtomHybridization/indices_one_hot.txt +++ /dev/null @@ -1,6 +0,0 @@ -SP3 -SP2 -S -SP3D -SP -SP3D2 diff --git a/chebai_graph/preprocessing/bin/sider/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/AtomNodeLevel/indices_one_hot.txt deleted file mode 100644 index 2d65776..0000000 --- a/chebai_graph/preprocessing/bin/sider/AtomNodeLevel/indices_one_hot.txt +++ /dev/null @@ -1,3 +0,0 @@ -atom_node_lvl -fg_node_lvl -graph_node_level diff --git a/chebai_graph/preprocessing/bin/sider/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/AtomNumHs/indices_one_hot.txt deleted file mode 100644 index df4a3e9..0000000 --- a/chebai_graph/preprocessing/bin/sider/AtomNumHs/indices_one_hot.txt +++ /dev/null @@ -1,5 +0,0 @@ -3 -1 -0 -2 -4 diff --git a/chebai_graph/preprocessing/bin/sider/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/AtomType/indices_one_hot.txt deleted file mode 100644 index 384eb37..0000000 --- a/chebai_graph/preprocessing/bin/sider/AtomType/indices_one_hot.txt +++ /dev/null @@ -1,40 +0,0 @@ -6 -8 -9 -16 -7 -35 -17 -11 -64 -15 -24 -53 -47 -49 -1 -25 -29 -20 -62 -3 -19 -78 -39 -43 -31 -26 -5 -12 -38 -79 -33 -30 -56 -98 -34 -32 -27 -88 -81 -57 diff --git a/chebai_graph/preprocessing/bin/sider/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/BondLevel/indices_one_hot.txt deleted file mode 100644 index 6b48d66..0000000 --- a/chebai_graph/preprocessing/bin/sider/BondLevel/indices_one_hot.txt +++ /dev/null @@ -1,4 +0,0 @@ -within_atoms_lvl -atom_fg_lvl -within_fg_lvl -to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/sider/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/BondType/indices_one_hot.txt deleted file mode 100644 index f14a241..0000000 --- a/chebai_graph/preprocessing/bin/sider/BondType/indices_one_hot.txt +++ /dev/null @@ -1,4 +0,0 @@ -SINGLE -DOUBLE -AROMATIC -TRIPLE diff --git a/chebai_graph/preprocessing/bin/sider/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/NumAtomBonds/indices_one_hot.txt deleted file mode 100644 index dfb5c6d..0000000 --- a/chebai_graph/preprocessing/bin/sider/NumAtomBonds/indices_one_hot.txt +++ /dev/null @@ -1,6 +0,0 @@ -1 -2 -3 -4 -0 -6 diff --git a/chebai_graph/preprocessing/bin/tox21/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/AtomCharge/indices_one_hot.txt deleted file mode 100644 index a8bd8c3..0000000 --- a/chebai_graph/preprocessing/bin/tox21/AtomCharge/indices_one_hot.txt +++ /dev/null @@ -1,6 +0,0 @@ -0 -1 --1 -2 -3 --2 diff --git a/chebai_graph/preprocessing/bin/tox21/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/AtomFunctionalGroup/indices_one_hot.txt deleted file mode 100644 index 5aa0016..0000000 --- a/chebai_graph/preprocessing/bin/tox21/AtomFunctionalGroup/indices_one_hot.txt +++ /dev/null @@ -1,115 +0,0 @@ -alkyl -hemiacetal -ether -graph_fg -RING_18 -NO_FG -RING_6 -ester -tertiary_carbon -secondary_amine -ketone -carboxyl -RING_5 -amide -primary_amine -RING_4 -carboxylate -alkene -nitro -imide -4_ammonium_ion -bromo -chloro -hydroxyl -sulfonyl -fluoro -quaternary_carbon -alkene_carbon -RING_3 -alkyne -secondary_ketimine -nitrile -isocyanate -azo -tertiary_amine -trifluoromethyl -aldehyde -thiolester -sulfide -nitroso -phosphate -phosphoryl -RING_7 -iodo -acetal -carbodithio -carbonate_ester -RING_21 -difluoromethyl -disulfide -RING_25 -nitrate -carbamate -secondary_aldimine -sulfonic_acid -haloformyl -amidine -sulfhydryl -phosphono -dibromomethyl -RING_23 -trichloromethyl -sulfonate_ester -RING_40 -RING_24 -RING_15 -primary_ketimine -RING_29 -thioketone -aldoxime -RING_16 -nitrosooxy -sulfino -RING_36 -peroxy -RING_14 -ketoxime -thiocyanate -RING_8 -RING_17 -sulfinyl -dichloromethyl -primary_aldimine -tribromomethyl -RING_19 -RING_9 -silyl_ether -phosphodiester -RING_33 -azide -RING_27 -RING_26 -RING_20 -phosphino -RING_12 -trimethylsilyl -carboxylic_anhydride -orthoester -RING_10 -thionoester -isothiocyanate -borono -chlorobromomethyl -dichlorosilane -hydroperoxy -RING_35 -RING_30 -RING_28 -fluorochloromethyl -RING_13 -bromodichloromethyl -ketal -difluorochloromethyl -bromodifluoromethyl -RING_11 diff --git a/chebai_graph/preprocessing/bin/tox21/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/AtomHybridization/indices_one_hot.txt deleted file mode 100644 index cdc6985..0000000 --- a/chebai_graph/preprocessing/bin/tox21/AtomHybridization/indices_one_hot.txt +++ /dev/null @@ -1,7 +0,0 @@ -SP3 -SP2 -SP -S -SP3D2 -SP3D -UNSPECIFIED diff --git a/chebai_graph/preprocessing/bin/tox21/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/AtomNodeLevel/indices_one_hot.txt deleted file mode 100644 index 2d65776..0000000 --- a/chebai_graph/preprocessing/bin/tox21/AtomNodeLevel/indices_one_hot.txt +++ /dev/null @@ -1,3 +0,0 @@ -atom_node_lvl -fg_node_lvl -graph_node_level diff --git a/chebai_graph/preprocessing/bin/tox21/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/AtomNumHs/indices_one_hot.txt deleted file mode 100644 index 5a694e8..0000000 --- a/chebai_graph/preprocessing/bin/tox21/AtomNumHs/indices_one_hot.txt +++ /dev/null @@ -1,6 +0,0 @@ -3 -2 -1 -0 -4 -6 diff --git a/chebai_graph/preprocessing/bin/tox21/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/AtomType/indices_one_hot.txt deleted file mode 100644 index 64e481b..0000000 --- a/chebai_graph/preprocessing/bin/tox21/AtomType/indices_one_hot.txt +++ /dev/null @@ -1,50 +0,0 @@ -6 -8 -7 -16 -15 -35 -17 -9 -56 -30 -26 -80 -53 -28 -5 -51 -14 -47 -11 -78 -40 -79 -4 -29 -20 -34 -81 -50 -82 -23 -32 -38 -48 -46 -60 -27 -83 -64 -25 -19 -33 -22 -3 -12 -42 -66 -24 -1 -49 -70 diff --git a/chebai_graph/preprocessing/bin/tox21/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/BondLevel/indices_one_hot.txt deleted file mode 100644 index 6b48d66..0000000 --- a/chebai_graph/preprocessing/bin/tox21/BondLevel/indices_one_hot.txt +++ /dev/null @@ -1,4 +0,0 @@ -within_atoms_lvl -atom_fg_lvl -within_fg_lvl -to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/tox21/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/BondType/indices_one_hot.txt deleted file mode 100644 index c819d38..0000000 --- a/chebai_graph/preprocessing/bin/tox21/BondType/indices_one_hot.txt +++ /dev/null @@ -1,4 +0,0 @@ -SINGLE -AROMATIC -DOUBLE -TRIPLE diff --git a/chebai_graph/preprocessing/bin/tox21/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/NumAtomBonds/indices_one_hot.txt deleted file mode 100644 index 21b19a8..0000000 --- a/chebai_graph/preprocessing/bin/tox21/NumAtomBonds/indices_one_hot.txt +++ /dev/null @@ -1,7 +0,0 @@ -1 -2 -3 -4 -0 -5 -6 diff --git a/chebai_graph/preprocessing/bin/toxcast/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/AtomCharge/indices_one_hot.txt deleted file mode 100644 index c857ab5..0000000 --- a/chebai_graph/preprocessing/bin/toxcast/AtomCharge/indices_one_hot.txt +++ /dev/null @@ -1,6 +0,0 @@ -0 --1 -1 -2 -3 --2 diff --git a/chebai_graph/preprocessing/bin/toxcast/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/AtomFunctionalGroup/indices_one_hot.txt deleted file mode 100644 index 500ebf5..0000000 --- a/chebai_graph/preprocessing/bin/toxcast/AtomFunctionalGroup/indices_one_hot.txt +++ /dev/null @@ -1,117 +0,0 @@ -ester -RING_6 -sulfonyl -alkyl -NO_FG -4_ammonium_ion -graph_fg -sulfide -primary_amine -carboxyl -chloro -quaternary_carbon -hydroxyl -secondary_amine -fluoro -ketone -tertiary_carbon -ether -alkene -nitro -tertiary_amine -aldehyde -RING_5 -alkyne -nitrile -RING_7 -alkene_carbon -carbonate_ester -trifluoromethyl -imide -sulfonic_acid -primary_ketimine -disulfide -sulfonate_ester -amide -nitroso -iodo -RING_3 -carboxylate -RING_8 -RING_20 -acetal -bromo -carbamate -RING_12 -ketoxime -RING_18 -silyl_ether -RING_4 -dichloromethyl -sulfinyl -aldoxime -thioketone -secondary_ketimine -isothiocyanate -isocyanate -RING_14 -sulfhydryl -phosphono -carbodithioic_acid -secondary_aldimine -haloformyl -phosphoryl -azo -nitrate -RING_35 -thiocyanate -hemiacetal -thiolester -difluorochloromethyl -dibromomethyl -difluoromethyl -phosphino -RING_29 -chlorobromomethyl -RING_9 -tribromomethyl -trichloromethyl -RING_28 -RING_16 -nitrosooxy -bromodichloromethyl -carbodithio -phosphodiester -RING_23 -phosphate -carboxylic_anhydride -RING_13 -RING_24 -RING_30 -RING_11 -RING_17 -azide -RING_40 -primary_aldimine -RING_15 -peroxy -RING_33 -trimethylsilyl -sulfino -ketal -RING_10 -thionoester -borono -amidine -orthoester -RING_27 -RING_26 -RING_38 -hydroperoxy -RING_21 -RING_19 -dichlorosilane -diiodomethyl -fluorochloromethyl -bromodifluoromethyl -RING_36 diff --git a/chebai_graph/preprocessing/bin/toxcast/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/AtomHybridization/indices_one_hot.txt deleted file mode 100644 index 5d039f5..0000000 --- a/chebai_graph/preprocessing/bin/toxcast/AtomHybridization/indices_one_hot.txt +++ /dev/null @@ -1,7 +0,0 @@ -SP2 -SP3 -SP -S -UNSPECIFIED -SP3D2 -SP3D diff --git a/chebai_graph/preprocessing/bin/toxcast/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/AtomNodeLevel/indices_one_hot.txt deleted file mode 100644 index 2d65776..0000000 --- a/chebai_graph/preprocessing/bin/toxcast/AtomNodeLevel/indices_one_hot.txt +++ /dev/null @@ -1,3 +0,0 @@ -atom_node_lvl -fg_node_lvl -graph_node_level diff --git a/chebai_graph/preprocessing/bin/toxcast/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/AtomNumHs/indices_one_hot.txt deleted file mode 100644 index ed1daa5..0000000 --- a/chebai_graph/preprocessing/bin/toxcast/AtomNumHs/indices_one_hot.txt +++ /dev/null @@ -1,5 +0,0 @@ -0 -1 -3 -2 -4 diff --git a/chebai_graph/preprocessing/bin/toxcast/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/AtomType/indices_one_hot.txt deleted file mode 100644 index 9c59d72..0000000 --- a/chebai_graph/preprocessing/bin/toxcast/AtomType/indices_one_hot.txt +++ /dev/null @@ -1,53 +0,0 @@ -8 -6 -7 -16 -17 -9 -50 -19 -22 -14 -11 -53 -35 -33 -49 -78 -15 -51 -12 -3 -83 -28 -26 -13 -20 -47 -29 -5 -1 -30 -38 -80 -27 -34 -79 -56 -81 -24 -23 -21 -64 -70 -60 -46 -32 -4 -82 -66 -25 -48 -63 -40 -42 diff --git a/chebai_graph/preprocessing/bin/toxcast/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/BondLevel/indices_one_hot.txt deleted file mode 100644 index 6b48d66..0000000 --- a/chebai_graph/preprocessing/bin/toxcast/BondLevel/indices_one_hot.txt +++ /dev/null @@ -1,4 +0,0 @@ -within_atoms_lvl -atom_fg_lvl -within_fg_lvl -to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/toxcast/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/BondType/indices_one_hot.txt deleted file mode 100644 index f14a241..0000000 --- a/chebai_graph/preprocessing/bin/toxcast/BondType/indices_one_hot.txt +++ /dev/null @@ -1,4 +0,0 @@ -SINGLE -DOUBLE -AROMATIC -TRIPLE diff --git a/chebai_graph/preprocessing/bin/toxcast/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/NumAtomBonds/indices_one_hot.txt deleted file mode 100644 index a61fff5..0000000 --- a/chebai_graph/preprocessing/bin/toxcast/NumAtomBonds/indices_one_hot.txt +++ /dev/null @@ -1,7 +0,0 @@ -2 -1 -3 -4 -0 -6 -5 From fe4d2f047e9909a8a70a3df06e2e35933d31ce03 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sun, 11 Oct 2026 12:01:34 +0200 Subject: [PATCH 33/58] fixes --- chebai_graph/loss/msp.py | 43 ++++++++++++++--------- chebai_graph/models/architectures/base.py | 18 +++++----- 2 files changed, 36 insertions(+), 25 deletions(-) diff --git a/chebai_graph/loss/msp.py b/chebai_graph/loss/msp.py index 20055b8..67caa4c 100644 --- a/chebai_graph/loss/msp.py +++ b/chebai_graph/loss/msp.py @@ -2,7 +2,7 @@ Multi-level Self-supervised Pre-training from Zang, X., Zhao, X. & Tang, B. Hierarchical Molecular Graph Self-Supervised Learning for property prediction. Commun Chem 6, 34 (2023). https://doi.org/10.1038/s42004-023-00825-5 -The files containes modified code from: +This file contains modified code from: - https://github.com/ZangXuan/HiMol/blob/main/pretrain.py """ @@ -141,28 +141,28 @@ def _metadata_targets(metadata: dict, device: torch.device): atom_vocab = list(range(1, 119)) # Atomic numbers from 1 to 118 bond_vocab = ["SINGLE", "DOUBLE", "AROMATIC", "TRIPLE", "DATIVE", "UNSPECIFIED"] atom_mol = metadata["nodes"]["atom_nodes"] - atom_targets = torch.tensor( + atom_type_targets = torch.tensor( [atom_vocab.index(atom.GetAtomicNum()) for atom in atom_mol.GetAtoms()], dtype=torch.long, device=device, ) within_atoms = metadata["edges"]["within_atoms_lvl"] bond_indices = [] - bond_targets = [] + bond_type_targets = [] for bond in within_atoms.GetBonds(): begin = bond.GetBeginAtomIdx() end = bond.GetEndAtomIdx() bond_class = bond_vocab.index(str(bond.GetBondType())) bond_indices.extend(((begin, end), (end, begin))) - bond_targets.extend((bond_class, bond_class)) + bond_type_targets.extend((bond_class, bond_class)) edge_index = ( torch.tensor(bond_indices, dtype=torch.long, device=device).t() if bond_indices else torch.empty((2, 0), dtype=torch.long, device=device) ) return ( - atom_targets, - torch.tensor(bond_targets, dtype=torch.long, device=device), + atom_type_targets, + torch.tensor(bond_type_targets, dtype=torch.long, device=device), edge_index, ) @@ -183,6 +183,7 @@ def forward( ] graph_node_embeddings = self._super(graph_data, node_embeddings, graph_nodes) + # Number of atoms prediction loss atom_counts = torch.stack( [ (is_atom_node & (batch_index == graph_index)).sum() @@ -202,15 +203,16 @@ def forward( for graph_index, nodes in enumerate(graph_nodes): metadata = self._molecule_metadata(custom_metadata, graph_index) mol = metadata["nodes"]["atom_nodes"] - atom_targets, bond_targets, local_edges = self._metadata_targets( + atom_type_targets, bond_type_targets, local_edges = self._metadata_targets( metadata, node_embeddings.device ) - # Bond existence prediction + + # Bond existence prediction loss atoms = nodes[is_atom_node[nodes]] - if atoms.numel() != atom_targets.numel(): + if atoms.numel() != atom_type_targets.numel(): raise ValueError( f"Graph {graph_index} has {atoms.numel()} atom nodes but metadata has " - f"{atom_targets.numel()} atoms" + f"{atom_type_targets.numel()} atoms" ) atom_embeddings = self.feat_drop(self.bond_if_proj(node_embeddings[atoms])) bond_counts.append(mol.GetNumBonds()) @@ -227,6 +229,7 @@ def forward( if local_edges.numel(): bond_target[local_edges[0], local_edges[1]] = 1 bond_if_loss = bond_if_loss + self.bond_if_loss(bond_pred, bond_target) + if loss_kwargs.get("compute_metrics", True) and local_edges.numel(): target_cpu = bond_target.flatten().detach().cpu() pred_cpu = bond_pred.flatten().detach().sigmoid().cpu() @@ -234,12 +237,15 @@ def forward( auc_values.append(roc_auc_score(target_cpu, pred_cpu)) ap_values.append(average_precision_score(target_cpu, pred_cpu)) + # Atom type prediction loss atom_pred = self.atom_type_s(node_embeddings[atoms]) atom_type_loss = atom_type_loss + self.atom_type_loss( - atom_pred, atom_targets + atom_pred, atom_type_targets ) - atom_correct += int((atom_pred.argmax(1) == atom_targets).sum()) - atom_total += atom_targets.numel() + atom_correct += int((atom_pred.argmax(1) == atom_type_targets).sum()) + atom_total += atom_type_targets.numel() + + # Bond type prediction loss if local_edges.numel(): edge_rep = torch.cat( [atom_embeddings[local_edges[0]], atom_embeddings[local_edges[1]]], @@ -247,11 +253,14 @@ def forward( ) bond_pred_type = self.bond_type_s(edge_rep) bond_type_loss = bond_type_loss + self.bond_type_loss( - bond_pred_type, bond_targets + bond_pred_type, bond_type_targets ) - bond_correct += int((bond_pred_type.argmax(1) == bond_targets).sum()) - bond_total += bond_targets.numel() + bond_correct += int( + (bond_pred_type.argmax(1) == bond_type_targets).sum() + ) + bond_total += bond_type_targets.numel() + # Number of bonds prediction loss bond_counts = torch.tensor( bond_counts, dtype=graph_node_embeddings.dtype, @@ -259,8 +268,8 @@ def forward( ) bond_count_pred = self.bond_num_s(graph_node_embeddings).squeeze(-1) bond_count_loss = self.count_loss(bond_count_pred, bond_counts) - divisor = max(len(graph_nodes), 1) + divisor = max(len(graph_nodes), 1) losses = torch.stack( [ bond_if_loss / divisor, diff --git a/chebai_graph/models/architectures/base.py b/chebai_graph/models/architectures/base.py index 9b5fd81..b6f6e26 100644 --- a/chebai_graph/models/architectures/base.py +++ b/chebai_graph/models/architectures/base.py @@ -123,7 +123,9 @@ def __init__( def _load_pretrained_gnn(self, checkpoint_path: str) -> None: """Load encoder weights while leaving the task-specific head untouched.""" - checkpoint = torch.load(checkpoint_path, weights_only=False) + checkpoint = torch.load( + checkpoint_path, weights_only=False, map_location=self.device + ) state_dict = checkpoint.get("state_dict", checkpoint) if not isinstance(state_dict, dict): @@ -155,13 +157,13 @@ def _load_pretrained_gnn(self, checkpoint_path: str) -> None: key: value for key, value in state_dict.items() if key in gnn_keys } skipped_edge_weights = [] - for key in list(pretrained_gnn_state): - if ( - key.endswith(".lin_edge.weight") - and pretrained_gnn_state[key].shape != gnn_state[key].shape - ): - skipped_edge_weights.append(key) - del pretrained_gnn_state[key] + # for key in list(pretrained_gnn_state): + # if ( + # key.endswith(".lin_edge.weight") + # and pretrained_gnn_state[key].shape != gnn_state[key].shape + # ): + # skipped_edge_weights.append(key) + # del pretrained_gnn_state[key] load_result = self.gnn.load_state_dict(pretrained_gnn_state, strict=False) allowed_missing = set(skipped_edge_weights) From 2a8048f20f620a27b7f993ef1924161ed7141243 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Wed, 14 Oct 2026 19:24:11 +0200 Subject: [PATCH 34/58] use existing vocab for msp --- chebai_graph/loss/msp.py | 41 +++++++++++-------- .../preprocessing/property_encoder.py | 3 ++ configs/model/pretraining/gat_msp.yml | 4 +- 3 files changed, 27 insertions(+), 21 deletions(-) diff --git a/chebai_graph/loss/msp.py b/chebai_graph/loss/msp.py index 67caa4c..8117497 100644 --- a/chebai_graph/loss/msp.py +++ b/chebai_graph/loss/msp.py @@ -10,6 +10,12 @@ from sklearn.metrics import average_precision_score, roc_auc_score from torch_geometric.data import Data +from chebai_graph.preprocessing.properties.augmented_properties import ( + AugAtomType, + AugBondType, +) +from chebai_graph.preprocessing.property_encoder import IndexEncoder + class MSPPretrainingLoss(torch.nn.Module): """HiMol-style molecular structure pretraining loss. @@ -22,19 +28,15 @@ class MSPPretrainingLoss(torch.nn.Module): def __init__( self, hidden_size: int, - atom_num_classes: int = 118, - bond_num_classes: int = 4, - atom_target_start: int = 0, - atom_target_width: int = 1, - bond_target_start: int = 0, - bond_target_width: int = 1, dropout: float = 0.2, ) -> None: super().__init__() - self.atom_target_start = atom_target_start - self.atom_target_width = atom_target_width - self.bond_target_start = bond_target_start - self.bond_target_width = bond_target_width + self.atom_type_encoder = IndexEncoder( + property=AugAtomType(), use_one_hot_cache=True, eval=True + ) + self.bond_type_encoder = IndexEncoder( + property=AugBondType(), use_one_hot_cache=True, eval=True + ) # Bond existence projection self.bond_if_proj = torch.nn.Sequential( @@ -53,13 +55,13 @@ def __init__( self.bond_type_s = torch.nn.Sequential( torch.nn.Linear(2 * hidden_size, hidden_size), torch.nn.ReLU(), - torch.nn.Linear(hidden_size, bond_num_classes), + torch.nn.Linear(hidden_size, len(self.bond_type_encoder.cache)), ) # Atom type classifier self.atom_type_s = torch.nn.Sequential( torch.nn.Linear(hidden_size, hidden_size), torch.nn.ReLU(), - torch.nn.Linear(hidden_size, atom_num_classes), + torch.nn.Linear(hidden_size, len(self.atom_type_encoder.cache)), ) count_hidden = max(hidden_size // 4, 1) @@ -136,13 +138,16 @@ def _molecule_metadata(custom_metadata, graph_index: int) -> dict: ) return metadata - @staticmethod - def _metadata_targets(metadata: dict, device: torch.device): - atom_vocab = list(range(1, 119)) # Atomic numbers from 1 to 118 - bond_vocab = ["SINGLE", "DOUBLE", "AROMATIC", "TRIPLE", "DATIVE", "UNSPECIFIED"] + def _metadata_targets(self, metadata: dict, device: torch.device): + # atom_vocab = list(range(1, 119)) # Atomic numbers from 1 to 118 + # bond_vocab = ["SINGLE", "DOUBLE", "AROMATIC", "TRIPLE", "DATIVE", "UNSPECIFIED"] atom_mol = metadata["nodes"]["atom_nodes"] + atom_type_targets = torch.tensor( - [atom_vocab.index(atom.GetAtomicNum()) for atom in atom_mol.GetAtoms()], + [ + self.atom_type_encoder.encode(str(atom.GetAtomicNum())) + for atom in atom_mol.GetAtoms() + ], dtype=torch.long, device=device, ) @@ -152,7 +157,7 @@ def _metadata_targets(metadata: dict, device: torch.device): for bond in within_atoms.GetBonds(): begin = bond.GetBeginAtomIdx() end = bond.GetEndAtomIdx() - bond_class = bond_vocab.index(str(bond.GetBondType())) + bond_class = self.bond_type_encoder.encode(str(bond.GetBondType())) bond_indices.extend(((begin, end), (end, begin))) bond_type_targets.extend((bond_class, bond_class)) edge_index = ( diff --git a/chebai_graph/preprocessing/property_encoder.py b/chebai_graph/preprocessing/property_encoder.py index 38cb279..73becc3 100644 --- a/chebai_graph/preprocessing/property_encoder.py +++ b/chebai_graph/preprocessing/property_encoder.py @@ -104,6 +104,7 @@ def __init__(self, property, indices_dir: str | None = None, **kwargs) -> None: if indices_dir is None: indices_dir = os.path.dirname(inspect.getfile(self.__class__)) self.dirname = indices_dir + self._use_one_hot_cache = kwargs.get("use_one_hot_cache", False) # load already existing cache with open(self.index_path, "r") as pk: self.cache: dict[str, int] = { @@ -117,6 +118,8 @@ def __init__(self, property, indices_dir: str | None = None, **kwargs) -> None: @property def name(self) -> str: """Name of this encoder.""" + if self._use_one_hot_cache: + return "one_hot" return "index" @property diff --git a/configs/model/pretraining/gat_msp.yml b/configs/model/pretraining/gat_msp.yml index 4f4d617..76c5a60 100644 --- a/configs/model/pretraining/gat_msp.yml +++ b/configs/model/pretraining/gat_msp.yml @@ -4,8 +4,6 @@ init_args: class_path: chebai_graph.loss.msp.MSPPretrainingLoss init_args: hidden_size: 512 - atom_num_classes: 11 - bond_num_classes: 4 optimizer_kwargs: lr: 1e-3 config: @@ -13,7 +11,7 @@ init_args: hidden_channels: 256 out_channels: 512 num_layers: 4 - edge_dim: 10 # number of bond properties + edge_dim: 12 # number of bond properties heads: 8 # the number of heads should be divisible by output channels (hidden channels if output channel not given) v2: True # This uses `torch_geometric.nn.conv.GATv2Conv` convolution layers, set False to use `GATConv` dropout: 0.5 From 05e3e9e9f3703723e61e0eb152ced66cc211b8d7 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Fri, 16 Oct 2026 16:02:26 +0200 Subject: [PATCH 35/58] Reapply "Separate Token files based on data type (#47)" This reverts commit 40076b9ce70257470804328f6ba7bc494ed27244. --- .../AtomCharge/indices_one_hot.txt | 0 .../AtomFunctionalGroup/indices_one_hot.txt | 0 .../AtomHybridization/indices_one_hot.txt | 0 .../AtomNodeLevel/indices_one_hot.txt | 0 .../{ => chebi}/AtomNumHs/indices_one_hot.txt | 0 .../{ => chebi}/AtomType/indices_one_hot.txt | 0 .../{ => chebi}/BondLevel/indices_one_hot.txt | 0 .../{ => chebi}/BondType/indices_one_hot.txt | 0 .../NumAtomBonds/indices_one_hot.txt | 0 chebai_graph/preprocessing/datasets/base.py | 2 +- chebai_graph/preprocessing/datasets/utils.py | 10 ++- .../properties/augmented_properties.py | 40 +++++++---- chebai_graph/preprocessing/properties/base.py | 14 ++-- .../preprocessing/properties/properties.py | 72 ++++++++++++------- .../preprocessing/property_encoder.py | 16 +++-- 15 files changed, 103 insertions(+), 51 deletions(-) rename chebai_graph/preprocessing/bin/{ => chebi}/AtomCharge/indices_one_hot.txt (100%) rename chebai_graph/preprocessing/bin/{ => chebi}/AtomFunctionalGroup/indices_one_hot.txt (100%) rename chebai_graph/preprocessing/bin/{ => chebi}/AtomHybridization/indices_one_hot.txt (100%) rename chebai_graph/preprocessing/bin/{ => chebi}/AtomNodeLevel/indices_one_hot.txt (100%) rename chebai_graph/preprocessing/bin/{ => chebi}/AtomNumHs/indices_one_hot.txt (100%) rename chebai_graph/preprocessing/bin/{ => chebi}/AtomType/indices_one_hot.txt (100%) rename chebai_graph/preprocessing/bin/{ => chebi}/BondLevel/indices_one_hot.txt (100%) rename chebai_graph/preprocessing/bin/{ => chebi}/BondType/indices_one_hot.txt (100%) rename chebai_graph/preprocessing/bin/{ => chebi}/NumAtomBonds/indices_one_hot.txt (100%) diff --git a/chebai_graph/preprocessing/bin/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/chebi/AtomCharge/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/AtomCharge/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/chebi/AtomCharge/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/chebi/AtomFunctionalGroup/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/AtomFunctionalGroup/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/chebi/AtomFunctionalGroup/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/chebi/AtomHybridization/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/AtomHybridization/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/chebi/AtomHybridization/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/chebi/AtomNodeLevel/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/AtomNodeLevel/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/chebi/AtomNodeLevel/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/chebi/AtomNumHs/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/AtomNumHs/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/chebi/AtomNumHs/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/chebi/AtomType/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/AtomType/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/chebi/AtomType/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/chebi/BondLevel/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/BondLevel/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/chebi/BondLevel/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/chebi/BondType/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/BondType/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/chebi/BondType/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/chebi/NumAtomBonds/indices_one_hot.txt similarity index 100% rename from chebai_graph/preprocessing/bin/NumAtomBonds/indices_one_hot.txt rename to chebai_graph/preprocessing/bin/chebi/NumAtomBonds/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/datasets/base.py b/chebai_graph/preprocessing/datasets/base.py index 7513f76..d94a333 100644 --- a/chebai_graph/preprocessing/datasets/base.py +++ b/chebai_graph/preprocessing/datasets/base.py @@ -50,7 +50,7 @@ def __init__( super().__init__(**kwargs) # atom_properties and bond_properties are given as lists containing class_paths if properties is not None: - properties = [resolve_property(prop) for prop in properties] + properties = [resolve_property(prop, self.data_type) for prop in properties] properties = self._sort_properties(properties) else: properties = [] diff --git a/chebai_graph/preprocessing/datasets/utils.py b/chebai_graph/preprocessing/datasets/utils.py index 3ec6515..79f2fb3 100644 --- a/chebai_graph/preprocessing/datasets/utils.py +++ b/chebai_graph/preprocessing/datasets/utils.py @@ -4,7 +4,9 @@ from chebai_graph.preprocessing.properties import MolecularProperty -def resolve_property(property: str | MolecularProperty) -> MolecularProperty: +def resolve_property( + property: str | MolecularProperty, data_type: str +) -> MolecularProperty: """ Resolves a molecular property specification (either as a class instance or class path string) into a MolecularProperty instance. @@ -19,6 +21,8 @@ def resolve_property(property: str | MolecularProperty) -> MolecularProperty: property (str | MolecularProperty): The property to resolve. Can be a class instance, a fully qualified class name (e.g. "module.ClassName"), or a class name assumed to be in `chebai_graph.preprocessing.properties`. + data_type (str): The data type associated with the property. This used to determine or set + tokens file path for the property if applicable. Returns: MolecularProperty: An instance of the resolved MolecularProperty. @@ -37,7 +41,7 @@ def resolve_property(property: str | MolecularProperty) -> MolecularProperty: module_name = property[:last_dot] class_name = property[last_dot + 1 :] module = importlib.import_module(module_name) - return getattr(module, class_name)() + return getattr(module, class_name)(data_type=data_type) except ValueError: # if only a class name is given, assume the module is chebai_graph.processing.properties - return getattr(graph_properties, property)() + return getattr(graph_properties, property)(data_type=data_type) diff --git a/chebai_graph/preprocessing/properties/augmented_properties.py b/chebai_graph/preprocessing/properties/augmented_properties.py index f5f7b1d..6f112ea 100644 --- a/chebai_graph/preprocessing/properties/augmented_properties.py +++ b/chebai_graph/preprocessing/properties/augmented_properties.py @@ -23,14 +23,17 @@ class AtomNodeLevel(AllNodeTypeProperty): - def __init__(self, encoder: PropertyEncoder | None = None): + def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: """ Initialize AtomNodeLevel with an optional encoder. Args: encoder (PropertyEncoder | None): Property encoder to use. Defaults to OneHotEncoder. """ - super().__init__(encoder or OneHotEncoder(self)) + data_type = kwargs.get("data_type") + super().__init__( + encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + ) def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> str | int | bool: """ @@ -46,14 +49,17 @@ def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> str | int | bool: class AtomFunctionalGroup(FGNodeTypeProperty): - def __init__(self, encoder: PropertyEncoder | None = None): + def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: """ Initialize AtomFunctionalGroup with an optional encoder. Args: encoder (PropertyEncoder | None): Property encoder to use. Defaults to OneHotEncoder. """ - super().__init__(encoder or OneHotEncoder(self)) + data_type = kwargs.get("data_type") + super().__init__( + encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + ) def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> str | int | bool: """ @@ -69,14 +75,17 @@ def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> str | int | bool: class AtomRingSize(FGNodeTypeProperty): - def __init__(self, encoder: PropertyEncoder | None = None): + def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: """ Initialize AtomRingSize with an optional encoder. Args: encoder (PropertyEncoder | None): Property encoder to use. Defaults to OneHotEncoder. """ - super().__init__(encoder or OneHotEncoder(self)) + data_type = kwargs.get("data_type") + super().__init__( + encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + ) def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> int: """ @@ -113,14 +122,14 @@ def _check_modify_atom_prop_value( class IsHydrogenBondDonorFG(FGNodeTypeProperty): - def __init__(self, encoder: PropertyEncoder | None = None): + def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: """ Initialize IsHydrogenBondDonorFG with an optional encoder. Args: encoder (PropertyEncoder | None): Property encoder to use. Defaults to BoolEncoder. """ - super().__init__(encoder or BoolEncoder(self)) + super().__init__(encoder=encoder or BoolEncoder(self), **kwargs) # fmt: off # https://github.com/thaonguyen217/farm_molecular_representation/blob/main/src/(6)gen_FG_KG.py#L26-L31 self._hydrogen_bond_donor: set[str] = { @@ -146,14 +155,14 @@ def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> bool: class IsHydrogenBondAcceptorFG(FGNodeTypeProperty): - def __init__(self, encoder: PropertyEncoder | None = None): + def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: """ Initialize IsHydrogenBondAcceptorFG with an optional encoder. Args: encoder (PropertyEncoder | None): Property encoder to use. Defaults to BoolEncoder. """ - super().__init__(encoder or BoolEncoder(self)) + super().__init__(encoder=encoder or BoolEncoder(self), **kwargs) # fmt: off # https://github.com/thaonguyen217/farm_molecular_representation/blob/main/src/(6)gen_FG_KG.py#L33-L39 self._hydrogen_bond_acceptor: set[str] = { @@ -180,13 +189,13 @@ def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> bool: class IsFGAlkyl(FGNodeTypeProperty): - def __init__(self, encoder: PropertyEncoder | None = None): + def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: """ Args: encoder (PropertyEncoder | None): Optional encoder to use for this property. Defaults to BoolEncoder if not provided. """ - super().__init__(encoder or BoolEncoder(self)) + super().__init__(encoder=encoder or BoolEncoder(self), **kwargs) def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> int: """ @@ -323,12 +332,15 @@ class AugAtomAromaticity(AugNodeValueDefaulter, pr.AtomAromaticity): class BondLevel(AugmentedBondProperty): - def __init__(self, encoder: PropertyEncoder | None = None): + def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: """ Args: encoder (PropertyEncoder | None): Optional encoder to use. Defaults to OneHotEncoder. """ - super().__init__(encoder or OneHotEncoder(self)) + data_type = kwargs.get("data_type") + super().__init__( + encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + ) def get_bond_value(self, bond: Chem.rdchem.Bond | dict) -> str: """ diff --git a/chebai_graph/preprocessing/properties/base.py b/chebai_graph/preprocessing/properties/base.py index ced740d..1f2e0aa 100644 --- a/chebai_graph/preprocessing/properties/base.py +++ b/chebai_graph/preprocessing/properties/base.py @@ -31,10 +31,16 @@ class MolecularProperty(ABC): Defaults to IndexEncoder if not provided. """ - def __init__(self, encoder: PropertyEncoder | None = None) -> None: + def __init__( + self, + data_type: str, + encoder: PropertyEncoder | None = None, + ) -> None: + assert data_type is not None, "data_type must be provided for MolecularProperty" if encoder is None: - encoder = IndexEncoder(self) + encoder = IndexEncoder(self, data_type=data_type) self.encoder: PropertyEncoder = encoder + self._data_type = data_type @property def name(self) -> str: @@ -174,8 +180,8 @@ class AugAtomType(FrozenPropertyAlias, AtomType): ... ValueError: If new tokens are added to the frozen encoder during processing. """ - def __init__(self, encoder: PropertyEncoder | None = None) -> None: - super().__init__(encoder) + def __init__(self, encoder: PropertyEncoder, **kwargs) -> None: + super().__init__(encoder=encoder, **kwargs) # Lock the encoder's cache to prevent adding new tokens if hasattr(self.encoder, "cache") and isinstance(self.encoder.cache, dict): self.encoder.cache = MappingProxyType(self.encoder.cache) diff --git a/chebai_graph/preprocessing/properties/properties.py b/chebai_graph/preprocessing/properties/properties.py index 2154f9c..87112da 100644 --- a/chebai_graph/preprocessing/properties/properties.py +++ b/chebai_graph/preprocessing/properties/properties.py @@ -19,8 +19,11 @@ class AtomType(AtomProperty): Uses a one-hot encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None) -> None: - super().__init__(encoder or OneHotEncoder(self)) + def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: + data_type = kwargs.get("data_type") + super().__init__( + encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + ) def get_atom_value(self, atom: Chem.rdchem.Atom) -> int: """ @@ -42,8 +45,11 @@ class NumAtomBonds(AtomProperty): Uses a one-hot encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None) -> None: - super().__init__(encoder or OneHotEncoder(self)) + def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: + data_type = kwargs.get("data_type") + super().__init__( + encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + ) def get_atom_value(self, atom: Chem.rdchem.Atom) -> int: """ @@ -65,8 +71,11 @@ class AtomCharge(AtomProperty): Uses a one-hot encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None) -> None: - super().__init__(encoder or OneHotEncoder(self)) + def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: + data_type = kwargs.get("data_type") + super().__init__( + encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + ) def get_atom_value(self, atom: Chem.rdchem.Atom) -> int: """ @@ -88,8 +97,11 @@ class AtomChirality(AtomProperty): Uses a one-hot encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None) -> None: - super().__init__(encoder or OneHotEncoder(self)) + def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: + data_type = kwargs.get("data_type") + super().__init__( + encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + ) def get_atom_value(self, atom: Chem.rdchem.Atom) -> Chem.rdchem.ChiralType: """ @@ -111,8 +123,11 @@ class AtomHybridization(AtomProperty): Uses a one-hot encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None) -> None: - super().__init__(encoder or OneHotEncoder(self)) + def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: + data_type = kwargs.get("data_type") + super().__init__( + encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + ) def get_atom_value(self, atom: Chem.rdchem.Atom) -> Chem.rdchem.HybridizationType: """ @@ -134,8 +149,11 @@ class AtomNumHs(AtomProperty): Uses a one-hot encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None) -> None: - super().__init__(encoder or OneHotEncoder(self)) + def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: + data_type = kwargs.get("data_type") + super().__init__( + encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + ) def get_atom_value(self, atom: Chem.rdchem.Atom) -> int: """ @@ -157,8 +175,8 @@ class AtomAromaticity(AtomProperty): Uses a boolean encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None) -> None: - super().__init__(encoder or BoolEncoder(self)) + def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: + super().__init__(encoder=encoder or BoolEncoder(self), **kwargs) def get_atom_value(self, atom: Chem.rdchem.Atom) -> bool: """ @@ -180,8 +198,8 @@ class BondAromaticity(BondProperty): Uses a boolean encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None) -> None: - super().__init__(encoder or BoolEncoder(self)) + def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: + super().__init__(encoder=encoder or BoolEncoder(self), **kwargs) def get_bond_value(self, bond: Chem.rdchem.Bond) -> bool: """ @@ -203,8 +221,11 @@ class BondType(BondProperty): Uses a one-hot encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None) -> None: - super().__init__(encoder or OneHotEncoder(self)) + def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: + data_type = kwargs.get("data_type") + super().__init__( + encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + ) def get_bond_value(self, bond: Chem.rdchem.Bond) -> Chem.rdchem.BondType: """ @@ -226,8 +247,8 @@ class BondInRing(BondProperty): Uses a boolean encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None) -> None: - super().__init__(encoder or BoolEncoder(self)) + def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: + super().__init__(encoder=encoder or BoolEncoder(self), **kwargs) def get_bond_value(self, bond: Chem.rdchem.Bond) -> bool: """ @@ -249,8 +270,11 @@ class MoleculeNumRings(MoleculeProperty): Uses a one-hot encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None) -> None: - super().__init__(encoder or OneHotEncoder(self)) + def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: + data_type = kwargs.get("data_type") + super().__init__( + encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + ) def get_property_value(self, mol: Chem.rdchem.Mol) -> list[int]: """ @@ -272,8 +296,8 @@ class RDKit2DNormalized(MoleculeProperty): Uses an identity encoder by default. """ - def __init__(self, encoder: PropertyEncoder | None = None) -> None: - super().__init__(encoder or AsIsEncoder(self)) + def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: + super().__init__(encoder=encoder or AsIsEncoder(self), **kwargs) self.generator_normalized = rdNormalizedDescriptors.RDKit2DNormalized() # Create a dummy molecule (e.g., methane) to extract the length of descriptor vector dummy_mol = Chem.MolFromSmiles("C") diff --git a/chebai_graph/preprocessing/property_encoder.py b/chebai_graph/preprocessing/property_encoder.py index 73becc3..2689a63 100644 --- a/chebai_graph/preprocessing/property_encoder.py +++ b/chebai_graph/preprocessing/property_encoder.py @@ -99,13 +99,16 @@ class IndexEncoder(PropertyEncoder): **kwargs: Additional keyword arguments. """ - def __init__(self, property, indices_dir: str | None = None, **kwargs) -> None: + def __init__( + self, property, data_type: str, indices_dir: str | None = None, **kwargs + ) -> None: super().__init__(property, **kwargs) if indices_dir is None: indices_dir = os.path.dirname(inspect.getfile(self.__class__)) self.dirname = indices_dir self._use_one_hot_cache = kwargs.get("use_one_hot_cache", False) # load already existing cache + self._data_type = data_type with open(self.index_path, "r") as pk: self.cache: dict[str, int] = { token.strip(): idx for idx, token in enumerate(pk) @@ -129,12 +132,15 @@ def index_path(self) -> str: Returns: Path to index file. """ + assert self._data_type is not None, "data_type must be set for IndexEncoder" index_path = os.path.join( - self.dirname, "bin", self.property.name, f"indices_{self.name}.txt" - ) - os.makedirs( - os.path.join(self.dirname, "bin", self.property.name), exist_ok=True + self.dirname, + "bin", + self._data_type, + self.property.name, + f"indices_{self.name}.txt", ) + os.makedirs(os.path.dirname(index_path), exist_ok=True) if not os.path.exists(index_path): with open(index_path, "x"): pass From 30422b09fdddaca46ca3a848a161c7666a92d268 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Fri, 16 Oct 2026 16:43:06 +0200 Subject: [PATCH 36/58] commit tokens of all dataset and _shared --- .../_shared/AtomCharge/indices_one_hot.txt | 13 ++ .../AtomFunctionalGroup/indices_one_hot.txt | 160 ++++++++++++++++++ .../AtomHybridization/indices_one_hot.txt | 7 + .../_shared/AtomNodeLevel/indices_one_hot.txt | 3 + .../bin/_shared/AtomNumHs/indices_one_hot.txt | 8 + .../bin/_shared/AtomType/indices_one_hot.txt | 119 +++++++++++++ .../bin/_shared/BondLevel/indices_one_hot.txt | 4 + .../bin/_shared/BondType/indices_one_hot.txt | 6 + .../_shared/NumAtomBonds/indices_one_hot.txt | 13 ++ .../bin/bace/AtomCharge/indices_one_hot.txt | 3 + .../AtomFunctionalGroup/indices_one_hot.txt | 44 +++++ .../AtomHybridization/indices_one_hot.txt | 3 + .../bace/AtomNodeLevel/indices_one_hot.txt | 3 + .../bin/bace/AtomNumHs/indices_one_hot.txt | 4 + .../bin/bace/AtomType/indices_one_hot.txt | 8 + .../bin/bace/BondLevel/indices_one_hot.txt | 4 + .../bin/bace/BondType/indices_one_hot.txt | 4 + .../bin/bace/NumAtomBonds/indices_one_hot.txt | 4 + .../bin/bbbp/AtomCharge/indices_one_hot.txt | 4 + .../AtomFunctionalGroup/indices_one_hot.txt | 80 +++++++++ .../AtomHybridization/indices_one_hot.txt | 4 + .../bbbp/AtomNodeLevel/indices_one_hot.txt | 3 + .../bin/bbbp/AtomNumHs/indices_one_hot.txt | 4 + .../bin/bbbp/AtomType/indices_one_hot.txt | 13 ++ .../bin/bbbp/BondLevel/indices_one_hot.txt | 4 + .../bin/bbbp/BondType/indices_one_hot.txt | 4 + .../bin/bbbp/NumAtomBonds/indices_one_hot.txt | 5 + .../clin_tox/AtomCharge/indices_one_hot.txt | 6 + .../AtomFunctionalGroup/indices_one_hot.txt | 84 +++++++++ .../AtomHybridization/indices_one_hot.txt | 7 + .../AtomNodeLevel/indices_one_hot.txt | 3 + .../clin_tox/AtomNumHs/indices_one_hot.txt | 4 + .../bin/clin_tox/AtomType/indices_one_hot.txt | 30 ++++ .../clin_tox/BondLevel/indices_one_hot.txt | 4 + .../bin/clin_tox/BondType/indices_one_hot.txt | 5 + .../clin_tox/NumAtomBonds/indices_one_hot.txt | 6 + .../bin/hiv/AtomCharge/indices_one_hot.txt | 9 + .../AtomFunctionalGroup/indices_one_hot.txt | 122 +++++++++++++ .../hiv/AtomHybridization/indices_one_hot.txt | 7 + .../bin/hiv/AtomNodeLevel/indices_one_hot.txt | 3 + .../bin/hiv/AtomNumHs/indices_one_hot.txt | 5 + .../bin/hiv/AtomType/indices_one_hot.txt | 54 ++++++ .../bin/hiv/BondLevel/indices_one_hot.txt | 4 + .../bin/hiv/BondType/indices_one_hot.txt | 5 + .../bin/hiv/NumAtomBonds/indices_one_hot.txt | 11 ++ .../bin/muv/AtomCharge/indices_one_hot.txt | 3 + .../AtomFunctionalGroup/indices_one_hot.txt | 68 ++++++++ .../muv/AtomHybridization/indices_one_hot.txt | 3 + .../bin/muv/AtomNodeLevel/indices_one_hot.txt | 3 + .../bin/muv/AtomNumHs/indices_one_hot.txt | 4 + .../bin/muv/AtomType/indices_one_hot.txt | 7 + .../bin/muv/BondLevel/indices_one_hot.txt | 4 + .../bin/muv/BondType/indices_one_hot.txt | 4 + .../bin/muv/NumAtomBonds/indices_one_hot.txt | 4 + .../bin/pcba/AtomCharge/indices_one_hot.txt | 8 + .../AtomFunctionalGroup/indices_one_hot.txt | 127 ++++++++++++++ .../AtomHybridization/indices_one_hot.txt | 7 + .../pcba/AtomNodeLevel/indices_one_hot.txt | 3 + .../bin/pcba/AtomNumHs/indices_one_hot.txt | 6 + .../bin/pcba/AtomType/indices_one_hot.txt | 44 +++++ .../bin/pcba/BondLevel/indices_one_hot.txt | 4 + .../bin/pcba/BondType/indices_one_hot.txt | 4 + .../bin/pcba/NumAtomBonds/indices_one_hot.txt | 6 + .../bin/sider/AtomCharge/indices_one_hot.txt | 6 + .../AtomFunctionalGroup/indices_one_hot.txt | 93 ++++++++++ .../AtomHybridization/indices_one_hot.txt | 6 + .../sider/AtomNodeLevel/indices_one_hot.txt | 3 + .../bin/sider/AtomNumHs/indices_one_hot.txt | 5 + .../bin/sider/AtomType/indices_one_hot.txt | 40 +++++ .../bin/sider/BondLevel/indices_one_hot.txt | 4 + .../bin/sider/BondType/indices_one_hot.txt | 4 + .../sider/NumAtomBonds/indices_one_hot.txt | 6 + .../bin/tox21/AtomCharge/indices_one_hot.txt | 6 + .../AtomFunctionalGroup/indices_one_hot.txt | 115 +++++++++++++ .../AtomHybridization/indices_one_hot.txt | 7 + .../tox21/AtomNodeLevel/indices_one_hot.txt | 3 + .../bin/tox21/AtomNumHs/indices_one_hot.txt | 6 + .../bin/tox21/AtomType/indices_one_hot.txt | 50 ++++++ .../bin/tox21/BondLevel/indices_one_hot.txt | 4 + .../bin/tox21/BondType/indices_one_hot.txt | 4 + .../tox21/NumAtomBonds/indices_one_hot.txt | 7 + .../toxcast/AtomCharge/indices_one_hot.txt | 6 + .../AtomFunctionalGroup/indices_one_hot.txt | 117 +++++++++++++ .../AtomHybridization/indices_one_hot.txt | 7 + .../toxcast/AtomNodeLevel/indices_one_hot.txt | 3 + .../bin/toxcast/AtomNumHs/indices_one_hot.txt | 5 + .../bin/toxcast/AtomType/indices_one_hot.txt | 53 ++++++ .../bin/toxcast/BondLevel/indices_one_hot.txt | 4 + .../bin/toxcast/BondType/indices_one_hot.txt | 4 + .../toxcast/NumAtomBonds/indices_one_hot.txt | 7 + 90 files changed, 1784 insertions(+) create mode 100644 chebai_graph/preprocessing/bin/_shared/AtomCharge/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/_shared/AtomFunctionalGroup/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/_shared/AtomHybridization/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/_shared/AtomNodeLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/_shared/AtomNumHs/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/_shared/AtomType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/_shared/BondLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/_shared/BondType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/_shared/NumAtomBonds/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bace/AtomCharge/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bace/AtomFunctionalGroup/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bace/AtomHybridization/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bace/AtomNodeLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bace/AtomNumHs/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bace/AtomType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bace/BondLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bace/BondType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bace/NumAtomBonds/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bbbp/AtomCharge/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bbbp/AtomFunctionalGroup/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bbbp/AtomHybridization/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bbbp/AtomNodeLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bbbp/AtomNumHs/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bbbp/AtomType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bbbp/BondLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bbbp/BondType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/bbbp/NumAtomBonds/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/clin_tox/AtomCharge/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/clin_tox/AtomFunctionalGroup/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/clin_tox/AtomHybridization/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/clin_tox/AtomNodeLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/clin_tox/AtomNumHs/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/clin_tox/AtomType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/clin_tox/BondLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/clin_tox/BondType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/clin_tox/NumAtomBonds/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/hiv/AtomCharge/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/hiv/AtomFunctionalGroup/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/hiv/AtomHybridization/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/hiv/AtomNodeLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/hiv/AtomNumHs/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/hiv/AtomType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/hiv/BondLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/hiv/BondType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/hiv/NumAtomBonds/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/muv/AtomCharge/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/muv/AtomFunctionalGroup/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/muv/AtomHybridization/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/muv/AtomNodeLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/muv/AtomNumHs/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/muv/AtomType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/muv/BondLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/muv/BondType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/muv/NumAtomBonds/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/pcba/AtomCharge/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/pcba/AtomFunctionalGroup/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/pcba/AtomHybridization/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/pcba/AtomNodeLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/pcba/AtomNumHs/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/pcba/AtomType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/pcba/BondLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/pcba/BondType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/pcba/NumAtomBonds/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/sider/AtomCharge/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/sider/AtomFunctionalGroup/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/sider/AtomHybridization/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/sider/AtomNodeLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/sider/AtomNumHs/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/sider/AtomType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/sider/BondLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/sider/BondType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/sider/NumAtomBonds/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/tox21/AtomCharge/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/tox21/AtomFunctionalGroup/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/tox21/AtomHybridization/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/tox21/AtomNodeLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/tox21/AtomNumHs/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/tox21/AtomType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/tox21/BondLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/tox21/BondType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/tox21/NumAtomBonds/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/toxcast/AtomCharge/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/toxcast/AtomFunctionalGroup/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/toxcast/AtomHybridization/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/toxcast/AtomNodeLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/toxcast/AtomNumHs/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/toxcast/AtomType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/toxcast/BondLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/toxcast/BondType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/toxcast/NumAtomBonds/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/_shared/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/_shared/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..b0f2214 --- /dev/null +++ b/chebai_graph/preprocessing/bin/_shared/AtomCharge/indices_one_hot.txt @@ -0,0 +1,13 @@ +-1 +1 +0 +2 +3 +-4 +-2 +4 +-3 +5 +6 +7 +-5 diff --git a/chebai_graph/preprocessing/bin/_shared/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/_shared/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..273cc9b --- /dev/null +++ b/chebai_graph/preprocessing/bin/_shared/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,160 @@ +NO_FG +graph_fg +alkyl +4_ammonium_ion +fluoro +chloro +ether +RING_6 +nitrile +primary_amine +secondary_amine +tertiary_carbon +RING_5 +alkene +RING_7 +tertiary_amine +ketone +hydroxyl +sulfide +sulfonyl +amide +ester +carboxyl +quaternary_carbon +secondary_ketimine +alkyne +alkene_carbon +nitro +trifluoromethyl +RING_4 +secondary_aldimine +primary_ketimine +acetal +carbonate_ester +azo +carbamate +bromo +RING_3 +phosphono +RING_40 +imide +RING_19 +nitroso +RING_8 +sulfhydryl +aldehyde +ketoxime +borono +difluoromethyl +thiolester +aldoxime +carboxylate +thiocyanate +isothiocyanate +phosphate +phosphodiester +iodo +sulfonic_acid +carbodithio +sulfinyl +azide +thioketone +RING_23 +trimethylsilyl +RING_21 +sulfonate_ester +RING_16 +RING_15 +RING_9 +RING_18 +RING_20 +RING_32 +RING_14 +RING_22 +nitrate +primary_aldimine +disulfide +RING_10 +hemiacetal +dichloromethyl +trichloromethyl +haloformyl +hydroperoxy +hemiketal +isonitrile +RING_24 +RING_11 +RING_12 +nitrosooxy +RING_33 +phosphino +sulfino +RING_17 +cyanate +RING_13 +carbothioic_O-acid +RING_36 +RING_72 +hydrazone +RING_31 +RING_25 +RING_28 +RING_35 +RING_34 +RING_42 +RING_38 +RING_29 +RING_26 +RING_51 +RING_75 +RING_27 +RING_45 +RING_60 +RING_30 +RING_43 +RING_41 +RING_56 +RING_81 +RING_39 +RING_37 +RING_47 +RING_49 +RING_44 +phosphoryl +RING_54 +RING_53 +RING_63 +carboxylic_anhydride +carbothioic_S-acid +silyl_ether +borinate +carbodithioic_acid +peroxy +tribromomethyl +chlorobromomethyl +amidine +dibromomethyl +isocyanate +bromodichloromethyl +ketal +boronate +borino +chloroiodomethyl +bromoiodomethyl +thionoester +fluorochloromethyl +RING_68 +bromodifluoromethyl +RING_58 +difluorochloromethyl +thial +RING_52 +RING_48 +RING_62 +RING_164 +RING_71 +RING_46 +orthoester +RING_55 +triiodomethyl diff --git a/chebai_graph/preprocessing/bin/_shared/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/_shared/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..dfc8e85 --- /dev/null +++ b/chebai_graph/preprocessing/bin/_shared/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,7 @@ +UNSPECIFIED +SP3 +S +SP3D2 +SP3D +SP +SP2 diff --git a/chebai_graph/preprocessing/bin/_shared/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/_shared/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/_shared/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/_shared/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/_shared/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..5ad8538 --- /dev/null +++ b/chebai_graph/preprocessing/bin/_shared/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,8 @@ +0 +3 +2 +4 +1 +5 +6 +8 diff --git a/chebai_graph/preprocessing/bin/_shared/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/_shared/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..a4766ec --- /dev/null +++ b/chebai_graph/preprocessing/bin/_shared/AtomType/indices_one_hot.txt @@ -0,0 +1,119 @@ +0 +13 +9 +11 +7 +47 +6 +4 +19 +50 +51 +17 +12 +8 +20 +27 +38 +81 +3 +30 +29 +39 +55 +25 +56 +70 +88 +40 +37 +26 +82 +78 +80 +16 +33 +1 +35 +15 +5 +44 +53 +85 +79 +24 +14 +87 +32 +2 +48 +63 +76 +92 +23 +64 +57 +34 +22 +28 +74 +42 +46 +52 +83 +62 +49 +58 +71 +65 +67 +77 +31 +59 +75 +45 +36 +54 +18 +10 +86 +43 +84 +118 +90 +73 +113 +114 +115 +116 +117 +112 +72 +104 +41 +105 +106 +107 +108 +109 +110 +111 +21 +103 +60 +61 +66 +68 +69 +89 +91 +93 +94 +95 +96 +97 +98 +99 +100 +101 +102 diff --git a/chebai_graph/preprocessing/bin/_shared/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/_shared/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..b389485 --- /dev/null +++ b/chebai_graph/preprocessing/bin/_shared/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +atom_fg_lvl +to_graphNode_lvl +within_atoms_lvl +within_fg_lvl diff --git a/chebai_graph/preprocessing/bin/_shared/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/_shared/BondType/indices_one_hot.txt new file mode 100644 index 0000000..f36bdf7 --- /dev/null +++ b/chebai_graph/preprocessing/bin/_shared/BondType/indices_one_hot.txt @@ -0,0 +1,6 @@ +DATIVE +SINGLE +AROMATIC +TRIPLE +DOUBLE +UNSPECIFIED diff --git a/chebai_graph/preprocessing/bin/_shared/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/_shared/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..22da8da --- /dev/null +++ b/chebai_graph/preprocessing/bin/_shared/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,13 @@ +0 +1 +2 +4 +5 +3 +6 +8 +7 +10 +12 +11 +9 diff --git a/chebai_graph/preprocessing/bin/bace/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..d7e52d1 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/AtomCharge/indices_one_hot.txt @@ -0,0 +1,3 @@ +0 +1 +-1 diff --git a/chebai_graph/preprocessing/bin/bace/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..b621c2a --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,44 @@ +RING_6 +difluoromethyl +RING_5 +ether +alkyne +alkyl +primary_amine +fluoro +amide +graph_fg +NO_FG +alkene +tertiary_carbon +alkene_carbon +nitrile +trifluoromethyl +ketone +chloro +quaternary_carbon +sulfonyl +hydroxyl +bromo +secondary_amine +ester +aldoxime +sulfide +iodo +tertiary_amine +RING_4 +secondary_ketimine +ketoxime +RING_3 +carboxylate +acetal +sulfinyl +RING_7 +RING_16 +RING_15 +nitro +primary_ketimine +sulfonate_ester +RING_8 +RING_14 +RING_13 diff --git a/chebai_graph/preprocessing/bin/bace/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..7c99da3 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,3 @@ +SP2 +SP3 +SP diff --git a/chebai_graph/preprocessing/bin/bace/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/bace/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..c4b216a --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,4 @@ +0 +1 +3 +2 diff --git a/chebai_graph/preprocessing/bin/bace/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..6da4980 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/AtomType/indices_one_hot.txt @@ -0,0 +1,8 @@ +6 +7 +8 +9 +17 +16 +35 +53 diff --git a/chebai_graph/preprocessing/bin/bace/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/bace/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/BondType/indices_one_hot.txt new file mode 100644 index 0000000..f1a2ce5 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/BondType/indices_one_hot.txt @@ -0,0 +1,4 @@ +SINGLE +AROMATIC +TRIPLE +DOUBLE diff --git a/chebai_graph/preprocessing/bin/bace/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bace/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..cf286d6 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bace/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,4 @@ +3 +2 +4 +1 diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..7b6a00a --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/AtomCharge/indices_one_hot.txt @@ -0,0 +1,4 @@ +0 +1 +-1 +2 diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..d2ee211 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,80 @@ +alkyl +ester +chloro +RING_6 +tertiary_amine +graph_fg +tertiary_carbon +primary_amine +dichloromethyl +sulfonyl +ketone +secondary_amine +hydroxyl +ether +quaternary_carbon +amide +NO_FG +carboxyl +nitro +alkene +alkene_carbon +secondary_ketimine +secondary_aldimine +carboxylate +iodo +aldoxime +bromo +alkyne +fluoro +carbamate +trifluoromethyl +sulfide +trichloromethyl +hemiacetal +primary_ketimine +azide +4_ammonium_ion +sulfonate_ester +ketoxime +nitroso +sulfhydryl +difluoromethyl +fluorochloromethyl +chlorobromomethyl +phosphono +carbodithio +disulfide +difluorochloromethyl +carbonate_ester +phosphate +RING_5 +thiolester +RING_7 +nitrile +phosphoryl +aldehyde +sulfinyl +thioketone +acetal +RING_4 +sulfonic_acid +imide +RING_14 +RING_24 +amidine +RING_3 +RING_9 +RING_19 +RING_16 +RING_8 +isonitrile +RING_23 +RING_12 +RING_21 +RING_17 +RING_30 +RING_33 +RING_36 +azo +RING_15 diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..0e4a9aa --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,4 @@ +SP3 +SP2 +S +SP diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..35f5032 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,4 @@ +2 +0 +1 +3 diff --git a/chebai_graph/preprocessing/bin/bbbp/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..9500570 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/AtomType/indices_one_hot.txt @@ -0,0 +1,13 @@ +6 +17 +8 +7 +16 +11 +53 +35 +9 +1 +15 +5 +20 diff --git a/chebai_graph/preprocessing/bin/bbbp/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/bbbp/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/BondType/indices_one_hot.txt new file mode 100644 index 0000000..a134437 --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/BondType/indices_one_hot.txt @@ -0,0 +1,4 @@ +DOUBLE +SINGLE +AROMATIC +TRIPLE diff --git a/chebai_graph/preprocessing/bin/bbbp/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/bbbp/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..5d07eee --- /dev/null +++ b/chebai_graph/preprocessing/bin/bbbp/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,5 @@ +2 +4 +1 +3 +0 diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..374de68 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/AtomCharge/indices_one_hot.txt @@ -0,0 +1,6 @@ +0 +-1 +1 +3 +2 +-2 diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..36bc343 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,84 @@ +alkyl +chloro +RING_6 +ether +secondary_amine +fluoro +graph_fg +tertiary_carbon +hydroxyl +ketone +sulfonyl +carboxylate +RING_5 +NO_FG +amide +sulfide +RING_4 +quaternary_carbon +alkyne +alkene +primary_amine +carboxyl +trifluoromethyl +ester +alkene_carbon +RING_20 +sulfhydryl +nitrile +RING_3 +RING_29 +RING_7 +4_ammonium_ion +acetal +RING_9 +iodo +RING_38 +sulfinyl +difluoromethyl +imide +RING_16 +tertiary_amine +nitro +RING_8 +secondary_ketimine +nitrate +sulfonic_acid +carbamate +secondary_aldimine +bromo +thioketone +RING_14 +azo +RING_33 +trichloromethyl +RING_15 +ketoxime +phosphodiester +RING_24 +RING_22 +aldehyde +RING_18 +RING_21 +nitroso +dichloromethyl +carbonate_ester +silyl_ether +thiolester +borono +aldoxime +sulfonate_ester +RING_23 +RING_19 +RING_25 +carbodithio +disulfide +RING_36 +RING_31 +fluorochloromethyl +RING_12 +phosphate +azide +chlorobromomethyl +bromodifluoromethyl +phosphino diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..a96b9b3 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,7 @@ +SP3 +SP2 +SP +S +SP3D2 +UNSPECIFIED +SP3D diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..1127527 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,4 @@ +3 +0 +2 +1 diff --git a/chebai_graph/preprocessing/bin/clin_tox/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..2c8a324 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/AtomType/indices_one_hot.txt @@ -0,0 +1,30 @@ +6 +17 +7 +8 +9 +16 +15 +53 +80 +1 +13 +35 +79 +0 +27 +22 +78 +33 +34 +30 +14 +5 +29 +24 +81 +43 +83 +25 +20 +26 diff --git a/chebai_graph/preprocessing/bin/clin_tox/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/clin_tox/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/BondType/indices_one_hot.txt new file mode 100644 index 0000000..bdc90d7 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/BondType/indices_one_hot.txt @@ -0,0 +1,5 @@ +SINGLE +AROMATIC +DOUBLE +TRIPLE +DATIVE diff --git a/chebai_graph/preprocessing/bin/clin_tox/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/clin_tox/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..3c2db51 --- /dev/null +++ b/chebai_graph/preprocessing/bin/clin_tox/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,6 @@ +1 +3 +2 +4 +0 +6 diff --git a/chebai_graph/preprocessing/bin/hiv/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..22aacd7 --- /dev/null +++ b/chebai_graph/preprocessing/bin/hiv/AtomCharge/indices_one_hot.txt @@ -0,0 +1,9 @@ +0 +1 +-1 +3 +-4 +-3 +-2 +2 +-5 diff --git a/chebai_graph/preprocessing/bin/hiv/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..e5b33d4 --- /dev/null +++ b/chebai_graph/preprocessing/bin/hiv/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,122 @@ +alkyl +RING_6 +chloro +NO_FG +secondary_amine +graph_fg +carboxyl +primary_amine +sulfonic_acid +nitro +ketone +hydroxyl +sulfonyl +ester +nitroso +sulfhydryl +tertiary_carbon +tertiary_amine +dichloromethyl +nitrile +primary_ketimine +alkyne +alkene +ether +iodo +sulfide +secondary_aldimine +phosphoryl +thioketone +alkene_carbon +secondary_ketimine +nitrate +4_ammonium_ion +ketoxime +amide +amidine +trichloromethyl +fluoro +quaternary_carbon +hydrazone +disulfide +thiocyanate +carbamate +trifluoromethyl +phosphate +sulfonate_ester +sulfinyl +azo +bromo +imide +carbodithioic_acid +carbodithio +sulfino +silyl_ether +aldehyde +phosphono +azide +aldoxime +phosphino +trimethylsilyl +thionoester +thiolester +isonitrile +acetal +difluorochloromethyl +dibromomethyl +ketal +diiodomethyl +borono +isothiocyanate +difluoromethyl +tribromomethyl +carbonate_ester +hemiacetal +phosphodiester +carboxylate +orthoester +isocyanate +primary_aldimine +hemiketal +borinate +RING_5 +cyanate +RING_4 +RING_24 +RING_3 +thial +RING_7 +RING_12 +RING_9 +RING_16 +hydroperoxy +RING_15 +RING_18 +RING_8 +RING_35 +RING_10 +RING_19 +RING_17 +RING_29 +RING_23 +RING_27 +RING_11 +RING_14 +RING_20 +RING_33 +RING_13 +RING_31 +RING_21 +RING_28 +RING_22 +RING_30 +RING_34 +RING_54 +RING_38 +RING_36 +RING_26 +RING_32 +carboxylic_anhydride +RING_25 +carbothioic_S-acid +RING_40 diff --git a/chebai_graph/preprocessing/bin/hiv/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..cdc6985 --- /dev/null +++ b/chebai_graph/preprocessing/bin/hiv/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,7 @@ +SP3 +SP2 +SP +S +SP3D2 +SP3D +UNSPECIFIED diff --git a/chebai_graph/preprocessing/bin/hiv/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/hiv/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/hiv/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..df4a3e9 --- /dev/null +++ b/chebai_graph/preprocessing/bin/hiv/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,5 @@ +3 +1 +0 +2 +4 diff --git a/chebai_graph/preprocessing/bin/hiv/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..95c9cba --- /dev/null +++ b/chebai_graph/preprocessing/bin/hiv/AtomType/indices_one_hot.txt @@ -0,0 +1,54 @@ +6 +17 +8 +7 +15 +16 +33 +53 +9 +35 +14 +11 +34 +79 +50 +80 +32 +51 +82 +5 +19 +3 +83 +29 +67 +44 +65 +74 +27 +55 +25 +26 +75 +42 +78 +47 +24 +46 +89 +20 +31 +12 +30 +28 +77 +1 +40 +81 +45 +22 +52 +64 +92 +23 diff --git a/chebai_graph/preprocessing/bin/hiv/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/hiv/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/hiv/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/BondType/indices_one_hot.txt new file mode 100644 index 0000000..58c0a5c --- /dev/null +++ b/chebai_graph/preprocessing/bin/hiv/BondType/indices_one_hot.txt @@ -0,0 +1,5 @@ +SINGLE +DOUBLE +AROMATIC +TRIPLE +DATIVE diff --git a/chebai_graph/preprocessing/bin/hiv/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/hiv/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..88ddb48 --- /dev/null +++ b/chebai_graph/preprocessing/bin/hiv/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,11 @@ +1 +2 +3 +4 +0 +5 +6 +9 +7 +8 +10 diff --git a/chebai_graph/preprocessing/bin/muv/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..d7e52d1 --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/AtomCharge/indices_one_hot.txt @@ -0,0 +1,3 @@ +0 +1 +-1 diff --git a/chebai_graph/preprocessing/bin/muv/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..c55dc92 --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,68 @@ +primary_amine +RING_6 +ketone +carboxyl +alkyl +tertiary_carbon +secondary_amine +graph_fg +hydroxyl +alkene +thioketone +bromo +fluoro +ester +ether +tertiary_amine +NO_FG +chloro +sulfonyl +difluoromethyl +quaternary_carbon +nitro +amide +alkyne +trifluoromethyl +aldoxime +alkene_carbon +nitrile +sulfide +secondary_aldimine +ketoxime +carbodithio +secondary_ketimine +thionoester +thiolester +sulfinyl +carbamate +thiocyanate +aldehyde +sulfonate_ester +amidine +trichloromethyl +sulfhydryl +sulfonic_acid +RING_5 +nitroso +acetal +carbonate_ester +imide +azo +primary_ketimine +dichloromethyl +RING_7 +RING_3 +RING_8 +hydrazone +RING_4 +RING_15 +RING_12 +azide +RING_13 +RING_18 +difluorochloromethyl +RING_14 +RING_9 +RING_11 +disulfide +4_ammonium_ion diff --git a/chebai_graph/preprocessing/bin/muv/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..7c99da3 --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,3 @@ +SP2 +SP3 +SP diff --git a/chebai_graph/preprocessing/bin/muv/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/muv/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..52db5ae --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,4 @@ +2 +1 +0 +3 diff --git a/chebai_graph/preprocessing/bin/muv/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..c36ddc7 --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/AtomType/indices_one_hot.txt @@ -0,0 +1,7 @@ +7 +6 +8 +35 +9 +16 +17 diff --git a/chebai_graph/preprocessing/bin/muv/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/muv/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/BondType/indices_one_hot.txt new file mode 100644 index 0000000..f14a241 --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/BondType/indices_one_hot.txt @@ -0,0 +1,4 @@ +SINGLE +DOUBLE +AROMATIC +TRIPLE diff --git a/chebai_graph/preprocessing/bin/muv/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/muv/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..94ebaf9 --- /dev/null +++ b/chebai_graph/preprocessing/bin/muv/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,4 @@ +1 +2 +3 +4 diff --git a/chebai_graph/preprocessing/bin/pcba/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..7897486 --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/AtomCharge/indices_one_hot.txt @@ -0,0 +1,8 @@ +0 +1 +-1 +3 +2 +4 +5 +6 diff --git a/chebai_graph/preprocessing/bin/pcba/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..89e00f3 --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,127 @@ +alkyl +RING_6 +sulfide +RING_5 +ketone +trifluoromethyl +secondary_amine +graph_fg +NO_FG +amide +sulfonyl +ether +quaternary_carbon +tertiary_carbon +imide +carboxyl +nitrile +ester +primary_amine +alkyne +fluoro +tertiary_amine +nitro +RING_7 +chloro +hydroxyl +RING_3 +alkene +secondary_aldimine +bromo +thioketone +sulfinyl +RING_4 +secondary_ketimine +RING_12 +primary_ketimine +alkene_carbon +carbamate +aldehyde +phosphoryl +trimethylsilyl +carboxylate +ketoxime +RING_23 +difluoromethyl +sulfonate_ester +RING_14 +iodo +thiocyanate +RING_8 +disulfide +carbonate_ester +RING_19 +sulfhydryl +nitroso +acetal +aldoxime +borono +RING_16 +phosphono +RING_9 +RING_15 +hydrazone +RING_11 +4_ammonium_ion +thiolester +trichloromethyl +carbodithio +azo +azide +sulfonic_acid +dichloromethyl +phosphodiester +borinate +RING_13 +RING_18 +thial +RING_21 +RING_10 +isocyanate +isonitrile +RING_22 +RING_36 +RING_20 +phosphate +amidine +dibromomethyl +RING_17 +peroxy +RING_40 +silyl_ether +thionoester +nitrate +haloformyl +hemiacetal +RING_25 +isothiocyanate +difluorochloromethyl +RING_24 +nitrosooxy +RING_32 +carbodithioic_acid +hydroperoxy +sulfino +RING_26 +tribromomethyl +RING_30 +RING_29 +diiodomethyl +phosphino +chlorobromomethyl +RING_33 +RING_38 +RING_34 +bromodifluoromethyl +carboxylic_anhydride +bromodichloromethyl +RING_27 +dichlorosilane +RING_42 +ketal +RING_28 +RING_56 +RING_31 +RING_35 +fluorochloromethyl +RING_54 diff --git a/chebai_graph/preprocessing/bin/pcba/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..fb3232b --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,7 @@ +SP3 +SP2 +SP +SP3D +S +UNSPECIFIED +SP3D2 diff --git a/chebai_graph/preprocessing/bin/pcba/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/pcba/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..cad2399 --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,6 @@ +3 +1 +0 +2 +6 +4 diff --git a/chebai_graph/preprocessing/bin/pcba/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..166ffe2 --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/AtomType/indices_one_hot.txt @@ -0,0 +1,44 @@ +6 +16 +7 +8 +9 +17 +35 +15 +14 +53 +5 +29 +79 +38 +50 +30 +33 +34 +48 +51 +1 +26 +82 +27 +78 +80 +3 +28 +22 +56 +44 +11 +13 +46 +25 +83 +47 +60 +24 +64 +32 +49 +74 +31 diff --git a/chebai_graph/preprocessing/bin/pcba/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/pcba/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/BondType/indices_one_hot.txt new file mode 100644 index 0000000..c819d38 --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/BondType/indices_one_hot.txt @@ -0,0 +1,4 @@ +SINGLE +AROMATIC +DOUBLE +TRIPLE diff --git a/chebai_graph/preprocessing/bin/pcba/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/pcba/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..2efbda4 --- /dev/null +++ b/chebai_graph/preprocessing/bin/pcba/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,6 @@ +1 +2 +3 +4 +0 +5 diff --git a/chebai_graph/preprocessing/bin/sider/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..6f89a5d --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/AtomCharge/indices_one_hot.txt @@ -0,0 +1,6 @@ +0 +1 +-1 +3 +2 +-2 diff --git a/chebai_graph/preprocessing/bin/sider/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..6fbc2c2 --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,93 @@ +alkyl +RING_6 +carboxyl +alkene +fluoro +sulfinyl +RING_5 +graph_fg +tertiary_carbon +secondary_amine +ketone +primary_amine +primary_ketimine +NO_FG +sulfide +hydroxyl +sulfonic_acid +bromo +4_ammonium_ion +chloro +quaternary_carbon +ether +carboxylate +tertiary_amine +amide +RING_3 +trifluoromethyl +ester +nitro +sulfonyl +iodo +secondary_ketimine +RING_7 +aldehyde +phosphono +alkyne +nitrile +carbamate +alkene_carbon +carbonate_ester +RING_4 +RING_18 +RING_24 +RING_14 +imide +sulfonate_ester +RING_81 +RING_20 +RING_8 +nitroso +RING_29 +RING_16 +RING_12 +RING_21 +sulfhydryl +thiolester +difluoromethyl +thioketone +RING_25 +RING_78 +RING_68 +RING_77 +nitrate +azo +RING_36 +phosphate +ketoxime +secondary_aldimine +dichloromethyl +carboxylic_anhydride +RING_33 +borono +RING_15 +diiodomethyl +RING_23 +carbodithio +disulfide +acetal +RING_40 +phosphodiester +fluorochloromethyl +RING_22 +RING_10 +RING_9 +azide +phosphino +RING_31 +RING_28 +RING_53 +RING_56 +RING_41 +RING_43 +RING_35 diff --git a/chebai_graph/preprocessing/bin/sider/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..e9958eb --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,6 @@ +SP3 +SP2 +S +SP3D +SP +SP3D2 diff --git a/chebai_graph/preprocessing/bin/sider/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/sider/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..df4a3e9 --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,5 @@ +3 +1 +0 +2 +4 diff --git a/chebai_graph/preprocessing/bin/sider/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..384eb37 --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/AtomType/indices_one_hot.txt @@ -0,0 +1,40 @@ +6 +8 +9 +16 +7 +35 +17 +11 +64 +15 +24 +53 +47 +49 +1 +25 +29 +20 +62 +3 +19 +78 +39 +43 +31 +26 +5 +12 +38 +79 +33 +30 +56 +98 +34 +32 +27 +88 +81 +57 diff --git a/chebai_graph/preprocessing/bin/sider/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/sider/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/BondType/indices_one_hot.txt new file mode 100644 index 0000000..f14a241 --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/BondType/indices_one_hot.txt @@ -0,0 +1,4 @@ +SINGLE +DOUBLE +AROMATIC +TRIPLE diff --git a/chebai_graph/preprocessing/bin/sider/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/sider/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..dfb5c6d --- /dev/null +++ b/chebai_graph/preprocessing/bin/sider/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,6 @@ +1 +2 +3 +4 +0 +6 diff --git a/chebai_graph/preprocessing/bin/tox21/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..a8bd8c3 --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/AtomCharge/indices_one_hot.txt @@ -0,0 +1,6 @@ +0 +1 +-1 +2 +3 +-2 diff --git a/chebai_graph/preprocessing/bin/tox21/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..5aa0016 --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,115 @@ +alkyl +hemiacetal +ether +graph_fg +RING_18 +NO_FG +RING_6 +ester +tertiary_carbon +secondary_amine +ketone +carboxyl +RING_5 +amide +primary_amine +RING_4 +carboxylate +alkene +nitro +imide +4_ammonium_ion +bromo +chloro +hydroxyl +sulfonyl +fluoro +quaternary_carbon +alkene_carbon +RING_3 +alkyne +secondary_ketimine +nitrile +isocyanate +azo +tertiary_amine +trifluoromethyl +aldehyde +thiolester +sulfide +nitroso +phosphate +phosphoryl +RING_7 +iodo +acetal +carbodithio +carbonate_ester +RING_21 +difluoromethyl +disulfide +RING_25 +nitrate +carbamate +secondary_aldimine +sulfonic_acid +haloformyl +amidine +sulfhydryl +phosphono +dibromomethyl +RING_23 +trichloromethyl +sulfonate_ester +RING_40 +RING_24 +RING_15 +primary_ketimine +RING_29 +thioketone +aldoxime +RING_16 +nitrosooxy +sulfino +RING_36 +peroxy +RING_14 +ketoxime +thiocyanate +RING_8 +RING_17 +sulfinyl +dichloromethyl +primary_aldimine +tribromomethyl +RING_19 +RING_9 +silyl_ether +phosphodiester +RING_33 +azide +RING_27 +RING_26 +RING_20 +phosphino +RING_12 +trimethylsilyl +carboxylic_anhydride +orthoester +RING_10 +thionoester +isothiocyanate +borono +chlorobromomethyl +dichlorosilane +hydroperoxy +RING_35 +RING_30 +RING_28 +fluorochloromethyl +RING_13 +bromodichloromethyl +ketal +difluorochloromethyl +bromodifluoromethyl +RING_11 diff --git a/chebai_graph/preprocessing/bin/tox21/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..cdc6985 --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,7 @@ +SP3 +SP2 +SP +S +SP3D2 +SP3D +UNSPECIFIED diff --git a/chebai_graph/preprocessing/bin/tox21/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/tox21/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..5a694e8 --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,6 @@ +3 +2 +1 +0 +4 +6 diff --git a/chebai_graph/preprocessing/bin/tox21/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..64e481b --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/AtomType/indices_one_hot.txt @@ -0,0 +1,50 @@ +6 +8 +7 +16 +15 +35 +17 +9 +56 +30 +26 +80 +53 +28 +5 +51 +14 +47 +11 +78 +40 +79 +4 +29 +20 +34 +81 +50 +82 +23 +32 +38 +48 +46 +60 +27 +83 +64 +25 +19 +33 +22 +3 +12 +42 +66 +24 +1 +49 +70 diff --git a/chebai_graph/preprocessing/bin/tox21/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/tox21/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/BondType/indices_one_hot.txt new file mode 100644 index 0000000..c819d38 --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/BondType/indices_one_hot.txt @@ -0,0 +1,4 @@ +SINGLE +AROMATIC +DOUBLE +TRIPLE diff --git a/chebai_graph/preprocessing/bin/tox21/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/tox21/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..21b19a8 --- /dev/null +++ b/chebai_graph/preprocessing/bin/tox21/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,7 @@ +1 +2 +3 +4 +0 +5 +6 diff --git a/chebai_graph/preprocessing/bin/toxcast/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..c857ab5 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/AtomCharge/indices_one_hot.txt @@ -0,0 +1,6 @@ +0 +-1 +1 +2 +3 +-2 diff --git a/chebai_graph/preprocessing/bin/toxcast/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..500ebf5 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,117 @@ +ester +RING_6 +sulfonyl +alkyl +NO_FG +4_ammonium_ion +graph_fg +sulfide +primary_amine +carboxyl +chloro +quaternary_carbon +hydroxyl +secondary_amine +fluoro +ketone +tertiary_carbon +ether +alkene +nitro +tertiary_amine +aldehyde +RING_5 +alkyne +nitrile +RING_7 +alkene_carbon +carbonate_ester +trifluoromethyl +imide +sulfonic_acid +primary_ketimine +disulfide +sulfonate_ester +amide +nitroso +iodo +RING_3 +carboxylate +RING_8 +RING_20 +acetal +bromo +carbamate +RING_12 +ketoxime +RING_18 +silyl_ether +RING_4 +dichloromethyl +sulfinyl +aldoxime +thioketone +secondary_ketimine +isothiocyanate +isocyanate +RING_14 +sulfhydryl +phosphono +carbodithioic_acid +secondary_aldimine +haloformyl +phosphoryl +azo +nitrate +RING_35 +thiocyanate +hemiacetal +thiolester +difluorochloromethyl +dibromomethyl +difluoromethyl +phosphino +RING_29 +chlorobromomethyl +RING_9 +tribromomethyl +trichloromethyl +RING_28 +RING_16 +nitrosooxy +bromodichloromethyl +carbodithio +phosphodiester +RING_23 +phosphate +carboxylic_anhydride +RING_13 +RING_24 +RING_30 +RING_11 +RING_17 +azide +RING_40 +primary_aldimine +RING_15 +peroxy +RING_33 +trimethylsilyl +sulfino +ketal +RING_10 +thionoester +borono +amidine +orthoester +RING_27 +RING_26 +RING_38 +hydroperoxy +RING_21 +RING_19 +dichlorosilane +diiodomethyl +fluorochloromethyl +bromodifluoromethyl +RING_36 diff --git a/chebai_graph/preprocessing/bin/toxcast/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..5d039f5 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,7 @@ +SP2 +SP3 +SP +S +UNSPECIFIED +SP3D2 +SP3D diff --git a/chebai_graph/preprocessing/bin/toxcast/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/toxcast/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..ed1daa5 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,5 @@ +0 +1 +3 +2 +4 diff --git a/chebai_graph/preprocessing/bin/toxcast/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..9c59d72 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/AtomType/indices_one_hot.txt @@ -0,0 +1,53 @@ +8 +6 +7 +16 +17 +9 +50 +19 +22 +14 +11 +53 +35 +33 +49 +78 +15 +51 +12 +3 +83 +28 +26 +13 +20 +47 +29 +5 +1 +30 +38 +80 +27 +34 +79 +56 +81 +24 +23 +21 +64 +70 +60 +46 +32 +4 +82 +66 +25 +48 +63 +40 +42 diff --git a/chebai_graph/preprocessing/bin/toxcast/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/toxcast/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/BondType/indices_one_hot.txt new file mode 100644 index 0000000..f14a241 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/BondType/indices_one_hot.txt @@ -0,0 +1,4 @@ +SINGLE +DOUBLE +AROMATIC +TRIPLE diff --git a/chebai_graph/preprocessing/bin/toxcast/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/toxcast/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..a61fff5 --- /dev/null +++ b/chebai_graph/preprocessing/bin/toxcast/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,7 @@ +2 +1 +3 +4 +0 +6 +5 From 71b8427e1c3745c58550562bd1055c93c6249675 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Sun, 18 Oct 2026 07:45:42 +0200 Subject: [PATCH 37/58] zinc15 250K samples tokens --- .../bin/zinc15/AtomCharge/indices_one_hot.txt | 3 + .../AtomFunctionalGroup/indices_one_hot.txt | 72 +++++++++++++++++++ .../AtomHybridization/indices_one_hot.txt | 3 + .../zinc15/AtomNodeLevel/indices_one_hot.txt | 3 + .../bin/zinc15/AtomNumHs/indices_one_hot.txt | 4 ++ .../bin/zinc15/AtomType/indices_one_hot.txt | 11 +++ .../bin/zinc15/BondLevel/indices_one_hot.txt | 4 ++ .../bin/zinc15/BondType/indices_one_hot.txt | 4 ++ .../zinc15/NumAtomBonds/indices_one_hot.txt | 4 ++ .../final/chebi50_final_augmented.yml | 2 + 10 files changed, 110 insertions(+) create mode 100644 chebai_graph/preprocessing/bin/zinc15/AtomCharge/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/zinc15/AtomFunctionalGroup/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/zinc15/AtomHybridization/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/zinc15/AtomNodeLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/zinc15/AtomNumHs/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/zinc15/AtomType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/zinc15/BondLevel/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/zinc15/BondType/indices_one_hot.txt create mode 100644 chebai_graph/preprocessing/bin/zinc15/NumAtomBonds/indices_one_hot.txt diff --git a/chebai_graph/preprocessing/bin/zinc15/AtomCharge/indices_one_hot.txt b/chebai_graph/preprocessing/bin/zinc15/AtomCharge/indices_one_hot.txt new file mode 100644 index 0000000..89a8b18 --- /dev/null +++ b/chebai_graph/preprocessing/bin/zinc15/AtomCharge/indices_one_hot.txt @@ -0,0 +1,3 @@ +0 +-1 +1 diff --git a/chebai_graph/preprocessing/bin/zinc15/AtomFunctionalGroup/indices_one_hot.txt b/chebai_graph/preprocessing/bin/zinc15/AtomFunctionalGroup/indices_one_hot.txt new file mode 100644 index 0000000..bc1b07e --- /dev/null +++ b/chebai_graph/preprocessing/bin/zinc15/AtomFunctionalGroup/indices_one_hot.txt @@ -0,0 +1,72 @@ +alkyl +secondary_amine +trifluoromethyl +quaternary_carbon +amide +sulfide +graph_fg +RING_6 +RING_5 +ketone +RING_3 +tertiary_carbon +fluoro +ether +imide +primary_amine +NO_FG +RING_4 +alkene_carbon +alkene +carboxyl +hydroxyl +sulfonyl +ester +tertiary_amine +alkyne +nitrile +secondary_aldimine +chloro +RING_7 +nitro +difluoromethyl +thioketone +secondary_ketimine +bromo +sulfinyl +primary_ketimine +nitroso +phosphate +RING_8 +aldehyde +sulfonate_ester +trimethylsilyl +carbamate +sulfhydryl +acetal +sulfonic_acid +RING_10 +thiolester +phosphoryl +iodo +carbonate_ester +nitrate +4_ammonium_ion +azo +amidine +bromodifluoromethyl +aldoxime +RING_9 +ketal +ketoxime +difluorochloromethyl +trichloromethyl +hydrazone +phosphono +RING_12 +carbodithio +RING_11 +thionoester +azide +RING_16 +phosphino diff --git a/chebai_graph/preprocessing/bin/zinc15/AtomHybridization/indices_one_hot.txt b/chebai_graph/preprocessing/bin/zinc15/AtomHybridization/indices_one_hot.txt new file mode 100644 index 0000000..a4a9930 --- /dev/null +++ b/chebai_graph/preprocessing/bin/zinc15/AtomHybridization/indices_one_hot.txt @@ -0,0 +1,3 @@ +SP3 +SP2 +SP diff --git a/chebai_graph/preprocessing/bin/zinc15/AtomNodeLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/zinc15/AtomNodeLevel/indices_one_hot.txt new file mode 100644 index 0000000..2d65776 --- /dev/null +++ b/chebai_graph/preprocessing/bin/zinc15/AtomNodeLevel/indices_one_hot.txt @@ -0,0 +1,3 @@ +atom_node_lvl +fg_node_lvl +graph_node_level diff --git a/chebai_graph/preprocessing/bin/zinc15/AtomNumHs/indices_one_hot.txt b/chebai_graph/preprocessing/bin/zinc15/AtomNumHs/indices_one_hot.txt new file mode 100644 index 0000000..3ed60bd --- /dev/null +++ b/chebai_graph/preprocessing/bin/zinc15/AtomNumHs/indices_one_hot.txt @@ -0,0 +1,4 @@ +3 +1 +0 +2 diff --git a/chebai_graph/preprocessing/bin/zinc15/AtomType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/zinc15/AtomType/indices_one_hot.txt new file mode 100644 index 0000000..81eeb68 --- /dev/null +++ b/chebai_graph/preprocessing/bin/zinc15/AtomType/indices_one_hot.txt @@ -0,0 +1,11 @@ +6 +7 +9 +8 +16 +17 +35 +15 +14 +5 +53 diff --git a/chebai_graph/preprocessing/bin/zinc15/BondLevel/indices_one_hot.txt b/chebai_graph/preprocessing/bin/zinc15/BondLevel/indices_one_hot.txt new file mode 100644 index 0000000..6b48d66 --- /dev/null +++ b/chebai_graph/preprocessing/bin/zinc15/BondLevel/indices_one_hot.txt @@ -0,0 +1,4 @@ +within_atoms_lvl +atom_fg_lvl +within_fg_lvl +to_graphNode_lvl diff --git a/chebai_graph/preprocessing/bin/zinc15/BondType/indices_one_hot.txt b/chebai_graph/preprocessing/bin/zinc15/BondType/indices_one_hot.txt new file mode 100644 index 0000000..f14a241 --- /dev/null +++ b/chebai_graph/preprocessing/bin/zinc15/BondType/indices_one_hot.txt @@ -0,0 +1,4 @@ +SINGLE +DOUBLE +AROMATIC +TRIPLE diff --git a/chebai_graph/preprocessing/bin/zinc15/NumAtomBonds/indices_one_hot.txt b/chebai_graph/preprocessing/bin/zinc15/NumAtomBonds/indices_one_hot.txt new file mode 100644 index 0000000..207d1d2 --- /dev/null +++ b/chebai_graph/preprocessing/bin/zinc15/NumAtomBonds/indices_one_hot.txt @@ -0,0 +1,4 @@ +1 +2 +4 +3 diff --git a/configs/data/augmented/final/chebi50_final_augmented.yml b/configs/data/augmented/final/chebi50_final_augmented.yml index 3ff3232..1b711f0 100644 --- a/configs/data/augmented/final/chebi50_final_augmented.yml +++ b/configs/data/augmented/final/chebi50_final_augmented.yml @@ -1,5 +1,7 @@ class_path: chebai_graph.preprocessing.datasets.ChEBI50_GNWA_FGWA_WFGE_AsPerNodeType init_args: + splits_file_path: data/chebi_v252/ChEBI50/processed/splits.csv + chebi_version: 252 properties: # All Node type properties - chebai_graph.preprocessing.properties.AtomNodeLevel From aac73508ab7147a83092614555b69420c216141e Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sun, 18 Oct 2026 08:39:20 +0200 Subject: [PATCH 38/58] change data_type -> tokens_directory --- chebai_graph/loss/msp.py | 10 ++++- chebai_graph/preprocessing/datasets/base.py | 4 +- chebai_graph/preprocessing/datasets/utils.py | 10 ++--- .../properties/augmented_properties.py | 20 ++++++---- chebai_graph/preprocessing/properties/base.py | 10 +++-- .../preprocessing/properties/properties.py | 40 +++++++++++-------- .../preprocessing/property_encoder.py | 10 +++-- 7 files changed, 64 insertions(+), 40 deletions(-) diff --git a/chebai_graph/loss/msp.py b/chebai_graph/loss/msp.py index 8117497..b6e91ac 100644 --- a/chebai_graph/loss/msp.py +++ b/chebai_graph/loss/msp.py @@ -32,10 +32,16 @@ def __init__( ) -> None: super().__init__() self.atom_type_encoder = IndexEncoder( - property=AugAtomType(), use_one_hot_cache=True, eval=True + property=AugAtomType(), + use_one_hot_cache=True, + eval=True, + tokens_directory="_shared", ) self.bond_type_encoder = IndexEncoder( - property=AugBondType(), use_one_hot_cache=True, eval=True + property=AugBondType(), + use_one_hot_cache=True, + eval=True, + tokens_directory="_shared", ) # Bond existence projection diff --git a/chebai_graph/preprocessing/datasets/base.py b/chebai_graph/preprocessing/datasets/base.py index d94a333..76fedf2 100644 --- a/chebai_graph/preprocessing/datasets/base.py +++ b/chebai_graph/preprocessing/datasets/base.py @@ -50,7 +50,9 @@ def __init__( super().__init__(**kwargs) # atom_properties and bond_properties are given as lists containing class_paths if properties is not None: - properties = [resolve_property(prop, self.data_type) for prop in properties] + properties = [ + resolve_property(prop, self.tokens_directory) for prop in properties + ] properties = self._sort_properties(properties) else: properties = [] diff --git a/chebai_graph/preprocessing/datasets/utils.py b/chebai_graph/preprocessing/datasets/utils.py index 79f2fb3..fa42078 100644 --- a/chebai_graph/preprocessing/datasets/utils.py +++ b/chebai_graph/preprocessing/datasets/utils.py @@ -5,7 +5,7 @@ def resolve_property( - property: str | MolecularProperty, data_type: str + property: str | MolecularProperty, tokens_directory: str ) -> MolecularProperty: """ Resolves a molecular property specification (either as a class instance or class path string) @@ -21,8 +21,8 @@ def resolve_property( property (str | MolecularProperty): The property to resolve. Can be a class instance, a fully qualified class name (e.g. "module.ClassName"), or a class name assumed to be in `chebai_graph.preprocessing.properties`. - data_type (str): The data type associated with the property. This used to determine or set - tokens file path for the property if applicable. + tokens_directory (str): The directory where the tokens for the property are stored. This is used to determine or set + the tokens file path for the property if applicable. Returns: MolecularProperty: An instance of the resolved MolecularProperty. @@ -41,7 +41,7 @@ def resolve_property( module_name = property[:last_dot] class_name = property[last_dot + 1 :] module = importlib.import_module(module_name) - return getattr(module, class_name)(data_type=data_type) + return getattr(module, class_name)(tokens_directory=tokens_directory) except ValueError: # if only a class name is given, assume the module is chebai_graph.processing.properties - return getattr(graph_properties, property)(data_type=data_type) + return getattr(graph_properties, property)(tokens_directory=tokens_directory) diff --git a/chebai_graph/preprocessing/properties/augmented_properties.py b/chebai_graph/preprocessing/properties/augmented_properties.py index 6f112ea..07d514d 100644 --- a/chebai_graph/preprocessing/properties/augmented_properties.py +++ b/chebai_graph/preprocessing/properties/augmented_properties.py @@ -30,9 +30,10 @@ def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: Args: encoder (PropertyEncoder | None): Property encoder to use. Defaults to OneHotEncoder. """ - data_type = kwargs.get("data_type") + tokens_directory = kwargs.get("tokens_directory") super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + encoder=encoder or OneHotEncoder(self, tokens_directory=tokens_directory), + **kwargs, ) def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> str | int | bool: @@ -56,9 +57,10 @@ def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: Args: encoder (PropertyEncoder | None): Property encoder to use. Defaults to OneHotEncoder. """ - data_type = kwargs.get("data_type") + tokens_directory = kwargs.get("tokens_directory") super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + encoder=encoder or OneHotEncoder(self, tokens_directory=tokens_directory), + **kwargs, ) def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> str | int | bool: @@ -82,9 +84,10 @@ def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: Args: encoder (PropertyEncoder | None): Property encoder to use. Defaults to OneHotEncoder. """ - data_type = kwargs.get("data_type") + tokens_directory = kwargs.get("tokens_directory") super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + encoder=encoder or OneHotEncoder(self, tokens_directory=tokens_directory), + **kwargs, ) def get_atom_value(self, atom: Chem.rdchem.Atom | dict) -> int: @@ -337,9 +340,10 @@ def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: Args: encoder (PropertyEncoder | None): Optional encoder to use. Defaults to OneHotEncoder. """ - data_type = kwargs.get("data_type") + tokens_directory = kwargs.get("tokens_directory") super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + encoder=encoder or OneHotEncoder(self, tokens_directory=tokens_directory), + **kwargs, ) def get_bond_value(self, bond: Chem.rdchem.Bond | dict) -> str: diff --git a/chebai_graph/preprocessing/properties/base.py b/chebai_graph/preprocessing/properties/base.py index 1f2e0aa..bf70665 100644 --- a/chebai_graph/preprocessing/properties/base.py +++ b/chebai_graph/preprocessing/properties/base.py @@ -33,14 +33,16 @@ class MolecularProperty(ABC): def __init__( self, - data_type: str, + tokens_directory: str, encoder: PropertyEncoder | None = None, ) -> None: - assert data_type is not None, "data_type must be provided for MolecularProperty" + assert tokens_directory is not None, ( + "tokens_directory must be provided for MolecularProperty" + ) if encoder is None: - encoder = IndexEncoder(self, data_type=data_type) + encoder = IndexEncoder(self, tokens_directory=tokens_directory) self.encoder: PropertyEncoder = encoder - self._data_type = data_type + self.tokens_directory = tokens_directory @property def name(self) -> str: diff --git a/chebai_graph/preprocessing/properties/properties.py b/chebai_graph/preprocessing/properties/properties.py index 87112da..0e91d57 100644 --- a/chebai_graph/preprocessing/properties/properties.py +++ b/chebai_graph/preprocessing/properties/properties.py @@ -20,9 +20,10 @@ class AtomType(AtomProperty): """ def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") + tokens_directory = kwargs.get("tokens_directory") super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + encoder=encoder or OneHotEncoder(self, tokens_directory=tokens_directory), + **kwargs, ) def get_atom_value(self, atom: Chem.rdchem.Atom) -> int: @@ -46,9 +47,10 @@ class NumAtomBonds(AtomProperty): """ def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") + tokens_directory = kwargs.get("tokens_directory") super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + encoder=encoder or OneHotEncoder(self, tokens_directory=tokens_directory), + **kwargs, ) def get_atom_value(self, atom: Chem.rdchem.Atom) -> int: @@ -72,9 +74,10 @@ class AtomCharge(AtomProperty): """ def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") + tokens_directory = kwargs.get("tokens_directory") super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + encoder=encoder or OneHotEncoder(self, tokens_directory=tokens_directory), + **kwargs, ) def get_atom_value(self, atom: Chem.rdchem.Atom) -> int: @@ -98,9 +101,10 @@ class AtomChirality(AtomProperty): """ def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") + tokens_directory = kwargs.get("tokens_directory") super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + encoder=encoder or OneHotEncoder(self, tokens_directory=tokens_directory), + **kwargs, ) def get_atom_value(self, atom: Chem.rdchem.Atom) -> Chem.rdchem.ChiralType: @@ -124,9 +128,10 @@ class AtomHybridization(AtomProperty): """ def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") + tokens_directory = kwargs.get("tokens_directory") super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + encoder=encoder or OneHotEncoder(self, tokens_directory=tokens_directory), + **kwargs, ) def get_atom_value(self, atom: Chem.rdchem.Atom) -> Chem.rdchem.HybridizationType: @@ -150,9 +155,10 @@ class AtomNumHs(AtomProperty): """ def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") + tokens_directory = kwargs.get("tokens_directory") super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + encoder=encoder or OneHotEncoder(self, tokens_directory=tokens_directory), + **kwargs, ) def get_atom_value(self, atom: Chem.rdchem.Atom) -> int: @@ -222,9 +228,10 @@ class BondType(BondProperty): """ def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") + tokens_directory = kwargs.get("tokens_directory") super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + encoder=encoder or OneHotEncoder(self, tokens_directory=tokens_directory), + **kwargs, ) def get_bond_value(self, bond: Chem.rdchem.Bond) -> Chem.rdchem.BondType: @@ -271,9 +278,10 @@ class MoleculeNumRings(MoleculeProperty): """ def __init__(self, encoder: PropertyEncoder | None = None, **kwargs) -> None: - data_type = kwargs.get("data_type") + tokens_directory = kwargs.get("tokens_directory") super().__init__( - encoder=encoder or OneHotEncoder(self, data_type=data_type), **kwargs + encoder=encoder or OneHotEncoder(self, tokens_directory=tokens_directory), + **kwargs, ) def get_property_value(self, mol: Chem.rdchem.Mol) -> list[int]: diff --git a/chebai_graph/preprocessing/property_encoder.py b/chebai_graph/preprocessing/property_encoder.py index 2689a63..1173590 100644 --- a/chebai_graph/preprocessing/property_encoder.py +++ b/chebai_graph/preprocessing/property_encoder.py @@ -100,7 +100,7 @@ class IndexEncoder(PropertyEncoder): """ def __init__( - self, property, data_type: str, indices_dir: str | None = None, **kwargs + self, property, tokens_directory: str, indices_dir: str | None = None, **kwargs ) -> None: super().__init__(property, **kwargs) if indices_dir is None: @@ -108,7 +108,7 @@ def __init__( self.dirname = indices_dir self._use_one_hot_cache = kwargs.get("use_one_hot_cache", False) # load already existing cache - self._data_type = data_type + self.tokens_directory = tokens_directory with open(self.index_path, "r") as pk: self.cache: dict[str, int] = { token.strip(): idx for idx, token in enumerate(pk) @@ -132,11 +132,13 @@ def index_path(self) -> str: Returns: Path to index file. """ - assert self._data_type is not None, "data_type must be set for IndexEncoder" + assert self.tokens_directory is not None, ( + "tokens_directory must be set for IndexEncoder" + ) index_path = os.path.join( self.dirname, "bin", - self._data_type, + self.tokens_directory, self.property.name, f"indices_{self.name}.txt", ) From 11edbc098c02aa8c1feaf53beefad7abf83e23fe Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sun, 18 Oct 2026 08:46:24 +0200 Subject: [PATCH 39/58] for msp loss set tokens dir to shared --- chebai_graph/loss/msp.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/chebai_graph/loss/msp.py b/chebai_graph/loss/msp.py index b6e91ac..7973698 100644 --- a/chebai_graph/loss/msp.py +++ b/chebai_graph/loss/msp.py @@ -29,19 +29,21 @@ def __init__( self, hidden_size: int, dropout: float = 0.2, + tokens_directory: str = "_shared", ) -> None: super().__init__() + self.tokens_directory = tokens_directory self.atom_type_encoder = IndexEncoder( - property=AugAtomType(), + property=AugAtomType(tokens_directory=self.tokens_directory), use_one_hot_cache=True, eval=True, - tokens_directory="_shared", + tokens_directory=self.tokens_directory, ) self.bond_type_encoder = IndexEncoder( - property=AugBondType(), + property=AugBondType(tokens_directory=self.tokens_directory), use_one_hot_cache=True, eval=True, - tokens_directory="_shared", + tokens_directory=self.tokens_directory, ) # Bond existence projection From 114567c784e95782479222a0dd1a3462a79f395c Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sun, 18 Oct 2026 13:09:10 +0200 Subject: [PATCH 40/58] fix permanenet del of of custom data for collator: for remaining epochs --- chebai_graph/preprocessing/collate.py | 8 ++++++-- tests/unit/test_collate.py | 13 +++++++++++++ 2 files changed, 19 insertions(+), 2 deletions(-) diff --git a/chebai_graph/preprocessing/collate.py b/chebai_graph/preprocessing/collate.py index 51b1498..0497554 100644 --- a/chebai_graph/preprocessing/collate.py +++ b/chebai_graph/preprocessing/collate.py @@ -21,7 +21,9 @@ def __call__(self, data): # and unlabelled batches otherwise make PyG collate a partial `y` key. merged_data = [] for row in data: - features = row["features"] + # Dataset samples may be cached and reused across epochs. Collation + # removes temporary attributes below, so never mutate the source graph. + features = row["features"].clone() if "y" in features: del features.y merged_data.append(features) @@ -34,7 +36,9 @@ def __call__(self, data): if "edge_attr" not in store: store["edge_attr"] = torch.tensor([]) if "custom_metadata" in store: - graph_metadata.append(store.pop("custom_metadata", None)) + graph_metadata.append(store.pop("custom_metadata")) + else: + pass custom_metadata.append(graph_metadata) # Ensure all attributes are float tensors to prevent torch.cat dtype issues diff --git a/tests/unit/test_collate.py b/tests/unit/test_collate.py index d55ad9a..cb91498 100644 --- a/tests/unit/test_collate.py +++ b/tests/unit/test_collate.py @@ -43,6 +43,19 @@ def test_mixed_missing_labels_do_not_enter_graph_collate(self): ) self.assertNotIn("y", batch.x[0]) + def test_custom_metadata_survives_repeated_collation(self): + metadata = {"nodes": {"atom_nodes": [6]}, "edges": {}} + graph = self._graph() + graph.custom_metadata = metadata + row = {"features": graph, "labels": None, "ident": "a"} + + first_batch = GraphCollator()([row]) + second_batch = GraphCollator()([row]) + + self.assertEqual(first_batch.additional_fields["custom_metadata"], [[metadata]]) + self.assertEqual(second_batch.additional_fields["custom_metadata"], [[metadata]]) + self.assertIn("custom_metadata", graph) + if __name__ == "__main__": unittest.main() From d78322f43f8f876f474cb23f2b5b3d9ca17dc637 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sun, 18 Oct 2026 20:32:14 +0200 Subject: [PATCH 41/58] print for loaded pretrained ckpt --- chebai_graph/models/architectures/base.py | 1 + tests/unit/test_collate.py | 4 +++- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/chebai_graph/models/architectures/base.py b/chebai_graph/models/architectures/base.py index b6f6e26..67f753b 100644 --- a/chebai_graph/models/architectures/base.py +++ b/chebai_graph/models/architectures/base.py @@ -181,6 +181,7 @@ def _load_pretrained_gnn(self, checkpoint_path: str) -> None: "Skipped incompatible pretrained edge weights: " + ", ".join(sorted(skipped_edge_weights)) ) + print(f"Loaded pretrained GNN weights from {checkpoint_path!r} successfully.") @abstractmethod def _get_gnn(self, config: dict) -> torch.nn.Module: diff --git a/tests/unit/test_collate.py b/tests/unit/test_collate.py index cb91498..8c81979 100644 --- a/tests/unit/test_collate.py +++ b/tests/unit/test_collate.py @@ -53,7 +53,9 @@ def test_custom_metadata_survives_repeated_collation(self): second_batch = GraphCollator()([row]) self.assertEqual(first_batch.additional_fields["custom_metadata"], [[metadata]]) - self.assertEqual(second_batch.additional_fields["custom_metadata"], [[metadata]]) + self.assertEqual( + second_batch.additional_fields["custom_metadata"], [[metadata]] + ) self.assertIn("custom_metadata", graph) From 8af009e9dd88841414a3c92a5f4b55174be33af1 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Tue, 20 Oct 2026 10:27:24 +0200 Subject: [PATCH 42/58] update visualization script --- .gitignore | 1 + .../final/chebi50_final_augmented.yml | 2 - results/combine_json_file.py | 37 + results/visualize_augmented_molecule.py | 60 +- tutorials/classwise_stats.ipynb | 738 ++++-------------- 5 files changed, 224 insertions(+), 614 deletions(-) create mode 100644 results/combine_json_file.py diff --git a/.gitignore b/.gitignore index 2a04373..b46a5ba 100644 --- a/.gitignore +++ b/.gitignore @@ -174,6 +174,7 @@ electra_pretrained.ckpt *.out *.sh *.ckpt +*.csv !configs/data/augmented/aug-ablation/*.sh !.vscode/launch.json diff --git a/configs/data/augmented/final/chebi50_final_augmented.yml b/configs/data/augmented/final/chebi50_final_augmented.yml index 1b711f0..3ff3232 100644 --- a/configs/data/augmented/final/chebi50_final_augmented.yml +++ b/configs/data/augmented/final/chebi50_final_augmented.yml @@ -1,7 +1,5 @@ class_path: chebai_graph.preprocessing.datasets.ChEBI50_GNWA_FGWA_WFGE_AsPerNodeType init_args: - splits_file_path: data/chebi_v252/ChEBI50/processed/splits.csv - chebi_version: 252 properties: # All Node type properties - chebai_graph.preprocessing.properties.AtomNodeLevel diff --git a/results/combine_json_file.py b/results/combine_json_file.py new file mode 100644 index 0000000..41da467 --- /dev/null +++ b/results/combine_json_file.py @@ -0,0 +1,37 @@ +import json + +import pandas as pd + +file1 = "/home/aditya/python-chebai/data/chebi_v252/ChEBI50/processed/gatv2_baseline_classes_test.json" +file2 = "/home/aditya/python-chebai/data/chebi_v252/ChEBI50/processed/gatv2_aug_classes_test.json" + +# Output CSV file +output_file = "combined_f1_scores.csv" + +# Load JSON data +with open(file1, "r") as f: + data1 = json.load(f) + +with open(file2, "r") as f: + data2 = json.load(f) + +# Create a combined set of labels from both files +labels = sorted(set(data1.keys()) | set(data2.keys())) + +# Create output rows +rows = [] +for label in labels: + rows.append( + { + "label": label, + file1: data1.get(label, {}).get("f1", None), + file2: data2.get(label, {}).get("f1", None), + } + ) + +# Create DataFrame and save to CSV +df = pd.DataFrame(rows) +df.to_csv(output_file, index=False) + +print(f"CSV file created: {output_file}") +print(df) diff --git a/results/visualize_augmented_molecule.py b/results/visualize_augmented_molecule.py index fcc406b..cfdd8cc 100644 --- a/results/visualize_augmented_molecule.py +++ b/results/visualize_augmented_molecule.py @@ -1,9 +1,26 @@ +try: + from plotly import graph_objects as go +except ImportError as e: + raise ImportError( + "Plotly is required for 3D plotting. Install it with `pip install plotly`." + ) from e + +try: + import matplotlib + + assert matplotlib.__version__ >= "3.10.9", "Matplotlib >= 3.10.9 required" +except (ImportError, AssertionError) as e: + raise ImportError( + "Matplotlib is required for 2D plotting. Install it with `pip install matplotlib`." + ) from e + import io -import matplotlib +matplotlib.use("Agg") import matplotlib.pyplot as plt import networkx as nx import numpy as np +from chebi_utils.read_molecule import smiles_or_inchi_to_mol from jsonargparse import CLI from PIL import Image from rdkit.Chem import AllChem, BondType, Mol, rdDepictor @@ -23,8 +40,6 @@ GN_WithAtoms_FG_WithAtoms_NoFGE, ) -matplotlib.use("TkAgg") - EDGE_COLOR_MAP = { k.WITHIN_ATOMS_EDGE: "#1f77b4", k.ATOM_FG_EDGE: "#9467bd", @@ -241,13 +256,16 @@ def _draw_simple(G: nx.Graph) -> None: plt.show() -def _draw_3d(G: nx.Graph, mol: Mol) -> None: +def _draw_3d( + G: nx.Graph, mol: Mol, output_html: str = "augmented_molecule_3d.html" +) -> None: """ Visualize the graph in 3D using Plotly. Args: G (nx.Graph): Augmented molecular graph. mol (Chem.Mol): RDKit molecule object for 3D coordinates. + output_html (str): Path where the interactive Plotly figure is saved. Raises: ImportError: If Plotly is not installed. @@ -312,7 +330,7 @@ def _draw_3d(G: nx.Graph, mol: Mol) -> None: ] ) graph_pos = { - graph_node: (graph_pos_arr[:, 0].mean(), graph_pos_arr[:, 1].mean(), 4) + graph_node: (graph_pos_arr[:, 0].mean(), graph_pos_arr[:, 1].mean(), -2) } pos = {**atom_pos, **fg_pos, **graph_pos} @@ -343,6 +361,7 @@ def _draw_3d(G: nx.Graph, mol: Mol) -> None: z=zs, mode="lines", line=dict(color=EDGE_COLOR_MAP[edge_type], width=4), + opacity=0.5, name=edge_type, hoverinfo="none", ) @@ -360,15 +379,25 @@ def _draw_3d(G: nx.Graph, mol: Mol) -> None: x=pos_x, y=pos_y, z=pos_z, - mode="markers+text", + mode="markers", marker=dict(size=8, color=node_colors, opacity=0.9), - text=node_names, - textposition="top center", hovertext=node_ids, hoverinfo="text", ) - fig = go.Figure(data=edge_traces + [node_trace]) + label_trace = go.Scatter3d( + x=pos_x, + y=pos_y, + z=pos_z, + mode="text", + text=node_names, + textposition="top center", + textfont=dict(color="black", size=12), + hoverinfo="none", + showlegend=False, + ) + + fig = go.Figure(data=edge_traces + [node_trace, label_trace]) fig.update_layout( title="3D Augmented Molecule Graph", showlegend=True, @@ -379,6 +408,7 @@ def _draw_3d(G: nx.Graph, mol: Mol) -> None: ), margin=dict(l=0, r=0, b=0, t=40), ) + # fig.write_html(output_html) fig.show() @@ -387,6 +417,7 @@ def plot_augmented_graph( augmented_molecule: dict, mol: Mol, plot_type: str, + output_html: str = "augmented_molecule_3d.html", ) -> None: """ Main plotting function to visualize the augmented graph. @@ -396,6 +427,7 @@ def plot_augmented_graph( augmented_molecule (dict): Augmented Molecule. mol (Chem.Mol): RDKit molecule object. plot_type (str): One of ["simple", "h", "3d"]. + output_html (str): Path where the 3D Plotly figure is saved. """ G = _create_graph( edge_index, augmented_molecule["nodes"], augmented_molecule["edges"] @@ -406,7 +438,7 @@ def plot_augmented_graph( elif plot_type == "simple": _draw_simple(G) elif plot_type == "3d": - _draw_3d(G, mol) + _draw_3d(G, mol, output_html) else: raise ValueError(f"Unknown plot type: {plot_type}") @@ -471,6 +503,7 @@ def plot( smiles: str = "OC(=O)c1ccccc1O", plot_type: str = "simple", reader: str = "w_fge_w_gn", + output_html: str = "augmented_molecule_3d.html", ) -> None: """ Plot an augmented molecular graph from SMILES. @@ -491,9 +524,10 @@ def plot( - 'gnwa_fgwa_wfge': Graph node connected to atoms and FG nodes connected to atoms, with FG edges. - 'gn_wall_fgwa_nfge': Graph node connected to all atoms and FG nodes, no FG edges. - 'gn_wall_fgwa_wfge': Graph node connected to all atoms and FG nodes, with FG edges. + output_html (str): Output path for the interactive 3D HTML figure. """ fg_reader = READER[reader]() - mol = fg_reader._smiles_to_mol(smiles) + mol = smiles_or_inchi_to_mol(smiles) if mol is None: raise ValueError(f"Invalid SMILES: {smiles}") @@ -502,7 +536,9 @@ def plot( return edge_index, augmented_molecule = fg_reader._create_augmented_graph(mol) - plot_augmented_graph(edge_index, augmented_molecule, mol, plot_type) + plot_augmented_graph( + edge_index, augmented_molecule, mol, plot_type, output_html + ) if __name__ == "__main__": diff --git a/tutorials/classwise_stats.ipynb b/tutorials/classwise_stats.ipynb index 072a128..48972e7 100644 --- a/tutorials/classwise_stats.ipynb +++ b/tutorials/classwise_stats.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "code", - "execution_count": 21, + "execution_count": null, "id": "07f2eb38", "metadata": {}, "outputs": [], @@ -12,172 +12,46 @@ }, { "cell_type": "code", - "execution_count": 10, + "execution_count": null, "id": "eb64bfee", "metadata": {}, - "outputs": [ - { - "data": { - "text/html": [ - "
\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
labelres_baseline_propsres_wfgn_wfge_wgn_wprops_apoolgatv2_wfgn_wfge_wgn_wprops_apoolgatv2 - baselineres - baseline
01461800.00.00000.00000.00000.0000
1600040.00.11760.00000.00000.1176
2167330.00.23530.00000.00000.2353
3222990.00.28570.00000.00000.2857
4178150.00.12500.11110.11110.1250
\n", - "
" - ], - "text/plain": [ - " label res_baseline_props res_wfgn_wfge_wgn_wprops_apool \\\n", - "0 146180 0.0 0.0000 \n", - "1 60004 0.0 0.1176 \n", - "2 16733 0.0 0.2353 \n", - "3 22299 0.0 0.2857 \n", - "4 17815 0.0 0.1250 \n", - "\n", - " gatv2_wfgn_wfge_wgn_wprops_apool gatv2 - baseline res - baseline \n", - "0 0.0000 0.0000 0.0000 \n", - "1 0.0000 0.0000 0.1176 \n", - "2 0.0000 0.0000 0.2353 \n", - "3 0.0000 0.0000 0.2857 \n", - "4 0.1111 0.1111 0.1250 " - ] - }, - "execution_count": 10, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ - "df = pd.read_excel(\"G:\\github-aditya0by0\\python-chebai-graph\\classwise_f1_score.xlsx\")\n", + "df = pd.read_csv(\"/home/aditya/python-chebai-graph/combined_f1_scores.csv\")\n", + "df[\"aug-baseline\"] = df[\"augmented\"] - df[\"baseline\"]\n", "df.head()" ] }, { "cell_type": "code", - "execution_count": 18, + "execution_count": null, "id": "14c10198", "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "(1528, 6)" - ] - }, - "execution_count": 18, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "df.shape" ] }, { "cell_type": "code", - "execution_count": 13, + "execution_count": null, "id": "32062eb6", "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "Index(['label', 'res_baseline_props', 'res_wfgn_wfge_wgn_wprops_apool',\n", - " 'gatv2_wfgn_wfge_wgn_wprops_apool', 'gatv2 - baseline ',\n", - " 'res - baseline'],\n", - " dtype='object')" - ] - }, - "execution_count": 13, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "df.columns" ] }, { "cell_type": "code", - "execution_count": 14, + "execution_count": null, "id": "5cd9265a", "metadata": {}, "outputs": [], "source": [ - "BASELINE = \"res_baseline_props\"\n", - "RES = \"res_wfgn_wfge_wgn_wprops_apool\"\n", - "GATv2 = \"gatv2_wfgn_wfge_wgn_wprops_apool\"\n", - "GATv2_BASELINE_DIFF = \"gatv2 - baseline \"\n", - "RES_BASELINE_DIFF = \"res - baseline\"" + "BASELINE = \"baseline\"\n", + "AUGMENTED = \"augmented\"\n", + "DIFF_COL = \"aug-baseline\"" ] }, { @@ -192,45 +66,24 @@ }, { "cell_type": "code", - "execution_count": 16, + "execution_count": null, "id": "b38fa321", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Mean improvement over baseline:\n", - " res - baseline 0.054853\n", - "gatv2 - baseline 0.068136\n", - "dtype: float64\n", - "\n", - "Median improvement over baseline:\n", - " res - baseline 0.03330\n", - "gatv2 - baseline 0.04535\n", - "dtype: float64\n", - "\n", - "Variation in improvement:\n", - " res - baseline 0.144831\n", - "gatv2 - baseline 0.139299\n", - "dtype: float64\n" - ] - } - ], + "outputs": [], "source": [ "# mean improvement\n", - "mean_diff = df[[RES_BASELINE_DIFF, GATv2_BASELINE_DIFF]].mean()\n", + "mean_diff = df[[DIFF_COL]].mean()\n", "print(\"Mean improvement over baseline:\\n\", mean_diff)\n", "print()\n", "\n", "# median improvement\n", - "median_diff = df[[RES_BASELINE_DIFF, GATv2_BASELINE_DIFF]].median()\n", + "median_diff = df[[DIFF_COL]].median()\n", "print(\"Median improvement over baseline:\\n\", median_diff)\n", "print()\n", "\n", "\n", "# standard deviation of improvements\n", - "std_diff = df[[RES_BASELINE_DIFF, GATv2_BASELINE_DIFF]].std()\n", + "std_diff = df[[DIFF_COL]].std()\n", "print(\"Variation in improvement:\\n\", std_diff)" ] }, @@ -244,24 +97,13 @@ }, { "cell_type": "code", - "execution_count": 19, + "execution_count": null, "id": "d046db01", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Labels improved in ResGated Final model over baseline: 1010\n", - "Labels improved in GATv2 Final model over baseline: 1091\n" - ] - } - ], + "outputs": [], "source": [ - "res_improvement = (df[RES_BASELINE_DIFF] > 0).sum()\n", - "gat_v2_improvement = (df[GATv2_BASELINE_DIFF] > 0).sum()\n", + "gat_v2_improvement = (df[DIFF_COL] > 0).sum()\n", "\n", - "print(f\"Labels improved in ResGated Final model over baseline: {res_improvement}\")\n", "print(f\"Labels improved in GATv2 Final model over baseline: {gat_v2_improvement}\")" ] }, @@ -275,79 +117,18 @@ }, { "cell_type": "code", - "execution_count": 27, - "id": "3c3972a1", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "best_model\n", - "gatv2_wfgn_wfge_wgn_wprops_apool 678\n", - "res_wfgn_wfge_wgn_wprops_apool 524\n", - "res_baseline_props 326\n", - "Name: count, dtype: int64\n" - ] - } - ], - "source": [ - "df[\"best_model\"] = df[[BASELINE, RES, GATv2]].idxmax(axis=1)\n", - "print(df[\"best_model\"].value_counts())" - ] - }, - { - "cell_type": "code", - "execution_count": 32, + "execution_count": null, "id": "02ecb78e", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "\n", - "Top 10 classes with highest improvement for res - baseline:\n", - " 25 0.8889\n", - "24 0.8889\n", - "37 0.7692\n", - "34 0.7143\n", - "35 0.7143\n", - "29 0.7143\n", - "77 0.6889\n", - "41 0.6487\n", - "36 0.6400\n", - "23 0.6364\n", - "Name: res - baseline, dtype: float64\n", - "\n", - "Top 10 classes with highest improvement for gatv2 - baseline :\n", - " 25 0.8889\n", - "24 0.8889\n", - "37 0.7692\n", - "34 0.7143\n", - "35 0.7143\n", - "29 0.7143\n", - "77 0.6889\n", - "41 0.6487\n", - "36 0.6400\n", - "23 0.6364\n", - "Name: res - baseline, dtype: float64\n" - ] - } - ], + "outputs": [], "source": [ "# Classes with highest difference\n", - "top_diff_res = df[RES_BASELINE_DIFF].sort_values(ascending=False).head(10)\n", + "top_diff_res = df[DIFF_COL].sort_values(ascending=False).head(10)\n", "print(\n", - " f\"\\nTop 10 classes with highest improvement for {RES_BASELINE_DIFF}:\\n\",\n", + " f\"\\nTop 10 classes with highest improvement for {DIFF_COL}:\\n\",\n", " top_diff_res,\n", ")\n", - "\n", - "top_diff_gat = df[GATv2_BASELINE_DIFF].sort_values(ascending=False).head(10)\n", - "print(\n", - " f\"\\nTop 10 classes with highest improvement for {GATv2_BASELINE_DIFF}:\\n\",\n", - " top_diff_res,\n", - ")" + "\n" ] }, { @@ -355,397 +136,154 @@ "execution_count": null, "id": "524a01d8", "metadata": {}, - "outputs": [ - { - "data": { - "text/html": [ - "
\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
labelres_baseline_propsres_wfgn_wfge_wgn_wprops_apoolgatv2_wfgn_wfge_wgn_wprops_apoolgatv2 - baselineres - baselinebest_model
01461800.00.00.00000.00000.0res_baseline_props
9257110.00.00.25000.25000.0gatv2_wfgn_wfge_wgn_wprops_apool
10390930.00.00.25000.25000.0gatv2_wfgn_wfge_wgn_wprops_apool
1841940.00.00.46150.46150.0gatv2_wfgn_wfge_wgn_wprops_apool
19176080.00.00.46150.46150.0gatv2_wfgn_wfge_wgn_wprops_apool
........................
1523766291.01.01.00000.00000.0res_baseline_props
1524787761.01.01.00000.00000.0res_baseline_props
1525792041.01.01.00000.00000.0res_baseline_props
15261327421.01.01.00000.00000.0res_baseline_props
15271366371.01.01.00000.00000.0res_baseline_props
\n", - "

87 rows × 7 columns

\n", - "
" - ], - "text/plain": [ - " label res_baseline_props res_wfgn_wfge_wgn_wprops_apool \\\n", - "0 146180 0.0 0.0 \n", - "9 25711 0.0 0.0 \n", - "10 39093 0.0 0.0 \n", - "18 4194 0.0 0.0 \n", - "19 17608 0.0 0.0 \n", - "... ... ... ... \n", - "1523 76629 1.0 1.0 \n", - "1524 78776 1.0 1.0 \n", - "1525 79204 1.0 1.0 \n", - "1526 132742 1.0 1.0 \n", - "1527 136637 1.0 1.0 \n", - "\n", - " gatv2_wfgn_wfge_wgn_wprops_apool gatv2 - baseline res - baseline \\\n", - "0 0.0000 0.0000 0.0 \n", - "9 0.2500 0.2500 0.0 \n", - "10 0.2500 0.2500 0.0 \n", - "18 0.4615 0.4615 0.0 \n", - "19 0.4615 0.4615 0.0 \n", - "... ... ... ... \n", - "1523 1.0000 0.0000 0.0 \n", - "1524 1.0000 0.0000 0.0 \n", - "1525 1.0000 0.0000 0.0 \n", - "1526 1.0000 0.0000 0.0 \n", - "1527 1.0000 0.0000 0.0 \n", - "\n", - " best_model \n", - "0 res_baseline_props \n", - "9 gatv2_wfgn_wfge_wgn_wprops_apool \n", - "10 gatv2_wfgn_wfge_wgn_wprops_apool \n", - "18 gatv2_wfgn_wfge_wgn_wprops_apool \n", - "19 gatv2_wfgn_wfge_wgn_wprops_apool \n", - "... ... \n", - "1523 res_baseline_props \n", - "1524 res_baseline_props \n", - "1525 res_baseline_props \n", - "1526 res_baseline_props \n", - "1527 res_baseline_props \n", - "\n", - "[87 rows x 7 columns]" - ] - }, - "execution_count": 42, - "metadata": {}, - "output_type": "execute_result" - } - ], + "outputs": [], "source": [ "# Classes with minimal or no difference\n", "minimal_diff_threshold = 0.01 # you can adjust\n", - "minimal_diff = df[df[RES_BASELINE_DIFF].abs() <= minimal_diff_threshold]\n", + "minimal_diff = df[df[DIFF_COL].abs() <= minimal_diff_threshold]\n", "minimal_diff" ] }, { "cell_type": "code", - "execution_count": 37, + "execution_count": null, "id": "5f15e872", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "\n", - "Classes where baseline performs better than our model:\n", - " label res_baseline_props res_wfgn_wfge_wgn_wprops_apool \\\n", - "45 59999 0.1176 0.0541 \n", - "47 59266 0.1176 0.1143 \n", - "50 25095 0.1333 0.0000 \n", - "70 48927 0.2000 0.0000 \n", - "72 83943 0.2000 0.0000 \n", - "... ... ... ... \n", - "1508 24846 1.0000 0.9333 \n", - "1509 26187 1.0000 0.9333 \n", - "1510 78231 1.0000 0.9333 \n", - "1511 26377 1.0000 0.9565 \n", - "1512 38834 1.0000 0.9677 \n", - "\n", - " gatv2_wfgn_wfge_wgn_wprops_apool gatv2 - baseline res - baseline \\\n", - "45 0.1250 0.0074 -0.0635 \n", - "47 0.5385 0.4209 -0.0033 \n", - "50 0.2667 0.1334 -0.1333 \n", - "70 0.0000 -0.2000 -0.2000 \n", - "72 0.2000 0.0000 -0.2000 \n", - "... ... ... ... \n", - "1508 1.0000 0.0000 -0.0667 \n", - "1509 1.0000 0.0000 -0.0667 \n", - "1510 1.0000 0.0000 -0.0667 \n", - "1511 1.0000 0.0000 -0.0435 \n", - "1512 1.0000 0.0000 -0.0323 \n", - "\n", - " best_model \n", - "45 gatv2_wfgn_wfge_wgn_wprops_apool \n", - "47 gatv2_wfgn_wfge_wgn_wprops_apool \n", - "50 gatv2_wfgn_wfge_wgn_wprops_apool \n", - "70 res_baseline_props \n", - "72 res_baseline_props \n", - "... ... \n", - "1508 res_baseline_props \n", - "1509 res_baseline_props \n", - "1510 res_baseline_props \n", - "1511 res_baseline_props \n", - "1512 res_baseline_props \n", - "\n", - "[431 rows x 7 columns]\n" - ] - } - ], + "outputs": [], "source": [ "# Classes where baseline performs better\n", - "baseline_better = df[df[RES_BASELINE_DIFF] < 0]\n", + "baseline_better = df[df[DIFF_COL] < 0]\n", "print(\"\\nClasses where baseline performs better than our model:\\n\", baseline_better)" ] }, { "cell_type": "code", - "execution_count": 38, + "execution_count": null, "id": "0c4292bf", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "\n", - "Classes where both baseline and res_wfgn_wfge_wgn_wprops_apool have F1 < 0.5:\n", - " label res_baseline_props res_wfgn_wfge_wgn_wprops_apool \\\n", - "0 146180 0.0000 0.0000 \n", - "1 60004 0.0000 0.1176 \n", - "2 16733 0.0000 0.2353 \n", - "3 22299 0.0000 0.2857 \n", - "4 17815 0.0000 0.1250 \n", - ".. ... ... ... \n", - "265 32877 0.4615 0.4706 \n", - "271 139358 0.4667 0.4000 \n", - "283 35346 0.4762 0.4706 \n", - "284 35507 0.4762 0.4737 \n", - "295 24586 0.4878 0.4762 \n", - "\n", - " gatv2_wfgn_wfge_wgn_wprops_apool gatv2 - baseline res - baseline \\\n", - "0 0.0000 0.0000 0.0000 \n", - "1 0.0000 0.0000 0.1176 \n", - "2 0.0000 0.0000 0.2353 \n", - "3 0.0000 0.0000 0.2857 \n", - "4 0.1111 0.1111 0.1250 \n", - ".. ... ... ... \n", - "265 0.5333 0.0718 0.0091 \n", - "271 0.5517 0.0850 -0.0667 \n", - "283 0.3529 -0.1233 -0.0056 \n", - "284 0.3871 -0.0891 -0.0025 \n", - "295 0.6667 0.1789 -0.0116 \n", - "\n", - " best_model \n", - "0 res_baseline_props \n", - "1 res_wfgn_wfge_wgn_wprops_apool \n", - "2 res_wfgn_wfge_wgn_wprops_apool \n", - "3 res_wfgn_wfge_wgn_wprops_apool \n", - "4 res_wfgn_wfge_wgn_wprops_apool \n", - ".. ... \n", - "265 gatv2_wfgn_wfge_wgn_wprops_apool \n", - "271 gatv2_wfgn_wfge_wgn_wprops_apool \n", - "283 res_baseline_props \n", - "284 res_baseline_props \n", - "295 gatv2_wfgn_wfge_wgn_wprops_apool \n", - "\n", - "[148 rows x 7 columns]\n" - ] - } - ], + "outputs": [], "source": [ "# Classes where both baseline and our model perform worse\n", "# Assuming \"worse\" means below a threshold, e.g., f1 < 0.5\n", "threshold = 0.5\n", - "both_worse = df[(df[BASELINE] < threshold) & (df[RES] < threshold)]\n", - "print(f\"\\nClasses where both baseline and {RES} have F1 < {threshold}:\\n\", both_worse)" + "both_worse = df[(df[BASELINE] < threshold) & (df[AUGMENTED] < threshold)]\n", + "print(f\"\\nClasses where both baseline and Augmented have F1 < {threshold}:\\n\", both_worse)" ] }, { "cell_type": "code", - "execution_count": 40, + "execution_count": null, "id": "8948b7bb", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "\n", - "Classes where baseline performs better than our model:\n", - " label res_baseline_props res_wfgn_wfge_wgn_wprops_apool \\\n", - "45 59999 0.1176 0.0541 \n", - "47 59266 0.1176 0.1143 \n", - "50 25095 0.1333 0.0000 \n", - "70 48927 0.2000 0.0000 \n", - "72 83943 0.2000 0.0000 \n", - "... ... ... ... \n", - "1508 24846 1.0000 0.9333 \n", - "1509 26187 1.0000 0.9333 \n", - "1510 78231 1.0000 0.9333 \n", - "1511 26377 1.0000 0.9565 \n", - "1512 38834 1.0000 0.9677 \n", - "\n", - " gatv2_wfgn_wfge_wgn_wprops_apool gatv2 - baseline res - baseline \\\n", - "45 0.1250 0.0074 -0.0635 \n", - "47 0.5385 0.4209 -0.0033 \n", - "50 0.2667 0.1334 -0.1333 \n", - "70 0.0000 -0.2000 -0.2000 \n", - "72 0.2000 0.0000 -0.2000 \n", - "... ... ... ... \n", - "1508 1.0000 0.0000 -0.0667 \n", - "1509 1.0000 0.0000 -0.0667 \n", - "1510 1.0000 0.0000 -0.0667 \n", - "1511 1.0000 0.0000 -0.0435 \n", - "1512 1.0000 0.0000 -0.0323 \n", - "\n", - " best_model \n", - "45 gatv2_wfgn_wfge_wgn_wprops_apool \n", - "47 gatv2_wfgn_wfge_wgn_wprops_apool \n", - "50 gatv2_wfgn_wfge_wgn_wprops_apool \n", - "70 res_baseline_props \n", - "72 res_baseline_props \n", - "... ... \n", - "1508 res_baseline_props \n", - "1509 res_baseline_props \n", - "1510 res_baseline_props \n", - "1511 res_baseline_props \n", - "1512 res_baseline_props \n", - "\n", - "[431 rows x 7 columns]\n" - ] - } - ], + "outputs": [], "source": [ "# Classes where baseline performs better\n", - "baseline_better = df[df[RES_BASELINE_DIFF] < 0]\n", + "baseline_better = df[df[DIFF_COL] < 0]\n", "print(\"\\nClasses where baseline performs better than our model:\\n\", baseline_better)" ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "0bc0f4c1", + "metadata": {}, + "outputs": [], + "source": [ + "len(df)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "61ab61c7", + "metadata": {}, + "outputs": [], + "source": [ + "print(len(df[df[DIFF_COL] >= 0.1]))\n", + "print(len(df[df[DIFF_COL] >= 0.01]))\n", + "\n", + "print(len(df[abs(df[DIFF_COL]) < 0.01]))\n", + "\n", + "print(len(df[df[DIFF_COL] <= -0.1]))\n", + "print(len(df[df[DIFF_COL] <= -0.01]))\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "828c727f", + "metadata": {}, + "outputs": [], + "source": [ + "1005 + 221 + 283" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7b5476b8", + "metadata": {}, + "outputs": [], + "source": [ + "len(df[df[DIFF_COL] == 0.0])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2caab4d5", + "metadata": {}, + "outputs": [], + "source": [ + "len(df[(df[DIFF_COL] == 0.0) & ~((df[BASELINE] == 1.0) & (df[AUGMENTED] == 1.0))])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fb50e15a", + "metadata": {}, + "outputs": [], + "source": [ + "\n", + "for label in [57880, 28802, 49172, 140949,59268, 59412,25627]:\n", + " print(label, label in df[df[DIFF_COL] >= 0.1][\"label\"].tolist())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bba6bb1a", + "metadata": {}, + "outputs": [], + "source": [ + "df[df[DIFF_COL] >= 0.1]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "20e2ea95", + "metadata": {}, + "outputs": [], + "source": [ + "df[(df[BASELINE] <= 0.1) & (df[AUGMENTED] >= 0.8)]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b9e043cf", + "metadata": {}, + "outputs": [], + "source": [ + "df.iloc[df[DIFF_COL].sort_values(ascending=False).head(10).index]" + ] } ], "metadata": { "kernelspec": { - "display_name": "gnn3", + "display_name": "python-chebai-graph (3.10.20)", "language": "python", "name": "python3" }, @@ -759,7 +297,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.10.16" + "version": "3.10.20" } }, "nbformat": 4, From 2b7554f82a03c3582765b10748d6c5bb2fb958c2 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Tue, 20 Oct 2026 13:22:57 +0200 Subject: [PATCH 43/58] fix prediction pipeline issue for graph --- chebai_graph/preprocessing/collate.py | 7 ++++++- chebai_graph/preprocessing/datasets/base.py | 4 ++-- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/chebai_graph/preprocessing/collate.py b/chebai_graph/preprocessing/collate.py index 0497554..1f5dcca 100644 --- a/chebai_graph/preprocessing/collate.py +++ b/chebai_graph/preprocessing/collate.py @@ -23,7 +23,12 @@ def __call__(self, data): for row in data: # Dataset samples may be cached and reused across epochs. Collation # removes temporary attributes below, so never mutate the source graph. - features = row["features"].clone() + if isinstance(row["features"], tuple): + # If tuple, then metadata is in second element. Move it to the first element for consistency. + row["features"][0]["custom_metadata"] = row["features"][1] + features = row["features"][0].clone() + else: + features = row["features"].clone() if "y" in features: del features.y merged_data.append(features) diff --git a/chebai_graph/preprocessing/datasets/base.py b/chebai_graph/preprocessing/datasets/base.py index 76fedf2..70d62ca 100644 --- a/chebai_graph/preprocessing/datasets/base.py +++ b/chebai_graph/preprocessing/datasets/base.py @@ -190,7 +190,7 @@ def _after_setup(self, **kwargs) -> None: self._setup_properties() super()._after_setup(**kwargs) - def _preprocess_smiles_for_pred( + def _preprocess_molecule_for_pred( self, idx, raw_data: str | Chem.Mol, model_hparams: Optional[dict] = None ) -> Optional[dict]: """Preprocess prediction data.""" @@ -646,7 +646,7 @@ def _merge_props_into_base( is_atom_node=is_atom_node, is_fg_node=is_fg_node, is_graph_node=is_graph_node, - custom_metadata=geom_data.custom_metadata, + custom_metadata=getattr(geom_data, "custom_metadata", None), ) def _build_node_property_tensor( From 5a9b75784ab842ab5235b38567676405f7a13c77 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sat, 24 Oct 2026 11:33:32 +0200 Subject: [PATCH 44/58] molnet final script --- .gitignore | 1 + configs/data/augmented/final/molnet.sh | 153 +++++++++++++++++++++++++ 2 files changed, 154 insertions(+) create mode 100644 configs/data/augmented/final/molnet.sh diff --git a/.gitignore b/.gitignore index b46a5ba..21f9d60 100644 --- a/.gitignore +++ b/.gitignore @@ -178,3 +178,4 @@ electra_pretrained.ckpt !configs/data/augmented/aug-ablation/*.sh !.vscode/launch.json +!configs/data/augmented/final/*.sh \ No newline at end of file diff --git a/configs/data/augmented/final/molnet.sh b/configs/data/augmented/final/molnet.sh new file mode 100644 index 0000000..ed7474e --- /dev/null +++ b/configs/data/augmented/final/molnet.sh @@ -0,0 +1,153 @@ +#!/bin/bash + +# ============================================================ +# Local single-GPU adaptation of ../aug-ablation-script.sh +# Run this script using: `bash configs/data/augmented/final/molnet.sh` +# +# The original is a Slurm job-array script. This version runs the +# six MolNet data configurations x three seeds serially on one GPU. +# +# Changes vs. the Slurm version: +# - Removed all #SBATCH scheduling directives. +# - Fixed hardcoded /home/staff/a/akhedekar paths -> local user. +# - Replaced SLURM_ARRAY_TASK_ID array fan-out with a serial loop +# over [seeds] x [data configs]. +# +# Requires (verified on this machine): +# - local venv: /home/aditya/python-chebai-graph/.venv +# - chebai CLI: /home/aditya/python-chebai (python -m chebai) +# - chebai_graph src: /home/aditya/python-chebai-graph. +# - wandb credentials in ~/.netrc (or run with a local logger). +# ============================================================ + +set -x +set -euo pipefail + +# ============================================================ +# Local paths +# ============================================================ +CHEBAI_DIR="/home/aditya/python-chebai" +CHEBAI_GRAPH_DIR="/home/aditya/python-chebai-graph" + + +# ============================================================ +# Seeds +# ============================================================ +SEEDS=(0 42 12345) + +# ============================================================ +# Activate Python environment +# ============================================================ +source "$CHEBAI_GRAPH_DIR/.venv/bin/activate" + +# ============================================================ +# Set working directory (chebai CLI looks up configs/... from cwd) +# ============================================================ +cd "$CHEBAI_DIR" +export SSL_CERT_FILE=$(python -m certifi) + +echo "============================================================" +echo "Node: $(hostname)" +echo "Date: $(date)" +echo "GPU:" +nvidia-smi +echo "============================================================" + +# ============================================================ +# Serial loop over all seeds and data configurations +# ============================================================ +DATASETS=(BBBP BACE SIDER ClinTox Tox21 ToxCast) + +for SEED in "${SEEDS[@]}"; do + for DATASET in "${DATASETS[@]}"; do + case "$DATASET" in + BBBP) + TRAINER="configs/training/binary_trainer.yml" + METRICS="configs/metrics/binary-f1-roc-auc.yml" + DROPOUT="0.5" + LR_GNN="0.0005" + LR_LINEAR="0.001" + ;; + BACE) + TRAINER="configs/training/binary_trainer.yml" + METRICS="configs/metrics/binary-f1-roc-auc.yml" + DROPOUT="0.5" + LR_GNN="0.001" + LR_LINEAR="0.001" + ;; + SIDER) + TRAINER="configs/training/default_trainer.yml" + METRICS="configs/metrics/micro-macro-f1-roc-auc.yml" + DROPOUT="0.5" + LR_GNN="0.0005" + LR_LINEAR="0.001" + ;; + ClinTox) + TRAINER="configs/training/default_trainer.yml" + METRICS="configs/metrics/micro-macro-f1-roc-auc.yml" + DROPOUT="0.5" + LR_GNN="0.001" + LR_LINEAR="0.001" + ;; + Tox21) + TRAINER="configs/training/default_trainer.yml" + METRICS="configs/metrics/micro-macro-f1-roc-auc.yml" + DROPOUT="0.7" + LR_GNN="0.0005" + LR_LINEAR="0.001" + ;; + ToxCast) + TRAINER="configs/training/default_trainer.yml" + METRICS="configs/metrics/micro-macro-f1-roc-auc.yml" + DROPOUT="0.7" + LR_GNN="0.001" + LR_LINEAR="0.001" + ;; + *) + echo "Unsupported dataset: $DATASET" >&2 + exit 1 + ;; + esac + + DATASET_LOWER=$(printf '%s' "$DATASET" | tr '[:upper:]' '[:lower:]') + RUN_NAME="${DATASET_LOWER}_gat_amg_augmented_s${SEED}" + + echo "============================================================" + echo "Starting run: $RUN_NAME" + echo "Dataset: $DATASET" + echo "Seed: $SEED" + echo "============================================================" + + python -m chebai fit \ + --trainer="$TRAINER" \ + --trainer.logger=configs/training/wandb_logger.yml \ + "--trainer.logger.init_args.tags=[\"augmented_paper\",\"aug-ablation\",\"$DATASET_LOWER\"]" \ + --trainer.callbacks=configs/training/roc-auc_callbacks.yml \ + --model="$CHEBAI_GRAPH_DIR/configs/model/augmented/pooling/amg_pool/gat.yml" \ + --model.train_metrics="$METRICS" \ + --model.test_metrics="$METRICS" \ + --model.val_metrics="$METRICS" \ + --data="$CHEBAI_GRAPH_DIR/configs/data/augmented/final/${DATASET}_final_augmented.yml" \ + --data.init_args.batch_size=32 \ + --data.init_args.num_workers=10 \ + --model.pass_loss_kwargs=false \ + --model.criterion=configs/loss/bce_unweighted.yml \ + --trainer.logger.init_args.name="$RUN_NAME" \ + --data.init_args.splits_file_path="data/${DATASET}:MNClassification/${DATASET}/processed/splits.csv" \ + --model.config.dropout="$DROPOUT" \ + --model.optimizer_kwargs.lr_gnn="$LR_GNN" \ + --model.optimizer_kwargs.lr_linear="$LR_LINEAR" \ + --seed_everything="$SEED" \ + --model.pretrained_checkpoint=/home/aditya/python-chebai-graph/msp_best_epoch=92_train_loss=0.0009.ckpt \ + --data.init_args.tokens_directory=_shared \ + + + echo "============================================================" + echo "Finished run: $RUN_NAME" + echo "============================================================" + done +done + +echo "============================================================" +echo "All experiments finished successfully" +echo "============================================================" From 42bdbf83b44bb326ad5e9ef28eef495e79772209 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Sun, 25 Oct 2026 07:47:31 +0100 Subject: [PATCH 45/58] add molnet script to run all classif dataset and report avg test scores for each --- chebai_graph/models/architectures/base.py | 5 +- configs/data/augmented/final/molnet.sh | 124 +++++++++++++++++- .../model/augmented/pooling/amg_pool/gat.yml | 1 - .../augmented/pooling/no_pooling/gat.yml | 1 - 4 files changed, 123 insertions(+), 8 deletions(-) diff --git a/chebai_graph/models/architectures/base.py b/chebai_graph/models/architectures/base.py index 67f753b..faa3bc7 100644 --- a/chebai_graph/models/architectures/base.py +++ b/chebai_graph/models/architectures/base.py @@ -77,7 +77,8 @@ def __init__(self, config: dict, **kwargs) -> None: self.hidden_channels = int(config["hidden_channels"]) self.out_channels = int(config["out_channels"]) self.edge_dim = int(config["edge_dim"]) # number of bond properties - self.dropout = float(config["dropout"]) + self.dropout = float(config.get("dropout", 0.0)) + print("----------Dropout set to: ", self.dropout) class GraphNetWrapper(GraphBaseNet, ABC): @@ -113,7 +114,7 @@ def __init__( if self.use_batch_norm: self.batch_norm = torch.nn.BatchNorm1d(self.lin_input_dim) - lin_hidden_dim = kwargs.get("lin_hidden_dim", gnn_out_dim) + lin_hidden_dim = kwargs.get("lin_hidden_dim", gnn_out_dim // 2) self.lin_sequential: torch.nn.Sequential = self._get_linear_module_list( n_linear_layers=n_linear_layers, in_dim=self.lin_input_dim, diff --git a/configs/data/augmented/final/molnet.sh b/configs/data/augmented/final/molnet.sh index ed7474e..0d59c98 100644 --- a/configs/data/augmented/final/molnet.sh +++ b/configs/data/augmented/final/molnet.sh @@ -20,7 +20,6 @@ # - wandb credentials in ~/.netrc (or run with a local logger). # ============================================================ -set -x set -euo pipefail # ============================================================ @@ -28,6 +27,15 @@ set -euo pipefail # ============================================================ CHEBAI_DIR="/home/aditya/python-chebai" CHEBAI_GRAPH_DIR="/home/aditya/python-chebai-graph" +LOG_DIR="$CHEBAI_DIR/logs" +LOG_FILE="$LOG_DIR/molnet_$(date +%Y%m%d_%H%M%S).txt" +mkdir -p "$LOG_DIR" + +# Keep output visible in the terminal while recording stdout, stderr, and +# Bash trace output in a timestamped text file. +exec > >(tee -a "$LOG_FILE") 2>&1 +set -x +echo "Writing MolNet run log to: $LOG_FILE" # ============================================================ @@ -57,6 +65,9 @@ echo "============================================================" # Serial loop over all seeds and data configurations # ============================================================ DATASETS=(BBBP BACE SIDER ClinTox Tox21 ToxCast) +RESULTS_FILE="$CHEBAI_DIR/logs/molnet_test_results.tsv" +mkdir -p "$(dirname "$RESULTS_FILE")" +printf 'dataset\tseed\tmetric\tvalue\tcheckpoint\n' > "$RESULTS_FILE" for SEED in "${SEEDS[@]}"; do for DATASET in "${DATASETS[@]}"; do @@ -64,6 +75,9 @@ for SEED in "${SEEDS[@]}"; do BBBP) TRAINER="configs/training/binary_trainer.yml" METRICS="configs/metrics/binary-f1-roc-auc.yml" + CALLBACKS="configs/training/binary_callbacks.yml" + CHECKPOINT_PREFIX="best_roc-auc_" + METRICS_TO_REPORT=(test_roc-auc test_f1) DROPOUT="0.5" LR_GNN="0.0005" LR_LINEAR="0.001" @@ -71,6 +85,9 @@ for SEED in "${SEEDS[@]}"; do BACE) TRAINER="configs/training/binary_trainer.yml" METRICS="configs/metrics/binary-f1-roc-auc.yml" + CALLBACKS="configs/training/binary_callbacks.yml" + CHECKPOINT_PREFIX="best_roc-auc_" + METRICS_TO_REPORT=(test_roc-auc test_f1) DROPOUT="0.5" LR_GNN="0.001" LR_LINEAR="0.001" @@ -78,6 +95,9 @@ for SEED in "${SEEDS[@]}"; do SIDER) TRAINER="configs/training/default_trainer.yml" METRICS="configs/metrics/micro-macro-f1-roc-auc.yml" + CALLBACKS="configs/training/roc-auc_callbacks.yml" + CHECKPOINT_PREFIX="best_himol-marco-roc-auc_" + METRICS_TO_REPORT=(test_himol-marco-roc-auc test_micro-f1 test_macro-f1) DROPOUT="0.5" LR_GNN="0.0005" LR_LINEAR="0.001" @@ -85,6 +105,9 @@ for SEED in "${SEEDS[@]}"; do ClinTox) TRAINER="configs/training/default_trainer.yml" METRICS="configs/metrics/micro-macro-f1-roc-auc.yml" + CALLBACKS="configs/training/roc-auc_callbacks.yml" + CHECKPOINT_PREFIX="best_himol-marco-roc-auc_" + METRICS_TO_REPORT=(test_himol-marco-roc-auc test_micro-f1 test_macro-f1) DROPOUT="0.5" LR_GNN="0.001" LR_LINEAR="0.001" @@ -92,6 +115,9 @@ for SEED in "${SEEDS[@]}"; do Tox21) TRAINER="configs/training/default_trainer.yml" METRICS="configs/metrics/micro-macro-f1-roc-auc.yml" + CALLBACKS="configs/training/roc-auc_callbacks.yml" + CHECKPOINT_PREFIX="best_himol-marco-roc-auc_" + METRICS_TO_REPORT=(test_himol-marco-roc-auc test_micro-f1 test_macro-f1) DROPOUT="0.7" LR_GNN="0.0005" LR_LINEAR="0.001" @@ -99,6 +125,9 @@ for SEED in "${SEEDS[@]}"; do ToxCast) TRAINER="configs/training/default_trainer.yml" METRICS="configs/metrics/micro-macro-f1-roc-auc.yml" + CALLBACKS="configs/training/roc-auc_callbacks.yml" + CHECKPOINT_PREFIX="best_himol-marco-roc-auc_" + METRICS_TO_REPORT=(test_himol-marco-roc-auc test_micro-f1 test_macro-f1) DROPOUT="0.7" LR_GNN="0.001" LR_LINEAR="0.001" @@ -118,12 +147,15 @@ for SEED in "${SEEDS[@]}"; do echo "Seed: $SEED" echo "============================================================" + RUN_MARKER=$(mktemp) + CALLBACK_ARGS=(--trainer.callbacks="$CALLBACKS") + python -m chebai fit \ --trainer="$TRAINER" \ --trainer.logger=configs/training/wandb_logger.yml \ - "--trainer.logger.init_args.tags=[\"augmented_paper\",\"aug-ablation\",\"$DATASET_LOWER\"]" \ - --trainer.callbacks=configs/training/roc-auc_callbacks.yml \ - --model="$CHEBAI_GRAPH_DIR/configs/model/augmented/pooling/amg_pool/gat.yml" \ + "--trainer.logger.init_args.tags=[\"augmented_paper\", \"classif-head\", \"aug-ablation\",\"$DATASET_LOWER\"]" \ + "${CALLBACK_ARGS[@]}" \ + --model="$CHEBAI_GRAPH_DIR/configs/model/augmented/pooling/no_pooling/gat.yml" \ --model.train_metrics="$METRICS" \ --model.test_metrics="$METRICS" \ --model.val_metrics="$METRICS" \ @@ -140,8 +172,72 @@ for SEED in "${SEEDS[@]}"; do --seed_everything="$SEED" \ --model.pretrained_checkpoint=/home/aditya/python-chebai-graph/msp_best_epoch=92_train_loss=0.0009.ckpt \ --data.init_args.tokens_directory=_shared \ + --model.init_args.n_linear_layers=2 \ + + CHECKPOINT_PATH=$( + find "$CHEBAI_DIR/logs" -type f \ + -name "${CHECKPOINT_PREFIX}*.ckpt" \ + -newer "$RUN_MARKER" -print | sort | tail -n 1 + ) + rm -f "$RUN_MARKER" + if [[ -z "$CHECKPOINT_PATH" ]]; then + echo "Could not find best checkpoint for $RUN_NAME" >&2 + exit 1 + fi + + TEST_NAME="${RUN_NAME}_test" + TEST_MARKER=$(mktemp) + python -m chebai test \ + --trainer="$TRAINER" \ + --trainer.logger=configs/training/csv_logger.yml \ + --trainer.logger.init_args.name="$TEST_NAME" \ + --trainer.devices=1 \ + --trainer.num_nodes=1 \ + --ckpt_path="$CHECKPOINT_PATH" \ + --model="$CHEBAI_GRAPH_DIR/configs/model/augmented/pooling/no_pooling/gat.yml" \ + --model.test_metrics="$METRICS" \ + --data="$CHEBAI_GRAPH_DIR/configs/data/augmented/final/${DATASET}_final_augmented.yml" \ + --data.init_args.batch_size=32 \ + --data.init_args.num_workers=10 \ + --model.pass_loss_kwargs=false \ + --model.criterion=configs/loss/bce_unweighted.yml \ + --data.init_args.splits_file_path="data/${DATASET}:MNClassification/${DATASET}/processed/splits.csv" \ + --data.init_args.tokens_directory=_shared \ + --model.init_args.n_linear_layers=2 \ + TEST_METRICS_FILE=$( + find "$CHEBAI_DIR/logs/$TEST_NAME" -type f \ + -name metrics.csv -newer "$TEST_MARKER" -print | sort | tail -n 1 + ) + rm -f "$TEST_MARKER" + if [[ ! -f "$TEST_METRICS_FILE" ]]; then + echo "Could not find fresh test metrics for $TEST_NAME" >&2 + exit 1 + fi + METRIC_VALUES=$(python - "$TEST_METRICS_FILE" "${METRICS_TO_REPORT[@]}" <<'PY' +import csv +import sys +metrics_path, *metric_names = sys.argv[1:] +values = {metric_name: None for metric_name in metric_names} +with open(metrics_path, newline="") as metrics_file: + for row in csv.DictReader(metrics_file): + for metric_name in metric_names: + if row.get(metric_name): + values[metric_name] = row[metric_name] +for metric_name in metric_names: + value = values[metric_name] + if value is None: + raise SystemExit(f"Metric {metric_name!r} not found in {metrics_path}") + print(f"{metric_name}\t{value}") +PY +) + while IFS=$'\t' read -r TEST_METRIC TEST_VALUE; do + printf '%s\t%s\t%s\t%s\t%s\n' \ + "$DATASET" "$SEED" "$TEST_METRIC" "$TEST_VALUE" "$CHECKPOINT_PATH" \ + >> "$RESULTS_FILE" + echo "Test $TEST_METRIC: $TEST_VALUE" + done <<< "$METRIC_VALUES" echo "============================================================" echo "Finished run: $RUN_NAME" echo "============================================================" @@ -151,3 +247,23 @@ done echo "============================================================" echo "All experiments finished successfully" echo "============================================================" + +python - "$RESULTS_FILE" <<'PY' +import csv +import math +import statistics +import sys +from collections import defaultdict + +results = defaultdict(list) +with open(sys.argv[1], newline="") as results_file: + for row in csv.DictReader(results_file, delimiter="\t"): + results[(row["dataset"], row["metric"])].append(float(row["value"])) + +print("\nDataset test performance (best validation checkpoint per seed)") +print("dataset\tmetric\tmean\tsample_sd\tn") +for (dataset, metric), values in sorted(results.items()): + mean = statistics.mean(values) + sample_sd = statistics.stdev(values) if len(values) > 1 else math.nan + print(f"{dataset}\t{metric}\t{mean:.6f}\t{sample_sd:.6f}\t{len(values)}") +PY diff --git a/configs/model/augmented/pooling/amg_pool/gat.yml b/configs/model/augmented/pooling/amg_pool/gat.yml index f88159d..186fbc8 100644 --- a/configs/model/augmented/pooling/amg_pool/gat.yml +++ b/configs/model/augmented/pooling/amg_pool/gat.yml @@ -10,5 +10,4 @@ init_args: edge_dim: 12 # number of bond properties heads: 8 # the number of heads should be divisible by output channels (hidden channels if output channel not given) v2: True # This uses `torch_geometric.nn.conv.GATv2Conv` convolution layers, set False to use `GATConv` - dropout: 0 n_linear_layers: 1 diff --git a/configs/model/augmented/pooling/no_pooling/gat.yml b/configs/model/augmented/pooling/no_pooling/gat.yml index f9d3575..55b9cfc 100644 --- a/configs/model/augmented/pooling/no_pooling/gat.yml +++ b/configs/model/augmented/pooling/no_pooling/gat.yml @@ -10,5 +10,4 @@ init_args: edge_dim: 12 # number of bond properties heads: 8 # the number of heads should be divisible by output channels (hidden channels if output channel not given) v2: True # This uses `torch_geometric.nn.conv.GATv2Conv` convolution layers, set False to use `GATConv` - dropout: 0 n_linear_layers: 1 From b47a2dc0f011b51d4593436c164d2508038760f6 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Mon, 26 Oct 2026 08:16:14 +0100 Subject: [PATCH 46/58] update changes --- chebai_graph/models/architectures/base.py | 4 +- configs/data/augmented/final/molnet.sh | 3 +- .../model/augmented/pooling/amg_pool/gine.yml | 1 - results/molecule_distance_statistics.py | 164 ++++++++++++++++++ 4 files changed, 169 insertions(+), 3 deletions(-) create mode 100644 results/molecule_distance_statistics.py diff --git a/chebai_graph/models/architectures/base.py b/chebai_graph/models/architectures/base.py index faa3bc7..f24079a 100644 --- a/chebai_graph/models/architectures/base.py +++ b/chebai_graph/models/architectures/base.py @@ -101,6 +101,7 @@ def __init__( n_linear_layers (int): Number of linear layers. **kwargs: Additional arguments. """ + lin_hidden_dim = kwargs.pop("lin_hidden_dim", None) super().__init__(**kwargs) self.gnn = self._get_gnn(config) if pretrained_checkpoint is not None: @@ -113,8 +114,9 @@ def __init__( self.use_batch_norm = use_batch_norm if self.use_batch_norm: self.batch_norm = torch.nn.BatchNorm1d(self.lin_input_dim) + if lin_hidden_dim is None: + lin_hidden_dim = gnn_out_dim // 2 - lin_hidden_dim = kwargs.get("lin_hidden_dim", gnn_out_dim // 2) self.lin_sequential: torch.nn.Sequential = self._get_linear_module_list( n_linear_layers=n_linear_layers, in_dim=self.lin_input_dim, diff --git a/configs/data/augmented/final/molnet.sh b/configs/data/augmented/final/molnet.sh index 0d59c98..cb9af03 100644 --- a/configs/data/augmented/final/molnet.sh +++ b/configs/data/augmented/final/molnet.sh @@ -20,7 +20,8 @@ # - wandb credentials in ~/.netrc (or run with a local logger). # ============================================================ -set -euo pipefail +set -Eeuo pipefail +trap 'status=$?; echo "ERROR: command failed with status ${status} at line ${BASH_LINENO[0]}: ${BASH_COMMAND}" >&2; exit "$status"' ERR # ============================================================ # Local paths diff --git a/configs/model/augmented/pooling/amg_pool/gine.yml b/configs/model/augmented/pooling/amg_pool/gine.yml index 773088a..a40a4bd 100644 --- a/configs/model/augmented/pooling/amg_pool/gine.yml +++ b/configs/model/augmented/pooling/amg_pool/gine.yml @@ -8,6 +8,5 @@ init_args: out_channels: 512 num_layers: 4 edge_dim: 12 # number of bond properties - dropout: 0 train_eps: true n_linear_layers: 1 diff --git a/results/molecule_distance_statistics.py b/results/molecule_distance_statistics.py new file mode 100644 index 0000000..337516b --- /dev/null +++ b/results/molecule_distance_statistics.py @@ -0,0 +1,164 @@ +#!/usr/bin/env python3 +"""Compute molecular graph-diameter statistics from a saved PyTorch dataset.""" + +from __future__ import annotations + +from rdkit import RDLogger + +# Cached molecules may have been serialized with an older RDKit release. +RDLogger.DisableLog("rdApp.warning") + +import argparse +import json +import statistics +from collections.abc import Iterable +from pathlib import Path +from typing import Any + +import torch +from rdkit import Chem +from rdkit.Chem import Draw, Mol, rdmolops + + +def extract_molecule(sample: Any) -> Mol: + """Extract the atom-level RDKit molecule from one saved dataset sample.""" + try: + molecule = sample["features"][1]["nodes"]["atom_nodes"] + except (IndexError, KeyError, TypeError) as error: + raise ValueError( + "Sample does not contain features[1]['nodes']['atom_nodes']" + ) from error + + if not isinstance(molecule, Mol): + raise TypeError( + "features[1]['nodes']['atom_nodes'] is not an RDKit molecule: " + f"{type(molecule).__name__}" + ) + return molecule + + +def molecule_diameter(molecule: Mol) -> int: + """Return the largest finite shortest-path distance in the molecule.""" + if molecule.GetNumAtoms() == 0: + raise ValueError("Cannot calculate a diameter for a molecule with no atoms") + + distances = rdmolops.GetDistanceMatrix(molecule) + components = rdmolops.GetMolFrags(molecule, asMols=False) + return max( + int(distances[atom_i][atom_j]) + for component in components + for atom_i in component + for atom_j in component + ) + + +def draw_molecule(molecule: Mol, output_path: Path) -> None: + """Draw an RDKit molecule and save it as an image.""" + output_path.parent.mkdir(parents=True, exist_ok=True) + Draw.MolToFile(molecule, str(output_path), size=(800, 600)) + + +def iter_samples(dataset: Any) -> Iterable[tuple[int, Any]]: + """Yield dataset samples while preserving their instance indices.""" + if isinstance(dataset, dict): + yield from dataset.items() + return + + try: + yield from enumerate(dataset) + except TypeError as error: + raise TypeError( + "The loaded .pt file must contain an iterable dataset" + ) from error + + +def calculate_statistics( + dataset: Any, +) -> tuple[list[dict[str, Any]], dict[str, Any], Mol]: + """Calculate per-instance molecular diameters and aggregate statistics.""" + distances: list[dict[str, Any]] = [] + molecules: list[Mol] = [] + for index, sample in iter_samples(dataset): + molecule = extract_molecule(sample) + molecules.append(molecule) + distances.append( + { + "index": int(index), + "num_atoms": molecule.GetNumAtoms(), + "largest_distance": molecule_diameter(molecule), + "smiles": Chem.MolToSmiles(molecule), + } + ) + + if not distances: + raise ValueError("The loaded dataset contains no samples") + + values = [item["largest_distance"] for item in distances] + largest_index = values.index(max(values)) + largest_molecule = distances[largest_index] + largest_rdkit_molecule = molecules[largest_index] + summary: dict[str, Any] = { + "num_molecules": len(values), + "min": min(values), + "max": max(values), + "mean": statistics.mean(values), + "median": statistics.median(values), + "stdev": statistics.stdev(values) if len(values) > 1 else 0.0, + "largest_distance_index": largest_molecule["index"], + "largest_distance_smiles": largest_molecule["smiles"], + } + return distances, summary, largest_rdkit_molecule + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Compute the largest topological (bond-path) distance between any " + "two atoms for every molecule in a saved .pt dataset." + ) + ) + parser.add_argument("data_path", type=Path, help="Path to the saved data.pt file") + parser.add_argument( + "--output", + type=Path, + help="Optional JSON file for per-molecule distances and aggregate statistics", + ) + parser.add_argument( + "--image-output", + type=Path, + default=Path("largest_distance_molecule.png"), + help=( + "PNG file for the molecule with the largest shortest path " + "(default: largest_distance_molecule.png)" + ), + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + dataset = torch.load(args.data_path, weights_only=False) + distances, summary, largest_molecule = calculate_statistics(dataset) + + print(f"Molecules: {summary['num_molecules']}") + print(f"Largest distance (min): {summary['min']}") + print(f"Largest distance (max): {summary['max']}") + print(f"Largest distance (mean): {summary['mean']:.3f}") + print(f"Largest distance (median): {summary['median']:.3f}") + print(f"Largest distance (stdev): {summary['stdev']:.3f}") + print( + "Molecule with largest shortest path " + f"(distance={summary['max']}, index={summary['largest_distance_index']}): " + f"{summary['largest_distance_smiles']}" + ) + draw_molecule(largest_molecule, args.image_output) + print(f"Molecule drawing written to {args.image_output}") + + if args.output is not None: + result = {"summary": summary, "molecules": distances} + args.output.write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8") + print(f"Results written to {args.output}") + + +if __name__ == "__main__": + main() From 7afe9bd1884e2fefa212e4f081dadf8d5caa9254 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Mon, 26 Oct 2026 18:17:32 +0100 Subject: [PATCH 47/58] update changes --- chebai_graph/models/architectures/gine.py | 86 ++++++++++++++++--- configs/data/augmented/final/molnet.sh | 12 +-- .../augmented/pooling/amg_pool/himol_gine.yml | 14 +++ 3 files changed, 94 insertions(+), 18 deletions(-) create mode 100644 configs/model/augmented/pooling/amg_pool/himol_gine.yml diff --git a/chebai_graph/models/architectures/gine.py b/chebai_graph/models/architectures/gine.py index f9e89a3..95696dc 100644 --- a/chebai_graph/models/architectures/gine.py +++ b/chebai_graph/models/architectures/gine.py @@ -1,7 +1,7 @@ from typing import Any, Final from torch import Tensor -from torch.nn import ELU +from torch.nn import ELU, ReLU from torch_geometric import nn as tgnn from torch_geometric.data import Data as GraphData from torch_geometric.nn.conv import MessagePassing @@ -65,6 +65,55 @@ def init_conv( return tgnn.GINEConv(mlp, **kwargs) +class HiMolGINEModel(BasicGNN): + """ + A GIN-based GNN model based on PyG's BasicGNN, using GINEConv layers so that + edge (bond) features are incorporated into the message-passing step. + + See: + + Attributes: + supports_edge_weight (bool): Indicates edge weights are not supported. + supports_edge_attr (bool): Indicates edge attributes are supported. + supports_norm_batch (bool): Indicates if batch normalization is supported. + """ + + supports_edge_weight: Final[bool] = False + supports_edge_attr: Final[bool] = True + supports_norm_batch: Final[bool] + + def init_conv( + self, in_channels: int | tuple[int, int], out_channels: int, **kwargs: Any + ) -> MessagePassing: + """ + Initializes a GINEConv layer. + + The inner network is a 2-layer MLP (Linear -> act -> Linear, no + activation on the last layer), matching both + `torch_geometric.nn.models.GIN.init_conv` and the GIN paper's + message-transform MLP. `edge_dim` (passed via **kwargs) lets + GINEConv linearly project bond features onto the node feature + space before adding them into the neighbor messages. + + Args: + in_channels (int or Tuple[int, int]): Number of input channels. + out_channels (int): Number of output channels. + **kwargs: Additional keyword arguments for the convolution layer + (e.g. `edge_dim`, `train_eps`). + + Returns: + MessagePassing: A GINEConv layer instance. + """ + mlp = MLP( + [in_channels, 2 * out_channels, out_channels], + act=self.act, + act_first=self.act_first, + norm=self.norm, + norm_kwargs=self.norm_kwargs, + ) + return tgnn.GINEConv(mlp, **kwargs) + + class GINEConvNetBase(GraphModelBase): """ Base model class for applying GINEConv layers to graph-structured data. @@ -85,20 +134,33 @@ class GINEConvNetBase(GraphModelBase): """ def __init__(self, config: dict[str, Any], **kwargs: Any): + use_himol_model = bool(config.pop("use_himol_model", False)) + super().__init__(config=config, **kwargs) self.activation = ELU() # Instantiate ELU once for reuse. self.train_eps = bool(config.get("train_eps", True)) - - self.gine: BasicGNN = GINEModel( - in_channels=self.in_channels, - hidden_channels=self.hidden_channels, - out_channels=self.out_channels, - num_layers=self.num_layers, - dropout=self.dropout, - edge_dim=self.edge_dim, - train_eps=self.train_eps, - act=self.activation, - ) + if use_himol_model: + self.gine: BasicGNN = HiMolGINEModel( + in_channels=self.in_channels, + hidden_channels=self.hidden_channels, + out_channels=self.out_channels, + num_layers=self.num_layers, + dropout=self.dropout, + edge_dim=self.edge_dim, + train_eps=self.train_eps, + act=ReLU(), + ) + else: + self.gine: BasicGNN = GINEModel( + in_channels=self.in_channels, + hidden_channels=self.hidden_channels, + out_channels=self.out_channels, + num_layers=self.num_layers, + dropout=self.dropout, + edge_dim=self.edge_dim, + train_eps=self.train_eps, + act=self.activation, + ) def forward(self, batch: dict[str, Any]) -> Tensor: """ diff --git a/configs/data/augmented/final/molnet.sh b/configs/data/augmented/final/molnet.sh index cb9af03..f7c40ed 100644 --- a/configs/data/augmented/final/molnet.sh +++ b/configs/data/augmented/final/molnet.sh @@ -154,9 +154,9 @@ for SEED in "${SEEDS[@]}"; do python -m chebai fit \ --trainer="$TRAINER" \ --trainer.logger=configs/training/wandb_logger.yml \ - "--trainer.logger.init_args.tags=[\"augmented_paper\", \"classif-head\", \"aug-ablation\",\"$DATASET_LOWER\"]" \ + "--trainer.logger.init_args.tags=[\"augmented_paper\",\"$DATASET_LOWER\"]" \ "${CALLBACK_ARGS[@]}" \ - --model="$CHEBAI_GRAPH_DIR/configs/model/augmented/pooling/no_pooling/gat.yml" \ + --model="$CHEBAI_GRAPH_DIR/configs/model/augmented/pooling/amg_pool/gat.yml" \ --model.train_metrics="$METRICS" \ --model.test_metrics="$METRICS" \ --model.val_metrics="$METRICS" \ @@ -171,9 +171,10 @@ for SEED in "${SEEDS[@]}"; do --model.optimizer_kwargs.lr_gnn="$LR_GNN" \ --model.optimizer_kwargs.lr_linear="$LR_LINEAR" \ --seed_everything="$SEED" \ - --model.pretrained_checkpoint=/home/aditya/python-chebai-graph/msp_best_epoch=92_train_loss=0.0009.ckpt \ --data.init_args.tokens_directory=_shared \ - --model.init_args.n_linear_layers=2 \ + # --model.pretrained_checkpoint=/home/aditya/python-chebai-graph/msp_best_epoch=92_train_loss=0.0009.ckpt \ + + CHECKPOINT_PATH=$( find "$CHEBAI_DIR/logs" -type f \ @@ -195,7 +196,7 @@ for SEED in "${SEEDS[@]}"; do --trainer.devices=1 \ --trainer.num_nodes=1 \ --ckpt_path="$CHECKPOINT_PATH" \ - --model="$CHEBAI_GRAPH_DIR/configs/model/augmented/pooling/no_pooling/gat.yml" \ + --model="$CHEBAI_GRAPH_DIR/configs/model/augmented/pooling/amg_pool/gat.yml" \ --model.test_metrics="$METRICS" \ --data="$CHEBAI_GRAPH_DIR/configs/data/augmented/final/${DATASET}_final_augmented.yml" \ --data.init_args.batch_size=32 \ @@ -204,7 +205,6 @@ for SEED in "${SEEDS[@]}"; do --model.criterion=configs/loss/bce_unweighted.yml \ --data.init_args.splits_file_path="data/${DATASET}:MNClassification/${DATASET}/processed/splits.csv" \ --data.init_args.tokens_directory=_shared \ - --model.init_args.n_linear_layers=2 \ TEST_METRICS_FILE=$( find "$CHEBAI_DIR/logs/$TEST_NAME" -type f \ diff --git a/configs/model/augmented/pooling/amg_pool/himol_gine.yml b/configs/model/augmented/pooling/amg_pool/himol_gine.yml new file mode 100644 index 0000000..4a22416 --- /dev/null +++ b/configs/model/augmented/pooling/amg_pool/himol_gine.yml @@ -0,0 +1,14 @@ +class_path: chebai_graph.models.GINEAMGPoolGraphPred +init_args: + optimizer_kwargs: + lr: 1e-3 + config: + in_channels: 203 # number of node/atom properties + hidden_channels: 512 + out_channels: 512 + num_layers: 5 + edge_dim: 12 # number of bond properties + train_eps: true + use_himol_model: true + n_linear_layers: 2 + lin_hidden_dim: 256 From 2d9ff97c63d74ab7c46dde490474de1b706597e1 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Tue, 27 Oct 2026 07:23:13 +0100 Subject: [PATCH 48/58] pre-commit fix --- .gitignore | 2 +- tutorials/classwise_stats.ipynb | 12 ++++++------ 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/.gitignore b/.gitignore index 21f9d60..491079d 100644 --- a/.gitignore +++ b/.gitignore @@ -178,4 +178,4 @@ electra_pretrained.ckpt !configs/data/augmented/aug-ablation/*.sh !.vscode/launch.json -!configs/data/augmented/final/*.sh \ No newline at end of file +!configs/data/augmented/final/*.sh diff --git a/tutorials/classwise_stats.ipynb b/tutorials/classwise_stats.ipynb index 48972e7..a64d8e0 100644 --- a/tutorials/classwise_stats.ipynb +++ b/tutorials/classwise_stats.ipynb @@ -127,8 +127,7 @@ "print(\n", " f\"\\nTop 10 classes with highest improvement for {DIFF_COL}:\\n\",\n", " top_diff_res,\n", - ")\n", - "\n" + ")" ] }, { @@ -167,7 +166,9 @@ "# Assuming \"worse\" means below a threshold, e.g., f1 < 0.5\n", "threshold = 0.5\n", "both_worse = df[(df[BASELINE] < threshold) & (df[AUGMENTED] < threshold)]\n", - "print(f\"\\nClasses where both baseline and Augmented have F1 < {threshold}:\\n\", both_worse)" + "print(\n", + " f\"\\nClasses where both baseline and Augmented have F1 < {threshold}:\\n\", both_worse\n", + ")" ] }, { @@ -205,7 +206,7 @@ "print(len(df[abs(df[DIFF_COL]) < 0.01]))\n", "\n", "print(len(df[df[DIFF_COL] <= -0.1]))\n", - "print(len(df[df[DIFF_COL] <= -0.01]))\n" + "print(len(df[df[DIFF_COL] <= -0.01]))" ] }, { @@ -245,8 +246,7 @@ "metadata": {}, "outputs": [], "source": [ - "\n", - "for label in [57880, 28802, 49172, 140949,59268, 59412,25627]:\n", + "for label in [57880, 28802, 49172, 140949, 59268, 59412, 25627]:\n", " print(label, label in df[df[DIFF_COL] >= 0.1][\"label\"].tolist())" ] }, From e8dbbfc051708bdefc1bcd19817a57e042cc0498 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Tue, 27 Oct 2026 07:31:25 +0100 Subject: [PATCH 49/58] pre-commit fix --- results/molecule_distance_statistics.py | 10 ++++------ results/visualize_augmented_molecule.py | 7 ------- 2 files changed, 4 insertions(+), 13 deletions(-) diff --git a/results/molecule_distance_statistics.py b/results/molecule_distance_statistics.py index 337516b..e1b7901 100644 --- a/results/molecule_distance_statistics.py +++ b/results/molecule_distance_statistics.py @@ -3,11 +3,6 @@ from __future__ import annotations -from rdkit import RDLogger - -# Cached molecules may have been serialized with an older RDKit release. -RDLogger.DisableLog("rdApp.warning") - import argparse import json import statistics @@ -16,9 +11,12 @@ from typing import Any import torch -from rdkit import Chem +from rdkit import Chem, RDLogger from rdkit.Chem import Draw, Mol, rdmolops +# Cached molecules may have been serialized with an older RDKit release. +RDLogger.DisableLog("rdApp.warning") + def extract_molecule(sample: Any) -> Mol: """Extract the atom-level RDKit molecule from one saved dataset sample.""" diff --git a/results/visualize_augmented_molecule.py b/results/visualize_augmented_molecule.py index cfdd8cc..76bd49d 100644 --- a/results/visualize_augmented_molecule.py +++ b/results/visualize_augmented_molecule.py @@ -1,10 +1,3 @@ -try: - from plotly import graph_objects as go -except ImportError as e: - raise ImportError( - "Plotly is required for 3D plotting. Install it with `pip install plotly`." - ) from e - try: import matplotlib From aafc917232fb5d2e6933923ed6a888db16eba264 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Tue, 27 Oct 2026 07:59:58 +0100 Subject: [PATCH 50/58] remove duplicates introduced in https://github.com/ChEB-AI/python-chebai-graph/pull/51 Reference: Also be PR can be abandoned and closed https://github.com/ChEB-AI/python-chebai-graph/pull/53 --- chebai_graph/preprocessing/datasets/chebi.py | 642 ------------------- 1 file changed, 642 deletions(-) diff --git a/chebai_graph/preprocessing/datasets/chebi.py b/chebai_graph/preprocessing/datasets/chebi.py index 31850b2..7a68dc6 100644 --- a/chebai_graph/preprocessing/datasets/chebi.py +++ b/chebai_graph/preprocessing/datasets/chebi.py @@ -46,648 +46,6 @@ def __init__(self, **kwargs): super().__init__(**kwargs) -class DataPropertiesSetter(ChEBIOverX, ABC): - """Mixin for adding molecular property encodings to graph-based ChEBI datasets.""" - - READER = GraphPropertyReader - - def __init__( - self, - properties: list | None = None, - transform: Callable | None = None, - **kwargs, - ): - """ - Initialize GraphPropertiesMixIn. - - Args: - properties: Optional list of MolecularProperty class paths or instances. - transform: Optional transformation applied to each data sample. - """ - super().__init__(**kwargs) - # atom_properties and bond_properties are given as lists containing class_paths - if properties is not None: - properties = [resolve_property(prop) for prop in properties] - properties = self._sort_properties(properties) - else: - properties = [] - self.properties: list[MolecularProperty] = properties - assert isinstance(self.properties, list) and all( - isinstance(p, MolecularProperty) for p in self.properties - ) - self.transform = transform - - def _sort_properties( - self, properties: list[MolecularProperty] - ) -> list[MolecularProperty]: - return sorted(properties, key=lambda prop: self.get_property_path(prop)) - - def _setup_properties(self) -> None: - """ - Process and cache molecular properties to disk. - - Returns: - None - """ - raw_data = [] - os.makedirs(self.processed_properties_dir, exist_ok=True) - - try: - file_names = self.processed_main_file_names - except NotImplementedError: - file_names = self.raw_file_names - - for file in file_names: - # processed_dir_main only exists for ChEBI datasets - path = os.path.join( - ( - self.processed_dir_main - if hasattr(self, "processed_dir_main") - else self.raw_dir - ), - file, - ) - raw_data += list(self._load_dict(path)) - - idents = [row["ident"] for row in raw_data] - features = [row["features"] for row in raw_data] - - # use vectorized version of encode function, apply only if value is present - def enc_if_not_none(encode, value): - return ( - [encode(v) for v in value] - if value is not None and len(value) > 0 - else None - ) - - if any( - not os.path.isfile(self.get_property_path(property)) - for property in self.properties - ): - # augment molecule graph if possible (this would also happen for the properties if needed, but this avoids redundancy) - if isinstance(self.reader, _AugmentorReader): - returned_results = [] - for mol in features: - try: - r = self.reader._create_augmented_graph(mol) - except Exception: - r = None - returned_results.append(r) - mols = [ - augmented_mol[1] if augmented_mol is not None else None - for augmented_mol in returned_results - ] - else: - mols = features - - for property in self.properties: - if not os.path.isfile(self.get_property_path(property)): - rank_zero_info(f"Processing property {property.name}") - # read all property values first, then encode - rank_zero_info(f"\tReading property values of {property.name}...") - property_values = [ - self.reader.read_property(mol, property) - if mol is not None - else None - for mol in tqdm.tqdm(mols) - ] - rank_zero_info(f"\tEncoding property values of {property.name}...") - property.encoder.on_start(property_values=property_values) - encoded_values = [ - enc_if_not_none(property.encoder.encode, value) - for value in tqdm.tqdm(property_values) - ] - assert len(encoded_values) == len(idents) == len(features) - torch.save( - [ - { - property.name: property.encoder.compress( - torch.cat(feat) - ), - "ident": id, - } - for feat, id in zip(encoded_values, idents) - if feat is not None - ], - self.get_property_path(property), - ) - property.on_finish() - - @property - def processed_properties_dir(self) -> str: - return os.path.join(self.processed_dir, "properties") - - def get_property_path(self, property: MolecularProperty) -> str: - """ - Construct the cache path for a given molecular property. - - Args: - property: Instance of a MolecularProperty. - - Returns: - Path to the cached property file. - """ - return os.path.join( - self.processed_properties_dir, - f"{property.name}_{property.encoder.name}.pt", - ) - - def _after_setup(self, **kwargs) -> None: - """ - Finalize setup after ensuring properties are processed. - - Args: - **kwargs: Additional keyword arguments passed to superclass. - - Returns: - None - """ - self._setup_properties() - super()._after_setup(**kwargs) - - def _preprocess_molecule_for_pred( - self, idx, raw_data: str | Chem.Mol, model_hparams: Optional[dict] = None - ) -> Optional[dict]: - """Preprocess prediction data.""" - # Add dummy labels because the collate function requires them. - # Note: If labels are set to `None`, the collator will insert a `non_null_labels` entry into `loss_kwargs`, - # which later causes `_get_prediction_and_labels` method in the prediction pipeline to treat the data as empty. - result = self.reader.to_data( - {"id": f"smiles_{idx}", "features": raw_data, "labels": [1, 2]} - ) - # _read_data can return an updated version of the input data (e.g. augmented molecule dict) along with the GeomData object - if isinstance(result["features"], tuple): - result["features"], raw_data = result["features"] - if result is None or result["features"] is None: - return None - for property in self.properties: - property.encoder.eval = True - property_value = self.reader.read_property(raw_data, property) - if property_value is None or len(property_value) == 0: - encoded_value = None - else: - encoded_value = torch.stack( - [property.encoder.encode(v) for v in property_value] - ) - if len(encoded_value.shape) == 3: - encoded_value = encoded_value.squeeze(0) - result[property.name] = encoded_value - - result["features"] = self._prediction_merge_props_into_base_wrapper( - result, model_hparams - ) - - # apply transformation, e.g. masking for pretraining task - if self.transform is not None: - result["features"] = self.transform(result["features"]) - - return result - - def _prediction_merge_props_into_base_wrapper( - self, row: pd.Series | dict, model_hparams: Optional[dict] = None - ) -> GeomData: - """ - Wrapper to merge properties into base features for prediction. - - Args: - row: A dictionary or pd.Series containing 'features' and encoded properties. - Returns: - A GeomData object with merged features. - """ - return self._merge_props_into_base(row) - - -class GraphPropertiesMixIn(DataPropertiesSetter, ABC): - def __init__( - self, - properties=None, - transform=None, - pad_node_features: int = None, - pad_edge_features: int = None, - distribution: str = "normal", - **kwargs, - ): - super().__init__(properties, transform, **kwargs) - self.pad_edge_features = int(pad_edge_features) if pad_edge_features else None - self.pad_node_features = int(pad_node_features) if pad_node_features else None - if self.pad_node_features or self.pad_edge_features: - assert ( - distribution is not None - and distribution in RandomFeatureInitializationReader.DISTRIBUTIONS - ), ( - "When using padding for features, a valid distribution must be specified." - ) - self.distribution = distribution - if self.pad_node_features: - print( - f"[Info] Node-level features will be padded with random" - f"{self.pad_node_features} values from {self.distribution} distribution." - ) - if self.pad_edge_features: - print( - f"[Info] Edge-level features will be padded with random" - f"{self.pad_edge_features} values from {self.distribution} distribution." - ) - - if self.properties: - print( - f"Data module uses these properties (ordered): {', '.join([str(p) for p in self.properties])}" - ) - - def _merge_props_into_base(self, row: pd.Series | dict) -> GeomData: - """ - Merge encoded molecular properties into the GeomData object. - - Args: - row: A dictionary containing 'features' and encoded properties. - - Returns: - A GeomData object with merged features. - """ - if isinstance(row["features"], tuple): - geom_data, _ = row[ - "features" - ] # ignore additional returned data from _read_data (e.g. augmented molecule dict) - else: - geom_data = row["features"] - assert isinstance(geom_data, GeomData) - edge_attr = geom_data.edge_attr - x = geom_data.x - molecule_attr = torch.empty((1, 0)) - - for property in self.properties: - property_values = row[f"{property.name}"] - if isinstance(property_values, torch.Tensor): - if len(property_values.size()) == 0: - property_values = property_values.unsqueeze(0) - if len(property_values.size()) == 1: - property_values = property_values.unsqueeze(1) - else: - property_values = torch.zeros( - (0, property.encoder.get_encoding_length()) - ) - - if isinstance(property, AtomProperty): - x = torch.cat([x, property_values], dim=1) - elif isinstance(property, BondProperty): - # Concat/Duplicate properties values for undirected graph as `edge_index` has first src to tgt edges, then tgt to src edges - edge_attr = torch.cat( - [edge_attr, torch.cat([property_values, property_values], dim=0)], - dim=1, - ) - elif isinstance(property, MoleculeProperty): - molecule_attr = torch.cat([molecule_attr, property_values], dim=1) - else: - raise TypeError(f"Unsupported property type: {type(property).__name__}") - - if self.pad_node_features: - padding_values = torch.empty((x.shape[0], self.pad_node_features)) - RandomFeatureInitializationReader.random_gni( - padding_values, self.distribution - ) - x = torch.cat([x, padding_values], dim=1) - - if self.pad_edge_features: - padding_values = torch.empty((edge_attr.shape[0], self.pad_edge_features)) - RandomFeatureInitializationReader.random_gni( - padding_values, self.distribution - ) - edge_attr = torch.cat([edge_attr, padding_values], dim=1) - - return GeomData( - x=x, - edge_index=geom_data.edge_index, - edge_attr=edge_attr, - molecule_attr=molecule_attr, - ) - - def load_processed_data( - self, kind: Optional[str] = None, filename: Optional[str] = None - ) -> list[dict]: - """ - Load dataset and merge cached properties into base features. - - Args: - filename: The path to the file to load. - - Returns: - List of data entries, each a dictionary. - """ - base_data = super().load_processed_data(kind, filename) - base_df = pd.DataFrame(base_data) - base_df["ident"] = base_df["ident"].astype(str) - - for property in self.properties: - property_data = torch.load( - self.get_property_path(property), weights_only=False - ) - for entry in property_data: - entry[property.name] = property.encoder.decompress(entry[property.name]) - if len(property_data[0][property.name].shape) > 1: - property.encoder.set_encoding_length( - property_data[0][property.name].shape[1] - ) - - property_df = pd.DataFrame(property_data) - property_df.rename( - columns={property.name: f"{property.name}"}, inplace=True - ) - property_df["ident"] = property_df["ident"].astype(str) - base_df = base_df.merge(property_df, on="ident", how="left") - - base_df["features"] = base_df.apply( - lambda row: self._merge_props_into_base(row), axis=1 - ) - - # apply transformation, e.g. masking for pretraining task - if self.transform is not None: - base_df["features"] = base_df["features"].apply(self.transform) - - prop_lengths = [ - (prop.name, prop.encoder.get_encoding_length()) for prop in self.properties - ] - - # -------------------------- Count total node properties - n_node_properties = sum( - p.encoder.get_encoding_length() - for p in self.properties - if isinstance(p, AtomProperty) - ) - - in_channels_str = "" - if self.pad_node_features: - n_node_properties += self.pad_node_features - in_channels_str += f" (with {self.pad_node_features} padded random values from {self.distribution} distribution)" - - in_channels_str = f"in_channels: {n_node_properties}" + in_channels_str - - # -------------------------- Count total edge properties - n_edge_properties = sum( - p.encoder.get_encoding_length() - for p in self.properties - if isinstance(p, BondProperty) - ) - edge_dim_str = "" - if self.pad_edge_features: - n_edge_properties += self.pad_edge_features - edge_dim_str += f" (with {self.pad_edge_features} padded random values from {self.distribution} distribution)" - - edge_dim_str = f"edge_dim: {n_edge_properties}" + edge_dim_str - - rank_zero_info( - f"Finished loading dataset from properties.\nEncoding lengths: {prop_lengths}\n" - f"Use following values for given parameters for model configuration: \n\t" - f"{in_channels_str} \n\t" - f"{edge_dim_str} \n\t" - f"n_molecule_properties: {sum(p.encoder.get_encoding_length() for p in self.properties if isinstance(p, MoleculeProperty))}" - ) - - return base_df[base_data[0].keys()].to_dict("records") - - -class GraphPropAsPerNodeType(DataPropertiesSetter, ABC): - def __init__(self, properties=None, transform=None, **kwargs): - super().__init__(properties, transform, **kwargs) - # Sort properties so that AllNodeTypeProperty instances come first, rest of the properties order remain same - first = self._sort_properties( - [prop for prop in self.properties if isinstance(prop, AllNodeTypeProperty)] - ) - rest = self._sort_properties( - [ - prop - for prop in self.properties - if not isinstance(prop, AllNodeTypeProperty) - ] - ) - self.properties = first + rest - print( - "Properties are sorted so that `AllNodeTypeProperty` properties are first in sequence and rest of the order remains same\n", - f"Data module uses these properties (ordered): {', '.join([str(p) for p in self.properties])}", - ) - - def load_processed_data( - self, kind: Optional[str] = None, filename: Optional[str] = None - ) -> list[dict]: - """ - Load dataset and merge cached properties into base features. - - Args: - filename: The path to the file to load. - - Returns: - List of data entries, each a dictionary. - """ - base_data = super().load_processed_data(kind, filename) - base_df = pd.DataFrame(base_data) - base_df["ident"] = base_df["ident"].astype(str) - props_categories = { - "AllNodeTypeProperties": [], - "FGNodeTypeProperties": [], - "AtomNodeTypeProperties": [], - "GraphNodeTypeProperties": [], - "BondProperties": [], - } - n_atom_node_properties, n_fg_node_properties = 0, 0 - n_bond_properties, n_graph_node_properties = 0, 0 - prop_lengths = [] - for prop in self.properties: - prop_length = prop.encoder.get_encoding_length() - prop_name = prop.name - prop_lengths.append((prop_name, prop_length)) - if isinstance(prop, AllNodeTypeProperty): - n_atom_node_properties += prop_length - n_fg_node_properties += prop_length - n_graph_node_properties += prop_length - props_categories["AllNodeTypeProperties"].append(prop_name) - elif isinstance(prop, FGNodeTypeProperty): - n_fg_node_properties += prop_length - props_categories["FGNodeTypeProperties"].append(prop_name) - elif isinstance(prop, AtomNodeTypeProperty): - n_atom_node_properties += prop_length - props_categories["AtomNodeTypeProperties"].append(prop_name) - elif isinstance(prop, BondProperty): - n_bond_properties += prop_length - props_categories["BondProperties"].append(prop_name) - elif isinstance(prop, MoleculeProperty): - # molecule props will be used as graph node props - n_graph_node_properties += prop_length - props_categories["GraphNodeTypeProperties"].append(prop_name) - else: - raise TypeError(f"Unsupported property type: {type(prop).__name__}") - - n_node_properties = max( - n_atom_node_properties, n_fg_node_properties, n_graph_node_properties - ) - rank_zero_info( - f"\nFinished loading dataset from properties.\nEncoding lengths: {prop_lengths}\n\n" - f"Properties Categories:\n{pformat(props_categories)}\n\n" - f"n_atom_node_properties: {n_atom_node_properties}, " - f"n_fg_node_properties: {n_fg_node_properties}, " - f"n_bond_properties: {n_bond_properties}, " - f"n_graph_node_properties: {n_graph_node_properties}\n\n" - f"Use following values for given parameters for model configuration: \n\t" - f"in_channels: {n_node_properties}, edge_dim: {n_bond_properties}, n_molecule_properties: 0\n" - ) - - for property in self.properties: - rank_zero_info(f"Loading property {property.name}...") - property_data = torch.load( - self.get_property_path(property), weights_only=False - ) - for entry in property_data: - entry[property.name] = property.encoder.decompress(entry[property.name]) - if len(property_data[0][property.name].shape) > 1: - property.encoder.set_encoding_length( - property_data[0][property.name].shape[1] - ) - - property_df = pd.DataFrame(property_data) - property_df["ident"] = property_df["ident"].astype(str) - property_df.rename( - columns={property.name: f"{property.name}"}, inplace=True - ) - base_df = base_df.merge(property_df, on="ident", how="left") - - base_df["features"] = base_df.apply( - lambda row: self._merge_props_into_base( - row, - max_len_node_properties=n_node_properties, - ), - axis=1, - ) - - # apply transformation, e.g. masking for pretraining task - if self.transform is not None: - base_df["features"] = base_df["features"].apply(self.transform) - - return base_df[base_data[0].keys()].to_dict("records") - - def _merge_props_into_base( - self, row: pd.Series, max_len_node_properties: int - ) -> GeomData: - """ - Merge encoded molecular properties into the GeomData object. - - Args: - row: A dictionary containing 'features' and encoded properties. - - Returns: - A GeomData object with merged features. - """ - geom_data = row["features"] - if geom_data is None: - return None - if isinstance(geom_data, tuple): - geom_data = geom_data[ - 0 - ] # ignore additional returned data from _read_data (e.g. augmented molecule dict) - assert isinstance(geom_data, GeomData) - - is_atom_node = geom_data.is_atom_node - assert is_atom_node is not None, "`is_atom_node` must be set in the geom_data" - is_graph_node = geom_data.is_graph_node - assert is_graph_node is not None, "`is_graph_node` must be set in the geom_data" - - is_fg_node = ~is_atom_node & ~is_graph_node - num_nodes = geom_data.x.size(0) - edge_attr = geom_data.edge_attr - - # Initialize node feature matrix - assert max_len_node_properties is not None, ( - "Maximum len of node properties should not be None" - ) - x = torch.zeros((num_nodes, max_len_node_properties)) - - # Track column offsets for each node type - atom_offset, fg_offset, graph_offset = 0, 0, 0 - - for property in self.properties: - property_values = row[f"{property.name}"].to(dtype=torch.float32) - if isinstance(property_values, torch.Tensor): - if len(property_values.size()) == 0: - property_values = property_values.unsqueeze(0) - if len(property_values.size()) == 1: - property_values = property_values.unsqueeze(1) - else: - property_values = torch.zeros( - (0, property.encoder.get_encoding_length()) - ) - - enc_len = property_values.shape[1] - # -------------- Node properties --------------- - if isinstance(property, AllNodeTypeProperty): - x[:, atom_offset : atom_offset + enc_len] = property_values - atom_offset += enc_len - fg_offset += enc_len - graph_offset += enc_len - - elif isinstance(property, AtomNodeTypeProperty): - x[is_atom_node, atom_offset : atom_offset + enc_len] = property_values[ - is_atom_node - ] - atom_offset += enc_len - - elif isinstance(property, FGNodeTypeProperty): - x[is_fg_node, fg_offset : fg_offset + enc_len] = property_values[ - is_fg_node - ] - fg_offset += enc_len - - elif isinstance(property, MoleculeProperty): - x[is_graph_node, graph_offset : graph_offset + enc_len] = ( - property_values - ) - graph_offset += enc_len - - # ------------- Bond Properties -------------- - elif isinstance(property, BondProperty): - # Concat/Duplicate properties values for undirected graph as `edge_index` has first src to tgt edges, then tgt to src edges - edge_attr = torch.cat( - [edge_attr, torch.cat([property_values, property_values], dim=0)], - dim=1, - ) - else: - raise TypeError(f"Unsupported property type: {type(property).__name__}") - - total_used_columns = max(atom_offset, fg_offset, graph_offset) - assert total_used_columns <= max_len_node_properties, ( - f"Used {total_used_columns} columns, but max allowed is {max_len_node_properties}" - ) - - return GeomData( - x=x, - edge_index=geom_data.edge_index, - edge_attr=edge_attr, - molecule_attr=torch.empty((1, 0)), # empty as not used for this class - is_atom_node=is_atom_node, - is_fg_node=is_fg_node, - is_graph_node=is_graph_node, - ) - - def _prediction_merge_props_into_base_wrapper( - self, row: pd.Series | dict, model_hparams: Optional[dict] = None - ) -> GeomData: - """ - Wrapper to merge properties into base features for prediction. - - Args: - row: A dictionary or pd.Series containing 'features' and encoded properties. - Returns: - A GeomData object with merged features. - """ - if ( - model_hparams is None - or "in_channels" not in model_hparams["config"] - or model_hparams["config"]["in_channels"] is None - ): - raise ValueError( - f"model_hparams must be provided for data class: {self.__class__.__name__}" - f" which should contain 'in_channels' key with valid value in 'config' dictionary." - ) - max_len_node_properties = int(model_hparams["config"]["in_channels"]) - return self._merge_props_into_base(row, max_len_node_properties) - - class ChEBI50_StaticGNI(DataPropertiesSetter, ChEBIOver50): READER = RandomFeatureInitializationReader From 2c2607310addeadf8bcebd9e82ee60ff8057ff66 Mon Sep 17 00:00:00 2001 From: aditya0by0 Date: Tue, 27 Oct 2026 08:17:05 +0100 Subject: [PATCH 51/58] pretraining scripts for zinc_msp_pretraining and final-chebi.sh have been added. The .gitignore file has been updated to include the new pretraining scripts. --- .gitignore | 1 + configs/data/augmented/final/final-chebi.sh | 110 ++++++++++++++++++ .../pretraining/zinc_msp_pretraining.sh | 65 +++++++++++ 3 files changed, 176 insertions(+) create mode 100644 configs/data/augmented/final/final-chebi.sh create mode 100644 configs/data/augmented/pretraining/zinc_msp_pretraining.sh diff --git a/.gitignore b/.gitignore index 491079d..cc9e754 100644 --- a/.gitignore +++ b/.gitignore @@ -179,3 +179,4 @@ electra_pretrained.ckpt !configs/data/augmented/aug-ablation/*.sh !.vscode/launch.json !configs/data/augmented/final/*.sh +!configs/data/augmented/pretraining/*.sh diff --git a/configs/data/augmented/final/final-chebi.sh b/configs/data/augmented/final/final-chebi.sh new file mode 100644 index 0000000..4b2cb84 --- /dev/null +++ b/configs/data/augmented/final/final-chebi.sh @@ -0,0 +1,110 @@ +#!/bin/bash + +#SBATCH --time=48:00:00 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=32 +#SBATCH --threads-per-core=1 +#SBATCH --mem=256000 + +#SBATCH --partition=gpu +#SBATCH --gres=gpu:1 +#SBATCH --constraint="A100|H100.80gb" + +# 3 experiments, maximum 3 running at the same time +#SBATCH --array=0-2%3 + +#SBATCH --job-name=chebi_gatv2_final + +# Separate log files for every array task +#SBATCH --output=chebi_gatv2_final_%A_%a.out +#SBATCH --error=chebi_gatv2_final_%A_%a.err + + +# ============================================================ +# Shell settings +# ============================================================ +set -x +set -euo pipefail +# ============================================================ +# Define seeds +# ============================================================ + +SEEDS=(0 42 12345) + +if [[ -z "${SEEDS[$SLURM_ARRAY_TASK_ID]}" ]]; then + echo "ERROR: Invalid array task ID: $SLURM_ARRAY_TASK_ID" + exit 1 +fi + +SEED=${SEEDS[$SLURM_ARRAY_TASK_ID]} +RUN_NAME="chebi_gatv2_final_${SEED}" + +# ============================================================ +# Print job information +# ============================================================ + +echo "============================================================" +echo "Job ID: $SLURM_JOB_ID" +echo "Array Job ID: $SLURM_ARRAY_JOB_ID" +echo "Array Task ID: $SLURM_ARRAY_TASK_ID" +echo "Seed: $SEED" +echo "Run name: $RUN_NAME" +echo "Node: $(hostname)" +echo "Date: $(date)" +echo "============================================================" + +echo "GPU information:" +nvidia-smi + +echo "============================================================" + +# ============================================================ +# Temporary directory +# ============================================================ + +export TMPDIR=/home/staff/a/akhedekar/atmp_dir/ + +# ============================================================ +# Activate Python environment +# ============================================================ + +source /home/staff/a/akhedekar/python-chebai-graph/.venv/bin/activate + +# ============================================================ +# Set working directory +# ============================================================ + +CHEBAI_DIR="/home/staff/a/akhedekar/python-chebai" +cd "$CHEBAI_DIR" +export SSL_CERT_FILE=$(python -m certifi) + +# ============================================================ +# Run training +# ============================================================ + +python -m chebai fit \ + --trainer=configs/training/default_trainer.yml \ + --trainer.logger=configs/training/wandb_logger.yml \ + --model=../python-chebai-graph/configs/model/augmented/pooling/amg_pool/gat.yml \ + --model.train_metrics=configs/metrics/micro-macro-f1.yml \ + --model.test_metrics=configs/metrics/micro-macro-f1.yml \ + --model.val_metrics=configs/metrics/micro-macro-f1.yml \ + --data=../python-chebai-graph/configs/data/augmented/final/chebi50_final_augmented.yml \ + --data.init_args.batch_size=64 \ + --data.init_args.num_workers=10 \ + --data.init_args.chebi_version=252 \ + --model.pass_loss_kwargs=false \ + --trainer.min_epochs=200 \ + --trainer.max_epochs=200 \ + --model.criterion=configs/loss/bce_weighted.yml \ + --trainer.logger.init_args.name="$RUN_NAME" \ + --data.init_args.splits_file_path=data/chebi_v252/ChEBI50/processed/splits.csv \ + --seed_everything="$SEED" \ + --trainer.logger.init_args.tags='["augmented_paper", "chebi", "final:gnwa_fgwa_wfge"]' \ + +echo "============================================================" +echo "Training finished" +echo "Seed: $SEED" +echo "Node: $(hostname)" +echo "Date: $(date)" +echo "============================================================" diff --git a/configs/data/augmented/pretraining/zinc_msp_pretraining.sh b/configs/data/augmented/pretraining/zinc_msp_pretraining.sh new file mode 100644 index 0000000..0c3c2e9 --- /dev/null +++ b/configs/data/augmented/pretraining/zinc_msp_pretraining.sh @@ -0,0 +1,65 @@ +#!/bin/bash + +#SBATCH --time=48:00:00 +#SBATCH --ntasks=1 +#SBATCH --cpus-per-task=32 +#SBATCH --threads-per-core=1 +#SBATCH --mem=256000 +#SBATCH --partition=gpu +#SBATCH --gres=gpu:3 +#SBATCH --constraint="A100|H100.80gb" +#SBATCH --job-name=zinc +#SBATCH --output=zinc.out +#SBATCH --error=zinc.err + + +# ============================================================ +# Shell settings +# ============================================================ +set -x +# ============================================================ +# Define seeds +# ============================================================ + +echo "GPU information:" +nvidia-smi + +echo "============================================================" + +# ============================================================ +# Temporary directory +# ============================================================ + +export TMPDIR=/home/staff/a/akhedekar/atmp_dir/ + +# ============================================================ +# Activate Python environment +# ============================================================ + +source /home/staff/a/akhedekar/python-chebai-graph/.venv/bin/activate + +# ============================================================ +# Set working directory +# ============================================================ + +CHEBAI_DIR="/home/staff/a/akhedekar/python-chebai" +cd "$CHEBAI_DIR" +export SSL_CERT_FILE=$(python -m certifi) + +# ============================================================ +# Run training +# ============================================================ + +python -m chebai fit \ + --trainer=configs/training/zinc_pretrainer.yml \ + --trainer.limit_val_batches=0 \ + --trainer.logger=configs/training/wandb_logger.yml \ + --trainer.logger.init_args.tags='["augmented_paper","zinc:pretraining","final:gnwa_fgwa_wfge"]' \ + --trainer.logger.init_args.name=zinc15pretrain_gat_amg_augmented \ + --model=../python-chebai-graph/configs/model/pretraining/gat_msp.yml \ + --data=../python-chebai-graph/configs/data/augmented/pretraining/zinc15.yml \ + --data.init_args.batch_size=32 \ + --model.pass_loss_kwargs=true \ + --data.init_args.num_workers=10 \ + --trainer.logger.init_args.name=zinc15pretrain_gat_amg_augmented \ + --data.init_args.splits_file_path=data/Zinc15:Pretrain/Zinc15/processed/splits.csv \ From 6d5f8912eab5c6b9cff83e7bb1a554d5de63c00f Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Tue, 27 Oct 2026 08:22:38 +0100 Subject: [PATCH 52/58] update readme --- README.md | 97 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 97 insertions(+) diff --git a/README.md b/README.md index b2a8ec0..52937cf 100644 --- a/README.md +++ b/README.md @@ -76,6 +76,103 @@ The list can be found in the `configs/data/chebi50_graph_properties.yml` file. python -m chebai fit --trainer=configs/training/default_trainer.yml --trainer.logger=configs/training/csv_logger.yml --model=../python-chebai-graph/configs/model/gnn_res_gated.yml --model.train_metrics=configs/metrics/micro-macro-f1.yml --model.test_metrics=configs/metrics/micro-macro-f1.yml --model.val_metrics=configs/metrics/micro-macro-f1.yml --data=../python-chebai-graph/configs/data/chebi50_graph_properties.yml --data.init_args.batch_size=128 --trainer.accumulate_grad_batches=4 --data.init_args.num_workers=10 --model.pass_loss_kwargs=false --data.init_args.chebi_version=241 --trainer.min_epochs=200 --trainer.max_epochs=200 --model.criterion=configs/loss/bce_weighted.yml ``` +## Experiment Scripts + +The scripts below reproduce the augmented-graph experiments introduced in this +work. They expect both repositories to be checked out as described in +[Recommended Folder Structure](#recommended-folder-structure), and they run +`python -m chebai` from the `python-chebai` checkout. Update the paths, +environment activation, and logging configuration in the scripts for a +different machine. + +### Augmentation ablations + +Each ablation script evaluates the same ten graph-construction variants: +`baseline`, `FGN`, `FGN+E`, `FGN+WGN`, `FGN+E+WGN`, `WGN`, +`gn_wall_fgwa_nfge`, `gn_wall_fgwa_wfge`, `gnwa_fgwa_nfge`, and +`gnwa_fgwa_wfge`. The experiments use seeds `0`, `42`, and `12345`, so each +dataset runs 30 configurations in total. The variants compare whether +functional-group nodes (FGN), functional-group edges (FGE), and the graph node +(WGN) are included, and whether their connections use the corresponding +weighted or unweighted schemes. + +- [`aug-ablation-chebi.sh`](configs/data/augmented/aug-ablation/aug-ablation-chebi.sh) + is a Slurm job-array runner for ChEBI50. It schedules the 10 configurations + across the three seeds, trains the baseline GAT, and uses the ChEBI v252 + split. +- [`aug-ablation-tox21.sh`](configs/data/augmented/aug-ablation/aug-ablation-tox21.sh) + runs the same 30-experiment matrix for Tox21 serially on one local GPU. It + reports micro/macro F1 and ROC-AUC using the Tox21 split. +- [`aug-ablation-toxcast.sh`](configs/data/augmented/aug-ablation/aug-ablation-toxcast.sh) + runs the corresponding ToxCast matrix serially on one local GPU. It uses + ToxCast-specific input dimensions and learning rates and reports micro/macro + F1 and ROC-AUC. + +Run a local ablation script from this repository with, for example: + +```bash +bash configs/data/augmented/aug-ablation/aug-ablation-tox21.sh +``` + +The ChEBI script is intended for Slurm and should be submitted with +`sbatch`. The Tox21 and ToxCast scripts run every seed and configuration in +sequence, rather than requiring Slurm. + +### MoleculeNet datasets + +[MoleculeNet](https://moleculenet.org/) is a collection of standard molecular +machine-learning benchmarks. These datasets cover binary classification tasks +with different label counts, class imbalance, and chemical domains, making +them useful for testing whether an architecture generalizes beyond ChEBI50. +The augmented-graph runner currently evaluates six MoleculeNet datasets: +**BBBP**, **BACE**, **SIDER**, **ClinTox**, **Tox21**, and **ToxCast**. Their +dataset configurations are in +`configs/data/augmented/final/_final_augmented.yml`. + +[`molnet.sh`](configs/data/augmented/final/molnet.sh) runs each of the six +datasets for seeds `0`, `42`, and `12345` on one local GPU. For each dataset +and seed it: + +1. trains the augmented atom-motif-graph pooling GAT; +2. selects the best validation checkpoint using ROC-AUC; +3. evaluates that checkpoint on the test split; and +4. writes per-run metrics and checkpoint paths to + `logs/molnet_test_results.tsv`, followed by mean and sample standard + deviation summaries. + +The script selects binary or multilabel trainers and metrics per dataset and +uses dataset-specific dropout and learning rates. Run it with: + +```bash +bash configs/data/augmented/final/molnet.sh +``` + +### Final ChEBI50 training + +[`final-chebi.sh`](configs/data/augmented/final/final-chebi.sh) is a Slurm +job-array script that trains the final augmented-graph GAT on ChEBI50 v252 for +the three seeds. It uses the `gnwa_fgwa_wfge` configuration, the +atom-motif-graph pooling architecture, weighted binary cross-entropy, and the +standard micro/macro F1 metrics. Submit it with: + +```bash +sbatch configs/data/augmented/final/final-chebi.sh +``` + +### ZINC15 pretraining + +[`zinc_msp_pretraining.sh`](configs/data/augmented/pretraining/zinc_msp_pretraining.sh) +trains the augmented GAT with the masked-subgraph prediction (MSP) +pretraining objective on ZINC15. It uses the dedicated ZINC pretrainer, +`gat_msp.yml`, and the augmented ZINC15 data configuration. Validation batches +are disabled because this is a pretraining run; the resulting checkpoint can +then be used to initialize downstream augmented-graph experiments. The script +requests three GPUs on Slurm and is submitted with: + +```bash +sbatch configs/data/augmented/pretraining/zinc_msp_pretraining.sh +``` + ## Augmented Graphs _See thesis related to this work [here](https://www.uni-osnabrueck.de/fileadmin/informatik/Arbeitsgruppen/Hybride_KI/mt_aditya_khedekar.pdf)_. From 90ee70e3157f6c456e6a3bd9cd7c8d742d23b6e0 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Tue, 27 Oct 2026 08:29:16 +0100 Subject: [PATCH 53/58] update model hyperparams section --- README.md | 84 +++++++++++++++++++++++++++++++++++++------------------ 1 file changed, 57 insertions(+), 27 deletions(-) diff --git a/README.md b/README.md index 52937cf..e9f35fa 100644 --- a/README.md +++ b/README.md @@ -196,45 +196,75 @@ python -m chebai fit --trainer=configs/training/default_trainer.yml --trainer.lo ### Model Hyperparameters -#### **GAT Architecture** - -To use a GAT-based model, choose **one** of the following configs: +The model paths below are relative to the `python-chebai` working directory +used by the training commands. In this repository they correspond to the +following configuration files: -- **Standard Pooling**: `--model=../python-chebai-graph/configs/model/gat.yml` - > Standard pooling sums the learned representations from all the nodes to produce a single representation which is used for classification. - -- **Atom-Augmented Node Pooling**: `--model=../python-chebai-graph/configs/model/gat_aug_aagpool.yml` - > With this pooling stratergy, the learned representations are first separated into **two distinct sets**: those from atom nodes and those from all artificial nodes (both functional groups and the graph node). The representations within each set are aggregated separately (using summation) to yield two distinct single vectors. These two resulting vectors are then concatenated before being passed to the classification layer. +#### **GAT Architecture** -- **Atom–Motif–Graph Node Pooling**: `--model=../python-chebai-graph/configs/model/gat_aug_amgpool.yml` - > This approach employs a finer granularity of separation, distinguishing learned representations into **three distinct sets**: atom nodes, Functional Group (FG) nodes, and the single graph node. Summation is performed separately on the atom node set and the FG node set, yielding two vectors. These two vectors are then concatenated along with the single vector corresponding to the graph node before the final linear layer. +Choose one of the following GAT configurations: + +- **Standard pooling**: + [`configs/model/baselines/gat.yml`](configs/model/baselines/gat.yml) + (`--model=../python-chebai-graph/configs/model/baselines/gat.yml`). + `GATGraphPred` pools all node representations into one graph + representation. This configuration is used for the ablation baselines. +- **Atom-augmented pooling**: + [`configs/model/augmented/pooling/aa_pool/gat.yml`](configs/model/augmented/pooling/aa_pool/gat.yml). + `GATAAPoolGraphPred` aggregates atom-node representations separately from + artificial-node representations, then concatenates the two pooled vectors. +- **Atom-motif-graph pooling**: + [`configs/model/augmented/pooling/amg_pool/gat.yml`](configs/model/augmented/pooling/amg_pool/gat.yml). + `GATAMGPoolGraphPred` separately aggregates atom nodes, functional-group + nodes, and the graph node before classification. This is the configuration + used by the final augmented experiments. +- **Augmented graphs without separate pooling**: + [`configs/model/augmented/pooling/no_pooling/gat.yml`](configs/model/augmented/pooling/no_pooling/gat.yml). + This uses `GATGraphPred` with augmented node and edge features while pooling + all nodes together. + +The augmented GAT configurations use 203 node features and 12 edge features; +the baseline configuration uses 161 node features and 8 edge features. All +current GAT configurations default to four message-passing layers, 256 hidden +channels, 512 output channels, eight attention heads, and GATv2 +(`v2: True`). The number of heads must divide the output or hidden channel +size. These values can be overridden, for example: -#### GAT-specific hyperparameters +```bash +--model.config.num_layers=5 +--model.config.heads=4 +--model.config.v2=False +--model.config.dropout=0.1 +``` -- **Number of message-passing layers**: `--model.config.num_layers=5` (default: 4) -- **Attention heads**: `--model.config.heads=4` (default: 8) - > **Note**: The number of heads should be divisible by the output channels (or hidden channels if output channels are not specified). +Set `v2=False` to use the original `GATConv`; the default `v2=True` uses +`GATv2Conv`, which provides dynamic rather than static attention. For details, +see the [GATv2 paper](https://arxiv.org/abs/2105.14491). -- **To Use different GAT versions**: - - **Use GAT**: `--model.config.v2=False` +#### **ResGated Architecture** - - **Use GATv2**: `--model.config.v2=True` (__default__) - > **Note**: GATv2 addresses the limitation of static attention in GAT by introducing a dynamic attention mechanism. For further details, please refer to the [original GATv2 paper](https://arxiv.org/abs/2105.14491). +Choose one of the following ResGated configurations: -#### **ResGated Architecture** +- **Atom-motif-graph pooling**: + [`configs/model/res_aug_amgpool.yml`](configs/model/res_aug_amgpool.yml). +- **Atom-augmented pooling**: + [`configs/model/res_aug_aapool.yml`](configs/model/res_aug_aapool.yml). -To use a ResGated GNN model, choose **one** of the following configs: +Both augmented ResGated configurations use 203 node features, 12 edge +features, four message-passing layers, 256 hidden channels, and 512 output +channels. The ResGated convolution does not expose GAT attention-head or +`v2` settings. -- **Atom–Motif–Graph Node Pooling**: `--model=../python-chebai-graph/configs/model/res_aug_amgpool.yml` -- **Atom-Augmented Node Pooling**: `--model=../python-chebai-graph/configs/model/res_aug_aagpool.yml` -- **Standard Pooling**: `--model=../python-chebai-graph/configs/model/resgated.yml` +#### **Shared hyperparameters** -#### **Common Hyperparameters** +These settings are available for both GAT and ResGated model configurations: -These can be used for both GAT and ResGated architectures: +- **Dropout**: `--model.config.dropout=0.1` (default: `0`) +- **Number of final linear layers**: `--model.n_linear_layers=2` + (default: `1`) -- **Dropout**: `--model.config.dropout=0.1` (default: 0) -- **Number of final linear layers**: `--model.n_linear_layers=2` (default: 1) +The complete defaults for each architecture are defined directly in the +linked YAML files above. ## Random Node Initialization From f5aa7d43fbf98dff765d988232752545f37dc5fc Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Tue, 27 Oct 2026 08:36:35 +0100 Subject: [PATCH 54/58] update augmented --- README.md | 29 +++++++++++++++++++++++++---- 1 file changed, 25 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index e9f35fa..05f1f5d 100644 --- a/README.md +++ b/README.md @@ -174,14 +174,35 @@ sbatch configs/data/augmented/pretraining/zinc_msp_pretraining.sh ``` ## Augmented Graphs -_See thesis related to this work [here](https://www.uni-osnabrueck.de/fileadmin/informatik/Arbeitsgruppen/Hybride_KI/mt_aditya_khedekar.pdf)_. + Graph Neural Networks (GNNs) often fail to explicitly leverage the chemically meaningful substructures present within molecules (i.e. **functional groups (FGs)**). To make this implicit information explicitly accessible to GNNs, we augment molecular graphs with **artificial nodes** that represent these substructures. The resulting graph are referred to as **augmented graphs**. > Note: Rings are also treated as functional groups in our work. -In these augmented graphs, each functional group node is connected to the atoms that constitute the group. Additionally, two functional group nodes are connected if any atom belonging to one group shares a bond with an atom from the other group. We further introduce a **graph node**, an extra node connected to all functional group nodes. +In these augmented graphs, each functional group node is connected to the atoms +that constitute the group. Additionally, two functional group nodes are +connected if any atom belonging to one group shares a bond with an atom from +the other group. We also introduce a **graph node**, which is connected to all +atom nodes. Rings are treated as functional groups. + +### Best configuration: `GNWA_FGWA_WFGE` + +Among the connection schemes evaluated, **GNWA_FGWA_WFGE** delivered the +strongest performance. The name describes the complete augmented topology: + +- **GNWA** (*Graph Node With Atom connections*): the graph node is connected + to every atom node. +- **FGWA** (*Functional Group Nodes With Atom connections*): every functional + group node is connected to the atoms that belong to that group. +- **WFGE** (*With Functional Group Edges*): functional group nodes are + connected when their member atoms are connected by a molecular bond. -Among all the connection schemes we evaluated, this configuration delivered the strongest performance. We denote it using the abbreviation **WFG_WFGE_WGN** in our work and is shown in below figure. +This configuration is used by the final ChEBI50 and MoleculeNet experiments. +Its ChEBI50 data definition is +[`chebi50_final_augmented.yml`](configs/data/augmented/final/chebi50_final_augmented.yml), +and its atom-motif-graph pooling model is +[`amg_pool/gat.yml`](configs/model/augmented/pooling/amg_pool/gat.yml). +The topology is illustrated below. mol_to_aug_mol @@ -191,7 +212,7 @@ Among all the connection schemes we evaluated, this configuration delivered the Below is the command for the model and data configuration that achieved the best classification performance using augmented graphs. ```bash -python -m chebai fit --trainer=configs/training/default_trainer.yml --trainer.logger=configs/training/wandb_logger.yml --model=../python-chebai-graph/configs/model/gat_aug_amgpool.yml --model.train_metrics=configs/metrics/micro-macro-f1.yml --model.test_metrics=configs/metrics/micro-macro-f1.yml --model.val_metrics=configs/metrics/micro-macro-f1.yml --data=../python-chebai-graph/configs/data/chebi50_aug_prop_as_per_node.yml --data.init_args.batch_size=128 --trainer.accumulate_grad_batches=4 --data.init_args.num_workers=10 --model.pass_loss_kwargs=false --data.init_args.chebi_version=241 --trainer.min_epochs=200 --trainer.max_epochs=200 --model.criterion=configs/loss/bce_weighted.yml --trainer.logger.init_args.name=gatv2_amg_s0 +python -m chebai fit --trainer=configs/training/default_trainer.yml --trainer.logger=configs/training/wandb_logger.yml --model=../python-chebai-graph/configs/model/augmented/pooling/amg_pool/gat.yml --model.train_metrics=configs/metrics/micro-macro-f1.yml --model.test_metrics=configs/metrics/micro-macro-f1.yml --model.val_metrics=configs/metrics/micro-macro-f1.yml --data=../python-chebai-graph/configs/data/augmented/properties/gnwa_fgwa_wfge/chebi50_all_props.yml --data.init_args.batch_size=64 --trainer.accumulate_grad_batches=1 --data.init_args.num_workers=10 --model.pass_loss_kwargs=false --data.init_args.chebi_version=252 --trainer.min_epochs=200 --trainer.max_epochs=200 --model.criterion=configs/loss/bce_unweighted.yml --trainer.logger.init_args.name=gnwa_fgwa_wfge_final_42 --model.init_args.optimizer_kwargs.lr=0.002 --seed_everything=42 --data.init_args.splits_file_path=data/chebi_v252/ChEBI50/processed/splits.csv --trainer.logger.init_args.tags=["augmented_paper","final:gnwa_fgwa_wfge"] ``` ### Model Hyperparameters From dcb90dfcf9b9b8a668bbd6ba78bde0ac70d62fcc Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Tue, 27 Oct 2026 08:40:29 +0100 Subject: [PATCH 55/58] other info --- README.md | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/README.md b/README.md index 05f1f5d..399a4d7 100644 --- a/README.md +++ b/README.md @@ -147,6 +147,14 @@ uses dataset-specific dropout and learning rates. Run it with: bash configs/data/augmented/final/molnet.sh ``` +For comparisons with published MoleculeNet results, use scaffold splits for +BBBP and BACE and random splits for Tox21, ToxCast, SIDER, and ClinTox. The +reported benchmark metric is the mean test ROC-AUC over three independent +runs; the runner also records the F1 metrics selected for each dataset. The +MoleculeNet experiments can be run with or without initializing the model from +the ZINC15 MSP checkpoint, allowing the contribution of pretraining to be +measured separately. + ### Final ChEBI50 training [`final-chebi.sh`](configs/data/augmented/final/final-chebi.sh) is a Slurm @@ -185,6 +193,12 @@ connected if any atom belonging to one group shares a bond with an atom from the other group. We also introduce a **graph node**, which is connected to all atom nodes. Rings are treated as functional groups. +Functional-group nodes encode both the presence and the identity of the +detected group, rather than representing an unlabeled fragment. The graph node +acts as a global or master node: it provides a shortcut for information +exchange between distant parts of a molecule and collects a whole-molecule +representation. + ### Best configuration: `GNWA_FGWA_WFGE` Among the connection schemes evaluated, **GNWA_FGWA_WFGE** delivered the @@ -204,6 +218,13 @@ and its atom-motif-graph pooling model is [`amg_pool/gat.yml`](configs/model/augmented/pooling/amg_pool/gat.yml). The topology is illustrated below. +The final augmented representation initializes atoms, bonds, functional-group +nodes, and the graph node with chemically meaningful, node-type-specific +properties. This is different from the isolated topology ablations, where +artificial nodes and added edges use zero initialization so that the effect of +connectivity can be measured independently of feature initialization and +type-specific pooling. + mol_to_aug_mol
@@ -262,6 +283,12 @@ Set `v2=False` to use the original `GATConv`; the default `v2=True` uses `GATv2Conv`, which provides dynamic rather than static attention. For details, see the [GATv2 paper](https://arxiv.org/abs/2105.14491). +The augmentation and property initialization do not add trainable parameters +to the core GNN. Type-specific pooling can increase the size of the +classification head, but the main performance gain is intended to come from +the augmented molecular representation and its chemically meaningful +features. + #### **ResGated Architecture** Choose one of the following ResGated configurations: From 493fe4c8d7182fea33e7a760b66a9b6bd215f41f Mon Sep 17 00:00:00 2001 From: Aditya Khedekar <65857172+aditya0by0@users.noreply.github.com> Date: Tue, 27 Oct 2026 08:42:37 +0100 Subject: [PATCH 56/58] Replace image in README with new graph illustration --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 399a4d7..cc30202 100644 --- a/README.md +++ b/README.md @@ -225,7 +225,7 @@ artificial nodes and added edges use zero initialization so that the effect of connectivity can be measured independently of feature initialization and type-specific pooling. -mol_to_aug_mol +mol graph drawio

From 4759b2dd9c50396232a244beb88b3e06f444b270 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Wed, 28 Oct 2026 19:44:17 +0100 Subject: [PATCH 57/58] add batch norm param --- chebai_graph/models/architectures/base.py | 1 + chebai_graph/models/architectures/gat.py | 1 + chebai_graph/models/architectures/gine.py | 2 ++ chebai_graph/models/architectures/resgated.py | 1 + configs/model/augmented/pooling/amg_pool/himol_gine.yml | 1 + 5 files changed, 6 insertions(+) diff --git a/chebai_graph/models/architectures/base.py b/chebai_graph/models/architectures/base.py index f24079a..4217761 100644 --- a/chebai_graph/models/architectures/base.py +++ b/chebai_graph/models/architectures/base.py @@ -78,6 +78,7 @@ def __init__(self, config: dict, **kwargs) -> None: self.out_channels = int(config["out_channels"]) self.edge_dim = int(config["edge_dim"]) # number of bond properties self.dropout = float(config.get("dropout", 0.0)) + self.norm = config.get("norm", None) print("----------Dropout set to: ", self.dropout) diff --git a/chebai_graph/models/architectures/gat.py b/chebai_graph/models/architectures/gat.py index 13dcabc..f3dae9f 100644 --- a/chebai_graph/models/architectures/gat.py +++ b/chebai_graph/models/architectures/gat.py @@ -42,6 +42,7 @@ def __init__(self, config: dict, **kwargs): heads=self.heads, v2=self.v2, act=self.activation, + norm=self.norm, **local_kwargs, ) diff --git a/chebai_graph/models/architectures/gine.py b/chebai_graph/models/architectures/gine.py index 95696dc..b7087d7 100644 --- a/chebai_graph/models/architectures/gine.py +++ b/chebai_graph/models/architectures/gine.py @@ -148,6 +148,7 @@ def __init__(self, config: dict[str, Any], **kwargs: Any): dropout=self.dropout, edge_dim=self.edge_dim, train_eps=self.train_eps, + norm=self.norm, act=ReLU(), ) else: @@ -160,6 +161,7 @@ def __init__(self, config: dict[str, Any], **kwargs: Any): edge_dim=self.edge_dim, train_eps=self.train_eps, act=self.activation, + norm=self.norm, ) def forward(self, batch: dict[str, Any]) -> Tensor: diff --git a/chebai_graph/models/architectures/resgated.py b/chebai_graph/models/architectures/resgated.py index 521a853..a9e4663 100644 --- a/chebai_graph/models/architectures/resgated.py +++ b/chebai_graph/models/architectures/resgated.py @@ -65,6 +65,7 @@ def __init__(self, config: dict[str, Any], **kwargs: Any): num_layers=self.num_layers, edge_dim=self.edge_dim, act=self.activation, + norm=self.norm, ) def forward(self, batch: dict[str, Any]) -> Tensor: diff --git a/configs/model/augmented/pooling/amg_pool/himol_gine.yml b/configs/model/augmented/pooling/amg_pool/himol_gine.yml index 4a22416..10559ac 100644 --- a/configs/model/augmented/pooling/amg_pool/himol_gine.yml +++ b/configs/model/augmented/pooling/amg_pool/himol_gine.yml @@ -10,5 +10,6 @@ init_args: edge_dim: 12 # number of bond properties train_eps: true use_himol_model: true + norm: batch_norm n_linear_layers: 2 lin_hidden_dim: 256 From 399bade4e82f3bce35672bd78a1aa5c3c5e38e99 Mon Sep 17 00:00:00 2001 From: aditya0b0 Date: Mon, 2 Nov 2026 09:26:52 +0100 Subject: [PATCH 58/58] script for amg analysis --- chebai_graph/models/architectures/base.py | 1 + chebai_graph/models/pooling.py | 23 +- scripts/analyze_amgpool.py | 311 ++++++++++++++++++++++ 3 files changed, 325 insertions(+), 10 deletions(-) create mode 100644 scripts/analyze_amgpool.py diff --git a/chebai_graph/models/architectures/base.py b/chebai_graph/models/architectures/base.py index 4217761..15f3775 100644 --- a/chebai_graph/models/architectures/base.py +++ b/chebai_graph/models/architectures/base.py @@ -80,6 +80,7 @@ def __init__(self, config: dict, **kwargs) -> None: self.dropout = float(config.get("dropout", 0.0)) self.norm = config.get("norm", None) print("----------Dropout set to: ", self.dropout) + print("----------GNN Norm set to: ", self.norm) class GraphNetWrapper(GraphBaseNet, ABC): diff --git a/chebai_graph/models/pooling.py b/chebai_graph/models/pooling.py index a1ee442..df5d558 100644 --- a/chebai_graph/models/pooling.py +++ b/chebai_graph/models/pooling.py @@ -103,24 +103,22 @@ def _get_lin_seq_input_dim(self, gnn_out_dim: int) -> int: """ return gnn_out_dim + gnn_out_dim + gnn_out_dim - def forward(self, batch: dict) -> torch.Tensor: + def pool_representations( + self, batch: dict + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ - Forward pass for pooling node embeddings. + Return atom-, functional-group-, and graph-level pooled embeddings. Steps: 1. Identify graph, atom, and functional group nodes. 2. Aggregate embeddings for each node type separately. - 3. Concatenate: - - Atom nodes vector - - Functional group nodes vector - - Graph node vector - 4. Pass the concatenated vector through the linear sequential block. Args: batch (dict): Batch containing graph data and features. Returns: - torch.Tensor: Output tensor after pooling and linear transformation. + tuple[torch.Tensor, torch.Tensor, torch.Tensor]: Pooled embeddings + in atom, functional-group, and graph order. """ graph_data = batch["features"][0] assert isinstance(graph_data, GraphData) @@ -145,7 +143,12 @@ def forward(self, batch: dict) -> torch.Tensor: atoms_vec = scatter_add(atoms_embeddings, atoms_batch, dim=0) fg_nodes_vec = scatter_add(fg_nodes_embeddings, fg_nodes_batch, dim=0) - # Concatenate all - graph_vector = torch.cat([atoms_vec, fg_nodes_vec, graph_node_vec], dim=1) + return atoms_vec, fg_nodes_vec, graph_node_vec + + def forward(self, batch: dict) -> torch.Tensor: + """Pool the three representation levels and apply the classifier.""" + representations = self.pool_representations(batch) + # Keep this order aligned with the classifier input layout. + graph_vector = torch.cat(representations, dim=1) return self.lin_sequential(graph_vector) diff --git a/scripts/analyze_amgpool.py b/scripts/analyze_amgpool.py new file mode 100644 index 0000000..05f3b21 --- /dev/null +++ b/scripts/analyze_amgpool.py @@ -0,0 +1,311 @@ +#!/usr/bin/env python3 +"""Measure representation scales and classifier contributions in AMGPool models. + +Example (preferred; uses hyperparameters stored in the checkpoint): + python scripts/analyze_amgpool.py \ + --checkpoint path/to/model.ckpt \ + --output-dir analysis/bbbp + +Explicit model and data configs can also be supplied together when the +checkpoint does not contain Lightning hyperparameters. +""" + +from __future__ import annotations + +import argparse +import copy +import csv +import importlib +import json +from pathlib import Path +from typing import Any + +import torch +import yaml +from chebai.models.base import ChebaiBaseNet +from chebai.preprocessing.datasets.base import XYBaseDataModule +from lightning.pytorch.cli import instantiate_module + +LEVELS = ("atom", "motif", "graph") + + +def _import_class(path: str) -> type: + module_name, class_name = path.rsplit(".", 1) + return getattr(importlib.import_module(module_name), class_name) + + +def _instantiate(config: dict[str, Any]) -> Any: + class_path = config.get("class_path") + if not class_path: + raise ValueError("Expected a config with a 'class_path' entry.") + init_args = config.get("init_args", {}) + return _import_class(class_path)(**init_args) + + +def _load_yaml(path: Path) -> dict[str, Any]: + with path.open() as handle: + config = yaml.safe_load(handle) + if not isinstance(config, dict): + raise ValueError(f"{path} must contain a YAML mapping.") + return config + + +def _load_checkpoint(model: torch.nn.Module, path: Path) -> None: + checkpoint = torch.load(path, map_location="cpu", weights_only=False) + state_dict = checkpoint.get("state_dict", checkpoint) + if not isinstance(state_dict, dict): + raise ValueError(f"{path} does not contain a state dictionary.") + + model_state = model.state_dict() + candidates = [] + for prefix in ("", "model.", "module."): + candidates.append( + { + key.removeprefix(prefix): value + for key, value in state_dict.items() + if key.startswith(prefix) + } + ) + matching = max( + (candidate for candidate in candidates if candidate), + key=lambda candidate: len(set(candidate) & set(model_state)), + default={}, + ) + filtered = {key: value for key, value in matching.items() if key in model_state} + if not filtered: + raise ValueError(f"{path} contains no weights matching the configured model.") + result = model.load_state_dict(filtered, strict=False) + if result.missing_keys: + raise RuntimeError( + f"Checkpoint {path} is missing model weights: " + + ", ".join(sorted(result.missing_keys)) + ) + + +def _load_experiment( + checkpoint_path: Path, +) -> tuple[ChebaiBaseNet, XYBaseDataModule]: + checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + try: + data_hparams = copy.deepcopy(checkpoint["datamodule_hyper_parameters"]) + model_hparams = copy.deepcopy(checkpoint["hyper_parameters"]) + except KeyError as error: + raise ValueError( + "Checkpoint must contain Lightning model and datamodule hyperparameters " + "when --model-config and --data-config are omitted." + ) from error + + data_hparams.pop("_instantiator", None) + model_hparams.pop("_instantiator", None) + data = instantiate_module(XYBaseDataModule, data_hparams) + + load_kwargs = ( + {"pretrained_checkpoint": None} + if "pretrained_checkpoint" in model_hparams + else {} + ) + model = ChebaiBaseNet.load_from_checkpoint( + checkpoint_path, + map_location="cpu", + **load_kwargs, + ) + return model, data + + +def _prepare_model_batch(model: ChebaiBaseNet, batch: Any, batch_index: int) -> dict: + """Apply the repository's standard XYData-to-model batch conversion.""" + if not hasattr(batch, "to") or not hasattr(batch, "x"): + raise TypeError( + "The test dataloader must return an XYData batch with an x attribute." + ) + return model._process_batch(batch.to(model.device), batch_index) + + +def _write_csv(path: Path, rows: list[dict[str, Any]]) -> None: + if not rows: + return + with path.open("w", newline="") as handle: + writer = csv.DictWriter(handle, fieldnames=list(rows[0])) + writer.writeheader() + writer.writerows(rows) + + +def _plot_summary(summary: dict[str, Any], output_dir: Path) -> None: + try: + import matplotlib.pyplot as plt + except ImportError: + return + + metrics = { + "embedding_rms": [summary["embedding"][level]["rms"] for level in LEVELS], + "weight_fraction": [summary["classifier"][level]["alpha"] for level in LEVELS], + "logit_fraction": [summary["contribution"][level]["beta"] for level in LEVELS], + } + fig, axes = plt.subplots(1, 3, figsize=(12, 4)) + for axis, (title, values) in zip(axes, metrics.items()): + axis.bar(LEVELS, values) + axis.set_title(title.replace("_", " ")) + axis.tick_params(axis="x", rotation=25) + fig.tight_layout() + fig.savefig(output_dir / "representation_summary.png", dpi=180) + plt.close(fig) + + +def analyse(args: argparse.Namespace) -> dict[str, Any]: + model, data = _load_experiment( + Path(args.checkpoint), + ) + model.to(args.device).eval() + + linear_layers = [ + layer for layer in model.lin_sequential if isinstance(layer, torch.nn.Linear) + ] + if len(linear_layers) != 1: + raise ValueError( + "This analysis requires a single-layer linear classifier; " + f"found {len(linear_layers)} linear layers." + ) + classifier = linear_layers[0] + out_dir = Path(args.output_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + loader = data.test_dataloader() + if isinstance(loader, (list, tuple)): + if len(loader) != 1: + raise ValueError("Multiple dataloaders are not supported by this script.") + loader = loader[0] + + sums = {level: {"std": 0.0, "rms": 0.0, "l2": 0.0} for level in LEVELS} + contribution_sum = {level: 0.0 for level in LEVELS} + class_contribution = torch.zeros((classifier.out_features, 3), dtype=torch.float64) + ablation_sum = torch.zeros(3, dtype=torch.float64) + sample_count = 0 + batch_limit = args.max_batches + + with torch.inference_mode(): + for batch_index, raw_batch in enumerate(loader): + if batch_limit is not None and batch_index >= batch_limit: + break + batch = _prepare_model_batch(model, raw_batch, batch_index) + representations = model.pool_representations(batch) + batch_size = representations[0].shape[0] + sample_count += batch_size + blocks = [] + for index, (level, representation) in enumerate( + zip(LEVELS, representations) + ): + representation_cpu = representation.detach().cpu() + sums[level]["std"] += ( + float(representation_cpu.std().item()) * batch_size + ) + sums[level]["rms"] += ( + float(representation_cpu.square().mean().sqrt().item()) * batch_size + ) + sums[level]["l2"] += ( + float(representation_cpu.norm(dim=1).mean().item()) * batch_size + ) + start = index * representation.shape[1] + weight = classifier.weight[:, start : start + representation.shape[1]] + contribution = representation @ weight.t() + contribution_sum[level] += ( + float(contribution.norm(dim=1).mean().item()) * batch_size + ) + class_contribution[:, index] += ( + contribution.abs().sum(dim=0).detach().cpu().double() + ) + blocks.append(representation) + full_logits = classifier(torch.cat(blocks, dim=1)) + for index in range(3): + ablated = list(blocks) + ablated[index] = torch.zeros_like(ablated[index]) + ablation_sum[index] += ( + (full_logits - classifier(torch.cat(ablated, dim=1))) + .abs() + .sum() + .detach() + .cpu() + .double() + ) + + if sample_count == 0: + raise RuntimeError("The selected split produced no samples.") + widths = [classifier.in_features // len(LEVELS)] * len(LEVELS) + if classifier.in_features % len(LEVELS): + raise ValueError( + "The classifier input dimension is not divisible into three AMGPool blocks." + ) + weight_norms = [] + for index, level in enumerate(LEVELS): + width = widths[index] + weight_norms.append( + float( + classifier.weight[:, index * width : (index + 1) * width].norm().item() + ) + ) + weight_total = sum(weight_norms) + contribution_total = sum(contribution_sum.values()) + if weight_total == 0 or contribution_total == 0: + raise RuntimeError( + "The checkpoint has zero total classifier weight or logit contribution; " + "normalized shares are undefined." + ) + summary = { + "checkpoint": str(Path(args.checkpoint).resolve()), + "split": "test", + "samples": sample_count, + "embedding": { + level: { + metric: value / sample_count for metric, value in sums[level].items() + } + for level in LEVELS + }, + "classifier": { + level: { + "frobenius_norm": weight_norms[index], + "alpha": weight_norms[index] / weight_total, + } + for index, level in enumerate(LEVELS) + }, + "contribution": { + level: { + "mean_logit_l2": contribution_sum[level] / sample_count, + "beta": contribution_sum[level] / contribution_total, + } + for level in LEVELS + }, + "ablation_logit_change_mean": { + level: float( + ablation_sum[index].item() / sample_count / classifier.out_features + ) + for index, level in enumerate(LEVELS) + }, + } + class_rows = [ + { + "class_index": index, + **{ + f"{level}_mean_abs_logit_contribution": float( + class_contribution[index, level_index].item() / sample_count + ) + for level_index, level in enumerate(LEVELS) + }, + } + for index in range(classifier.out_features) + ] + (out_dir / "summary.json").write_text(json.dumps(summary, indent=2) + "\n") + _write_csv(out_dir / "class_contributions.csv", class_rows) + _plot_summary(summary, out_dir) + return summary + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--checkpoint", required=True, type=Path) + parser.add_argument("--output-dir", type=Path, default=Path("amgpool_analysis")) + parser.add_argument("--device", default="cpu") + parser.add_argument("--max-batches", type=int, default=None) + analyse(parser.parse_args()) + + +if __name__ == "__main__": + main()