Function minimum_lddt_symmetry_coords

src/boltz/data/feature/symmetry.py:239–377

(
    coords: torch.Tensor,
    feats: dict,
    index_batch: int,
    **args_rmsd,
)

def minimum_lddt_symmetry_coords(
    coords: torch.Tensor,
    feats: dict,
    index_batch: int,
    **args_rmsd,
):
    all_coords = feats["all_coords"][index_batch].unsqueeze(0).to(coords)
    all_resolved_mask = (
        feats["all_resolved_mask"][index_batch].to(coords).to(torch.bool)
    )
    crop_to_all_atom_map = (
        feats["crop_to_all_atom_map"][index_batch].to(coords).to(torch.long)
    )
    chain_symmetries = feats["chain_symmetries"][index_batch]
    amino_acids_symmetries = feats["amino_acids_symmetries"][index_batch]
    ligand_symmetries = feats["ligand_symmetries"][index_batch]

    dmat_predicted = torch.cdist(
        coords[:, : len(crop_to_all_atom_map)], coords[:, : len(crop_to_all_atom_map)]
    )

    # Check best symmetry on chain swap
    best_true_coords = None
    best_lddt = 0
    for c in chain_symmetries:
        true_all_coords = all_coords.clone()
        true_all_resolved_mask = all_resolved_mask.clone()
        for start1, end1, start2, end2, chainidx1, chainidx2 in c:
            true_all_coords[:, start1:end1] = all_coords[:, start2:end2]
            true_all_resolved_mask[start1:end1] = all_resolved_mask[start2:end2]
        true_coords = true_all_coords[:, crop_to_all_atom_map]
        true_resolved_mask = true_all_resolved_mask[crop_to_all_atom_map]
        dmat_true = torch.cdist(true_coords, true_coords)
        pair_mask = (
            true_resolved_mask[:, None]
            * true_resolved_mask[None, :]
            * (1 - torch.eye(len(true_resolved_mask))).to(true_resolved_mask)
        )

        lddt = lddt_dist(
            dmat_predicted, dmat_true, pair_mask, cutoff=15.0, per_atom=False
        )[0]
        lddt = lddt.item()

        if lddt > best_lddt:
            best_lddt = lddt
            best_true_coords = true_coords
            best_true_resolved_mask = true_resolved_mask

    # atom symmetries (nucleic acid and protein residues), resolved greedily without recomputing alignment
    true_coords = best_true_coords.clone()
    true_resolved_mask = best_true_resolved_mask.clone()
    for symmetric_amino_or_lig in amino_acids_symmetries + ligand_symmetries:
        for c in symmetric_amino_or_lig:
            # starting from greedy best, try to swap the atoms
            new_true_coords = true_coords.clone()
            new_true_resolved_mask = true_resolved_mask.clone()
            indices = []