Blame for model/boltzgen/data/write/mmcif.py

import re
import string
from collections import defaultdict

import gemmi

from boltzgen.data import const
from boltzgen.data.data import Structure, elem_from_name


def _create_gemmi_model(model_num: int) -> gemmi.Model:
    """Create a gemmi.Model compatible with both gemmi 0.6.x and 0.7.x.

    In gemmi 0.7.0, the Model constructor changed from taking a string name
    to taking an integer number (Model.name -> Model.num).

    Args:
        model_num: The model number (typically 1).

    Returns:
        A gemmi.Model instance.
    """
    try:
        # gemmi 0.7.x: Model(int) constructor with Model.num property
        return gemmi.Model(model_num)
    except TypeError:
        # gemmi 0.6.x: Model(str) constructor with Model.name property
        return gemmi.Model(str(model_num))


def to_mmcif(
    structure: Structure,
    write_fake_atoms: bool = False,
    poly_seq_scheme: bool = True,
    plddt_cols: bool = True,
    design_coloring: bool = False,
    color_features=None,
) -> str:
    """
    Convert a `Structure` object into an mmCIF-formatted string using Gemmi.

    Args:
        structure (Structure): Boltz Structure object.
        write_fake_atoms (bool): If True, includes fake atoms
        poly_seq_scheme (bool): If True, writes polymer sequence information into the
            `_pdbx_poly_seq_scheme` table.
        plddt_cols (bool): If True, writes pLDDT scores into the `_ma_qa_metric` and
            `_ma_qa_metric_local` tables.
        design_coloring (bool): If True, adds a custom `design_coloring scheme.'
        color_features (np.ndarray): Array of coloring features for design coloring.

    Returns:
        str: A string representation of the structure in mmCIF format.
    """

    gemmi_struct = gemmi.Structure()
    gemmi_struct.name = "model"
    model = _create_gemmi_model(1)

    chain_to_entity_id = {}
    sequence_to_entity_id = {}
    entity_counter = 1
    chain_names = [re.sub(r"\d+", "", c["name"]) for c in structure.chains]
    chain_id_pool = list(reversed(string.ascii_uppercase)) + list(
        reversed(string.digits)
    )
    used_names = []
    old_to_new_chainid = {}
    for chain in structure.chains:
        old_chainid = chain["name"].item()

        new_chainid = re.sub(r"\d+", "", old_chainid)
        if new_chainid in used_names:
            # Find next unused chain ID from the pool
            for candidate in chain_id_pool:
                if candidate not in chain_names and candidate not in used_names:
                    new_chainid = candidate
                    break
        old_to_new_chainid[old_chainid] = new_chainid
        used_names.append(new_chainid)

        residues = structure.residues[
            chain["res_idx"] : chain["res_idx"] + chain["res_num"]
        ]

        sequence = "".join([res["name"].item() for res in residues])
        chain_type = chain["mol_type"].item()
        if sequence not in sequence_to_entity_id:
            sequence_to_entity_id[sequence] = str(entity_counter)

            entity = gemmi.Entity(str(entity_counter))

            if chain_type == const.chain_type_ids["NONPOLYMER"]:
                entity.entity_type = gemmi.EntityType.NonPolymer
            else:
                entity.entity_type = gemmi.EntityType.Polymer
                if chain_type == const.chain_type_ids["DNA"]:
                    entity.polymer_type = gemmi.PolymerType.Dna
                elif chain_type == const.chain_type_ids["RNA"]:
                    entity.polymer_type = gemmi.PolymerType.Rna
                elif chain_type == const.chain_type_ids["PROTEIN"]:
                    entity.polymer_type = gemmi.PolymerType.PeptideL

                # Seqres that also contains residues with missing coordinates
                entity.full_sequence = [res["name"].item() for res in residues]

            gemmi_struct.entities.append(entity)
            entity_counter += 1

        chain_to_entity_id[new_chainid] = sequence_to_entity_id[sequence]

    label_seq_dict = defaultdict(list)
    for chain in structure.chains:
        old_chainid = chain["name"].item()
        chain_id = old_to_new_chainid[old_chainid]

        gemmi_chain = gemmi.Chain(chain_id)

        residues = structure.residues[
            chain["res_idx"] : chain["res_idx"] + chain["res_num"]
        ]

        polymer = False
        for entity in gemmi_struct.entities:
            if (
                entity.name == chain_to_entity_id[chain_id]
                and entity.entity_type == gemmi.EntityType.Polymer
            ):
                polymer = True
                break

        label_seq_ids = []
        for res in residues:
            # add label_seq_ids if polymer and new entity.
            label_seq_ids.append(res["res_idx"].item() + 1)
            # Missing residues are in the seqres but not in the residue table
            if not res["is_present"]:
                continue

            res_name = res["name"].item()
            seqid = gemmi.SeqId(res["res_idx"].item() + 1, " ")
            gemmi_res = gemmi.Residue()
            gemmi_res.name = res_name
            gemmi_res.seqid = seqid
            gemmi_res.label_seq = res["res_idx"].item() + 1
            gemmi_res.subchain = chain_id

            if chain_id in chain_to_entity_id:
                gemmi_res.entity_id = chain_to_entity_id[chain_id]

            atoms = structure.atoms[res["atom_idx"] : res["atom_idx"] + res["atom_num"]]
            coords = structure.coords["coords"][
                res["atom_idx"] : res["atom_idx"] + res["atom_num"]
            ]

            for atom, coord in zip(atoms, coords):
                # Skip missing atoms
                if not atom["is_present"]:
                    continue

                atom_name = atom["name"].item()
                element = elem_from_name(atom_name, res_name)

                # Skip fake atoms
                if not write_fake_atoms and (
                    const.fake_element.upper() in atom_name
                    or const.mask_element.upper() in atom_name
                ):
                    assert (
                        element == const.fake_element or element == const.mask_element
                    ), "Atom name not consistent with element for possible fake atom."
                    continue
                pos = gemmi.Position(coord[0], coord[1], coord[2])
                gemmi_atom = gemmi.Atom()
                gemmi_atom.name = atom_name
                gemmi_atom.pos = pos
                gemmi_atom.occ = 1.0
                gemmi_atom.b_iso = atom["bfactor"].item()
                gemmi_atom.element = gemmi.Element(element)
                gemmi_res.add_atom(gemmi_atom)
            gemmi_chain.add_residue(gemmi_res)
        model.add_chain(gemmi_chain)

        if chain_to_entity_id[chain_id] not in label_seq_dict.keys() and polymer:
            label_seq_dict[chain_to_entity_id[chain_id]] = label_seq_ids

    # Add covalent,Disulfide bonds to struct_conn
    cov_count = 0
    disulf_count = 0
    for bond in structure.bonds:
        if bond["type"] == const.bond_type_ids["COVALENT"]:
            res1 = structure.residues[bond["res_1"]]
            res2 = structure.residues[bond["res_2"]]

            chain1 = structure.chains[bond["chain_1"]]
            chain2 = structure.chains[bond["chain_2"]]
            atom1 = structure.atoms[bond["atom_1"]]
            atom2 = structure.atoms[bond["atom_2"]]

            p1_chain_id = old_to_new_chainid[chain1["name"].item()]
            p2_chain_id = old_to_new_chainid[chain2["name"].item()]
            p1_seq_id = res1["res_idx"].item() + 1
            p2_seq_id = res2["res_idx"].item() + 1

            p1_atom_name = atom1["name"].item()
            p2_atom_name = atom2["name"].item()

            con = gemmi.Connection()
            if p1_atom_name == "SG" and p2_atom_name == "SG":
                disulf_count += 1
                con.type = gemmi.ConnectionType.Disulf
                con.name = f"disulf{disulf_count}"
            else:
                cov_count += 1
                con.type = gemmi.ConnectionType.Covale
                con.name = f"covale{cov_count}"

            con.partner1 = gemmi.AtomAddress(
                p1_chain_id,
                gemmi.SeqId(str(p1_seq_id)),
                res1["name"].item(),
                p1_atom_name,
            )
            con.partner2 = gemmi.AtomAddress(
                p2_chain_id,
                gemmi.SeqId(str(p2_seq_id)),
                res2["name"].item(),
                p2_atom_name,
            )

            con.asu = gemmi.Asu.Same
            gemmi_struct.connections.append(con)

    gemmi_struct.add_model(model)
    doc = gemmi_struct.make_mmcif_document()
    block = doc.sole_block()

    # entity_poly_seq generated from make_mmcif_document by default start from index 1
    # we manually fix this index mismatch to use real res_idx
    entity_poly_seq_block_idx = doc[0].get_index("_entity_poly_seq.entity_id")
    eps_loop = doc[0][entity_poly_seq_block_idx].loop
    entity_poly_seq_label_id_idx = eps_loop.tags.index("_entity_poly_seq.num")
    entity_poly_seq_entity_id_idx = eps_loop.tags.index("_entity_poly_seq.entity_id")

    curr_entity_id, curr_resid = None, 0
    for i in range(eps_loop.length()):
        entity_id = eps_loop[i, entity_poly_seq_entity_id_idx]
        if entity_id != curr_entity_id:
            curr_entity_id = entity_id
            curr_resid = 0
        eps_loop[i, entity_poly_seq_label_id_idx] = str(label_seq_dict[entity_id][curr_resid])
        curr_resid += 1
    

    struct_asym_loop = block.init_loop("_struct_asym.", ["id", "entity_id"])
    for chain_id, entity_id in sorted(chain_to_entity_id.items()):
        struct_asym_loop.add_row([chain_id, entity_id])

    # Include poly_seq_scheme table
    if poly_seq_scheme:
        add_poly_seq_scheme_cols(structure, block, chain_to_entity_id, label_seq_dict)

    if plddt_cols:
        add_plddt_cols(structure, block)

    if design_coloring:
        add_design_cols(structure, block, color_features)

    # remove _chem_comp records because they are empty and then just cause problems with visualization softwares
    block_string = doc.as_string()
    pattern = r"(loop_\n_chem_comp[\s\S]*?)(?=loop_\n)"
    block_string = re.sub(pattern, "", block_string)

    return block_string


