Add typehints

gauche/data_featuriser/featurisation.py

@@ -1,6 +1,7 @@
 import numpy as np
 from rdkit.Chem import MolFromSmiles, AllChem, Descriptors
 import pandas as pd
+from typing import List, Optional
 from rxnfp.transformer_fingerprints import (
     get_default_model_and_tokenizer,
     RXNBERTFingerprintGenerator,
@@ -10,10 +11,12 @@
 import selfies as sf
 import graphein.molecule as gm
 from rdkit.Chem import rdMolDescriptors
+from graphein.molecule.config import MoleculeGraphConfig
+import networkx as nx
 
 
 # Reactions
-def one_hot(df):
+def one_hot(df: pd.DataFrame) -> np.ndarray:
     """
     Builds reaction representation as a bit vector which indicates whether
     a certain condition, reagent, reactant etc. is present in the reaction.
@@ -28,7 +31,7 @@ def one_hot(df):
     return df_ohe.to_numpy(dtype=np.float64)
 
 
-def rxnfp(reaction_smiles):
+def rxnfp(reaction_smiles: List[str]) -> np.ndarray:
     """
     https://rxn4chemistry.github.io/rxnfp/
 
@@ -44,7 +47,7 @@ def rxnfp(reaction_smiles):
     return np.array(rxnfps, dtype=np.float64)
 
 
-def drfp(reaction_smiles, nBits=2048):
+def drfp(reaction_smiles: List[str], nBits: int = 2048) -> np.ndarray:
     """
     https://github.com/reymond-group/drfp
 
@@ -59,7 +62,7 @@ def drfp(reaction_smiles, nBits=2048):
 
 
 # Molecules
-def fingerprints(smiles, bond_radius=3, nBits=2048):
+def fingerprints(smiles: List[str], bond_radius: int = 3, nBits: int = 2048) -> np.ndarray:
     rdkit_mols = [MolFromSmiles(smiles) for smiles in smiles]
     fps = [
         AllChem.GetMorganFingerprintAsBitVect(mol, bond_radius, nBits=nBits)
@@ -69,7 +72,7 @@ def fingerprints(smiles, bond_radius=3, nBits=2048):
 
 
 # auxiliary function to calculate the fragment representation of a molecule
-def fragments(smiles):
+def fragments(smiles: List[str]) -> np.ndarray:
     # descList[115:] contains fragment-based features only
     # (https://www.rdkit.org/docs/source/rdkit.Chem.Fragments.html)
     # Update: in the new RDKit version the indices are [124:]
@@ -87,7 +90,7 @@ def fragments(smiles):
 
 
 # auxiliary function to calculate bag of character representation of a molecular string
-def bag_of_characters(smiles, max_ngram=5, selfies=False):
+def bag_of_characters(smiles: List[str], max_ngram: int = 5, selfies: bool = False) -> np.ndarray:
     if selfies:  # convert SMILES to SELFIES
         strings = [sf.encoder(smiles[i]) for i in range(len(smiles))]
     else:  # otherwise stick with SMILES
@@ -100,13 +103,13 @@ def bag_of_characters(smiles, max_ngram=5, selfies=False):
     return cv.fit_transform(strings).toarray()
 
 
-def graphs(smiles, graphein_config=None):
+def graphs(smiles: List[str], graphein_config: Optional[MoleculeGraphConfig] = None) -> List[nx.Graph]:
     return [
         gm.construct_graph(smiles=i, config=graphein_config) for i in smiles
     ]
 
 
-def mqn_features(smiles):
+def mqn_features(smiles: List[str]) -> np.ndarray:
     """
     Builds molecular representation as a vector of Molecular Quantum Numbers.
 

gauche/dataloader/data_utils.py

@@ -1,25 +1,40 @@
 """
-Utility functions for molecular data
+Utility functions for molecular data.
 """
-
+import numpy as np
+from typing import Optional Tuple
 from sklearn.decomposition import PCA
 from sklearn.preprocessing import StandardScaler
 
 
 def transform_data(
-    X_train, y_train, X_test, y_test, n_components=None, use_pca=False
-):
+    X_train: np.ndarray,
+    y_train: np.ndarray,
+    X_test: np.ndarray,
+    y_test: np.ndarray,
+    n_components: Optional[int] = None,
+    use_pca: bool = False
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, StandardScaler]:
     """
-    Apply feature scaling, dimensionality reduction to the data. Return the standardised and low-dimensional train and
+    Apply feature scaling, dimensionality reduction to the data.
+
+    Returns the standardised and low-dimensional train and
     test sets together with the scaler object for the target values.
 
     :param X_train: input train data
+    :type X_train: np.ndarray
     :param y_train: train labels
+    :type y_train: np.ndarray
     :param X_test: input test data
+    :type X_test: np.ndarray
     :param y_test: test labels
-    :param n_components: number of principal components to keep when use_pca = True
-    :param use_pca: Whether or not to use PCA
+    :type y_test: np.ndarray
+    :param n_components: number of principal components to keep when ``use_pca=True``
+    :type n_components: int, optional. Default is ``None``.
+    :param use_pca: Whether or not to use PCA.
+    :type use_pca: bool
     :return: X_train_scaled, y_train_scaled, X_test_scaled, y_test_scaled, y_scaler
+    :rtype: Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, sklearn.preprocessing.StandardScaler]
     """
 
     x_scaler = StandardScaler()

