GPSE Implementation

[semihcanturk]

Graph Positional and Structural Encoder implementation as per #8310. Adapted from the original repository: https://github.com/G-Taxonomy-Workgroup/GPSE. This version is a standalone implementation that is decoupled from GraphGym, and thus aims for better accessibility and a smoother integration into PyG. While the priority of this PR is to enable loading and using pre-trained models in plug-and-play fashion, it also includes the custom loss function used to train the model. Nevertheless, it might be easier to use the original repository for pre-training and fine-tuning new GPSE models for the time being.

This PR includes the following:

• GPSE: The main GPSE module, that generates learned encodings for input graphs.
• Several helper classes (FeatureEncoder, GNNStackStage, IdentityHead, GNNInductiveHybridMultiHead, ResGatedGCNConvGraphGymLayer, Linear, MLP, GeneralMultiLayer, GeneralLayer, BatchNorm1dNode, BatchNorm1dEdge, VirtualNodePatchSingleton) and wrapper functions (GNNPreMP, GNNLayer), all adapted from their GraphGym versions for compatibility and enabling the loading of weights pre-trained using the GraphGym/original version.
• The class method GPSE.from_pretrained() that returns a model with pre-trained weights from the original repository/Zenodo files.
• GPSENodeEncoder, a helper linear/MLP encoder that takes the GPSE encodings precomputed asbatch.pestat_GPSE in the input graphs, maps them to a desired dimension and appends them to node features.
• precompute_GPSE , a function that takes in a GPSE model and a dataset, and precomputes GPSE encodings in-place for a given dataset using the helper function gpse_process_batch.
• The transform AddGPSE, which in similar fashion to AddLaplacianEigenvectorPE and AddRandomWalkPE adds the GPSE encodings to a given graph using the helper function gpse_process
• The testing modules test/test_gpse.py and test/test_add_gpse.py.
• The loss function gpse_loss and helper functions cosim_col_sep and process_batch_idx used in GPSE training.
• A comprehensive example in examples/gpse.py using the ZINC dataset. Two different ways of using GPSE to generate encodings is demonstrated: (a) as a function that adds encodings in-place through precompute_GPSE, (b) as a pre-transform. [EDIT: To be added as a separate PR]

[codecov]

Codecov Report

Attention: Patch coverage is 64.41860% with 153 lines in your changes are missing coverage. Please review.

Project coverage is 89.23%. Comparing base (ed17034) to head (f24acbc).

❗ Current head f24acbc differs from pull request most recent head 9f000d0. Consider uploading reports for the commit 9f000d0 to get more accurate results

Files	Patch %	Lines
torch_geometric/nn/models/gpse.py	62.68%	153 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #9018      +/-   ##
==========================================
+ Coverage   88.46%   89.23%   +0.76%     
==========================================
  Files         470      472       +2     
  Lines       30189    30594     +405     
==========================================
+ Hits        26708    27301     +593     
+ Misses       3481     3293     -188     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.