# Haiku Geometric
[**Overview**](#overview)
| [**Installation**](#installation)
| [**Quickstart**](#quickstart)
| [**Examples**](#examples)
| [**Documentation**](https://haiku-geometric.readthedocs.io/en/latest/)
| [**License**](#license)
[](https://haiku-geometric.readthedocs.io/en/latest/?badge=latest)
[](https://github.com/alexOarga/haiku-geometric/actions/workflows/python-app.yml)
## Overview
Haiku Geometric is a collection of graph neural networks (GNNs) implemented using [JAX](https://jax.readthedocs.io/en/latest/notebooks/quickstart.html). It tries to provide **object-oriented** and **easy-to-use** modules for GNNs.
Haiku Geometric is built on top of [Haiku](https://github.com/deepmind/dm-haiku) and [Jraph](https://github.com/deepmind/jraph).
It is deeply inspired by [PyTorch Geometric](https://github.com/pyg-team/pytorch_geometric).
In most cases, Haiku Geometric tries to replicate the API of PyTorch Geometric to allow code sharing between the two.
Haiku Geometric is still under development and I would advise against using it in production.
## Installation
Haiku Geometric can be installed from source:
```bash
pip install git+https://github.com/alexOarga/haiku-geometric.git
```
Alternatively, you can install Haiku Geometric using pip:
```bash
pip install haiku-geometric
```
## Quickstart
For instance, we can create a simple graph convolutional network (GCN) of 2 layers
as follows:
```python
import jax
import haiku as hk
from haiku_geometric.nn import GCNConv
class GCN(hk.Module):
def __init__(self, hidden_channels, out_channels):
super().__init__()
self.conv1 = GCNConv(hidden_channels)
self.conv2 = GCNConv(hidden_channels)
self.linear = hk.Linear(out_channels)
def __call__(self, nodes,senders, receivers):
x = self.conv1(nodes, senders, receivers)
x = jax.nn.relu(x)
x = self.conv2(x, senders, receivers)
x = self.linear(nodes)
return x
def forward(nodes, senders, receivers):
gcn = GCN(16, 7)
return gcn(nodes, senders, receivers)
```
The GNN that we have defined is a Haiku Module.
To convert our module in a function that can be used with JAX, we transform
it using `hk.transform` as described in the
[Haiku documentation](https://dm-haiku.readthedocs.io/en/latest/).
```python
model = hk.transform(forward)
model = hk.without_apply_rng(model)
rng = jax.random.PRNGKey(42)
params = model.init(rng, nodes=nodes, senders=senders, receivers=receivers)
```
We can now run a forward pass on the model:
```python
output = model.apply(params=params, nodes=nodes, senders=senders, receivers=receivers)
```
## Documentation
The documentation for Haiku Geometric can be found [here](https://haiku-geometric.readthedocs.io/en/latest/).
## Examples
Haiku Geometric comes with a few examples that showcase the usage of the library.
The following examples are available:
| | Link |
|-----------------------------------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| Quickstart Example | [](https://colab.research.google.com/github/alexOarga/haiku-geometric/blob/main/docs/source/notebooks/1_quickstart.ipynb) |
| Graph Convolution Networks with Karate Club dataset | [](https://colab.research.google.com/github/alexOarga/haiku-geometric/blob/main/examples/GCNConv_karate_club.ipynb) |
| Graph Attention Networks with CORA dataset | [](https://colab.research.google.com/github/alexOarga/haiku-geometric/blob/main/examples/GATConv_CORA.ipynb) |
| TopKPooling and GraphConv with PROTEINS dataset | [](https://colab.research.google.com/github/alexOarga/haiku-geometric/blob/main/examples/TopKPooling_GraphConv_PROTEINS.ipynb) |
## Implemented GNNs modules
Currently, Haiku Geometric includes the following GNN modules:
| Model | Description |
|---------------------------------------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------|
| [GCNConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.GCNConv) | Graph convolution layer from the [Semi-Supervised Classification with Graph Convolutional Networks](https://arxiv.org/abs/1609.02907) paper. |
| [GATConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.GATConv) | Graph attention layer from the [Graph Attention Networks](https://arxiv.org/abs/1710.10903) paper. |
| [SAGEConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.SAGEConv) | Graph convolution layer from the [Inductive Representation Learning on Large Graphs](https://arxiv.org/abs/1706.02216) paper. |
| [GINConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.GINConv) | Graph isomorphism network layer from the [How Powerful are Graph Neural Networks?](https://arxiv.org/abs/1810.00826) paper. |
| [GINEConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.GINEConv) | Graph isomorphism network layer from the [Strategies for Pre-training Graph Neural Networks](https://arxiv.org/abs/1905.12265) paper. |
| [GraphConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.GraphConv) | Graph convolution layer from the [Weisfeiler and Leman Go Neural: Higher-order Graph Neural Networks](https://arxiv.org/abs/1810.02244) paper. |
| [GeneralConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.GeneralConv) | A general GNN layer adapted from the [Design Space for Graph Neural Networks](https://arxiv.org/abs/2011.08843) paper. |
| [GatedGraphConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.GatedGraphConv) | Graph convolution layer from the [Gated Graph Sequence Neural Networks](https://arxiv.org/abs/1511.05493) paper. |
| [EdgeConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.EdgeConv) | Edge convolution layer from the [Dynamic Graph CNN for Learning on Point Clouds](https://arxiv.org/abs/1801.07829) paper. |
| [PNAConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.PNAConv) | Propagation Network layer from the [Principal Neighbourhood Aggregation for Graph Nets](https://arxiv.org/abs/2004.05718) paper. |
| [MetaLayer](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.MetaLayer) | Meta layer from the [Relational Inductive Biases, Deep Learning, and Graph Networks](https://arxiv.org/abs/1806.01261) paper. |
| [GPSLayer](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.GPSLayer) | Graph layer from the [Recipe for a General, Powerful, Scalable Graph Transformer](https://arxiv.org/abs/2205.12454) paper. |
## Implemented positional encodings
The following positional encodings are currently available:
| Model | Description |
|---------------------------------------------------------------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------|
| [LaplacianEncoder](https://haiku-geometric.readthedocs.io/en/latest/modules/posenc.html#haiku_geometric.posenc.LaplacianEncoder) | Laplacian positional encoding from the [Rethinking Graph Transformers with Spectral Attention](https://arxiv.org/pdf/2106.03893) paper. |
| [MagLaplacianEncoder](https://haiku-geometric.readthedocs.io/en/latest/modules/posenc.html#haiku_geometric.posenc.MagLaplacianEncoder) | Magnetic Laplacian positional encoding from the [Transformers Meet Directed Graphs](https://arxiv.org/pdf/2302.00049) paper. |
## Issues
If you encounter any issue, please [open an issue](https://github.com/alexOarga/haiku-geometric/issues/new).
## Running tests
Haiku Geometric can be tested using `pytest` by running the following command:
```bash
python -m pytest test/
```
## License
[](https://opensource.org/licenses/Apache-2.0)
Raw data
{
"_id": null,
"home_page": "https://github.com/alexOarga/haiku-geometric",
"name": "haiku-geometric",
"maintainer": "",
"docs_url": null,
"requires_python": "",
"maintainer_email": "",
"keywords": "",
"author": "Alex Oarga",
"author_email": "alex718123@gmail.com",
"download_url": "https://files.pythonhosted.org/packages/63/31/6291b06bc36873b79803dc9eaea88e0a8bdc5062b9595c41dc5b8b82f178/haiku_geometric-0.0.5.tar.gz",
"platform": null,
"description": "# Haiku Geometric\n\n[**Overview**](#overview)\n| [**Installation**](#installation)\n| [**Quickstart**](#quickstart)\n| [**Examples**](#examples)\n| [**Documentation**](https://haiku-geometric.readthedocs.io/en/latest/)\n| [**License**](#license)\n\n[](https://haiku-geometric.readthedocs.io/en/latest/?badge=latest)\n[](https://github.com/alexOarga/haiku-geometric/actions/workflows/python-app.yml)\n\n\n## Overview\n\nHaiku Geometric is a collection of graph neural networks (GNNs) implemented using [JAX](https://jax.readthedocs.io/en/latest/notebooks/quickstart.html). It tries to provide **object-oriented** and **easy-to-use** modules for GNNs.\n\nHaiku Geometric is built on top of [Haiku](https://github.com/deepmind/dm-haiku) and [Jraph](https://github.com/deepmind/jraph).\nIt is deeply inspired by [PyTorch Geometric](https://github.com/pyg-team/pytorch_geometric). \nIn most cases, Haiku Geometric tries to replicate the API of PyTorch Geometric to allow code sharing between the two.\n\nHaiku Geometric is still under development and I would advise against using it in production.\n\n## Installation\n\nHaiku Geometric can be installed from source:\n\n```bash\npip install git+https://github.com/alexOarga/haiku-geometric.git\n```\n\nAlternatively, you can install Haiku Geometric using pip:\n```bash\npip install haiku-geometric\n```\n\n## Quickstart\n\nFor instance, we can create a simple graph convolutional network (GCN) of 2 layers \nas follows:\n```python\nimport jax\nimport haiku as hk\nfrom haiku_geometric.nn import GCNConv\n\nclass GCN(hk.Module):\n def __init__(self, hidden_channels, out_channels):\n super().__init__()\n self.conv1 = GCNConv(hidden_channels)\n self.conv2 = GCNConv(hidden_channels)\n self.linear = hk.Linear(out_channels)\n\n def __call__(self, nodes,senders, receivers):\n x = self.conv1(nodes, senders, receivers)\n x = jax.nn.relu(x)\n x = self.conv2(x, senders, receivers)\n x = self.linear(nodes)\n return x\n\ndef forward(nodes, senders, receivers):\n gcn = GCN(16, 7)\n return gcn(nodes, senders, receivers)\n```\n\nThe GNN that we have defined is a Haiku Module. \nTo convert our module in a function that can be used with JAX, we transform\nit using `hk.transform` as described in the \n[Haiku documentation](https://dm-haiku.readthedocs.io/en/latest/).\n\n```python\nmodel = hk.transform(forward)\nmodel = hk.without_apply_rng(model)\nrng = jax.random.PRNGKey(42)\nparams = model.init(rng, nodes=nodes, senders=senders, receivers=receivers)\n```\n\nWe can now run a forward pass on the model:\n```python\noutput = model.apply(params=params, nodes=nodes, senders=senders, receivers=receivers)\n```\n\n## Documentation\n\nThe documentation for Haiku Geometric can be found [here](https://haiku-geometric.readthedocs.io/en/latest/).\n\n## Examples\n\nHaiku Geometric comes with a few examples that showcase the usage of the library.\nThe following examples are available:\n\n| | Link |\n|-----------------------------------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|\n| Quickstart Example | [](https://colab.research.google.com/github/alexOarga/haiku-geometric/blob/main/docs/source/notebooks/1_quickstart.ipynb) |\n| Graph Convolution Networks with Karate Club dataset | [](https://colab.research.google.com/github/alexOarga/haiku-geometric/blob/main/examples/GCNConv_karate_club.ipynb) |\n| Graph Attention Networks with CORA dataset | [](https://colab.research.google.com/github/alexOarga/haiku-geometric/blob/main/examples/GATConv_CORA.ipynb) |\n| TopKPooling and GraphConv with PROTEINS dataset | [](https://colab.research.google.com/github/alexOarga/haiku-geometric/blob/main/examples/TopKPooling_GraphConv_PROTEINS.ipynb) |\n\n\n## Implemented GNNs modules\n\nCurrently, Haiku Geometric includes the following GNN modules:\n\n| Model | Description |\n|---------------------------------------------------------------------------------------------------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------|\n| [GCNConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.GCNConv) | Graph convolution layer from the [Semi-Supervised Classification with Graph Convolutional Networks](https://arxiv.org/abs/1609.02907) paper. |\n| [GATConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.GATConv) | Graph attention layer from the [Graph Attention Networks](https://arxiv.org/abs/1710.10903) paper. |\n| [SAGEConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.SAGEConv) | Graph convolution layer from the [Inductive Representation Learning on Large Graphs](https://arxiv.org/abs/1706.02216) paper. |\n| [GINConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.GINConv) | Graph isomorphism network layer from the [How Powerful are Graph Neural Networks?](https://arxiv.org/abs/1810.00826) paper. |\n| [GINEConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.GINEConv) | Graph isomorphism network layer from the [Strategies for Pre-training Graph Neural Networks](https://arxiv.org/abs/1905.12265) paper. |\n| [GraphConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.GraphConv) | Graph convolution layer from the [Weisfeiler and Leman Go Neural: Higher-order Graph Neural Networks](https://arxiv.org/abs/1810.02244) paper. |\n| [GeneralConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.GeneralConv) | A general GNN layer adapted from the [Design Space for Graph Neural Networks](https://arxiv.org/abs/2011.08843) paper. |\n| [GatedGraphConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.GatedGraphConv) | Graph convolution layer from the [Gated Graph Sequence Neural Networks](https://arxiv.org/abs/1511.05493) paper. |\n| [EdgeConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.EdgeConv) | Edge convolution layer from the [Dynamic Graph CNN for Learning on Point Clouds](https://arxiv.org/abs/1801.07829) paper. |\n| [PNAConv](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.PNAConv) | Propagation Network layer from the [Principal Neighbourhood Aggregation for Graph Nets](https://arxiv.org/abs/2004.05718) paper. |\n| [MetaLayer](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.MetaLayer) | Meta layer from the [Relational Inductive Biases, Deep Learning, and Graph Networks](https://arxiv.org/abs/1806.01261) paper. |\n| [GPSLayer](https://haiku-geometric.readthedocs.io/en/latest/modules/nn.html#haiku_geometric.nn.conv.GPSLayer) | Graph layer from the [Recipe for a General, Powerful, Scalable Graph Transformer](https://arxiv.org/abs/2205.12454) paper. |\n\n## Implemented positional encodings\n\nThe following positional encodings are currently available:\n\n| Model | Description |\n|---------------------------------------------------------------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------|\n| [LaplacianEncoder](https://haiku-geometric.readthedocs.io/en/latest/modules/posenc.html#haiku_geometric.posenc.LaplacianEncoder) | Laplacian positional encoding from the [Rethinking Graph Transformers with Spectral Attention](https://arxiv.org/pdf/2106.03893) paper. |\n| [MagLaplacianEncoder](https://haiku-geometric.readthedocs.io/en/latest/modules/posenc.html#haiku_geometric.posenc.MagLaplacianEncoder) | Magnetic Laplacian positional encoding from the [Transformers Meet Directed Graphs](https://arxiv.org/pdf/2302.00049) paper. |\n\n## Issues\n\nIf you encounter any issue, please [open an issue](https://github.com/alexOarga/haiku-geometric/issues/new).\n\n## Running tests\n\nHaiku Geometric can be tested using `pytest` by running the following command:\n\n```bash\npython -m pytest test/\n```\n\n## License\n\n[](https://opensource.org/licenses/Apache-2.0)\n\n\n",
"bugtrack_url": null,
"license": "Apache-2.0 license",
"summary": "",
"version": "0.0.5",
"project_urls": {
"Homepage": "https://github.com/alexOarga/haiku-geometric"
},
"split_keywords": [],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "5a03f4644548af56930076d20b128a82b4cd7854faf1770d2e7379a58efff802",
"md5": "60cce99b62affcc9b613b8cd06448bd6",
"sha256": "e0f2365719a5c7db71aec1e5a3a33694876eeacd4da1b65ded56bc624355b387"
},
"downloads": -1,
"filename": "haiku_geometric-0.0.5-py3-none-any.whl",
"has_sig": false,
"md5_digest": "60cce99b62affcc9b613b8cd06448bd6",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": null,
"size": 75449,
"upload_time": "2023-09-26T22:24:10",
"upload_time_iso_8601": "2023-09-26T22:24:10.802597Z",
"url": "https://files.pythonhosted.org/packages/5a/03/f4644548af56930076d20b128a82b4cd7854faf1770d2e7379a58efff802/haiku_geometric-0.0.5-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "63316291b06bc36873b79803dc9eaea88e0a8bdc5062b9595c41dc5b8b82f178",
"md5": "e21ff4377ad6d8f7fbec13c23e56ed7f",
"sha256": "3d3e9cb8412f1d0f4cb5e95d98d0c86d3c3b87401bfb64f67fdeff9b7b7963d7"
},
"downloads": -1,
"filename": "haiku_geometric-0.0.5.tar.gz",
"has_sig": false,
"md5_digest": "e21ff4377ad6d8f7fbec13c23e56ed7f",
"packagetype": "sdist",
"python_version": "source",
"requires_python": null,
"size": 118114,
"upload_time": "2023-09-26T22:24:15",
"upload_time_iso_8601": "2023-09-26T22:24:15.954502Z",
"url": "https://files.pythonhosted.org/packages/63/31/6291b06bc36873b79803dc9eaea88e0a8bdc5062b9595c41dc5b8b82f178/haiku_geometric-0.0.5.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2023-09-26 22:24:15",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "alexOarga",
"github_project": "haiku-geometric",
"travis_ci": false,
"coveralls": false,
"github_actions": true,
"requirements": [],
"lcname": "haiku-geometric"
}
JSON
