mindscience.e3nn.so2_conv.SO3Rotation
- class mindscience.e3nn.so2_conv.SO3Rotation(lmax, irreps_in, irreps_out)[源代码]
处理球面谐波不可约表示的 SO(3) 旋转的类。
- 参数:
lmax (int) - 考虑的最大角动量阶数。
irreps_in (Union[str, Irreps]) - 输入不可约表示。
irreps_out (Union[str, Irreps]) - 输出不可约表示。
样例:
>>> from mindscience.e3nn.so2_conv import SO3Rotation >>> rot = SO3Rotation(lmax=2, irreps_in="1x0e + 1x1o", irreps_out="1x1o") >>> wigner, wigner_inv = rot.set_wigner(rot_mat3x3) >>> rotated = rot.rotate(embedding, wigner)
- static narrow(inputs, axis, start, length)[源代码]
沿指定维度对张量进行切片。
- 参数:
inputs (Tensor) - 待切片的张量。
axis (int) - 进行切片的维度。
start (int) - 切片起始位置索引。
length (int) - 切片包含的元素数量。
- 返回:
Tensor,切片后的张量。
- rotate(embedding, wigner)[源代码]
按提供的 Wigner-D 矩阵对嵌入张量进行 SO(3) 旋转。
- 参数:
embedding (Tensor) - 待旋转的输入张量,形状为
(..., irreps_in.dim),包含待旋转的球谐系数。wigner (tuple[Tensor]) - l = 0 … lmax 的 Wigner-D 矩阵元组,每个形状为
(..., 2l+1, 2l+1)。
- 返回:
tuple[Tensor],与 irreps_in 中每个 irrep 对应的一组旋转后张量,每个形状为
(..., mul, 2l+1)。
- rotate_inv(embedding, wigner_inv)[源代码]
使用提供的逆 Wigner-D 矩阵对嵌入张量进行逆旋转。
- 参数:
embedding (tuple[Tensor]) - 与 irreps_out 中每个 irrep 对应的一组张量,每个形状为
(..., mul, 2l+1)。wigner_inv (tuple[Tensor]) - l = 0 … lmax 的逆(转置)Wigner-D 矩阵元组,每个形状为
(..., 2l+1, 2l+1)。
- 返回:
Tensor,逆旋转并拼接后的输出张量,形状为
(..., irreps_out.dim)。
- static rotation_to_wigner_d_matrix(edge_rot_mat, start_lmax, end_lmax)[源代码]
将批量 \(3 \times 3\) 旋转矩阵转换为指定角动量范围的 Wigner-D 矩阵。
- 参数:
edge_rot_mat (Tensor) - 旋转矩阵批次,形状为
(..., 3, 3)。start_lmax (int) - 最小角动量阶数。
end_lmax (int) - 最大角动量阶数。
- 返回:
list[Tensor],对应从 start_lmax 到 end_lmax 的各阶 Wigner-D 矩阵,每个张量形状为
(..., 2l+1, 2l+1)。