Extract a protein structure

"""Extract fold structure data from the Atlas Lance folds dataset.

Reads one protein's row from a Lance folds dataset and decodes it into the
columns the structure API returns: ptm, mean_plddt, per_residue_plddt, pdb.

The dataset URI is a local path or an s3:// URI to a `.lance` folds dataset.
"""

import inspect
import io

import brotli
import lance
import msgpack
import numpy as np
from esm.utils.structure.protein_chain import ProteinChain

# Logical column names -> Lance dataset columns. "pdb" is decoded from the
# "structure_blob" column; the rest are pre-extracted scalar/binary columns.
LANCE_COLUMNS = ["ptm", "mean_plddt", "per_residue_plddt", "structure_blob"]


def decode_structure_row(row: dict) -> dict:
    """Decode a Lance folds row into the columns returned by the structure API.

    Returns a dict with keys: ptm, mean_plddt, per_residue_plddt, pdb.
    """
    # ptm / mean_plddt: pre-extracted scalar columns, returned as-is.
    ptm = row.get("ptm")
    mean_plddt = row.get("mean_plddt")

    # per_residue_plddt: npz float16 array blob -> list[float].
    per_residue_plddt = None
    if row.get("per_residue_plddt") is not None:
        arr = np.load(io.BytesIO(row["per_residue_plddt"]))["arr"]
        per_residue_plddt = arr.astype(float).tolist()

    # pdb: brotli decompress -> msgpack unpack -> ProteinChain -> PDB string.
    pdb = None
    if row.get("structure_blob"):
        state_dict = msgpack.loads(brotli.decompress(row["structure_blob"]), strict_map_key=False)
        valid_keys = set(inspect.signature(ProteinChain.__init__).parameters) - {"self", "mmcif"}
        filtered = {k: v for k, v in state_dict.items() if k in valid_keys}
        # The blob stores chain_id/entity_id as per-residue int arrays, but the
        # pip esm package expects chain_id as a scalar string. Convert accordingly.
        if hasattr(filtered.get("chain_id"), "dtype") and filtered["chain_id"].dtype.kind == "i":
            unique = np.unique(filtered["chain_id"])
            filtered["chain_id"] = chr(ord("A") + int(unique[0]))
        if hasattr(filtered.get("entity_id"), "dtype") and filtered["entity_id"].dtype.kind == "i":
            filtered["entity_id"] = None
        pdb = ProteinChain.from_state_dict(filtered).to_pdb_string()

    return {
        "ptm": ptm,
        "mean_plddt": mean_plddt,
        "per_residue_plddt": per_residue_plddt,
        "pdb": pdb,
    }


def extract_structure(dataset_uri: str, protein_hash: str, region: str = "us-west-2") -> dict | None:
    """Read one protein's row from a Lance folds dataset and decode all API columns.

    Args:
        dataset_uri: Local path or s3:// URI to the Lance folds dataset.
        protein_hash: 32-char hex MD5 hash of the protein sequence.
        region: AWS region for S3 datasets.

    Returns:
        Dict with ptm/mean_plddt/per_residue_plddt/pdb, or None if not found.
    """
    storage_options = {"region": region}
    # S3 Express one-zone buckets (those with a "--...--x-s3" suffix) need vhost-style requests.
    if dataset_uri.startswith("s3://") and "--" in dataset_uri.split("/")[2]:
        storage_options["virtual_hosted_style_request"] = "true"

    ds = lance.dataset(dataset_uri, storage_options=storage_options)
    table = ds.to_table(
        filter=f"protein_hash = '{protein_hash}'",
        columns=LANCE_COLUMNS,
    )
    if table.num_rows == 0:
        return None
    return decode_structure_row(table.to_pylist()[0])


if __name__ == "__main__":
    import sys

    uri, protein_hash = sys.argv[1], sys.argv[2]
    result = extract_structure(uri, protein_hash)
    assert result is not None
    print("ptm:", result["ptm"])
    print("mean_plddt:", result["mean_plddt"])
    print("per_residue_plddt (first 5):", result["per_residue_plddt"][:5])
    print("pdb (first line):", result["pdb"].splitlines()[0])