RFDiffusion: SE(3) Diffusion for Protein Backbones

import torch
import numpy as np

def skew_symmetric(v):
    """
    Construct the 3x3 skew-symmetric matrix [v]_x from a 3D vector v.

    For a vector v = (v1, v2, v3), the skew-symmetric matrix satisfies
    [v]_x @ w = v x w (cross product) for any vector w.

    Args:
        v: [..., 3] axis vectors (e.g., rotation axes for each residue)

    Returns:
        K: [..., 3, 3] skew-symmetric matrices
    """
    batch_shape = v.shape[:-1]
    zero = torch.zeros(*batch_shape, device=v.device)

    return torch.stack([
        torch.stack([zero,    -v[..., 2],  v[..., 1]], dim=-1),
        torch.stack([v[..., 2],  zero,    -v[..., 0]], dim=-1),
        torch.stack([-v[..., 1], v[..., 0],  zero   ], dim=-1)
    ], dim=-2)


def axis_angle_to_rotation_matrix(axis_angle):
    """
    Convert axis-angle vectors to rotation matrices via the Rodrigues formula.

    Each residue's orientation perturbation can be expressed as a rotation
    by angle omega about axis n_hat. This function converts that compact
    representation into the full 3x3 matrix needed for coordinate transforms.

    Args:
        axis_angle: [..., 3] vectors where direction = axis, magnitude = angle

    Returns:
        R: [..., 3, 3] rotation matrices
    """
    angle = axis_angle.norm(dim=-1, keepdim=True)  # rotation angle omega
    axis = axis_angle / (angle + 1e-8)              # unit axis n_hat

    K = skew_symmetric(axis)  # [..., 3, 3]

    I = torch.eye(3, device=axis_angle.device).expand(
        *axis_angle.shape[:-1], 3, 3
    )
    sin_angle = torch.sin(angle).unsqueeze(-1)
    cos_angle = torch.cos(angle).unsqueeze(-1)

    # Rodrigues: R = I + sin(omega) * K + (1 - cos(omega)) * K^2
    R = I + sin_angle * K + (1 - cos_angle) * torch.einsum(
        '...ij,...jk->...ik', K, K
    )
    return R

def quaternion_to_rotation_matrix(q):
    """
    Convert unit quaternions to rotation matrices.

    Quaternions avoid gimbal lock and provide numerically stable
    interpolation between orientations --- critical when iteratively
    updating residue frames during diffusion.

    Args:
        q: [..., 4] quaternions in (w, x, y, z) convention

    Returns:
        R: [..., 3, 3] rotation matrices
    """
    # Normalize to unit quaternion (guards against numerical drift)
    q = q / (q.norm(dim=-1, keepdim=True) + 1e-8)
    w, x, y, z = q.unbind(-1)

    R = torch.stack([
        torch.stack([1 - 2*y*y - 2*z*z, 2*x*y - 2*w*z, 2*x*z + 2*w*y], dim=-1),
        torch.stack([2*x*y + 2*w*z, 1 - 2*x*x - 2*z*z, 2*y*z - 2*w*x], dim=-1),
        torch.stack([2*x*z - 2*w*y, 2*y*z + 2*w*x, 1 - 2*x*x - 2*y*y], dim=-1)
    ], dim=-2)
    return R


def rotation_matrix_to_quaternion(R):
    """
    Convert rotation matrices to unit quaternions.

    This uses the Shepperd method based on the matrix trace. The
    implementation handles the common case where the trace is positive;
    production code should also handle degenerate cases.

    Args:
        R: [..., 3, 3] rotation matrices

    Returns:
        q: [..., 4] unit quaternions (w, x, y, z)
    """
    trace = R[..., 0, 0] + R[..., 1, 1] + R[..., 2, 2]

    w = torch.sqrt(torch.clamp(1 + trace, min=1e-8)) / 2
    x = (R[..., 2, 1] - R[..., 1, 2]) / (4 * w + 1e-8)
    y = (R[..., 0, 2] - R[..., 2, 0]) / (4 * w + 1e-8)
    z = (R[..., 1, 0] - R[..., 0, 1]) / (4 * w + 1e-8)

    q = torch.stack([w, x, y, z], dim=-1)
    return q / (q.norm(dim=-1, keepdim=True) + 1e-8)

