Fixing unittest + mkdocs errors

.github/workflows/test.yml

@@ -16,7 +16,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.09", "3.10", "3.11"]
         pytorch-version: ["2.0"]
 
     runs-on: "ubuntu-latest"

docs/api/graphium.features.md

@@ -5,37 +5,8 @@ Feature extraction and manipulation
 === "Contents"
 
     * [Featurizer](#featurizer)
-    * [Positional Encoding](#positional-encoding)
-    * [Properties](#properties)
-    * [Spectral PE](#spectral-pe)
-    * [Random Walk PE](#random-walk-pe)
-    * [NMP](#nmp)
 
 ## Featurizer
 ------------
 ::: graphium.features.featurizer
 
-
-## Positional Encoding
-------------
-::: graphium.features.positional_encoding
-
-
-## Properties
-------------
-::: graphium.features.properties
-
-
-## Spectral PE
-------------
-::: graphium.features.spectral
-
-
-## Random Walk PE
-------------
-::: graphium.features.rw
-
-
-## NMP
-------------
-::: graphium.features.nmp

docs/api/graphium.utils.md

@@ -46,10 +46,6 @@ module for utility functions
 ::: graphium.utils.mup
 
 
-## Read File
-----------------
-::: graphium.utils.read_file
-
 ## Safe Run
 ----------------
 ::: graphium.utils.safe_run

graphium/features/README.md

@@ -7,8 +7,5 @@
 ## What is in this folder? 
 
 - ✅ `featurizer.py`: featurization code for the molecules, adding node, edge and graph features to the mol object
-- `nmp.py`: check if a string can be converted to float, helper function for featurization
-- `positional_encoding.py`: code for computing all raw positional and structural encoding of the graph, see `graph_positional_encoder` function
-- `properties.py`: code for computing properties of the molecule
-- `rw.py`: code for computing random walk positional encoding
-- `spectral.py`: code for computing the spectral positional encoding such as the Laplacian eigenvalues and eigenvectors
+
+Positional encodings, and atom/bond features (`nmp.py`) have been moved to the `/graphium_cpp` folder.

graphium/graphium_cpp/one_hot.cpp

@@ -28,7 +28,7 @@ class OneHotLookup {
     }
 };
 
-// This list of elements matches ATOM_LIST in graphium/features/nmp.py
+// This list of elements matches ATOM_LIST in older file graphium/features/nmp.py
 constexpr size_t atomicNumList[] = {
     6 -1, // C
     7 -1, // N
@@ -81,7 +81,7 @@ constexpr size_t degreeCount = 5;
 constexpr size_t valenceCount = 7;
 
 // Reverse alphabetical order, excluding "OTHER",
-// matching HYBRIDIZATION_LIST in graphium/features/nmp.py
+// matching HYBRIDIZATION_LIST in older file graphium/features/nmp.py
 constexpr size_t hybridizationList[] = {
     RDKit::Atom::HybridizationType::UNSPECIFIED,
     RDKit::Atom::HybridizationType::SP3D2,
@@ -142,7 +142,7 @@ constexpr ElementType atomicNumToType[] = {
 };
 constexpr size_t typeCount = ElementType::NUM_ELEMENT_TYPES;
 
-// This matches BOND_TYPES in graphium/features/nmp.py
+// This matches BOND_TYPES in older file graphium/features/nmp.py
 constexpr size_t bondTypeList[] = {
     RDKit::Bond::BondType::SINGLE,
     RDKit::Bond::BondType::DOUBLE,
@@ -152,7 +152,7 @@ constexpr size_t bondTypeList[] = {
 constexpr size_t bondTypeCount = std::extent<decltype(bondTypeList)>::value;
 constexpr OneHotLookup<22, bondTypeCount> bondTypeLookup(bondTypeList);
 
-// This matches BOND_STEREO in graphium/features/nmp.py
+// This matches BOND_STEREO in older file graphium/features/nmp.py
 constexpr size_t bondStereoList[] = {
     RDKit::Bond::BondStereo::STEREONONE,
     RDKit::Bond::BondStereo::STEREOANY,

graphium/nn/base_layers.py

@@ -190,7 +190,7 @@ def forward(
         # key_padding_mask: [batch, 1, 1, nodes]
         if key_padding_mask is not None:
             masked_attn_weights = attn_weights.masked_fill(
-                key_padding_mask.unsqueeze(1).unsqueeze(2),
+                key_padding_mask.unsqueeze(1).unsqueeze(2).bool(), # The mask is cast to float somewhere in TransformerEncoder
                 key_padding_mask_value,
             )
         else:

tests/test_ipu_dataloader.py

@@ -26,7 +26,6 @@
 
 import torch
 from torch.utils.data.dataloader import default_collate
-from lightning_graphcore import IPUStrategy
 
 
 def random_packing(num_nodes, batch_size):
@@ -119,6 +118,8 @@ def test_poptorch_simple_deviceiterations_gradient_accumulation(self):
         Test a simple version of the device-iterations and gradient accumulation
         to make sure that the dataloader and models handle them correcly.
         """
+        from lightning_graphcore import IPUStrategy
+
 
         with patch("poptorch.ipuHardwareIsAvailable", return_value=True):
             with patch("lightning_graphcore.accelerator._IPU_AVAILABLE", new=True):

tests/test_ipu_to_dense_batch.py

@@ -16,21 +16,8 @@
 import torch
 from torch_geometric.data import Data, Batch
 from graphium.ipu.to_dense_batch import to_dense_batch
-from warnings import warn
-
-
-# General imports
-import yaml
-import unittest as ut
-import numpy as np
-from copy import deepcopy
-from warnings import warn
-from lightning import Trainer, LightningModule
-from lightning_graphcore import IPUStrategy
-from functools import partial
 
 import torch
-from torch.utils.data.dataloader import default_collate
 
 # Current library imports
 from graphium.config._loader import load_datamodule, load_metrics, load_architecture, load_accelerator

tests/test_positional_encoders.py

@@ -187,8 +187,8 @@ def test_laplacian_eigvec_with_encoder(self):
                         data_dict = {
                             # "feat": tensors[2],
                             # "edge_feat": tensors[3],
-                            "laplacian_eigval": tensors[4],
-                            "laplacian_eigvec": tensors[5],
+                            "laplacian_eigval": tensors[4].float(),
+                            "laplacian_eigvec": tensors[5].float(),
                         }
                         # Create the PyG graph object `Data`
                         data = Data(