gauche/dataloader/dataloader.py

@@ -54,7 +54,7 @@ def featurize(self, representation):
         raise NotImplementedError
 
     def split_and_scale(
-        self, test_size=0.2, scale_labels=True, scale_features=False
+        self, test_size: float = 0.2, scale_labels: bool = True, scale_features: bool = False
     ):
         """Splits the data into training and testing sets.
 

gauche/dataloader/mol_prop.py

@@ -16,7 +16,8 @@
 
 from gauche.dataloader import DataLoader
 from rdkit.Chem import MolFromSmiles
-
+from graphein.molecule.config import MoleculeGraphConfig
+from typing import Optional
 
 class DataLoaderMP(DataLoader):
     def __init__(self):
@@ -41,7 +42,7 @@ def labels(self):
     def labels(self, value):
         self._labels = value
 
-    def validate(self, drop=True):
+    def validate(self, drop: bool = True):
         """Checks if the features are valid SMILES strings and (potentially)
         drops the entries that are not.
 
@@ -68,11 +69,11 @@ def validate(self, drop=True):
 
     def featurize(
         self,
-        representation,
-        bond_radius=3,
-        nBits=2048,
-        graphein_config=None,
-        max_ngram=5,
+        representation: str,
+        bond_radius: int = 3,
+        nBits: int = 2048,
+        graphein_config: Optional[MoleculeGraphConfig] = None,
+        max_ngram: int = 5,
     ):
         """Transforms SMILES into the specified molecular representation.
 
@@ -137,7 +138,7 @@ def featurize(
                 f"Choose between {valid_representations}."
             )
 
-    def load_benchmark(self, benchmark, path):
+    def load_benchmark(self, benchmark: str, path: str):
         """Loads features and labels from one of the included benchmark datasets
         and feeds them into the DataLoader.
 
@@ -146,6 +147,7 @@ def load_benchmark(self, benchmark, path):
         :type benchmark: str
         :param path: the path to the dataset in csv format
         :type path: str
+        :raises ValueError: If an unsupported benchmark is provided.
         """
 
         benchmarks = {

gauche/dataloader/reaction_loader.py

@@ -30,13 +30,14 @@ def labels(self, value):
     def validate(self, drop=True):
         invalid_idx = []
 
-    def featurize(self, representation, nBits=2048):
+    def featurize(self, representation: str, nBits: int = 2048):
         """Transforms reactions into the specified representation.
 
         :param representation: the desired reaction representation, one of [ohe, rxnfp, drfp, bag_of_smiles]
         :type representation: str
         :param nBits: int giving the bit vector length for drfp representation. Default is 2048
         :type nBits: int
+        :raises ValueError: If unsupported ``representation`` is provided.
         """
 
         valid_representations = [
@@ -64,7 +65,7 @@ def featurize(self, representation, nBits=2048):
                 f"Choose between {valid_representations}."
             )
 
-    def load_benchmark(self, benchmark, path):
+    def load_benchmark(self, benchmark: str, path: str):
 
         """Loads features and labels from one of the included benchmark datasets
                 and feeds them into the DataLoader.
@@ -75,6 +76,7 @@ def load_benchmark(self, benchmark, path):
         :type benchmark: str
         :param path: the path to the dataset in csv format
         :type path: str
+        :raises ValueError: If unsupported ``benchmark`` is provided.
         """
 
         benchmarks = {

gauche/kernels/fingerprint_kernels/base_fingerprint_kernel.py

@@ -34,7 +34,7 @@ def __init__(self, postprocess_script=default_postprocess_script):
         super().__init__()
         self._postprocess = postprocess_script
 
-    def _sim(self, x1, x2, postprocess, x1_eq_x2=False, metric="tanimoto"):
+    def _sim(self, x1: torch.Tensor, x2: torch.Tensor, postprocess: bool, x1_eq_x2: bool = False, metric: str = "tanimoto") -> torch.Tensor:
         r"""
         Computes the similarity between x1 and x2
         Args:
@@ -82,18 +82,18 @@ def __init__(self, metric="", **kwargs):
         super().__init__(**kwargs)
         self.metric = metric
 