def add_boltzgen_metadata(structure, block, old_to_new_chainid):
    cols = [
        "asym_id",  # Chain ID
        "res_idx_in_this_mmcif",  # Reindexed res_idx found in this mmcif file
        "res_idx_physical",  # True res_idx
        "design_mask",  # 1 for designed residues 0 for non-designed residues
    ]
    custom_loop = block.init_loop("_boltzgen_metadata.", cols)

    for chain in structure.chains:
        chain_id = old_to_new_chainid[chain["name"].item()]
        residues = structure.residues[
            chain["res_idx"] : chain["res_idx"] + chain["res_num"]
        ]
        for res_idx, pysical_idx in enumerate(residues, 1):
            seq_id = res_idx + 1
            if not res["is_present"]:
                continue
            mon_id = res["name"].item()
            score = res.get("custom_score", 0.0)  # Example access to your metadata
            custom_loop.add_row(
                [
                    chain_id,
                    str(seq_id),
                    f"{score:.3f}",
                    mon_id,
                ]
            )


def add_design_cols(structure, block, colors):
    plddt_loop = block.init_loop(
        "_ma_qa_metric.",
        [
            "id",
            "name",
            "description",
            "type",
            "mode",
            "type_other_details",
            "software_group_id",
        ],
    )
    plddt_loop.add_row(["1", "pLDDT", "'Predicted lddt'", "pLDDT", "local", ".", "."])

    plddt_cols = [
        "ordinal_id",
        "model_id",
        "label_asym_id",
        "label_seq_id",
        "label_comp_id",
        "metric_id",
        "metric_value",
    ]
    plddt_loop = block.init_loop("_ma_qa_metric_local.", plddt_cols)
    global_res_idx = -1
    for chain in structure.chains:
        chain_name_str = re.sub(r"\d+", "", chain["name"].item())
        chain_id = chain_name_str

        residues = structure.residues[
            chain["res_idx"] : chain["res_idx"] + chain["res_num"]
        ]

        for _, res in enumerate(residues, 1):
            global_res_idx += 1
            if not res["is_present"]:
                continue
            mon_id = res["name"].item()
            design_label = colors[global_res_idx]  # [plddt_loop.length()]
            plddt_loop.add_row(
                [
                    str(plddt_loop.length() + 1),  # ordinal_id
                    "1",  # model_id
                    chain_id,  # label_asym_id
                    str(res["res_idx"].item() + 1),
                    mon_id,  # comp_id
                    "1",  # metric_id
                    f"{design_label * 100:.3f}",  # metric_value, [0, 100]
                ]
            )


