Request to Implement GeoGNN

[jasona445]

🚀 The feature, motivation and pitch

I'd like to implement the paper Geometry-enhanced molecular representation learning for property prediction by Fang et al.

In short, they propose a novel architecture called GeoGNN which encodes additional spatial information about molecular geometries by modeling both atom to bond and bond to bond angle relations. They do so by creating two graphs, one for each relation type as shown below

The message passing for graph $G$ (the atom-bond graph) looks like

\begin{aligned}
    \mathbf{a^{(k)}_{u}} &= AGG^{(k)}_G\left( \{ (\mathbf{h^{(k-1)}_u}, \mathbf{h^{(k-1)}_v}, \mathbf{h^{(k-1)}_{uv}} : v \in \mathit{N}(u) \}\right)\\
    \mathbf{h^{(k)}_{u}} &= COMBINE^{(k)}_G (\mathbf{h^{(k-1)}_{u}}, \mathbf{a^{(k)}_{u}})
\end{aligned}

and for graph $H$ (the bond-bond angle graph) it looks like

\begin{aligned}
    \mathbf{a^{(k)}_{uv}} &= AGG^{(k)}_H\left( \{ (\mathbf{h^{(k-1)}_{uv}}, \mathbf{h^{(k-1)}_{uw}},  \mathbf{x_{wuv}}) : w \in \mathit{N}(u) \} \right. \notag \\ 
    & \qquad \qquad \qquad \left. \cup \{ (\mathbf{h^{(k-1)}_{uv}}, \mathbf{h^{(k-1)}_{vw}},  \mathbf{x_{uvw}}) : w \in \mathit{N}(v) \}\right)\\
    \mathbf{h^{(k)}_{uv}} &= COMBINE^{(k)}_H (\mathbf{h^{(k-1)}_{uv}}, \mathbf{a^{(k)}_{uv}})\\
\end{aligned}

They show it does quite well on the MoleculeNet dataset. Would it be possible to contribute it to PyG?

Alternatives

There's a repo here written in paddle.

Additional context

No response

[rusty1s]

Yes, this is a welcome contribution :)