Source code for boltz_data.mol._to._to_bzcif

from itertools import groupby

import gemmi
import numpy as np

from boltz_data.mol._mol import BZMol




def bzcif_from_bzmol(bzmol: BZMol, /) -> gemmi.cif.Block:
    """Convert a BZMol to a CIF block representation."""
    bzmol_dict = bzmol.model_dump()

    keys = sorted(bzmol_dict)

    block = gemmi.cif.Block("test")

    for prefix, subkeys in groupby(keys, lambda key: key.split("_")[0]):
        columns = {}
        for subkey in subkeys:
            colname = subkey.split("_", 1)[1]
            value = bzmol_dict[subkey]
            if value is None:
                continue
            if len(value.shape) == 1:
                if value.dtype == np.bool_:
                    columns[colname] = ["Y" if item else "N" for item in value]
                else:
                    columns[colname] = value.tolist()
            else:
                for axis in range(value.shape[1]):
                    columns[f"{colname}[{axis}]"] = value[..., axis].tolist()

        loop = block.init_loop(f"_{prefix}.", list(columns))
        for row in zip(*columns.values(), strict=False):
            loop.add_row([str(v) for v in row])

    return block