def add_plddt_cols(structure, block):
    plddt_loop = block.init_loop(
        "_ma_qa_metric.",
        [
            "id",
            "name",
            "description",
            "type",
            "mode",
            "type_other_details",
            "software_group_id",
        ],
    )
    plddt_loop.add_row(["1", "pLDDT", "'Predicted lddt'", "pLDDT", "local", ".", "."])

    plddt_cols = [
        "ordinal_id",
        "model_id",
        "label_asym_id",
        "label_seq_id",
        "label_comp_id",
        "metric_id",
        "metric_value",
    ]
    plddt_loop = block.init_loop("_ma_qa_metric_local.", plddt_cols)
    for chain in structure.chains:
        if chain["mol_type"].item() == const.chain_type_ids["NONPOLYMER"]:
            continue

        chain_name_str = re.sub(r"\d+", "", chain["name"].item())
        chain_id = chain_name_str

        residues = structure.residues[
            chain["res_idx"] : chain["res_idx"] + chain["res_num"]
        ]

        for _, res in enumerate(residues, 1):
            if not res["is_present"]:
                continue
            mon_id = res["name"].item()
            ref_atom_idx = res["atom_idx"]
            plddt_score = structure.atoms[ref_atom_idx]["bfactor"].item()
            plddt_loop.add_row(
                [
                    str(plddt_loop.length() + 1),  # ordinal_id
                    "1",  # model_id
                    chain_id,  # label_asym_id
                    str(res["res_idx"].item() + 1),
                    mon_id,  # comp_id
                    "1",  # metric_id
                    f"{plddt_score * 100:.3f}",  # metric_value, [0, 100]
                ]
            )