class RigidTransform:
    """
    Rigid body transformation (rotation + translation) for protein residues.

    Each residue in a protein backbone is described by a frame T = (R, t),
    where R is a 3x3 rotation matrix defining the local coordinate system
    and t is a 3D vector giving the C-alpha position.
    """

    def __init__(self, rotations, translations):
        """
        Args:
            rotations:    [..., 3, 3] rotation matrices (one per residue)
            translations: [..., 3]    C-alpha positions (one per residue)
        """
        self.rots = rotations
        self.trans = translations

    @classmethod
    def identity(cls, batch_shape, device='cpu'):
        """Create identity frames (no rotation, origin position)."""
        rots = torch.eye(3, device=device).expand(*batch_shape, 3, 3).clone()
        trans = torch.zeros(*batch_shape, 3, device=device)
        return cls(rots, trans)

    def compose(self, other):
        """
        Compose two transformations: self followed by other.
        Result: (R1 @ R2, R1 @ t2 + t1)
        """
        new_rots = torch.einsum('...ij,...jk->...ik', self.rots, other.rots)
        new_trans = (
            torch.einsum('...ij,...j->...i', self.rots, other.trans)
            + self.trans
        )
        return RigidTransform(new_rots, new_trans)

    def apply(self, points):
        """Apply transformation to points: R @ x + t."""
        return (
            torch.einsum('...ij,...j->...i', self.rots, points) + self.trans
        )

    def invert(self):
        """Compute inverse transformation: (R^T, -R^T @ t)."""
        inv_rots = self.rots.transpose(-1, -2)
        inv_trans = -torch.einsum('...ij,...j->...i', inv_rots, self.trans)
        return RigidTransform(inv_rots, inv_trans)

    def to_tensor_7(self):
        """Convert to 7D representation (quaternion + translation)."""
        quat = rotation_matrix_to_quaternion(self.rots)
        return torch.cat([quat, self.trans], dim=-1)

    @classmethod
    def from_tensor_7(cls, tensor):
        """Create from 7D representation (quaternion + translation)."""
        quat = tensor[..., :4]
        trans = tensor[..., 4:]
        rots = quaternion_to_rotation_matrix(quat)
        return cls(rots, trans)

def sample_igso3(shape, sigma, device='cpu'):
    """
    Sample rotation matrices from the Isotropic Gaussian on SO(3).

    This is the rotational analogue of sampling from a Gaussian in
    Euclidean space. Each sample is a random rotation whose angle
    (roughly) follows a Gaussian with standard deviation sigma,
    applied about a uniformly random axis.

    Args:
        shape:  batch dimensions (e.g., (L,) for L residues)
        sigma:  noise level (scalar or broadcastable tensor);
                small sigma -> small perturbations, large sigma -> near-uniform
        device: torch device

    Returns:
        rotations: [*shape, 3, 3] sampled rotation matrices
    """
    # Step 1: Sample rotation angle |omega| ~ folded normal(0, sigma^2)
    omega = torch.abs(torch.randn(*shape, device=device) * sigma)

    # Step 2: Sample rotation axis uniformly on the unit sphere S^2
    axis = torch.randn(*shape, 3, device=device)
    axis = axis / (axis.norm(dim=-1, keepdim=True) + 1e-8)

    # Step 3: Convert to rotation matrix via Rodrigues formula
    rotations = axis_angle_to_rotation_matrix(axis * omega.unsqueeze(-1))

    return rotations

def diffuse_translations(translations, t, noise_schedule):
    """
    Add Gaussian noise to C-alpha positions.

    Args:
        translations: [..., 3] clean C-alpha positions
        t:            [...] diffusion timesteps in [0, 1]
        noise_schedule: object providing alpha_bar(t)

    Returns:
        noisy_translations: [..., 3]
        noise:              [..., 3] the sampled Gaussian noise
    """
    alpha_bar = noise_schedule.alpha_bar(t)  # [..., 1]
    noise = torch.randn_like(translations)
    noisy = torch.sqrt(alpha_bar) * translations + torch.sqrt(1 - alpha_bar) * noise
    return noisy, noise

