Plan to add an e3nn model

[mariogeiger]

🚀 Feature

Dear @rusty1s, What do you think about adding a Model similar to Schnet that has equivariance using higher dimensional representations. To do so I extracted the essential part of e3nn into this repo and copied your implementation of Schnet and changed that architecture with mine.

Motivation

Adding our equivariant model and part of its framework to your nice framework.

Additional context

• e3nn_little is still quite some amount of lines of code, how open are you to integrating stuff? Shall I reduce even more (might be possible)
• I still have to improve the architecture, so far it's not at the same perf as Schnet
• The model is here

[rusty1s]

That sounds like a great contribution. Thanks you very much in advance.

I'm definitely open to integrating. As far as I can see, the amount of dependencies is manageable. One could even think of making some of the dependencies optional (like we do in DimeNet).

On the other hand, there are a lot of utility files/functions. Currently, I'm not sure how to handle those properly. Dependent on how specific those are, one could imagine including them in torch_geometric.nn.models.e3nn or in torch_geometric.utils.e3. WDYT?

As an alternative, I see that e3nn is installable via pip. It might be an option to require people to install this package, while having access to a simple model and example in PyTorch Geometric.