This repository contains an implementation of a multiscale operator learning framework for modelling and forecasting complex social systems. The framework learns multiscale dynamics and forecasts the evolution of a complex system given an initial adjacency matrix
A brief overview of important modules in this repository are:
train.py
- Data loading, LR scheduling, graph sampling, and logging of training data.
nn/models/renonet.py
- Contains a module of the framework shown below, as well as vmapped and serial loss functions for optimizing the loss shown below.
nn/models/models.py
- Modules for the encoder and renormalization networks (GCN, HGCN) and decoder networks (MLP, Transformer, DeepOnet).
lib/graph_utils.py
- Utilities for sampling, padding, and manipulating graphs.
lib/positional_encoding.py
- Functions for computing positional encoding (node2vec, random walk PE, and laplacian eigenvector PE).
nn/manifolds/
- Manifold definitions for hyperbolic layers. Ported from the original pytorch code (HGCN) to JAX.
- Includes Euclidean, Poincaré, and Hyperboloid manifolds.
![renonet](https://private-user-images.githubusercontent.com/50005216/278086383-f1f3b8a8-45cd-4c43-a76a-5978f7df79f1.png?jwt=eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJnaXRodWIuY29tIiwiYXVkIjoicmF3LmdpdGh1YnVzZXJjb250ZW50LmNvbSIsImtleSI6ImtleTUiLCJleHAiOjE3MjUwOTA1MzEsIm5iZiI6MTcyNTA5MDIzMSwicGF0aCI6Ii81MDAwNTIxNi8yNzgwODYzODMtZjFmM2I4YTgtNDVjZC00YzQzLWE3NmEtNTk3OGY3ZGY3OWYxLnBuZz9YLUFtei1BbGdvcml0aG09QVdTNC1ITUFDLVNIQTI1NiZYLUFtei1DcmVkZW50aWFsPUFLSUFWQ09EWUxTQTUzUFFLNFpBJTJGMjAyNDA4MzElMkZ1cy1lYXN0LTElMkZzMyUyRmF3czRfcmVxdWVzdCZYLUFtei1EYXRlPTIwMjQwODMxVDA3NDM1MVomWC1BbXotRXhwaXJlcz0zMDAmWC1BbXotU2lnbmF0dXJlPWI3MDk4OGFhNDhkN2VkNTdlMWI2ZjA0ZTFlYmYxOTIwYTlkMThhNDIyZmRmNzViZTY4NTI4Y2ZmY2M1ZTAyNTEmWC1BbXotU2lnbmVkSGVhZGVycz1ob3N0JmFjdG9yX2lkPTAma2V5X2lkPTAmcmVwb19pZD0wIn0.bcFjWVt_y94-kVnYe3TYOKCCwSH137kvqzeCbxxOaCI)