def diffuse_rotations(rotations, t, noise_schedule):
    """
    Add IGSO(3) noise to residue orientations.

    Unlike translations (where noise is added), rotational noise is
    composed (multiplied) because SO(3) is a multiplicative group.

    Args:
        rotations:      [..., 3, 3] clean rotation matrices
        t:              [...] diffusion timesteps in [0, 1]
        noise_schedule: object providing sigma_rot(t)

    Returns:
        noisy_rotations: [..., 3, 3]
        noise_rotations: [..., 3, 3] the IGSO(3) samples that were applied
    """
    sigma = noise_schedule.sigma_rot(t)

    # Sample noise rotation from IGSO(3)
    noise_rot = sample_igso3(rotations.shape[:-2], sigma, rotations.device)

    # Compose: R_t = R_noise @ R_0
    noisy_rotations = torch.einsum('...ij,...jk->...ik', noise_rot, rotations)

    return noisy_rotations, noise_rot

class SE3DiffusionSchedule:
    """
    Noise schedule for SE(3) diffusion on protein frames.

    The translation schedule uses the cosine form from Nichol & Dhariwal (2021),
    which avoids the abrupt noise increase at the end of a linear schedule.
    The rotation schedule is linear in the IGSO(3) concentration parameter.
    """

    def __init__(self, T=1000, trans_sigma_max=10.0, rot_sigma_max=1.5):
        """
        Args:
            T:              number of diffusion timesteps
            trans_sigma_max: maximum translation noise (Angstroms)
            rot_sigma_max:  maximum rotation noise (radians, roughly)
        """
        self.T = T
        self.trans_sigma_max = trans_sigma_max
        self.rot_sigma_max = rot_sigma_max

    def alpha_bar(self, t):
        """Cosine schedule for translations: alpha_bar(t) in [0, 1]."""
        s = 0.008  # small offset to prevent alpha_bar(0) = 1 exactly
        f_t = torch.cos((t + s) / (1 + s) * np.pi / 2) ** 2
        f_0 = np.cos(s / (1 + s) * np.pi / 2) ** 2
        return f_t / f_0

    def sigma_rot(self, t):
        """Linear schedule for rotations: sigma(t) = t * sigma_max."""
        return t * self.rot_sigma_max

    def sample_timestep(self, batch_size, device='cpu'):
        """Sample uniform random timesteps in [0, 1]."""
        return torch.rand(batch_size, device=device)

def diffuse_frames(frames, t, schedule):
    """
    Apply the SE(3) forward diffusion process to protein frames.

    Each residue frame T_i = (R_i, t_i) is independently noised:
    - Positions receive Gaussian noise in R^3
    - Orientations receive IGSO(3) noise on SO(3)

    Args:
        frames:   RigidTransform with L residue frames
        t:        [L] or scalar timestep
        schedule: SE3DiffusionSchedule

    Returns:
        noisy_frames: RigidTransform (the noised structure)
        noise:        dict with 'trans' and 'rot' noise components
    """
    noisy_trans, trans_noise = diffuse_translations(
        frames.trans, t.unsqueeze(-1), schedule
    )
    noisy_rots, rot_noise = diffuse_rotations(
        frames.rots, t, schedule
    )

    noisy_frames = RigidTransform(noisy_rots, noisy_trans)
    noise = {'trans': trans_noise, 'rot': rot_noise}
    return noisy_frames, noise

