Upload folder using huggingface_hub

6aab6b3 verified about 2 months ago

16.8 kB

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
                ]
            )

	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
	]
	)