SO(3) interpolation and tangent target

Your task

Function signature

so3_geodesic(
    R0: torch.Tensor,
    R1: torch.Tensor,
    t: torch.Tensor,
    valid_mask: torch.Tensor,
    diffuse_mask: torch.Tensor,
) -> dict[str, torch.Tensor]

Worked example

Half of a quarter-turn

import torch

R0 = torch.tensor([[[[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]]], dtype=torch.float64, requires_grad=True)
R1 = torch.tensor([[[[0.0, -1.0, 0.0], [1.0, 0.0, 0.0], [0.0, 0.0, 1.0]]]], dtype=torch.float64, requires_grad=True)
t = torch.tensor([0.5], dtype=torch.float64, requires_grad=True)
valid_mask = torch.tensor([[True]], dtype=torch.bool)
diffuse_mask = torch.tensor([[True]], dtype=torch.bool)

result = so3_geodesic(R0=R0, R1=R1, t=t, valid_mask=valid_mask, diffuse_mask=diffuse_mask)

Expected return

{
    'R_t': torch.tensor([[[[0.7071067811865475, -0.7071067811865475, 0.0], [0.7071067811865475, 0.7071067811865475, 0.0], [0.0, 0.0, 1.0]]]], dtype=torch.float64),
    'body_velocity': torch.tensor([[[0.0, 0.0, 1.5707963267948966]]], dtype=torch.float64),
    'tangent_residual': torch.tensor([[[0.0, 0.0, 0.7853981633974482]]], dtype=torch.float64),
}

Parameters

Return

dict[str, torch.Tensor]

Required behavior

• Given \mathrm{R0},\mathrm{R1} [B,N,3,3], t [B] in [0,1], valid_mask and diffuse_mask [B,N], define \omega=\operatorname{Log}(\mathrm{R0}^{\top} \mathrm{R1}), R_{t}=\mathrm{R0} \operatorname{Exp}(t\cdot \omega).
• Return R_t, \mathrm{body\_velocity}=\omega and \mathrm{tangent\_residual}=\operatorname{Log}(R_{t}^{\top} \mathrm{R1}). These are LOCAL/BODY tangent coordinates: right-multiplication applies the increment, and the matrix derivative is R_t[\omega]_\times.
• On the chosen shortest geodesic the residual is (1-t)\cdot \omega, but compute the log of R_{t}^{\top} R1 for its specified ambient gradient. Do not divide by 1-t; the residual is zero at t=1 and the path velocity remains \omega. Valid fixed frames are R1 at all times, with both tangent outputs zero.
• Invalid frames are identity with zero tangents; replace inactive matrices by identity before logarithms. Return complete frames, not update matrices. No keys. All masks are bool: valid_mask=True denotes a real residue/atom; diffuse_mask=True denotes a movable one. Their conjunction selects updates.
• B,N (or A) are positive; an all-false mask, not a zero-length tensor, represents an empty example. All supplied numbers, including padding, are finite and share dtype/device; floating tensors are float32 or float64. Rotations act on column vectors and map local coordinates to the global frame.
• For v=(x,y,z), \operatorname{hat}(v)=\begin{pmatrix}0&-z&y\\z&0&-x\\-y&x&0\end{pmatrix}. \operatorname{Exp}(v)=I+A\operatorname{hat}(v)+B\operatorname{hat}(v)^2, where q=v\cdot v, A=\frac{\sin\sqrt q}{\sqrt q}, B=\frac{1-\cos\sqrt q}{q}.
• For q<10^{-4} use A=1-\frac q6+\frac{q^2}{120}-\frac{q^3}{5040} and B=\frac12-\frac q{24}+\frac{q^2}{720}-\frac{q^3}{40320}; at zero \operatorname{Exp}(0)=I.
• For \operatorname{Log}(R), set v=\frac12\operatorname{vee}(R-R^\top), q=v\cdot v, s=\sqrt{q}, c=\frac{\operatorname{tr}R-1}{2}, \theta=\operatorname{atan2}(s,c). Ordinarily \operatorname{Log}(R)=\frac{\theta}{s}v. If q<10^{-8} and c>0 use v\left(1+\frac q6+\frac{3q^2}{40}+\frac{5q^3}{112}\right).
• If q<10^{-6} and c<0, recover an axis from P=\frac{R+R^\top-2cI}{2(1-c)}: choose the largest diagonal (\mathrm{ties} x, \mathrm{then} y, \mathrm{then} z), take its positive square root, and divide its row by that component; normalize the axis.
• Orient it to have nonnegative dot with v when q>10^{-20}; otherwise keep the chosen component positive. Return \theta\cdot \mathrm{axis}.
• For dormant branches clamp 1-c below by 10^{-6} and diagonal square-root arguments below by 10^{-12}; these floors do not alter the selected near-\pi axis. Evaluate every inactive square root/division safely before torch.where. R must be a proper \operatorname{SO}(3) matrix; no projection/repair is requested.
• At \pi axis sign is inherently ambiguous: \operatorname{exp}(\operatorname{log}(R)) is the geometric invariant. No global continuous logarithm exists at the \pi cut locus.

Why this matters

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def so3_geodesic(
    R0: torch.Tensor,  # (B, N, 3, 3)
    R1: torch.Tensor,  # (B, N, 3, 3)
    t: torch.Tensor,  # (B,)
    valid_mask: torch.Tensor,  # (B, N)
    diffuse_mask: torch.Tensor,  # (B, N)
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    R_t: (B, N, 3, 3).
    body_velocity: (B, N, 3).
    tangent_residual: (B, N, 3)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement so3_geodesic")

Sources and further reading

• microsoft/protein-frame-flow · so3_utils.py
• microsoft/protein-frame-flow · interpolant.py
• RosettaCommons/RFdiffusion2 · noisers.py