scripts/_verify_precompute.py

"""
Verify a precomputed h5ad + manifest edit for one dataset against the
precompute-cache verification checklist:

  1. .uns provenance round-trip (plain strings after write_h5ad/read_h5ad)
  2. dataset_describe(): loading_plan step 1 == decoupler_load_url_counts,
     validation.valid == True
  3. decoupler_load_url_counts() smoke test against the new HF url
  4. dataset_validate_manifest_against_data(): overall_valid == True

Usage: .venv/bin/python scripts/_verify_precompute.py <dataset_id>
"""
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

import anndata as ad

from src.tools.bulk_rnaseq import decoupler_load_url_counts
from src.tools.dataset_tools import dataset_describe, dataset_validate_manifest_against_data


def main():
    dataset_id = sys.argv[1]

    print(f"=== {dataset_id} ===\n")

    # 1. .uns provenance round-trip
    tmp_path = ROOT / "tmp" / "datasets" / f"{dataset_id}.h5ad"
    adata = ad.read_h5ad(tmp_path)
    print(f"local h5ad shape: {adata.shape}")
    print("precompute_* .uns keys (post round-trip):")
    for k in sorted(adata.uns):
        if k.startswith("precompute_"):
            v = adata.uns[k]
            print(f"  {k}: {v!r}  (type={type(v).__name__})")

    # 2. dataset_describe
    print("\n--- dataset_describe ---")
    desc = dataset_describe(dataset_id)
    print("validation:", desc["validation"])
    print("loading_plan:")
    for step in desc["loading_plan"]:
        print(" ", step)

    step1 = desc["loading_plan"][0]
    url = step1["key_args"]["url_or_path"]
    feature_id_type = step1["key_args"].get("feature_id_type", "gene_symbol")
    if feature_id_type not in ("gene_symbol", "ensembl_gene_id", "entrez_id"):
        feature_id_type = "gene_symbol"  # placeholder; .h5ad fast path doesn't use it

    # 3. smoke test: decoupler_load_url_counts against the new HF url
    print(f"\n--- decoupler_load_url_counts smoke test: {url} ---")
    result = decoupler_load_url_counts(
        url_or_path=url, feature_id_type=feature_id_type, out_prefix=f"verify_{dataset_id}"
    )
    for k in ("message", "n_obs", "n_vars", "shape", "obs_columns", "var_index_sample", "output_path"):
        print(f"  {k}: {result.get(k)}")

    # 4. dataset_validate_manifest_against_data
    print("\n--- dataset_validate_manifest_against_data ---")
    val = dataset_validate_manifest_against_data(dataset_id, result["output_path"])
    print(val)


if __name__ == "__main__":
    main()