AF3 Rigid3Array类介绍

AlphaFold3  geometry.rigid_matrix_vector 模块的 Rigid3Array类表示 3D 刚体变换，用于描述 旋转 + 平移 这种常见的空间操作，在 AlphaFold3 里主要用于 蛋白质结构的旋转对齐、坐标变换、骨架参考系生成。Rigid3Array 是一个用于表示三维刚性变换（即特殊欧几里得群 SE(3) 的元素）的类。

主要组件

• rotation: rotation_matrix.Rot3Array
3x3 旋转矩阵，负责旋转对象。

• translation: vector.Vec3Array
3D 平移向量，负责平移对象。

这个类基本上在实现 特殊欧几里得群 SE(3) 里的变换操作（3D 空间的刚体运动）。

源代码：

@dataclasses.dataclass(frozen=True)
class Rigid3Array:
    """Rigid Transformation, i.e. element of special euclidean group."""

    rotation: rotation_matrix.Rot3Array
    translation: vector.Vec3Array

    def __matmul__(self, other: Rigid3Array) -> Rigid3Array:
        new_rotation = self.rotation @ other.rotation  # __matmul__
        new_translation = self.apply_to_point(other.translation)
        return Rigid3Array(new_rotation, new_translation)

    def __getitem__(self, index) -> Rigid3Array:
        return Rigid3Array(
            self.rotation[index],
            self.translation[index],
        )

    def __mul__(self, other: torch.Tensor) -> Rigid3Array:
        return Rigid3Array(
            self.rotation * other,
            self.translation * other,
        )

    def map_tensor_fn(self, fn) -> Rigid3Array:
        return Rigid3Array(
            self.rotation.map_tensor_fn(fn),
            self.translation.map_tensor_fn(fn),
        )

    def inverse(self) -> Rigid3Array:
        """Return Rigid3Array corresponding to inverse transform."""
        inv_rotation = self.rotation.inverse()
        inv_translation = inv_rotation.apply_to_point(-self.translation)
        return Rigid3Array(inv_rotation, inv_translation)

    def apply_to_point(self, point: vector.Vec3Array) -> vector.Vec3Array:
        """Apply Rigid3Array transform to point."""
        return self.rotation.apply_to_point(point) + self.translation

    def apply(self, point: torch.Tensor) -> torch.Tensor:
        return self.apply_to_point(vector.Vec3Array.from_array(point)).to_tensor()