e3nn/e3nn

adding scalar edge feature in Rs_in❓ [QUESTION]

Chen-Cai-OSU opened this issue · 24 comments

Chen-Cai-OSU commented

Hello,

Previously I have node features (of dimension=5), so I am using Rs_in = [(5, 0, 1)] . Recently I get some new scalar features on edges, I was wondering how to specify these extra edge features in Rs_in? Thank you very much!

Chen-Cai-OSU commented

In other words, previously I have 3 dim features (coordinate difference) on edges. Now I have one extra scalar feature on the edges. I was wondering how to handle this?

mariogeiger commented

Hello,

You have to modify the code here https://github.com/e3nn/e3nn/blob/main/e3nn/nn/models/gate_points_2102.py#L314

Something like

edge_attr = torch.cat([edge_attr, your_own_attributes], dim=1)

And in the same time change irreps_edge_attr into Irreps.spherical_harmonics(lmax) + Irreps("10 x 0e")

Chen-Cai-OSU commented

Thanks for the quick response. I am using a quite old version of equivariant nets at master branch. https://github.com/e3nn/e3nn/blob/master/e3nn/networks/point.py#L62

I plan to switch to the main branch in order to handle the extra edge features. What is the major change from master to the main branch? Is there any backward compatibility issue I should be aware of?

blondegeek commented
mariogeiger commented

For instance what you want to do would have been non trivial with the old version

Chen-Cai-OSU commented

Thank you both! I am going to switch to the new version. Will let you know if I have any questions.

Chen-Cai-OSU commented

Hello,

I have some quick questions/comments:

  1. I saw in many cases there are three arguments for input irreps: a) irresp_in b) irreps_node_attr and c) irreps_edge_attr. What is the relation between a) and b,c?

  2. In the example, teris_gate.py, the API of Gate class is quite complicated. I tried to understand it via documents but it is quite difficult to me. If Gate is an nonlinear activation function, maybe in the example it can be replaced by a simpler nonlinear function? This will make it more accessible to new users

  3. I suggest to simplify the code (maybe into less than 50 lines?) in examples and add an example handling qm9 types of dataset. (dataset with both node and edge features). Since this can not be done easily in the previous version of e3nn, I feel adding an example will be good to showcase the flexibility of new version.

blondegeek commented
Chen-Cai-OSU commented

Hello,

Again thank you both for the quick response. I would like to provide some feedback on the current API design. I am a computer scientist who is quite enthusiastic about equivariant nets.

Conceptually, when I am thinking about equivariant nets, I thought there are mainly the following components

  1. linear equivairant layer

  2. nonlienar equivairant layer

  3. equivariant attention layer (this is probably can be considered in the future)

So far it is no easy to build 1) and 2) easily. By easy I mean to achieve this in 5 lines of code.

  • For linear equivariant layer, I thought there are mainly the following arguments user need to provide: 1) irreps of nodes/edges and output irreps, 2) how to parametrize the radial function with a small neural nets
  • For the nonlinear layer, there are many options. Gated nonlinearity is quite complex and IMO the simpler one such as norm nonlinearity can be implemented first

Of course, there are massive amounts of details (such as CG coeffients) I didn’t mention but as a user, I think all they need to do is to import the right layer and specify the right irreps.

I have to say the current examples is quite difficult to the new user. I spend a couple of months on the mathematics of equivariant nets. I watched most of the talks and went through Tess’s IPAM tutorial. But I still find it difficult to start to use the new version of e3nn. A disclaimer: I only glanced through examples and the examples in model/nn and documents.

Some suggestions:

  • there are too many details exposed to the user in the forward method. I feel much work can be abstracted into nn equivariant layers and users just need to call the right layer in the forward method. For example, the scatter function in torch_scatter can be only used in linear equivariant layers, and users only need to call the linear equivariant layers. I quickly checked torch_geometric and this also seems to be the case.
  • I find the number of arguments to initializae the network class in nn/model is overwhelming. I don’t have a good solution though.
  • add more examples such as qm9, equivariant autoencoder, spherical mnist, use equivariant nets to parameterized Hamiltonian, etc. I believe this is crucial for the wide use of e3nn.

I hope feedbacks above can be at least marginally helpful. I am also thinking about which part I can contribute. Will let you later!

mariogeiger commented

Thanks a lot for the feedback.

I understand to concern of the end-user that don't want to dig into the technicalities.

The problem is that there is a very wide spectrum of possibilities between

  • our aim to make the library general
  • make the library accessible

There is many different possible architecture and we don't pretend to know what architecture is best for the end user. So if we provide a network with everything preset it's cool for the new user but can be misleading because it's not at all the only and best option.

The solution I see is that I try to modularize thing is layer of complexity.
Based on your concern I'm now working on gate_points_2103

  • I will convert Network into MessagePassing that will be generic
  • then I will implement very simple networks that will call the generic message passing

I know it probably not 100% align with your vision but I hope it will go a little bit into the good direction.

mariogeiger commented

Have a look at this network for newbies: https://docs.e3nn.org/en/latest/api/nn/models/v2103.html (might need for readthedocs to compile the last version)

image

Chen-Cai-OSU commented