def compute_invariant_edge_features(frames, edge_index, max_dist=20.0):
    """
    Compute SE(3)-invariant features for pairs of residues.

    These features describe the geometric relationship between two
    residues without reference to any global coordinate system.

    Args:
        frames:    RigidTransform [N] (one frame per residue)
        edge_index: [2, E] pairs of interacting residues
        max_dist:  maximum distance for radial basis encoding (Angstroms)

    Returns:
        edge_features: [E, edge_dim] invariant features
    """
    src, dst = edge_index

    # --- Distance (invariant under all of SE(3)) ---
    pos_src = frames.trans[src]
    pos_dst = frames.trans[dst]
    distances = (pos_src - pos_dst).norm(dim=-1, keepdim=True)

    # Encode distance with radial basis functions
    rbf_dist = rbf_encode(distances, n_bins=16, max_val=max_dist)

    # --- Relative rotation angle (invariant under global rotation) ---
    # R_rel = R_dst^T @ R_src captures how one frame is rotated
    # relative to the other, independent of global orientation.
    R_rel = torch.einsum(
        'eij,ejk->eik',
        frames.rots[dst].transpose(-1, -2),
        frames.rots[src]
    )
    trace = R_rel[:, 0, 0] + R_rel[:, 1, 1] + R_rel[:, 2, 2]
    rot_angle = torch.acos(torch.clamp((trace - 1) / 2, -1, 1))

    # Combine into feature vector
    edge_features = torch.cat([
        rbf_dist,
        rot_angle.unsqueeze(-1),
        torch.sin(rot_angle).unsqueeze(-1),
        torch.cos(rot_angle).unsqueeze(-1)
    ], dim=-1)

    return edge_features


def rbf_encode(distances, n_bins=16, max_val=20.0):
    """
    Radial basis function encoding of distances.

    Transforms a scalar distance into a vector of Gaussian bumps centered
    at evenly spaced values. This gives the network a richer representation
    of distance than a single scalar.

    Args:
        distances: [..., 1] pairwise distances in Angstroms
        n_bins:    number of Gaussian basis functions
        max_val:   maximum distance to encode

    Returns:
        rbf: [..., n_bins] encoded distances
    """
    centers = torch.linspace(0, max_val, n_bins, device=distances.device)
    gamma = 1.0 / (max_val / n_bins)
    return torch.exp(-gamma * (distances - centers) ** 2)

class SE3EquivariantConv(nn.Module):
    """
    SE(3)-equivariant graph convolution layer for protein residues.

    Messages between residues are computed from invariant features
    (distances, relative orientations). Position updates are predicted
    in the local coordinate frame of each source residue, then
    transformed to global coordinates. This ensures equivariance:
    rotating the input rotates the output by the same amount.
    """

    def __init__(self, node_dim, edge_dim, hidden_dim=128):
        super().__init__()
        self.node_dim = node_dim
        self.edge_dim = edge_dim

        # Process invariant edge features
        self.edge_mlp = nn.Sequential(
            nn.Linear(edge_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim)
        )

        # Compute scalar messages (invariant)
        self.message_mlp = nn.Sequential(
            nn.Linear(node_dim * 2 + hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, node_dim)
        )

        # Predict position displacements in local frames (equivariant)
        self.pos_mlp = nn.Sequential(
            nn.Linear(node_dim * 2 + hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 3)
        )

    def forward(self, node_features, frames, edge_index, edge_features):
        """
        Args:
            node_features: [N, node_dim] per-residue features
            frames:        RigidTransform [N] per-residue frames
            edge_index:    [2, E] edges (pairs of interacting residues)
            edge_features: [E, edge_dim] invariant edge features

        Returns:
            updated_features:  [N, node_dim] updated per-residue features
            position_updates:  [N, 3] equivariant position updates
        """
        src, dst = edge_index
        N = node_features.shape[0]

        # Embed edge features
        edge_emb = self.edge_mlp(edge_features)

        # Build message inputs
        src_feat = node_features[src]
        dst_feat = node_features[dst]
        msg_input = torch.cat([src_feat, dst_feat, edge_emb], dim=-1)

        # Compute scalar messages and aggregate
        messages = self.message_mlp(msg_input)
        agg = torch.zeros(N, self.node_dim, device=messages.device)
        agg.scatter_add_(0, dst.unsqueeze(-1).expand(-1, self.node_dim), messages)

        # Predict position updates in SOURCE residue's local frame
        pos_local = self.pos_mlp(msg_input)  # [E, 3]

        # Transform to global coordinates using source frame's rotation
        src_frames = RigidTransform(frames.rots[src], frames.trans[src])
        pos_global = src_frames.apply(pos_local) - src_frames.trans

        # Aggregate position updates at destination residues
        pos_agg = torch.zeros(N, 3, device=pos_global.device)
        pos_agg.scatter_add_(0, dst.unsqueeze(-1).expand(-1, 3), pos_global)

        return node_features + agg, pos_agg

