| 108 | |
| 109 | |
| 110 | class Sphere(nn.Module): |
| 111 | |
| 112 | def __init__(self, l=2): |
| 113 | super(Sphere, self).__init__() |
| 114 | self.l = l |
| 115 | |
| 116 | def forward(self, edge_vec): |
| 117 | edge_sh = self._spherical_harmonics(self.l, edge_vec[..., 0], edge_vec[..., 1], edge_vec[..., 2]) |
| 118 | return edge_sh |
| 119 | |
| 120 | @staticmethod |
| 121 | def _spherical_harmonics(lmax: int, x: torch.Tensor, y: torch.Tensor, z: torch.Tensor) -> torch.Tensor: |
| 122 | |
| 123 | sh_1_0, sh_1_1, sh_1_2 = x, y, z |
| 124 | |
| 125 | if lmax == 1: |
| 126 | return torch.stack([sh_1_0, sh_1_1, sh_1_2], dim=-1) |
| 127 | |
| 128 | sh_2_0 = math.sqrt(3.0) * x * z |
| 129 | sh_2_1 = math.sqrt(3.0) * x * y |
| 130 | y2 = y.pow(2) |
| 131 | x2z2 = x.pow(2) + z.pow(2) |
| 132 | sh_2_2 = y2 - 0.5 * x2z2 |
| 133 | sh_2_3 = math.sqrt(3.0) * y * z |
| 134 | sh_2_4 = math.sqrt(3.0) / 2.0 * (z.pow(2) - x.pow(2)) |
| 135 | |
| 136 | if lmax == 2: |
| 137 | return torch.stack([sh_1_0, sh_1_1, sh_1_2, sh_2_0, sh_2_1, sh_2_2, sh_2_3, sh_2_4], dim=-1) |
| 138 | |
| 139 | |
| 140 | class VecLayerNorm(nn.Module): |