@mariogeiger Thanks a lot! Your speed of refactoring code is amazing. The code in v2103 is much easier to follow than the previous version. I also agree with what you said about offering users the flexibility to build their own network easily, instead of providing a few suboptimal ones.

  • Is the MessagePassing class in gate_points_message_passing going to be a base class? As pytorch_geometric use a similar MessagePassing as the base class and various convolution layer (GCNconv, sageconv...) only modify a few methods, it is probably worth considering following the similar convention in e3nn. https://pytorch-geometric.readthedocs.io/en/latest/notes/create_gnn.html?highlight=MessagePassing

  • maybe Convolution class in points_convolution.py can be merged with MessagePassing class? I am not too sure about this one.

  • what is the edge length embedding in the forward method of SimpleNetwork class doing?

mariogeiger commented

Convolution is of the form of the MessagePassing class of pytorch_geometric. The MessagePassing class of pytorch_geometric is not providing much, it's just introducing an elegant notation that I prefer not to use.

The code:

edge_length_embedding = soft_one_hot_linspace(
            edge_length,
            0.0,
            self.max_radius,
            self.number_of_basis,
            base='cosine',  # the cosine basis with endpoint = True goes to zero at max_radius
            endpoint=False,  # no need for an additional smooth cutoff
        ).mul(self.number_of_basis**0.5)

project the edge length on a basis of gaussians like that:
image

It is then provided to MessagePassing who gives it to Convolution under the argument named edge_scalars,

weight = self.fc(edge_scalars)
...
edge_features = self.tp(node_features[edge_src], edge_attr, weight)
Chen-Cai-OSU commented

Just stumbled on https://github.com/QUVA-Lab/e2cnn. It has very detailed documents and examples, which I feel might also be useful to the design of e3nn as well.

Chen-Cai-OSU commented

Hello,

I think I am still having trouble on using updated e3nn. Right now I don't care too much about the architecture, as long as it is equivalent and has the right input/output irreps then I am fine. But even achieving this is difficult. The trouble I am having is

  • following your suggestion, I am using Network in gate_points_2012.py. I still don't understand what is irreps_node_attr and irreps_edge_attr and their relation to irreps_in?
  • My output is a symmetric 3x3 matrix. Previously I am using CartesianTensor(torch.rand(3, 3), 'ij=ji').to_irrep_transformation(), to get the Rs_out and Q (a matrix to allow me convert irrep_label to cart_label). What is the such method in the new CartesianTensor(...) class?

Specifically, I have some graph (not point clouds) with both node and edge features. Besides the node positions, all features are scalars. The output irreps in old notation is [(1, 0, 1), (1, 2, 1)].

Is it possible to provide an example for handling such data? I think it covers some important user cases where 1) output is not a scalar 2) there are edge features of multiple types (both scalars and the position difference, and previously this can not be handled?). I have spent two days browsing the documents and code but I still don't think I can figure this out.

Chen-Cai-OSU commented

The general comments I had on the current API of e3nn. Again, this is just my experience as a user.

  • I think there are many design choices from e2cnn can be borrowed to e3nn. Especially, the nn module IMO should be organized into different types of equivariant layers.

  • The documents is hard to understand and confusing sometimes. For example, I can’t understand the example o3.ElementwiseTensorProduct('5x0e + 5x1e', '4x0e + 6x1e').visualize()

    https://docs.e3nn.org/en/latest/api/o3/o3_tp.html. What does those 1x0e and 1x1e comes from?. Another example is e3nn.o3.irreps at https://docs.e3nn.org/en/latest/api/o3/o3_irreps.html. The dim argument does not seem to be used in the examples.

My suggestions

  • add a caveat section in the README detailing what parts of the library is under construction.
  • introduce more examples. IMO this is probably the most important thing for the new user.

Again, I hope my complaints can play a constructive role in the development of e3nn. I have great respect for your work.

mariogeiger commented 
x = o3.ReducedTensorProducts("ij=ji", i="1o")
x.change_of_basis
mariogeiger commented

o3.ElementwiseTensorProduct needs to split inputs into similar multiplicities because of the way TensorProduct works.

mariogeiger commented

add a caveat section in the README detailing what parts of the library is under construction.

The hole library 🤣

  • 😄1
mariogeiger commented

I will write later a network with

irreps_node_input : input messages for the message passing
irreps_node_attr : attributes that are constant layer-wise
irreps_edge_attr : extra (in addition to spherical harmonics) attributes for the edges

doc here

blondegeek commented

Hi @Chen-Cai-OSU,

I agree with you that e2cnn has more complete documentation. But, I'd be careful in comparing e2cnn to e3nn. The aims and therefore challenges of these libraries are different.

From a mathematical aspect, there are many aspects of going to SO(2) -> SO(3) that substantially complicate things -- namely requiring the TensorProduct (because our irreps are multi-dimensional) which is a challenging operation to not only articulate but also optimize using out of the box PyTorch operations.

Additionally, I believe (and I could be wrong about this) the focus of e2cnn is computer vision tasks -- whereas with e3nn we are mostly serving a community of scientists (physicists, chemists, biologists) that are often dealing with non-standard data-types and tasks (e.g. point clouds of atoms with connectivity defined in different ways depending on the chemistry). This is why we have spent less time creating complete documentation and more time creating what we view as the most flexible modules. Flexibility and user-friendliness are often at odds with each other and is something we are still figuring out how to balance.

This is not an excuse for not having documentation, but rather to hopefully explain how our priorities may differ from other repositories. We of course would like to improve our documentation and so I thank you for your feedback, but please be patient with us as we are juggling many different types of demands (e.g. improving initialization, representing scientific data that require additional infrastructure, creating well-behaved algorithms for manipulating geometry).

Best,
Tess

  • 👍1
Chen-Cai-OSU commented

Thank you, Tess. I totally understand it. Will wait patiently for future improvement!