Datasets¶

from pandora.datasets import curate_structure
from pandora.metadata import collect_metadata
from pandora.parsing import mmcif_to_structure
from pandora.schemas.dataset import DatasetCurationPolicy, QualityRules

structure, _, _ = mmcif_to_structure("datasets/dev/mmcif/104m.cif")
metadata = collect_metadata(structure)

policy = DatasetCurationPolicy(
    policy_id="c1", policy_name="Default", policy_version="1.0.0"
)
curated, exclusions, provenance = curate_structure(structure, metadata, policy)
print("default:", curated is not None, exclusions)

strict = policy.model_copy(
    update={"quality_rules": QualityRules(max_resolution=1.0)}
)
curated, exclusions, provenance = curate_structure(structure, metadata, strict)
print("strict:", curated is not None)
for e in exclusions:
    print(f"  {e.reason_code} - {e.message}")

from pandora.datasets import curate_structure
from pandora.metadata import collect_metadata
from pandora.parsing import mmcif_to_structure
from pandora.schemas.dataset import DatasetCurationPolicy, QualityRules

# 22jy: a cryo-EM structure at 2.2 A.
structure, _, _ = mmcif_to_structure("datasets/dev/mmcif/22jy.cif")
metadata = collect_metadata(structure)
print(metadata.quality.experimental_method, metadata.quality.resolution)

rules = {
    "one limit, 3.5": QualityRules(max_resolution=3.5),
    "X-ray 2.5, cryo-EM 2.0": QualityRules(
        max_resolution_by_method={
            "X-RAY DIFFRACTION": 2.5,
            "ELECTRON MICROSCOPY": 2.0,
        }
    ),
}
for label, quality_rules in rules.items():
    policy = DatasetCurationPolicy(
        policy_id="c",
        policy_name="c",
        policy_version="1.0.0",
        quality_rules=quality_rules,
    )
    curated, exclusions, _ = curate_structure(structure, metadata, policy)
    print(f"{label}: kept={curated is not None}", end="")
    print("".join(f" ({e.reason_code}: {e.message})" for e in exclusions))

from pandora.datasets import curate_structure
from pandora.metadata import collect_metadata
from pandora.parsing import mmcif_to_structure
from pandora.schemas.dataset import DatasetCurationPolicy, QualityRules


def check(entry_id: str, quality_rules: QualityRules) -> str:
    structure, _, _ = mmcif_to_structure(f"datasets/dev/mmcif/{entry_id}.cif")
    policy = DatasetCurationPolicy(
        policy_id="c",
        policy_name="c",
        policy_version="1.0.0",
        quality_rules=quality_rules,
    )
    curated, exclusions, _ = curate_structure(
        structure, collect_metadata(structure), policy
    )
    if curated is not None:
        return "kept"
    return f"{exclusions[-1].reason_code} ({exclusions[-1].message})"


# 1a08 is X-ray. It reports Rmerge but neither Rsym nor Rfree.
structure, _, _ = mmcif_to_structure("datasets/dev/mmcif/1a08.cif")
quality = collect_metadata(structure).quality
print(
    f"1a08: r_free={quality.r_free} r_work={quality.r_work} "
    f"r_sym={quality.r_sym} r_merge={quality.r_merge}"
)

# Rsym falls back to Rmerge.
print("max_r_sym=0.05:", check("1a08", QualityRules(max_r_sym=0.05)))
# A missing Rfree is kept by default...
print("max_r_free=0.25:", check("1a08", QualityRules(max_r_free=0.25)))
# ...or excluded if you ask for it.
print(
    "max_r_free=0.25, exclude nulls:",
    check(
        "1a08",
        QualityRules(max_r_free=0.25, null_rfactor_behavior="exclude"),
    ),
)
# R-factor rules only check X-ray entries, so cryo-EM 22jy passes.
print(
    "22jy, same rule:",
    check(
        "22jy",
        QualityRules(max_r_free=0.25, null_rfactor_behavior="exclude"),
    ),
)

from pandora.canonicalisation import canonicalise_structure
from pandora.datasets import curate_structure
from pandora.parsing import mmcif_to_structure
from pandora.schemas.canonicalisation import (
    ModifiedResidueRules,
    canonicalisationPolicy,
)
from pandora.schemas.dataset import DatasetCurationPolicy, QualityRules


def curate(structure, quality_rules: QualityRules) -> str:
    policy = DatasetCurationPolicy(
        policy_id="c",
        policy_name="c",
        policy_version="1.0.0",
        quality_rules=quality_rules,
    )
    curated, exclusions, _ = curate_structure(structure, None, policy)
    return "kept" if curated else exclusions[-1].message


reject = QualityRules(exclude_nonstandard_residues=True)

# 1a08 contains ACE and DIP caps and FTY, a modified tyrosine.
structure, _, _ = mmcif_to_structure("datasets/dev/mmcif/1a08.cif")
print("1a08:", curate(structure, reject))
allow = reject.model_copy(
    update={"allowed_nonstandard_residues": ["ACE", "DIP", "FTY"]}
)
print("1a08, allow-list:", curate(structure, allow))

# 1b6w contains MSE. Mapping it to MET in canonicalisation makes it
# standard by the time curation runs.
structure, _, _ = mmcif_to_structure("datasets/dev/mmcif/1b6w.cif")
print("1b6w as deposited:", curate(structure, reject))
canonical, _, _ = canonicalise_structure(
    structure,
    canonicalisationPolicy(
        policy_id="p",
        policy_name="p",
        policy_version="1.0.0",
        modified_residue_rules=ModifiedResidueRules(strategy="map_to_parent"),
    ),
)
print("1b6w after MSE -> MET:", curate(canonical, reject))