-    def forward(self, x1, x2, **params):
+    def forward(self, x1: torch.Tensor, x2: torch.Tensor, **params) -> torch.Tensor:
         return self.covar_dist(x1, x2, **params)
 
     def covar_dist(
         self,
-        x1,
-        x2,
-        last_dim_is_batch=False,
+        x1: torch.Tensor,
+        x2: torch.Tensor,
+        last_dim_is_batch: bool = False,
         dist_postprocess_func=default_postprocess_script,
-        postprocess=True,
+        postprocess: bool = True,
         **params,
-    ):
+    ) -> torch.Tensor:
         r"""
         This is a helper method for computing the bit vector similarity between
         all pairs of points in x1 and x2.

gauche/kernels/fingerprint_kernels/tanimoto_kernel.py

@@ -47,7 +47,7 @@ def __init__(self, **kwargs):
         super(TanimotoKernel, self).__init__(**kwargs)
         self.metric = "tanimoto"
 
-    def forward(self, x1, x2, diag=False, **params):
+    def forward(self, x1: torch.Tensor, x2: torch.Tensor, diag: bool = False, **params) -> torch.Tensor:
         if diag:
             assert x1.size() == x2.size() and torch.equal(x1, x2)
             return torch.ones(

gauche/kernels/gnn_kernels/pretrained_kernel.py

@@ -55,7 +55,7 @@
 masked_bond_token = 5  # bond type for masked edges
 
 
-def mol_to_pyg(mol):
+def mol_to_pyg(mol: Chem.Mol) -> Data:
     """
     A featuriser that accepts an rdkit mol instance and
     converts it to a PyTorch Geometric data object that
@@ -117,7 +117,7 @@ class GINConv(MessagePassing):
     edge information by concatenating edge embeddings.
     """
 
-    def __init__(self, emb_dim, aggr="add"):
+    def __init__(self, emb_dim: int, aggr: str = "add"):
         """
         Initialise GIN convolutional layer.
         Args:
@@ -139,7 +139,7 @@ def __init__(self, emb_dim, aggr="add"):
 
         self.aggr = aggr
 
-    def forward(self, x, edge_index, edge_attr):
+    def forward(self, x: torch.Tensor, edge_index: torch.Tensor, edge_attr: torch.Tensor) -> torch.Tensor:
         """
         Message passing and aggregation function
         of the adapted GIN convolutional layer.
@@ -169,10 +169,10 @@ def forward(self, x, edge_index, edge_attr):
         ) + self.edge_embedding2(edge_attr[:, 1])
         return self.propagate(edge_index, x=x, edge_attr=edge_embeddings)
 
-    def message(self, x_j, edge_attr):
+    def message(self, x_j: torch.Tensor, edge_attr: torch.Tensor) -> torch.Tensor:
         return x_j + edge_attr
 
-    def update(self, aggr_out):
+    def update(self, aggr_out: torch.Tensor) -> torch.Tensor:
         return self.mlp(aggr_out)
 
 
@@ -182,7 +182,7 @@ class GCNConv(MessagePassing):
     edge information by concatenating edge embeddings.
     """
 
-    def __init__(self, emb_dim, aggr="add"):
+    def __init__(self, emb_dim: int, aggr: str = "add"):
         super(GCNConv, self).__init__(aggr=aggr)
 
         self.linear = torch.nn.Linear(emb_dim, emb_dim)
@@ -193,7 +193,7 @@ def __init__(self, emb_dim, aggr="add"):
         torch.nn.init.xavier_uniform_(self.edge_embedding2.weight.data)
 
     @staticmethod
-    def norm(edge_index, num_nodes, dtype):
+    def norm(edge_index: torch.Tensor, num_nodes: int, dtype) -> torch.Tensor:
         """
         Symmetric normalisation step of the adjacency matrix A:
         .. math::
@@ -217,7 +217,7 @@ def norm(edge_index, num_nodes, dtype):
         deg_inv_sqrt[deg_inv_sqrt == float("inf")] = 0
         return deg_inv_sqrt[row] * edge_weight * deg_inv_sqrt[col]
 
-    def forward(self, x, edge_index, edge_attr):
+    def forward(self, x: torch.Tensor, edge_index: torch.Tensor, edge_attr: torch.Tensor) -> torch.Tensor:
         """
         Message passing and aggregation function
         of the adapted GCN convolutional layer.
@@ -261,7 +261,7 @@ class GNN(torch.nn.Module):
     Combine multiple GNN layers into a network.
     """
 
-    def __init__(self, num_layers=5, embed_dim=300, gnn_type="gin"):
+    def __init__(self, num_layers: int = 5, embed_dim: int = 300, gnn_type: str = "gin"):
         """
         Compose convolution layers into GNN. Pretrained parameters
         exist for a 5-layer network with 300 hidden units.
@@ -300,7 +300,7 @@ def __init__(self, num_layers=5, embed_dim=300, gnn_type="gin"):
         for layer in range(self.num_layers):
             self.batch_norms.append(torch.nn.BatchNorm1d(self.embed_dim))
 
-    def load_pretrained(self, pretrain_type, device):
+    def load_pretrained(self, pretrain_type: str, device):
         """
         Checks if a pretrained parameter set for the specified
         pretraining procedure exists and updates the
@@ -328,7 +328,7 @@ def load_pretrained(self, pretrain_type, device):
             )
             self.load_state_dict(pretrained_state_dict)
 
-    def forward(self, x, edge_index, edge_attr):
+    def forward(self, x: torch.Tensor, edge_index: torch.Tensor, edge_attr: torch.Tensor) -> torch.Tensor:
         """
         Forward function of the GNN class that takes a PyTorch geometric
         representation of a molecule or a batch of molecules