class RFDiffusionBlock(nn.Module):
    """
    One block of the RFDiffusion denoising network.

    Each block updates three representations:
    1. Per-residue (node) features via Invariant Point Attention
    2. Pairwise features via triangular updates
    3. Residue frames via predicted rotational and translational updates
    """

    def __init__(self, node_dim=256, pair_dim=128, n_heads=8):
        super().__init__()

        # Invariant Point Attention (from AlphaFold2's structure module)
        self.node_attn = InvariantPointAttention(node_dim, pair_dim, n_heads)

        # Triangular update on pair features (geometric consistency)
        self.pair_update = TriangularUpdate(pair_dim)

        # Frame update layer (predicts rotation + translation corrections)
        self.frame_update = FrameUpdateLayer(node_dim)

        # Timestep conditioning
        self.time_embed = nn.Sequential(
            nn.Linear(256, node_dim),
            nn.SiLU(),
            nn.Linear(node_dim, node_dim)
        )

    def forward(self, node_features, pair_features, frames, t_embed):
        """
        Args:
            node_features: [L, node_dim] per-residue features
            pair_features: [L, L, pair_dim] pairwise relationship features
            frames:        RigidTransform [L] current residue frames
            t_embed:       [256] sinusoidal timestep embedding

        Returns:
            Updated node_features, pair_features, and frames
        """
        # Inject timestep information into node features
        time_bias = self.time_embed(t_embed)
        node_features = node_features + time_bias.unsqueeze(0)

        # Update node features using Invariant Point Attention
        node_features = self.node_attn(node_features, pair_features, frames)

        # Update pair features with triangular consistency
        pair_features = self.pair_update(pair_features)

        # Predict and apply frame updates
        frame_updates = self.frame_update(node_features)
        frames = update_frames(frames, frame_updates)

        return node_features, pair_features, frames

class FrameUpdateLayer(nn.Module):
    """
    Predict per-residue frame corrections (rotation + translation).

    The rotation update is predicted as an axis-angle vector (3D),
    converted to a rotation matrix, and composed with the current frame.
    The translation update is a 3D displacement added to the current position.
    """

    def __init__(self, node_dim):
        super().__init__()
        self.norm = nn.LayerNorm(node_dim)
        self.to_update = nn.Linear(node_dim, 6)  # 3 rotation + 3 translation
        self.rot_scale = 0.1    # keep rotation updates small
        self.trans_scale = 1.0  # translation updates in Angstroms

    def forward(self, node_features):
        """
        Args:
            node_features: [L, node_dim] per-residue features

        Returns:
            dict with 'rot' ([L, 3] axis-angle) and 'trans' ([L, 3]) updates
        """
        x = self.norm(node_features)
        updates = self.to_update(x)

        rot_update = updates[:, :3] * self.rot_scale
        trans_update = updates[:, 3:] * self.trans_scale

        return {'rot': rot_update, 'trans': trans_update}


def update_frames(frames, updates):
    """
    Apply predicted updates to residue frames by composition.

    The update is treated as a small rigid-body transformation that is
    composed (not added) with the current frame. This respects the
    group structure of SE(3).
    """
    rot_update = axis_angle_to_rotation_matrix(updates['rot'])
    trans_update = updates['trans']

    update_transform = RigidTransform(rot_update, trans_update)

    # Compose: T_new = T_current o T_update
    return frames.compose(update_transform)

class MotifConditioning:
    """
    Condition the diffusion process on a fixed structural motif.

    During each denoising step, motif residue frames are clamped to their
    ground-truth values. The model learns to generate scaffold residues
    that are geometrically compatible with the fixed motif.
    """

    def __init__(self, motif_positions, motif_frames):
        """
        Args:
            motif_positions: list of int, residue indices to hold fixed
            motif_frames:    RigidTransform for the motif residues
        """
        self.motif_pos = set(motif_positions)
        self.motif_frames = motif_frames

    def apply(self, frames, t):
        """
        Replace motif positions with ground-truth frames.

        Called at each denoising step to enforce the motif constraint.
        """
        for i, pos in enumerate(self.motif_pos):
            frames.rots[pos] = self.motif_frames.rots[i]
            frames.trans[pos] = self.motif_frames.trans[i]
        return frames

    def get_conditioning_mask(self, L):
        """
        Return a binary mask: 1 for motif (fixed), 0 for scaffold (generated).
        """
        mask = torch.zeros(L)
        for pos in self.motif_pos:
            mask[pos] = 1
        return mask