def add_poly_seq_scheme_cols(structure, block, chain_to_entity_id, label_seq_dict):
    poly_seq_scheme_cols = [
        "asym_id",
        "entity_id",
        "seq_id",
        "mon_id",
        "pdb_seq_num",
        "auth_seq_num",
        "pdb_mon_id",
        "auth_mon_id",
        "pdb_strand_id",
        "pdb_ins_code",
        "hetero",
    ]
    poly_seq_loop = block.init_loop("_pdbx_poly_seq_scheme.", poly_seq_scheme_cols)

    for chain in structure.chains:
        if chain["mol_type"].item() == const.chain_type_ids["NONPOLYMER"]:
            continue

        chain_name_str = re.sub(r"\d+", "", chain["name"].item())
        chain_id = chain_name_str
        entity_id = chain_to_entity_id[chain_id]

        residues = structure.residues[
            chain["res_idx"] : chain["res_idx"] + chain["res_num"]
        ]

        # Use enumerate to get the sequential 1-based seq_id
        for seq_id, res in enumerate(residues, 1):
            mon_id = res["name"].item()
            auth_seq_num = str(res["res_idx"].item() + 1)

            # For missing residues, many fields should be '?'
            if res["is_present"]:
                pdb_seq_num = auth_seq_num
                pdb_mon_id = mon_id
                auth_mon_id = mon_id
            else:
                pdb_seq_num = auth_seq_num
                pdb_mon_id = "?"
                auth_mon_id = "?"
            poly_seq_loop.add_row(
                [
                    chain_id,  # asym_id
                    entity_id,  # entity_id
                    str(label_seq_dict[entity_id][seq_id - 1]),
                    mon_id,  # mon_id
                    pdb_seq_num,  # pdb_seq_num (use auth if present)
                    auth_seq_num,  # auth_seq_num (your res_idx)
                    pdb_mon_id,  # pdb_mon_id
                    auth_mon_id,  # auth_mon_id
                    chain_id,  # pdb_strand_id
                    ".",  # pdb_ins_code
                    "n",  # hetero
                ]
            )