Struct Affine3D 

pub struct Affine3D {
    pub rot: Tensor,
    pub trans: Tensor,
}

Per-residue local reference frame: rotation matrix + translation.

rot: (*, 3, 3) — orthonormal rotation matrix, columns are the frame axes. trans: (*, 3) — CA position in global coordinates.

Fields

rot: Tensor

trans: Tensor

Implementations

impl Affine3D

pub fn new(rot: Tensor, trans: Tensor) -> Self

pub fn apply_rot(rot: &Tensor, v: &Tensor) -> Result<Tensor>

Rotate vectors from local frame to global frame.

rot: (*, 3, 3), v: (*, H, 3) → (*, H, 3).

In row-vector convention: v @ rot^T.

pub fn apply_rot_inv(rot: &Tensor, v: &Tensor) -> Result<Tensor>

Rotate vectors from global frame to local frame (inverse rotation).

rot: (*, 3, 3), v: (*, H, 3) → (*, H, 3).

For orthogonal matrices: inverse = transpose, so v @ rot.

pub fn apply(&self, v: &Tensor) -> Result<Tensor>

Apply full affine transform (rotate + translate) to v of shape (*, H, 3).

Returns (*, H, 3) in global coordinates.

pub fn build_affine3d_from_coordinates( coords: &Tensor, ) -> Result<(Self, Tensor)>

Build per-residue local frames from backbone atom coordinates.

coords: (B, L, 3, 3) — atom order is (N, CA, C).

Frame construction follows AlphaFold2/ESM convention (Graham-Schmidt on CA-C and CA-N):

• x-axis: normalized CA - C
• y-axis: normalized component of N - CA perpendicular to x-axis
• z-axis: x × y
• origin: CA position

Returns (Affine3D, mask) where mask (B, L) is true (u8=1) where both backbone vectors have non-trivial length (atoms are present).