Queries and ideas: whether rascaline.torch.LodeSphericalExpansion() supports GPUs, and the format of the output aligned with e3nn

[clecust]

Idea: the format of the output of rascaline.torch.LodeSphericalExpansion() is aligned with e3nn:

I am interested in combining rascaline.torch.LodeSphericalExpansion() with e3nn(https://github.com/e3nn/e3nn). Both methods are based on spherical harmonics and satisfy the equivariant property. However, inconsistencies in output formats arise, especially when max_radial is greater than 1.

Here is the output format of e3nn:

Irreps("1x0e + 1x1o + 1x0e + 1x1o")

This is another example of e3nn:

Irreps("2x0e + 2x1o ")

Aligning the output formats, particularly for higher-order equivariant models, would enhance compatibility between rascaline and e3nn frameworks. Your assistance in resolving these issues would be highly beneficial, not only for my use case but also for encouraging more users to seamlessly combine these powerful frameworks.

Thank you for your consideration.

[Luthaf]

I split the GPU support discussion in a separate issue (#271).

---

Regarding the output format: rascaline uses metatensor to store it's output, the main reason for this is that e3nn-style storage format is very wasteful, and stores a lot of explicit zeros (all the out-of-diagonal blocks in your plots). These extra explicit zeros add up and fill the RAM/GPU memory, limiting what one can do with some given hardware.

For this reason, it is very unlikely we will change rascaline output format, but we can and probably should provide some ways to adapt between metatensor and e3nn. We would need two functions for this:

from metatensor import TensorMap

def metatentor_to_e3nn(tensor: TensorMap) -> (Irreps, torch.Tensor):
    ...

def e3nn_to_metatensor(irreps: Irreps, data: torch.Tensor) -> TensorMap:
    ...

I'm not fully familiar with e3nn, but is the 2x in 2x1o typically used to indicate 2 radial channels? If so, doing the conversions above should be fairly easy! Also, how does e3nn represents different samples/batches? As an extra first dimension in the tensor, the plots above representing the second dimension?

For reference, metatensor format represent the 0o, 1o, ... part as "keys" of the TensorMap (named spherical_harmonics_l in rascaline), and stores the radial channels as "properties" of the different blocks in the TensorMap.

[clecust]

I split the GPU support discussion in a separate issue (#271).

Regarding the output format: rascaline uses metatensor to store it's output, the main reason for this is that e3nn-style storage format is very wasteful, and stores a lot of explicit zeros (all the out-of-diagonal blocks in your plots). These extra explicit zeros add up and fill the RAM/GPU memory, limiting what one can do with some given hardware.

For this reason, it is very unlikely we will change rascaline output format, but we can and probably should provide some ways to adapt between metatensor and e3nn. We would need two functions for this:

from metatensor import TensorMap

def metatentor_to_e3nn(tensor: TensorMap) -> (Irreps, torch.Tensor):
    ...

def e3nn_to_metatensor(irreps: Irreps, data: torch.Tensor) -> TensorMap:
    ...

I'm not fully familiar with e3nn, but is the 2x in 2x1o typically used to indicate 2 radial channels? If so, doing the conversions above should be fairly easy! Also, how does e3nn represents different samples/batches? As an extra first dimension in the tensor, the plots above representing the second dimension?

For reference, metatensor format represent the 0o, 1o, ... part as "keys" of the TensorMap (named spherical_harmonics_l in rascaline), and stores the radial channels as "properties" of the different blocks in the TensorMap.

Thank you very much for your quick reply． Regarding the above, I'm sorry I didn't describe it clearly．

Let me briefly answer your confusion: the above figure shows the Wigner-D matrix in the output space; as for the dimensions of e3nn, they are very close to those of the rascaline output (block.values), except that the components of $Y(l>0)$ are arranged in a different order in the multiple (radial) channels. (Also, e3nn's default Euler angle convention is changed from the standard ZYZ to YXY, and for first-order spherical harmonic functions, the component ordering is [x,y,z] instead of [y,z,x]; https://docs.e3nn.org/en/latest/guide/change_of_basis.html)

I went through the Wigner-D matrix of the output space to get an idea of the output format of rascaline and e3nn, because I wanted to try to combine the two software.
I used torch.matmul(torch.linalg.LA.pinv(y_inital), y_rotated).T to get the Wigner-D matrix of rascaline as follows(the rotations matrix are randomly generated and do not correspond to the values in the image above).

For the lode set by the descriptors of the two radial channels and the highest 1st order spherical harmonic function, I surmised based on Wigner-D that its output format should be $Y_{n=0,l=0}^{m=0},Y_{n=1,l=0}^{m=0},Y_{n=0,l=1}^{m=-1},Y_{n=1,l=1}^{m=-1},Y_{n=0,l=1}^{m=0},Y_{n=1,l=1}^{m=0}...$

e3nn can provide two similar expressions, but with different output formats (presumably based on the figure from the original comment), especially $Y(l>0)$, including Irreps("1x0e + 1x1o + 1x0e + 1x1o ") and Irreps("2x0e + 2x1o ").　

So it might be worth considering to add an official interface to achieve compatibility between e3nn and rascaline.
Thank you very much for your contribution and reply.

[Luthaf]

to get the Wigner-D matrix of rascaline as follows(the rotations matrix are randomly generated and do not correspond to the values in the image above).

Could you share the full script you used to generate these plots?

For the lode set by the descriptors of the two radial channels and the highest 1st order spherical harmonic function, I surmised based on Wigner-D that its output format should be ...

I think (but I need to check with the script you used) that this is an artefact of how you constructed a single matrix from the block structure of metatensor TensorMap. For example, the code you initially shared

lrvp = lr.keys_to_samples("species_center")
lrvp = lrvp.components_to_properties(["spherical_harmonics_m"])
lrvp = lrvp.keys_to_properties(["species_neighbor", "spherical_harmonics_l"])

Will mix species_neighbor and spherical_harmonics_m which you typically don't want. More generally, metatensor/rascaline store the data for different λ separate, in 3D tensors with the shape [atoms, μ, radial-channels], and components_to_properties just reshape that to [atoms, μ * radial-channels], which would explain the pattern you see in the Wigner-D matrix.

You can triple check this by looking at the metadata attached to the different axes of the array: lrvp.block().samples for the first dimension; lrvp.block().properties for the last dimension and lrvp.block().components for any middle dimension.

---

Also, e3nn's default Euler angle convention is changed from the standard ZYZ to YXY, and for first-order spherical harmonic functions, the component ordering is [x,y,z] instead of [y,z,x]; https://docs.e3nn.org/en/latest/guide/change_of_basis.html

On this front, rascaline uses the same convention as wikipedia for spherical harmonics: https://en.wikipedia.org/wiki/Spherical_harmonics#Real_form; and the ZYZ convention for Euler angles.