class SelfConditionedRFDiffusion(nn.Module):
    """
    RFDiffusion with self-conditioning.

    At each denoising step, the model can optionally receive its own
    previous prediction as input. This helps maintain consistency
    across steps and improves sample quality.
    """

    def __init__(self, base_model):
        super().__init__()
        self.base_model = base_model
        # Project 7D frame representation (4 quat + 3 trans) to features
        self.self_cond_proj = nn.Linear(7, 128)

    def forward(self, noisy_frames, t, prev_pred=None):
        if prev_pred is not None:
            self_cond = self.self_cond_proj(prev_pred.to_tensor_7())
        else:
            self_cond = torch.zeros(noisy_frames.trans.shape[0], 128)

        return self.base_model(noisy_frames, t, self_cond)

def sample_with_guidance(model, initial_frames, condition,
                         guidance_scale=2.0, steps=100):
    """
    Generate a protein backbone with classifier-free guidance.

    Higher guidance_scale produces structures that more strongly
    satisfy the conditioning constraints, but with less diversity.

    Args:
        model:          trained RFDiffusion model
        initial_frames: starting frames (typically random noise)
        condition:      conditioning information (motif, target, etc.)
        guidance_scale: s > 1 for stronger conditioning
        steps:          number of denoising steps
    """
    frames = initial_frames
    schedule = SE3DiffusionSchedule()

    for step in range(steps, 0, -1):
        t = torch.tensor([step / steps])

        # Conditional prediction (with design constraints)
        pred_cond = model(frames, t, condition)

        # Unconditional prediction (no constraints)
        pred_uncond = model(frames, t, None)

        # Guided prediction: extrapolate in the conditioning direction
        pred = pred_uncond + guidance_scale * (pred_cond - pred_uncond)

        # Take one denoising step
        frames = denoise_step(frames, pred, t, schedule)

    return frames

def rfdiffusion_loss(pred_frames, true_frames, t, loss_weights=None):
    """
    Compute the RFDiffusion training loss.

    The loss has two components:
    - Translation loss: MSE in Euclidean space (Angstroms^2)
    - Rotation loss: squared geodesic distance on SO(3) (radians^2)

    Args:
        pred_frames: RigidTransform, model predictions
        true_frames: RigidTransform, ground truth from PDB
        t:           timestep (for potential weighting)
        loss_weights: dict with 'trans' and 'rot' weights

    Returns:
        total_loss: scalar
        losses:     dict with individual loss components
    """
    losses = {}

    # --- Translation loss: L2 distance between C-alpha positions ---
    trans_loss = (
        (pred_frames.trans - true_frames.trans) ** 2
    ).sum(dim=-1).mean()
    losses['trans'] = trans_loss

    # --- Rotation loss: geodesic distance on SO(3) ---
    # R_diff = R_pred @ R_true^T gives the relative rotation
    R_diff = torch.einsum(
        '...ij,...kj->...ik',
        pred_frames.rots,
        true_frames.rots
    )
    trace = R_diff[:, 0, 0] + R_diff[:, 1, 1] + R_diff[:, 2, 2]
    rot_angle = torch.acos(
        torch.clamp((trace - 1) / 2, -1 + 1e-7, 1 - 1e-7)
    )
    rot_loss = (rot_angle ** 2).mean()
    losses['rot'] = rot_loss

    # --- Combine ---
    if loss_weights is None:
        loss_weights = {'trans': 1.0, 'rot': 1.0}

    total_loss = sum(
        loss_weights.get(k, 1.0) * v for k, v in losses.items()
    )
    return total_loss, losses

