Question about the edge_rot_mat

[HeegerGao]

Hi, thanks for your wonderful work and code! I am new to the equivariant learning area, and I am trying to understand each step of your code. I suspect this function (the function for calculating the edge rotation matrix) would break the equivariance during the edge-degree embedding layer.

In this function, you set the original edge_distance_vec as the final x_axis, and randomly select the y-axis and generate the z_axis according to the other two axes. I think this random operation is not correct.

For example, suppose I send two same molecules into the network, where the coordinates of one molecule are translated relative to the other molecule, but not rotated. If the network is SE(3)-equivariant, the final feature should be the same for these two molecules. However, during the above function, they will have different edge rotation matrices on the same edge since the randomly selected y_axis, and the corresponding Wigner-D matrix would be different, thus during the edge-degree embedding layer, the embedding of the two molecules are different since this step uses the Wigner-D matrix. This breaks the equivariance for all features with type>0.

I am not sure if I am correct. Looking forward to your opinion on this issue.