One Lie-group Euler step for frame flow

Your task

Function signature

lie_euler(
    x_t: torch.Tensor,
    x1_hat: torch.Tensor,
    R_t: torch.Tensor,
    R1_hat: torch.Tensor,
    t: torch.Tensor,
    dt: torch.Tensor,
    valid_mask: torch.Tensor,
    diffuse_mask: torch.Tensor,
    rotation_schedule: str = 'linear',
    exp_rate: float = 10.0,
) -> tuple[torch.Tensor, torch.Tensor]

Worked example

Advance only part of the remaining path

Parameters

Return

tuple[torch.Tensor, torch.Tensor]

Required behavior

• Inputs x_t,x1_hat [B,N,3], R_t,R1_hat [B,N,3,3], t,dt [B], and two bool masks.
• For 0\le t<1 and 0\le \mathrm{dt}\le 1-t, f=\mathrm{dt}/(1-t), x_{\mathrm{next}}=x_{t}+f\cdot (\mathrm{x1\_hat}-x_{t}). Evaluate this as (1-f)\cdot x_{t}+f\cdot \mathrm{x1\_hat} to avoid overflowing a residual whose endpoints are finite.
• Set \omega=\operatorname{Log}(R_{t}^{\top} \mathrm{R1\_hat}) in the LOCAL/BODY frame and R_{\mathrm{next}}=R_{t} \operatorname{Exp}(g\cdot \omega). rotation_schedule='linear' (default) uses g=f.
• The alternate 'exp' uses g=\mathrm{dt}\cdot \mathrm{exp\_rate}, \mathrm{exp\_rate}>0 default 10; translation still uses the linear factor.
• This exponential rotation scaling corresponds to \mathrm{kappa}(t)=1-\operatorname{exp}(-\mathrm{exp\_rate}\cdot t), whose kappa'/(1-\mathrm{kappa}) is constant; it need not hit the predicted rotational endpoint at t=1.
• In the linear variant \mathrm{dt}=1-t reaches both supplied endpoints exactly (up to roundoff) for proper rotations away from the cut locus. At t=1 only \mathrm{dt}=0 is valid: return the current valid frames unchanged, and never evaluate 1/(1-t).
• Valid fixed residues always retain their CURRENT x_t/R_t; invalid positions become zero and invalid rotations identity. This is a positive-time flow step; do not insert the minus sign of a reverse-time diffusion integrator. No \epsilon, stochastic noise or post-update orthogonalization.
• The only keys are rotation_schedule and exp_rate (used only for 'exp'). 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 lie_euler(
    x_t: torch.Tensor,  # (B, N, 3)
    x1_hat: torch.Tensor,  # (B, N, 3)
    R_t: torch.Tensor,  # (B, N, 3, 3)
    R1_hat: torch.Tensor,  # (B, N, 3, 3)
    t: torch.Tensor,  # (B,)
    dt: torch.Tensor,  # (B,)
    valid_mask: torch.Tensor,  # (B, N)
    diffuse_mask: torch.Tensor,  # (B, N)
    rotation_schedule: str = 'linear',
    exp_rate: float = 10.0,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return tuple[torch.Tensor, torch.Tensor]:
    
    x_next: (B, N, 3).
    R_next: (B, N, 3, 3)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement lie_euler")

Sources and further reading

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