from pandora.datasets import curate_structure
from pandora.parsing import mmcif_to_structure
from pandora.schemas.dataset import DatasetCurationPolicy, QualityRules

# 1aui chain A misses 95 residues in its middle; chain B is nearly whole.
structure, _, _ = mmcif_to_structure("datasets/dev/mmcif/1aui.cif")


def curate(quality_rules: QualityRules) -> None:
    policy = DatasetCurationPolicy(
        policy_id="c",
        policy_name="c",
        policy_version="1.0.0",
        quality_rules=quality_rules,
    )
    curated, exclusions, _ = curate_structure(structure, None, policy)
    for e in exclusions:
        print(f"  {e.chain_id or 'entry'}: {e.reason_code} - {e.message}")
    if curated is not None:
        chains = sorted(
            {a.label_asym_id for a in curated.atoms if a.label_seq_id}
        )
        print(f"  kept polymer chains: {chains}")


print("max 30% missing tails, 10% missing middle:")
curate(
    QualityRules(max_missing_tail_fraction=0.3, max_missing_middle_fraction=0.1)
)
print("the same, but the entry needs two chains:")
curate(
    QualityRules(
        max_missing_tail_fraction=0.3,
        max_missing_middle_fraction=0.1,
        min_polymer_chains=2,
    )
)

from pandora.datasets import curate_structure
from pandora.parsing import mmcif_to_structure
from pandora.schemas.dataset import DatasetCurationPolicy, QualityRules


def curate(entry_id: str, quality_rules: QualityRules) -> None:
    structure, _, _ = mmcif_to_structure(f"datasets/dev/mmcif/{entry_id}.cif")
    policy = DatasetCurationPolicy(
        policy_id="c",
        policy_name="c",
        policy_version="1.0.0",
        quality_rules=quality_rules,
    )
    curated, exclusions, _ = curate_structure(structure, None, policy)
    print(f"{entry_id}: kept={curated is not None}")
    for e in exclusions:
        print(f"  {e.chain_id or 'entry'}: {e.reason_code} - {e.message}")


# Whole entries: 1aui has 4,832 atoms.
curate("1aui", QualityRules(max_atoms=4000))
# Single chains: 1aui chain A has 378 observed residues, chain B 165.
curate("1aui", QualityRules(max_chain_length=300))

from pandora.datasets import curate_structure
from pandora.metadata import collect_metadata
from pandora.parsing import mmcif_to_structure
from pandora.schemas.dataset import DatasetCurationPolicy, OrganismRules

policy = DatasetCurationPolicy(
    policy_id="c",
    policy_name="Human only",
    policy_version="1.0.0",
    organism_rules=OrganismRules(include_taxa=["9606"]),
)
for entry_id in ["1aui", "1ayi"]:
    structure, _, _ = mmcif_to_structure(f"datasets/dev/mmcif/{entry_id}.cif")
    metadata = collect_metadata(structure)
    curated, exclusions, _ = curate_structure(structure, metadata, policy)
    organisms = sorted({t.organism_scientific for t in metadata.taxonomies})
    outcome = "kept" if curated else exclusions[-1].reason_code
    print(f"{entry_id} {organisms}: {outcome}")

from pandora.datasets import curate_structure
from pandora.parsing import mmcif_to_structure
from pandora.schemas.dataset import ContentRules, DatasetCurationPolicy

structure, _, _ = mmcif_to_structure("datasets/dev/mmcif/104m.cif")

policy = DatasetCurationPolicy(
    policy_id="c",
    policy_name="Protein only",
    policy_version="1.0.0",
    content_rules=ContentRules(
        keep_ligands=False, keep_waters=False, keep_ions=False
    ),
)
curated, exclusions, _ = curate_structure(structure, None, policy)
print(len(structure.atoms), "->", len(curated.atoms), "atoms;", exclusions)

from pandora.datasets import curate_structure
from pandora.metadata import collect_metadata
from pandora.parsing import mmcif_to_structure
from pandora.schemas.dataset import DatasetCurationPolicy, QualityRules

policy = DatasetCurationPolicy(
    policy_id="c",
    policy_name="Sharp, complete chains",
    policy_version="1.0.0",
    quality_rules=QualityRules(
        max_resolution=2.5,
        max_missing_middle_fraction=0.1,
    ),
)

kept = {}
exclusions = []
for entry_id in ["104m", "1aui", "1a08", "1p58"]:
    structure, _, _ = mmcif_to_structure(f"datasets/dev/mmcif/{entry_id}.cif")
    curated, records, _ = curate_structure(
        structure, collect_metadata(structure), policy
    )
    exclusions.extend(records)  # chain records from kept entries too
    if curated is not None:
        kept[entry_id] = curated

entries_out = [e for e in exclusions if e.chain_id is None]
chains_out = [e for e in exclusions if e.chain_id is not None]
print("kept:", sorted(kept))
print("entries excluded:", [(e.entry_id, e.reason_code) for e in entries_out])
print(
    "chains removed:",
    [(e.entry_id, e.chain_id, e.reason_code) for e in chains_out],
)

from pandora.datasets import deduplicate_structures
from pandora.schemas.dataset import DeduplicationRules

duplicate, _, _ = mmcif_to_structure("datasets/dev/mmcif/104m.cif")
duplicate_canonical, _, _ = canonicalise_structure(duplicate, policy)

retained, removed, provenance = deduplicate_structures(
    [canonical, duplicate_canonical], DeduplicationRules(enabled=True)
)
print(
    len(retained),
    "retained,",
    len(removed),
    "removed:",
    [r.reason_code for r in removed],
)
# 1 retained, 1 removed: ['DUPLICATE']