def training_step(model, batch, optimizer, schedule):
    """
    One step of RFDiffusion training.

    1. Sample a random noise level t
    2. Add noise to the ground-truth protein structure
    3. Feed the noisy structure to the model
    4. Compute loss between prediction and ground truth
    5. Update model weights
    """
    true_frames = batch['frames']  # Ground truth from PDB
    L = true_frames.trans.shape[0]  # Number of residues

    # Sample a random timestep uniformly in [0, 1]
    t = schedule.sample_timestep(1).expand(L)

    # Apply the forward diffusion process
    noisy_frames, noise = diffuse_frames(true_frames, t, schedule)

    # Model predicts the clean structure from the noisy input
    pred_frames = model(noisy_frames, t)

    # Compute and minimize the loss
    loss, loss_dict = rfdiffusion_loss(pred_frames, true_frames, t)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    return loss_dict

@torch.no_grad()
def sample_rfdiffusion(model, L, schedule, n_steps=100, condition=None):
    """
    Generate a novel protein backbone using RFDiffusion.

    Starting from random noise, iteratively denoise to produce a
    coherent protein backbone with L residues.

    Args:
        model:     trained RFDiffusion model
        L:         desired number of residues
        schedule:  SE3DiffusionSchedule
        n_steps:   number of denoising steps (more steps = higher quality)
        condition: optional conditioning (MotifConditioning, etc.)

    Returns:
        frames: RigidTransform with the generated backbone
    """
    # Initialize: identity frames + maximum noise = random structure
    frames = RigidTransform.identity((L,), device='cuda')
    t_max = torch.ones(L, device='cuda')
    frames, _ = diffuse_frames(frames, t_max, schedule)

    # Self-conditioning state
    prev_pred = None

    # Denoise from t=1 (pure noise) to t=0 (clean structure)
    timesteps = torch.linspace(1, 0, n_steps + 1)

    for i in range(n_steps):
        t_now = timesteps[i]
        t_next = timesteps[i + 1]
        t_tensor = torch.full((L,), t_now, device='cuda')

        # Model predicts the clean structure
        pred_frames = model(frames, t_tensor, condition, prev_pred)

        # Self-conditioning: with 50% probability, feed prediction
        # back to the model at the next step
        if torch.rand(1) < 0.5:
            prev_pred = pred_frames

        # Move one step closer to the clean structure
        if t_next > 0:
            frames = denoise_and_renoise(
                frames, pred_frames, t_now, t_next, schedule
            )
        else:
            # Final step: use the prediction directly
            frames = pred_frames

        # Enforce conditioning constraints
        if condition is not None:
            frames = condition.apply(frames, t_next)

    return frames

def denoise_and_renoise(noisy_frames, pred_clean, t_now, t_next, schedule):
    """
    Take one denoising step from t_now to t_next on SE(3).

    For translations: use the standard DDPM update rule in R^3.
    For rotations: interpolate on SO(3) toward the predicted clean rotation.

    Args:
        noisy_frames: current noisy frames at timestep t_now
        pred_clean:   model's prediction of the clean frames
        t_now:        current timestep
        t_next:       target timestep (t_next < t_now)
        schedule:     SE3DiffusionSchedule
    """
    # --- Translations: DDPM update in R^3 ---
    alpha_now = schedule.alpha_bar(torch.tensor(t_now))
    alpha_next = schedule.alpha_bar(torch.tensor(t_next))

    # Estimate the noise that was added
    noise_est = (
        (noisy_frames.trans - torch.sqrt(alpha_now) * pred_clean.trans)
        / torch.sqrt(1 - alpha_now)
    )

    # Posterior mean at t_next
    trans_mean = (
        torch.sqrt(alpha_next) * pred_clean.trans
        + torch.sqrt(1 - alpha_next) * noise_est
    )

    # Add stochastic noise (except at the final step)
    if t_next > 0:
        noise = torch.randn_like(trans_mean)
        beta = 1 - alpha_next / alpha_now
        trans_next = trans_mean + torch.sqrt(beta) * noise
    else:
        trans_next = trans_mean

    # --- Rotations: interpolate on SO(3) via SLERP ---
    sigma_now = schedule.sigma_rot(torch.tensor(t_now))
    sigma_next = schedule.sigma_rot(torch.tensor(t_next))

    # Interpolation factor: how far to move toward the clean rotation
    interp_factor = sigma_next / (sigma_now + 1e-8)
    rots_next = slerp_rotation(
        noisy_frames.rots, pred_clean.rots, 1 - interp_factor
    )

    return RigidTransform(rots_next, trans_next)