> ## Documentation Index
> Fetch the complete documentation index at: https://docs.biohub.ai/llms.txt
> Use this file to discover all available pages before exploring further.

# Extract a protein structure

> Read one protein from the Atlas folds dataset and decode its structure and confidence values.

Read one protein from the Atlas folds dataset and decode its structure and confidence values.

<a href="data:text/x-python;base64,IiIiRXh0cmFjdCBmb2xkIHN0cnVjdHVyZSBkYXRhIGZyb20gdGhlIEF0bGFzIExhbmNlIGZvbGRzIGRhdGFzZXQuCgpSZWFkcyBvbmUgcHJvdGVpbidzIHJvdyBmcm9tIGEgTGFuY2UgZm9sZHMgZGF0YXNldCBhbmQgZGVjb2RlcyBpdCBpbnRvIHRoZQpjb2x1bW5zIHRoZSBzdHJ1Y3R1cmUgQVBJIHJldHVybnM6IHB0bSwgbWVhbl9wbGRkdCwgcGVyX3Jlc2lkdWVfcGxkZHQsIHBkYi4KClRoZSBkYXRhc2V0IFVSSSBpcyBhIGxvY2FsIHBhdGggb3IgYW4gczM6Ly8gVVJJIHRvIGEgYC5sYW5jZWAgZm9sZHMgZGF0YXNldC4KIiIiCgppbXBvcnQgaW5zcGVjdAppbXBvcnQgaW8KCmltcG9ydCBicm90bGkKaW1wb3J0IGxhbmNlCmltcG9ydCBtc2dwYWNrCmltcG9ydCBudW1weSBhcyBucApmcm9tIGVzbS51dGlscy5zdHJ1Y3R1cmUucHJvdGVpbl9jaGFpbiBpbXBvcnQgUHJvdGVpbkNoYWluCgojIExvZ2ljYWwgY29sdW1uIG5hbWVzIC0+IExhbmNlIGRhdGFzZXQgY29sdW1ucy4gInBkYiIgaXMgZGVjb2RlZCBmcm9tIHRoZQojICJzdHJ1Y3R1cmVfYmxvYiIgY29sdW1uOyB0aGUgcmVzdCBhcmUgcHJlLWV4dHJhY3RlZCBzY2FsYXIvYmluYXJ5IGNvbHVtbnMuCkxBTkNFX0NPTFVNTlMgPSBbInB0bSIsICJtZWFuX3BsZGR0IiwgInBlcl9yZXNpZHVlX3BsZGR0IiwgInN0cnVjdHVyZV9ibG9iIl0KCgpkZWYgZGVjb2RlX3N0cnVjdHVyZV9yb3cocm93OiBkaWN0KSAtPiBkaWN0OgogICAgIiIiRGVjb2RlIGEgTGFuY2UgZm9sZHMgcm93IGludG8gdGhlIGNvbHVtbnMgcmV0dXJuZWQgYnkgdGhlIHN0cnVjdHVyZSBBUEkuCgogICAgUmV0dXJucyBhIGRpY3Qgd2l0aCBrZXlzOiBwdG0sIG1lYW5fcGxkZHQsIHBlcl9yZXNpZHVlX3BsZGR0LCBwZGIuCiAgICAiIiIKICAgICMgcHRtIC8gbWVhbl9wbGRkdDogcHJlLWV4dHJhY3RlZCBzY2FsYXIgY29sdW1ucywgcmV0dXJuZWQgYXMtaXMuCiAgICBwdG0gPSByb3cuZ2V0KCJwdG0iKQogICAgbWVhbl9wbGRkdCA9IHJvdy5nZXQoIm1lYW5fcGxkZHQiKQoKICAgICMgcGVyX3Jlc2lkdWVfcGxkZHQ6IG5weiBmbG9hdDE2IGFycmF5IGJsb2IgLT4gbGlzdFtmbG9hdF0uCiAgICBwZXJfcmVzaWR1ZV9wbGRkdCA9IE5vbmUKICAgIGlmIHJvdy5nZXQoInBlcl9yZXNpZHVlX3BsZGR0IikgaXMgbm90IE5vbmU6CiAgICAgICAgYXJyID0gbnAubG9hZChpby5CeXRlc0lPKHJvd1sicGVyX3Jlc2lkdWVfcGxkZHQiXSkpWyJhcnIiXQogICAgICAgIHBlcl9yZXNpZHVlX3BsZGR0ID0gYXJyLmFzdHlwZShmbG9hdCkudG9saXN0KCkKCiAgICAjIHBkYjogYnJvdGxpIGRlY29tcHJlc3MgLT4gbXNncGFjayB1bnBhY2sgLT4gUHJvdGVpbkNoYWluIC0+IFBEQiBzdHJpbmcuCiAgICBwZGIgPSBOb25lCiAgICBpZiByb3cuZ2V0KCJzdHJ1Y3R1cmVfYmxvYiIpOgogICAgICAgIHN0YXRlX2RpY3QgPSBtc2dwYWNrLmxvYWRzKGJyb3RsaS5kZWNvbXByZXNzKHJvd1sic3RydWN0dXJlX2Jsb2IiXSksIHN0cmljdF9tYXBfa2V5PUZhbHNlKQogICAgICAgIHZhbGlkX2tleXMgPSBzZXQoaW5zcGVjdC5zaWduYXR1cmUoUHJvdGVpbkNoYWluLl9faW5pdF9fKS5wYXJhbWV0ZXJzKSAtIHsic2VsZiIsICJtbWNpZiJ9CiAgICAgICAgZmlsdGVyZWQgPSB7azogdiBmb3IgaywgdiBpbiBzdGF0ZV9kaWN0Lml0ZW1zKCkgaWYgayBpbiB2YWxpZF9rZXlzfQogICAgICAgICMgVGhlIGJsb2Igc3RvcmVzIGNoYWluX2lkL2VudGl0eV9pZCBhcyBwZXItcmVzaWR1ZSBpbnQgYXJyYXlzLCBidXQgdGhlCiAgICAgICAgIyBwaXAgZXNtIHBhY2thZ2UgZXhwZWN0cyBjaGFpbl9pZCBhcyBhIHNjYWxhciBzdHJpbmcuIENvbnZlcnQgYWNjb3JkaW5nbHkuCiAgICAgICAgaWYgaGFzYXR0cihmaWx0ZXJlZC5nZXQoImNoYWluX2lkIiksICJkdHlwZSIpIGFuZCBmaWx0ZXJlZFsiY2hhaW5faWQiXS5kdHlwZS5raW5kID09ICJpIjoKICAgICAgICAgICAgdW5pcXVlID0gbnAudW5pcXVlKGZpbHRlcmVkWyJjaGFpbl9pZCJdKQogICAgICAgICAgICBmaWx0ZXJlZFsiY2hhaW5faWQiXSA9IGNocihvcmQoIkEiKSArIGludCh1bmlxdWVbMF0pKQogICAgICAgIGlmIGhhc2F0dHIoZmlsdGVyZWQuZ2V0KCJlbnRpdHlfaWQiKSwgImR0eXBlIikgYW5kIGZpbHRlcmVkWyJlbnRpdHlfaWQiXS5kdHlwZS5raW5kID09ICJpIjoKICAgICAgICAgICAgZmlsdGVyZWRbImVudGl0eV9pZCJdID0gTm9uZQogICAgICAgIHBkYiA9IFByb3RlaW5DaGFpbi5mcm9tX3N0YXRlX2RpY3QoZmlsdGVyZWQpLnRvX3BkYl9zdHJpbmcoKQoKICAgIHJldHVybiB7CiAgICAgICAgInB0bSI6IHB0bSwKICAgICAgICAibWVhbl9wbGRkdCI6IG1lYW5fcGxkZHQsCiAgICAgICAgInBlcl9yZXNpZHVlX3BsZGR0IjogcGVyX3Jlc2lkdWVfcGxkZHQsCiAgICAgICAgInBkYiI6IHBkYiwKICAgIH0KCgpkZWYgZXh0cmFjdF9zdHJ1Y3R1cmUoZGF0YXNldF91cmk6IHN0ciwgcHJvdGVpbl9oYXNoOiBzdHIsIHJlZ2lvbjogc3RyID0gInVzLXdlc3QtMiIpIC0+IGRpY3QgfCBOb25lOgogICAgIiIiUmVhZCBvbmUgcHJvdGVpbidzIHJvdyBmcm9tIGEgTGFuY2UgZm9sZHMgZGF0YXNldCBhbmQgZGVjb2RlIGFsbCBBUEkgY29sdW1ucy4KCiAgICBBcmdzOgogICAgICAgIGRhdGFzZXRfdXJpOiBMb2NhbCBwYXRoIG9yIHMzOi8vIFVSSSB0byB0aGUgTGFuY2UgZm9sZHMgZGF0YXNldC4KICAgICAgICBwcm90ZWluX2hhc2g6IDMyLWNoYXIgaGV4IE1ENSBoYXNoIG9mIHRoZSBwcm90ZWluIHNlcXVlbmNlLgogICAgICAgIHJlZ2lvbjogQVdTIHJlZ2lvbiBmb3IgUzMgZGF0YXNldHMuCgogICAgUmV0dXJuczoKICAgICAgICBEaWN0IHdpdGggcHRtL21lYW5fcGxkZHQvcGVyX3Jlc2lkdWVfcGxkZHQvcGRiLCBvciBOb25lIGlmIG5vdCBmb3VuZC4KICAgICIiIgogICAgc3RvcmFnZV9vcHRpb25zID0geyJyZWdpb24iOiByZWdpb259CiAgICAjIFMzIEV4cHJlc3Mgb25lLXpvbmUgYnVja2V0cyAodGhvc2Ugd2l0aCBhICItLS4uLi0teC1zMyIgc3VmZml4KSBuZWVkIHZob3N0LXN0eWxlIHJlcXVlc3RzLgogICAgaWYgZGF0YXNldF91cmkuc3RhcnRzd2l0aCgiczM6Ly8iKSBhbmQgIi0tIiBpbiBkYXRhc2V0X3VyaS5zcGxpdCgiLyIpWzJdOgogICAgICAgIHN0b3JhZ2Vfb3B0aW9uc1sidmlydHVhbF9ob3N0ZWRfc3R5bGVfcmVxdWVzdCJdID0gInRydWUiCgogICAgZHMgPSBsYW5jZS5kYXRhc2V0KGRhdGFzZXRfdXJpLCBzdG9yYWdlX29wdGlvbnM9c3RvcmFnZV9vcHRpb25zKQogICAgdGFibGUgPSBkcy50b190YWJsZSgKICAgICAgICBmaWx0ZXI9ZiJwcm90ZWluX2hhc2ggPSAne3Byb3RlaW5faGFzaH0nIiwKICAgICAgICBjb2x1bW5zPUxBTkNFX0NPTFVNTlMsCiAgICApCiAgICBpZiB0YWJsZS5udW1fcm93cyA9PSAwOgogICAgICAgIHJldHVybiBOb25lCiAgICByZXR1cm4gZGVjb2RlX3N0cnVjdHVyZV9yb3codGFibGUudG9fcHlsaXN0KClbMF0pCgoKaWYgX19uYW1lX18gPT0gIl9fbWFpbl9fIjoKICAgIGltcG9ydCBzeXMKCiAgICB1cmksIHByb3RlaW5faGFzaCA9IHN5cy5hcmd2WzFdLCBzeXMuYXJndlsyXQogICAgcmVzdWx0ID0gZXh0cmFjdF9zdHJ1Y3R1cmUodXJpLCBwcm90ZWluX2hhc2gpCiAgICBhc3NlcnQgcmVzdWx0IGlzIG5vdCBOb25lCiAgICBwcmludCgicHRtOiIsIHJlc3VsdFsicHRtIl0pCiAgICBwcmludCgibWVhbl9wbGRkdDoiLCByZXN1bHRbIm1lYW5fcGxkZHQiXSkKICAgIHByaW50KCJwZXJfcmVzaWR1ZV9wbGRkdCAoZmlyc3QgNSk6IiwgcmVzdWx0WyJwZXJfcmVzaWR1ZV9wbGRkdCJdWzo1XSkKICAgIHByaW50KCJwZGIgKGZpcnN0IGxpbmUpOiIsIHJlc3VsdFsicGRiIl0uc3BsaXRsaW5lcygpWzBdKQo=" download="extract_fold_structure.py">Download extract\_fold\_structure.py</a>

```python theme={null}
"""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])
```