Higher Order Tensor for NequIP Edge Feature?

[blannek]

The edge feature suggested for NequIP as in the paper the tutorial is the distance vectors, scaled by cutoff. I intend to put more information into the edge feature, which is a second-order tensor, a 3x3 matrix. I wonder whether it's feasible, and if so, is there any requirement for this tensor?

[mariogeiger]

In nequip there is also

1. a set of radial basis functions
2. Spherical harmonics typically up to l=3

And I'm sorry to disappoint you but you need to know that the spherical harmonics l=0 and l=2 together contain the exact same information as in the 3x3 outer product. And changing the format will likely have no effect.

[blannek]

Thanks a lot for your response! There seems a bit misunderstanding. I could understand the edge vector is further expanded by the spherical harmonics, whose weight are defined by the radial basis functions. I am not talking about the embedding of edge features, instead I am focusing on the feature itself. My question is whether the input edge vector could be extended to the form other from '1o'? Say, apart from the distance vector between atoms, I want to add a weight into the vector so that it becomes '1o + 0e' (in my case I have a 3x3 symmetric matrix), then further inputed into NequIP layer, expanded by the spherical harmonics. Would that break any rules in NequIP, or in other words, does NequIP assume that the input edge vector must be '1o'? I know that it assumes the edge vector ranges from 0.0 to 1.0, I am just concerning something similar to that.

And I'm sorry to disappoint you but you need to know that the spherical harmonics l=0 and l=2 together contain the exact same information as in the 3x3 outer product. And changing the format will likely have no effect.

And a point here, it seems in current NequIP layer the spherical harmonics only expand from l=1 to max_ell. Does it implicitly include the l=0, or to some degree they are equivalent?

[mariogeiger]

Ok sorry for my erroneous assumption.

The line https://github.com/mariogeiger/nequip-jax/blob/main/nequip_jax/nequip.py#L127 is here to replace the tensor product with the spherical harmonic l=0 which is just 1. Instead of doing x*1 I just put x.

If you want to add extra edge features it should be quite simple. I would recommend creating an edge features tensor containing everything:

• the spherical harmonics of the relative vectors
• along with the decomposition of your 3x3 symmetry matrix into an l=0 and a l=2 irrep.

Then at line 128 take the tensor product of the node features with the edge features.

[blannek]

I see the point. Thanks a lot for the reponse!

[blannek]

Sorry for reopening this, but I have encountered another issue. Now the case is the extra edge feature creates a different graph topology as direct coordinates, as the extra edge features on some edges are just 0. Also the coodinate information is already implicitly contained in the extra edge features. Therefore I just decided only to use my 3x3 symmetric traceless matrix (decomposed to '2e' irrep as I found it traceless) as my edge feature. I am wondering whether this is rational? Seems it's not breaking any equivariance, but it seems reducing the complexity of tensor product a lot (as I also find spherical harmonic only accept 1o input). In general, should I do this, or just keeping all the 0 edge features and combined with coordinate vectors?

[blannek]

After attempted, I decided to use spherical harmonic vectors along with the extra 2e irrep. Please ignore the above question.

Now a short question, in radial MLP, do I also need to put the extra information into the lengths = e3nn.norm(vectors).array, with vectors replaced by a concatenated irrep as lengths = e3nn.norm(e3nn.concatenate((vectors, tensors)), per_irrep=False).array, or just use the radial function purely of $r_{ij}$?

[mariogeiger]

Yes you can. I typically concatenate non scalar edge features with the sh and concatenate the scalar edge features with the radial embeddings.