diff --git a/genomics/.env.example b/genomics/.env.example new file mode 100644 index 0000000..1f7a3c5 --- /dev/null +++ b/genomics/.env.example @@ -0,0 +1,34 @@ +# ProGenome genomics: configuration template. +# cp .env.example .env # then fill in what you use; .env is git-ignored and read by config.py and every script +# Alternative location: ~/.progenome.env (same format, kept outside the repository). +# Precedence: variables exported in your shell > genomics/.env > ~/.progenome.env > defaults in config.py. +# Keep comments on their own lines. Check what is in effect with: make config + +# --- LLM decoder (make decode, run_v2.sh) -------------------------------------------------------- +# Key from https://build.nvidia.com (sign in, any model page, "Get API Key"); starts with nvapi- +NVIDIA_API_KEY=nvapi-REPLACE_ME +# Model id on the endpoint (default: Nemotron 3 Super) +NIM_MODEL=nvidia/nemotron-3-super-120b-a12b +# Any OpenAI-compatible chat-completions endpoint (a local NIM container, vLLM, ...) +NIM_URL=https://integrate.api.nvidia.com/v1/chat/completions + +# --- Neo4j browser (make neo4j-load) ---------------------------------------------------------------- +# Defaults match docker-compose.yml; change both if you change one +NEO4J_URI=bolt://localhost:7687 +NEO4J_USER=neo4j +NEO4J_PASSWORD=progenome + +# --- Data and defaults ------------------------------------------------------------------------------ +# Where fetch_data.sh downloads the HaploGraph from +HAPLOBLOCKS_BASE=https://data.haploblocks.org +# Chromosome for every stage (also: CHROM=chr21 make run) +CHROM=chr22 +# Optional: move the 15 MB of data and the ~400 MB of outputs elsewhere +# PROGENOME_DATA_DIR=/data/progenome/data +# PROGENOME_OUTPUTS_DIR=/data/progenome/outputs + +# --- NVIDIA Brev (make brev) ------------------------------------------------------------------------- +# The CLI stores its own login (brev login --api-key ...); these only choose the instance +BREV_INSTANCE=progenome-gpu +BREV_TYPE=g2-standard-4:nvidia-l4:1 +# A100 example: BREV_INSTANCE=progenome-a100 BREV_TYPE=a100-80gb.1x diff --git a/genomics/.gitignore b/genomics/.gitignore new file mode 100644 index 0000000..0dcb2fd --- /dev/null +++ b/genomics/.gitignore @@ -0,0 +1,9 @@ +__pycache__/ +.pytest_cache/ +.venv/ +*.pyc +# downloaded inputs (re-fetch with fetch_data.sh) and generated outputs +data/ +outputs_brev/ +outputs/ +.env diff --git a/genomics/Dockerfile b/genomics/Dockerfile new file mode 100644 index 0000000..31000f9 --- /dev/null +++ b/genomics/Dockerfile @@ -0,0 +1,34 @@ +# ProGenome genomics pipeline: knowledge graph -> co-occurrence analysis -> GNN (PyTorch Geometric). +# CUDA image for NVIDIA Brev / any GPU box; the same image runs on CPU when no GPU is present. +# +# Build from the REPOSITORY ROOT so the two proteomics inputs are baked in (make docker does this): +# docker build -f genomics/Dockerfile -t progenome-genomics . +# Run (data/ and outputs/ bind-mounted; .env passed for the decoder): +# docker run --rm --gpus all --env-file genomics/.env -v $PWD/genomics/data:/app/genomics/data -v $PWD/genomics/outputs:/app/genomics/outputs progenome-genomics run_all.sh +# docker run --rm --gpus all -v ... progenome-genomics run_v2.sh +# docker run --rm --gpus all -v ... progenome-genomics -c "python train_gnn.py --target ancestry --init node2vec" +ARG BASE=pytorch/pytorch:2.14.0-cuda12.6-cudnn9-runtime +FROM ${BASE} +ARG PYG_WHEELS=https://data.pyg.org/whl/torch-2.14.0+cu126.html + +ENV PYTHONUNBUFFERED=1 PIP_NO_CACHE_DIR=1 PIP_BREAK_SYSTEM_PACKAGES=1 NX_CUGRAPH_AUTOCONFIG=True +RUN apt-get update && apt-get install -y --no-install-recommends curl ca-certificates && rm -rf /var/lib/apt/lists/* + +WORKDIR /app/genomics +COPY genomics/requirements.txt . +RUN pip install -r requirements.txt \ + && (pip install pyg_lib -f ${PYG_WHEELS} \ + || echo "no pyg_lib wheel for this torch build -> Node2Vec falls back to SVD") \ + && (pip install nx-cugraph-cu12 --extra-index-url https://pypi.nvidia.com \ + || echo "nx-cugraph not installed -> NetworkX runs on CPU") \ + && (pip install torch-tensorrt --extra-index-url https://download.pytorch.org/whl/cu126 \ + || echo "torch-tensorrt not installed -> infer.py --compile tensorrt falls back to inductor") + +COPY proteomics/uniprot_chr22.bed /app/proteomics/uniprot_chr22.bed +COPY proteomics/synthetic_proteomics_chr22/gene_symbol_cache.csv /app/proteomics/synthetic_proteomics_chr22/gene_symbol_cache.csv +COPY genomics/ /app/genomics/ +RUN python -m pytest tests -q + +# data/ and outputs/ are bind-mounted at run time (see docker-compose.yml) +ENTRYPOINT ["bash"] +CMD ["run_all.sh"] diff --git a/genomics/Dockerfile.dockerignore b/genomics/Dockerfile.dockerignore new file mode 100644 index 0000000..88a371a --- /dev/null +++ b/genomics/Dockerfile.dockerignore @@ -0,0 +1,15 @@ +# Build context is the repository root (docker build -f genomics/Dockerfile ..); keep the context small. +# BuildKit reads .dockerignore next to the Dockerfile. +* +!genomics/** +!proteomics/uniprot_chr22.bed +!proteomics/synthetic_proteomics_chr22/gene_symbol_cache.csv +genomics/data +genomics/outputs +genomics/outputs_brev +genomics/.venv +genomics/.env +genomics/docs/report/*.pdf +genomics/docs/report/*.docx +**/__pycache__ +**/.pytest_cache diff --git a/genomics/Makefile b/genomics/Makefile new file mode 100644 index 0000000..c645cac --- /dev/null +++ b/genomics/Makefile @@ -0,0 +1,93 @@ +# Clone-and-go entry points for the genomics pipeline. `make help` lists them. +CHROM ?= chr22 +INIT ?= svd +PY := .venv/bin/python + +.PHONY: help setup data kg analysis baseline graph gnn embeddings run run-v2 eda decode federated test docker docker-run docker-run-v2 docker-federated docker-shell config neo4j neo4j-load neo4j-down brev report clean + +help: ## show this help + @grep -E '^[a-z-]+:.*##' $(MAKEFILE_LIST) | awk -F':.*## ' '{printf " make %-12s %s\n", $$1, $$2}' + +setup: ## create .venv with torch (CPU or CUDA) + pinned deps, run unit tests + bash setup.sh + +data: ## download the HaploGraph, phenotypes and block stats for $(CHROM) + bash fetch_data.sh $(CHROM) + +kg: data ## build the knowledge graph (tables, sparse carriers, PyG HeteroData) + $(PY) build_kg.py --chrom $(CHROM) + +analysis: kg ## cluster/edge/block vs phenotype co-occurrence tests + plots + $(PY) cooccurrence_analysis.py --chrom $(CHROM) + +baseline: kg ## logistic-regression baseline, writes the shared train/val/test split + $(PY) baseline.py --chrom $(CHROM) + +graph: analysis ## NetworkX stats, GraphML export and plots of the graph + $(PY) graph_explore.py --chrom $(CHROM) + +gnn: baseline ## train the hetero-GNN on ancestry, population and sex (control) + for t in ancestry population sex; do $(PY) train_gnn.py --chrom $(CHROM) --target $$t --init $(INIT); done + +embeddings: gnn ## evaluate/plot SVD and GNN embeddings + $(PY) embeddings.py --chrom $(CHROM) + +run: ## the whole pipeline (same as run_all.sh) + PYTHON=$(PY) CHROM=$(CHROM) INIT=$(INIT) bash run_all.sh + +run-v2: run ## schema v2: proteomics on 1000G IDs, genes/proteins, genome+proteome GNN, EDA, decoder dry-run + PYTHON=$(PY) CHROM=$(CHROM) bash run_v2.sh + +eda: ## exploratory data analysis report -> outputs/eda/$(CHROM)/EDA.md + $(PY) eda.py --chrom $(CHROM) + +decode: ## GraphRAG insight for one person via NVIDIA NIM (needs NVIDIA_API_KEY; WHO=HG00103) + $(PY) graphrag_decoder.py --chrom $(CHROM) --individual $(or $(WHO),HG00103) --run phenotype_both_raw + +federated: ## NVFlare FedAvg over 3 sites (simulator), score the global model centrally, compare with each site alone + $(PY) federated/job.py --chrom $(CHROM) --rounds $(or $(ROUNDS),10) --local-epochs $(or $(LOCAL_EPOCHS),5) + $(PY) federated/evaluate_global.py --chrom $(CHROM) + $(PY) federated/local_only.py --chrom $(CHROM) --steps $$(( $(or $(ROUNDS),10) * $(or $(LOCAL_EPOCHS),5) )) + +test: ## unit tests + $(PY) -m pytest tests -q + +DOCKER_RUN = docker run --rm $$(command -v nvidia-smi >/dev/null 2>&1 && echo --gpus all) $$( [ -f .env ] && echo --env-file .env ) \ + -e CHROM=$(CHROM) -e INIT=$(INIT) -v $(PWD)/data:/app/genomics/data -v $(PWD)/outputs:/app/genomics/outputs progenome-genomics + +docker: ## build the CUDA image from the repository root (runs on CPU too); proteomics inputs are baked in + docker build -f Dockerfile -t progenome-genomics .. + +docker-run: ## v1 pipeline inside the image (GPU if available; .env passed if present) + $(DOCKER_RUN) run_all.sh + +docker-run-v2: ## v2 chain inside the image (after docker-run) + $(DOCKER_RUN) run_v2.sh + +docker-federated: ## NVFlare FedAvg inside the image (after docker-run-v2) + $(DOCKER_RUN) -c "python federated/job.py --chrom $(CHROM) --rounds $(or $(ROUNDS),30) --local-epochs $(or $(LOCAL_EPOCHS),5) && python federated/evaluate_global.py --chrom $(CHROM) && python federated/local_only.py --chrom $(CHROM) --steps $$(( $(or $(ROUNDS),30) * $(or $(LOCAL_EPOCHS),5) ))" + +docker-shell: ## interactive shell inside the image with data/ and outputs/ mounted + $(DOCKER_RUN:--rm=--rm -it) + +config: ## show the effective configuration (.env / ~/.progenome.env / defaults), secrets masked + $(PY) config.py + +neo4j: ## start Neo4j community at http://localhost:7474 (neo4j / progenome) + docker compose up -d neo4j + +neo4j-load: neo4j ## load the knowledge graph into Neo4j + $(PY) neo4j_load.py --chrom $(CHROM) + +neo4j-down: ## stop Neo4j (keeps its data volume) + docker compose down + +brev: ## create/use a Brev GPU instance, build there, run with Node2Vec + TensorRT benchmark, copy outputs back + bash brev_deploy.sh + +report: ## rebuild docs/report (figures -> .tex + .docx -> .pdf); needs node with the docx package and tectonic + $(PY) docs/report/make_report_figures.py + cd docs/report && node build_report.js && (command -v tectonic >/dev/null && tectonic ProGenome_KT.tex || echo "tectonic not installed: compile ProGenome_KT.tex with pdflatex/xelatex") + +clean: ## remove generated outputs (keeps downloaded data) + rm -rf outputs diff --git a/genomics/README.md b/genomics/README.md new file mode 100644 index 0000000..045246b --- /dev/null +++ b/genomics/README.md @@ -0,0 +1,322 @@ +# genomics/ — haploblock knowledge graph → phenotypes → proteomics → GNN → LLM + +[![Python](https://img.shields.io/badge/Python-3.10--3.13-3776AB?logo=python&logoColor=white)](https://www.python.org/) +[![PyTorch](https://img.shields.io/badge/PyTorch-2.14-EE4C2C?logo=pytorch&logoColor=white)](https://pytorch.org/) +[![PyTorch Geometric](https://img.shields.io/badge/PyTorch%20Geometric-2.8-3C2179)](https://pyg.org/) +[![CUDA](https://img.shields.io/badge/CUDA-12.6-76B900?logo=nvidia&logoColor=white)](https://developer.nvidia.com/cuda-toolkit) +[![Docker](https://img.shields.io/badge/Docker-Containerized-2496ED?logo=docker&logoColor=white)](Dockerfile) +[![NVIDIA FLARE](https://img.shields.io/badge/NVIDIA%20FLARE-2.9%20FedAvg-76B900?logo=nvidia&logoColor=white)](https://github.com/NVIDIA/NVFlare) +[![NVIDIA NIM](https://img.shields.io/badge/NVIDIA%20NIM-Nemotron%203%20Super-76B900?logo=nvidia&logoColor=white)](https://build.nvidia.com/) +[![NVIDIA Brev](https://img.shields.io/badge/NVIDIA%20Brev-A100%2080GB-76B900?logo=nvidia&logoColor=white)](https://brev.nvidia.com/) +[![Neo4j](https://img.shields.io/badge/Neo4j-5.26%20community-008CC1?logo=neo4j&logoColor=white)](https://neo4j.com/) +[![NetworkX](https://img.shields.io/badge/NetworkX-3.6%20%2B%20nx--cugraph-1B6AC6)](https://networkx.org/) +[![Data](https://img.shields.io/badge/Data-1000G%20HaploGraph%20chr22-0E7C7B)](https://data.haploblocks.org/haplograph/1000G/) +[![Tests](https://img.shields.io/badge/tests-11%20passing-brightgreen?logo=pytest&logoColor=white)](tests/) +[![Hackathon](https://img.shields.io/badge/Nordic%20Biobank%20x%20NVIDIA-Federated%20Learning%20Hackathon%202026-5A9E3F)](https://github.com/collaborativebioinformatics/ProGenome) + +The genome side of ProGenome, plus the join to proteomics. It takes the **published 1000 Genomes +HaploGraph** (built for this hackathon at Rigshospitalet, ), +connects it to the **real 1000G labels** (ancestry, population, sex) and to **per-site proteomics keyed by the +same sample IDs**, trains a heterogeneous **PyTorch Geometric GNN**, and decodes a person's graph neighbourhood +with an **LLM (NVIDIA NIM)** into a cited insight. Everything runs from a clone; chr22 takes ~10 min on a laptop CPU +and ~2 min on an A100. + +Architecture page (data flow, schema, federated topology, stack): `docs/architecture.html` · +Mermaid source for Lucidchart / GitHub: `docs/architecture.mmd`. + +## Quick start + +```bash +git clone https://github.com/collaborativebioinformatics/ProGenome.git +cd ProGenome && git checkout modelling && cd genomics +make setup # .venv with torch (CPU, or CUDA if nvidia-smi works) + pinned deps, runs the unit tests (~30 s with uv) +make run # v1: fetch -> knowledge graph -> co-occurrence analysis -> baseline -> graph plots -> GNN -> embeddings (~6 min CPU) +make run-v2 # v2: proteomics on 1000G IDs -> genes/proteins in the graph -> EDA -> genome+proteome GNN -> ridge -> decoder (~19 min CPU) +make federated ROUNDS=30 LOCAL_EPOCHS=5 # NVFlare FedAvg over 3 sites + central scoring + each site alone (~4 min CPU) +make decode WHO=HG00103 # LLM insight for one person via NVIDIA NIM (needs an API key, step 4 below) +make neo4j-load # browse the graph at http://localhost:7474 (neo4j / progenome) +make docker # CUDA image (pytorch 2.14 + cu12.6, PyG 2.8, pyg-lib, nx-cugraph, torch-tensorrt); runs on CPU too +make brev # create/use a Brev GPU instance, build there, run everything, copy outputs back +make help # every target +``` + +## Implementation guide: from clone to every result + +Everything below was executed end-to-end on 18 Sept 2026 from an empty folder on a laptop (Apple M2, CPU only, 30 min) +and on an A100 through the Docker image (14 min). Numbers you should see are in the *Results* section; the tolerance +between runs is stated there. + +### 0. What you need + +| Need | Details | +|---|---| +| OS, Python | macOS or Linux; Python 3.10-3.13 (3.13 tested). `uv` is optional and makes `make setup` take 30 s instead of minutes. | +| Tools | git, curl, make, bash. `md5sum` or macOS `md5` for the download check. | +| Resources | ~8 GB RAM, ~2 GB disk (15 MB download, ~400 MB of outputs). No GPU needed for chr22. | +| Network | data.haploblocks.org (the graph), download.pytorch.org and PyPI (setup), integrate.api.nvidia.com (decoder only). | +| Optional GPU | NVIDIA driver with CUDA 12.6: `setup.sh` picks CUDA wheels automatically. Docker + NVIDIA Container Toolkit for the image. | +| Optional cloud | NVIDIA Brev CLI and account for `make brev`. | +| Optional LLM | An NVIDIA API key (free tier is enough) for `make decode`; without it the decoder prints the prompt and stops. | + +### 1. Clone and set up + +```bash +git clone https://github.com/collaborativebioinformatics/ProGenome.git +cd ProGenome && git checkout modelling && cd genomics +make setup +``` + +`setup.sh` creates `./.venv`, installs torch 2.14.0 from the CPU index (or cu126 if `nvidia-smi` works), the pinned +`requirements.txt`, then runs the 11 unit tests and prints the torch / PyG / CUDA versions. Force a variant with +`TORCH_INDEX=cpu bash setup.sh` or `PYTHON=python3.12 bash setup.sh`. Every later command uses `.venv/bin/python` through +the Makefile, so nothing needs activating. + +### 2. Genome side (v1): `make run` + +Runs `fetch_data.sh` (md5-verified download of the chr22 HaploGraph, phenotypes and block statistics into `data/`), +`build_kg.py`, `cooccurrence_analysis.py`, `baseline.py`, `graph_explore.py`, `train_gnn.py` for ancestry, population +and sex, and `embeddings.py`. Check these files when it finishes: + +| File | What to expect | +|---|---| +| `outputs/kg/chr22/summary.json` | 2,548 individuals, 6,551 kept clusters, 2,365,574 CARRIES, 187,030 CO_OCCURS, 0 edges dropped | +| `outputs/cooccurrence/chr22/summary.json` | 6,470 ancestry-associated clusters, 0 sex-associated, edge cosine 0.86 vs 0.22 | +| `outputs/baseline/chr22/metrics.json` | test balanced accuracy ancestry 0.977, population 0.614, sex 0.463 | +| `outputs/gnn/chr22/ancestry_svd/metrics.json` | 0.974 (population_svd 0.437, sex_svd ~0.50); `history.csv`, `test_predictions.csv`, embeddings | +| `outputs/embeddings/chr22/embedding_quality.json` | silhouette by ancestry 0.06 (SVD) -> 0.70 (GNN) | + +Options: `CHROM=chr21 make run` (any chromosome on the server), `INIT=raw make gnn` (raw carrier row; population 0.611), +`INIT=node2vec` (needs pyg-lib, i.e. the GPU image; `--node2vec-epochs`, default 50). + +### 3. Proteomics integration (v2): `make run-v2` + +Runs `proteomics_synth_1000g.py` (synthetic proteomics on the real 1000G ids, 3 sites, `ground_truth.json`), +`build_kg_v2.py`, `eda.py`, `train_gnn_v2.py` for genome / proteome / both (SVD and raw input), the site and ancestry +controls, `proteome_linear_baseline.py`, and the decoder in dry-run mode (plus one real call if a key is configured). + +| File | What to expect | +|---|---| +| `outputs/kg/chr22/summary_v2.json` | 458 genes, 460 proteins, 1,063 block-gene overlaps, 1,065,712 MEASURED edges | +| `outputs/eda/chr22/EDA.md` | the 8-section report with tables and plots | +| `outputs/gnn_v2/chr22/phenotype_{genome,proteome,both}_svd/metrics.json` | AUC ~0.64 / 0.96 / 0.99 | +| `outputs/gnn_v2/chr22/phenotype_both_raw/metrics.json` | AUC ~0.99, `saliency.hits_in_ground_truth` 4 of 20, `saliency_top100.csv` | +| `outputs/gnn_v2/chr22/site_both_svd/metrics.json` | balanced accuracy ~0.30-0.34 (chance 0.33: the batch control) | +| `outputs/gnn_v2/chr22/proteome_ridge_baseline/metrics.json` | 4 of 20 cis proteins with test R2 > 0.1, 0 of 440 others | + +### 4. Configuration, credentials and the LLM decoder (NVIDIA NIM) + +All settings enter the code through **`config.py`**; scripts never read credentials on their own. Precedence: +variables exported in your shell > `genomics/.env` > `~/.progenome.env` > defaults. + +```bash +cp .env.example .env # template with every variable and a comment; .env is git-ignored +make config # prints what is in effect and where each value came from (secrets masked) +``` + +| Variable | Used by | Default | +|---|---|---| +| `NVIDIA_API_KEY` | `graphrag_decoder.py` (`make decode`, the real call in `run_v2.sh`) | none: required for the decoder | +| `NIM_MODEL`, `NIM_URL` | decoder model id and endpoint (any OpenAI-compatible chat server) | `nvidia/nemotron-3-super-120b-a12b`, `https://integrate.api.nvidia.com/v1/chat/completions` | +| `NEO4J_URI`, `NEO4J_USER`, `NEO4J_PASSWORD` | `neo4j_load.py` | `bolt://localhost:7687`, `neo4j`, `progenome` (matches docker-compose.yml) | +| `HAPLOBLOCKS_BASE`, `CHROM` | `fetch_data.sh`, every stage | `https://data.haploblocks.org`, `chr22` | +| `PROGENOME_DATA_DIR`, `PROGENOME_OUTPUTS_DIR` | move data / outputs elsewhere | `genomics/data`, `genomics/outputs` | +| `BREV_INSTANCE`, `BREV_TYPE` | `brev_deploy.sh` | `progenome-gpu`, `g2-standard-4:nvidia-l4:1` | + +Getting and using the LLM key: + +1. Sign in at , open any model page, click *Get API Key*; it starts with `nvapi-`. The free tier + covers this (one call is ~5,600 tokens, 4-13 s). +2. Put it in `.env` (`NVIDIA_API_KEY=nvapi-...`) or in `~/.progenome.env` outside the repository (`chmod 600`). +3. `make decode WHO=HG00103` (any 1000G id in the graph; people in the held-out split also get the GNN's prediction). + Output: `outputs/graphrag/chr22/HG00103_insight.json` (summary, ancestry and phenotype assessments, genome-proteome + links, caveats, cited ids, citation check, token usage) and `HG00103_context.json` (exactly what the model was given). + `.venv/bin/python graphrag_decoder.py --individual HG00103 --dry-run` prints the prompt without calling anything; + `--reasoning low|medium|high` turns on Nemotron's separate reasoning field (slower); `--run` picks another GNN run folder. +4. The call is `POST $NIM_URL` with `Authorization: Bearer $NVIDIA_API_KEY` and body `{model, messages, temperature 0.2, + max_tokens 4000, reasoning_effort "none"}`. A local NIM container, vLLM or any OpenAI-compatible server works by + changing `NIM_URL` and `NIM_MODEL`; a local NIM on the GPU box keeps patient context on site. + +### 5. Federated learning (NVFlare): `make federated ROUNDS=30 LOCAL_EPOCHS=5` + +Runs `federated/job.py` (FedAvgRecipe, 3 simulated sites as threads, only weights exchanged), `federated/evaluate_global.py` +(scores `FL_global_model.pt` on the same 376 held-out people as the central model) and `federated/local_only.py` +(each site alone, same number of steps). Expect in `outputs/federated/chr22/evaluation.json` a global AUC of 0.987-0.998 +against a central 0.992-0.995, and in `local_only_vs_federated.json` site-alone AUCs of 0.986-0.998. The NVFlare +workspace (logs, per-round metrics, global models) is under `outputs/federated/chr22/workspace/`. Ten rounds are not +enough (AUC 0.955); thirty converge. To run across real machines use NVFlare POC mode with the same `client.py` and +`model.py`; each site needs the shared graph files (`outputs/kg/chr22/`), its own people and the split file. + +### 6. Browse the graph: `make neo4j-load` + +Starts Neo4j 5.26 community with docker compose and loads the graph (a few minutes for the 2.4 M CARRIES edges). Open + (user `neo4j`, password `progenome`) and try +`MATCH (i:Individual {id:'HG00096'})-[:CARRIES]->(c:Cluster)-[:IN_BLOCK]->(b:Block) RETURN i,c,b LIMIT 50`. +`make neo4j-down` stops it and keeps the volume. + +### 7. GPU: Docker image and Brev + +* **Docker is the whole solution in one image.** `make docker` builds `progenome-genomics` from the repository root + (`docker build -f genomics/Dockerfile ..`) on `pytorch/pytorch:2.14.0-cuda12.6-cudnn9-runtime` with PyG 2.8, pyg-lib, + nx-cugraph, torch-tensorrt and NVFlare; the two proteomics inputs are baked in and the 11 unit tests run inside the build + (~10 GB, ~4 min on an A100). Only `data/` (15 MB, downloaded on first run) and `outputs/` are bind-mounted, and `.env` + is passed in when present. Then, on a Linux box with the NVIDIA Container Toolkit (or CPU-only without `--gpus`): + ```bash + make docker-run # v1 chain make docker-run-v2 # v2 chain incl. decoder if .env has a key + make docker-federated # NVFlare + scoring make docker-shell # interactive shell in the image + docker run --rm --gpus all -v $PWD/data:/app/genomics/data -v $PWD/outputs:/app/genomics/outputs progenome-genomics -c "python train_gnn.py --target ancestry --init node2vec" + ``` + `docker compose up -d neo4j` adds the graph browser next to it. Not in the image: the report toolchain (Node.js, tectonic) + and the Brev CLI. On Apple silicon the image builds under emulation and runs CPU-only: use `make setup` there instead. +* **Brev**: `brew install brevdev/homebrew-brev/brev` (or the installer at ), + `brev login --api-key `, then `make brev`. `brev_deploy.sh` creates the + instance if missing (default `g2-standard-4:nvidia-l4:1`, an L4 on GCP; `BREV_INSTANCE=progenome-a100 BREV_TYPE=a100-80gb.1x` + for an A100 where your account has a provider), waits for it, uploads the code, builds the image natively, runs the v1 + chain and the inference benchmark, and copies everything to `outputs_brev//`. Stop billing with + `brev stop ` (delete with `brev delete`). +* **Inference benchmark**: `.venv/bin/python infer.py --run ancestry_svd --compile inductor` (eager vs `torch.compile`, + reports the max logit difference); `--compile tensorrt` requires torch-tensorrt (in the image) and did not finish on this + hetero-GNN. + +### 8. The report and the docs + +**Presenting or reviewing? Start with [RESULTS.md](RESULTS.md)**: methods and results on one GitHub page with every figure. +`docs/report/ProGenome_KT.pdf` (+ `.docx`, `.tex`) is the full knowledge-transfer document; `docs/DEEP_DIVE.md` has the +formulas and shapes, `docs/METHODS.md` the short account, `docs/architecture.html` / `.mmd` the diagrams. `make report` +regenerates the figures from `outputs*/` and rebuilds the three formats (needs Node.js with the `docx` package on +`NODE_PATH`, and `tectonic` or another LaTeX engine). + +### 9. Troubleshooting + +| Symptom | Cause and fix | +|---|---| +| `no pyg_lib wheel for this torch build -> Node2Vec falls back to SVD` | Expected on CPU / macOS; `--init node2vec` needs the GPU image. SVD is the default and the better start anyway. | +| torch will not install | Python must be 3.10-3.13: `PYTHON=python3.12 bash setup.sh`; CPU-only box: `TORCH_INDEX=cpu`. | +| `MISMATCH` in the md5 check | A partial download: `rm -rf data && make data`. | +| `NVIDIA_API_KEY is not set` | `cp .env.example .env` and fill it in (or `~/.progenome.env`, or export it); `make config` shows what is loaded. | +| Decoder: `model did not return JSON` | Rare; the raw reply is saved as `_raw.txt`; rerun, or use `--reasoning none` (the default). | +| `brev create` fails with `cloudCredId or workspaceGroupId must be specified` | That instance type has no provider on your account; use the default L4 or a type listed by `brev ls --types`. | +| `brev exec` hangs | Not used: the deploy script talks to the box over `ssh ` (alias written by `brev refresh`). | +| Docker `externally-managed-environment` | Already handled (`PIP_BREAK_SYSTEM_PACKAGES=1` in the Dockerfile). | +| Docker on Apple silicon is very slow / no GPU | Expected (amd64 emulation); use the venv path locally and the image on Linux/Brev. | +| NVFlare run shows `status: None` | Normal for the simulator; look at `evaluation.json` and the workspace logs instead. | +| Different numbers than documented | Counts and statistics must be identical; model scores vary by up to ~0.01 AUC between runs (GPU kernels, FedAvg). | + +## How the pieces connect + +The HaploGraph file `nodes.csv.gz` has **one row per haploblock cluster and one 0/1 column per 1000G individual**. +`phenotypes_real.csv` and the proteomics matrices are keyed by the **same sample IDs** (`HG00096`, …). That single +string is the whole join: + +``` +Individual ─CARRIES→ Cluster ─IN_BLOCK→ Block ─OVERLAPS→ Gene ─ENCODES→ Protein +Individual ─MEASURED{log2, z}→ Protein Cluster ─CO_OCCURS{weight, lift}→ Cluster +``` + +Phenotypes are *properties* of the Individual node — training targets, never neighbours — so the GNN cannot read +the label off the graph. Individuals being their own node type is what makes the graph federatable: the +cluster/block/gene/protein graph is public and identical at every site; a site holds only its Individual nodes and +their CARRIES / MEASURED edges. + +| node / edge (chr22) | from | count | +|---|---|---| +| `Individual` {ancestry, population, sex, site, phenotype} | `phenotypes_real.csv` (2,503 labelled) + proteomics metadata | 2,548 | +| `Cluster` {support, block stats} | `nodes.csv.gz`, symmetric support filter ≥ 25 | 6,551 | +| `Block` {length, n_clusters, entropy, dominance} | `block_stats.tsv` | 669 | +| `Gene` · `Protein` | `../proteomics/uniprot_chr22.bed` (isoforms collapsed) | 458 · 460 | +| `CARRIES` | `nodes.csv.gz` | 2,365,574 | +| `CO_OCCURS` (weight, lift ≥ 5) | `edges_lift_above_threshold.csv.gz` | 187,030 | +| `IN_BLOCK` · `NEXT_BLOCK` · `OVERLAPS` · `ENCODES` | ids / coordinates | 6,551 · 668 · 1,063 · 460 | +| `MEASURED` (harmonised: robust z per site × protein; LOD-missing = no edge) | proteomics matrices | 1,065,712 | + +**Controls** are part of the design: `sex` (chr22 is autosomal → must be chance) and `site` (sites are +mixed-ancestry batches → must be chance). Every model is scored on the same seeded 70/15/15 split over individuals. + +## Results (chr22, seed 42, held-out test set) + +Graph ↔ phenotype (`cooccurrence_analysis.py`): 98.8 % of clusters are ancestry-associated (FDR 5 %), **0 %** +sex-associated; co-occurring clusters share ancestry profiles (cosine 0.86 vs 0.22 for shuffled pairs). + +| target | classes | logistic regression | GNN | +|---|---|---|---| +| ancestry (real) | 5 | 0.977 bal-acc | 0.974 (SVD init) | +| population (real) | 26 | 0.614 | 0.611 (`--init raw`) | +| sex — control | 2 | 0.463 | 0.503 | + +Genome + proteome (`train_gnn_v2.py`, synthetic phenotype with saved ground truth; AUC on the same test people): + +| modality | AUC | balanced accuracy | +|---|---|---| +| genome only | 0.60 | 0.57 | +| proteome only (MLP) | 0.96 | 0.96 | +| **genome + proteome (graph)** | **0.99** | **0.97** | +| site — batch control | — | 0.30 (chance 0.33) | + +Genome → proteome: the cis effects are in the data (per-cluster/protein r² up to 0.54, `eda.py` §8) and a per-protein +ridge recovers the strong ones (4/20 cis proteins with test R² > 0.1, 0/440 others; `proteome_linear_baseline.py`), +but the GNN's 64-d embedding does not — the GNN is the integration/embedding tool, cis discovery wants sparse +per-protein models or a pQTL edge prior (UKB-PPP on AWS Open Data is the planned source). + +Embeddings: the ancestry-GNN's 64-d space has silhouette 0.70 by ancestry vs 0.06 for plain SVD (`embeddings.py`). +SVD initialisation is essential (free learned embeddings: 0.585). Node2Vec (`--init node2vec`, pyg-lib, 50 pretraining +epochs, A100) is a working but weaker start: ancestry 0.950, population 0.292, sex 0.489; the same 0.70 silhouette. + +Compute (A100 80 GB via Brev): GNN epoch 0.10 s (0.95 s on an M2 CPU); inference over the full graph 36 ms eager → +**4.9 ms with `torch.compile`** (7.5×, identical logits). Torch-TensorRT compilation of this scatter-heavy hetero-GNN +did not finish in 3 h and is not reported. + +Decoder (`graphrag_decoder.py`, Nemotron 3 Super, `reasoning_effort=none`): ~13 s per person, 25 cited ids, 0 invented. + +Federated (`federated/job.py`, NVFlare 2.9 FedAvg, 3 mixed-ancestry sites, each training only on its own people; +only weights exchanged; `make federated`): after 30 rounds × 5 local epochs the global model scores **AUC 0.998 / +balanced accuracy 0.963** on the same 376 held-out people as the central model (0.992 / 0.969); per site 1.000 / +0.999 / 0.997. Ten rounds were not enough (0.955); thirty converge. ~4 min on the M2 CPU. +Site-alone comparison (`federated/local_only.py`, same 150 steps per site): a lone site reaches own-test AUC 0.998 / 0.986 / 0.988 and transfers to the other sites at worst 0.949 / 0.969 / 0.997; the federated model scores 0.999 / 1.000 / 0.968 on the same people. With only 50 steps per site (10 rounds) lone sites reach 0.88-0.92 while the federated model reaches +0.99. Federation costs nothing when a site has enough data and budget, helps when it does not, and in both cases solves the +constraint (one model trained on everyone, no row leaves a site). A larger gain is expected with ancestry-pure sites, the next experiment. + +Reproduction (18 Sept): a simulated fresh clone on the laptop (`make setup && make run && make run-v2 && make federated`) +and the Docker image rebuilt on the A100 both reproduce every count exactly and every score within run-to-run noise: +phenotype AUC both 0.994 (laptop) / 0.995 (A100), federated 0.996 / 0.987 vs central 0.994 / 0.995, +inference 36.3 -> 4.9 ms. Details: `docs/report/ProGenome_KT.pdf`, section 13. + +## Pipeline + +| step | script | writes | +|---|---|---| +| 1 | `fetch_data.sh` | `data/` — HaploGraph files, phenotypes, block stats (md5-verified) | +| 2 | `build_kg.py` (`haplokg.py`) | `outputs/kg//` tables, sparse `carries.npz`, PyG `hetero.pt` | +| 3 | `cooccurrence_analysis.py` | `outputs/cooccurrence//` | +| 4 | `baseline.py` | `outputs/baseline//` + shared split `outputs/splits/` | +| 5 | `graph_explore.py` (NetworkX; cuGraph via `nx-cugraph` in the image) | `outputs/graph//` stats, GraphML, plots | +| 6 | `train_gnn.py` (`--init svd|node2vec|learned|raw`) | `outputs/gnn//_/` | +| 7 | `embeddings.py` | `outputs/embeddings//` | +| v2.1 | `proteomics_synth_1000g.py` | `outputs/proteomics_synth//` (+ `ground_truth.json`) | +| v2.2 | `build_kg_v2.py` (`haplokg_proteins.py`) | genes/proteins/measured tables, `hetero_v2.pt` | +| v2.3 | `eda.py` | `outputs/eda//EDA.md` + tables + plots | +| v2.4 | `train_gnn_v2.py` (`--modality genome|proteome|both`, `--target phenotype|site|ancestry|sex|proteome`) | `outputs/gnn_v2//` | +| v2.5 | `proteome_linear_baseline.py` | `outputs/gnn_v2//proteome_ridge_baseline/` | +| v2.6 | `graphrag_decoder.py` | `outputs/graphrag//_insight.json` | +| v2.7 | `federated/job.py` → `federated/evaluate_global.py` | `outputs/federated//` NVFlare workspace, `evaluation.json` | +| v2.8 | `federated/local_only.py` | `outputs/federated//local_only_vs_federated.json` (each site alone vs the global model) | +| — | `infer.py` (`--compile none|inductor|tensorrt`) | inference benchmark | +| — | `neo4j_load.py` + `docker-compose.yml` | Neo4j browser | +| — | `Dockerfile`, `run_all.sh`, `run_v2.sh`, `brev_deploy.sh`, `Makefile`, `setup.sh` | packaging / deploy | + +Tests: `make test` (toy graph: parsing, filtering, edge remapping, protein layer, HeteroData). + +## GNN + +`HeteroConv`, 2 layers, hidden 64, LayerNorm + residual, dropout 0.3: `SAGEConv` on carries / in_block / next_block / +overlaps / encodes (both directions), `GraphConv` with edge weights on co_occurs (normalised log lift) and measured +(harmonised z). Individual input = SVD-32 of the carrier matrix (or the raw carrier row) + the harmonised protein +vector and its observed mask; clusters and blocks add their z-scored statistics. Class-weighted cross-entropy, +Adam 5e-3, early stopping on validation balanced accuracy. With `--init raw`, a gradient saliency ranks clusters and +is scored against the ground-truth causal clusters (precision@20 = 0.20 vs 0.06 by chance). + +## Next + +* **Federated, next level:** weight FedAvg by site size (`aggregation_weights`), per-site held-out reporting only + (no central test set), NVFlare POC mode across real machines (the L4 and A100 as two sites). +* **Real proteomics:** Wu et al. 2013 LCL proteomics (95 HapMap individuals with 1000G IDs) for a real join; UKB-PPP + cis-pQTLs (AWS Open Data) as `Cluster → Protein` propensity edges. +* More chromosomes: `CHROM=chr21 make run`. diff --git a/genomics/RESULTS.md b/genomics/RESULTS.md new file mode 100644 index 0000000..d150271 --- /dev/null +++ b/genomics/RESULTS.md @@ -0,0 +1,573 @@ +# ProGenome genomics: Methods and Results (chromosome 22) + +*Person-level knowledge graph of haploblock clusters, genes and proteins; a heterogeneous GNN encoder; an LLM decoder; +federated training with NVIDIA FLARE.* This page renders on GitHub and is meant to be presented from directly. Every +number comes from files under `genomics/outputs*/` (chr22, seed 42); the diagrams are Mermaid (rendered by GitHub), the +plots are in `docs/report/figures/` and are regenerated by `make report`. The companion documents are the [README](README.md) (implementation guide), the +[knowledge-transfer report](docs/report/ProGenome_KT.pdf) (38 pages) and the [deep dive](docs/DEEP_DIVE.md) (formulas). + +**In one sentence:** we joined the published 1000 Genomes HaploGraph to real labels and to per-site proteomics through +the sample id, trained a graph neural network on the resulting person-level graph, showed that genome + proteome through +the graph predicts a phenotype better than either alone (AUC 0.60 / 0.96 / 0.99) while negative controls stay at chance, +turned the model's outputs into cited per-person reports with an NVIDIA NIM language model, and trained the same model +federated across three sites with no loss (AUC 0.987-0.998 vs central 0.992-0.995). + +--- + +## 1. Goal + +The team README asks three questions. This is how they were operationalised and what was achieved: + +| RQ | Question | Operational form | Status | +|---|---|---|---| +| 1 | How can haploblock genomics connect to genes and proteomics in a graph model? | one typed graph: Individual, Cluster, Block, Gene, Protein; seven edge types; joined by the 1000G sample id | done (schema v2, browsable in Neo4j) | +| 2 | Can a GNN combine genomic and proteomic information to identify disease-related phenotype clusters? | ablation genome / proteome / both on a synthetic phenotype with saved ground truth; controls; saliency against the planted causal clusters | done: 0.60 / 0.96 / 0.99 AUC; controls at chance; top-3 salient clusters all causal | +| 3 | Can it be trained across institutions without transferring individual-level data? | NVFlare FedAvg, three sites, only weights exchanged, scored on the same held-out people as the central model | done in simulation: 0.987-0.998 vs 0.992-0.995 | + +![System architecture](docs/report/figures/architecture_slide.png) + +*Figure 0. System architecture: private hospital sites on the left, the shared reference graph and the GNN encoder trained through the NVFlare server in the middle, the model's outputs and the LLM decoder on the right.* + +**v1: genome graph to phenotypes** (`make run`) + +```mermaid +flowchart TB + A["data.haploblocks.org
HaploGraph chr22 (pre-built)"] --> B["fetch_data.sh
md5-verified download"] --> C["build_kg.py
sparse carrier matrix, PyG HeteroData"] --> D["cooccurrence_analysis.py
cluster / edge vs phenotype"] --> E["baseline.py
logistic regression, shared split"] --> F["train_gnn.py
hetero-GNN encoder"] --> G["embeddings.py
silhouette, kNN, PCA plots"] +``` + +**v2: proteomics, integration, decoder, federated** (`make run-v2`, `make federated`) + +```mermaid +flowchart TB + H["proteomics_synth_1000g.py
3 sites, 1000G ids, ground truth"] --> I["build_kg_v2.py
hetero.pt from v1 + genes, proteins, harmonised MEASURED edges"] --> J["eda.py
EDA report"] --> K["train_gnn_v2.py
genome / proteome / both, controls, saliency"] --> L["proteome_linear_baseline.py
ridge cis test"] --> M["graphrag_decoder.py
NIM LLM, cited insight"] --> N["federated/job.py
NVFlare FedAvg, 3 sites + site-alone"] +``` + +*Figure 1. The two pipeline chains. Each box is one script with a Makefile target; outputs land under `outputs//chr22/`.* + +--- + +## 2. Data + +| Data | Source | Content | Role | +|---|---|---|---| +| HaploGraph nodes | data.haploblocks.org/haplograph/1000G/chr22 | 248,254 haplotype clusters x 2,548 people, 0/1 carrier matrix | person-to-cluster edges (CARRIES) | +| HaploGraph edges | same | 187,030 cluster pairs with weight and lift >= 5 | cluster-to-cluster edges (CO_OCCURS) | +| Block statistics | same | 669 recombination-defined blocks, 17.1-50.2 Mb, length, entropy, dominance | Block nodes | +| Phenotypes | 1000G panel, republished with the graph | ancestry (5), population (26), sex for 2,503 people | labels on the person node; never edges | +| Gene / protein BED | UCSC UniProt track (team) | 917 isoform rows -> 458 genes, 460 proteins | Gene, Protein nodes; OVERLAPS, ENCODES | +| Proteomics (synthetic, joinable) | `proteomics_synth_1000g.py` | 2,503 real 1000G ids, 3 mixed-ancestry sites, 460 proteins, phenotype driven by 20 causal clusters, cis effects, batch shift, missingness; `ground_truth.json` | MEASURED edges, person features, site and phenotype labels, ground truth | +| Proteomics (team, SITE ids) | `proteomics/` on main | 4,000 patients, not 1000G ids | validated the proteomics side; cannot join the genome | + +![Genome to graph](docs/report/figures/genome_to_graph.png) + +*Figure 2. Upstream (haploblocks.org): blocks from recombination peaks, phased haplotypes per person and block, MMseqs2 +clusters, the carrier matrix; our symmetric support filter keeps 6,551 clusters and reproduces the edge file's node set +exactly.* + +```mermaid +flowchart LR + subgraph S ["Data sources"] + S1["HaploGraph nodes.csv.gz
who carries which cluster"] + S2["HaploGraph edges + block_stats
co-occurrence, block statistics"] + S3["phenotypes_real.csv
ancestry, population, sex"] + S4["uniprot_chr22.bed
genes, proteins, coordinates"] + S5["proteomics matrices, 3 sites
log2 intensity per protein"] + S6["proteomics metadata
site, age, sex, case/control"] + end + subgraph K ["Knowledge graph (hetero_v2.pt)"] + K1["Graph structure
Individual, Cluster, Block, Gene, Protein
CARRIES, CO_OCCURS, IN_BLOCK, NEXT_BLOCK,
OVERLAPS, ENCODES, MEASURED (harmonised z)"] + K2["Node features
cluster and block statistics;
person = SVD-32 (label-free) or raw row,
plus protein z and observed mask"] + K3["Labels on the person node
ancestry, population, sex, site, phenotype
targets and ground truth only:
never an edge, never a feature"] + end + E["GNN encoder
HeteroConv x 2, hidden 64
SAGEConv + edge-weighted GraphConv
trained on TRAIN labels, selected on VAL,
reported on TEST (376 people)"] + subgraph O ["What the trained model gives"] + O1["class probabilities per person"] + O2["64-d embedding per person and cluster"] + O3["saliency per haploblock cluster"] + O4["model weights (state dict)"] + O5["predictions for all 2,548 people (infer.py)"] + end + subgraph U ["Consumers"] + U1["evaluation against ground truth
(test labels, ground_truth.json)"] + U2["LLM decoder (NIM GraphRAG)
cited insight per person"] + U3["NVFlare FedAvg
only weights leave a site"] + U4["plots, Neo4j, CSVs"] + end + S1 --> K1 + S2 --> K1 + S4 --> K1 + S1 --> K2 + S5 --> K2 + S3 --> K3 + S6 --> K3 + K1 -- "messages" --> E + K2 -- "inputs" --> E + K3 -- "loss" --> E + E --> O1 + E --> O2 + E --> O3 + E --> O4 + E --> O5 + O1 --> U1 + O3 --> U1 + O2 --> U2 + O3 --> U2 + O4 --> U3 + O5 --> U4 +``` + +*Figure 3. Where each data source enters and what the trained model produces. Phenotypes are targets and ground truth +only; the encoder can reach a label only through the loss on training people.* + +--- + +## 3. Methods + +### 3.1 Knowledge graph + +The HaploGraph has one node type (cluster) with people as a feature vector. We made the person a node type of its own, +which is what lets labels and protein measurements attach to a person and what makes the graph federatable: the +cluster / block / gene / protein layer is public and identical everywhere, a site holds only its people and their edges. + +| Node | Count | Edge | Count | +|---|---|---|---| +| Individual {ancestry, population, sex, site, phenotype} | 2,548 | Individual CARRIES Cluster | 2,365,574 | +| Cluster {support, block statistics} | 6,551 | Cluster CO_OCCURS {weight, lift} Cluster | 187,030 | +| Block {length, n_clusters, entropy, dominance} | 669 | Cluster IN_BLOCK Block; Block NEXT_BLOCK Block | 6,551; 668 | +| Gene | 458 | Block OVERLAPS Gene; Gene ENCODES Protein | 1,063; 460 | +| Protein | 460 | Individual MEASURED {log2, z} Protein (missing = no edge) | 1,065,712 | + +Cluster filter: keep a cluster if 25 <= carriers <= N - 25 (248,254 -> 6,551; 0 of 187,030 edges lose an endpoint). +Protein levels are harmonised per site and protein (robust z = (x - median) / (1.4826 MAD)) before entering the graph; +the between-site shift goes from 0.48 to 0.00 log2 while 227 proteins keep a phenotype association at FDR 5 %. + +```mermaid +flowchart LR + I["Individual (2,548)
ancestry, population, sex,
site, phenotype: labels, never neighbours"] + C["Cluster (6,551)
support, block statistics, SVD-32"] + B["Block (669)
length, n_clusters, entropy, dominance"] + G["Gene (458)
coordinates"] + P["Protein (460)
coordinates, isoforms"] + I -- "CARRIES 2,365,574" --> C + C -- "CO_OCCURS (weight, lift) 187,030" --> C + C -- "IN_BLOCK 6,551" --> B + B -- "NEXT_BLOCK 668" --> B + B -- "OVERLAPS 1,063" --> G + G -- "ENCODES 460" --> P + I -- "MEASURED (log2, z) 1,065,712; missing = no edge" --> P +``` + +*Figure 4. The schema with chr22 counts. Labels live on the Individual node, never as neighbours (a label neighbour would +let a two-layer GNN read the answer and report a meaningless 100 %).* + +### 3.2 Statistics before modelling + +For every cluster and label a 2 x k contingency test (Cramer's V = sqrt(chi2 / N)) with Benjamini-Hochberg FDR, all +6,551 tests in one sparse matrix product; for every edge the cosine similarity of the endpoints' ancestry-enrichment +profiles against a degree-preserving shuffled null. + +### 3.3 Embeddings and the GNN encoder + +Starting vectors: cluster and block statistics; for people and clusters a label-free 32-component SVD of the carrier +matrix (rows of U S and V S share one space). Alternatives compared: Node2Vec (pyg-lib, 50 epochs), free learned +embeddings, the raw carrier row. + +Encoder: PyTorch Geometric `HeteroConv`, two layers, hidden 64; `SAGEConv` (neighbour mean) on carries, in_block, +next_block, overlaps, encodes in both directions; `GraphConv` with edge weights on co_occurs (normalised log lift) and +measured (harmonised z, signed); LayerNorm + residual + ReLU + dropout 0.3 per layer; linear head on the person vector. +Class-weighted cross-entropy on training people only, Adam 5e-3, weight decay 5e-4, early stopping on validation +balanced accuracy (patience 30). Full batch: one epoch is one pass over the graph (0.95 s on an M2 CPU, 0.10 s on an A100). + +### 3.4 Evaluation protocol and controls + +One stratified 70/15/15 split over the 2,503 labelled people (seed 42), written once and reused by every model, so all +held-out numbers are on the same 376 people. Two negative controls are part of the design: **sex** (chr22 is autosomal, +so any signal is noise) and **site** (sites are mixed-ancestry random partitions, so any signal is batch). Metrics: +balanced accuracy (mean per-class recall), ROC-AUC for binary targets, macro-F1, silhouette and 5-NN accuracy for +embeddings, precision at 20 for saliency, R2 per protein for regression. + +### 3.5 Synthetic proteome with ground truth + +Per person (real 1000G id, real sex, random age, site by shuffling within ancestry): case/control logit = weighted sum over +20 causal clusters (beta ~ N(0, 1.5)) + small age term, calibrated to 38 % cases; each causal cluster shifts one protein +in its own block (cis, beta ~ N(0, 1)); 15 % of proteins respond to the phenotype (beta ~ N(0, 0.5)); age, sex, site +shift (sd 0.3) and biological + technical noise; missingness rising toward the detection limit (MNAR). The first version +shifted every protein with the phenotype and every model scored 1.0; the signal was made sparse on purpose. + +### 3.6 Decoder (GraphRAG with NVIDIA NIM) + +For one person, deterministic retrieval from the graph (profile with the true phenotype withheld, the GNN prediction, +salient clusters carried, the 8 most ancestry-informative clusters with block / genes / proteins, the 8 most extreme +protein levels with their encoding block, the 5 nearest people in the embedding) -> JSON (~4.5 k tokens) -> +`nvidia/nemotron-3-super-120b-a12b` on the OpenAI-compatible NIM endpoint (`reasoning_effort=none`, temperature 0.2) +under a system prompt that forbids new entities and requires every id verbatim -> cited ids validated against the context. + +### 3.7 Federated learning (NVIDIA FLARE 2.9) + +`FedAvgRecipe` + `SimEnv`, three sites (835 / 835 / 833 people, mixed ancestry). Each client builds its site subgraph +(`HeteroData.subgraph` over its own people; public node types kept whole), receives the global weights, evaluates them on +its own validation and test people, trains five full-batch epochs with local class weights and sends back weights, +metrics and the step count; the server averages (step-weighted) and keeps the best model by validation balanced +accuracy. Per round per site only one state dict (~0.5 M floats) and five scalars cross the boundary. The global model +is then scored centrally on the same 376 held-out people as the central model, and `local_only.py` trains each site alone +for the same 150 steps as the reference. + +```mermaid +flowchart TB + SHARED["Shared and public at every site
Cluster, Block, Gene, Protein graph (CO_OCCURS, IN_BLOCK, NEXT_BLOCK, OVERLAPS, ENCODES)
and the encoder weights"] + subgraph SITES ["Three hospital sites: Individual nodes, labels, CARRIES and MEASURED edges never leave"] + direction LR + A["Site 1
835 people
774,753 CARRIES, 355,294 MEASURED"] + B["Site 2
835 people"] + C["Site 3
833 people"] + end + SRV["NVFlare server, FedAvg
averages the three state dicts, keeps the best by validation balanced accuracy,
returns the global model; 30 rounds x 5 local epochs"] + SHARED -.-> SITES + A <-- "weights up, global model back" --> SRV + B <-- "weights up, global model back" --> SRV + C <-- "weights up, global model back" --> SRV +``` + +*Figure 5. Three sites with private people and edges, one shared public graph, weights only to the server.* + +### 3.8 Compute and packaging + +Python 3.13, PyTorch 2.14, PyG 2.8, pyg-lib, scikit-learn, scipy sparse, NetworkX (+ nx-cugraph), Neo4j 5.26; Docker +image `pytorch/pytorch:2.14.0-cuda12.6-cudnn9-runtime` + PyG + nx-cugraph + torch-tensorrt + NVFlare, built from the +repository root with the proteomics inputs baked in and the unit tests run at build time; NVIDIA Brev L4 and A100 80 GB; +`config.py` + `.env.example` as the single credential entry point; Makefile targets for every stage. + +--- + +## 4. Results + +### 4.1 The graph carries phenotype information, and the control is clean + +| Test | Ancestry | Population | Sex (control) | +|---|---|---|---| +| Clusters associated at FDR 5 % | 6,470 / 6,551 (98.8 %) | 6,378 (97.4 %) | **0** | +| Median / max Cramer's V | 0.20 / 0.81 | 0.24 / - | 0.013 / 0.03 | + +Edges: cosine similarity of endpoint ancestry profiles 0.86 (real) vs 0.22 (shuffled); 86 % of edges join clusters +enriched in the same ancestry (42 % expected); similarity rises with lift (0.84 -> 0.90 by quartile). The co-occurrence +graph is largely population structure: long-range co-inheritance within ancestries, not physical linkage. + +![Cramer's V](docs/report/figures/cramers_v.png) + +*Figure 6. Per-cluster association with each label. Ancestry and population carry signal; sex does not.* + +![Edge similarity](docs/report/figures/edge_similarity.png) + +*Figure 7. Co-occurring clusters share ancestry-enrichment profiles; shuffled pairs do not.* + +![Informative clusters](docs/report/figures/informative_clusters_heatmap.png) + +*Figure 8. The 30 most ancestry-informative clusters and the fraction of each ancestry that carries them.* + +### 4.2 Real labels: the GNN ties the linear baseline on accuracy and the control stays at chance + +| Target (real labels) | Classes | Logistic regression | GNN (SVD input) | +|---|---|---|---| +| ancestry | 5 | 0.977 | 0.974 | +| population | 26 | 0.614 | 0.437 (0.611 with the raw carrier row) | +| sex, negative control | 2 | 0.463 | 0.503 | + +Balanced accuracy on the 376 held-out people. On chr22 alone ancestry is almost a linear function of which clusters a +person carries, so the GNN cannot beat logistic regression on accuracy here; its value is the embedding space (4.3) and +the integration (4.4). + +![Baseline vs GNN](docs/report/figures/baseline_vs_gnn.png) + +*Figure 9. Real labels: logistic regression versus the GNN; sex at chance for both.* + +![Confusion matrices](docs/report/figures/confusion_matrices.png) + +*Figure 10. Ground truth versus prediction on the held-out people. Ancestry: 368 / 376 correct, the eight errors all +between AMR and EUR (admixture). Synthetic phenotype: genome + proteome through the graph 367 / 376 with one false case; +genome alone 202 / 376. Sex: a coin flip, as it must be.* + +### 4.3 Embeddings: what the GNN adds + +| Space | Silhouette by ancestry | 5-NN balanced accuracy | +|---|---|---| +| SVD-32 (label-free start) | 0.06 | 0.90 | +| GNN hidden layer, ancestry run | **0.70** | **0.97** | +| GNN with free learned embeddings | 0.33 | 0.61 | + +![Embedding of people](docs/report/figures/embedding_individuals.png) + +*Figure 11. The GNN's 64-d embedding of people in two dimensions: five ancestries separate, AMR spread between EUR and +AFR as admixture predicts; coloured by sex the same points are fully mixed.* + +![Embedding quality](docs/report/figures/embedding_quality.png) + +*Figure 12. Silhouette and nearest-neighbour accuracy for SVD and for each GNN's hidden layer.* + +![Starting embeddings](docs/report/figures/init_comparison.png) + +*Figure 13. The same GNN with different starting embeddings. SVD is the best start (ancestry 0.974); Node2Vec with 50 +pretraining epochs is a working but weaker alternative (0.950); free learned embeddings fail (0.585); the raw carrier +row is best for the 26-class population target (0.611).* + +### 4.4 Integration: genome + proteome through the graph beats either alone + +Synthetic phenotype with saved ground truth; same 376 held-out people; AUC. + +| Model | Input to the person node | Relations | AUC | Balanced accuracy | +|---|---|---|---|---| +| genome only (graph) | SVD-32 or raw carrier row | genome relations | 0.60-0.65 | 0.57-0.62 | +| proteome only (MLP, no graph) | harmonised z + observed mask | none | 0.96 | 0.96 | +| **genome + proteome (graph)** | both | all twelve | **0.99** | **0.95-0.97** | +| site, batch control (full graph) | both | all | - | 0.30-0.34 (chance 0.33) | +| ancestry (full graph) | both | all | - | 0.89-0.90 | + +Saliency (gradient of the case score w.r.t. the carrier row, averaged over test cases): the top three clusters are all +among the 20 planted causal clusters out of 6,551, four of the top twenty are (0.06 expected by chance). + +Genome -> proteome, kept as a negative result: regressing the 460 protein z-scores from the person embedding gives R2 ~ 0 +even for cis-affected proteins, while a per-protein ridge on the carrier row recovers 4 of 20 cis proteins at test +R2 > 0.1 and 0 of 440 others; the data's own ceiling (r2 between carrying the causal cluster and its protein) is 0.54 +max, 0.17 mean. The GNN is the integration and embedding tool; cis discovery wants sparse per-protein models or a +pQTL edge prior. + +![Results by modality](docs/report/figures/results_modalities.png) + +*Figure 14. Held-out AUC and balanced accuracy by modality; the site control sits at chance.* + +![Training curves](docs/report/figures/training_curves.png) + +*Figure 15. Training loss and validation balanced accuracy per epoch; early stopping picks the best validation epoch.* + +![Saliency](docs/report/figures/saliency_top20.png) + +*Figure 16. Cluster saliency of the combined model: ground-truth causal clusters (teal) among the top 20.* + +![Ridge](docs/report/figures/ridge_r2.png) + +*Figure 17. Per-protein ridge from the genome: only the strong cis effects are recoverable.* + +![Ground truth](docs/report/figures/ground_truth_effects.png) + +*Figure 18. What was planted: causal-cluster effects on the phenotype, cis effects on proteins, phenotype effects on the +70 responsive proteins, carrier frequencies of the causal clusters.* + +### 4.5 Decoder: from numbers to a cited report + +The encoder's output for a person is numbers: a class score, 64 coordinates, a ranking of clusters. The decoder turns +them into something a research team can read, without letting the language model invent anything. Worked example, +HG00103 (EUR, GBR, male, site 3, age 29, 994 clusters carried; true label control). + +**What the GNN alone gives for HG00103** (all retrieved deterministically from the graph and the trained model): + +| Retrieved fact | Value | +|---|---| +| GNN prediction | control (truth: control) | +| Globally salient clusters this person carries | chr22_46902935-46974137_cluster186 (saliency 0.110), chr22_46974137-47051453_cluster43 (0.055), chr22_40032702-40132216_cluster25 (0.046) | +| Most ancestry-informative clusters carried, with the genes and proteins in their block | chr22_40032702-40132216_cluster151 (EAS-enriched; TNRC6B, Q9UPQ9); chr22_35738379-36060254_cluster545 (EAS; RBFOX2, O43251); chr22_26024448-26060666_cluster35 (EUR; MYO18B, Q8IUG5) | +| Most extreme protein levels (harmonised z) | P07305 / H1-0 +2.75; Q9UFD9 / RIMBP3 -2.65; Q15735 / INPP5J +2.53; P53673 / CRYBA4 +2.52; Q04917 / YWHAH +2.37 | +| Nearest people in the 64-d embedding | HG01767 (EUR, IBS), NA12842 (EUR, CEU), HG00351 (EUR, FIN), HG04002 (SAS, ITU), HG00336 (EUR, FIN) | + +**What the decoder returns** (Nemotron 3 Super on NVIDIA NIM, 13.3 s, 4,555 prompt tokens, 1,020 completion tokens): + +| Field | Text returned | +|---|---| +| summary | Individual HG00103 is of EUR ancestry (GBR population) and carries 994 haploblock clusters. The GNN phenotype prediction for this individual is 'control', which is a synthetic case/control label used for pipeline testing. Several ancestry-informative clusters are carried, showing high enrichment in various populations including EUR, EAS, and SAS. | +| ancestry_assessment | The individual's ancestry is EUR (European) as per the 1000G panel data, with population label GBR. This reflects population structure, not a medical finding. | +| phenotype_assessment | The GNN-predicted phenotype is 'control', which is a synthetic case/control label used to test the pipeline; it does not reflect a real medical condition. | +| caveats | Ancestry labels reflect population structure and are not medical findings. The phenotype is synthetic. Not all carried clusters have gene or protein annotations. The GNN prediction is based on the embedding and may not correlate with observed protein levels. | +| citation check | 25 ids cited, 0 unknown (every id exists in the retrieved context) | + +Genome-to-proteome links the decoder wrote, each traceable to graph ids: + +| Cluster carried | Block | Gene in the block | Protein | Observation | +|---|---|---|---|---| +| chr22_40032702-40132216_cluster151 | chr22_40032702-40132216 | TNRC6B | Q9UPQ9 | carried by individual, enriched in EAS | +| chr22_35738379-36060254_cluster545 | chr22_35738379-36060254 | RBFOX2 | O43251 | carried by individual, enriched in EAS | +| chr22_26024448-26060666_cluster35 | chr22_26024448-26060666 | MYO18B | Q8IUG5 | carried by individual, enriched in EUR | + +A second person, HG00096 (not in the held-out split, run on a fresh clone): the decoder reported that no prediction +exists for them rather than inventing one, backed the EUR label with carrier fractions it read from the context +("chr22_26024448-26060666_cluster35 at 0.952 ... consistent with the EUR ancestry label"), listed six genome-to-proteome +links including CELSR1 / Q9NYQ6 and PPARA / Q6NVV7, and cited 24 ids with 0 unknown in 10.8 s. + +**What the LLM adds, and what it does not.** It adds the sentence layer: one paragraph a clinician or biologist can read, +the genome-to-protein chain (cluster carried, block, gene, protein) spelled out per person, the right hedges attached +automatically (population structure is not a medical finding, the phenotype is synthetic), graceful handling of missing +information, and traceability, because every claim carries an id that the code verifies against the retrieved facts. +It does not make predictions, does not add outside knowledge, and cannot cite anything the graph did not provide; the +wording and the number of ids it chooses to cite vary between calls, the citation check does not. + +![Person neighbourhood](docs/report/figures/person_neighbourhood.png) + +*Figure 19. What the decoder retrieves for one person: clusters with blocks, genes and proteins, extreme protein levels, +nearest neighbours in the embedding.* + +### 4.6 Federated: no loss, no data movement + +| Model | Balanced accuracy | AUC | Evaluated on | +|---|---|---|---| +| federated, 10 rounds x 5 epochs | 0.89 | 0.955 | the same 376 held-out people | +| federated, 30 rounds x 5 epochs | 0.963-0.971 | **0.987-0.998** (three runs) | the same 376 held-out people | +| central model (train_gnn_v2, both / raw) | 0.953-0.969 | 0.992-0.995 | the same 376 held-out people | +| federated model per site (30 rounds) | 0.98 / 0.98-0.99 / 0.93 | 0.999 / 1.000 / 0.968-0.997 | each site's own held-out people | + +**Each site alone versus the federated model.** `federated/local_only.py` trains a model at each site on that site's +people only, for the same number of optimizer steps the federated clients used, and scores it on the same held-out people +as the federated global model. Two budgets were run on the A100: 150 steps per site (30 rounds x 5 epochs) and 50 steps +(10 rounds x 5 epochs). AUC / balanced accuracy on each site's own held-out people: + +| Site (held-out n) | Alone, 150 steps | Federated, 150 steps | Alone, 50 steps | Federated, 50 steps | +|---|---|---|---|---| +| SITE1 (111) | 0.998 / 0.952 | 0.999 / 0.983 | 0.881 / 0.543 | 0.998 / 0.950 | +| SITE2 (136) | 0.986 / 0.852 | 1.000 / 0.991 | 0.902 / 0.772 | 0.991 / 0.907 | +| SITE3 (129) | 0.988 / 0.923 | 0.968 / 0.933 | 0.915 / 0.830 | 0.989 / 0.865 | +| **Average over sites** | **0.991 / 0.909** | **0.989 / 0.969** | **0.899 / 0.715** | **0.993 / 0.907** | +| Pooled 376 people, federated global model | - | 0.987 / 0.967 | - | 0.991 / 0.901 | +| Pooled 376 people, central model (all data in one place) | 0.995 / 0.960 | | 0.997 / 0.959 | | + +A lone site's model also has to work on other hospitals' patients. Test AUC of each model on each site's held-out +people (150 steps; rows = where the model was trained, columns = whose patients it is scored on): + +| Model trained at | on SITE1 people | on SITE2 people | on SITE3 people | +|---|---|---|---| +| SITE1 alone | 0.998 | 0.982 | 0.949 | +| SITE2 alone | 0.993 | 0.986 | 0.969 | +| SITE3 alone | 0.997 | 0.998 | 0.988 | +| **Federated global model** | **0.999** | **1.000** | **0.968** | + +Reading. With a full budget a lone site is nearly as good as the federated model on AUC (0.991 vs 0.989 on average) +but clearly worse on balanced accuracy (0.909 vs 0.969), because a single site's threshold is calibrated on its own +case mix. With the smaller budget the gap opens: lone sites average 0.899 AUC and 0.715 balanced accuracy, the federated +model 0.993 and 0.907, because the averaged weights have effectively seen every site's people. Federation therefore costs +nothing when a site has enough data and budget, helps clearly when it does not, and in both cases solves the constraint: +one model trained on everyone with no row leaving a site. A larger gain is expected with ancestry-pure sites. + +![FedAvg rounds](docs/report/figures/federated_rounds.png) + +*Figure 20. Global-model AUC and validation balanced accuracy at each site per round, against the central model.* + +![Federated vs central](docs/report/figures/federated_vs_central.png) + +*Figure 21. Central versus federated (10 and 30 rounds) on the same held-out people, and per site.* + +![Site alone vs federated](docs/report/figures/federated_site_alone.png) + +*Figure 22. Each site trained alone versus the federated global model on the same per-site held-out people.* + +### 4.7 Compute + +| Quantity | Laptop (M2 CPU) | A100 80 GB | +|---|---|---| +| GNN epoch (full graph) | 0.95 s | 0.10 s | +| whole chain from an empty folder | ~30 min | ~14 min incl. image build | +| full-graph inference, eager -> torch.compile | - | 36.3 ms -> 4.9 ms (7.4x, logits equal to 2e-6) | +| Torch-TensorRT | - | did not finish compiling this scatter-heavy hetero-GNN in 3 h; not claimed | + +![Compute](docs/report/figures/compute_benchmarks.png) + +*Figure 23. Epoch time laptop vs A100; inference eager vs torch.compile.* + +### 4.8 Reproduction + +Before the commit the chain was re-run from the committable files as a fresh clone on the laptop and through the rebuilt +Docker image on the A100. Every count is identical; every score is within run-to-run noise. + +| Quantity | Documented | Fresh clone, laptop | Rebuilt image, A100 | +|---|---|---|---| +| clusters / CARRIES / CO_OCCURS / MEASURED | 6,551 / 2,365,574 / 187,030 / 1,065,712 | identical | identical | +| ancestry-associated clusters / sex | 6,470 / 0 | 6,470 / 0 | 6,470 / 0 | +| baseline ancestry / population / sex | 0.977 / 0.614 / 0.463 | 0.977 / 0.614 / 0.463 | 0.981 / 0.614 / 0.463 | +| GNN ancestry / population / sex (SVD) | 0.974 / 0.437 / 0.503 | 0.974 / 0.437 / 0.503 | Node2Vec: 0.950 / 0.292 / 0.489 | +| phenotype AUC genome / proteome / both | 0.64 / 0.96 / 0.99 | 0.65 / 0.96 / 0.99 | 0.64 / 0.96 / 0.99 | +| site control (chance 0.33) | 0.30 | 0.34 | 0.30 | +| saliency hits in top 20; ridge cis / other | 4; 4 / 0 | 4; 4 / 0 | 4; 4 / 0 | +| federated AUC / central AUC | 0.998 / 0.992 | 0.996 / 0.994 | 0.987 / 0.995 | +| decoder cited / invented | 25 / 0 | 4 / 0 | dry run | +| unit tests | 11 pass | 11 pass | 11 pass (in the build) | + +--- + +## 5. Exploratory data analysis (selected) + +![Populations](docs/report/figures/eda_populations.png) + +*Figure 24. 1000G individuals per population, coloured by continental ancestry (2,503 labelled: AFR 660, EAS 504, +EUR 503, SAS 489, AMR 347).* + +![Blocks](docs/report/figures/eda_blocks.png) + +*Figure 25. Haploblocks: length (median 29.7 kb), clusters per block versus length, entropy along the chromosome.* + +![Clusters](docs/report/figures/eda_clusters.png) + +*Figure 26. Kept clusters: carriers per cluster and clusters carried per person (~928) by ancestry.* + +![Co-occurrence](docs/report/figures/eda_cooccurrence.png) + +*Figure 27. Co-occurrence edges: lift (median 5.5, max 76), distance between endpoints (tens of Mb), degree (max 358; the +top hubs are AFR-enriched rare clusters, the mega-hub artefact).* + +![Graph region](docs/report/figures/graph_region.png) + +*Figure 28. The densest published island of chr22: clusters coloured by the ancestry they are enriched in.* + +![Edge positions](docs/report/figures/graph_edge_positions.png) + +*Figure 29. All 187,030 edges by the positions of their endpoints: block structure and long-range population structure.* + +![Proteomics](docs/report/figures/eda_proteomics.png) + +*Figure 30. Synthetic proteomics: intensity distribution, missingness rising for low-abundance proteins (7.4 % overall, +up to 16 %), batch effect between sites before and after the harmoniser.* + +![Sites](docs/report/figures/sites_composition.png) + +*Figure 31. The three federated sites: people per site by ancestry (mixed by design), case prevalence, missing values.* + +![Batch effect](docs/report/figures/sites_batch_effect.png) + +*Figure 32. Per-site protein medians before and after harmonisation: the batch offsets vanish.* + +--- + +## 6. What is claimed and what is not + +* The case/control phenotype is synthetic: the integration numbers show the pipeline recovers a planted signal, not + biology. Ancestry, population and sex results are on real labels. +* Chromosome 22 only; nothing in the code is chromosome-specific (`CHROM=chr21 make run`). +* On single-label accuracy the GNN ties, not beats, logistic regression; its value is the embedding space and the + integration. +* The GNN embedding does not recover single-cluster cis effects; a per-protein ridge does for the strong ones. +* Federated evaluation reuses the central held-out people; a real deployment reports per site only. Federation is run in + the NVFlare simulator, not across real machines. +* TensorRT is not part of the inference claim; torch.compile is. + +## 7. How this fits with the team's proteomics work + +The team's proteomics work on `main` (synthetic proteomics generator, protein filtering and classification baselines, +the protein-centred knowledge graph, the federated comparison and `methods_and_results.md`) covers the proteome side and +the federated logistic setting. `genomics/` adds the genome side: the person-level knowledge graph built on the 1000 +Genomes ids, the graph neural network, the LLM decoder and the NVFlare run. The two share the same HaploGraph edge file, +the same UniProt gene BED and the same three-site design, and they are complementary rather than overlapping: nothing in +`genomics/` touches a path on `main`, and for the manuscript the proteomics analysis and the graph model plug into the +same methods and results structure. + +## 8. Next steps + +1. Ancestry-pure sites: the federated experiment where FedAvg must fight client drift and site-alone models fail to + transfer. +2. Real proteomics: Wu et al. 2013 (95 HapMap LCLs with 1000G ids) as a real fourth site; UKB-PPP cis-pQTLs (AWS Open + Data) as Cluster -> Protein propensity edges. +3. NVFlare POC mode across real machines (the L4 and the A100 as two sites); per-site held-out reporting only. +4. More chromosomes, then the whole genome, with the same code. + +## 9. Reproduce + +```bash +git clone https://github.com/collaborativebioinformatics/ProGenome.git && cd ProGenome && git checkout modelling && cd genomics +make setup && make run && make run-v2 && make federated ROUNDS=30 LOCAL_EPOCHS=5 # ~30 min on a laptop CPU +cp .env.example .env # add NVIDIA_API_KEY for the decoder, then: make decode WHO=HG00103 +make report # regenerates every figure on this page and the PDF / DOCX / LaTeX report +``` + +The [README](README.md) has the full implementation guide, the expected output of every stage and a troubleshooting table. diff --git a/genomics/baseline.py b/genomics/baseline.py new file mode 100644 index 0000000..b3fc8aa --- /dev/null +++ b/genomics/baseline.py @@ -0,0 +1,100 @@ +#!/usr/bin/env python3 +"""Baseline: L2 logistic regression on the sparse carrier matrix. + +Individuals are the rows, kept haploblock clusters the 0/1 columns. One model +per target (ancestry, population, sex); C is chosen on the validation split; +metrics are reported on the untouched test split. Sex is the negative control +(autosomal chromosome -> expect chance level). The split is saved so that the +GNN is scored on exactly the same people. +""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np +import pandas as pd +from sklearn.linear_model import LogisticRegression +from sklearn.metrics import accuracy_score, balanced_accuracy_score, f1_score + +import haplokg + +TARGETS = ("ancestry", "population", "sex") + + +def metrics(y_true, y_pred) -> dict: + return { + "accuracy": float(accuracy_score(y_true, y_pred)), + "balanced_accuracy": float(balanced_accuracy_score(y_true, y_pred)), + "macro_f1": float(f1_score(y_true, y_pred, average="macro", zero_division=0)), + "n": int(len(y_true)), + } + + +def main() -> int: + here = Path(__file__).resolve().parent + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--chrom", default="chr22") + parser.add_argument("--kg-dir", type=Path, default=None) + parser.add_argument("--out-dir", type=Path, default=None) + parser.add_argument("--split-path", type=Path, default=None) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--C-grid", type=float, nargs="+", default=[0.01, 0.1, 1.0]) + args = parser.parse_args() + kg_dir = args.kg_dir or here / "outputs" / "kg" / args.chrom + out_dir = args.out_dir or here / "outputs" / "baseline" / args.chrom + split_path = args.split_path or here / "outputs" / "splits" / args.chrom / f"split_seed{args.seed}.csv" + out_dir.mkdir(parents=True, exist_ok=True) + + kg = haplokg.load_kg(kg_dir) + ind, clusters, X = kg["individuals"], kg["clusters"], kg["carries"].astype(np.float32) + split = haplokg.load_or_make_split(kg, split_path, seed=args.seed) + masks = {name: split == name for name in ("train", "val", "test")} + print({k: int(v.sum()) for k, v in masks.items()}, "unlabelled:", int((split == "unlabelled").sum())) + + report = {"chrom": args.chrom, "seed": args.seed, "n_features": int(X.shape[1]), "targets": {}} + for target in TARGETS: + y = ind[f"{target}_code"].to_numpy() + ok = y >= 0 # a few individuals lack a label for this target + tr, va, te = masks["train"] & ok, masks["val"] & ok, masks["test"] & ok + classes = kg["label_maps"][target] + + majority = np.bincount(y[tr]).argmax() + chance = metrics(y[te], np.full(te.sum(), majority)) + + best = None + for C in args.C_grid: + model = LogisticRegression(C=C, max_iter=5000, random_state=args.seed) + model.fit(X[tr], y[tr]) + val = metrics(y[va], model.predict(X[va])) + if best is None or val["balanced_accuracy"] > best[1]["balanced_accuracy"]: + best = (C, val, model) + C, val, model = best + test = metrics(y[te], model.predict(X[te])) + report["targets"][target] = {"C": C, "val": val, "test": test, "majority_class_test": chance, "n_classes": len(classes)} + print(f"{target:11s} C={C:<5} val bal-acc={val['balanced_accuracy']:.3f} TEST acc={test['accuracy']:.3f} " + f"bal-acc={test['balanced_accuracy']:.3f} macro-F1={test['macro_f1']:.3f} (majority-class acc={chance['accuracy']:.3f})") + + pd.DataFrame({ + "individual_id": ind.loc[te, "individual_id"].to_numpy(), + "true": [classes[i] for i in y[te]], + "pred": [classes[i] for i in model.predict(X[te])], + }).to_csv(out_dir / f"{target}_test_predictions.csv", index=False) + + # which clusters drive each class (largest positive coefficients) + coef = model.coef_ if model.coef_.shape[0] > 1 else np.vstack([-model.coef_[0], model.coef_[0]]) + rows = [] + for k, name in enumerate(classes): + for j in np.argsort(coef[k])[::-1][:10]: + rows.append({"class": name, "cluster_id": clusters.loc[j, "cluster_id"], "coef": float(coef[k, j]), + "support": int(clusters.loc[j, "support"])}) + pd.DataFrame(rows).to_csv(out_dir / f"{target}_top_clusters_per_class.csv", index=False) + + (out_dir / "metrics.json").write_text(json.dumps(report, indent=2)) + print(f"wrote {out_dir}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/genomics/brev_deploy.sh b/genomics/brev_deploy.sh new file mode 100755 index 0000000..6e4c610 --- /dev/null +++ b/genomics/brev_deploy.sh @@ -0,0 +1,54 @@ +#!/usr/bin/env bash +# Deploy and run the pipeline on an NVIDIA Brev GPU instance, then pull the outputs back. +# bash brev_deploy.sh # instance $BREV_INSTANCE (default progenome-gpu); created if missing +# BREV_INSTANCE=progenome-a100 bash brev_deploy.sh +# STAGE=infer bash brev_deploy.sh # only the inference benchmark on an existing run +# INIT=node2vec bash brev_deploy.sh # Node2Vec instead of SVD starting embeddings (slightly weaker: ancestry 0.95 vs 0.97) +# Needs: brev CLI logged in (brev login --api-key ...). `brev refresh` writes an SSH alias named after the +# instance; this script drives the box with plain ssh/scp through that alias (brev exec is interactive-prone). +set -euo pipefail +cd "$(dirname "$0")" +. ./load_env.sh # BREV_INSTANCE / BREV_TYPE / INIT from .env +INSTANCE="${BREV_INSTANCE:-progenome-gpu}" +TYPE="${BREV_TYPE:-g2-standard-4:nvidia-l4:1}" # L4 24 GB (works without extra cloud credentials) +STAGE="${STAGE:-all}" # all | infer +INIT="${INIT:-svd}" # svd is the documented best; INIT=node2vec needs pyg_lib (in the image), 50 pretraining epochs +REMOTE=/home/ubuntu/progenome +SSH="ssh -o BatchMode=yes -o StrictHostKeyChecking=accept-new -o ServerAliveInterval=30" +SCP="scp -o BatchMode=yes -o StrictHostKeyChecking=accept-new" + +if ! brev ls 2>/dev/null | grep -qE "^\s*${INSTANCE}\s"; then + echo "== creating ${INSTANCE} (${TYPE})" + brev create "${INSTANCE}" --type "${TYPE}" --min-disk 100 --timeout 900 +fi +echo "== waiting for ${INSTANCE} shell" +until brev ls 2>/dev/null | grep -qE "^\s*${INSTANCE}\s+RUNNING\s+\S+\s+READY"; do sleep 15; done +brev refresh >/dev/null 2>&1 || true +$SSH "${INSTANCE}" 'nvidia-smi -L; docker --version' + +if [ "$STAGE" = all ]; then + echo "== upload code + data" + # genomics/ plus the two small proteomics inputs the v2 schema needs (protein BED, gene symbols) + tar czf /tmp/genomics_upload.tgz --exclude=outputs --exclude=outputs_brev --exclude='__pycache__' --exclude=.venv --exclude=.pytest_cache \ + -C .. genomics proteomics/uniprot_chr22.bed proteomics/synthetic_proteomics_chr22/gene_symbol_cache.csv + $SCP /tmp/genomics_upload.tgz "${INSTANCE}:/tmp/genomics_upload.tgz" + $SSH "${INSTANCE}" "mkdir -p ${REMOTE} && tar xzf /tmp/genomics_upload.tgz -C ${REMOTE} && ls ${REMOTE}/genomics | head -3" + + echo "== build image on the GPU box (native, no emulation)" + $SSH "${INSTANCE}" "cd ${REMOTE} && docker build -f genomics/Dockerfile -t progenome-genomics . 2>&1 | grep -E '^#[0-9]+ (DONE|ERROR)|passed|failed|error|Successfully|naming to' | tail -20" + + echo "== run the full pipeline on the GPU (INIT=${INIT})" + $SSH "${INSTANCE}" "cd ${REMOTE}/genomics && mkdir -p outputs && docker run --rm --gpus all -e INIT=${INIT} \ + -v ${REMOTE}/genomics/data:/app/genomics/data -v ${REMOTE}/genomics/outputs:/app/genomics/outputs progenome-genomics run_all.sh" +fi + +echo "== inference benchmark: eager vs Torch-TensorRT" +RUN_NAME=$($SSH "${INSTANCE}" "ls ${REMOTE}/genomics/outputs/gnn/chr22 2>/dev/null | grep -E '^ancestry_(node2vec|svd)$' | head -1") +$SSH "${INSTANCE}" "cd ${REMOTE}/genomics && docker run --rm --gpus all \ + -v ${REMOTE}/genomics/data:/app/genomics/data -v ${REMOTE}/genomics/outputs:/app/genomics/outputs progenome-genomics \ + -c 'python infer.py --run ${RUN_NAME} --compile none && python infer.py --run ${RUN_NAME} --compile inductor && python infer.py --run ${RUN_NAME} --compile tensorrt --precision fp16'" + +echo "== copy outputs back to outputs_brev/${INSTANCE}/" +mkdir -p "outputs_brev/${INSTANCE}" +$SCP -r "${INSTANCE}:${REMOTE}/genomics/outputs/." "outputs_brev/${INSTANCE}/" +echo "done. Stop billing when finished: brev stop ${INSTANCE} (delete: brev delete ${INSTANCE})" diff --git a/genomics/build_kg.py b/genomics/build_kg.py new file mode 100644 index 0000000..088f43f --- /dev/null +++ b/genomics/build_kg.py @@ -0,0 +1,66 @@ +#!/usr/bin/env python3 +"""Build the haploblock knowledge graph for one chromosome. + + python build_kg.py --chrom chr22 --min-support 25 + +Reads data/ (see fetch_data.sh) and writes outputs/kg//: + individuals.csv clusters.csv blocks.csv co_occurs.csv next_block.csv + carries.npz sparse 0/1 matrix, individuals x kept clusters + label_maps.json class order for ancestry / population / sex codes + hetero.pt torch_geometric HeteroData (load with torch.load(..., weights_only=False)) + summary.json +""" +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import haplokg + + +def main() -> int: + here = Path(__file__).resolve().parent + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--chrom", default="chr22") + parser.add_argument("--data-dir", type=Path, default=here / "data") + parser.add_argument("--out-dir", type=Path, default=None, help="default outputs/kg/") + parser.add_argument("--min-support", type=int, default=25, + help="keep clusters carried by >= this many individuals (default 25, same as the edge file)") + parser.add_argument("--no-symmetric", action="store_true", + help="also keep near-universal clusters (default drops clusters with < min-support non-carriers)") + parser.add_argument("--chunksize", type=int, default=8192) + parser.add_argument("--no-torch", action="store_true", help="skip writing hetero.pt") + args = parser.parse_args() + out_dir = args.out_dir or here / "outputs" / "kg" / args.chrom + + graph_dir = args.data_dir / "haplograph" / args.chrom + t0 = time.time() + print(f"[1/4] streaming {graph_dir / 'nodes.csv.gz'} ...", flush=True) + cluster_ids, block_ids, individual_ids, matrix = haplokg.read_node_matrix(graph_dir / "nodes.csv.gz", args.chunksize) + print(f" {matrix.shape[0]:,} clusters x {matrix.shape[1]:,} individuals, nnz={matrix.nnz:,} ({time.time()-t0:.0f}s)") + + print("[2/4] loading phenotypes, block stats, edges", flush=True) + phenotypes = haplokg.load_phenotypes(args.data_dir / "haplograph" / "phenotypes_real.csv") + block_stats = haplokg.load_block_stats(args.data_dir / "haploblocks" / "block_stats.tsv", args.chrom) + edges = haplokg.load_edges(graph_dir / "edges_lift_above_threshold.csv.gz") + + print(f"[3/4] building tables (min_support={args.min_support}, symmetric={not args.no_symmetric})", flush=True) + tables = haplokg.build_tables(cluster_ids, block_ids, individual_ids, matrix, phenotypes, block_stats, edges, + min_support=args.min_support, symmetric=not args.no_symmetric) + summary = haplokg.save_tables(tables, out_dir) + + if not args.no_torch: + import torch + print("[4/4] writing HeteroData", flush=True) + data = haplokg.to_hetero_data(tables) + torch.save(data, out_dir / "hetero.pt") + print(data) + print(json.dumps(summary, indent=2)) + print(f"wrote {out_dir} ({time.time()-t0:.0f}s total)") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/genomics/build_kg_v2.py b/genomics/build_kg_v2.py new file mode 100644 index 0000000..eacd3c7 --- /dev/null +++ b/genomics/build_kg_v2.py @@ -0,0 +1,50 @@ +#!/usr/bin/env python3 +"""Schema v2: add genes, proteins and per-individual proteomics to the genome knowledge graph. + + python build_kg_v2.py --chrom chr22 # uses outputs/proteomics_synth/ + python build_kg_v2.py --chrom chr22 --measured my_measured_long.csv --metadata my_samples.csv + +Requires build_kg.py to have run. Writes genes.csv, proteins.csv, block_gene.csv, gene_protein.csv, +measured.csv, abundance_z.npz, abundance_observed.npz, individuals_v2.csv, label_maps_v2.json and +hetero_v2.pt into outputs/kg//. +""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import haplokg +import haplokg_proteins as hp + + +def main() -> int: + here = Path(__file__).resolve().parent + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--chrom", default="chr22") + parser.add_argument("--bed", type=Path, default=here.parent / "proteomics" / "uniprot_chr22.bed") + parser.add_argument("--symbols", type=Path, default=here.parent / "proteomics" / "synthetic_proteomics_chr22" / "gene_symbol_cache.csv") + parser.add_argument("--measured", type=Path, default=None, help="long csv: individual_id, protein_id, log2_intensity, site") + parser.add_argument("--metadata", type=Path, default=None, help="csv: sample_id, site, age, sex, phenotype") + args = parser.parse_args() + kg_dir = here / "outputs" / "kg" / args.chrom + synth = here / "outputs" / "proteomics_synth" / args.chrom + measured = args.measured or synth / "measured_long.csv" + metadata = args.metadata or synth / "sample_metadata.csv" + + kg = haplokg.load_kg(kg_dir) + t = hp.build_protein_tables(kg, args.bed, measured, metadata, args.symbols) + summary = hp.save_protein_tables(t, kg_dir) + + import torch + data = torch.load(kg_dir / "hetero.pt", weights_only=False) + data = hp.extend_hetero_data(data, t) + torch.save(data, kg_dir / "hetero_v2.pt") + print(data) + print(json.dumps(summary, indent=2)) + print(f"wrote {kg_dir}/hetero_v2.pt") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/genomics/config.py b/genomics/config.py new file mode 100644 index 0000000..5b522bf --- /dev/null +++ b/genomics/config.py @@ -0,0 +1,106 @@ +#!/usr/bin/env python3 +"""One place where credentials and settings enter the code. + +Precedence (first wins): variables already exported in the shell > genomics/.env > ~/.progenome.env > defaults. +Copy .env.example to .env and fill it in; .env is git-ignored. Every script that needs a credential or an endpoint +imports `settings` from here instead of reading os.environ itself. + + python config.py # prints the effective settings with secrets masked, and where each came from +""" +from __future__ import annotations + +import os +import sys +from dataclasses import dataclass, field, fields +from pathlib import Path + +GENOMICS = Path(__file__).resolve().parent +ENV_FILES = (GENOMICS / ".env", Path.home() / ".progenome.env") +SECRET_KEYS = {"NVIDIA_API_KEY", "NEO4J_PASSWORD", "BREV_API_KEY"} + + +def _parse_env_file(path: Path) -> dict[str, str]: + """KEY=VALUE lines; 'export KEY=VALUE' accepted; blank lines and full-line comments ignored; quotes stripped.""" + out: dict[str, str] = {} + for raw in path.read_text().splitlines(): + line = raw.strip() + if not line or line.startswith("#") or "=" not in line: + continue + key, value = line.split("=", 1) + key = key.removeprefix("export ").strip() + value = value.strip() + if value[:1] in ("'", '"') and value[-1:] == value[:1]: + value = value[1:-1] + if key: + out[key] = value + return out + + +def load_env_files() -> dict[str, str]: + """Put file values into os.environ for keys that are not already set; return {key: source} for reporting.""" + source: dict[str, str] = {} + for path in ENV_FILES: + if not path.exists(): + continue + for key, value in _parse_env_file(path).items(): + if key not in os.environ: + os.environ[key] = value + source[key] = str(path) + return source + + +_SOURCES = load_env_files() + + +def _env(key: str, default: str) -> str: + return os.environ.get(key, default) + + +@dataclass(frozen=True) +class Settings: + # LLM decoder (NVIDIA NIM or any OpenAI-compatible chat endpoint) + nvidia_api_key: str = field(default_factory=lambda: _env("NVIDIA_API_KEY", "")) + nim_model: str = field(default_factory=lambda: _env("NIM_MODEL", "nvidia/nemotron-3-super-120b-a12b")) + nim_url: str = field(default_factory=lambda: _env("NIM_URL", "https://integrate.api.nvidia.com/v1/chat/completions")) + # Neo4j browser + neo4j_uri: str = field(default_factory=lambda: _env("NEO4J_URI", "bolt://localhost:7687")) + neo4j_user: str = field(default_factory=lambda: _env("NEO4J_USER", "neo4j")) + neo4j_password: str = field(default_factory=lambda: _env("NEO4J_PASSWORD", "progenome")) + # data source and defaults + haploblocks_base: str = field(default_factory=lambda: _env("HAPLOBLOCKS_BASE", "https://data.haploblocks.org")) + chrom: str = field(default_factory=lambda: _env("CHROM", "chr22")) + # NVIDIA Brev (the CLI keeps its own login; these only pick the instance) + brev_instance: str = field(default_factory=lambda: _env("BREV_INSTANCE", "progenome-gpu")) + brev_type: str = field(default_factory=lambda: _env("BREV_TYPE", "g2-standard-4:nvidia-l4:1")) + # paths (relative to this folder unless overridden) + data_dir: Path = field(default_factory=lambda: Path(_env("PROGENOME_DATA_DIR", str(GENOMICS / "data")))) + outputs_dir: Path = field(default_factory=lambda: Path(_env("PROGENOME_OUTPUTS_DIR", str(GENOMICS / "outputs")))) + + +settings = Settings() + + +def require(env_key: str) -> str: + """Return a setting that must be present, or exit with instructions instead of a traceback.""" + value = os.environ.get(env_key, "") + if not value or value.endswith("REPLACE_ME"): + raise SystemExit(f"{env_key} is not set. Copy genomics/.env.example to genomics/.env (or ~/.progenome.env) and fill it in, " + f"or export {env_key}=... ; see README step 4.") + return value + + +def describe() -> str: + lines = [] + for f_ in fields(Settings): + key = f_.name.upper() if not f_.name.endswith("_dir") else f"PROGENOME_{f_.name.upper()}" + value = getattr(settings, f_.name) + shown = ("***" + str(value)[-4:] if value else "(not set)") if key in SECRET_KEYS else str(value) + origin = _SOURCES.get(key, "exported" if key in os.environ else "default") + lines.append(f"{key:24s} {shown:60s} [{origin}]") + return "\n".join(lines) + + +if __name__ == "__main__": + print(f"env files read: {[str(p) for p in ENV_FILES if p.exists()] or 'none'}") + print(describe()) + sys.exit(0) diff --git a/genomics/cooccurrence_analysis.py b/genomics/cooccurrence_analysis.py new file mode 100644 index 0000000..92c04fe --- /dev/null +++ b/genomics/cooccurrence_analysis.py @@ -0,0 +1,213 @@ +#!/usr/bin/env python3 +"""Does the haploblock-cluster graph co-occur with phenotypes? + +Three questions, all answered from the knowledge graph in outputs/kg/: + +1. cluster x phenotype For every cluster: is being a carrier independent of + ancestry / population / sex? (chi-square on the 2 x k + carrier-by-class table, Cramer's V, BH-FDR). Sex is the + negative control: chr22 is autosomal, so nothing should + pass. +2. edge x phenotype Do co-occurring clusters (lift edges) have more similar + ancestry profiles than random cluster pairs? If the + graph structure tracks population structure, yes. +3. block x phenotype Where along the chromosome is the graph most + ancestry-informative? (max Cramer's V per block) + +Writes CSV tables, a summary.json and three PNGs to outputs/cooccurrence//. +""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np +import pandas as pd +from scipy import sparse, stats + +import haplokg + +TARGETS = ("ancestry", "population", "sex") + + +def bh_fdr(p: np.ndarray) -> np.ndarray: + p = np.asarray(p, dtype=float) + n = len(p) + order = np.argsort(p) + ranked = p[order] * n / (np.arange(n) + 1) + q = np.minimum.accumulate(ranked[::-1])[::-1] + out = np.empty(n) + out[order] = np.clip(q, 0, 1) + return out + + +def association(carries: sparse.csr_matrix, codes: np.ndarray, n_classes: int) -> dict: + """Carrier-vs-class independence test for every cluster at once. + + carries: (n_labelled x clusters) 0/1; codes: class code per row (0..k-1). + Cramer's V for a 2 x k table is sqrt(chi2 / n) because min(2-1, k-1) = 1. + """ + n = carries.shape[0] + onehot = sparse.csr_matrix((np.ones(n), (np.arange(n), codes)), shape=(n, n_classes)) + carriers = np.asarray((onehot.T @ carries).todense(), dtype=float) # k x clusters + n_class = np.asarray(onehot.sum(axis=0)).ravel() # k + support = carriers.sum(axis=0) # clusters + non_carriers = n_class[:, None] - carriers + expected1 = np.outer(n_class, support) / n + expected0 = n_class[:, None] - expected1 + with np.errstate(divide="ignore", invalid="ignore"): + chi2 = np.nansum((carriers - expected1) ** 2 / expected1, axis=0) + \ + np.nansum((non_carriers - expected0) ** 2 / expected0, axis=0) + p = stats.chi2.sf(chi2, df=n_classes - 1) + frac = carriers / n_class[:, None] # P(carry | class) + overall = support / n + with np.errstate(divide="ignore", invalid="ignore"): + enrich = np.where(overall > 0, frac / overall, np.nan) # k x clusters + return { + "chi2": chi2, "p": p, "q": bh_fdr(p), "cramers_v": np.sqrt(chi2 / n), + "frac": frac, "enrich": enrich, "dominant": np.nanargmax(np.nan_to_num(enrich, nan=-1), axis=0), + "support": support, + } + + +def edge_profile_similarity(co: pd.DataFrame, enrich: np.ndarray, dominant: np.ndarray, rng: np.random.Generator) -> dict: + """Cosine similarity of the two endpoints' ancestry-deviation profiles, real edges vs shuffled targets.""" + deviation = np.nan_to_num(enrich.T - 1.0) # clusters x k, 0 = neutral + norm = np.linalg.norm(deviation, axis=1, keepdims=True) + unit = deviation / np.where(norm > 0, norm, 1.0) + src, dst = co["src"].to_numpy(), co["dst"].to_numpy() + perm = rng.permutation(dst) # keeps the degree of every source + cos_real = (unit[src] * unit[dst]).sum(1) + cos_null = (unit[src] * unit[perm]).sum(1) + same_real = dominant[src] == dominant[dst] + same_null = dominant[src] == dominant[perm] + rho, rho_p = stats.spearmanr(co["lift"].to_numpy(), cos_real) + return { + "n_edges": int(len(co)), + "profile_cosine_real_mean": float(cos_real.mean()), + "profile_cosine_null_mean": float(cos_null.mean()), + "same_dominant_ancestry_real": float(same_real.mean()), + "same_dominant_ancestry_null": float(same_null.mean()), + "spearman_lift_vs_profile_cosine": float(rho), + "spearman_p": float(rho_p), + "_cos_real": cos_real, "_cos_null": cos_null, + } + + +def main() -> int: + here = Path(__file__).resolve().parent + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--chrom", default="chr22") + parser.add_argument("--kg-dir", type=Path, default=None) + parser.add_argument("--out-dir", type=Path, default=None) + parser.add_argument("--fdr", type=float, default=0.05) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--no-plots", action="store_true") + args = parser.parse_args() + kg_dir = args.kg_dir or here / "outputs" / "kg" / args.chrom + out_dir = args.out_dir or here / "outputs" / "cooccurrence" / args.chrom + out_dir.mkdir(parents=True, exist_ok=True) + rng = np.random.default_rng(args.seed) + + kg = haplokg.load_kg(kg_dir) + ind, clusters, blocks, co = kg["individuals"], kg["clusters"], kg["blocks"], kg["co_occurs"] + labelled = ind["ancestry_code"].to_numpy() >= 0 + carries = kg["carries"][labelled] + summary = {"chrom": args.chrom, "n_individuals_labelled": int(labelled.sum()), "n_clusters": int(len(clusters)), "fdr": args.fdr} + + # ---- 1. cluster x phenotype ------------------------------------------------- + table = clusters[["cluster_idx", "cluster_id", "block_id", "block_idx", "support"]].copy() + results = {} + for target in TARGETS: + codes = ind.loc[labelled, f"{target}_code"].to_numpy() + classes = kg["label_maps"][target] + res = association(carries, codes, len(classes)) + results[target] = res + table[f"{target}_cramers_v"] = res["cramers_v"] + table[f"{target}_p"] = res["p"] + table[f"{target}_q"] = res["q"] + table[f"{target}_dominant"] = [classes[i] for i in res["dominant"]] + if target == "ancestry": + for i, name in enumerate(classes): + table[f"carrier_frac_{name}"] = res["frac"][i] + sig = res["q"] < args.fdr + summary[f"{target}_clusters_significant"] = int(sig.sum()) + summary[f"{target}_clusters_significant_frac"] = float(sig.mean()) + summary[f"{target}_cramers_v_median"] = float(np.median(res["cramers_v"])) + summary[f"{target}_cramers_v_p90"] = float(np.quantile(res["cramers_v"], 0.9)) + table.to_csv(out_dir / "cluster_phenotype_association.csv", index=False) + top = table.sort_values("ancestry_cramers_v", ascending=False).head(25) + top.to_csv(out_dir / "top_ancestry_clusters.csv", index=False) + + # ---- 2. edge x phenotype ---------------------------------------------------- + edge = edge_profile_similarity(co, results["ancestry"]["enrich"], results["ancestry"]["dominant"], rng) + cos_real, cos_null = edge.pop("_cos_real"), edge.pop("_cos_null") + summary["edges"] = edge + # by lift quartile + q = pd.qcut(co["lift"], 4, labels=["Q1 (lowest lift)", "Q2", "Q3", "Q4 (highest lift)"]) + by_lift = pd.DataFrame({"lift_quartile": q, "cos": cos_real}).groupby("lift_quartile", observed=True)["cos"].agg(["mean", "count"]).reset_index() + by_lift.to_csv(out_dir / "edge_similarity_by_lift_quartile.csv", index=False) + summary["edges"]["profile_cosine_by_lift_quartile"] = dict(zip(by_lift["lift_quartile"].astype(str), by_lift["mean"].round(4))) + + # ---- 3. block x phenotype --------------------------------------------------- + per_block = table.groupby("block_idx").agg( + n_clusters_kept=("cluster_idx", "size"), + max_ancestry_v=("ancestry_cramers_v", "max"), + n_ancestry_significant=("ancestry_q", lambda s: int((s < args.fdr).sum())), + max_sex_v=("sex_cramers_v", "max"), + ).reset_index() + per_block = blocks[["block_idx", "block_id", "start", "end", "block_length"]].merge(per_block, on="block_idx", how="left") + per_block.to_csv(out_dir / "block_informativeness.csv", index=False) + summary["blocks_with_kept_clusters"] = int(per_block["n_clusters_kept"].notna().sum()) + summary["blocks_with_significant_ancestry_cluster"] = int((per_block["n_ancestry_significant"].fillna(0) > 0).sum()) + + (out_dir / "summary.json").write_text(json.dumps(summary, indent=2, default=float)) + + # ---- plots ------------------------------------------------------------------ + if not args.no_plots: + import matplotlib + matplotlib.use("Agg") + import matplotlib.pyplot as plt + + colours = {"ancestry": "#1f6f8b", "population": "#6a4c93", "sex": "#c44536"} + fig, ax = plt.subplots(figsize=(7, 4)) + bins = np.linspace(0, max(0.05, table[[f"{t}_cramers_v" for t in TARGETS]].to_numpy().max()), 60) + for target in TARGETS: + ax.hist(table[f"{target}_cramers_v"], bins=bins, histtype="step", linewidth=1.8, + label=f"{target} ({summary[f'{target}_clusters_significant']:,} clusters FDR<{args.fdr})", color=colours[target]) + ax.set_xlabel("Cramér's V (carrier status vs phenotype)") + ax.set_ylabel("clusters") + ax.set_title(f"{args.chrom}: how strongly each haploblock cluster tracks a phenotype") + ax.legend(frameon=False) + fig.tight_layout(); fig.savefig(out_dir / "cramers_v_by_phenotype.png", dpi=150, bbox_inches="tight"); plt.close(fig) + + fig, ax = plt.subplots(figsize=(9, 3.6)) + pb = per_block.dropna(subset=["max_ancestry_v"]) + mid = (pb["start"] + pb["end"]) / 2e6 + ax.scatter(mid, pb["max_ancestry_v"], s=10, color=colours["ancestry"], label="ancestry (max V in block)") + ax.scatter(mid, pb["max_sex_v"], s=10, color=colours["sex"], alpha=0.7, label="sex (negative control)") + ax.set_xlabel(f"{args.chrom} position (Mb)"); ax.set_ylabel("max Cramér's V per block") + ax.set_title("Where along the chromosome the haploblock graph tracks ancestry") + ax.legend(frameon=False, loc="upper center", bbox_to_anchor=(0.5, -0.22), ncol=2) # below the axes, never on the points + fig.tight_layout(); fig.savefig(out_dir / "informativeness_along_chromosome.png", dpi=150, bbox_inches="tight"); plt.close(fig) + + fig, ax = plt.subplots(figsize=(7, 4)) + bins = np.linspace(-1, 1, 50) + ax.hist(cos_null, bins=bins, alpha=0.6, color="#9a9a9a", label=f"shuffled pairs (mean {edge['profile_cosine_null_mean']:.2f})") + ax.hist(cos_real, bins=bins, alpha=0.7, color=colours["ancestry"], label=f"lift edges (mean {edge['profile_cosine_real_mean']:.2f})") + ax.set_xlabel("cosine similarity of endpoint ancestry profiles"); ax.set_ylabel("edges") + ax.set_title("Co-occurring clusters share ancestry profiles") + ax.legend(frameon=False, loc="upper left") + fig.tight_layout(); fig.savefig(out_dir / "edge_ancestry_similarity.png", dpi=150, bbox_inches="tight"); plt.close(fig) + + print(json.dumps(summary, indent=2, default=float)) + print("\nTop ancestry-informative clusters:") + cols = ["cluster_id", "support", "ancestry_cramers_v", "ancestry_dominant"] + [c for c in table.columns if c.startswith("carrier_frac_")] + print(top[cols].head(10).to_string(index=False, float_format=lambda v: f"{v:.3f}")) + print(f"\nwrote {out_dir}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/genomics/docker-compose.yml b/genomics/docker-compose.yml new file mode 100644 index 0000000..00a9c52 --- /dev/null +++ b/genomics/docker-compose.yml @@ -0,0 +1,51 @@ +# Local graph browser for the haploblock knowledge graph. +# docker compose up -d neo4j # then open http://localhost:7474 (user neo4j / password progenome) +# python neo4j_load.py --chrom chr22 # loads outputs/kg/chr22 into it (a few minutes for all CARRIES edges) +# docker compose down # keeps the data volume; add -v to wipe it +services: + neo4j: + image: neo4j:5.26-community + container_name: progenome-neo4j + ports: + - "7474:7474" # browser + - "7687:7687" # bolt + environment: + NEO4J_AUTH: neo4j/progenome + NEO4J_server_memory_heap_initial__size: 1G + NEO4J_server_memory_heap_max__size: 2G + NEO4J_server_memory_pagecache_size: 1G + NEO4J_dbms_security_procedures_unrestricted: "apoc.*" + volumes: + - neo4j_data:/data + healthcheck: + test: ["CMD-SHELL", "wget -qO- http://localhost:7474 >/dev/null || exit 1"] + interval: 10s + timeout: 5s + retries: 20 + + # GPU pipeline image (see Dockerfile); run stages with e.g. + # docker compose run --rm pipeline python train_gnn.py --target ancestry + pipeline: + build: + context: .. + dockerfile: genomics/Dockerfile + image: progenome-genomics + env_file: + - path: .env + required: false + volumes: + - ./data:/app/genomics/data + - ./outputs:/app/genomics/outputs + environment: + NEO4J_URI: bolt://neo4j:7687 + NX_CUGRAPH_AUTOCONFIG: "True" + deploy: + resources: + reservations: + devices: + - driver: nvidia + count: all + capabilities: [gpu] + +volumes: + neo4j_data: {} diff --git a/genomics/docs/DEEP_DIVE.md b/genomics/docs/DEEP_DIVE.md new file mode 100644 index 0000000..8776144 --- /dev/null +++ b/genomics/docs/DEEP_DIVE.md @@ -0,0 +1,319 @@ +# ProGenome deep dive — every step, with the formulas, shapes and numbers + +Companion to `METHODS.md` (the short account) and `../README.md` (how to run). This file explains *how* +each stage works at the level of tensors and equations, and *why* it was built that way. chr22, seed 42. + +--- + +## 0. Vocabulary + +| term | meaning here | +|---|---| +| haploblock (block) | a recombination-defined stretch of chromosome; 669 on chr22, median 29.7 kb, tiling 17.1–50.2 Mb with no gaps | +| haplotype | one of a person's two copies of a block's sequence (phased: we know which variants sit on the same copy) | +| cluster | a group of haplotypes of one block that MMseqs2 called similar; a person *carries* a cluster if at least one of their two haplotypes is in it | +| carrier matrix | `M ∈ {0,1}^{N×C}`: `M[i,c] = 1` iff person `i` carries cluster `c`; N = 2,548 people, C = 6,551 kept clusters | +| support | number of carriers of a cluster, `Σ_i M[i,c]` | +| lift | `P(A∩B) / (P(A)·P(B))` for two clusters: >1 means they travel together more than chance | +| phenotype | a property of a person used as a label: real (ancestry, population, sex) or synthetic (case/control) | +| site | a hospital; owns its people's rows of `M` and their proteomics; never sees other sites' rows | + +--- + +## 1. Inputs, byte by byte + +### 1.1 `nodes.csv.gz` (HaploGraph) + +``` +id,high_dim_edge,HG00096,HG00097,...,NA21144 # 2 + 2,548 columns +chr22_17099658-17118145_cluster1,chr22_17099658-17118145,1,0,...,1 +``` + +248,254 rows (one per cluster of any size), 2,548 person columns. Uncompressed 1.3 GB, gzipped 9 MB. +Read in 8,192-row chunks with pandas, dtype `int8` for every person column, each chunk turned into a +`scipy.sparse.csr_matrix` and stacked (`haplokg.read_node_matrix`). Peak memory stays under 1 GB; the +full dense int8 matrix would be 632 MB and int64 would be 5 GB. + +### 1.2 The cluster filter + +`support[c] = Σ_i M[i,c]`; keep `c` iff `25 ≤ support[c] ≤ N − 25`. Rationale: a cluster carried by 2,540 +of 2,548 people is as uninformative as one carried by 8 — the "informative side" has 8 people either way +(this mirrors the HaploGraph's own `MIN_CLUSTER_SUPPORT=25` symmetric edge filter). 248,254 → 6,551. +176,903 of the dropped clusters are singletons (one haplotype). The filter reproduces the edge file's node +set exactly: 0 of 187,030 edges lose an endpoint. + +### 1.3 Edges + +`source,target,weight,lift`. `weight` = number of people carrying both; `lift` as above, ≥ 5 by +construction. Stored both directions in PyG (374,060 directed edges) with `edge_attr = [weight, lift]`. +We do **not** use the raw `edges.csv` (1 GB for chr6): its weights are dominated by one near-universal +haplotype ("mega-hub"), which the lift normalisation removes. + +### 1.4 Phenotypes + +`phenotypes_real.csv` long format `individual_id,phenotype,value,source` → pivoted wide. Labels encoded as +integers by sorted class name (`AFR=0, AMR=1, EAS=2, EUR=3, SAS=4`; `female=0, male=1`; 26 populations); +`-1` for the 45 people without labels (they stay in the graph as unlabelled nodes and are excluded from +every loss and metric). + +### 1.5 Proteins + +`uniprot_chr22.bed` (BED12 from the UCSC UniProt track): 917 isoform rows like `Q9BXF3-1`, `Q9BXF3-2`. +Isoforms are collapsed to the base accession (`Q9BXF3`), giving 460 proteins with the union span of their +isoforms; gene symbols from `gene_symbol_cache.csv` (UniProt lookup) give 458 genes. + +--- + +## 2. Graph schema and how each edge is computed + +| edge | computation | count | +|---|---|---| +| `Individual —CARRIES→ Cluster` | non-zeros of `M` after the filter (`carries = M[kept].T`), coordinates `(i, c)` | 2,365,574 | +| `Cluster —IN_BLOCK→ Block` | parsed from the cluster id (`chr22_-_clusterN` → block `chr22_-`) | 6,551 | +| `Block —NEXT_BLOCK→ Block` | blocks sorted by `start`; consecutive pairs | 668 | +| `Cluster —CO_OCCURS→ Cluster` | edge file, ids mapped to kept-cluster indices | 187,030 | +| `Block —OVERLAPS→ Gene` | `block.start ≤ gene.end` and `block.end ≥ gene.start+1` (BED is 0-based half-open; blocks 1-based inclusive) | 1,063 | +| `Gene —ENCODES→ Protein` | from the BED (one gene → up to 2 proteins) | 460 | +| `Individual —MEASURED→ Protein` | one edge per observed (person, protein) measurement; attributes `[z, log2]`; missing = no edge | 1,065,712 | + +Node features (all z-scored, NaN → 0): + +- `cluster.x ∈ ℝ^{6551×7}`: `log1p(support)`, `support/N`, block `log10(length)`, block entropy, block dominance, block `log1p(n_clusters)`, block singleton rate. +- `block.x ∈ ℝ^{669×5}`: the five block statistics. +- `gene.x`, `protein.x ∈ ℝ^{·×2}`: `log10(span)`, `log1p(n_proteins | n_isoforms)`. +- `individual.x`: set by the model (§5), not stored in the graph. + +The join key everywhere is the 1000G sample id. Phenotype labels are stored as `individual.y_` tensors — +node **properties**, never nodes or edges. This matters: if `Individual —HAS→ Phenotype` existed, a two-layer +GNN would read the label from its neighbour and report ~100 %. + +--- + +## 3. The statistics (what "co-occurrence with phenotypes" means numerically) + +### 3.1 Cluster × phenotype (`cooccurrence_analysis.py`) + +For each cluster `c` and a k-class label, the 2×k table of (carrier, non-carrier) × class. With one-hot +`A ∈ {0,1}^{N×k}` and the carrier matrix `M`: observed carriers per class `O₁ = Aᵀ M ∈ ℝ^{k×C}` in one sparse +product; non-carriers `O₀ = n_class − O₁`; expected `E₁ = n_class · support / N`, `E₀ = n_class − E₁`; +`χ² = Σ (O₁−E₁)²/E₁ + Σ (O₀−E₀)²/E₀` with `k−1` degrees of freedom; Cramér's `V = √(χ² / (N·min(1, k−1))) = √(χ²/N)` +because one dimension is binary. p-values → Benjamini–Hochberg q-values. A cluster's *dominant* ancestry is +the class with the largest enrichment `frac_class / frac_overall`. + +Result: ancestry 6,470/6,551 clusters at q < 0.05 (median V 0.20, max 0.81); population 6,378; **sex 0** +(median V 0.013, max 0.03). Sex is the negative control: any method that found sex signal on an autosome +would be fitting noise. + +### 3.2 Edge × phenotype + +For every lift edge, both endpoints' 5-vector of enrichment deviations `(frac_a/frac_overall − 1)` is +L2-normalised and the cosine similarity taken: mean **0.86** for real edges vs **0.22** for a null made by +permuting the target column (keeps each source's degree). 86 % of edges join two clusters with the same +dominant ancestry (42 % under the null). Similarity rises with lift (quartiles 0.84, 0.84, 0.85, 0.90). +Reading: the graph structure *is* population structure; the median endpoint distance is tens of Mb, i.e. +these are not physical linkage but co-inheritance within ancestries. + +--- + +## 4. Baseline: logistic regression on `M` + +`LogisticRegression(C ∈ {0.01, 0.1, 1}, max_iter=5000)` on the sparse float carrier matrix, one model per +label; C chosen on validation balanced accuracy; test scores: ancestry 0.977, population 0.614 (26 classes, +60–113 people each), sex 0.463. The split: `train_test_split` stratified on ancestry, 70/15/15 over the +2,503 labelled people, seed 42, written once to `outputs/splits/chr22/split_seed42.csv` and reused by every +model, so all held-out numbers are on the same 376 people. + +--- + +## 5. Embeddings: how a node becomes a vector + +### 5.1 Truncated SVD (`--init svd`, default) + +`M ≈ U Σ Vᵀ` with `k = 32` (`scipy.sparse.linalg.svds`): `U ∈ ℝ^{2548×32}`, `Σ` diagonal (top values 1025, 287, +198, 112, 101), `V ∈ ℝ^{6551×32}`. Person embedding `UΣ`, cluster embedding `VΣ`, each column z-scored. +Both live in one space: `(UΣ)(VΣ)ᵀ ≈ MΣ`, so a person is close to the clusters they carry and two people who +share clusters are close to each other. No labels are used, so nothing can leak into the test split. + +### 5.2 Node2Vec (`--init node2vec`) + +Random walks (length 20, 10 per node, context 10) on the homogeneous graph of people + clusters with +CARRIES and CO_OCCURS edges, skip-gram objective with 1 negative sample; `pyg-lib` provides the walks +(`torch_geometric.nn.Node2Vec` requires `pyg-lib ≥ 0.6`; `torch_cluster` is deprecated). Runs in the GPU image. + +### 5.3 Learned and raw + +`--init learned`: an `nn.Embedding(2548, 64)` per person — no prior, and it reached only 0.585 on ancestry. +`--init raw`: the person's own carrier row (6,551 zeros/ones) through the first linear layer — the most +information but 6,551 × 64 weights in that layer. + +### 5.4 What the GNN adds to an embedding + +`embeddings.py` scores each space by silhouette (how tight and separated the ancestry groups are) and by 5-NN +balanced accuracy on test people. SVD-32: silhouette 0.06, 5-NN 0.90. GNN hidden layer (ancestry run): +**0.70 / 0.97**. That 0.06 → 0.70 is the GNN's contribution: it reorganises the space so that groups are +compact, which is what downstream retrieval (nearest neighbours for the decoder) and clustering need. + +--- + +## 6. The GNN, layer by layer + +Hidden size `d = 64`, two layers. + +1. **Projection**: for each node type `t`, `h_t⁽⁰⁾ = W_t x_t + b_t`, `W_t ∈ ℝ^{d×in_t}`. + `in_individual` = 32 (SVD) or 6,551 (raw), + 460 + 460 in v2 (protein z and observed mask); + `in_cluster` = 7 (+32 with SVD), `in_block` = 5, `in_gene` = `in_protein` = 2. +2. **Message passing** (`HeteroConv`, one operator per relation, results summed per target type): + - `SAGEConv` on `carries`, `rev_carries`, `in_block`, `rev_in_block`, `next_block`, `overlaps`, `encodes` (each with its reverse): + `h_v' = W₁ h_v + W₂ · mean_{u ∈ N_r(v)} h_u` — the node keeps its own state and adds the mean of its neighbours under relation `r`. + - `GraphConv` on `co_occurs` and `measured`/`rev_measured`: + `h_v' = W₁ h_v + W₂ · mean_{u} w_{uv} h_u` with `w = log(lift)/max log(lift)` for co-occurrence and `w = z` (the harmonised protein level, signed) for measured — a protein that is high in this person pushes with positive weight, one that is low with negative weight. + - Sum over relations, then `h ← Dropout(ReLU(LayerNorm(h' + h)))` (residual keeps the projection signal alive through both layers). +3. **Head**: `logits = W_out h_individual⁽²⁾ + b`, `W_out ∈ ℝ^{k×64}`. + +What two rounds mean for a person `i`: round 1 pulls in the clusters `i` carries (and, in v2, `i`'s +proteins); round 2 pulls in what those clusters co-occur with, their blocks, and what genes/proteins sit in +those blocks. So `h_i⁽²⁾` summarises "my haplotypes, their population context, and my proteome" in 64 numbers. + +**Loss**: cross-entropy over training people only, with class weights `w_k = N_train / (k · n_k)` so rare +classes (AMR, small populations) are not ignored. **Optimiser**: Adam, lr 5·10⁻³, weight decay 5·10⁻⁴. +**Early stopping**: validation balanced accuracy, patience 30, best weights restored. Full-batch: one +epoch = one forward/backward over the whole graph (2.4 M + 0.37 M + 1.07 M edges): 0.95 s on the M2 CPU, +0.10 s on the A100. + +**Metrics**: accuracy; balanced accuracy = mean per-class recall (robust to the 26-class imbalance); +macro-F1; ROC-AUC for the binary phenotype; all computed only on test people with a label. + +Parameters: ≈105 k with SVD input; ≈0.5 M with raw input (dominated by the 6,551×64 projection). + +--- + +## 7. The synthetic proteome: exact generative model + +Per person `i` (real 1000G id, real sex, random age 18–85, site assigned by shuffling within each ancestry +and splitting into 3): + +- Causal clusters: 20 kept clusters drawn from blocks that encode a protein, with carrier frequency 5–60 %. + `β_c ~ N(0, 1.5²)`. Logit `η_i = Σ_c β_c M[i,c] + 0.02·(age_i − mean age)`; intercept set at the + `1−0.35` quantile of `η`; `case_i ~ Bernoulli(σ(η_i − intercept))` → 38 % cases. +- Protein `p` for person `i`: `y_ip = b_p + β^{age}_p·age_z + β^{sex}_p·sex_i + β^{pheno}_p·case_i + Σ_{c: cis(c)=p} β^{cis}_c M[i,c] + s_{site(i),p} + ε_bio + ε_tech` + with `b_p ~ U(6,16)` log2, `β^{age}, β^{sex} ~ N(0, 0.5²)`, `β^{pheno}_p ~ N(0, 0.5²)` for a random 15 % of proteins and 0 otherwise (70 responsive proteins), `β^{cis}_c ~ N(0, 1²)` for the one protein in the causal cluster's block, site shift `s ~ N(0, 0.3²)`, `ε_bio ~ N(0, 0.8²)`, `ε_tech ~ N(0, 0.3²)`. +- Missingness: `P(missing_ip) = 0.15 · (b_max − b_p)/(b_max − b_min)` — the least abundant proteins are missing most often (MNAR at the detection limit). + +Everything is written to `ground_truth.json`, so a model's saliency can be scored against the causal +clusters and a regression against the cis proteins. + +**Harmoniser** (`haplokg_proteins.harmonise`): within each `(site, protein)`, +`z = (y − median) / (1.4826 · MAD)`. Robust to outliers, computed from each site's own samples only, removes +`s_{site,p}` exactly (between-site median shift 0.48 → 0.00 log2 in the EDA) while preserving within-site +biology (227 proteins keep a phenotype association at FDR 5 %). + +--- + +## 8. Integration experiments: what each row of the table is + +| run | person input | relations | AUC | +|---|---|---|---| +| genome / svd | `UΣ` (32) | genome relations | 0.635 | +| genome / raw | carrier row (6,551) | genome relations | 0.601 | +| proteome | `[z, mask]` (920) → 2-layer MLP, **no graph** | — | 0.962 | +| both / svd | `[UΣ, z, mask]` | all 12 relations | 0.993 | +| both / raw | `[row, z, mask]` | all 12 relations | 0.992 | +| site (both) | as both | all | bal. acc. 0.30 (chance 0.33) | +| ancestry (both) | as both | all | bal. acc. 0.90 | + +Why genome-only is low: the phenotype's genomic part is 20 clusters with noisy logistic link → the Bayes +rate is far from 1; why proteome-only is high: 70 proteins each shift by ~N(0,0.5) with noise sd ~0.85, and +the MLP sums the evidence; why both is higher still: the graph brings the genomic evidence in and the +`measured` edges let protein evidence propagate. The site row is the guarantee that none of this is batch. + +**Saliency**: `∂ logit_case / ∂ x_individual`, averaged over test cases, restricted to the carrier-row +coordinates, ranked; 3–4 of the top 20 are ground-truth causal (chance 20 × 20/6551 = 0.06). + +**Genome → proteome**: predicting all 460 z-scores from `h_i⁽²⁾` gives R² ≈ 0 even for cis proteins, +although their linear ceiling (`r²` between carrying the causal cluster and the protein's z) is 0.54 for the +strongest and 0.17 on average. A per-protein ridge on the carrier row (`Ridge(alpha=1000)`, multi-output) +recovers 4/20 cis proteins at test R² > 0.1 and 0/440 others. Interpretation: a 64-d embedding trained for a +classification target does not preserve single-cluster cis effects; that needs sparse per-protein models or +a pQTL edge prior (UKB-PPP). + +--- + +## 9. Decoder: retrieval → prompt → validated JSON + +Retrieval for person `i` (all from the graph, deterministic): profile (ancestry, population, sex, site, age; +true phenotype withheld), GNN prediction from `test_predictions.csv`, the globally salient clusters `i` +carries, the 8 carried clusters with the highest ancestry Cramér's V (each with block, carrier fractions by +ancestry, genes and proteins in the block), the 8 proteins with the largest |z| (with their encoding block +and whether that block holds one of the notable clusters), and the 5 nearest labelled people by Euclidean +distance in `h⁽²⁾`. Serialised as JSON (~4.5 k tokens). + +Model: `nvidia/nemotron-3-super-120b-a12b` via the OpenAI-compatible NIM endpoint, `temperature 0.2`, +`max_tokens 4000`, `reasoning_effort="none"`. Without that flag the model emits its chain of thought in the +content and hits the token cap before the JSON (measured: 1,200 tokens of reasoning, no answer); +with `"low"` the reasoning goes to a separate field and the answer still arrives, at 2–4× the latency. +The system prompt forbids new entities and requires ids verbatim; `cited_ids` are checked against the +set of ids present in the context (25 cited, 0 unknown for HG00103, 13 s). + +--- + +## 10. Federated learning with NVFlare: the mechanics + +**Partition.** Site `s` gets the people with `site_code == s` (835 / 835 / 833, mixed ancestry). Its graph is +`HeteroData.subgraph({'individual': members})`: the person nodes are re-indexed to that site, their CARRIES +and MEASURED edges kept, every other node type (cluster, block, gene, protein) and their edges kept whole — +those are public. Site 1, for example: 835 people, 774,753 CARRIES, 355,294 MEASURED edges. Person input += raw carrier row + protein z + mask — computed from the site's own rows only; no cross-site preprocessing +(the harmoniser is already per site). Cluster/block/gene/protein features are public statistics. + +**Model config.** `job.py` derives every constructor value from the public graph (input dims, the 12 +relations, hidden 64, 2 layers, dropout 0.3) into `model_args.json`; the server builds +`model.ProGenomeGNN(config_path)` from it and each client builds the same, so state-dict keys and shapes +match by construction. + +**Round loop (Client API, `client.py`).** `flare.init()` → `flare.get_site_name()` → build the site graph +once → `while flare.is_running(): m = flare.receive(); model.load_state_dict(m.params); evaluate the received +global model on the site's own val/test people; if it is an evaluate-only task, send metrics; else train 5 +full-batch epochs (5 optimizer steps) on the site's training people with local class weights; send +FLModel(params=state_dict on CPU, metrics, meta[NUM_STEPS_CURRENT_ROUND]=5)`. + +**Server.** `FedAvgRecipe` (NVFlare 2.9, `nvflare.app_opt.pt.recipes.fedavg`): after each round the global +weights are the step-weighted average `θ ← Σ_s n_s θ_s / Σ_s n_s` with `n_s = NUM_STEPS_CURRENT_ROUND` +(equal here — full-batch — so a plain mean; a real deployment would weight by local sample count via +`aggregation_weights`). `key_metric="val_balanced_accuracy"` selects `best_FL_global_model.pt`; +`FL_global_model.pt` is the final round. Tensor-native transport: `server_expected_format=PYTORCH` + +`TensorDecomposer`. `SimEnv(num_clients=3)` runs the three sites as threads on one machine; +`recipe.execute(env)` materialises the job under `outputs/federated/chr22/workspace/`. + +**What crosses the site boundary.** Per round per site: one state dict (~0.5 M floats) and five scalars. +No rows of `M`, no protein values, no embeddings of people. + +**Result (10 rounds × 5 local epochs, ~70 s on the M2).** Global model scored centrally on the same 376 +held-out people as the central model: AUC **0.955** vs central 0.992; per site (their own test people) +AUC 0.977 / 0.968 / 0.938. Per-round curves show the global model still improving at round 9, i.e. 50 local +steps is short of the central run's ~80 epochs. **30 rounds × 5 local epochs (~4 min on the M2): AUC 0.998 / +balanced accuracy 0.963 centrally on the same 376 people — equal to the central model (0.992 / 0.969); per site +AUC 1.000 / 0.999 / 0.997.** Federated training loses nothing here because the sites are i.i.d. draws of the same +population (mixed ancestry by construction); with ancestry-pure sites the averaging would have to fight client drift. + +--- + +## 11. Compute and engineering facts + +- Image: `pytorch/pytorch:2.14.0-cuda12.6-cudnn9-runtime` (system-managed Python 3.12 → `PIP_BREAK_SYSTEM_PACKAGES=1`), torch_geometric 2.8.0.post1, pyg-lib 0.9 (the only extension PyG still ships; `torch_scatter` has no wheel for torch 2.14 and PyG no longer needs it), nx-cugraph (NetworkX dispatch to cuGraph with `NX_CUGRAPH_AUTOCONFIG=True`), torch-tensorrt 2.14. +- GPU: A100 80 GB (Crusoe, $1.98/h) — image built natively in ~4 min; epoch 0.10 s; whole v1 pipeline < 2 min. L4 24 GB (GCP) also ran it. +- Inference: eager 36 ms per full graph; `torch.compile` (inductor) 4.9 ms, max |Δlogit| 1.4·10⁻⁶; Torch-TensorRT (`torch.compile(backend="torch_tensorrt")`) did not finish partitioning/compiling this scatter-heavy hetero-GNN in 3 h and is reported as not applicable. +- Graph display: Neo4j 5.26 community (2,548 + 6,551 + 669 nodes, 2.37 M + 187 k + … relationships loaded with batched `UNWIND ... MERGE` in ~2 min); NetworkX 3.6 statistics and GraphML export. +- Packaging: `make setup/run/run-v2/eda/decode/federated/docker/brev`, pinned `requirements.txt`, 11 unit tests on a toy graph, secrets only via environment. + +--- + +## 12. Caveats, stated plainly + +1. The phenotype is synthetic. The genome ↔ proteome ↔ phenotype numbers show the *pipeline* recovers a + planted signal; they are not biology. Ancestry/population/sex results are on real labels. +2. chr22 only; ancestry on one chromosome is almost linearly separable, so the GNN cannot beat logistic + regression on accuracy there — its value is the embedding space and the integration. +3. The GNN does not recover cis genotype→protein effects; the ridge does for the strong ones. +4. Federated evaluation is on the same held-out people as central; in a real deployment each site would + report its own held-out score and no central evaluation would exist. +5. TensorRT is not part of the inference claim; `torch.compile` is. diff --git a/genomics/docs/METHODS.md b/genomics/docs/METHODS.md new file mode 100644 index 0000000..5f5aba4 --- /dev/null +++ b/genomics/docs/METHODS.md @@ -0,0 +1,198 @@ +# ProGenome — what the goal is and what we built (technical account) + +Written for the team and the manuscript's methods paragraphs. Every number is from `genomics/outputs*/` +on chromosome 22, seed 42, unless stated otherwise. + +## 1. The goal, precisely + +The README asks three questions. Translated into things a program can do: + +| RQ | question | operational form | status | +|---|---|---|---| +| 1 | How can haploblock genomics connect to genes and proteomics in one graph model? | a typed graph in which one **person** node links to the haplotype **clusters** they carry, clusters sit in **blocks**, blocks overlap **genes**, genes encode **proteins**, and the same person links to measured **protein levels** | built (v2 schema) | +| 2 | Can a GNN combine genomic + proteomic information to identify phenotype groups? | train a heterogeneous GNN on that graph to predict a phenotype from the person's neighbourhood; show that genome + proteome beats either alone, and that the embedding separates groups | built; result on synthetic ground truth | +| 3 | Can it be trained across institutions without moving individual-level data? | split the person nodes by site, keep the shared reference graph identical everywhere, exchange only model weights (FedAvg) | designed; next step (NVFlare) | + +The mission sentence — "each institution retains its individual-level data locally and trains the same +graph-based model; only model updates are exchanged" — is a statement about *where node types live*. That is +why the whole design starts from making the person a separate node type. + +## 2. The data and where it comes from + +**HaploGraph (data.haploblocks.org, built at MDxCORE/Rigshospitalet, September 2026, for this hackathon).** +Upstream pipeline (haploblocks.org, Kubica et al. 2025): recombination-rate peaks define *haploblocks* +(recombination-defined regions); each 1000 Genomes individual's two phased haplotypes are extracted per block +from the GRCh38 phased VCF; haplotype sequences of a block are clustered with MMseqs2; each haplotype gets a +*cluster id*; a *hash* encodes strand/chromosome/block/cluster/variants. We did **not** re-run those steps — +we consumed their published outputs and cross-checked them against each other (669 blocks in the boundaries +file = 669 in block_stats = 669 in the node matrix; per-block cluster counts identical for all 669; +248,254 clusters both ways). + +Files consumed for chr22: + +- `nodes.csv.gz` — 248,254 clusters × 2,548 individuals, 0/1 = individual carries that cluster (on either haplotype). +- `edges_lift_above_threshold.csv.gz` — 187,030 cluster–cluster edges with `weight` (individuals carrying both) and + `lift = P(A∩B) / (P(A)·P(B))`, pre-filtered to lift ≥ 5 (the raw `edges.csv` is dominated by one + population-frequency "mega-hub", per the HaploGraph README). +- `block_stats.tsv` — per block: coordinates, length, number of clusters, largest-cluster share (dominance), + Shannon entropy, singleton count. +- `phenotypes_real.csv` — ancestry (5 super-populations), population (26), sex, for 2,503 of the 2,548 people. + **1000G has no other phenotype.** Height (whiteboard) therefore had to become a simulated label. +- `proteomics/uniprot_chr22.bed` — 917 UniProt isoform rows → 460 proteins, 458 gene symbols, with genomic spans. +- Proteomics per site: today synthetic (below); planned real sources: Wu et al. 2013 (TMT proteomics on 95 HapMap + LCLs, 53 CEU/33 YRI/9 EAS, ids are 1000G ids) and UKB-PPP cis-pQTL summary statistics (AWS Open Data). + +Why chr22: smallest autosome with a complete HaploGraph (9 MB of graph), so the whole loop runs on a laptop in +minutes; nothing in the code is chromosome-specific (`CHROM=chr21 make run`). + +## 3. Knowledge-graph construction (`build_kg.py`, `haplokg.py`) + +1. Stream `nodes.csv.gz` in 8,192-row chunks as `int8`, convert each chunk to a CSR sparse matrix, stack → + clusters × individuals (1.3 GB CSV → 2.8 M non-zeros in memory). +2. **Cluster filter**: keep clusters with `25 ≤ carriers ≤ N−25` (symmetric, like the HaploGraph edge filter). + 248,254 → 6,551 clusters; 176,903 singletons and the near-universal clusters carry no population signal. + The filter reproduces the edge file's node set exactly (0 of 187,030 edges dropped). +3. Transpose → `carries` = individuals × kept clusters (2,548 × 6,551, 2,365,574 ones; ~928 per person = two + haplotypes × 669 blocks minus filtered clusters). +4. Join phenotypes on `individual_id`; encode labels as integers, −1 for the 45 unlabelled people. +5. Map edge endpoints to cluster indices; blocks sorted by position → `NEXT_BLOCK` edges. +6. PyG `HeteroData`: node types `individual`, `cluster` (7 z-scored features: log support, support fraction, + block log-length, entropy, dominance, log n_clusters, singleton rate), `block` (5 features); edge types + `carries`/`rev_carries`, `in_block`/`rev_in_block`, `co_occurs` (both directions, `edge_attr` = [weight, lift]), + `next_block` (both directions). Labels `y_ancestry`, `y_population`, `y_sex` on `individual`. + +**Where the phenotype connects:** the phenotype file and the node matrix share the sample id (`HG00096` is a +column header in one and a row key in the other). Phenotypes become node *properties*, never neighbours — if +the label were a neighbour, the GNN would read it directly (that is how a 100 % accuracy is produced by mistake). + +## 4. Statistics before any model (`cooccurrence_analysis.py`, `eda.py`) + +- For every cluster, a 2 × k contingency test of carrier-status vs ancestry / population / sex; Cramér's V = + √(χ²/N) for a 2 × k table; Benjamini–Hochberg FDR. Result: 6,470/6,551 clusters (98.8 %) are ancestry-associated + at FDR 5 % (median V 0.20, max 0.81); 6,378 population-associated; **0 sex-associated** (median V 0.013). + The sex result is the negative control: chr22 is autosomal, so a method that finds sex signal is broken. +- For every lift edge, the cosine similarity of the two endpoints' ancestry-enrichment profiles (5-vector of + carrier-fraction / overall-fraction − 1): 0.86 for real edges vs 0.22 for degree-preserving shuffled pairs; + 86 % of edges join clusters enriched in the same ancestry (42 % expected); higher-lift quartiles are more + similar (0.84 → 0.90). Interpretation: the co-occurrence graph is largely population structure — long-range + edges (median endpoint distance in the tens of Mb) are ancestry, not physical linkage. +- Graph statistics (`graph_explore.py`, NetworkX): 6,551 nodes, 187,030 edges, largest connected component + 4,327, 2,207 clusters without a lift edge, degree max 358; the top hubs are all AFR-enriched low-support + clusters — the mega-hub artefact per node. +- EDA (`outputs/eda/chr22/EDA.md`): provenance table, population/sex tables, block length (median 29.7 kb, + 5–95 % 8.9–151 kb), clusters per block (median 214), singleton rate (median 0.62), entropy along the + chromosome, support distributions, edge lift/distance/degree, gene/protein counts, proteomics missingness + (7.4 % overall, rising to 16 % for the least abundant proteins — MNAR at the detection limit), between-site + shift before/after harmonisation, and the cis genotype→protein correlations. + +## 5. Baseline (`baseline.py`) + +L2 logistic regression on the sparse 0/1 carrier matrix, one model per target, C chosen on validation, +scored on the untouched test split (stratified 70/15/15 over the 2,503 labelled people; the split file is +shared by every later model). Test balanced accuracy: ancestry 0.977, population 0.614 (26 classes), sex 0.463 +(chance). This sets the bar the GNN must at least match. + +## 6. Embeddings (`train_gnn.py --init`, `embeddings.py`) + +Nodes need a starting vector. Options implemented and compared: + +- **SVD-32** (default): truncated SVD of the carrier matrix `M ≈ UΣVᵀ`; rows of `UΣ` embed people, rows of + `VΣ` embed clusters, in one shared space. Unsupervised → cannot leak labels. +- **Node2Vec** on the individual–cluster + cluster–cluster graph (pyg-lib random walks; needs the GPU image). +- **learned**: a free `nn.Embedding` per person. +- **raw**: the person's own 6,551-long carrier row through a linear layer. + +Finding: with free learned embeddings the GNN reaches only 0.585 balanced accuracy on ancestry; with SVD +initialisation 0.974. The starting embedding matters more than the architecture. `embeddings.py` measures +each space by silhouette (by ancestry) and 5-NN accuracy: SVD-32 0.06 / 0.90; the ancestry-GNN's hidden layer +**0.70 / 0.97** — the GNN's real contribution is a much better-organised embedding, not higher accuracy. + +## 7. The GNN (`train_gnn.py`, `train_gnn_v2.py`) + +Per node type a linear projection to hidden size 64, then two rounds of `HeteroConv` (sum over relations): + +- `SAGEConv` (mean or sum aggregation) on `carries`, `rev_carries`, `in_block`, `rev_in_block`, `next_block`, + and in v2 `overlaps`, `encodes` (both directions); +- `GraphConv` with edge weights on `co_occurs` (weight = normalised log lift) and in v2 on `measured` / + `rev_measured` (weight = harmonised protein z-score). + +After each round: LayerNorm, residual connection, ReLU, dropout 0.3. A linear head maps the person's 64-d +vector to class logits. Loss: class-weighted cross-entropy on training people only; Adam 5e-3, weight decay +5e-4; early stopping on validation balanced accuracy (patience 30). Full-batch: one epoch is one pass over +the whole graph (2.4 M carries edges) — 0.95 s on an M2 CPU, 0.10 s on the A100. + +What a round of message passing means here: a person's vector becomes a summary of the clusters they carry; +a cluster's vector becomes a summary of its carriers, of the clusters it co-occurs with (weighted by lift) and +of its block; after two rounds a person "sees" the clusters that co-occur with theirs and the block context — +and in v2 the proteins they express and the genes those proteins come from. + +Results (test, balanced accuracy): ancestry 0.974 (baseline 0.977), population 0.437 with SVD-32 input → +0.611 with the raw carrier row (baseline 0.614), sex 0.503 (chance). Honest reading: on chr22 alone these +labels are essentially linear in the carrier matrix, so the GNN matches but does not beat logistic regression; +it wins on the embedding space. + +## 8. Schema v2: joining proteomics (`proteomics_synth_1000g.py`, `haplokg_proteins.py`, `build_kg_v2.py`) + +Synthetic proteomics keyed to the **real 1000G ids** (2,503 people, 460 chr22 proteins), with a saved ground +truth so recovery can be scored: 3 sites assigned at random *within* each ancestry (mixed-ancestry sites → +`site` is a pure batch label); age, sex (real), a binary phenotype whose logit is a weighted sum over 20 +"causal" clusters (+ small age term), calibrated to 38 % cases; each causal cluster also shifts one protein +encoded in its block (cis effect, β ~ N(0, 1)); ~15 % of proteins respond to the phenotype (β ~ N(0, 0.5)); +age and sex effects; per-site batch shift (sd 0.3 log2); missingness increasing toward the detection limit +(15 % for the least abundant protein). The first version shifted *all* proteins with the phenotype, which made +every model score 1.0 — the same failure mode as an earlier team result; that is why the signal is now sparse. + +Graph additions: `Gene` and `Protein` nodes from the UniProt BED (isoforms collapsed); `Block —OVERLAPS→ Gene` +by coordinate intersection (1,063 edges; 29 genes fall outside any block, 117 span two); `Gene —ENCODES→ Protein`; +`Individual —MEASURED{log2, z}→ Protein` (1,065,712 edges). **Harmoniser**: robust z per (site, protein), +`(x − median)/(1.4826·MAD)`, computed over each site's own samples before anything crosses sites; it removes the +between-site median shift entirely (0.48 → 0.00) while 227 proteins keep a phenotype association at FDR 5 %. +Detection-limit missingness stays missing — no edge — never imputed as zero. + +## 9. Integration experiments (`train_gnn_v2.py`, A100) + +Same split, same test people, held-out AUC for the synthetic phenotype: + +| modality | input to the person node | AUC | bal. acc. | +|---|---|---|---| +| genome | SVD-32 or raw carrier row; genome relations only | 0.60–0.63 | 0.57–0.62 | +| proteome | harmonised z (460) + observed mask (460); MLP, no graph | 0.96 | 0.96 | +| **both** | SVD/raw + z + mask; all relations incl. `measured`, `overlaps`, `encodes` | **0.99** | **0.97** | +| site (batch control, both) | — | — | 0.30 (chance 0.33) | +| ancestry (both) | — | — | 0.90 | + +So the graph adds information on top of the proteome. Saliency (gradient of the case logit w.r.t. the raw +carrier row, averaged over test cases) puts 3–4 of the 20 ground-truth causal clusters in its top 20 +(chance 1.2). Genome → proteome as a regression from the person's embedding fails (R² ≈ 0), while a +per-protein ridge on the carrier row recovers the strongest cis effects (4/20 cis proteins at test R² > 0.1, +0/440 others; ceiling from the data itself: r² up to 0.54). Conclusion for the paper: the GNN is the +integration/embedding tool; cis discovery wants sparse per-protein models or a pQTL edge prior. + +## 10. Decoder (`graphrag_decoder.py`) + +For one person, deterministic retrieval from the graph: profile (true phenotype withheld), the GNN prediction, +globally salient clusters they carry, their 8 most ancestry-informative clusters with block/genes/proteins, +their 8 most extreme protein levels with the encoding block and whether it holds a notable cluster, and their +5 nearest neighbours in the GNN embedding. Serialised as JSON (~4.5 k tokens) to an OpenAI-compatible NIM +endpoint — `nvidia/nemotron-3-super-120b-a12b`, `reasoning_effort=none` (without it the reasoning model thinks +inline and exhausts the token budget before the JSON) — under a system prompt that forbids inventing entities +and requires ids to be cited verbatim; the reply's `cited_ids` are checked against the context (25 cited, 0 +unknown for HG00103). Output: summary, ancestry and phenotype assessments, genome↔proteome links, caveats. + +## 11. Compute, packaging, deployment + +CUDA image `pytorch/pytorch:2.14.0-cuda12.6-cudnn9-runtime` + torch_geometric 2.8.0 + pyg-lib + nx-cugraph + +torch-tensorrt 2.14 (the PyG docs' deprecation of torch_cluster in favour of pyg-lib was the cause of the first +build failure). Runs on CPU too. Neo4j 5.26 community via docker compose for browsing the graph. Brev: +`brev_deploy.sh` creates/uses an instance, uploads code + data, builds natively, runs, copies outputs back; +executed on an L4 (GCP) and an A100 80 GB (Crusoe). Inference: 36 ms per full graph eager → 4.9 ms with +`torch.compile` (identical logits); Torch-TensorRT compilation of this scatter-heavy hetero-GNN did not +finish in 3 h and is not claimed. Everything is `make`-driven from a clone; secrets via environment only. + +## 12. What the federated step (RQ3) will do + +Sites = the 3 mixed-ancestry sites; each holds its Individual nodes with their CARRIES and MEASURED edges; +the cluster/block/gene/protein graph and the encoder weights are shared. NVFlare FedAvg: each round every +site trains `train_gnn_v2.py`'s model on its subgraph for a few epochs and sends weights; the server averages +by sample count and returns the global model. Report: central vs federated on the same held-out people, plus +the `site` control. Exchanging embeddings of people would be a leak; exchanging weights is not. diff --git a/genomics/docs/architecture.html b/genomics/docs/architecture.html new file mode 100644 index 0000000..362c89e --- /dev/null +++ b/genomics/docs/architecture.html @@ -0,0 +1,330 @@ + +ProGenome Architecture + + + +
+
+
ProGenome · Team #3 · Nordic Biobank × NVIDIA hackathon · branch modelling
+

ProGenome Architecture

+

One knowledge graph joins the public 1000 Genomes haploblock graph to per-site proteomics on the sample ID; a PyTorch Geometric GNN encodes it, an LLM decodes it, and only model weights ever cross a site boundary.

+
+ chr22 demo2,548 individuals · 6,551 clusters · 669 blocks + 2.37 M CARRIES · 187 k CO_OCCURS458 genes · 460 proteins · 1.07 M MEASURED + data: data.haploblocks.org +
+
+ +
+

Data flow: two lanes into one graph, one encoder, one decoder

+
+ + + + + + + + SOURCESGRAPH BUILD · SCHEMAEMBEDDINGS + GNN · PyTorch GeometricINFERENCEDECODER + + + + + + + + + + + HaploGraph · 1000G · GRCh38 + nodes.csv.gz 248,254 × 2,548 + edges (lift ≥ 5) 187,030 + block_stats · boundaries + phenotypes_real.csv (3 labels) + + build_kg.py → hetero.pt + pandas · scipy CSR int8 stream + Individual · Cluster · Block + CARRIES · CO_OCCURS · IN_BLOCK + NEXT_BLOCK · support ≥ 25 + + SVD-32 + scipy svds · shared space + Node2Vec + pyg-lib random walks · GPU + + + + + + + Neo4j 5.26 community + neo4j_load.py · :7474 + NetworkX · nx-cugraph + stats · GraphML · plots + + + + + + + Proteomics per site · 1000G IDs + uniprot_chr22.bed 460 proteins + site{1,2,3}_proteomics_log2.csv + synthetic + ground truth today + Wu 2013 · UKB-PPP pQTL next + + build_kg_v2.py → hetero_v2.pt + harmoniser: robust z / site + Gene · Protein · OVERLAPS + ENCODES · MEASURED {log2, z} + LOD missing → no edge + + + + + + + fetch_data.sh + measured_long + + + + joins on sample ID (HG00096 …) + + + + + + + carries.npz + + + + + + + + genome edges + cluster / block stats + + + measured · overlaps · encodes + z vector + + + + + + 32-dim init + + + + HeteroConv × 2 · hidden 64 + LayerNorm + residual · dropout 0.3 + SAGEConv (mean or sum) + carries · in_block · next_block + overlaps · encodes + GraphConv, edge-weighted + co_occurs · w = log lift + measured · w = harmonised z + individual input + SVD-32 | raw carriers (6,551) + + protein z & mask (460 each) + heads + phenotype: CE, class-weighted + proteome: masked MSE + controls: sex (autosome), site (batch) + Adam 5e-3 · early stop on val · seed 42 + + + + infer.py · CUDA + eager 36 ms + torch.compile 4.9 ms + Torch-TensorRT fp16 + per full graph · A100 80 GB + 7.5× · logits identical + + model.pt + + + + GraphRAG decoder + KG neighbourhood + + GNN logits + embeds + → NVIDIA NIM LLM + Nemotron / Llama endpoint + stretch: PyG G-Retriever + + + + + outputs/ per run + metrics.json · test_predictions.csv · history.csv + embedding_individual.npy · embedding_cluster.npy + + + + + RUNTIME + docker pytorch/pytorch:2.14.0-cuda12.6-cudnn9-runtime · torch_geometric 2.8.0 · pyg-lib · nx-cugraph · torch-tensorrt 2.14 + compose neo4j:5.26-community + pipeline · Brev A100 80 GB (Crusoe, $1.98/h) · L4 24 GB (GCP, $0.85/h) · Mac M2 CPU ~10 min + make setup | run | docker | docker-run | neo4j-load | brev · Makefile · setup.sh · run_all.sh · brev_deploy.sh · pytest + +
+
genome laneproteome laneshared / neutral
+
The two source lanes meet in build_kg_v2.py on the 1000 Genomes sample ID: the same string is a column header in nodes.csv.gz and a row key in phenotypes_real.csv and in the proteomics matrices. Phenotypes stay node properties (training targets), never neighbours, so the GNN cannot read the label off the graph.
+
+ +
+

Federated topology: what moves, what never leaves

+
+ + + + + + + + + weights only, after each local round — rows never leave + + + + Site 1 (private)Site 2 (private)Site 3 (private) + Individual nodes · labelsIndividual nodes · labelsIndividual nodes · labels + CARRIES edges (genotype)CARRIES edges (genotype)CARRIES edges (genotype) + MEASURED edges (proteome)MEASURED edges (proteome)MEASURED edges (proteome) + local training round on its subgraphlocal training round on its subgraphlocal training round on its subgraph + + + Shared, public: Cluster · Block · Gene · Protein graph + the GNN encoder weights + identical at every site (data.haploblocks.org); a new site maps its haplotypes to these clusters by BLAST against shared representatives + + + NVFlare server · FedAvg + average weights (by n) + send global model back + no rows, no embeddings of people + control: site stays at chance + + model + +
+
site-privateshared, public
+
Because individuals are their own node type, the site boundary is a clean cut through the graph: everything in the teal band is public and identical everywhere; everything inside the dashed boxes stays put. Mixed-ancestry sites make site a pure batch label, so a model that predicts site is learning batch, not biology.
+
+ +
+
+

Tech stack

+ + + + + + + + + + + + + + + +
LayerChoiceStatus
Graph datadata.haploblocks.org HaploGraph 1000G (39,141 blocks genome-wide; chr22 here)done
Graph buildpandas 3.0 · scipy 1.18 CSR int8 streaming · PyG HeteroDatadone
Schema v2Gene/Protein nodes from uniprot_chr22.bed; MEASURED with per-site robust-z harmoniserdone
Graph displayNeo4j 5.26 community (Docker, localhost:7474) · NetworkX 3.6 · nx-cugraph on GPUdone
EmbeddingsSVD-32 (scipy) · Node2Vec (pyg-lib) · GNN hidden layer 64-ddone
Modeltorch 2.14 · torch_geometric 2.8.0 · HeteroConv(SAGEConv, GraphConv) · class-weighted CE · masked MSEdone
Training computeBrev A100 80 GB (Crusoe) · L4 24 GB (GCP) · Mac M2 CPU for devdone
Inferencetorch.compile inductor 7.5× · Torch-TensorRT 2.14 fp16 (dense parts; PyTorch fallback for scatter)done
DecoderGraphRAG prompt → NVIDIA NIM (Nemotron / Llama) · stretch: PyG G-Retriever soft-prompt LLMnext
FederationNVFlare 2.9 Recipe API, FedAvg over 3 mixed-ancestry sitesnext
PackagingDockerfile (CUDA 12.6 base) · docker compose · Makefile · setup.sh · brev_deploy.sh · pinned requirements · pytestdone
+
+
+

Schema (chr22 counts)

+ + + + + + + + + + + + + +
Node / edgeFromCount
Individual {ancestry, population, sex, site, phenotype}phenotypes_real.csv + sample metadata2,548
Cluster {support, block stats, SVD}nodes.csv.gz, support ≥ 256,551
Block {length, entropy, dominance}block_stats.tsv669
Gene · Proteinuniprot_chr22.bed (isoforms collapsed)458 · 460
Individual —CARRIES→ Clusternodes.csv.gz columns2,365,574
Cluster —CO_OCCURS{weight, lift}→ Clusteredges_lift_above_threshold187,030
Cluster —IN_BLOCK→ Block · Block —NEXT_BLOCK→ Blockids · coordinates6,551 · 668
Block —OVERLAPS→ Gene · Gene —ENCODES→ Proteincoordinate intersection1,063 · 460
Individual —MEASURED{log2, z}→ Proteinproteomics matrices, harmonised1,065,123
+
+
+ +
+

Measured, not assumed

+
+
GNN epoch0.10 sA100 · vs 0.95 s on the M2 CPU (full batch, 2.4 M edges)
+
Inference, full graph4.9 mstorch.compile · 36 ms eager · max logit diff 1.4e-6
+
Ancestry (held-out)0.97balanced accuracy · logistic baseline 0.98 · sex control 0.50
+
Phenotype: proteome → + genome0.96 → 0.99AUC, held-out · genome alone 0.60 · synthetic ground truth
+
Site (batch) control0.303 sites, chance 0.33 · harmoniser removes the injected shift
+
Graph ↔ phenotype98.8 %of clusters ancestry-associated (FDR 5 %) · 0 % sex
+
+
+ +
Everything runs from a clone: make setup && make run on a laptop, make docker for the CUDA image, make neo4j-load to browse the graph, make brev to build and train on an NVIDIA Brev GPU and copy the outputs back. Numbers are chr22, seed 42, 70/15/15 stratified split; synthetic proteomics carries a saved ground truth so recovery can be scored.
+
diff --git a/genomics/docs/architecture.mmd b/genomics/docs/architecture.mmd new file mode 100644 index 0000000..3941e2a --- /dev/null +++ b/genomics/docs/architecture.mmd @@ -0,0 +1,67 @@ +%% ProGenome architecture — Mermaid source. +%% Lucidchart: Insert → Diagram as code → Mermaid → paste. GitHub renders it in a ```mermaid fence. +flowchart LR + subgraph SRC["Sources"] + HG["HaploGraph · 1000G · GRCh38
nodes.csv.gz 248,254 × 2,548
edges lift ≥ 5 · 187,030
block_stats · phenotypes_real.csv"] + PR["Proteomics per site · 1000G IDs
uniprot_chr22.bed · 460 proteins
site1..3_proteomics_log2.csv
synthetic + ground truth today · Wu 2013 / UKB-PPP next"] + end + + subgraph BUILD["Graph build · schema"] + KG1["build_kg.py → hetero.pt
pandas · scipy CSR int8
Individual · Cluster · Block
CARRIES 2.37M · CO_OCCURS 187k · IN_BLOCK · NEXT_BLOCK"] + KG2["build_kg_v2.py → hetero_v2.pt
harmoniser: robust z per site × protein
Gene 458 · Protein 460
OVERLAPS 1,063 · ENCODES 460 · MEASURED 1.07M"] + NEO["Neo4j 5.26 community
neo4j_load.py → localhost:7474"] + NX["NetworkX 3.6 / nx-cugraph
stats · GraphML · region plots"] + end + + subgraph EMB["Embeddings"] + SVD["SVD-32
scipy svds on carriers
shared individual/cluster space"] + N2V["Node2Vec
pyg-lib random walks · GPU"] + end + + subgraph GNN["GNN · PyTorch Geometric 2.8 · CUDA"] + ENC["HeteroConv × 2 · hidden 64
SAGEConv: carries · in_block · next_block · overlaps · encodes
GraphConv weighted: co_occurs (log lift) · measured (z)
input: SVD-32 or raw carrier row + protein z & mask
heads: phenotype CE · proteome masked MSE
controls: sex (autosome) · site (batch)"] + end + + subgraph INF["Inference"] + IE["infer.py
eager 36 ms → torch.compile 4.9 ms (7.5×)
Torch-TensorRT fp16 · A100 80 GB"] + end + + subgraph DEC["Decoder"] + RAG["GraphRAG
KG neighbourhood + GNN logits/embeddings
→ NVIDIA NIM LLM (Nemotron 3 Super)
stretch: PyG G-Retriever"] + OUT["Insights per patient
cited to cluster / block / protein ids"] + end + + HG -- "fetch_data.sh (md5)" --> KG1 + PR -- "measured_long.csv" --> KG2 + KG1 -- "joins on sample ID (HG00096…)" --> KG2 + KG1 -- "display before embedding" --> NEO + KG1 --> NX + KG1 -- "carries.npz" --> SVD + KG1 -- "carries + co_occurs" --> N2V + SVD -- "32-dim init" --> ENC + N2V -- "32-dim init" --> ENC + KG1 -- "genome edges + cluster/block stats" --> ENC + KG2 -- "measured · overlaps · encodes + z vector" --> ENC + ENC -- "model.pt · embeddings" --> IE + IE -- "logits + 64-d embeddings" --> RAG + RAG --> OUT + + subgraph RUN["Runtime"] + DK["Docker: pytorch/pytorch:2.14.0-cuda12.6-cudnn9-runtime
torch_geometric 2.8.0 · pyg-lib · nx-cugraph · torch-tensorrt 2.14
compose: neo4j:5.26-community + pipeline"] + HW["Brev A100 80 GB (Crusoe, $1.98/h) · L4 24 GB (GCP)
Mac M2 CPU for dev (~10 min end-to-end)
make setup | run | docker | neo4j-load | brev"] + end + + subgraph FED["Federated (NVFlare FedAvg)"] + S1["Site 1 (private)
Individual · CARRIES · MEASURED"] + S2["Site 2 (private)"] + S3["Site 3 (private)"] + SRV["NVFlare server
averages weights by n
returns global model"] + SHARED["Shared, public: Cluster · Block · Gene · Protein graph + encoder weights"] + end + S1 -- "weights only" --> SRV + S2 -- "weights only" --> SRV + S3 -- "weights only" --> SRV + SRV -. "global model" .-> S1 + SHARED --- S1 + SHARED --- S2 + SHARED --- S3 diff --git a/genomics/docs/report/ProGenome_KT.docx b/genomics/docs/report/ProGenome_KT.docx new file mode 100644 index 0000000..cabe49e Binary files /dev/null and b/genomics/docs/report/ProGenome_KT.docx differ diff --git a/genomics/docs/report/ProGenome_KT.pdf b/genomics/docs/report/ProGenome_KT.pdf new file mode 100644 index 0000000..ed3d76a Binary files /dev/null and b/genomics/docs/report/ProGenome_KT.pdf differ diff --git a/genomics/docs/report/ProGenome_KT.tex b/genomics/docs/report/ProGenome_KT.tex new file mode 100644 index 0000000..8b07218 --- /dev/null +++ b/genomics/docs/report/ProGenome_KT.tex @@ -0,0 +1,879 @@ +\documentclass[11pt,a4paper]{article} +\usepackage[margin=2.2cm]{geometry} +\usepackage[T1]{fontenc} +\usepackage[utf8]{inputenc} +\usepackage{lmodern} +\usepackage{graphicx} +\usepackage{booktabs} +\usepackage{longtable} +\usepackage{array} +\usepackage{enumitem} +\usepackage{hyperref} +\usepackage{xcolor} +\hypersetup{colorlinks=true, linkcolor=blue!50!black, urlcolor=blue!50!black} +\graphicspath{{figures/}} +\setlength{\parskip}{4pt} +\title{ProGenome: a federated workflow for genome-graph and proteomic integration\\[6pt]\large Complete knowledge-transfer document: problem, biology, data, pipeline, knowledge graph, models, federated learning, deployment and results} +\author{Team \#3, Nordic Biobank x NVIDIA Federated Learning Hackathon, Copenhagen, September 2026. Branch: modelling. Author of this build: Koushik Telaprolu (genomics/ pipeline), with the team's proteomics and README work referenced where used.} +\date{18 September 2026} +\begin{document} +\maketitle +\tableofcontents +\newpage + +\section{How to read this document} + +This is written for a teammate who joins today and knows nothing about the biology, the software, the models or the infrastructure. Part I is a primer that defines every concept used later. Part II is the whole workflow on a few pages: what goes in and what comes out of every stage, where each data source is used, what the trained model produces and how it is used, ground truth against prediction, where the language model sits, and how every item of the team README on branch main is covered. Part III is the project in depth, in the order the data flows: problem, data sources, data pipeline, exploratory analysis, knowledge graph, embeddings, the graph neural network (the encoder), the language-model decoder, the federated training, deployment, results and caveats. Part IV is how to run everything, a glossary and a repository map. Every number is measured on chromosome 22 with random seed 42 and comes from files under genomics/outputs; nothing is typed in from memory. + +\section{Part I: Primer} + +\subsection{Biology in ten minutes} + +DNA is a long text written in four letters (A, C, G, T). Humans have about 3 billion letters, packaged in 23 pairs of chromosomes; chromosome 22 is one of the smallest, about 50 million letters. Everyone carries two copies of each chromosome, one from each parent. A gene is a region of DNA that encodes a protein; proteins are the molecules that do the work in cells, and measuring how much of each protein a person has is called proteomics. + +Two people's DNA differs at roughly one letter in a thousand; those positions are variants. A haplotype is the specific combination of variants along one copy of a chromosome. Because DNA is inherited in chunks (recombination cuts and rejoins the parental copies at a limited number of places), neighbouring variants tend to travel together. A haploblock is a stretch of chromosome between recombination hotspots that is usually inherited as one unit. Phased data means we know, for every variant, which of the two copies it sits on, so each person's two haplotypes per block are known separately. + +The 1000 Genomes Project (1000G) sequenced 2,548 people from 26 populations grouped into five continental ancestries: AFR (African), AMR (admixed American), EAS (East Asian), EUR (European) and SAS (South Asian). It recorded only ancestry, population and sex about them, nothing clinical. A phenotype is any observable property of a person; in this project the real phenotypes are those three, and a clinical-looking case/control phenotype had to be simulated. + +\subsection{Software in ten minutes} + +Python is the language of every script here; pandas handles tables, scipy handles sparse matrices (mostly-zero grids stored compactly), PyTorch does neural networks and PyTorch Geometric (PyG) adds graph neural networks. Git tracks versions of the code; a branch is a parallel line of work (ours is called modelling). Docker packages the code and every library it needs into an image so it runs identically on any machine. A GPU is a processor built for many small parallel calculations; CUDA is NVIDIA's software that lets PyTorch use it. NVIDIA Brev rents GPU machines by the hour. Neo4j is a database made for graphs, with a browser to look at them. + +\subsection{Data science in ten minutes} + +Exploratory data analysis (EDA) means measuring and plotting the data before modelling: sizes, distributions, missing values, obvious structure. To judge a model honestly the people are split once into train (learn), validation (choose settings and when to stop) and test (report only once, at the end); here 70/15/15 percent, chosen at random but stratified so each ancestry is represented in every part. Accuracy is the fraction correct; balanced accuracy averages the accuracy per class so a rare class cannot be ignored; AUC (area under the ROC curve) is the probability that a random case is scored higher than a random control, 0.5 is guessing and 1.0 is perfect. A negative control is a target that must come out at chance level if the method is sound. + +\subsection{Machine learning in ten minutes} + +An embedding is a list of numbers (a vector) that represents an object so that similar objects get similar vectors. A graph is a set of nodes joined by edges; a knowledge graph is a graph whose nodes and edges have types and properties. A graph neural network (GNN) computes a vector for every node by repeatedly mixing each node's own vector with those of its neighbours (message passing); after two rounds a node's vector summarises its two-hop neighbourhood. Training means adjusting the network's weights so that a prediction made from those vectors matches known labels, measured by a loss; an optimiser (Adam) nudges the weights to reduce the loss, one epoch being one pass over the data. An encoder turns raw data into vectors; a decoder turns vectors (and retrieved facts) into an output, here text written by a large language model (LLM). Retrieval-augmented generation (RAG) gives the LLM the facts it should use so it does not invent them; GraphRAG retrieves those facts from a graph. + +\subsection{Federated learning in five minutes} + +Hospitals cannot pool patient data. Federated learning trains one model across sites without moving records: each site trains the shared model on its own data for a few steps, sends only the updated weights to a coordinator, the coordinator averages them (FedAvg) and sends the average back; repeat for several rounds. NVIDIA FLARE (NVFlare) is the framework that runs this loop; its simulator runs all sites on one machine so the workflow can be tested before real deployment. + +\section{Part II: The workflow on a few pages} + +\subsection{The story, stage by stage} + +Read the workflow as one sentence first: a public graph of haplotype clusters is joined to people by their sample id; each person's row of that graph, their phenotype labels and their protein measurements are attached to a person node; a graph neural network is trained to predict a label from the person's neighbourhood; the trained network yields predictions, an embedding per person and cluster, a saliency per cluster and a set of weights; the predictions and embeddings are checked against ground truth, a language model turns them into a cited report per person, and the weights are what a federated deployment exchanges instead of data. Every stage below is one script with one Makefile target. + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.98\linewidth]{data_flow_map.png} +\caption{Where each data source enters, what the trained model produces, and where each product goes.} +\end{figure} + +\begin{longtable}{p{0.132\linewidth}p{0.282\linewidth}p{0.357\linewidth}p{0.169\linewidth}} +\toprule +\textbf{Stage} & \textbf{What goes in} & \textbf{What comes out} & \textbf{Lands in} \\ +\midrule +\endhead +1 download & URLs on data.haploblocks.org & HaploGraph node and edge files, phenotypes, block statistics, md5-verified & data/ \\ +2 knowledge graph v1 & nodes.csv.gz (who carries which cluster), edges (which clusters co-occur), block statistics, phenotypes & Individual, Cluster, Block nodes; CARRIES, CO\_OCCURS, IN\_BLOCK, NEXT\_BLOCK edges; labels stored on the person nodes & outputs/kg/chr22/ (carries.npz, hetero.pt) \\ +3 statistics & the graph and the labels & for every cluster and edge, how strongly it tracks ancestry, population and sex (Cramer's V, FDR) & outputs/cooccurrence/chr22/ \\ +4 baseline and split & carrier matrix and labels & logistic-regression scores to beat; the single train/val/test split every later model reuses & outputs/baseline/, outputs/splits/ \\ +5 starting embeddings & carrier matrix only, no labels & a 32-number vector per person and per cluster (SVD), or Node2Vec, learned, or the raw row & computed inside training \\ +6 GNN v1 (encoder) & graph, starting vectors, and the labels of training people only & trained weights; class probabilities per person; a 64-number embedding per person and per cluster & outputs/gnn/chr22/\_/ \\ +7 embedding check & SVD and GNN embeddings, labels of test people & silhouette and nearest-neighbour accuracy, 2-D plots & outputs/embeddings/chr22/ \\ +v2.1 synthetic proteomics & the real 1000G ids, the carrier matrix, the gene BED & 3 per-site protein matrices, metadata (site, age, sex, case/control), ground\_truth.json & outputs/proteomics\_synth/chr22/ \\ +v2.2 knowledge graph v2 & graph v1, gene BED, proteomics & Gene and Protein nodes; OVERLAPS, ENCODES, MEASURED edges (harmonised z); site and phenotype labels on people & outputs/kg/chr22/hetero\_v2.pt \\ +v2.3 EDA & everything above & the eight-section report with tables and plots & outputs/eda/chr22/EDA.md \\ +v2.4 GNN v2 & graph v2; modality genome / proteome / both; target phenotype / site / ancestry / sex / proteome & metrics, test predictions, embeddings, saliency per cluster, training history & outputs/gnn\_v2/chr22/\_\_/ \\ +v2.5 ridge & carrier matrix, protein z-scores & per-protein test R-squared, split into cis and other proteins & outputs/gnn\_v2/chr22/proteome\_ridge\_baseline/ \\ +v2.6 decoder & one person: graph neighbourhood, GNN prediction and embedding neighbours, saliency & a cited JSON insight written by the NIM language model, with a citation check & outputs/graphrag/chr22/\_insight.json \\ +v2.7 federated & graph v2 split by site, the model definition & a global model trained without moving any person's data; its score on the same held-out people & outputs/federated/chr22/ \\ +inference & a trained run & predictions and embeddings for all 2,548 people; eager versus compiled timing & outputs/gnn/chr22//inference/ \\ +\bottomrule +\end{longtable} + +\subsection{Where each data source is used, and where it is not} + +The most common confusion is what the phenotypes and the proteomics do. The phenotypes are labels: they are stored on the person node as training targets and evaluation ground truth and never become an edge, a node feature or an input to the starting embeddings. The proteomics enters three times: as MEASURED edges from a person to the proteins observed in them, as part of the person's input vector, and, through its metadata, as two more labels (site and the synthetic case/control phenotype). The HaploGraph provides the structure everyone shares. + +\begin{longtable}{p{0.226\linewidth}p{0.470\linewidth}p{0.244\linewidth}} +\toprule +\textbf{Source} & \textbf{Used for} & \textbf{Never used for} \\ +\midrule +\endhead +HaploGraph nodes.csv.gz & CARRIES edges; the carrier matrix behind the SVD starting vectors, the logistic baseline, the raw input, the saliency and the ridge & labels \\ +HaploGraph edges\_lift\_above\_threshold.csv.gz & CO\_OCCURS edges with weight = normalised log lift & anything about people directly \\ +block\_stats.tsv, boundaries & Block node features; five of the seven Cluster features; NEXT\_BLOCK order; block-gene overlaps & labels \\ +phenotypes\_real.csv (ancestry, population, sex) & labels on Individual nodes: the training target for training people, the ground truth for validation and test people, the classes in the Cramer's V tests, the colours in plots & edges, node features, the SVD (which is label-free), the test split (only its stratification) \\ +uniprot\_chr22.bed & Gene and Protein nodes with span and isoform count; OVERLAPS by coordinate; ENCODES & labels \\ +proteomics matrices (per site, log2) & harmonised per site and protein into z; MEASURED edge weights; the person's input vector (z plus observed mask); the proteome-only MLP & cross-site normalisation (the harmoniser never sees two sites at once) \\ +proteomics metadata (site, age, sex, phenotype) & site = batch control and the federated partition; phenotype = the case/control target; age = decoder context & edges or node features \\ +ground\_truth.json & scoring only: which clusters are causal (saliency precision), which proteins are cis-affected (ridge R-squared) & any model input \\ +\bottomrule +\end{longtable} + +\subsection{What the trained model gives you, and what is done with it} + +A trained run is a folder with model.pt (the weights), metrics.json, history.csv, test\_predictions.csv, embedding\_individual.npy (and embedding\_cluster.npy in v1), and for the raw-input phenotype run saliency\_top100.csv. Five products come out of the encoder and each has a consumer: + +\begin{itemize}[leftmargin=1.4em] + \item Class probabilities per person: for every one of the 2,548 people, a probability per class (five ancestries, 26 populations, case/control). Consumed by the metrics against held-out labels, by infer.py which writes predictions\_all\_individuals.csv with a confidence per person, and by the decoder, which quotes the prediction for the person it describes. + \item A 64-number embedding per person and per cluster: the encoder's last hidden layer. Consumed by embeddings.py (silhouette, nearest-neighbour accuracy, PCA plots), by the decoder (the five nearest people in this space are part of the retrieved context) and available for any downstream clustering of people into phenotype groups. + \item A saliency per haploblock cluster: the gradient of the case score with respect to the person's carrier row, averaged over test cases. Consumed by the ground-truth check (how many of the top 20 are planted causal clusters) and by the decoder, which lists the globally salient clusters the person carries. + \item The weights themselves: in the federated setting the weights are the only thing that leaves a site, so the same model definition (federated/model.py) is what NVFlare averages. + \item A timed inference path: infer.py rebuilds the run's inputs, scores the whole graph in one pass (36 ms eager, 4.9 ms with torch.compile on the A100) and writes predictions and embeddings for everyone. +\end{itemize} + +\subsection{Ground truth versus prediction} + +Every model is scored on the same 376 held-out people that no model saw during training or model selection. For the real labels the ground truth is the 1000G panel; for the synthetic phenotype the ground truth is the label the generator drew from the planted causal clusters; for the saliency and the ridge it is the list of causal clusters and cis proteins in ground\_truth.json; for the federated model it is the same 376 people scored centrally. The confusion tables below are read row = truth, column = prediction. + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.95\linewidth]{confusion_matrices.png} +\caption{Ground truth versus prediction on the 376 held-out people: real ancestry (GNN, SVD input), the synthetic phenotype from genome plus proteome, from the genome alone, and sex, the negative control.} +\end{figure} + +\begin{longtable}{p{0.329\linewidth}p{0.122\linewidth}p{0.122\linewidth}p{0.122\linewidth}p{0.122\linewidth}p{0.122\linewidth}} +\toprule +\textbf{Ancestry, true / predicted} & \textbf{AFR} & \textbf{AMR} & \textbf{EAS} & \textbf{EUR} & \textbf{SAS} \\ +\midrule +\endhead +AFR (99) & 99 & 0 & 0 & 0 & 0 \\ +AMR (52) & 1 & 48 & 0 & 3 & 0 \\ +EAS (76) & 0 & 0 & 76 & 0 & 0 \\ +EUR (76) & 0 & 4 & 0 & 72 & 0 \\ +SAS (73) & 0 & 0 & 0 & 0 & 73 \\ +\bottomrule +\end{longtable} + +368 of 376 correct (balanced accuracy 0.974). The eight errors are all between AMR and EUR, which is what admixture predicts: the AMR panel populations carry European haplotypes. + +\begin{longtable}{p{0.432\linewidth}p{0.122\linewidth}p{0.122\linewidth}p{0.263\linewidth}} +\toprule +\textbf{Synthetic phenotype, true / predicted} & \textbf{case} & \textbf{control} & \textbf{correct} \\ +\midrule +\endhead +genome plus proteome (graph): case (136) & 128 & 8 & 367 / 376, AUC 0.992 \\ +genome plus proteome (graph): control (240) & 1 & 239 & \\ +proteome only (MLP): case (136) & 129 & 7 & 364 / 376, AUC 0.962 \\ +proteome only (MLP): control (240) & 5 & 235 & \\ +genome only (graph): case (136) & 95 & 41 & 202 / 376, AUC 0.601 \\ +genome only (graph): control (240) & 133 & 107 & \\ +sex, negative control: female (188) & 119 & 69 & 189 / 376, chance \\ +sex, negative control: male (188) & 118 & 70 & \\ +\bottomrule +\end{longtable} + +The genome alone gets little more than half right because the planted genomic signal is 20 clusters through a noisy logistic link; the proteome carries most of the signal; the graph that joins the two makes the fewest errors, with a single false case. Sex on an autosome is a coin flip, as it must be. + +\begin{longtable}{p{0.075\linewidth}p{0.470\linewidth}p{0.132\linewidth}p{0.263\linewidth}} +\toprule +\textbf{Rank} & \textbf{Cluster (saliency of the combined model)} & \textbf{Saliency} & \textbf{Planted causal cluster?} \\ +\midrule +\endhead +1 & chr22\_46902935-46974137\_cluster186 & 0.110 & yes \\ +2 & chr22\_44263132-44288496\_cluster19 & 0.080 & yes \\ +3 & chr22\_25066667-25206817\_cluster346 & 0.077 & yes \\ +4 & chr22\_46974137-47051453\_cluster43 & 0.055 & no \\ +5 & chr22\_44166469-44185056\_cluster231 & 0.054 & yes \\ +6 & chr22\_47257518-47285160\_cluster1 & 0.046 & no \\ +7 & chr22\_40032702-40132216\_cluster25 & 0.046 & no \\ +8 & chr22\_49408121-49430759\_cluster3 & 0.043 & no \\ +9-20 & twelve further clusters & 0.042-0.036 & no \\ +\bottomrule +\end{longtable} + +Four of the top twenty (and the top three outright) are among the 20 planted causal clusters out of 6,551; by chance 0.06 would be. For the genome-to-protein direction the ridge finds 4 of the 20 cis proteins with test R-squared above 0.1 and none of the 440 others. The federated global model, scored on the same 376 people, reaches AUC 0.998 and balanced accuracy 0.963 against the central model's 0.992 and 0.969. + +\subsection{Where the language model sits, and what it achieves} + +The language model is not part of training and makes no prediction. It sits after the encoder, once per person, as the decoder: graphrag\_decoder.py walks the graph around one person and collects facts (their ancestry-informative clusters with blocks, genes and proteins; their most extreme protein levels and the block that encodes each; the GNN's prediction; the salient clusters they carry; their five nearest neighbours in the embedding), serialises them as JSON and asks the NIM model (nvidia/nemotron-3-super-120b-a12b) to write a structured report using only those facts and citing every id verbatim. The code then checks each cited id against the context. What this achieves is the last step the README's mission implies, turning numbers into an insight a research team can read: a summary, an ancestry assessment, a phenotype assessment, explicit genome-to-proteome links and caveats, each traceable to graph ids. For HG00103 (EUR, GBR, 994 clusters carried, predicted control) the model answered in 13 seconds from a 4,555-token context with a 1,020-token reply, cited 25 ids and invented none; its links named, for example, cluster chr22\_40032702-40132216\_cluster151 in the block encoding TNRC6B (protein Q9UPQ9) and cluster chr22\_26024448-26060666\_cluster35 in the block encoding MYO18B (Q8IUG5), and its caveats stated that the phenotype is synthetic and that ancestry is population structure, not a medical finding. Without an API key the same script prints the exact prompt (dry run), so the retrieval can be inspected offline. + +\begin{longtable}{p{0.329\linewidth}p{0.611\linewidth}} +\toprule +\textbf{Retrieved fact for HG00103} & \textbf{Value} \\ +\midrule +\endhead +GNN prediction & control (truth: control) \\ +globally salient clusters carried & chr22\_46902935-46974137\_cluster186 (0.110), chr22\_46974137-47051453\_cluster43 (0.055), chr22\_40032702-40132216\_cluster25 (0.046) \\ +most ancestry-informative clusters, with block genes and proteins & cluster151 of block 40032702-40132216 (TNRC6B, Q9UPQ9); cluster545 of 35738379-36060254 (RBFOX2, O43251); cluster35 of 26024448-26060666 (MYO18B, Q8IUG5) \\ +most extreme protein levels (z) & H1-0 +2.75, RIMBP3 -2.65, INPP5J +2.53, CRYBA4 +2.52, YWHAH +2.37 \\ +nearest people in the embedding & HG01767 (EUR), NA12842 (EUR), HG00351 (EUR), HG04002 (SAS), HG00336 (EUR) \\ +\bottomrule +\end{longtable} + +\begin{longtable}{p{0.207\linewidth}p{0.733\linewidth}} +\toprule +\textbf{Decoder field} & \textbf{Text returned (Nemotron 3 Super, 13.3 s, 25 ids cited, 0 unknown)} \\ +\midrule +\endhead +summary & Individual HG00103 is of EUR ancestry (GBR population) and carries 994 haploblock clusters. The GNN phenotype prediction for this individual is 'control', which is a synthetic case/control label used for pipeline testing. Several ancestry-informative clusters are carried, showing high enrichment in various populations including EUR, EAS, and SAS. \\ +ancestry\_assessment & The individual's ancestry is EUR (European) as per the 1000G panel data, with population label GBR. This reflects population structure, not a medical finding. \\ +phenotype\_assessment & The GNN-predicted phenotype is 'control', which is a synthetic case/control label used to test the pipeline; it does not reflect a real medical condition. \\ +genome\_proteome\_links & cluster151 $\rightarrow$ block 40032702-40132216 $\rightarrow$ TNRC6B $\rightarrow$ Q9UPQ9 (carried, enriched in EAS); cluster545 $\rightarrow$ RBFOX2 $\rightarrow$ O43251; cluster35 $\rightarrow$ MYO18B $\rightarrow$ Q8IUG5 (enriched in EUR) \\ +caveats & ancestry is population structure, not a medical finding; the phenotype is synthetic; not all clusters have gene or protein annotations; the prediction comes from the embedding and may not match observed protein levels \\ +\bottomrule +\end{longtable} + +What the LLM adds is the sentence layer: a readable paragraph, the genome-to-protein chain spelled out per person, the right hedges attached automatically, graceful handling of missing information (for HG00096, who is not in the held-out split, it reported that no prediction exists instead of inventing one), and traceability, because every claim carries an id that the code verifies. It makes no prediction and adds no outside knowledge. + +\subsection{Connecting the dots to the README on branch main} + +The team README fixes a mission, three research questions, a chromosome-22 demo scope, five required datasets and a data-integration flowchart. Each maps to something concrete in genomics/: + +\begin{longtable}{p{0.282\linewidth}p{0.423\linewidth}p{0.235\linewidth}} +\toprule +\textbf{README item} & \textbf{What it became} & \textbf{Evidence} \\ +\midrule +\endhead +Mission: each institution keeps individual-level data locally, trains the same graph model, exchanges only model updates; a server aggregates and returns them & the person is its own node type, so a site holds only its people and their CARRIES and MEASURED edges while the cluster/block/gene/protein graph is public; NVFlare FedAvg with client.py sending a state dict only & outputs/federated/chr22/evaluation.json: AUC 0.998 vs central 0.992 \\ +Opening line: a variant-based phenotype-propensity reference graph combined with patient-specific proteomics, what can we learn? & reference graph = the shared cluster/block/gene/protein layer; patient-specific = Individual nodes with CARRIES and MEASURED edges; learned: integration beats either modality alone, cis effects need sparse models, federation costs nothing when sites are alike & sections 9, 11, 13 \\ +Background: one gene gives many protein products & 917 UniProt isoform rows collapsed to 460 proteins with an isoform count as a Protein feature; ENCODES keeps gene to protein explicit & haplokg\_proteins.load\_protein\_bed \\ +RQ1: connect haploblock genomics to genes and proteomic data in a graph model & schema v2, five node types, seven edge types, one join key; the README flowchart maps one-to-one: Participant = Individual, Haploblock hash = Cluster, Haploblock = Block, Encoded protein = Gene ENCODES Protein, Measured abundance = MEASURED edge & section 6; Neo4j browser \\ +RQ2: can a GNN combine genomic and proteomic information to identify disease-related phenotype clusters & yes, in both senses of cluster: groups of people (AUC 0.60 genome, 0.96 proteome, 0.99 both; embedding silhouette 0.70) and haploblock clusters tied to the phenotype (saliency top 3 all causal); the person-level graph, with every participant as a node, is what lets the GNN use both modalities & sections 8, 9; Part II ground truth tables \\ +RQ3 (aspirational): train across institutions without transferring individual-level data & done in NVFlare simulation with three sites; real-machine POC is the next step & section 11 \\ +Demo scope: chromosome 22 first & everything runs on chr22; nothing is chromosome-specific (CHROM=chr21 make run) & Makefile \\ +Dataset 1: haploblock BED and per-individual haploblock hashes & boundaries and block\_stats from haploblocks.org; the hashes were already clustered upstream into the HaploGraph clusters we consume, so no hash was recomputed & fetch\_data.sh \\ +Dataset 3: gene BED for chr22 & uniprot\_chr22.bed (team); Block OVERLAPS Gene by coordinate intersection, 1,063 edges, 29 genes outside every block & build\_kg\_v2.py \\ +Dataset 4: gene-to-protein mapping & the same BED: UniProt accession per gene, isoforms collapsed; 460 ENCODES edges & build\_kg\_v2.py \\ +Dataset 5: proteomic data for chr22 proteins & the team's 120-patient synthetic set validated the proteomics plumbing but cannot join the genome (its ids are not 1000G ids and its phenotype has no genomic cause); the joinable set on 2,503 real 1000G ids with ground truth replaces it for integration; real data (Wu 2013, UKB-PPP) is the planned next source & proteomics\_synth\_1000g.py \\ +Data-integration flowchart & implemented as a PyTorch Geometric HeteroData object plus a Neo4j load, with two additions the README did not list: CO\_OCCURS between clusters and NEXT\_BLOCK between blocks & hetero\_v2.pt; neo4j\_load.py \\ +\bottomrule +\end{longtable} + +\section{Part III: The project} + +\subsection{1. Problem statement and goal} + +The team's README (branch main) states the mission: develop a proof-of-concept workflow that integrates a known genome graph with proteomics, with a federated approach where each participating institution retains its individual-level data locally and trains the same graph-based model, exchanging only model updates with a coordinating server. Three research questions make it concrete: + +\begin{itemize}[leftmargin=1.4em] + \item RQ1: How can haploblock-based genomic information be connected to genes and proteomic data in a graph-based data model? + \item RQ2: Can a graph neural network combine genomic and proteomic information to identify disease-related phenotype clusters? + \item RQ3 (aspirational): Can a graph neural network trained across multiple institutions predict clinical outcomes without transferring individual-level data? +\end{itemize} + +Translated into engineering: build a typed graph in which a person links to the haplotype clusters they carry, clusters sit in blocks, blocks overlap genes, genes encode proteins, and the same person links to their measured protein levels (RQ1); train a GNN on it and show that genome plus proteome predicts a phenotype better than either alone while negative controls stay at chance (RQ2); split the people by site, share only the reference graph and the model weights, and show the federated model matches the central one (RQ3). The demo scope fixed by the team is chromosome 22. + +What was already available: the mentor (Ben Busby) pointed the team to haploblocks.org, whose data server publishes a ready-made graph of haplotype clusters for the 1000 Genomes people, built for this hackathon at Rigshospitalet's MDxCORE unit. That is what "we don't have to build the genome graph" meant; everything downstream of it is ours. + +\subsection{2. Data sources: what, where, why} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.95\linewidth]{genome_to_graph.png} +\caption{From genomes to a graph: blocks, phased haplotypes, MMseqs2 clusters, the carrier matrix and the filters that produce the 6,551 kept clusters.} +\end{figure} + +\begin{longtable}{p{0.188\linewidth}p{0.207\linewidth}p{0.282\linewidth}p{0.263\linewidth}} +\toprule +\textbf{Data} & \textbf{Where it comes from} & \textbf{What it contains} & \textbf{Why we use it} \\ +\midrule +\endhead +HaploGraph nodes (nodes.csv.gz) & data.haploblocks.org/haplograph/1000G/chr22 & 248,254 haplotype clusters x 2,548 people; 1 if the person carries the cluster & The genome graph's node features; becomes our person-to-cluster edges \\ +HaploGraph edges (edges\_lift\_above\_threshold.csv.gz) & same server & 187,030 cluster pairs that co-occur in people more than chance (lift $\geq$ 5), with weight and lift & Co-occurrence structure; the raw edges.csv is dominated by one near-universal haplotype and is not used \\ +Block statistics (block\_stats.tsv, boundaries) & same server & 669 blocks on chr22: coordinates, length, number of clusters, entropy, dominance, singletons & Block nodes and their features \\ +Phenotypes (phenotypes\_real.csv) & 1000G panel via IGSR, republished with the graph & ancestry, population, sex for 2,503 of the 2,548 people & The only real labels that exist; ancestry and population are targets, sex is the negative control \\ +Protein coordinates (uniprot\_chr22.bed) & UCSC UniProt track, prepared by Friederike & 917 protein isoform rows $\rightarrow$ 460 proteins, 458 genes, with genomic spans & Block-to-gene overlaps and gene-to-protein mapping in one file \\ +Synthetic proteomics, team version & proteomics/generate\_synthetic\_proteomics.py (Nolan; regenerated on main with 4,000 samples) & 3 sites, first 40 then 1,333-1,334 patients each (ids SITE1\_PT0001...), age, sex, case/control, log2 intensities for 460 proteins & The team's proteomics analyses on main run on it; for the genome join a version keyed to 1000G ids was needed (next row) \\ +Synthetic proteomics, joinable version & genomics/proteomics\_synth\_1000g.py (this work) & 2,503 real 1000G ids, 3 mixed-ancestry sites, phenotype driven by 20 causal clusters, cis effects, batch shift, missingness; ground truth saved & Lets the integration be scored against a known answer \\ +Planned real proteomics & Wu et al. 2013 (Nature): 95 HapMap LCLs with 1000G ids; UKB-PPP pQTL summary statistics on AWS Open Data & per-person protein levels; variant-to-protein effect sizes & Real join on the same ids; real cluster-to-protein propensity edges \\ +\bottomrule +\end{longtable} + +How the genome graph was made upstream (haploblocks.org pipeline, Kubica et al. 2025): (1) recombination-rate peaks define haploblocks; (2) each person's two phased haplotype sequences are cut out per block from the 1000G VCF; (3) all haplotypes of a block are merged into one FASTA; (4) MMseqs2 clusters near-identical haplotypes, giving each haplotype a cluster id; (5) a compact hash encodes strand, chromosome, block, cluster and variants. A sixth step (haploblock-graph-builder) produced the HaploGraph: nodes are clusters, the node feature is the 0/1 vector over people, edges join clusters that co-occur in the same people. We did not rerun these steps; we verified their outputs agree (669 blocks in every file, identical per-block cluster counts in all 669, 248,254 clusters both ways) and consumed them. + +\subsection{3. The data pipeline} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.95\linewidth]{workflow_pipeline.png} +\caption{The two pipeline chains (v1 genome graph, v2 proteomics integration), one script per stage.} +\end{figure} + +Everything is a numbered script under genomics/, driven by a Makefile so a fresh clone runs with two commands. Downloaded inputs and generated outputs are git-ignored and regenerated; secrets are read from environment variables and never stored in the repository. + +\begin{longtable}{p{0.066\linewidth}p{0.235\linewidth}p{0.423\linewidth}p{0.216\linewidth}} +\toprule +\textbf{Step} & \textbf{Script} & \textbf{What it does} & \textbf{Output} \\ +\midrule +\endhead +1 & fetch\_data.sh & downloads the chr22 HaploGraph files, phenotypes, block statistics; verifies md5 checksums & data/ \\ +2 & build\_kg.py (haplokg.py) & streams nodes.csv.gz as int8 in 8,192-row chunks into a sparse matrix; filters clusters; joins phenotypes; maps edges; writes tables and the PyG graph & outputs/kg/chr22/ (carries.npz, hetero.pt) \\ +3 & cooccurrence\_analysis.py & cluster x phenotype tests, edge x phenotype similarity, per-block informativeness, plots & outputs/cooccurrence/chr22/ \\ +4 & baseline.py & logistic regression on the carrier matrix; writes the shared train/val/test split & outputs/baseline/, outputs/splits/ \\ +5 & graph\_explore.py & NetworkX statistics, GraphML export, region and chromosome-wide plots & outputs/graph/chr22/ \\ +6 & train\_gnn.py & the genome-only GNN with SVD / Node2Vec / learned / raw inputs & outputs/gnn/chr22/ \\ +7 & embeddings.py & quality of SVD and GNN embeddings: silhouette, nearest-neighbour accuracy, PCA plots & outputs/embeddings/chr22/ \\ +v2.1 & proteomics\_synth\_1000g.py & synthetic proteomics on the real 1000G ids with saved ground truth & outputs/proteomics\_synth/chr22/ \\ +v2.2 & build\_kg\_v2.py (haplokg\_proteins.py) & adds genes, proteins, block-gene overlaps, harmonised measurements & outputs/kg/chr22/hetero\_v2.pt \\ +v2.3 & eda.py & the full exploratory report with tables and plots & outputs/eda/chr22/EDA.md \\ +v2.4 & train\_gnn\_v2.py & genome / proteome / both ablations, controls, saliency, proteome regression & outputs/gnn\_v2/chr22/ \\ +v2.5 & proteome\_linear\_baseline.py & per-protein ridge: can the genome predict each protein? & outputs/gnn\_v2/chr22/proteome\_ridge\_baseline/ \\ +v2.6 & graphrag\_decoder.py & graph retrieval + NVIDIA NIM LLM $\rightarrow$ cited insight per person & outputs/graphrag/chr22/ \\ +v2.7 & federated/job.py, client.py, model.py, evaluate\_global.py & NVFlare FedAvg over 3 sites and central scoring of the global model & outputs/federated/chr22/ \\ +- & infer.py & inference benchmark: eager vs torch.compile vs TensorRT & outputs/gnn/.../inference/ \\ +- & neo4j\_load.py, docker-compose.yml & loads the graph into Neo4j for browsing & localhost:7474 \\ +\bottomrule +\end{longtable} + +Memory arithmetic that shaped step 2: the node file is 1.3 GB as text; as a dense 64-bit matrix it would be 5 GB, as dense int8 632 MB, as a sparse matrix with 2.8 million non-zeros about 30 MB. Streaming chunks into sparse form keeps the whole build under 1 GB and 26 seconds on a laptop. + +The cluster filter: a cluster is kept if between 25 and N-25 people carry it (N = 2,548). A cluster carried by 2,540 people is as uninformative as one carried by 8; both have only 8 people on the informative side. This mirrors the HaploGraph's own symmetric edge filter and reproduces its node set exactly: 248,254 clusters become 6,551 and not one of the 187,030 edges loses an endpoint. 176,903 of the dropped clusters are singletons (one haplotype). + +\subsection{4. Exploratory data analysis} + +The EDA report (outputs/eda/chr22/EDA.md, generated by eda.py) has eight sections: provenance, individuals, haploblocks, clusters, co-occurrence, genes and proteins, proteomics, genome-proteome. The key measurements: + +\begin{longtable}{p{0.188\linewidth}p{0.752\linewidth}} +\toprule +\textbf{Aspect} & \textbf{Measurement} \\ +\midrule +\endhead +Individuals & 2,548 people; 2,503 labelled (AFR 660, EAS 504, EUR 503, SAS 489, AMR 347), 45 unlabelled kept as nodes; 26 populations of 61-113 people; sex balanced within each ancestry; split train 1,752 / val 375 / test 376 \\ +Blocks & 669 tiling 17.1-50.2 Mb with no gaps; length median 29.7 kb (5-95\%: 8.9-151 kb, longest 775 kb); clusters per block median 214, max 2,806; singleton rate median 0.62; entropy median 3.15; longer blocks hold more clusters (Spearman 0.37) \\ +Clusters & 6,551 kept of 248,254; each person carries \textasciitilde{}928 (two haplotypes x 669 blocks minus filtered); the person x cluster matrix is 14\% dense \\ +Co-occurrence edges & 187,030 with lift $\geq$ 5 (median 5.5, max 76); only 121 join clusters of the same block; median distance between endpoints tens of Mb; 4,344 clusters have at least one edge; degree up to 358; the top hubs are all AFR-enriched rare clusters (the mega-hub artefact) \\ +Genes / proteins & 458 genes, 460 proteins, 1,063 block-gene overlaps; 29 genes outside every block (chromosome ends and gaps); 117 genes span two blocks; up to 14 genes in one block \\ +Proteomics (synthetic) & 2,503 people x 460 proteins, 3 sites of \textasciitilde{}835; log2 range 4.9-17.0; 7.4\% missing overall, up to 16\% for the least abundant proteins (missing-not-at-random at the detection limit); between-site shift 0.48 log2 before harmonisation, 0.00 after; 227 proteins keep a phenotype association at FDR 5\% \\ +Genome-proteome & for the 20 ground-truth causal clusters, the correlation between carrying the cluster and its cis protein gives r-squared up to 0.54, mean 0.17, 10 of 20 above 0.1 - the ceiling any genome-to-proteome model can reach on this data \\ +\bottomrule +\end{longtable} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.90\linewidth]{eda_populations.png} +\caption{1000G individuals per population, coloured by continental ancestry.} +\end{figure} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.95\linewidth]{eda_blocks.png} +\caption{Haploblocks on chr22: length distribution, clusters per block versus length, Shannon entropy of clusters along the chromosome.} +\end{figure} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.90\linewidth]{eda_clusters.png} +\caption{Kept clusters: carriers per cluster (log scale) and clusters carried per person by ancestry.} +\end{figure} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.95\linewidth]{eda_cooccurrence.png} +\caption{Co-occurrence edges: lift, distance between endpoints, degree.} +\end{figure} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.95\linewidth]{eda_proteomics.png} +\caption{Synthetic proteomics: intensity distribution, missingness rising for low-abundance proteins, batch effect between sites before and after the harmoniser.} +\end{figure} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.95\linewidth]{eda_blocks_populations.png} +\caption{Blocks with one dominant haplotype versus many rare ones (left); sex within each of the 26 populations (right).} +\end{figure} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.95\linewidth]{sites_composition.png} +\caption{The three federated sites: people per site by ancestry (mixed by design), case prevalence per site, missing protein values per site.} +\end{figure} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.95\linewidth]{sites_batch_effect.png} +\caption{Per-site protein medians before and after the harmoniser: the batch offsets vanish.} +\end{figure} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.95\linewidth]{ground_truth_effects.png} +\caption{Ground truth of the synthetic proteome: causal-cluster effects on the phenotype, cis effects on proteins, phenotype effects on the 70 responsive proteins, carrier frequencies of the causal clusters.} +\end{figure} + +\subsection{5. Statistics: does the graph co-occur with phenotypes?} + +Before any model we asked whether the graph carries phenotype information at all. For every cluster and each label we built the 2 x k table of carrier status versus class and computed a chi-square test with Cramer's V (a 0-1 effect size; for a 2 x k table V = sqrt(chi-square / N)), then Benjamini-Hochberg false-discovery-rate correction. All 6,551 tests run in one sparse matrix product. + +\begin{itemize}[leftmargin=1.4em] + \item Ancestry: 6,470 of 6,551 clusters (98.8\%) are associated at FDR 5\%; median V 0.20, maximum 0.81. + \item Population: 6,378 clusters (97.4\%). + \item Sex: 0 clusters; median V 0.013, maximum 0.03. This is the negative control: chr22 is autosomal, so a method that found sex signal would be fitting noise. +\end{itemize} + +For the edges we compared the two endpoint clusters' ancestry-enrichment profiles: cosine similarity 0.86 for real edges versus 0.22 for degree-preserving shuffled pairs; 86\% of edges join clusters enriched in the same ancestry (42\% expected), and the similarity rises with lift (0.84 in the lowest lift quartile to 0.90 in the highest). Interpretation: the co-occurrence graph is largely population structure, long-range co-inheritance within ancestries rather than physical linkage. + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.85\linewidth]{cramers_v.png} +\caption{How strongly each cluster tracks a phenotype: ancestry and population carry signal, sex (the control) does not.} +\end{figure} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.95\linewidth]{informativeness.png} +\caption{Maximum Cramer's V per block along chr22 for ancestry versus sex.} +\end{figure} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.85\linewidth]{edge_similarity.png} +\caption{Co-occurring clusters share ancestry profiles: real lift edges versus shuffled pairs.} +\end{figure} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.60\linewidth]{informative_clusters_heatmap.png} +\caption{The 30 most ancestry-informative clusters and the fraction of each ancestry that carries them.} +\end{figure} + +\subsection{6. The knowledge graph} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.95\linewidth]{schema_diagram.png} +\caption{Node and edge types of the knowledge graph with chr22 counts.} +\end{figure} + +The HaploGraph has a single node type (cluster) and carries people only as a feature vector. We turned that vector into a second node type, the person, because a person is what phenotypes and protein measurements belong to, and what a hospital owns. The join is the 1000G sample id: the same string (for example HG00096) is a column header in nodes.csv.gz, a row key in phenotypes\_real.csv and a column in the proteomics matrices. + +\begin{longtable}{p{0.141\linewidth}p{0.376\linewidth}p{0.113\linewidth}p{0.310\linewidth}} +\toprule +\textbf{Node type} & \textbf{Properties} & \textbf{Count} & \textbf{Source} \\ +\midrule +\endhead +Individual & ancestry, population, sex, site, phenotype, age & 2,548 & nodes.csv.gz columns + phenotypes + proteomics metadata \\ +Cluster & support, block statistics, SVD vector & 6,551 & nodes.csv.gz rows after the filter \\ +Block & length, n\_clusters, entropy, dominance, singleton rate & 669 & block\_stats.tsv \\ +Gene & coordinates, number of proteins & 458 & uniprot\_chr22.bed \\ +Protein & coordinates, number of isoforms & 460 & uniprot\_chr22.bed \\ +\bottomrule +\end{longtable} + +\begin{longtable}{p{0.329\linewidth}p{0.470\linewidth}p{0.141\linewidth}} +\toprule +\textbf{Edge type} & \textbf{Meaning} & \textbf{Count} \\ +\midrule +\endhead +Individual -CARRIES$\rightarrow$ Cluster & the person carries this haplotype cluster (the 1s of the matrix) & 2,365,574 \\ +Cluster -CO\_OCCURS\{weight, lift\}$\rightarrow$ Cluster & the two clusters co-occur in people more than chance & 187,030 \\ +Cluster -IN\_BLOCK$\rightarrow$ Block; Block -NEXT\_BLOCK$\rightarrow$ Block & position on the chromosome & 6,551; 668 \\ +Block -OVERLAPS$\rightarrow$ Gene; Gene -ENCODES$\rightarrow$ Protein & coordinate intersection; protein product & 1,063; 460 \\ +Individual -MEASURED\{log2, z\}$\rightarrow$ Protein & one edge per observed protein level; missing values create no edge & 1,065,712 \\ +\bottomrule +\end{longtable} + +Two design rules matter. First, phenotypes are properties of the person node, never nodes or edges: if a Phenotype node were connected to the person, a two-layer GNN would read the label from its neighbour and report a meaningless 100\%. Second, because people are their own node type, the site boundary is a clean cut: the cluster, block, gene and protein graph is public and identical at every site, while a site holds only its people and their CARRIES and MEASURED edges. + +The graph is stored as a PyTorch Geometric HeteroData object (hetero.pt, hetero\_v2.pt) for modelling and loaded into Neo4j community edition (docker compose up neo4j; neo4j\_load.py) for browsing at localhost:7474, where a query such as MATCH (i:Individual \{id:'HG00096'\})-[:CARRIES]$\rightarrow$(c:Cluster)-[:IN\_BLOCK]$\rightarrow$(b:Block) RETURN i,c,b shows a person's haplotypes with their blocks. NetworkX (with the nx-cugraph GPU backend in the image) produces statistics and GraphML for Gephi. + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.80\linewidth]{graph_region.png} +\caption{One region of chr22 (the densest published island): clusters as nodes coloured by the ancestry they are enriched in, co-occurrence edges weighted by lift.} +\end{figure} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.70\linewidth]{graph_edge_positions.png} +\caption{All 187,030 co-occurrence edges plotted by the positions of their two endpoints: block structure and long-range population structure.} +\end{figure} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.85\linewidth]{person_neighbourhood.png} +\caption{A real person's neighbourhood in the graph (HG00103), as the decoder retrieves it: clusters, blocks, genes, proteins, extreme protein levels and nearest neighbours in the embedding.} +\end{figure} + +\subsection{7. From graph to embeddings} + +A GNN needs a starting vector for every node. Cluster and block nodes use their statistics (support, block length, entropy, dominance, and so on), standardised. People and clusters together get a truncated singular value decomposition of the carrier matrix M (2,548 x 6,551): M is approximated as U S V-transpose with 32 components; the rows of U S are 32-number vectors for people and the rows of V S are 32-number vectors for clusters, in one shared space, so a person sits near the clusters they carry and near people with similar haplotypes. No labels are used, so nothing can leak into the test set. Alternatives implemented and compared: Node2Vec random-walk embeddings on the person-cluster graph (GPU, via pyg-lib; with 50 pretraining epochs it reaches 0.950 on ancestry against 0.974 for SVD, with the same 0.70 silhouette), free learned embeddings (0.585), and the raw 6,551-long carrier row. + +Where the phenotypes enter: only as training targets. The GNN is trained to predict the label from a person's neighbourhood; the loss reshapes all weights so the learned 64-number vectors separate the phenotype groups. The quality of a space is measured by the silhouette score (how compact and separated the groups are) and by a 5-nearest-neighbour classifier: plain SVD scores silhouette 0.06 and 5-NN accuracy 0.90 by ancestry; the GNN's hidden layer scores 0.70 and 0.97. That 0.06 to 0.70 is the GNN's contribution. With free learned embeddings instead of SVD initialisation the GNN reaches only 0.585 balanced accuracy on ancestry: the starting embedding matters more than the architecture. + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.95\linewidth]{embedding_individuals.png} +\caption{The GNN's 64-dimensional embedding of people projected to two dimensions: five ancestry groups separate cleanly (AMR spread between EUR and AFR, as admixture predicts); coloured by sex the same points are fully mixed.} +\end{figure} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.70\linewidth]{embedding_clusters.png} +\caption{The same model's embedding of clusters, coloured by the ancestry each cluster is enriched in.} +\end{figure} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.90\linewidth]{embedding_quality.png} +\caption{Embedding quality: silhouette by ancestry and 5-nearest-neighbour accuracy for SVD and for each GNN's hidden layer; sex stays at chance.} +\end{figure} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.90\linewidth]{init_comparison.png} +\caption{The same GNN with different starting embeddings for the person nodes: SVD-32, Node2Vec-32 (50 pretraining epochs, A100), free learned embeddings and the raw carrier row, against the logistic-regression line.} +\end{figure} + +\subsection{8. The encoder: a heterogeneous graph neural network} + +Architecture (PyTorch Geometric HeteroConv, two layers, hidden size 64): a linear projection per node type to 64 numbers; then two rounds of message passing in which, for every relation, SAGEConv adds the mean of a node's neighbours to its own state (carries, in\_block, next\_block, overlaps, encodes, each in both directions) and GraphConv adds an edge-weighted mean for co\_occurs (weight = normalised log lift) and measured (weight = the protein's harmonised z-score, signed, so a high protein pushes positively and a low one negatively); after each round LayerNorm, a residual connection, ReLU and dropout 0.3; a linear head from the person's 64 numbers to class scores. After round one a person has absorbed their clusters (and proteins); after round two, what those clusters co-occur with, their blocks, and the genes and proteins in those blocks. + +Training: class-weighted cross-entropy on training people only (weights inversely proportional to class size so AMR and small populations count); Adam with learning rate 0.005 and weight decay 0.0005; early stopping on validation balanced accuracy with patience 30, best weights restored. Full-batch: one epoch is one pass over the whole graph (about 2.4 million CARRIES, 0.37 million CO\_OCCURS and 1.07 million MEASURED edges), 0.95 s on an M2 laptop CPU and 0.10 s on an A100 GPU. About 105 thousand parameters with SVD input, about 500 thousand with the raw carrier row. + +\begin{longtable}{p{0.282\linewidth}p{0.141\linewidth}p{0.235\linewidth}p{0.282\linewidth}} +\toprule +\textbf{Target (real labels)} & \textbf{Classes} & \textbf{Logistic regression} & \textbf{GNN} \\ +\midrule +\endhead +ancestry & 5 & 0.977 & 0.974 (SVD input) \\ +population & 26 & 0.614 & 0.611 (raw input; 0.437 with SVD-32) \\ +sex (negative control) & 2 & 0.463 & 0.503 \\ +\bottomrule +\end{longtable} + +Reading: balanced accuracy on the held-out 376 people. On chr22 alone ancestry and population are almost linear functions of which clusters a person carries, so the GNN ties the linear baseline on accuracy; it wins on the embedding space and, as the next section shows, on integration. + +\subsubsection{Metrics, defined} + +\begin{itemize}[leftmargin=1.4em] + \item Accuracy: correct predictions divided by all predictions. Misleading when classes are unequal (predicting control for everyone scores 64\% on the phenotype). + \item Balanced accuracy: the mean over classes of the recall of that class (correct in class k divided by the number truly in class k). Chance is 1/k: 0.20 for ancestry, 0.038 for population, 0.5 for sex, 0.33 for site. + \item Macro-F1: for each class the harmonic mean of precision and recall, averaged over classes. + \item ROC-AUC: the probability that a randomly chosen case receives a higher case score than a randomly chosen control; 0.5 is chance, 1.0 is a perfect ranking. Only defined for two classes. + \item R-squared per protein: 1 minus the residual sum of squares divided by the total sum of squares, on test people with an observed value (at least five); negative means worse than predicting the mean. + \item Silhouette: for each person, the mean distance to their own group minus the mean distance to the nearest other group, scaled to -1..1; averaged. Higher means tighter, better separated groups in the embedding. + \item 5-nearest-neighbour accuracy: label a test person by majority vote of the five nearest training people in the embedding; balanced accuracy of that vote. + \item Precision at 20: of the twenty clusters with the highest saliency, the fraction that are planted causal clusters; chance is 20 / 6,551 x 20 = 0.06. + \item Cramer's V: sqrt(chi-square / N) for a 2 x k table; 0 means the cluster is independent of the label, 1 means it determines it. + \item Early stopping and model selection use the validation people only; every number reported in this document is on the test people. +\end{itemize} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.70\linewidth]{baseline_vs_gnn.png} +\caption{Real labels: logistic regression versus the GNN; sex, the negative control, stays at chance for both.} +\end{figure} + +\subsection{9. Integrating the proteome (schema v2)} + +The synthetic proteome we generated is keyed to the real 1000G ids and has a saved ground truth (ground\_truth.json). Per person: a site assigned at random within each ancestry (so site is a pure batch effect), the real sex, a random age; a case/control phenotype whose log-odds is a weighted sum over 20 causal haploblock clusters plus a small age term, calibrated to 38\% cases; each causal cluster also shifts one protein encoded in its own block (a cis effect); 15\% of proteins respond to the phenotype, all respond to age and sex; a per-site batch shift; and missingness that increases toward the detection limit. The first version shifted every protein with the phenotype and every model scored 1.0, the same trap an earlier team result fell into, so the signal was made sparse. + +Before entering the graph, protein levels pass a harmoniser: within each site and protein, z = (value - median) / (1.4826 x median absolute deviation). Computed from each site's own samples, it removes the between-site shift completely (0.48 to 0.00 log2) while 227 proteins keep their phenotype association. Missing values stay missing; they simply create no MEASURED edge and are never imputed as zero. + +\begin{longtable}{p{0.263\linewidth}p{0.254\linewidth}p{0.188\linewidth}p{0.094\linewidth}p{0.141\linewidth}} +\toprule +\textbf{Model} & \textbf{Input to the person node} & \textbf{Graph relations} & \textbf{AUC} & \textbf{Balanced accuracy} \\ +\midrule +\endhead +genome only & SVD-32 or raw carrier row & genome relations & 0.60-0.63 & 0.57-0.62 \\ +proteome only (MLP, no graph) & harmonised z + observed mask (920 numbers) & none & 0.96 & 0.96 \\ +genome + proteome (graph) & both & all twelve relations & 0.99 & 0.97 \\ +site (batch control, full graph) & both & all & - & 0.30 (chance 0.33) \\ +ancestry (full graph) & both & all & - & 0.90 \\ +\bottomrule +\end{longtable} + +The graph adds signal on top of the proteome, and the site control shows none of it is batch. A gradient saliency on the raw carrier input ranks clusters by their influence on the case score: 3-4 of the top 20 are ground-truth causal clusters (1.2 expected by chance). One negative result is kept deliberately: predicting the proteome from the genome embedding gives R-squared near zero even for cis-affected proteins, whereas a per-protein ridge regression on the carrier row recovers the strong cis effects (4 of 20 cis proteins with test R-squared above 0.1, 0 of 440 others). The GNN is the integration and embedding tool; discovering single-cluster cis effects needs sparse per-protein models or a prior from published protein-QTL data. + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.90\linewidth]{results_modalities.png} +\caption{Held-out AUC and balanced accuracy for the synthetic phenotype by modality; the site control sits at chance.} +\end{figure} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.95\linewidth]{training_curves.png} +\caption{Training loss and validation balanced accuracy per epoch for the v2 runs (early stopping picks the best validation epoch).} +\end{figure} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.95\linewidth]{saliency_top20.png} +\caption{Cluster saliency of the combined model: ground-truth causal clusters (teal) among the top 20.} +\end{figure} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.85\linewidth]{ridge_r2.png} +\caption{Per-protein ridge from the genome: only the strong cis effects are recoverable.} +\end{figure} + +\subsection{10. The decoder: GraphRAG with an NVIDIA NIM language model} + +For one person the decoder retrieves, deterministically and only from the graph: the profile (ancestry, population, sex, site, age; the true phenotype withheld), the GNN's prediction, the globally salient clusters the person carries, their eight most ancestry-informative clusters with block, genes and proteins, their eight most extreme protein levels with the encoding block and whether that block holds a notable cluster, and their five nearest people in the GNN embedding. This is serialised as JSON (about 4,500 tokens) and sent to nvidia/nemotron-3-super-120b-a12b through NVIDIA's OpenAI-compatible NIM endpoint with reasoning\_effort set to none (otherwise this reasoning model thinks inline and exhausts the token budget before answering) under a system prompt that forbids inventing entities and requires every id to be cited verbatim. The reply's cited ids are checked against the context before anything is written. For person HG00103 the model answered in 13 seconds with 25 cited ids and none unknown, producing a summary, ancestry and phenotype assessments, genome-proteome links (for example cluster chr22\_40032702-40132216\_cluster151 in the block encoding TNRC6B / Q9UPQ9) and caveats stating that the phenotype is synthetic and ancestry is population structure, not a medical finding. + +\subsection{11. Federated learning with NVFlare} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.95\linewidth]{federated_topology.png} +\caption{Federated topology: three sites with private people and edges, one shared public graph, weights only to the server.} +\end{figure} + +Partition: site s receives the people whose site code is s (835 / 835 / 833, mixed ancestry). Its graph is PyG's HeteroData.subgraph restricted to those people: their nodes are re-indexed, their CARRIES and MEASURED edges kept, and every other node type and edge kept whole because those are public. Site 1, for instance, holds 835 people, 774,753 CARRIES and 355,294 MEASURED edges. The person's input (carrier row plus harmonised protein vector and mask) is computed from the site's own rows; no cross-site preprocessing exists because the harmoniser is already per site. + +Mechanics (NVFlare 2.9, FedAvgRecipe with the PyTorch Client API): job.py derives every model constructor value from the public graph into model\_args.json so the server and all clients build byte-identical models. Each client runs flare.init(), builds its site graph once, then loops: receive the global weights, evaluate them on its own validation and test people, train five full-batch epochs on its own training people, send back the weights, the metrics and the number of optimizer steps. The server averages the weights (weighted by steps, equal here) and selects the best global model by validation balanced accuracy. What crosses the site boundary per round: one state dict of about half a million numbers and five scalars; no rows, no protein values, no embeddings of people. + +\begin{longtable}{p{0.376\linewidth}p{0.188\linewidth}p{0.141\linewidth}p{0.235\linewidth}} +\toprule +\textbf{Model} & \textbf{Balanced accuracy} & \textbf{AUC} & \textbf{Evaluated on} \\ +\midrule +\endhead +federated global model, 10 rounds x 5 epochs & 0.90 & 0.955 & the same 376 held-out people \\ +federated global model, 30 rounds x 5 epochs & 0.963 & 0.998 & the same 376 held-out people \\ +central model (train\_gnn\_v2, both/raw) & 0.969 & 0.992 & the same 376 held-out people \\ +federated model per site (30 rounds) & 0.983 / 0.981 / 0.933 & 1.000 / 0.999 / 0.997 & each site's own held-out people \\ +\bottomrule +\end{longtable} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.95\linewidth]{federated_rounds.png} +\caption{FedAvg convergence: the global model's AUC and validation balanced accuracy at each site, per round, against the central model.} +\end{figure} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.90\linewidth]{federated_vs_central.png} +\caption{Central versus federated (10 and 30 rounds) on the same held-out people, and the federated model on each site's own held-out people.} +\end{figure} + +Ten rounds (50 local steps) were short of the central run's \textasciitilde{}80 epochs; thirty rounds converge to the central model's level. Federated training loses nothing here because the sites are random draws of the same population by construction; with ancestry-pure sites the averaging would have to fight client drift, which is the next experiment. + +\subsubsection{With and without federation, on this data} + +federated/local\_only.py trains each site alone on its own people for the same 150 optimizer steps the federated clients used, then scores that lone model on its own held-out people and on the other sites' held-out people. On this synthetic data a lone site already does well, because 835 people and a strong proteome signal are enough: own-site AUC SITE1 0.998, SITE2 0.986, SITE3 0.988; the worst transfer of a lone model to another site's people is SITE1 0.949, SITE2 0.969, SITE3 0.997. The federated global model scores SITE1 0.999, SITE2 1.000, SITE3 0.968 on the same per-site people and 0.987 to 0.998 on the pooled held-out set across runs. With a smaller budget the picture changes: at 50 optimizer steps per site (a 10-round run inside the Docker image on the A100) a lone site reaches only 0.88 to 0.92 while the federated model reaches 0.989 to 0.998 on the same per-site people, because the averaged weights have effectively seen every site's people. So federation costs nothing when a site has enough data and budget and helps clearly when it does not; in both cases what it solves is the constraint, not the score: one shared model, trained on everyone, with no row leaving any site. A larger gain is expected when sites differ systematically (ancestry-pure hospitals, different protein panels), which is the ancestry-partitioned experiment listed under next steps. + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.80\linewidth]{federated_site_alone.png} +\caption{With and without federation: each site trained alone on its own people versus the federated global model, scored on the same held-out people per site (A100 re-run).} +\end{figure} + +\begin{longtable}{p{0.301\linewidth}p{0.160\linewidth}p{0.160\linewidth}p{0.160\linewidth}p{0.160\linewidth}} +\toprule +\textbf{Site (held-out n)} & \textbf{Alone, 150 steps} & \textbf{Federated, 150 steps} & \textbf{Alone, 50 steps} & \textbf{Federated, 50 steps} \\ +\midrule +\endhead +SITE1 (111) & 0.998 / 0.952 & 0.999 / 0.983 & 0.881 / 0.543 & 0.998 / 0.950 \\ +SITE2 (136) & 0.986 / 0.852 & 1.000 / 0.991 & 0.902 / 0.772 & 0.991 / 0.907 \\ +SITE3 (129) & 0.988 / 0.923 & 0.968 / 0.933 & 0.915 / 0.830 & 0.989 / 0.865 \\ +average over sites & 0.991 / 0.909 & 0.989 / 0.969 & 0.899 / 0.715 & 0.993 / 0.907 \\ +pooled 376 people, federated global model & - & 0.987 / 0.967 & - & 0.991 / 0.901 \\ +pooled 376 people, central model & 0.995 / 0.960 & & 0.997 / 0.959 & \\ +\bottomrule +\end{longtable} + +AUC / balanced accuracy on each site's own held-out people. Transfer of a lone model to other sites' people (150 steps, AUC): the SITE1 model scores 0.982 on SITE2 people and 0.949 on SITE3 people; SITE2: 0.993 and 0.969; SITE3: 0.997 and 0.998; the federated model 0.999 / 1.000 / 0.968 on the three sites. + +\subsubsection{Why federated, and what kind} + +The mission sentence of the README is a statement about where data lives: a hospital may compute on its patients but may not ship their rows. This is horizontal federated learning: every site has the same columns (the same graph schema, the same protein panel) and different rows (different people). It is not a split of the genome across sites (each site would then hold part of every person, which is vertical federation and a different problem) and not a way to add chromosomes: another chromosome is another shared reference graph, added at every site at once. What is gained is the ability to train on all 2,503 people while each site only ever reads its 835; what is lost, in this experiment, is nothing measurable, because the sites are alike. A real deployment would report each site's own held-out score and never assemble a central test set. + +\subsubsection{Mechanics in detail} + +\begin{itemize}[leftmargin=1.4em] + \item Model definition (federated/model.py): the same architecture as train\_gnn\_v2, copied into a self-contained file so the NVFlare server can import it without the pipeline. Its constructor reads model\_args.json (input sizes per node type, the twelve relations, hidden 64, two layers, dropout 0.3, mean aggregation) so server and clients build identical state dicts. + \item Job (federated/job.py): FedAvgRecipe(name, model class path and args, min\_clients 3, num\_rounds, train\_script client.py, train\_args, key\_metric val\_balanced\_accuracy, server\_expected\_format PYTORCH); add\_decomposers registers TensorDecomposer so tensors travel natively; add\_server\_file ships model.py to the server; SimEnv(num\_clients 3, workspace\_root) runs the three sites as threads on one machine; recipe.execute(env) writes the job and runs it. + \item Client (federated/client.py): flare.init(); the site name (site-1, site-2, site-3) selects the site code; HeteroData.subgraph keeps that site's people and their CARRIES and MEASURED edges and leaves the public node types whole; then the loop: flare.receive() gives the global weights and the round number; evaluate them on the site's validation and test people; if the task is evaluate-only, send metrics; otherwise train five full-batch epochs (five optimizer steps) with the site's own class weights and send FLModel(params = state dict on CPU, metrics, meta NUM\_STEPS\_CURRENT\_ROUND = 5). + \item Server: after each round the global weights become the step-weighted mean of the three state dicts (weights equal here because every site does five steps; aggregation\_weights can weight by site size); the best global model by validation balanced accuracy is kept as best\_FL\_global\_model.pt, the final one as FL\_global\_model.pt, both under the workspace's app\_server folder. + \item What crosses the boundary per round per site: one state dict of about 0.5 million floats and five scalars. No carrier row, no protein value, no person embedding. Weight updates can in principle leak information about training data; NVFlare offers differential privacy and homomorphic-encryption filters for that, and neither was needed for a simulation. + \item Evaluation (federated/evaluate\_global.py): loads FL\_global\_model.pt, rebuilds the full graph exactly as the central run did, scores the same 376 held-out people and each site's own held-out people, and writes evaluation.json next to the central numbers. + \item From simulation to real machines: NVFlare's POC mode starts a server and clients as separate processes (or machines) with the same job; each client would run client.py against its own kg-dir and split; the L4 and A100 instances could be two such sites. Adding a real fourth site means: its own people with 1000G-style ids, its own protein matrix keyed by those ids, its own harmoniser pass, and the shared graph files copied over. +\end{itemize} + +\subsection{12. System design, tech stack and deployment} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.98\linewidth]{architecture_slide.png} +\caption{System architecture on one page: private hospital sites, the shared reference graph, the GNN encoder trained through the NVFlare server, and the outputs feeding the LLM decoder.} +\end{figure} + +\begin{longtable}{p{0.169\linewidth}p{0.442\linewidth}p{0.329\linewidth}} +\toprule +\textbf{Layer} & \textbf{Choice} & \textbf{Notes} \\ +\midrule +\endhead +Language and data & Python 3.13, pandas 3.0, scipy 1.18 (sparse), scikit-learn 1.9 & pinned in requirements.txt \\ +Graph learning & PyTorch 2.14, PyTorch Geometric 2.8, pyg-lib (random walks for Node2Vec) & torch\_cluster is deprecated in favour of pyg-lib; that caused the first image-build failure \\ +Graph tooling & NetworkX 3.6, nx-cugraph (GPU dispatch), Neo4j 5.26 community in Docker & graph statistics, GraphML, browsing \\ +Container & pytorch/pytorch:2.14.0-cuda12.6-cudnn9-runtime base, PIP\_BREAK\_SYSTEM\_PACKAGES=1, torch-tensorrt 2.14 & 10.5 GB image; runs on CPU when no GPU is present \\ +Compute & Mac M2 CPU for development; NVIDIA Brev: L4 24 GB (GCP, \$0.85/h) and A100 80 GB (Crusoe, \$1.98/h) & brev\_deploy.sh creates or reuses an instance, uploads, builds natively, runs, copies outputs back \\ +Inference & torch.compile (inductor): 36 ms $\rightarrow$ 4.9 ms per full graph, logits equal to within 0.0001 & Torch-TensorRT could not compile this scatter-heavy GNN within 3 hours and is not claimed \\ +LLM & NVIDIA NIM, nvidia/nemotron-3-super-120b-a12b & key in the environment only; a local NIM container on the A100 would keep patient context on site \\ +Federated & NVFlare 2.9 FedAvgRecipe, SimEnv simulator; federated/local\_only.py for the site-alone comparison & 3 clients as threads on one machine; POC mode across real machines is the next step \\ +Packaging & Makefile (setup, run, run-v2, eda, decode, federated, docker, docker-run-v2, docker-federated, brev, report), setup.sh, run\_all.sh, run\_v2.sh, config.py + .env.example, 11 unit tests & clone-and-run on laptop or GPU; the image is the whole solution (data mounted, .env passed) \\ +\bottomrule +\end{longtable} + +Runtime, measured on 18 September from an empty folder (the verification runs of section 13): + +\begin{longtable}{p{0.470\linewidth}p{0.235\linewidth}p{0.235\linewidth}} +\toprule +\textbf{Stage} & \textbf{Laptop, Apple M2 CPU} & \textbf{A100 80 GB (Docker)} \\ +\midrule +\endhead +environment: venv or image, pinned dependencies, unit tests & 28 s (uv) & 3 min image build \\ +v1: download, knowledge graph, statistics, baseline, plots, three GNN runs, embedding quality & 6 min & 3 min \\ +v2: synthetic proteomics, graph v2, EDA, six GNN runs, ridge, decoder & 19 min & 4 min \\ +federated: 30 rounds x 5 local epochs, central scoring, site-alone comparison & 4 min & 2 min \\ +whole chain & about 30 min & about 14 min including the image build \\ +one GNN training run & 30 s to 3 min & 1 to 16 s \\ +one full-graph inference pass (2,548 people) & - & 36 ms eager, 4.9 ms compiled \\ +one decoder call (NIM, remote) & 4 to 13 s & same \\ +\bottomrule +\end{longtable} + +\begin{figure}[htbp] +\centering +\includegraphics[width=0.85\linewidth]{compute_benchmarks.png} +\caption{Training epoch time on the laptop CPU versus the A100, and full-graph inference eager versus torch.compile on the A100.} +\end{figure} + +\subsubsection{The Docker image, layer by layer} + +The Dockerfile starts from pytorch/pytorch:2.14.0-cuda12.6-cudnn9-runtime (PyTorch with CUDA 12.6 already inside), sets PIP\_BREAK\_SYSTEM\_PACKAGES=1 because the base image's Python is system-managed, installs curl, copies requirements.txt and installs the pinned libraries, then tries three optional extras and prints a clear fallback message if any is unavailable for this torch build: pyg\_lib from the PyG wheel index (Node2Vec random walks; without it the code falls back to SVD), nx-cugraph from NVIDIA's index (NetworkX dispatches to cuGraph on the GPU; without it NetworkX runs on the CPU), and torch-tensorrt from the PyTorch cu126 index. It copies the code, runs the unit tests as part of the build so a broken image cannot be produced, and defaults to running run\_all.sh. Data and outputs are bind-mounted at run time, so the image never contains data. docker-compose.yml adds Neo4j 5.26 community with a persistent volume and the same pipeline image with GPU reservation. + +\subsubsection{GPU deployment on Brev, step by step} + +brev\_deploy.sh needs a logged-in brev CLI (brev login --api-key). It creates the instance if it does not exist (default an L4 on GCP; BREV\_INSTANCE=progenome-a100 selects the A100), waits until the instance reports RUNNING and READY, refreshes the SSH alias, checks nvidia-smi and docker over ssh, uploads genomics/ plus the two small proteomics inputs as one tarball, builds the image natively on the GPU box (no emulation, about 4 minutes on the A100), runs the whole v1 pipeline inside the container with the GPU, runs the inference benchmark (eager, torch.compile, TensorRT attempt) and copies outputs/ back to outputs\_brev//. Billing continues until brev stop; the A100 is left running for the presentation. + +\subsubsection{Inference and what torch.compile does} + +infer.py rebuilds the exact inputs of a trained run, wraps the model so torch.compile sees plain tensors instead of dictionaries, times ten full-graph passes eager, then compiles with the inductor backend (kernel fusion and graph capture) or with the torch\_tensorrt backend and times again, and reports the maximum absolute logit difference so a speed-up cannot hide a numerical change. On the A100: 36.3 ms eager, 4.9 ms compiled (7.4x), maximum logit difference below 0.0001 across three runs, accuracy on all 2,503 labelled people 0.980 with SVD input. The TensorRT backend partitions the graph into dense parts it can compile (linear layers, norms) and scatter parts it cannot; on this hetero-GNN the compile step did not finish in three hours and is not part of any claim. + +\subsection{13. How the goals were met, and what is not claimed} + +\begin{longtable}{p{0.301\linewidth}p{0.470\linewidth}p{0.169\linewidth}} +\toprule +\textbf{Goal} & \textbf{Evidence} & \textbf{Status} \\ +\midrule +\endhead +RQ1: one graph joining haploblocks, genes, proteins and people & schema v2 with 5 node types and 7 edge types, 2,548 people joined by id, browsable in Neo4j & done \\ +RQ2: GNN combines genome and proteome & AUC 0.60 (genome) / 0.96 (proteome) / 0.99 (both); site control at chance; embedding silhouette 0.70; saliency finds causal clusters at 3x chance & done on synthetic ground truth \\ +RQ3: federated training without moving records & NVFlare FedAvg over 3 sites, AUC 0.998 vs central 0.992 on the same people & done in simulation \\ +Encoder $\rightarrow$ LLM decoder $\rightarrow$ insights & GraphRAG decoder with validated citations on Nemotron 3 Super & done \\ +Reproducible and deployable & Makefile, Docker image, Brev deployment on L4 and A100, tests & done \\ +\bottomrule +\end{longtable} + +\begin{itemize}[leftmargin=1.4em] + \item The case/control phenotype is synthetic: the integration numbers show the pipeline recovers a planted signal, not biology. Ancestry, population and sex results are on real labels. + \item Chromosome 22 only; the code is chromosome-agnostic (CHROM=chr21 make run). + \item On single-label accuracy the GNN ties, not beats, logistic regression; its value is the embedding space and the integration. + \item The GNN embedding does not recover single-cluster cis effects; a per-protein ridge does for the strong ones. + \item Federated evaluation reuses the central held-out people; a real deployment would report per site only. + \item TensorRT is not part of the inference claim; torch.compile is. +\end{itemize} + +\subsubsection{Verification and reproduction (18 September, before the commit)} + +Before committing, the whole chain was re-run twice from the files that would be committed: once as a simulated fresh clone on the laptop (only the tracked files copied to an empty folder, then make setup, make run, make run-v2, make federated, make test) and once on the A100 with the Docker image rebuilt from the current Dockerfile, which had not been built since the pyg\_lib change. Both runs reproduce every count exactly and every model score within run-to-run noise (GPU kernels and FedAvg are not bit-reproducible). The rebuilt image contains pyg\_lib, so Node2Vec ran natively for the first time; with its original five pretraining epochs it was clearly undertrained (ancestry 0.80), with fifty epochs the loss converges and it becomes a working but slightly weaker alternative to SVD, so SVD remains the default everywhere and the deploy script now defaults to it. The decoder was also re-called on the laptop run; the language model's wording and the number of ids it chooses to cite vary between calls, the citation check is what stays constant. + +\begin{longtable}{p{0.320\linewidth}p{0.207\linewidth}p{0.207\linewidth}p{0.207\linewidth}} +\toprule +\textbf{Quantity} & \textbf{Documented (original runs)} & \textbf{Fresh clone, laptop CPU} & \textbf{Rebuilt image, A100} \\ +\midrule +\endhead +kept clusters / CARRIES / CO\_OCCURS & 6,551 / 2,365,574 / 187,030 & identical & identical \\ +genes / proteins / MEASURED edges & 458 / 460 / 1,065,712 & identical & identical \\ +ancestry-associated clusters (FDR 5\%) / sex & 6,470 / 0 & 6,470 / 0 & 6,470 / 0 \\ +logistic baseline ancestry / population / sex & 0.977 / 0.614 / 0.463 & 0.977 / 0.614 / 0.463 & 0.981 / 0.614 / 0.463 \\ +GNN v1 ancestry / population / sex (SVD input) & 0.974 / 0.437 / 0.503 & 0.974 / 0.437 / 0.503 & run with Node2Vec instead (next row) \\ +GNN v1 with Node2Vec input, 50 epochs (A100 only) & not previously measured & - & 0.950 / 0.292 / 0.489 \\ +embedding silhouette by ancestry: SVD / GNN & 0.06 / 0.70 & 0.06 / 0.70 & 0.06 / 0.70 (Node2Vec-initialised GNN) \\ +phenotype AUC genome / proteome / both & 0.64 / 0.96 / 0.99 & 0.65 / 0.96 / 0.99 & 0.64 / 0.96 / 0.99 \\ +site control balanced accuracy (chance 0.33) & 0.30 & 0.34 & 0.30 \\ +saliency: causal clusters in top 20 & 4 & 4 & 4 \\ +ridge: cis proteins with R2 > 0.1 / others & 4 / 0 & 4 / 0 & 4 / 0 \\ +federated 30 rounds AUC / central AUC, same 376 people & 0.998 / 0.992 & 0.996 / 0.994 & 0.987 / 0.995 \\ +site-alone AUC on own test people vs federated model on the same people & not previously measured & - & 0.998 / 0.986 / 0.988 vs 0.999 / 1.000 / 0.968 \\ +decoder: cited ids / invented ids & 25 / 0 & 4 / 0 (4.2 s) & dry run (no key on the box) \\ +inference eager / torch.compile (A100) & 36 ms / 4.9 ms & - & 36.3 ms / 4.9 ms (7.45x, max logit diff 1.9e-06) \\ +unit tests & 11 pass & 11 pass (make setup) & 11 pass (docker build) \\ +\bottomrule +\end{longtable} + +\subsection{14. Team context} + +Friederike Duendar (lead) wrote the README's research questions, the UniProt gene BED and an R exploration of gene-block overlaps; Nolan Bruyat built the synthetic proteomics generator and its plots; Zillur Rahman built the proteomics-side analysis now on main (see below); Yan Zhou coordinates the manuscript (two introduction paragraphs, two methods paragraphs, one results paragraph); Anita Egebor and Alvaro Martinez Barrio contributed to the README and data access. The genomics/ pipeline described here is the modelling backbone into which those pieces plug. + +\subsubsection{What is on main since this branch was created, and how it relates} + +The team's proteomics work on main (the synthetic proteomics generator with 4,000 SITE-id samples, protein filtering and classification baselines, a protein-centred knowledge graph built from the HaploGraph edge list annotated with UniProt proteins, a federated comparison of feature sets, and methods\_and\_results.md) covers the proteome side and the federated logistic setting. genomics/ adds the genome side: the person-level knowledge graph on the 1000 Genomes ids, the graph neural network, the LLM decoder and the NVFlare run. The two share the same HaploGraph edge file, the same UniProt gene BED and the same three-site design, and they are complementary: nothing in genomics/ touches a path on main, so the branch merges cleanly, and for the manuscript the proteomics analysis and the graph model plug into the same methods and results structure. + +\section{Part IV: Reference} + +\subsection{Run it yourself} + +From a fresh clone of the repository on branch modelling, on a laptop or a GPU machine: + +\begin{small}\begin{verbatim} +git clone https://github.com/collaborativebioinformatics/ProGenome.git +cd ProGenome && git checkout modelling && cd genomics +make setup # .venv with torch (CPU, or CUDA if nvidia-smi works), pinned deps, unit tests +make run # v1: download -> graph -> statistics -> baseline -> plots -> GNN (ancestry, population, sex) -> embeddings +make run-v2 # v2: synthetic proteomics -> graph v2 -> EDA -> GNN ablations -> ridge -> decoder dry run +make federated ROUNDS=30 LOCAL_EPOCHS=5 # NVFlare FedAvg over 3 sites, then central scoring +make neo4j-load # browse the graph at http://localhost:7474 (neo4j / progenome) +make docker && make docker-run # the same v1 chain inside the CUDA image +BREV_INSTANCE=progenome-a100 make brev # build and run on the A100, copy outputs back +make decode WHO=HG00103 # LLM insight for one person (needs NVIDIA_API_KEY) +make test # 11 unit tests on a toy graph +make help # every target +\end{verbatim}\end{small} + +Single stages take a --chrom flag and, where relevant, --target, --modality, --init; for example .venv/bin/python train\_gnn\_v2.py --target phenotype --modality both --init raw reproduces the combined model with saliency, and .venv/bin/python infer.py --run ancestry\_svd --compile inductor reproduces the inference benchmark. Every output path is relative to genomics/, and data/, outputs/ and outputs\_brev/ are git-ignored. + +\subsection{Secrets and configuration} + +Configuration enters the code in one place, config.py, which reads genomics/.env, then \textasciitilde{}/.progenome.env, then defaults, with exported shell variables taking precedence; .env.example lists every variable with a comment (NVIDIA\_API\_KEY, NIM\_MODEL, NIM\_URL, NEO4J\_URI/USER/PASSWORD, HAPLOBLOCKS\_BASE, CHROM, data and output directories, BREV\_INSTANCE/TYPE) and make config prints what is in effect with secrets masked. The shell scripts source the same files through load\_env.sh, and the Docker targets pass .env into the container. The only secret is the NVIDIA API key; .env is git-ignored, nothing under the repository contains a key, and the decoder refuses to call the endpoint without one instead of falling back silently. Versions are pinned in requirements.txt (torch 2.14.0, torch\_geometric 2.8.0.post1, pandas 3.0.5, scipy 1.18.1, scikit-learn 1.9.1, networkx 3.6.1, neo4j 6.3.1, nvflare 2.9.0, pytest 9.1.1) and setup.sh installs torch from the CPU or cu126 index depending on whether a GPU is present. + +\subsection{Repository map (genomics/)} + +\begin{small}\begin{verbatim} +haplokg.py, build_kg.py knowledge graph v1 (people, clusters, blocks) +haplokg_proteins.py, build_kg_v2.py genes, proteins, harmonised measurements (v2) +proteomics_synth_1000g.py synthetic proteomics on 1000G ids with ground truth +cooccurrence_analysis.py, eda.py statistics and the EDA report +baseline.py logistic regression + the shared split +graph_explore.py, neo4j_load.py NetworkX statistics/plots, Neo4j loader +train_gnn.py, train_gnn_v2.py the GNN (genome; genome+proteome) +embeddings.py embedding quality and plots +proteome_linear_baseline.py per-protein ridge (genome -> proteome) +graphrag_decoder.py NIM LLM decoder +infer.py inference benchmark +federated/{model,client,job,evaluate_global}.py NVFlare FedAvg +Dockerfile, docker-compose.yml, Makefile, setup.sh, run_all.sh, run_v2.sh, brev_deploy.sh +docs/architecture.html (.mmd), METHODS.md, DEEP_DIVE.md, report/ +tests/ unit tests on a toy graph +\end{verbatim}\end{small} + +\subsection{Glossary} + +\begin{longtable}{p{0.235\linewidth}p{0.705\linewidth}} +\toprule +\textbf{Term} & \textbf{Meaning} \\ +\midrule +\endhead +Haploblock & a stretch of chromosome between recombination hotspots, inherited as a unit \\ +Haplotype & one copy's sequence of a block; each person has two per block \\ +Cluster & a group of near-identical haplotypes of one block (MMseqs2); a person carries it if either haplotype is in it \\ +Carrier matrix & people x clusters 0/1 matrix; the CARRIES edges \\ +Lift & how much more often two clusters co-occur than chance: P(A and B) / (P(A) P(B)) \\ +Cramer's V & effect size of a contingency test, 0 (independent) to 1 (perfectly associated) \\ +FDR & false discovery rate; Benjamini-Hochberg controls the expected fraction of false positives among findings \\ +SVD & singular value decomposition; here a 32-component factorisation of the carrier matrix giving embeddings \\ +GNN / message passing & neural network on a graph; each layer mixes a node's vector with its neighbours' \\ +SAGEConv / GraphConv & two PyG layer types: neighbour-mean aggregation; edge-weighted aggregation \\ +Balanced accuracy / AUC & mean per-class recall; probability a random case outscores a random control \\ +Harmoniser & per-site, per-protein robust z-score removing batch offsets \\ +MNAR & missing not at random; here low-abundance proteins go missing first \\ +FedAvg & federated averaging of model weights across sites \\ +NIM & NVIDIA Inference Microservice; an OpenAI-compatible LLM endpoint \\ +Brev & NVIDIA's GPU cloud; instances by the hour \\ +\bottomrule +\end{longtable} + +\end{document} diff --git a/genomics/docs/report/build_report.js b/genomics/docs/report/build_report.js new file mode 100644 index 0000000..4b05623 --- /dev/null +++ b/genomics/docs/report/build_report.js @@ -0,0 +1,590 @@ +// Builds the ProGenome knowledge-transfer document from ONE content source into LaTeX and DOCX. +// NODE_PATH= node build_report.js +// Outputs: ProGenome_KT.tex (compile with tectonic/pdflatex) and ProGenome_KT.docx, next to this file. +"use strict"; +const fs = require("fs"); +const path = require("path"); +const docx = require("docx"); + +const HERE = __dirname; +const FIG = path.join(HERE, "figures"); + +// ----------------------------------------------------------------------------- content +// Block types: h(level,text) p(text) b([items]) t(header,rows,widths?) f(file,caption,width) c(code) +const h = (level, text) => ({ k: "h", level, text }); +const p = (text) => ({ k: "p", text }); +const b = (items) => ({ k: "b", items }); +const t = (header, rows, widths) => ({ k: "t", header, rows, widths }); +const f = (file, caption, width = 0.95) => ({ k: "f", file, caption, width }); +const c = (code) => ({ k: "c", code }); + +const META = { + title: "ProGenome: a federated workflow for genome-graph and proteomic integration", + subtitle: "Complete knowledge-transfer document: problem, biology, data, pipeline, knowledge graph, models, federated learning, deployment and results", + team: "Team #3, Nordic Biobank x NVIDIA Federated Learning Hackathon, Copenhagen, September 2026. Branch: modelling. Author of this build: Koushik Telaprolu (genomics/ pipeline), with the team's proteomics and README work referenced where used.", + date: "18 September 2026", +}; + +const CONTENT = [ + h(1, "How to read this document"), + p("This is written for a teammate who joins today and knows nothing about the biology, the software, the models or the infrastructure. Part I is a primer that defines every concept used later. Part II is the whole workflow on a few pages: what goes in and what comes out of every stage, where each data source is used, what the trained model produces and how it is used, ground truth against prediction, where the language model sits, and how every item of the team README on branch main is covered. Part III is the project in depth, in the order the data flows: problem, data sources, data pipeline, exploratory analysis, knowledge graph, embeddings, the graph neural network (the encoder), the language-model decoder, the federated training, deployment, results and caveats. Part IV is how to run everything, a glossary and a repository map. Every number is measured on chromosome 22 with random seed 42 and comes from files under genomics/outputs; nothing is typed in from memory."), + + h(1, "Part I: Primer"), + h(2, "Biology in ten minutes"), + p("DNA is a long text written in four letters (A, C, G, T). Humans have about 3 billion letters, packaged in 23 pairs of chromosomes; chromosome 22 is one of the smallest, about 50 million letters. Everyone carries two copies of each chromosome, one from each parent. A gene is a region of DNA that encodes a protein; proteins are the molecules that do the work in cells, and measuring how much of each protein a person has is called proteomics."), + p("Two people's DNA differs at roughly one letter in a thousand; those positions are variants. A haplotype is the specific combination of variants along one copy of a chromosome. Because DNA is inherited in chunks (recombination cuts and rejoins the parental copies at a limited number of places), neighbouring variants tend to travel together. A haploblock is a stretch of chromosome between recombination hotspots that is usually inherited as one unit. Phased data means we know, for every variant, which of the two copies it sits on, so each person's two haplotypes per block are known separately."), + p("The 1000 Genomes Project (1000G) sequenced 2,548 people from 26 populations grouped into five continental ancestries: AFR (African), AMR (admixed American), EAS (East Asian), EUR (European) and SAS (South Asian). It recorded only ancestry, population and sex about them, nothing clinical. A phenotype is any observable property of a person; in this project the real phenotypes are those three, and a clinical-looking case/control phenotype had to be simulated."), + h(2, "Software in ten minutes"), + p("Python is the language of every script here; pandas handles tables, scipy handles sparse matrices (mostly-zero grids stored compactly), PyTorch does neural networks and PyTorch Geometric (PyG) adds graph neural networks. Git tracks versions of the code; a branch is a parallel line of work (ours is called modelling). Docker packages the code and every library it needs into an image so it runs identically on any machine. A GPU is a processor built for many small parallel calculations; CUDA is NVIDIA's software that lets PyTorch use it. NVIDIA Brev rents GPU machines by the hour. Neo4j is a database made for graphs, with a browser to look at them."), + h(2, "Data science in ten minutes"), + p("Exploratory data analysis (EDA) means measuring and plotting the data before modelling: sizes, distributions, missing values, obvious structure. To judge a model honestly the people are split once into train (learn), validation (choose settings and when to stop) and test (report only once, at the end); here 70/15/15 percent, chosen at random but stratified so each ancestry is represented in every part. Accuracy is the fraction correct; balanced accuracy averages the accuracy per class so a rare class cannot be ignored; AUC (area under the ROC curve) is the probability that a random case is scored higher than a random control, 0.5 is guessing and 1.0 is perfect. A negative control is a target that must come out at chance level if the method is sound."), + h(2, "Machine learning in ten minutes"), + p("An embedding is a list of numbers (a vector) that represents an object so that similar objects get similar vectors. A graph is a set of nodes joined by edges; a knowledge graph is a graph whose nodes and edges have types and properties. A graph neural network (GNN) computes a vector for every node by repeatedly mixing each node's own vector with those of its neighbours (message passing); after two rounds a node's vector summarises its two-hop neighbourhood. Training means adjusting the network's weights so that a prediction made from those vectors matches known labels, measured by a loss; an optimiser (Adam) nudges the weights to reduce the loss, one epoch being one pass over the data. An encoder turns raw data into vectors; a decoder turns vectors (and retrieved facts) into an output, here text written by a large language model (LLM). Retrieval-augmented generation (RAG) gives the LLM the facts it should use so it does not invent them; GraphRAG retrieves those facts from a graph."), + h(2, "Federated learning in five minutes"), + p("Hospitals cannot pool patient data. Federated learning trains one model across sites without moving records: each site trains the shared model on its own data for a few steps, sends only the updated weights to a coordinator, the coordinator averages them (FedAvg) and sends the average back; repeat for several rounds. NVIDIA FLARE (NVFlare) is the framework that runs this loop; its simulator runs all sites on one machine so the workflow can be tested before real deployment."), + + h(1, "Part II: The workflow on a few pages"), + h(2, "The story, stage by stage"), + p("Read the workflow as one sentence first: a public graph of haplotype clusters is joined to people by their sample id; each person's row of that graph, their phenotype labels and their protein measurements are attached to a person node; a graph neural network is trained to predict a label from the person's neighbourhood; the trained network yields predictions, an embedding per person and cluster, a saliency per cluster and a set of weights; the predictions and embeddings are checked against ground truth, a language model turns them into a cited report per person, and the weights are what a federated deployment exchanges instead of data. Every stage below is one script with one Makefile target."), + f("data_flow_map.png", "Where each data source enters, what the trained model produces, and where each product goes.", 0.98), + t(["Stage", "What goes in", "What comes out", "Lands in"], [ + ["1 download", "URLs on data.haploblocks.org", "HaploGraph node and edge files, phenotypes, block statistics, md5-verified", "data/"], + ["2 knowledge graph v1", "nodes.csv.gz (who carries which cluster), edges (which clusters co-occur), block statistics, phenotypes", "Individual, Cluster, Block nodes; CARRIES, CO_OCCURS, IN_BLOCK, NEXT_BLOCK edges; labels stored on the person nodes", "outputs/kg/chr22/ (carries.npz, hetero.pt)"], + ["3 statistics", "the graph and the labels", "for every cluster and edge, how strongly it tracks ancestry, population and sex (Cramer's V, FDR)", "outputs/cooccurrence/chr22/"], + ["4 baseline and split", "carrier matrix and labels", "logistic-regression scores to beat; the single train/val/test split every later model reuses", "outputs/baseline/, outputs/splits/"], + ["5 starting embeddings", "carrier matrix only, no labels", "a 32-number vector per person and per cluster (SVD), or Node2Vec, learned, or the raw row", "computed inside training"], + ["6 GNN v1 (encoder)", "graph, starting vectors, and the labels of training people only", "trained weights; class probabilities per person; a 64-number embedding per person and per cluster", "outputs/gnn/chr22/_/"], + ["7 embedding check", "SVD and GNN embeddings, labels of test people", "silhouette and nearest-neighbour accuracy, 2-D plots", "outputs/embeddings/chr22/"], + ["v2.1 synthetic proteomics", "the real 1000G ids, the carrier matrix, the gene BED", "3 per-site protein matrices, metadata (site, age, sex, case/control), ground_truth.json", "outputs/proteomics_synth/chr22/"], + ["v2.2 knowledge graph v2", "graph v1, gene BED, proteomics", "Gene and Protein nodes; OVERLAPS, ENCODES, MEASURED edges (harmonised z); site and phenotype labels on people", "outputs/kg/chr22/hetero_v2.pt"], + ["v2.3 EDA", "everything above", "the eight-section report with tables and plots", "outputs/eda/chr22/EDA.md"], + ["v2.4 GNN v2", "graph v2; modality genome / proteome / both; target phenotype / site / ancestry / sex / proteome", "metrics, test predictions, embeddings, saliency per cluster, training history", "outputs/gnn_v2/chr22/__/"], + ["v2.5 ridge", "carrier matrix, protein z-scores", "per-protein test R-squared, split into cis and other proteins", "outputs/gnn_v2/chr22/proteome_ridge_baseline/"], + ["v2.6 decoder", "one person: graph neighbourhood, GNN prediction and embedding neighbours, saliency", "a cited JSON insight written by the NIM language model, with a citation check", "outputs/graphrag/chr22/_insight.json"], + ["v2.7 federated", "graph v2 split by site, the model definition", "a global model trained without moving any person's data; its score on the same held-out people", "outputs/federated/chr22/"], + ["inference", "a trained run", "predictions and embeddings for all 2,548 people; eager versus compiled timing", "outputs/gnn/chr22//inference/"], + ], [0.14, 0.3, 0.38, 0.18]), + + h(2, "Where each data source is used, and where it is not"), + p("The most common confusion is what the phenotypes and the proteomics do. The phenotypes are labels: they are stored on the person node as training targets and evaluation ground truth and never become an edge, a node feature or an input to the starting embeddings. The proteomics enters three times: as MEASURED edges from a person to the proteins observed in them, as part of the person's input vector, and, through its metadata, as two more labels (site and the synthetic case/control phenotype). The HaploGraph provides the structure everyone shares."), + t(["Source", "Used for", "Never used for"], [ + ["HaploGraph nodes.csv.gz", "CARRIES edges; the carrier matrix behind the SVD starting vectors, the logistic baseline, the raw input, the saliency and the ridge", "labels"], + ["HaploGraph edges_lift_above_threshold.csv.gz", "CO_OCCURS edges with weight = normalised log lift", "anything about people directly"], + ["block_stats.tsv, boundaries", "Block node features; five of the seven Cluster features; NEXT_BLOCK order; block-gene overlaps", "labels"], + ["phenotypes_real.csv (ancestry, population, sex)", "labels on Individual nodes: the training target for training people, the ground truth for validation and test people, the classes in the Cramer's V tests, the colours in plots", "edges, node features, the SVD (which is label-free), the test split (only its stratification)"], + ["uniprot_chr22.bed", "Gene and Protein nodes with span and isoform count; OVERLAPS by coordinate; ENCODES", "labels"], + ["proteomics matrices (per site, log2)", "harmonised per site and protein into z; MEASURED edge weights; the person's input vector (z plus observed mask); the proteome-only MLP", "cross-site normalisation (the harmoniser never sees two sites at once)"], + ["proteomics metadata (site, age, sex, phenotype)", "site = batch control and the federated partition; phenotype = the case/control target; age = decoder context", "edges or node features"], + ["ground_truth.json", "scoring only: which clusters are causal (saliency precision), which proteins are cis-affected (ridge R-squared)", "any model input"], + ], [0.24, 0.5, 0.26]), + + h(2, "What the trained model gives you, and what is done with it"), + p("A trained run is a folder with model.pt (the weights), metrics.json, history.csv, test_predictions.csv, embedding_individual.npy (and embedding_cluster.npy in v1), and for the raw-input phenotype run saliency_top100.csv. Five products come out of the encoder and each has a consumer:"), + b(["Class probabilities per person: for every one of the 2,548 people, a probability per class (five ancestries, 26 populations, case/control). Consumed by the metrics against held-out labels, by infer.py which writes predictions_all_individuals.csv with a confidence per person, and by the decoder, which quotes the prediction for the person it describes.", + "A 64-number embedding per person and per cluster: the encoder's last hidden layer. Consumed by embeddings.py (silhouette, nearest-neighbour accuracy, PCA plots), by the decoder (the five nearest people in this space are part of the retrieved context) and available for any downstream clustering of people into phenotype groups.", + "A saliency per haploblock cluster: the gradient of the case score with respect to the person's carrier row, averaged over test cases. Consumed by the ground-truth check (how many of the top 20 are planted causal clusters) and by the decoder, which lists the globally salient clusters the person carries.", + "The weights themselves: in the federated setting the weights are the only thing that leaves a site, so the same model definition (federated/model.py) is what NVFlare averages.", + "A timed inference path: infer.py rebuilds the run's inputs, scores the whole graph in one pass (36 ms eager, 4.9 ms with torch.compile on the A100) and writes predictions and embeddings for everyone."]), + + h(2, "Ground truth versus prediction"), + p("Every model is scored on the same 376 held-out people that no model saw during training or model selection. For the real labels the ground truth is the 1000G panel; for the synthetic phenotype the ground truth is the label the generator drew from the planted causal clusters; for the saliency and the ridge it is the list of causal clusters and cis proteins in ground_truth.json; for the federated model it is the same 376 people scored centrally. The confusion tables below are read row = truth, column = prediction."), + f("confusion_matrices.png", "Ground truth versus prediction on the 376 held-out people: real ancestry (GNN, SVD input), the synthetic phenotype from genome plus proteome, from the genome alone, and sex, the negative control.", 0.95), + t(["Ancestry, true / predicted", "AFR", "AMR", "EAS", "EUR", "SAS"], [ + ["AFR (99)", "99", "0", "0", "0", "0"], ["AMR (52)", "1", "48", "0", "3", "0"], ["EAS (76)", "0", "0", "76", "0", "0"], ["EUR (76)", "0", "4", "0", "72", "0"], ["SAS (73)", "0", "0", "0", "0", "73"], + ], [0.35, 0.13, 0.13, 0.13, 0.13, 0.13]), + p("368 of 376 correct (balanced accuracy 0.974). The eight errors are all between AMR and EUR, which is what admixture predicts: the AMR panel populations carry European haplotypes."), + t(["Synthetic phenotype, true / predicted", "case", "control", "correct"], [ + ["genome plus proteome (graph): case (136)", "128", "8", "367 / 376, AUC 0.992"], ["genome plus proteome (graph): control (240)", "1", "239", ""], + ["proteome only (MLP): case (136)", "129", "7", "364 / 376, AUC 0.962"], ["proteome only (MLP): control (240)", "5", "235", ""], + ["genome only (graph): case (136)", "95", "41", "202 / 376, AUC 0.601"], ["genome only (graph): control (240)", "133", "107", ""], + ["sex, negative control: female (188)", "119", "69", "189 / 376, chance"], ["sex, negative control: male (188)", "118", "70", ""], + ], [0.46, 0.13, 0.13, 0.28]), + p("The genome alone gets little more than half right because the planted genomic signal is 20 clusters through a noisy logistic link; the proteome carries most of the signal; the graph that joins the two makes the fewest errors, with a single false case. Sex on an autosome is a coin flip, as it must be."), + t(["Rank", "Cluster (saliency of the combined model)", "Saliency", "Planted causal cluster?"], [ + ["1", "chr22_46902935-46974137_cluster186", "0.110", "yes"], ["2", "chr22_44263132-44288496_cluster19", "0.080", "yes"], ["3", "chr22_25066667-25206817_cluster346", "0.077", "yes"], + ["4", "chr22_46974137-47051453_cluster43", "0.055", "no"], ["5", "chr22_44166469-44185056_cluster231", "0.054", "yes"], ["6", "chr22_47257518-47285160_cluster1", "0.046", "no"], + ["7", "chr22_40032702-40132216_cluster25", "0.046", "no"], ["8", "chr22_49408121-49430759_cluster3", "0.043", "no"], ["9-20", "twelve further clusters", "0.042-0.036", "no"], + ], [0.08, 0.5, 0.14, 0.28]), + p("Four of the top twenty (and the top three outright) are among the 20 planted causal clusters out of 6,551; by chance 0.06 would be. For the genome-to-protein direction the ridge finds 4 of the 20 cis proteins with test R-squared above 0.1 and none of the 440 others. The federated global model, scored on the same 376 people, reaches AUC 0.998 and balanced accuracy 0.963 against the central model's 0.992 and 0.969."), + + h(2, "Where the language model sits, and what it achieves"), + p("The language model is not part of training and makes no prediction. It sits after the encoder, once per person, as the decoder: graphrag_decoder.py walks the graph around one person and collects facts (their ancestry-informative clusters with blocks, genes and proteins; their most extreme protein levels and the block that encodes each; the GNN's prediction; the salient clusters they carry; their five nearest neighbours in the embedding), serialises them as JSON and asks the NIM model (nvidia/nemotron-3-super-120b-a12b) to write a structured report using only those facts and citing every id verbatim. The code then checks each cited id against the context. What this achieves is the last step the README's mission implies, turning numbers into an insight a research team can read: a summary, an ancestry assessment, a phenotype assessment, explicit genome-to-proteome links and caveats, each traceable to graph ids. For HG00103 (EUR, GBR, 994 clusters carried, predicted control) the model answered in 13 seconds from a 4,555-token context with a 1,020-token reply, cited 25 ids and invented none; its links named, for example, cluster chr22_40032702-40132216_cluster151 in the block encoding TNRC6B (protein Q9UPQ9) and cluster chr22_26024448-26060666_cluster35 in the block encoding MYO18B (Q8IUG5), and its caveats stated that the phenotype is synthetic and that ancestry is population structure, not a medical finding. Without an API key the same script prints the exact prompt (dry run), so the retrieval can be inspected offline."), + + t(["Retrieved fact for HG00103", "Value"], [ + ["GNN prediction", "control (truth: control)"], + ["globally salient clusters carried", "chr22_46902935-46974137_cluster186 (0.110), chr22_46974137-47051453_cluster43 (0.055), chr22_40032702-40132216_cluster25 (0.046)"], + ["most ancestry-informative clusters, with block genes and proteins", "cluster151 of block 40032702-40132216 (TNRC6B, Q9UPQ9); cluster545 of 35738379-36060254 (RBFOX2, O43251); cluster35 of 26024448-26060666 (MYO18B, Q8IUG5)"], + ["most extreme protein levels (z)", "H1-0 +2.75, RIMBP3 -2.65, INPP5J +2.53, CRYBA4 +2.52, YWHAH +2.37"], + ["nearest people in the embedding", "HG01767 (EUR), NA12842 (EUR), HG00351 (EUR), HG04002 (SAS), HG00336 (EUR)"], + ], [0.35, 0.65]), + t(["Decoder field", "Text returned (Nemotron 3 Super, 13.3 s, 25 ids cited, 0 unknown)"], [ + ["summary", "Individual HG00103 is of EUR ancestry (GBR population) and carries 994 haploblock clusters. The GNN phenotype prediction for this individual is 'control', which is a synthetic case/control label used for pipeline testing. Several ancestry-informative clusters are carried, showing high enrichment in various populations including EUR, EAS, and SAS."], + ["ancestry_assessment", "The individual's ancestry is EUR (European) as per the 1000G panel data, with population label GBR. This reflects population structure, not a medical finding."], + ["phenotype_assessment", "The GNN-predicted phenotype is 'control', which is a synthetic case/control label used to test the pipeline; it does not reflect a real medical condition."], + ["genome_proteome_links", "cluster151 -> block 40032702-40132216 -> TNRC6B -> Q9UPQ9 (carried, enriched in EAS); cluster545 -> RBFOX2 -> O43251; cluster35 -> MYO18B -> Q8IUG5 (enriched in EUR)"], + ["caveats", "ancestry is population structure, not a medical finding; the phenotype is synthetic; not all clusters have gene or protein annotations; the prediction comes from the embedding and may not match observed protein levels"], + ], [0.22, 0.78]), + p("What the LLM adds is the sentence layer: a readable paragraph, the genome-to-protein chain spelled out per person, the right hedges attached automatically, graceful handling of missing information (for HG00096, who is not in the held-out split, it reported that no prediction exists instead of inventing one), and traceability, because every claim carries an id that the code verifies. It makes no prediction and adds no outside knowledge."), + h(2, "Connecting the dots to the README on branch main"), + p("The team README fixes a mission, three research questions, a chromosome-22 demo scope, five required datasets and a data-integration flowchart. Each maps to something concrete in genomics/:"), + t(["README item", "What it became", "Evidence"], [ + ["Mission: each institution keeps individual-level data locally, trains the same graph model, exchanges only model updates; a server aggregates and returns them", "the person is its own node type, so a site holds only its people and their CARRIES and MEASURED edges while the cluster/block/gene/protein graph is public; NVFlare FedAvg with client.py sending a state dict only", "outputs/federated/chr22/evaluation.json: AUC 0.998 vs central 0.992"], + ["Opening line: a variant-based phenotype-propensity reference graph combined with patient-specific proteomics, what can we learn?", "reference graph = the shared cluster/block/gene/protein layer; patient-specific = Individual nodes with CARRIES and MEASURED edges; learned: integration beats either modality alone, cis effects need sparse models, federation costs nothing when sites are alike", "sections 9, 11, 13"], + ["Background: one gene gives many protein products", "917 UniProt isoform rows collapsed to 460 proteins with an isoform count as a Protein feature; ENCODES keeps gene to protein explicit", "haplokg_proteins.load_protein_bed"], + ["RQ1: connect haploblock genomics to genes and proteomic data in a graph model", "schema v2, five node types, seven edge types, one join key; the README flowchart maps one-to-one: Participant = Individual, Haploblock hash = Cluster, Haploblock = Block, Encoded protein = Gene ENCODES Protein, Measured abundance = MEASURED edge", "section 6; Neo4j browser"], + ["RQ2: can a GNN combine genomic and proteomic information to identify disease-related phenotype clusters", "yes, in both senses of cluster: groups of people (AUC 0.60 genome, 0.96 proteome, 0.99 both; embedding silhouette 0.70) and haploblock clusters tied to the phenotype (saliency top 3 all causal); the person-level graph, with every participant as a node, is what lets the GNN use both modalities", "sections 8, 9; Part II ground truth tables"], + ["RQ3 (aspirational): train across institutions without transferring individual-level data", "done in NVFlare simulation with three sites; real-machine POC is the next step", "section 11"], + ["Demo scope: chromosome 22 first", "everything runs on chr22; nothing is chromosome-specific (CHROM=chr21 make run)", "Makefile"], + ["Dataset 1: haploblock BED and per-individual haploblock hashes", "boundaries and block_stats from haploblocks.org; the hashes were already clustered upstream into the HaploGraph clusters we consume, so no hash was recomputed", "fetch_data.sh"], + ["Dataset 3: gene BED for chr22", "uniprot_chr22.bed (team); Block OVERLAPS Gene by coordinate intersection, 1,063 edges, 29 genes outside every block", "build_kg_v2.py"], + ["Dataset 4: gene-to-protein mapping", "the same BED: UniProt accession per gene, isoforms collapsed; 460 ENCODES edges", "build_kg_v2.py"], + ["Dataset 5: proteomic data for chr22 proteins", "the team's 120-patient synthetic set validated the proteomics plumbing but cannot join the genome (its ids are not 1000G ids and its phenotype has no genomic cause); the joinable set on 2,503 real 1000G ids with ground truth replaces it for integration; real data (Wu 2013, UKB-PPP) is the planned next source", "proteomics_synth_1000g.py"], + ["Data-integration flowchart", "implemented as a PyTorch Geometric HeteroData object plus a Neo4j load, with two additions the README did not list: CO_OCCURS between clusters and NEXT_BLOCK between blocks", "hetero_v2.pt; neo4j_load.py"], + ], [0.3, 0.45, 0.25]), + + h(1, "Part III: The project"), + h(2, "1. Problem statement and goal"), + p("The team's README (branch main) states the mission: develop a proof-of-concept workflow that integrates a known genome graph with proteomics, with a federated approach where each participating institution retains its individual-level data locally and trains the same graph-based model, exchanging only model updates with a coordinating server. Three research questions make it concrete:"), + b(["RQ1: How can haploblock-based genomic information be connected to genes and proteomic data in a graph-based data model?", + "RQ2: Can a graph neural network combine genomic and proteomic information to identify disease-related phenotype clusters?", + "RQ3 (aspirational): Can a graph neural network trained across multiple institutions predict clinical outcomes without transferring individual-level data?"]), + p("Translated into engineering: build a typed graph in which a person links to the haplotype clusters they carry, clusters sit in blocks, blocks overlap genes, genes encode proteins, and the same person links to their measured protein levels (RQ1); train a GNN on it and show that genome plus proteome predicts a phenotype better than either alone while negative controls stay at chance (RQ2); split the people by site, share only the reference graph and the model weights, and show the federated model matches the central one (RQ3). The demo scope fixed by the team is chromosome 22."), + p("What was already available: the mentor (Ben Busby) pointed the team to haploblocks.org, whose data server publishes a ready-made graph of haplotype clusters for the 1000 Genomes people, built for this hackathon at Rigshospitalet's MDxCORE unit. That is what \"we don't have to build the genome graph\" meant; everything downstream of it is ours."), + + h(2, "2. Data sources: what, where, why"), + f("genome_to_graph.png", "From genomes to a graph: blocks, phased haplotypes, MMseqs2 clusters, the carrier matrix and the filters that produce the 6,551 kept clusters.", 0.95), + t(["Data", "Where it comes from", "What it contains", "Why we use it"], [ + ["HaploGraph nodes (nodes.csv.gz)", "data.haploblocks.org/haplograph/1000G/chr22", "248,254 haplotype clusters x 2,548 people; 1 if the person carries the cluster", "The genome graph's node features; becomes our person-to-cluster edges"], + ["HaploGraph edges (edges_lift_above_threshold.csv.gz)", "same server", "187,030 cluster pairs that co-occur in people more than chance (lift >= 5), with weight and lift", "Co-occurrence structure; the raw edges.csv is dominated by one near-universal haplotype and is not used"], + ["Block statistics (block_stats.tsv, boundaries)", "same server", "669 blocks on chr22: coordinates, length, number of clusters, entropy, dominance, singletons", "Block nodes and their features"], + ["Phenotypes (phenotypes_real.csv)", "1000G panel via IGSR, republished with the graph", "ancestry, population, sex for 2,503 of the 2,548 people", "The only real labels that exist; ancestry and population are targets, sex is the negative control"], + ["Protein coordinates (uniprot_chr22.bed)", "UCSC UniProt track, prepared by Friederike", "917 protein isoform rows -> 460 proteins, 458 genes, with genomic spans", "Block-to-gene overlaps and gene-to-protein mapping in one file"], + ["Synthetic proteomics, team version", "proteomics/generate_synthetic_proteomics.py (Nolan; regenerated on main with 4,000 samples)", "3 sites, first 40 then 1,333-1,334 patients each (ids SITE1_PT0001...), age, sex, case/control, log2 intensities for 460 proteins", "The team's proteomics analyses on main run on it; for the genome join a version keyed to 1000G ids was needed (next row)"], + ["Synthetic proteomics, joinable version", "genomics/proteomics_synth_1000g.py (this work)", "2,503 real 1000G ids, 3 mixed-ancestry sites, phenotype driven by 20 causal clusters, cis effects, batch shift, missingness; ground truth saved", "Lets the integration be scored against a known answer"], + ["Planned real proteomics", "Wu et al. 2013 (Nature): 95 HapMap LCLs with 1000G ids; UKB-PPP pQTL summary statistics on AWS Open Data", "per-person protein levels; variant-to-protein effect sizes", "Real join on the same ids; real cluster-to-protein propensity edges"], + ], [0.2, 0.22, 0.3, 0.28]), + p("How the genome graph was made upstream (haploblocks.org pipeline, Kubica et al. 2025): (1) recombination-rate peaks define haploblocks; (2) each person's two phased haplotype sequences are cut out per block from the 1000G VCF; (3) all haplotypes of a block are merged into one FASTA; (4) MMseqs2 clusters near-identical haplotypes, giving each haplotype a cluster id; (5) a compact hash encodes strand, chromosome, block, cluster and variants. A sixth step (haploblock-graph-builder) produced the HaploGraph: nodes are clusters, the node feature is the 0/1 vector over people, edges join clusters that co-occur in the same people. We did not rerun these steps; we verified their outputs agree (669 blocks in every file, identical per-block cluster counts in all 669, 248,254 clusters both ways) and consumed them."), + + h(2, "3. The data pipeline"), + f("workflow_pipeline.png", "The two pipeline chains (v1 genome graph, v2 proteomics integration), one script per stage.", 0.95), + p("Everything is a numbered script under genomics/, driven by a Makefile so a fresh clone runs with two commands. Downloaded inputs and generated outputs are git-ignored and regenerated; secrets are read from environment variables and never stored in the repository."), + t(["Step", "Script", "What it does", "Output"], [ + ["1", "fetch_data.sh", "downloads the chr22 HaploGraph files, phenotypes, block statistics; verifies md5 checksums", "data/"], + ["2", "build_kg.py (haplokg.py)", "streams nodes.csv.gz as int8 in 8,192-row chunks into a sparse matrix; filters clusters; joins phenotypes; maps edges; writes tables and the PyG graph", "outputs/kg/chr22/ (carries.npz, hetero.pt)"], + ["3", "cooccurrence_analysis.py", "cluster x phenotype tests, edge x phenotype similarity, per-block informativeness, plots", "outputs/cooccurrence/chr22/"], + ["4", "baseline.py", "logistic regression on the carrier matrix; writes the shared train/val/test split", "outputs/baseline/, outputs/splits/"], + ["5", "graph_explore.py", "NetworkX statistics, GraphML export, region and chromosome-wide plots", "outputs/graph/chr22/"], + ["6", "train_gnn.py", "the genome-only GNN with SVD / Node2Vec / learned / raw inputs", "outputs/gnn/chr22/"], + ["7", "embeddings.py", "quality of SVD and GNN embeddings: silhouette, nearest-neighbour accuracy, PCA plots", "outputs/embeddings/chr22/"], + ["v2.1", "proteomics_synth_1000g.py", "synthetic proteomics on the real 1000G ids with saved ground truth", "outputs/proteomics_synth/chr22/"], + ["v2.2", "build_kg_v2.py (haplokg_proteins.py)", "adds genes, proteins, block-gene overlaps, harmonised measurements", "outputs/kg/chr22/hetero_v2.pt"], + ["v2.3", "eda.py", "the full exploratory report with tables and plots", "outputs/eda/chr22/EDA.md"], + ["v2.4", "train_gnn_v2.py", "genome / proteome / both ablations, controls, saliency, proteome regression", "outputs/gnn_v2/chr22/"], + ["v2.5", "proteome_linear_baseline.py", "per-protein ridge: can the genome predict each protein?", "outputs/gnn_v2/chr22/proteome_ridge_baseline/"], + ["v2.6", "graphrag_decoder.py", "graph retrieval + NVIDIA NIM LLM -> cited insight per person", "outputs/graphrag/chr22/"], + ["v2.7", "federated/job.py, client.py, model.py, evaluate_global.py", "NVFlare FedAvg over 3 sites and central scoring of the global model", "outputs/federated/chr22/"], + ["-", "infer.py", "inference benchmark: eager vs torch.compile vs TensorRT", "outputs/gnn/.../inference/"], + ["-", "neo4j_load.py, docker-compose.yml", "loads the graph into Neo4j for browsing", "localhost:7474"], + ], [0.07, 0.25, 0.45, 0.23]), + p("Memory arithmetic that shaped step 2: the node file is 1.3 GB as text; as a dense 64-bit matrix it would be 5 GB, as dense int8 632 MB, as a sparse matrix with 2.8 million non-zeros about 30 MB. Streaming chunks into sparse form keeps the whole build under 1 GB and 26 seconds on a laptop."), + p("The cluster filter: a cluster is kept if between 25 and N-25 people carry it (N = 2,548). A cluster carried by 2,540 people is as uninformative as one carried by 8; both have only 8 people on the informative side. This mirrors the HaploGraph's own symmetric edge filter and reproduces its node set exactly: 248,254 clusters become 6,551 and not one of the 187,030 edges loses an endpoint. 176,903 of the dropped clusters are singletons (one haplotype)."), + + h(2, "4. Exploratory data analysis"), + p("The EDA report (outputs/eda/chr22/EDA.md, generated by eda.py) has eight sections: provenance, individuals, haploblocks, clusters, co-occurrence, genes and proteins, proteomics, genome-proteome. The key measurements:"), + t(["Aspect", "Measurement"], [ + ["Individuals", "2,548 people; 2,503 labelled (AFR 660, EAS 504, EUR 503, SAS 489, AMR 347), 45 unlabelled kept as nodes; 26 populations of 61-113 people; sex balanced within each ancestry; split train 1,752 / val 375 / test 376"], + ["Blocks", "669 tiling 17.1-50.2 Mb with no gaps; length median 29.7 kb (5-95%: 8.9-151 kb, longest 775 kb); clusters per block median 214, max 2,806; singleton rate median 0.62; entropy median 3.15; longer blocks hold more clusters (Spearman 0.37)"], + ["Clusters", "6,551 kept of 248,254; each person carries ~928 (two haplotypes x 669 blocks minus filtered); the person x cluster matrix is 14% dense"], + ["Co-occurrence edges", "187,030 with lift >= 5 (median 5.5, max 76); only 121 join clusters of the same block; median distance between endpoints tens of Mb; 4,344 clusters have at least one edge; degree up to 358; the top hubs are all AFR-enriched rare clusters (the mega-hub artefact)"], + ["Genes / proteins", "458 genes, 460 proteins, 1,063 block-gene overlaps; 29 genes outside every block (chromosome ends and gaps); 117 genes span two blocks; up to 14 genes in one block"], + ["Proteomics (synthetic)", "2,503 people x 460 proteins, 3 sites of ~835; log2 range 4.9-17.0; 7.4% missing overall, up to 16% for the least abundant proteins (missing-not-at-random at the detection limit); between-site shift 0.48 log2 before harmonisation, 0.00 after; 227 proteins keep a phenotype association at FDR 5%"], + ["Genome-proteome", "for the 20 ground-truth causal clusters, the correlation between carrying the cluster and its cis protein gives r-squared up to 0.54, mean 0.17, 10 of 20 above 0.1 - the ceiling any genome-to-proteome model can reach on this data"], + ], [0.2, 0.8]), + f("eda_populations.png", "1000G individuals per population, coloured by continental ancestry.", 0.9), + f("eda_blocks.png", "Haploblocks on chr22: length distribution, clusters per block versus length, Shannon entropy of clusters along the chromosome.", 0.95), + f("eda_clusters.png", "Kept clusters: carriers per cluster (log scale) and clusters carried per person by ancestry.", 0.9), + f("eda_cooccurrence.png", "Co-occurrence edges: lift, distance between endpoints, degree.", 0.95), + f("eda_proteomics.png", "Synthetic proteomics: intensity distribution, missingness rising for low-abundance proteins, batch effect between sites before and after the harmoniser.", 0.95), + f("eda_blocks_populations.png", "Blocks with one dominant haplotype versus many rare ones (left); sex within each of the 26 populations (right).", 0.95), + f("sites_composition.png", "The three federated sites: people per site by ancestry (mixed by design), case prevalence per site, missing protein values per site.", 0.95), + f("sites_batch_effect.png", "Per-site protein medians before and after the harmoniser: the batch offsets vanish.", 0.95), + f("ground_truth_effects.png", "Ground truth of the synthetic proteome: causal-cluster effects on the phenotype, cis effects on proteins, phenotype effects on the 70 responsive proteins, carrier frequencies of the causal clusters.", 0.95), + + h(2, "5. Statistics: does the graph co-occur with phenotypes?"), + p("Before any model we asked whether the graph carries phenotype information at all. For every cluster and each label we built the 2 x k table of carrier status versus class and computed a chi-square test with Cramer's V (a 0-1 effect size; for a 2 x k table V = sqrt(chi-square / N)), then Benjamini-Hochberg false-discovery-rate correction. All 6,551 tests run in one sparse matrix product."), + b(["Ancestry: 6,470 of 6,551 clusters (98.8%) are associated at FDR 5%; median V 0.20, maximum 0.81.", + "Population: 6,378 clusters (97.4%).", + "Sex: 0 clusters; median V 0.013, maximum 0.03. This is the negative control: chr22 is autosomal, so a method that found sex signal would be fitting noise."]), + p("For the edges we compared the two endpoint clusters' ancestry-enrichment profiles: cosine similarity 0.86 for real edges versus 0.22 for degree-preserving shuffled pairs; 86% of edges join clusters enriched in the same ancestry (42% expected), and the similarity rises with lift (0.84 in the lowest lift quartile to 0.90 in the highest). Interpretation: the co-occurrence graph is largely population structure, long-range co-inheritance within ancestries rather than physical linkage."), + f("cramers_v.png", "How strongly each cluster tracks a phenotype: ancestry and population carry signal, sex (the control) does not.", 0.85), + f("informativeness.png", "Maximum Cramer's V per block along chr22 for ancestry versus sex.", 0.95), + f("edge_similarity.png", "Co-occurring clusters share ancestry profiles: real lift edges versus shuffled pairs.", 0.85), + f("informative_clusters_heatmap.png", "The 30 most ancestry-informative clusters and the fraction of each ancestry that carries them.", 0.6), + + h(2, "6. The knowledge graph"), + f("schema_diagram.png", "Node and edge types of the knowledge graph with chr22 counts.", 0.95), + p("The HaploGraph has a single node type (cluster) and carries people only as a feature vector. We turned that vector into a second node type, the person, because a person is what phenotypes and protein measurements belong to, and what a hospital owns. The join is the 1000G sample id: the same string (for example HG00096) is a column header in nodes.csv.gz, a row key in phenotypes_real.csv and a column in the proteomics matrices."), + t(["Node type", "Properties", "Count", "Source"], [ + ["Individual", "ancestry, population, sex, site, phenotype, age", "2,548", "nodes.csv.gz columns + phenotypes + proteomics metadata"], + ["Cluster", "support, block statistics, SVD vector", "6,551", "nodes.csv.gz rows after the filter"], + ["Block", "length, n_clusters, entropy, dominance, singleton rate", "669", "block_stats.tsv"], + ["Gene", "coordinates, number of proteins", "458", "uniprot_chr22.bed"], + ["Protein", "coordinates, number of isoforms", "460", "uniprot_chr22.bed"], + ], [0.15, 0.4, 0.12, 0.33]), + t(["Edge type", "Meaning", "Count"], [ + ["Individual -CARRIES-> Cluster", "the person carries this haplotype cluster (the 1s of the matrix)", "2,365,574"], + ["Cluster -CO_OCCURS{weight, lift}-> Cluster", "the two clusters co-occur in people more than chance", "187,030"], + ["Cluster -IN_BLOCK-> Block; Block -NEXT_BLOCK-> Block", "position on the chromosome", "6,551; 668"], + ["Block -OVERLAPS-> Gene; Gene -ENCODES-> Protein", "coordinate intersection; protein product", "1,063; 460"], + ["Individual -MEASURED{log2, z}-> Protein", "one edge per observed protein level; missing values create no edge", "1,065,712"], + ], [0.35, 0.5, 0.15]), + p("Two design rules matter. First, phenotypes are properties of the person node, never nodes or edges: if a Phenotype node were connected to the person, a two-layer GNN would read the label from its neighbour and report a meaningless 100%. Second, because people are their own node type, the site boundary is a clean cut: the cluster, block, gene and protein graph is public and identical at every site, while a site holds only its people and their CARRIES and MEASURED edges."), + p("The graph is stored as a PyTorch Geometric HeteroData object (hetero.pt, hetero_v2.pt) for modelling and loaded into Neo4j community edition (docker compose up neo4j; neo4j_load.py) for browsing at localhost:7474, where a query such as MATCH (i:Individual {id:'HG00096'})-[:CARRIES]->(c:Cluster)-[:IN_BLOCK]->(b:Block) RETURN i,c,b shows a person's haplotypes with their blocks. NetworkX (with the nx-cugraph GPU backend in the image) produces statistics and GraphML for Gephi."), + f("graph_region.png", "One region of chr22 (the densest published island): clusters as nodes coloured by the ancestry they are enriched in, co-occurrence edges weighted by lift.", 0.8), + f("graph_edge_positions.png", "All 187,030 co-occurrence edges plotted by the positions of their two endpoints: block structure and long-range population structure.", 0.7), + f("person_neighbourhood.png", "A real person's neighbourhood in the graph (HG00103), as the decoder retrieves it: clusters, blocks, genes, proteins, extreme protein levels and nearest neighbours in the embedding.", 0.85), + + h(2, "7. From graph to embeddings"), + p("A GNN needs a starting vector for every node. Cluster and block nodes use their statistics (support, block length, entropy, dominance, and so on), standardised. People and clusters together get a truncated singular value decomposition of the carrier matrix M (2,548 x 6,551): M is approximated as U S V-transpose with 32 components; the rows of U S are 32-number vectors for people and the rows of V S are 32-number vectors for clusters, in one shared space, so a person sits near the clusters they carry and near people with similar haplotypes. No labels are used, so nothing can leak into the test set. Alternatives implemented and compared: Node2Vec random-walk embeddings on the person-cluster graph (GPU, via pyg-lib; with 50 pretraining epochs it reaches 0.950 on ancestry against 0.974 for SVD, with the same 0.70 silhouette), free learned embeddings (0.585), and the raw 6,551-long carrier row."), + p("Where the phenotypes enter: only as training targets. The GNN is trained to predict the label from a person's neighbourhood; the loss reshapes all weights so the learned 64-number vectors separate the phenotype groups. The quality of a space is measured by the silhouette score (how compact and separated the groups are) and by a 5-nearest-neighbour classifier: plain SVD scores silhouette 0.06 and 5-NN accuracy 0.90 by ancestry; the GNN's hidden layer scores 0.70 and 0.97. That 0.06 to 0.70 is the GNN's contribution. With free learned embeddings instead of SVD initialisation the GNN reaches only 0.585 balanced accuracy on ancestry: the starting embedding matters more than the architecture."), + f("embedding_individuals.png", "The GNN's 64-dimensional embedding of people projected to two dimensions: five ancestry groups separate cleanly (AMR spread between EUR and AFR, as admixture predicts); coloured by sex the same points are fully mixed.", 0.95), + f("embedding_clusters.png", "The same model's embedding of clusters, coloured by the ancestry each cluster is enriched in.", 0.7), + f("embedding_quality.png", "Embedding quality: silhouette by ancestry and 5-nearest-neighbour accuracy for SVD and for each GNN's hidden layer; sex stays at chance.", 0.9), + f("init_comparison.png", "The same GNN with different starting embeddings for the person nodes: SVD-32, Node2Vec-32 (50 pretraining epochs, A100), free learned embeddings and the raw carrier row, against the logistic-regression line.", 0.9), + + h(2, "8. The encoder: a heterogeneous graph neural network"), + p("Architecture (PyTorch Geometric HeteroConv, two layers, hidden size 64): a linear projection per node type to 64 numbers; then two rounds of message passing in which, for every relation, SAGEConv adds the mean of a node's neighbours to its own state (carries, in_block, next_block, overlaps, encodes, each in both directions) and GraphConv adds an edge-weighted mean for co_occurs (weight = normalised log lift) and measured (weight = the protein's harmonised z-score, signed, so a high protein pushes positively and a low one negatively); after each round LayerNorm, a residual connection, ReLU and dropout 0.3; a linear head from the person's 64 numbers to class scores. After round one a person has absorbed their clusters (and proteins); after round two, what those clusters co-occur with, their blocks, and the genes and proteins in those blocks."), + p("Training: class-weighted cross-entropy on training people only (weights inversely proportional to class size so AMR and small populations count); Adam with learning rate 0.005 and weight decay 0.0005; early stopping on validation balanced accuracy with patience 30, best weights restored. Full-batch: one epoch is one pass over the whole graph (about 2.4 million CARRIES, 0.37 million CO_OCCURS and 1.07 million MEASURED edges), 0.95 s on an M2 laptop CPU and 0.10 s on an A100 GPU. About 105 thousand parameters with SVD input, about 500 thousand with the raw carrier row."), + t(["Target (real labels)", "Classes", "Logistic regression", "GNN"], [ + ["ancestry", "5", "0.977", "0.974 (SVD input)"], + ["population", "26", "0.614", "0.611 (raw input; 0.437 with SVD-32)"], + ["sex (negative control)", "2", "0.463", "0.503"], + ], [0.3, 0.15, 0.25, 0.3]), + p("Reading: balanced accuracy on the held-out 376 people. On chr22 alone ancestry and population are almost linear functions of which clusters a person carries, so the GNN ties the linear baseline on accuracy; it wins on the embedding space and, as the next section shows, on integration."), + h(3, "Metrics, defined"), + b(["Accuracy: correct predictions divided by all predictions. Misleading when classes are unequal (predicting control for everyone scores 64% on the phenotype).", + "Balanced accuracy: the mean over classes of the recall of that class (correct in class k divided by the number truly in class k). Chance is 1/k: 0.20 for ancestry, 0.038 for population, 0.5 for sex, 0.33 for site.", + "Macro-F1: for each class the harmonic mean of precision and recall, averaged over classes.", + "ROC-AUC: the probability that a randomly chosen case receives a higher case score than a randomly chosen control; 0.5 is chance, 1.0 is a perfect ranking. Only defined for two classes.", + "R-squared per protein: 1 minus the residual sum of squares divided by the total sum of squares, on test people with an observed value (at least five); negative means worse than predicting the mean.", + "Silhouette: for each person, the mean distance to their own group minus the mean distance to the nearest other group, scaled to -1..1; averaged. Higher means tighter, better separated groups in the embedding.", + "5-nearest-neighbour accuracy: label a test person by majority vote of the five nearest training people in the embedding; balanced accuracy of that vote.", + "Precision at 20: of the twenty clusters with the highest saliency, the fraction that are planted causal clusters; chance is 20 / 6,551 x 20 = 0.06.", + "Cramer's V: sqrt(chi-square / N) for a 2 x k table; 0 means the cluster is independent of the label, 1 means it determines it.", + "Early stopping and model selection use the validation people only; every number reported in this document is on the test people."]), + f("baseline_vs_gnn.png", "Real labels: logistic regression versus the GNN; sex, the negative control, stays at chance for both.", 0.7), + + h(2, "9. Integrating the proteome (schema v2)"), + p("The synthetic proteome we generated is keyed to the real 1000G ids and has a saved ground truth (ground_truth.json). Per person: a site assigned at random within each ancestry (so site is a pure batch effect), the real sex, a random age; a case/control phenotype whose log-odds is a weighted sum over 20 causal haploblock clusters plus a small age term, calibrated to 38% cases; each causal cluster also shifts one protein encoded in its own block (a cis effect); 15% of proteins respond to the phenotype, all respond to age and sex; a per-site batch shift; and missingness that increases toward the detection limit. The first version shifted every protein with the phenotype and every model scored 1.0, the same trap an earlier team result fell into, so the signal was made sparse."), + p("Before entering the graph, protein levels pass a harmoniser: within each site and protein, z = (value - median) / (1.4826 x median absolute deviation). Computed from each site's own samples, it removes the between-site shift completely (0.48 to 0.00 log2) while 227 proteins keep their phenotype association. Missing values stay missing; they simply create no MEASURED edge and are never imputed as zero."), + t(["Model", "Input to the person node", "Graph relations", "AUC", "Balanced accuracy"], [ + ["genome only", "SVD-32 or raw carrier row", "genome relations", "0.60-0.63", "0.57-0.62"], + ["proteome only (MLP, no graph)", "harmonised z + observed mask (920 numbers)", "none", "0.96", "0.96"], + ["genome + proteome (graph)", "both", "all twelve relations", "0.99", "0.97"], + ["site (batch control, full graph)", "both", "all", "-", "0.30 (chance 0.33)"], + ["ancestry (full graph)", "both", "all", "-", "0.90"], + ], [0.28, 0.27, 0.2, 0.1, 0.15]), + p("The graph adds signal on top of the proteome, and the site control shows none of it is batch. A gradient saliency on the raw carrier input ranks clusters by their influence on the case score: 3-4 of the top 20 are ground-truth causal clusters (1.2 expected by chance). One negative result is kept deliberately: predicting the proteome from the genome embedding gives R-squared near zero even for cis-affected proteins, whereas a per-protein ridge regression on the carrier row recovers the strong cis effects (4 of 20 cis proteins with test R-squared above 0.1, 0 of 440 others). The GNN is the integration and embedding tool; discovering single-cluster cis effects needs sparse per-protein models or a prior from published protein-QTL data."), + f("results_modalities.png", "Held-out AUC and balanced accuracy for the synthetic phenotype by modality; the site control sits at chance.", 0.9), + f("training_curves.png", "Training loss and validation balanced accuracy per epoch for the v2 runs (early stopping picks the best validation epoch).", 0.95), + f("saliency_top20.png", "Cluster saliency of the combined model: ground-truth causal clusters (teal) among the top 20.", 0.95), + f("ridge_r2.png", "Per-protein ridge from the genome: only the strong cis effects are recoverable.", 0.85), + + h(2, "10. The decoder: GraphRAG with an NVIDIA NIM language model"), + p("For one person the decoder retrieves, deterministically and only from the graph: the profile (ancestry, population, sex, site, age; the true phenotype withheld), the GNN's prediction, the globally salient clusters the person carries, their eight most ancestry-informative clusters with block, genes and proteins, their eight most extreme protein levels with the encoding block and whether that block holds a notable cluster, and their five nearest people in the GNN embedding. This is serialised as JSON (about 4,500 tokens) and sent to nvidia/nemotron-3-super-120b-a12b through NVIDIA's OpenAI-compatible NIM endpoint with reasoning_effort set to none (otherwise this reasoning model thinks inline and exhausts the token budget before answering) under a system prompt that forbids inventing entities and requires every id to be cited verbatim. The reply's cited ids are checked against the context before anything is written. For person HG00103 the model answered in 13 seconds with 25 cited ids and none unknown, producing a summary, ancestry and phenotype assessments, genome-proteome links (for example cluster chr22_40032702-40132216_cluster151 in the block encoding TNRC6B / Q9UPQ9) and caveats stating that the phenotype is synthetic and ancestry is population structure, not a medical finding."), + + h(2, "11. Federated learning with NVFlare"), + f("federated_topology.png", "Federated topology: three sites with private people and edges, one shared public graph, weights only to the server.", 0.95), + p("Partition: site s receives the people whose site code is s (835 / 835 / 833, mixed ancestry). Its graph is PyG's HeteroData.subgraph restricted to those people: their nodes are re-indexed, their CARRIES and MEASURED edges kept, and every other node type and edge kept whole because those are public. Site 1, for instance, holds 835 people, 774,753 CARRIES and 355,294 MEASURED edges. The person's input (carrier row plus harmonised protein vector and mask) is computed from the site's own rows; no cross-site preprocessing exists because the harmoniser is already per site."), + p("Mechanics (NVFlare 2.9, FedAvgRecipe with the PyTorch Client API): job.py derives every model constructor value from the public graph into model_args.json so the server and all clients build byte-identical models. Each client runs flare.init(), builds its site graph once, then loops: receive the global weights, evaluate them on its own validation and test people, train five full-batch epochs on its own training people, send back the weights, the metrics and the number of optimizer steps. The server averages the weights (weighted by steps, equal here) and selects the best global model by validation balanced accuracy. What crosses the site boundary per round: one state dict of about half a million numbers and five scalars; no rows, no protein values, no embeddings of people."), + t(["Model", "Balanced accuracy", "AUC", "Evaluated on"], [ + ["federated global model, 10 rounds x 5 epochs", "0.90", "0.955", "the same 376 held-out people"], + ["federated global model, 30 rounds x 5 epochs", "0.963", "0.998", "the same 376 held-out people"], + ["central model (train_gnn_v2, both/raw)", "0.969", "0.992", "the same 376 held-out people"], + ["federated model per site (30 rounds)", "0.983 / 0.981 / 0.933", "1.000 / 0.999 / 0.997", "each site's own held-out people"], + ], [0.4, 0.2, 0.15, 0.25]), + f("federated_rounds.png", "FedAvg convergence: the global model's AUC and validation balanced accuracy at each site, per round, against the central model.", 0.95), + f("federated_vs_central.png", "Central versus federated (10 and 30 rounds) on the same held-out people, and the federated model on each site's own held-out people.", 0.9), + p("Ten rounds (50 local steps) were short of the central run's ~80 epochs; thirty rounds converge to the central model's level. Federated training loses nothing here because the sites are random draws of the same population by construction; with ancestry-pure sites the averaging would have to fight client drift, which is the next experiment."), + h(3, "With and without federation, on this data"), + p("federated/local_only.py trains each site alone on its own people for the same 150 optimizer steps the federated clients used, then scores that lone model on its own held-out people and on the other sites' held-out people. On this synthetic data a lone site already does well, because 835 people and a strong proteome signal are enough: own-site AUC SITE1 0.998, SITE2 0.986, SITE3 0.988; the worst transfer of a lone model to another site's people is SITE1 0.949, SITE2 0.969, SITE3 0.997. The federated global model scores SITE1 0.999, SITE2 1.000, SITE3 0.968 on the same per-site people and 0.987 to 0.998 on the pooled held-out set across runs. With a smaller budget the picture changes: at 50 optimizer steps per site (a 10-round run inside the Docker image on the A100) a lone site reaches only 0.88 to 0.92 while the federated model reaches 0.989 to 0.998 on the same per-site people, because the averaged weights have effectively seen every site's people. So federation costs nothing when a site has enough data and budget and helps clearly when it does not; in both cases what it solves is the constraint, not the score: one shared model, trained on everyone, with no row leaving any site. A larger gain is expected when sites differ systematically (ancestry-pure hospitals, different protein panels), which is the ancestry-partitioned experiment listed under next steps."), + f("federated_site_alone.png", "With and without federation: each site trained alone on its own people versus the federated global model, scored on the same held-out people per site (A100 re-run).", 0.8), + t(["Site (held-out n)", "Alone, 150 steps", "Federated, 150 steps", "Alone, 50 steps", "Federated, 50 steps"], [ + ["SITE1 (111)", "0.998 / 0.952", "0.999 / 0.983", "0.881 / 0.543", "0.998 / 0.950"], + ["SITE2 (136)", "0.986 / 0.852", "1.000 / 0.991", "0.902 / 0.772", "0.991 / 0.907"], + ["SITE3 (129)", "0.988 / 0.923", "0.968 / 0.933", "0.915 / 0.830", "0.989 / 0.865"], + ["average over sites", "0.991 / 0.909", "0.989 / 0.969", "0.899 / 0.715", "0.993 / 0.907"], + ["pooled 376 people, federated global model", "-", "0.987 / 0.967", "-", "0.991 / 0.901"], + ["pooled 376 people, central model", "0.995 / 0.960", "", "0.997 / 0.959", ""], + ], [0.32, 0.17, 0.17, 0.17, 0.17]), + p("AUC / balanced accuracy on each site's own held-out people. Transfer of a lone model to other sites' people (150 steps, AUC): the SITE1 model scores 0.982 on SITE2 people and 0.949 on SITE3 people; SITE2: 0.993 and 0.969; SITE3: 0.997 and 0.998; the federated model 0.999 / 1.000 / 0.968 on the three sites."), + + h(3, "Why federated, and what kind"), + p("The mission sentence of the README is a statement about where data lives: a hospital may compute on its patients but may not ship their rows. This is horizontal federated learning: every site has the same columns (the same graph schema, the same protein panel) and different rows (different people). It is not a split of the genome across sites (each site would then hold part of every person, which is vertical federation and a different problem) and not a way to add chromosomes: another chromosome is another shared reference graph, added at every site at once. What is gained is the ability to train on all 2,503 people while each site only ever reads its 835; what is lost, in this experiment, is nothing measurable, because the sites are alike. A real deployment would report each site's own held-out score and never assemble a central test set."), + h(3, "Mechanics in detail"), + b(["Model definition (federated/model.py): the same architecture as train_gnn_v2, copied into a self-contained file so the NVFlare server can import it without the pipeline. Its constructor reads model_args.json (input sizes per node type, the twelve relations, hidden 64, two layers, dropout 0.3, mean aggregation) so server and clients build identical state dicts.", + "Job (federated/job.py): FedAvgRecipe(name, model class path and args, min_clients 3, num_rounds, train_script client.py, train_args, key_metric val_balanced_accuracy, server_expected_format PYTORCH); add_decomposers registers TensorDecomposer so tensors travel natively; add_server_file ships model.py to the server; SimEnv(num_clients 3, workspace_root) runs the three sites as threads on one machine; recipe.execute(env) writes the job and runs it.", + "Client (federated/client.py): flare.init(); the site name (site-1, site-2, site-3) selects the site code; HeteroData.subgraph keeps that site's people and their CARRIES and MEASURED edges and leaves the public node types whole; then the loop: flare.receive() gives the global weights and the round number; evaluate them on the site's validation and test people; if the task is evaluate-only, send metrics; otherwise train five full-batch epochs (five optimizer steps) with the site's own class weights and send FLModel(params = state dict on CPU, metrics, meta NUM_STEPS_CURRENT_ROUND = 5).", + "Server: after each round the global weights become the step-weighted mean of the three state dicts (weights equal here because every site does five steps; aggregation_weights can weight by site size); the best global model by validation balanced accuracy is kept as best_FL_global_model.pt, the final one as FL_global_model.pt, both under the workspace's app_server folder.", + "What crosses the boundary per round per site: one state dict of about 0.5 million floats and five scalars. No carrier row, no protein value, no person embedding. Weight updates can in principle leak information about training data; NVFlare offers differential privacy and homomorphic-encryption filters for that, and neither was needed for a simulation.", + "Evaluation (federated/evaluate_global.py): loads FL_global_model.pt, rebuilds the full graph exactly as the central run did, scores the same 376 held-out people and each site's own held-out people, and writes evaluation.json next to the central numbers.", + "From simulation to real machines: NVFlare's POC mode starts a server and clients as separate processes (or machines) with the same job; each client would run client.py against its own kg-dir and split; the L4 and A100 instances could be two such sites. Adding a real fourth site means: its own people with 1000G-style ids, its own protein matrix keyed by those ids, its own harmoniser pass, and the shared graph files copied over."]), + + h(2, "12. System design, tech stack and deployment"), + f("architecture_slide.png", "System architecture on one page: private hospital sites, the shared reference graph, the GNN encoder trained through the NVFlare server, and the outputs feeding the LLM decoder.", 0.98), + t(["Layer", "Choice", "Notes"], [ + ["Language and data", "Python 3.13, pandas 3.0, scipy 1.18 (sparse), scikit-learn 1.9", "pinned in requirements.txt"], + ["Graph learning", "PyTorch 2.14, PyTorch Geometric 2.8, pyg-lib (random walks for Node2Vec)", "torch_cluster is deprecated in favour of pyg-lib; that caused the first image-build failure"], + ["Graph tooling", "NetworkX 3.6, nx-cugraph (GPU dispatch), Neo4j 5.26 community in Docker", "graph statistics, GraphML, browsing"], + ["Container", "pytorch/pytorch:2.14.0-cuda12.6-cudnn9-runtime base, PIP_BREAK_SYSTEM_PACKAGES=1, torch-tensorrt 2.14", "10.5 GB image; runs on CPU when no GPU is present"], + ["Compute", "Mac M2 CPU for development; NVIDIA Brev: L4 24 GB (GCP, $0.85/h) and A100 80 GB (Crusoe, $1.98/h)", "brev_deploy.sh creates or reuses an instance, uploads, builds natively, runs, copies outputs back"], + ["Inference", "torch.compile (inductor): 36 ms -> 4.9 ms per full graph, logits equal to within 0.0001", "Torch-TensorRT could not compile this scatter-heavy GNN within 3 hours and is not claimed"], + ["LLM", "NVIDIA NIM, nvidia/nemotron-3-super-120b-a12b", "key in the environment only; a local NIM container on the A100 would keep patient context on site"], + ["Federated", "NVFlare 2.9 FedAvgRecipe, SimEnv simulator; federated/local_only.py for the site-alone comparison", "3 clients as threads on one machine; POC mode across real machines is the next step"], + ["Packaging", "Makefile (setup, run, run-v2, eda, decode, federated, docker, docker-run-v2, docker-federated, brev, report), setup.sh, run_all.sh, run_v2.sh, config.py + .env.example, 11 unit tests", "clone-and-run on laptop or GPU; the image is the whole solution (data mounted, .env passed)"], + ], [0.18, 0.47, 0.35]), + p("Runtime, measured on 18 September from an empty folder (the verification runs of section 13):"), + t(["Stage", "Laptop, Apple M2 CPU", "A100 80 GB (Docker)"], [ + ["environment: venv or image, pinned dependencies, unit tests", "28 s (uv)", "3 min image build"], + ["v1: download, knowledge graph, statistics, baseline, plots, three GNN runs, embedding quality", "6 min", "3 min"], + ["v2: synthetic proteomics, graph v2, EDA, six GNN runs, ridge, decoder", "19 min", "4 min"], + ["federated: 30 rounds x 5 local epochs, central scoring, site-alone comparison", "4 min", "2 min"], + ["whole chain", "about 30 min", "about 14 min including the image build"], + ["one GNN training run", "30 s to 3 min", "1 to 16 s"], + ["one full-graph inference pass (2,548 people)", "-", "36 ms eager, 4.9 ms compiled"], + ["one decoder call (NIM, remote)", "4 to 13 s", "same"], + ], [0.5, 0.25, 0.25]), + f("compute_benchmarks.png", "Training epoch time on the laptop CPU versus the A100, and full-graph inference eager versus torch.compile on the A100.", 0.85), + h(3, "The Docker image, layer by layer"), + p("The Dockerfile starts from pytorch/pytorch:2.14.0-cuda12.6-cudnn9-runtime (PyTorch with CUDA 12.6 already inside), sets PIP_BREAK_SYSTEM_PACKAGES=1 because the base image's Python is system-managed, installs curl, copies requirements.txt and installs the pinned libraries, then tries three optional extras and prints a clear fallback message if any is unavailable for this torch build: pyg_lib from the PyG wheel index (Node2Vec random walks; without it the code falls back to SVD), nx-cugraph from NVIDIA's index (NetworkX dispatches to cuGraph on the GPU; without it NetworkX runs on the CPU), and torch-tensorrt from the PyTorch cu126 index. It copies the code, runs the unit tests as part of the build so a broken image cannot be produced, and defaults to running run_all.sh. Data and outputs are bind-mounted at run time, so the image never contains data. docker-compose.yml adds Neo4j 5.26 community with a persistent volume and the same pipeline image with GPU reservation."), + h(3, "GPU deployment on Brev, step by step"), + p("brev_deploy.sh needs a logged-in brev CLI (brev login --api-key). It creates the instance if it does not exist (default an L4 on GCP; BREV_INSTANCE=progenome-a100 selects the A100), waits until the instance reports RUNNING and READY, refreshes the SSH alias, checks nvidia-smi and docker over ssh, uploads genomics/ plus the two small proteomics inputs as one tarball, builds the image natively on the GPU box (no emulation, about 4 minutes on the A100), runs the whole v1 pipeline inside the container with the GPU, runs the inference benchmark (eager, torch.compile, TensorRT attempt) and copies outputs/ back to outputs_brev//. Billing continues until brev stop; the A100 is left running for the presentation."), + h(3, "Inference and what torch.compile does"), + p("infer.py rebuilds the exact inputs of a trained run, wraps the model so torch.compile sees plain tensors instead of dictionaries, times ten full-graph passes eager, then compiles with the inductor backend (kernel fusion and graph capture) or with the torch_tensorrt backend and times again, and reports the maximum absolute logit difference so a speed-up cannot hide a numerical change. On the A100: 36.3 ms eager, 4.9 ms compiled (7.4x), maximum logit difference below 0.0001 across three runs, accuracy on all 2,503 labelled people 0.980 with SVD input. The TensorRT backend partitions the graph into dense parts it can compile (linear layers, norms) and scatter parts it cannot; on this hetero-GNN the compile step did not finish in three hours and is not part of any claim."), + + h(2, "13. How the goals were met, and what is not claimed"), + t(["Goal", "Evidence", "Status"], [ + ["RQ1: one graph joining haploblocks, genes, proteins and people", "schema v2 with 5 node types and 7 edge types, 2,548 people joined by id, browsable in Neo4j", "done"], + ["RQ2: GNN combines genome and proteome", "AUC 0.60 (genome) / 0.96 (proteome) / 0.99 (both); site control at chance; embedding silhouette 0.70; saliency finds causal clusters at 3x chance", "done on synthetic ground truth"], + ["RQ3: federated training without moving records", "NVFlare FedAvg over 3 sites, AUC 0.998 vs central 0.992 on the same people", "done in simulation"], + ["Encoder -> LLM decoder -> insights", "GraphRAG decoder with validated citations on Nemotron 3 Super", "done"], + ["Reproducible and deployable", "Makefile, Docker image, Brev deployment on L4 and A100, tests", "done"], + ], [0.32, 0.5, 0.18]), + b(["The case/control phenotype is synthetic: the integration numbers show the pipeline recovers a planted signal, not biology. Ancestry, population and sex results are on real labels.", + "Chromosome 22 only; the code is chromosome-agnostic (CHROM=chr21 make run).", + "On single-label accuracy the GNN ties, not beats, logistic regression; its value is the embedding space and the integration.", + "The GNN embedding does not recover single-cluster cis effects; a per-protein ridge does for the strong ones.", + "Federated evaluation reuses the central held-out people; a real deployment would report per site only.", + "TensorRT is not part of the inference claim; torch.compile is."]), + + h(3, "Verification and reproduction (18 September, before the commit)"), + p("Before committing, the whole chain was re-run twice from the files that would be committed: once as a simulated fresh clone on the laptop (only the tracked files copied to an empty folder, then make setup, make run, make run-v2, make federated, make test) and once on the A100 with the Docker image rebuilt from the current Dockerfile, which had not been built since the pyg_lib change. Both runs reproduce every count exactly and every model score within run-to-run noise (GPU kernels and FedAvg are not bit-reproducible). The rebuilt image contains pyg_lib, so Node2Vec ran natively for the first time; with its original five pretraining epochs it was clearly undertrained (ancestry 0.80), with fifty epochs the loss converges and it becomes a working but slightly weaker alternative to SVD, so SVD remains the default everywhere and the deploy script now defaults to it. The decoder was also re-called on the laptop run; the language model's wording and the number of ids it chooses to cite vary between calls, the citation check is what stays constant."), + t(["Quantity", "Documented (original runs)", "Fresh clone, laptop CPU", "Rebuilt image, A100"], [ + ["kept clusters / CARRIES / CO_OCCURS", "6,551 / 2,365,574 / 187,030", "identical", "identical"], + ["genes / proteins / MEASURED edges", "458 / 460 / 1,065,712", "identical", "identical"], + ["ancestry-associated clusters (FDR 5%) / sex", "6,470 / 0", "6,470 / 0", "6,470 / 0"], + ["logistic baseline ancestry / population / sex", "0.977 / 0.614 / 0.463", "0.977 / 0.614 / 0.463", "0.981 / 0.614 / 0.463"], + ["GNN v1 ancestry / population / sex (SVD input)", "0.974 / 0.437 / 0.503", "0.974 / 0.437 / 0.503", "run with Node2Vec instead (next row)"], + ["GNN v1 with Node2Vec input, 50 epochs (A100 only)", "not previously measured", "-", "0.950 / 0.292 / 0.489"], + ["embedding silhouette by ancestry: SVD / GNN", "0.06 / 0.70", "0.06 / 0.70", "0.06 / 0.70 (Node2Vec-initialised GNN)"], + ["phenotype AUC genome / proteome / both", "0.64 / 0.96 / 0.99", "0.65 / 0.96 / 0.99", "0.64 / 0.96 / 0.99"], + ["site control balanced accuracy (chance 0.33)", "0.30", "0.34", "0.30"], + ["saliency: causal clusters in top 20", "4", "4", "4"], + ["ridge: cis proteins with R2 > 0.1 / others", "4 / 0", "4 / 0", "4 / 0"], + ["federated 30 rounds AUC / central AUC, same 376 people", "0.998 / 0.992", "0.996 / 0.994", "0.987 / 0.995"], + ["site-alone AUC on own test people vs federated model on the same people", "not previously measured", "-", "0.998 / 0.986 / 0.988 vs 0.999 / 1.000 / 0.968"], + ["decoder: cited ids / invented ids", "25 / 0", "4 / 0 (4.2 s)", "dry run (no key on the box)"], + ["inference eager / torch.compile (A100)", "36 ms / 4.9 ms", "-", "36.3 ms / 4.9 ms (7.45x, max logit diff 1.9e-06)"], + ["unit tests", "11 pass", "11 pass (make setup)", "11 pass (docker build)"] + ], [0.34, 0.22, 0.22, 0.22]), + h(2, "14. Team context"), + p("Friederike Duendar (lead) wrote the README's research questions, the UniProt gene BED and an R exploration of gene-block overlaps; Nolan Bruyat built the synthetic proteomics generator and its plots; Zillur Rahman built the proteomics-side analysis now on main (see below); Yan Zhou coordinates the manuscript (two introduction paragraphs, two methods paragraphs, one results paragraph); Anita Egebor and Alvaro Martinez Barrio contributed to the README and data access. The genomics/ pipeline described here is the modelling backbone into which those pieces plug."), + h(3, "What is on main since this branch was created, and how it relates"), + p("The team's proteomics work on main (the synthetic proteomics generator with 4,000 SITE-id samples, protein filtering and classification baselines, a protein-centred knowledge graph built from the HaploGraph edge list annotated with UniProt proteins, a federated comparison of feature sets, and methods_and_results.md) covers the proteome side and the federated logistic setting. genomics/ adds the genome side: the person-level knowledge graph on the 1000 Genomes ids, the graph neural network, the LLM decoder and the NVFlare run. The two share the same HaploGraph edge file, the same UniProt gene BED and the same three-site design, and they are complementary: nothing in genomics/ touches a path on main, so the branch merges cleanly, and for the manuscript the proteomics analysis and the graph model plug into the same methods and results structure."), + + h(1, "Part IV: Reference"), + h(2, "Run it yourself"), + p("From a fresh clone of the repository on branch modelling, on a laptop or a GPU machine:"), + c(`git clone https://github.com/collaborativebioinformatics/ProGenome.git +cd ProGenome && git checkout modelling && cd genomics +make setup # .venv with torch (CPU, or CUDA if nvidia-smi works), pinned deps, unit tests +make run # v1: download -> graph -> statistics -> baseline -> plots -> GNN (ancestry, population, sex) -> embeddings +make run-v2 # v2: synthetic proteomics -> graph v2 -> EDA -> GNN ablations -> ridge -> decoder dry run +make federated ROUNDS=30 LOCAL_EPOCHS=5 # NVFlare FedAvg over 3 sites, then central scoring +make neo4j-load # browse the graph at http://localhost:7474 (neo4j / progenome) +make docker && make docker-run # the same v1 chain inside the CUDA image +BREV_INSTANCE=progenome-a100 make brev # build and run on the A100, copy outputs back +make decode WHO=HG00103 # LLM insight for one person (needs NVIDIA_API_KEY) +make test # 11 unit tests on a toy graph +make help # every target`), + p("Single stages take a --chrom flag and, where relevant, --target, --modality, --init; for example .venv/bin/python train_gnn_v2.py --target phenotype --modality both --init raw reproduces the combined model with saliency, and .venv/bin/python infer.py --run ancestry_svd --compile inductor reproduces the inference benchmark. Every output path is relative to genomics/, and data/, outputs/ and outputs_brev/ are git-ignored."), + h(2, "Secrets and configuration"), + p("Configuration enters the code in one place, config.py, which reads genomics/.env, then ~/.progenome.env, then defaults, with exported shell variables taking precedence; .env.example lists every variable with a comment (NVIDIA_API_KEY, NIM_MODEL, NIM_URL, NEO4J_URI/USER/PASSWORD, HAPLOBLOCKS_BASE, CHROM, data and output directories, BREV_INSTANCE/TYPE) and make config prints what is in effect with secrets masked. The shell scripts source the same files through load_env.sh, and the Docker targets pass .env into the container. The only secret is the NVIDIA API key; .env is git-ignored, nothing under the repository contains a key, and the decoder refuses to call the endpoint without one instead of falling back silently. Versions are pinned in requirements.txt (torch 2.14.0, torch_geometric 2.8.0.post1, pandas 3.0.5, scipy 1.18.1, scikit-learn 1.9.1, networkx 3.6.1, neo4j 6.3.1, nvflare 2.9.0, pytest 9.1.1) and setup.sh installs torch from the CPU or cu126 index depending on whether a GPU is present."), + h(2, "Repository map (genomics/)"), + c(`haplokg.py, build_kg.py knowledge graph v1 (people, clusters, blocks) +haplokg_proteins.py, build_kg_v2.py genes, proteins, harmonised measurements (v2) +proteomics_synth_1000g.py synthetic proteomics on 1000G ids with ground truth +cooccurrence_analysis.py, eda.py statistics and the EDA report +baseline.py logistic regression + the shared split +graph_explore.py, neo4j_load.py NetworkX statistics/plots, Neo4j loader +train_gnn.py, train_gnn_v2.py the GNN (genome; genome+proteome) +embeddings.py embedding quality and plots +proteome_linear_baseline.py per-protein ridge (genome -> proteome) +graphrag_decoder.py NIM LLM decoder +infer.py inference benchmark +federated/{model,client,job,evaluate_global}.py NVFlare FedAvg +Dockerfile, docker-compose.yml, Makefile, setup.sh, run_all.sh, run_v2.sh, brev_deploy.sh +docs/architecture.html (.mmd), METHODS.md, DEEP_DIVE.md, report/ +tests/ unit tests on a toy graph`), + h(2, "Glossary"), + t(["Term", "Meaning"], [ + ["Haploblock", "a stretch of chromosome between recombination hotspots, inherited as a unit"], + ["Haplotype", "one copy's sequence of a block; each person has two per block"], + ["Cluster", "a group of near-identical haplotypes of one block (MMseqs2); a person carries it if either haplotype is in it"], + ["Carrier matrix", "people x clusters 0/1 matrix; the CARRIES edges"], + ["Lift", "how much more often two clusters co-occur than chance: P(A and B) / (P(A) P(B))"], + ["Cramer's V", "effect size of a contingency test, 0 (independent) to 1 (perfectly associated)"], + ["FDR", "false discovery rate; Benjamini-Hochberg controls the expected fraction of false positives among findings"], + ["SVD", "singular value decomposition; here a 32-component factorisation of the carrier matrix giving embeddings"], + ["GNN / message passing", "neural network on a graph; each layer mixes a node's vector with its neighbours'"], + ["SAGEConv / GraphConv", "two PyG layer types: neighbour-mean aggregation; edge-weighted aggregation"], + ["Balanced accuracy / AUC", "mean per-class recall; probability a random case outscores a random control"], + ["Harmoniser", "per-site, per-protein robust z-score removing batch offsets"], + ["MNAR", "missing not at random; here low-abundance proteins go missing first"], + ["FedAvg", "federated averaging of model weights across sites"], + ["NIM", "NVIDIA Inference Microservice; an OpenAI-compatible LLM endpoint"], + ["Brev", "NVIDIA's GPU cloud; instances by the hour"], + ], [0.25, 0.75]), +]; + +// ----------------------------------------------------------------------------- LaTeX +function tex(s) { + return String(s).replace(/\\/g, "\\textbackslash{}").replace(/([&%$#_{}])/g, "\\$1").replace(/~/g, "\\textasciitilde{}").replace(/\^/g, "\\textasciicircum{}") + .replace(/->/g, "$\\rightarrow$").replace(/>=/g, "$\\geq$").replace(/<=/g, "$\\leq$"); +} +function buildTex() { + const out = []; + out.push(`\\documentclass[11pt,a4paper]{article} +\\usepackage[margin=2.2cm]{geometry} +\\usepackage[T1]{fontenc} +\\usepackage[utf8]{inputenc} +\\usepackage{lmodern} +\\usepackage{graphicx} +\\usepackage{booktabs} +\\usepackage{longtable} +\\usepackage{array} +\\usepackage{enumitem} +\\usepackage{hyperref} +\\usepackage{xcolor} +\\hypersetup{colorlinks=true, linkcolor=blue!50!black, urlcolor=blue!50!black} +\\graphicspath{{figures/}} +\\setlength{\\parskip}{4pt} +\\title{${tex(META.title)}\\\\[6pt]\\large ${tex(META.subtitle)}} +\\author{${tex(META.team)}} +\\date{${tex(META.date)}} +\\begin{document} +\\maketitle +\\tableofcontents +\\newpage +`); + let figN = 0; + for (const blk of CONTENT) { + if (blk.k === "h") { + const cmd = blk.level === 1 ? "\\section" : blk.level === 2 ? "\\subsection" : "\\subsubsection"; + out.push(`${cmd}{${tex(blk.text)}}\n`); + } else if (blk.k === "p") { + out.push(`${tex(blk.text)}\n`); + } else if (blk.k === "b") { + out.push("\\begin{itemize}[leftmargin=1.4em]\n" + blk.items.map((i) => ` \\item ${tex(i)}`).join("\n") + "\n\\end{itemize}\n"); + } else if (blk.k === "t") { + const n = blk.header.length; + const widths = blk.widths || Array(n).fill(1 / n); + const spec = widths.map((w) => `p{${(w * 0.94).toFixed(3)}\\linewidth}`).join(""); + out.push(`\\begin{longtable}{${spec}}\n\\toprule\n${blk.header.map((x) => `\\textbf{${tex(x)}}`).join(" & ")} \\\\\n\\midrule\n\\endhead\n` + + blk.rows.map((r) => r.map(tex).join(" & ") + " \\\\").join("\n") + "\n\\bottomrule\n\\end{longtable}\n"); + } else if (blk.k === "f") { + figN += 1; + out.push(`\\begin{figure}[htbp]\n\\centering\n\\includegraphics[width=${blk.width.toFixed(2)}\\linewidth]{${blk.file}}\n\\caption{${tex(blk.caption)}}\n\\end{figure}\n`); + } else if (blk.k === "c") { + out.push("\\begin{small}\\begin{verbatim}\n" + blk.code + "\n\\end{verbatim}\\end{small}\n"); + } + } + out.push("\\end{document}\n"); + return out.join("\n"); +} + +// ----------------------------------------------------------------------------- DOCX +function pngSize(file) { + const buf = fs.readFileSync(file); + return { w: buf.readUInt32BE(16), h: buf.readUInt32BE(20) }; +} +function buildDocx() { + const { Document, Packer, Paragraph, TextRun, HeadingLevel, Table, TableRow, TableCell, WidthType, ImageRun, AlignmentType, + TableOfContents, LevelFormat, BorderStyle, ShadingType, PageBreak } = docx; + const PAGE_W = 11906, MARGIN = 1134, CONTENT_W = PAGE_W - 2 * MARGIN; // A4 in DXA + const children = []; + children.push(new Paragraph({ text: META.title, heading: HeadingLevel.TITLE })); + children.push(new Paragraph({ children: [new TextRun({ text: META.subtitle, italics: true, size: 24 })], spacing: { after: 200 } })); + children.push(new Paragraph({ children: [new TextRun({ text: META.team, size: 20 })] })); + children.push(new Paragraph({ children: [new TextRun({ text: META.date, size: 20 })], spacing: { after: 300 } })); + children.push(new Paragraph({ text: "Contents", heading: HeadingLevel.HEADING_1 })); + children.push(new TableOfContents("Contents", { hyperlink: true, headingStyleRange: "1-3" })); + children.push(new Paragraph({ children: [new PageBreak()] })); + + for (const blk of CONTENT) { + if (blk.k === "h") { + const lvl = blk.level === 1 ? HeadingLevel.HEADING_1 : blk.level === 2 ? HeadingLevel.HEADING_2 : HeadingLevel.HEADING_3; + if (blk.level === 1 && children.length > 8) children.push(new Paragraph({ children: [new PageBreak()] })); + children.push(new Paragraph({ text: blk.text, heading: lvl })); + } else if (blk.k === "p") { + children.push(new Paragraph({ children: [new TextRun({ text: blk.text })], spacing: { after: 140 } })); + } else if (blk.k === "b") { + for (const it of blk.items) children.push(new Paragraph({ children: [new TextRun({ text: it })], numbering: { reference: "bullets", level: 0 }, spacing: { after: 60 } })); + } else if (blk.k === "t") { + const n = blk.header.length; + const widths = (blk.widths || Array(n).fill(1 / n)).map((w) => Math.round(w * CONTENT_W)); + const cell = (text, bold, shade) => new TableCell({ + width: { size: 0, type: WidthType.DXA }, // overwritten below + shading: shade ? { type: ShadingType.CLEAR, fill: "E6ECEB", color: "auto" } : undefined, + margins: { top: 60, bottom: 60, left: 80, right: 80 }, + children: [new Paragraph({ children: [new TextRun({ text: String(text), bold: !!bold, size: 18 })] })], + }); + const mk = (cells, bold, shade) => new TableRow({ tableHeader: !!bold, children: cells.map((x, i) => { const cl = cell(x, bold, shade); cl.options.width = { size: widths[i], type: WidthType.DXA }; return cl; }) }); + const rows = [mk(blk.header, true, true), ...blk.rows.map((r) => mk(r, false, false))]; + children.push(new Table({ rows, columnWidths: widths, width: { size: CONTENT_W, type: WidthType.DXA }, + borders: { top: { style: BorderStyle.SINGLE, size: 4, color: "999999" }, bottom: { style: BorderStyle.SINGLE, size: 4, color: "999999" }, + left: { style: BorderStyle.NONE, size: 0 }, right: { style: BorderStyle.NONE, size: 0 }, + insideHorizontal: { style: BorderStyle.SINGLE, size: 2, color: "CCCCCC" }, insideVertical: { style: BorderStyle.NONE, size: 0 } } })); + children.push(new Paragraph({ spacing: { after: 160 } })); + } else if (blk.k === "f") { + const file = path.join(FIG, blk.file); + const { w, h } = pngSize(file); + const widthPx = Math.round(620 * blk.width); // ~6.5in printable at 96 dpi + const heightPx = Math.round(widthPx * h / w); + children.push(new Paragraph({ alignment: AlignmentType.CENTER, children: [new ImageRun({ type: "png", data: fs.readFileSync(file), transformation: { width: widthPx, height: heightPx } })] })); + children.push(new Paragraph({ alignment: AlignmentType.CENTER, children: [new TextRun({ text: blk.caption, italics: true, size: 18 })], spacing: { after: 200 } })); + } else if (blk.k === "c") { + for (const line of blk.code.split("\n")) children.push(new Paragraph({ children: [new TextRun({ text: line, font: "Courier New", size: 17 })], spacing: { after: 0 } })); + children.push(new Paragraph({ spacing: { after: 160 } })); + } + } + + const doc = new Document({ + creator: "ProGenome team", title: META.title, + styles: { default: { document: { run: { font: "Calibri", size: 22 } } }, + paragraphStyles: [{ id: "Heading1", name: "Heading 1", basedOn: "Normal", next: "Normal", quickFormat: true, run: { size: 32, bold: true, color: "17232A" }, paragraph: { spacing: { before: 360, after: 160 }, outlineLevel: 0 } }, + { id: "Heading2", name: "Heading 2", basedOn: "Normal", next: "Normal", quickFormat: true, run: { size: 26, bold: true, color: "0E7C7B" }, paragraph: { spacing: { before: 280, after: 120 }, outlineLevel: 1 } }, + { id: "Heading3", name: "Heading 3", basedOn: "Normal", next: "Normal", quickFormat: true, run: { size: 23, bold: true }, paragraph: { spacing: { before: 200, after: 80 }, outlineLevel: 2 } }] }, + numbering: { config: [{ reference: "bullets", levels: [{ level: 0, format: LevelFormat.BULLET, text: "\u2022", alignment: AlignmentType.LEFT, style: { paragraph: { indent: { left: 540, hanging: 270 } } } }] }] }, + features: { updateFields: true }, + sections: [{ properties: { page: { size: { width: PAGE_W, height: 16838 }, margin: { top: MARGIN, bottom: MARGIN, left: MARGIN, right: MARGIN } } }, children }], + }); + return Packer.toBuffer(doc); +} + +(async () => { + fs.writeFileSync(path.join(HERE, "ProGenome_KT.tex"), buildTex()); + fs.writeFileSync(path.join(HERE, "ProGenome_KT.docx"), await buildDocx()); + const words = CONTENT.filter((b) => b.k === "p").map((b) => b.text.split(/\s+/).length).reduce((a, b) => a + b, 0); + console.log(`wrote ProGenome_KT.tex and ProGenome_KT.docx (${CONTENT.length} blocks, ~${words} words of prose, ${CONTENT.filter((b) => b.k === "f").length} figures, ${CONTENT.filter((b) => b.k === "t").length} tables)`); +})(); diff --git a/genomics/docs/report/figures/architecture.png b/genomics/docs/report/figures/architecture.png new file mode 100644 index 0000000..7f48179 Binary files /dev/null and b/genomics/docs/report/figures/architecture.png differ diff --git a/genomics/docs/report/figures/architecture_slide.png b/genomics/docs/report/figures/architecture_slide.png new file mode 100644 index 0000000..c233d0b Binary files /dev/null and b/genomics/docs/report/figures/architecture_slide.png differ diff --git a/genomics/docs/report/figures/baseline_vs_gnn.png b/genomics/docs/report/figures/baseline_vs_gnn.png new file mode 100644 index 0000000..09c4274 Binary files /dev/null and b/genomics/docs/report/figures/baseline_vs_gnn.png differ diff --git a/genomics/docs/report/figures/compute_benchmarks.png b/genomics/docs/report/figures/compute_benchmarks.png new file mode 100644 index 0000000..991a7d4 Binary files /dev/null and b/genomics/docs/report/figures/compute_benchmarks.png differ diff --git a/genomics/docs/report/figures/confusion_matrices.png b/genomics/docs/report/figures/confusion_matrices.png new file mode 100644 index 0000000..5acd25b Binary files /dev/null and b/genomics/docs/report/figures/confusion_matrices.png differ diff --git a/genomics/docs/report/figures/cramers_v.png b/genomics/docs/report/figures/cramers_v.png new file mode 100644 index 0000000..18b5313 Binary files /dev/null and b/genomics/docs/report/figures/cramers_v.png differ diff --git a/genomics/docs/report/figures/data_flow_map.png b/genomics/docs/report/figures/data_flow_map.png new file mode 100644 index 0000000..bb52c1e Binary files /dev/null and b/genomics/docs/report/figures/data_flow_map.png differ diff --git a/genomics/docs/report/figures/eda_blocks.png b/genomics/docs/report/figures/eda_blocks.png new file mode 100644 index 0000000..9a2887d Binary files /dev/null and b/genomics/docs/report/figures/eda_blocks.png differ diff --git a/genomics/docs/report/figures/eda_blocks_populations.png b/genomics/docs/report/figures/eda_blocks_populations.png new file mode 100644 index 0000000..9d33591 Binary files /dev/null and b/genomics/docs/report/figures/eda_blocks_populations.png differ diff --git a/genomics/docs/report/figures/eda_clusters.png b/genomics/docs/report/figures/eda_clusters.png new file mode 100644 index 0000000..7e1ed9f Binary files /dev/null and b/genomics/docs/report/figures/eda_clusters.png differ diff --git a/genomics/docs/report/figures/eda_cooccurrence.png b/genomics/docs/report/figures/eda_cooccurrence.png new file mode 100644 index 0000000..b7a3213 Binary files /dev/null and b/genomics/docs/report/figures/eda_cooccurrence.png differ diff --git a/genomics/docs/report/figures/eda_populations.png b/genomics/docs/report/figures/eda_populations.png new file mode 100644 index 0000000..ff20357 Binary files /dev/null and b/genomics/docs/report/figures/eda_populations.png differ diff --git a/genomics/docs/report/figures/eda_proteomics.png b/genomics/docs/report/figures/eda_proteomics.png new file mode 100644 index 0000000..27b8983 Binary files /dev/null and b/genomics/docs/report/figures/eda_proteomics.png differ diff --git a/genomics/docs/report/figures/edge_similarity.png b/genomics/docs/report/figures/edge_similarity.png new file mode 100644 index 0000000..b47aebb Binary files /dev/null and b/genomics/docs/report/figures/edge_similarity.png differ diff --git a/genomics/docs/report/figures/embedding_clusters.png b/genomics/docs/report/figures/embedding_clusters.png new file mode 100644 index 0000000..b631eea Binary files /dev/null and b/genomics/docs/report/figures/embedding_clusters.png differ diff --git a/genomics/docs/report/figures/embedding_individuals.png b/genomics/docs/report/figures/embedding_individuals.png new file mode 100644 index 0000000..39d6bd4 Binary files /dev/null and b/genomics/docs/report/figures/embedding_individuals.png differ diff --git a/genomics/docs/report/figures/embedding_quality.png b/genomics/docs/report/figures/embedding_quality.png new file mode 100644 index 0000000..03ec501 Binary files /dev/null and b/genomics/docs/report/figures/embedding_quality.png differ diff --git a/genomics/docs/report/figures/federated_rounds.png b/genomics/docs/report/figures/federated_rounds.png new file mode 100644 index 0000000..20b1b55 Binary files /dev/null and b/genomics/docs/report/figures/federated_rounds.png differ diff --git a/genomics/docs/report/figures/federated_site_alone.png b/genomics/docs/report/figures/federated_site_alone.png new file mode 100644 index 0000000..e4ce7bc Binary files /dev/null and b/genomics/docs/report/figures/federated_site_alone.png differ diff --git a/genomics/docs/report/figures/federated_topology.png b/genomics/docs/report/figures/federated_topology.png new file mode 100644 index 0000000..e932a29 Binary files /dev/null and b/genomics/docs/report/figures/federated_topology.png differ diff --git a/genomics/docs/report/figures/federated_vs_central.png b/genomics/docs/report/figures/federated_vs_central.png new file mode 100644 index 0000000..a6e19aa Binary files /dev/null and b/genomics/docs/report/figures/federated_vs_central.png differ diff --git a/genomics/docs/report/figures/genome_to_graph.png b/genomics/docs/report/figures/genome_to_graph.png new file mode 100644 index 0000000..7f146c0 Binary files /dev/null and b/genomics/docs/report/figures/genome_to_graph.png differ diff --git a/genomics/docs/report/figures/graph_edge_positions.png b/genomics/docs/report/figures/graph_edge_positions.png new file mode 100644 index 0000000..97f3a44 Binary files /dev/null and b/genomics/docs/report/figures/graph_edge_positions.png differ diff --git a/genomics/docs/report/figures/graph_region.png b/genomics/docs/report/figures/graph_region.png new file mode 100644 index 0000000..9a071a2 Binary files /dev/null and b/genomics/docs/report/figures/graph_region.png differ diff --git a/genomics/docs/report/figures/ground_truth_effects.png b/genomics/docs/report/figures/ground_truth_effects.png new file mode 100644 index 0000000..8ccead0 Binary files /dev/null and b/genomics/docs/report/figures/ground_truth_effects.png differ diff --git a/genomics/docs/report/figures/informative_clusters_heatmap.png b/genomics/docs/report/figures/informative_clusters_heatmap.png new file mode 100644 index 0000000..45d4745 Binary files /dev/null and b/genomics/docs/report/figures/informative_clusters_heatmap.png differ diff --git a/genomics/docs/report/figures/informativeness.png b/genomics/docs/report/figures/informativeness.png new file mode 100644 index 0000000..20c993c Binary files /dev/null and b/genomics/docs/report/figures/informativeness.png differ diff --git a/genomics/docs/report/figures/init_comparison.png b/genomics/docs/report/figures/init_comparison.png new file mode 100644 index 0000000..4978e58 Binary files /dev/null and b/genomics/docs/report/figures/init_comparison.png differ diff --git a/genomics/docs/report/figures/person_neighbourhood.png b/genomics/docs/report/figures/person_neighbourhood.png new file mode 100644 index 0000000..7f2d166 Binary files /dev/null and b/genomics/docs/report/figures/person_neighbourhood.png differ diff --git a/genomics/docs/report/figures/results_modalities.png b/genomics/docs/report/figures/results_modalities.png new file mode 100644 index 0000000..cad25d8 Binary files /dev/null and b/genomics/docs/report/figures/results_modalities.png differ diff --git a/genomics/docs/report/figures/ridge_r2.png b/genomics/docs/report/figures/ridge_r2.png new file mode 100644 index 0000000..b950d64 Binary files /dev/null and b/genomics/docs/report/figures/ridge_r2.png differ diff --git a/genomics/docs/report/figures/saliency_top20.png b/genomics/docs/report/figures/saliency_top20.png new file mode 100644 index 0000000..838c815 Binary files /dev/null and b/genomics/docs/report/figures/saliency_top20.png differ diff --git a/genomics/docs/report/figures/schema_diagram.png b/genomics/docs/report/figures/schema_diagram.png new file mode 100644 index 0000000..098f431 Binary files /dev/null and b/genomics/docs/report/figures/schema_diagram.png differ diff --git a/genomics/docs/report/figures/sites_batch_effect.png b/genomics/docs/report/figures/sites_batch_effect.png new file mode 100644 index 0000000..db707d5 Binary files /dev/null and b/genomics/docs/report/figures/sites_batch_effect.png differ diff --git a/genomics/docs/report/figures/sites_composition.png b/genomics/docs/report/figures/sites_composition.png new file mode 100644 index 0000000..ae96e4b Binary files /dev/null and b/genomics/docs/report/figures/sites_composition.png differ diff --git a/genomics/docs/report/figures/training_curves.png b/genomics/docs/report/figures/training_curves.png new file mode 100644 index 0000000..894d5ce Binary files /dev/null and b/genomics/docs/report/figures/training_curves.png differ diff --git a/genomics/docs/report/figures/workflow_pipeline.png b/genomics/docs/report/figures/workflow_pipeline.png new file mode 100644 index 0000000..ccc37dd Binary files /dev/null and b/genomics/docs/report/figures/workflow_pipeline.png differ diff --git a/genomics/docs/report/make_report_figures.py b/genomics/docs/report/make_report_figures.py new file mode 100644 index 0000000..4e5d49d --- /dev/null +++ b/genomics/docs/report/make_report_figures.py @@ -0,0 +1,563 @@ +#!/usr/bin/env python3 +"""Extra figures for the KT report, all from files under genomics/outputs* (nothing hand-typed). + + python docs/report/make_report_figures.py # writes docs/report/figures/*.png +""" +from __future__ import annotations + +import json +import sys +from pathlib import Path + +import matplotlib +matplotlib.use("Agg") +import matplotlib.pyplot as plt +from matplotlib.patches import FancyBboxPatch, FancyArrowPatch +import networkx as nx +import numpy as np +import pandas as pd + +HERE = Path(__file__).resolve().parent +G = HERE.parents[1] # genomics/ +FIG = HERE / "figures" +FIG.mkdir(exist_ok=True) +sys.path.insert(0, str(G)) +import haplokg # noqa: E402 +import haplokg_proteins as hp # noqa: E402 + +ANC = {"AFR": "#d55e00", "AMR": "#cc79a7", "EAS": "#009e73", "EUR": "#0072b2", "SAS": "#e69f00"} +TEAL, AMBER, INK, GREY = "#0e7c7b", "#b26a0c", "#17232a", "#8a9599" +OUT = G / "outputs" +BREV = G / "outputs_brev" / "progenome-a100" + + +def box(ax, x, y, w, h, text, fc="#ffffff", ec=INK, fs=8.5, lw=1.2, pad=0.02): + """Rounded box; note the pad extends the drawn box by `pad` on every side, so leave step >= w + 2*pad + gap.""" + ax.add_patch(FancyBboxPatch((x, y), w, h, boxstyle=f"round,pad={pad},rounding_size={min(pad, 0.02)}", fc=fc, ec=ec, lw=lw)) + ax.text(x + w / 2, y + h / 2, text, ha="center", va="center", fontsize=fs, color=INK, wrap=True) + + +def arrow(ax, x1, y1, x2, y2, text=None, color=INK, fs=7.5): + ax.add_patch(FancyArrowPatch((x1, y1), (x2, y2), arrowstyle="-|>", mutation_scale=12, lw=1.2, color=color)) + if text: + ax.text((x1 + x2) / 2, (y1 + y2) / 2 + 0.018, text, ha="center", va="bottom", fontsize=fs, color=GREY) + + + +# --------------------------------------------------------------------- 0. figures produced by the pipeline scripts themselves: copy the current versions in +PIPELINE_FIGURES = { + "cramers_v.png": OUT / "cooccurrence/chr22/cramers_v_by_phenotype.png", + "edge_similarity.png": OUT / "cooccurrence/chr22/edge_ancestry_similarity.png", + "informativeness.png": OUT / "cooccurrence/chr22/informativeness_along_chromosome.png", + "embedding_individuals.png": OUT / "embeddings/chr22/individuals_gnn_ancestry_svd.png", + "embedding_clusters.png": OUT / "embeddings/chr22/clusters_gnn_ancestry_svd.png", + "eda_populations.png": OUT / "eda/chr22/plots/02_populations.png", + "eda_blocks.png": OUT / "eda/chr22/plots/03_blocks.png", + "eda_clusters.png": OUT / "eda/chr22/plots/04_clusters.png", + "eda_cooccurrence.png": OUT / "eda/chr22/plots/05_cooccurrence.png", + "eda_proteomics.png": OUT / "eda/chr22/plots/07_proteomics.png", + "graph_region.png": OUT / "graph/chr22/region_45035149-45534032.png", + "graph_edge_positions.png": OUT / "graph/chr22/edge_positions.png", +} + + +def sync_pipeline_figures(): + """Copy the plots that cooccurrence_analysis.py, embeddings.py, eda.py and graph_explore.py wrote under outputs/.""" + import shutil + for name, src in PIPELINE_FIGURES.items(): + if src.exists(): + shutil.copyfile(src, FIG / name) + else: + print(" (missing, kept previous copy)", src) + + +# --------------------------------------------------------------------- 1. workflow +def workflow(): + fig, ax = plt.subplots(figsize=(16, 5.6)); ax.set_xlim(0, 1); ax.set_ylim(0, 1); ax.axis("off") + W, STEP, PAD = 0.112, 0.141, 0.006 # drawn width = W + 2*PAD = 0.124 < STEP: no overlap, room for the arrow + stages = [("data.haploblocks.org\nHaploGraph chr22\n(pre-built graph)", TEAL), ("fetch_data.sh\nmd5-verified\ndownload", None), + ("build_kg.py\nsparse carrier matrix\nPyG HeteroData", None), ("cooccurrence_analysis.py\ncluster / edge\nvs phenotype", None), + ("baseline.py\nlogistic regression\nshared split", None), ("train_gnn.py\nhetero-GNN\nencoder", None), ("embeddings.py\nsilhouette / kNN\nPCA plots", None)] + for i, (txt, col) in enumerate(stages): + x = 0.012 + i * STEP + box(ax, x, 0.64, W, 0.22, txt, fc="#d9efee" if col else "#ffffff", ec=TEAL if col else INK, fs=7.6, pad=PAD) + if i < len(stages) - 1: arrow(ax, x + W + PAD, 0.75, x + STEP - PAD, 0.75) + ax.text(0.012, 0.93, "v1 genome graph -> phenotypes", fontsize=11, weight="bold", color=TEAL) + stages2 = [("proteomics_synth_1000g.py\n3 sites, 1000G ids\nground truth", AMBER), ("build_kg_v2.py\ngenes, proteins,\nharmonised MEASURED", AMBER), + ("eda.py\nfull EDA\nreport", None), ("train_gnn_v2.py\ngenome / proteome / both\ncontrols, saliency", None), + ("proteome_linear_baseline.py\nridge cis test", None), ("graphrag_decoder.py\nNIM LLM\ncited insight", None), ("federated/job.py\nNVFlare FedAvg\n3 sites + site-alone", None)] + for i, (txt, col) in enumerate(stages2): + x = 0.012 + i * STEP + box(ax, x, 0.18, W, 0.24, txt, fc="#f6e7cf" if col else "#ffffff", ec=AMBER if col else INK, fs=7.4, pad=PAD) + if i < len(stages2) - 1: arrow(ax, x + W + PAD, 0.30, x + STEP - PAD, 0.30) + ax.text(0.012, 0.5, "v2 + proteomics -> integration -> decoder -> federated", fontsize=11, weight="bold", color=AMBER) + xk = 0.012 + 2 * STEP + W / 2 # build_kg.py -> build_kg_v2.py + arrow(ax, xk, 0.64 - PAD, xk, 0.42 + PAD, "hetero.pt", TEAL) + ax.text(0.012, 0.05, "Every stage is one script with a Makefile target; outputs land under genomics/outputs//chr22/. The Docker image and brev_deploy.sh run the same chain on a GPU.", fontsize=8.5, color=GREY) + fig.savefig(FIG / "workflow_pipeline.png", dpi=150, bbox_inches="tight"); plt.close(fig) + + +# --------------------------------------------------------------------- 2. schema +def schema(): + """Five node boxes with gaps wide enough for the edge labels; labels sit above the arrows, never on them.""" + fig, ax = plt.subplots(figsize=(15, 5.4)); ax.set_xlim(0, 1); ax.set_ylim(0, 1); ax.axis("off") + W, H, Y = 0.13, 0.2, 0.42 + xs = [0.03, 0.235, 0.44, 0.645, 0.85] + names = ["Individual", "Cluster", "Block", "Gene", "Protein"] + counts = ["2,548", "6,551", "669", "458", "460"] + cols = ["#e8e6f7", "#d9efee", "#d9efee", "#f6e7cf", "#f6e7cf"] + props = ["ancestry, population, sex,\nsite, phenotype (labels)", "support, block statistics,\nSVD-32 vector", "length, n_clusters,\nentropy, dominance", "coordinates", "coordinates, isoforms"] + for x, n, c, col, pr in zip(xs, names, counts, cols, props): + box(ax, x, Y, W, H, f"{n}\n({c})", fc=col, fs=11, pad=0.008) + ax.text(x + W / 2, Y - 0.05, pr, ha="center", va="top", fontsize=8, color=GREY) + edges = ["CARRIES\n2,365,574", "IN_BLOCK\n6,551", "OVERLAPS\n1,063", "ENCODES\n460"] + for i, lbl in enumerate(edges): + x1, x2 = xs[i] + W + 0.008, xs[i + 1] - 0.008 + ax.annotate("", xy=(x2, Y + H / 2), xytext=(x1, Y + H / 2), arrowprops=dict(arrowstyle="-|>", lw=1.3, color=INK)) + ax.text((x1 + x2) / 2, Y + H / 2 + 0.05, lbl, ha="center", va="bottom", fontsize=8.5, color=INK) + # self relations drawn as loops above the box + for x, lbl in ((xs[1], "CO_OCCURS {weight, lift}\n187,030"), (xs[2], "NEXT_BLOCK\n668")): + ax.annotate("", xy=(x + W * 0.72, Y + H + 0.01), xytext=(x + W * 0.28, Y + H + 0.01), arrowprops=dict(arrowstyle="-|>", connectionstyle="arc3,rad=-1.3", lw=1.2, color=INK)) + ax.text(x + W / 2, Y + H + 0.2, lbl, ha="center", va="bottom", fontsize=8.5, color=INK) + # MEASURED arc below, Individual -> Protein + ax.annotate("", xy=(xs[4] + W / 2, Y - 0.01), xytext=(xs[0] + W / 2, Y - 0.01), arrowprops=dict(arrowstyle="-|>", connectionstyle="arc3,rad=0.28", lw=1.4, color=AMBER)) + ax.text(0.5, 0.05, "MEASURED {log2, z} 1,065,712 (Individual -> Protein: one edge per observed protein level; a missing value creates no edge)", ha="center", fontsize=9, color=AMBER) + ax.text(0.5, 0.97, "Knowledge-graph schema (chr22 counts). Labels live on the Individual node, never as neighbours.", ha="center", va="top", fontsize=11, weight="bold", color=INK) + fig.savefig(FIG / "schema_diagram.png", dpi=150, bbox_inches="tight"); plt.close(fig) + + +# --------------------------------------------------------------------- 3. genome -> clusters schematic +def genome_schematic(): + fig, ax = plt.subplots(figsize=(14, 6.4)); ax.set_xlim(0, 1); ax.set_ylim(0, 1); ax.axis("off") + H = dict(fontsize=10, weight="bold", color=INK) + ax.text(0.02, 0.965, "1. chromosome 22 is cut into haploblocks at recombination hotspots (669 blocks tiling 17.1-50.2 Mb)", **H) + ax.add_patch(FancyBboxPatch((0.03, 0.845), 0.94, 0.055, boxstyle="round,pad=0.004", fc="#dfe6e8", ec=GREY)) + for xx in np.linspace(0.03, 0.97, 9)[1:-1]: + ax.plot([xx, xx], [0.845, 0.9], color=INK, lw=1) + ax.text(0.2, 0.815, "block 1", ha="center", fontsize=8.5, color=GREY); ax.text(0.5, 0.815, "block k", ha="center", fontsize=8.5, color=GREY); ax.text(0.85, 0.815, "block 669", ha="center", fontsize=8.5, color=GREY) + ax.text(0.02, 0.73, "2. two phased haplotypes per person and block", **H) + for j, (name, seqs) in enumerate([("HG00096", ["ACGTTGCA...", "ACGATGCA..."]), ("HG00097", ["ACGATGCA...", "TCGTTGCA..."]), ("NA21144", ["ACGATGCA...", "ACGATGCA..."])]): + ax.text(0.03, 0.66 - j * 0.065, name, fontsize=9, family="monospace") + for k, sq in enumerate(seqs): + ax.text(0.13 + k * 0.17, 0.66 - j * 0.065, f"hap{k}: {sq}", fontsize=8.5, family="monospace", color=INK) + ax.text(0.53, 0.73, "3. MMseqs2 groups near-identical haplotypes into clusters", **H) + for j, (cl, members, col) in enumerate([("cluster 1", "HG00096/hap0, HG00097/hap0, ... (1,983 carriers)", TEAL), ("cluster 2", "HG00096/hap1, HG00097/hap0, NA21144/hap0+1, ... (1,142)", AMBER), ("cluster 3", "HG00097/hap1, ... (559)", "#7a5af8")]): + ax.add_patch(FancyBboxPatch((0.54, 0.645 - j * 0.065), 0.018, 0.035, boxstyle="round,pad=0.002", fc=col, ec=col)) + ax.text(0.57, 0.662 - j * 0.065, f"{cl}: {members}", fontsize=8.5, va="center") + ax.text(0.02, 0.42, "4. carrier matrix = the graph's person-to-cluster edges (1 if either haplotype is in the cluster)", **H) + hdr = ["", "cluster 1", "cluster 2", "cluster 3", "..."] + rows = [["HG00096", 1, 1, 0, "..."], ["HG00097", 1, 1, 1, "..."], ["NA21144", 0, 1, 0, "..."]] + for c_, hname in enumerate(hdr): + ax.text(0.07 + c_ * 0.1, 0.35, hname, fontsize=9, weight="bold", ha="center") + for r_, row in enumerate(rows): + for c_, v in enumerate(row): + ax.text(0.07 + c_ * 0.1, 0.29 - r_ * 0.06, str(v), fontsize=9, ha="center", family="monospace" if c_ else None, + color=(TEAL if v == 1 else GREY) if isinstance(v, int) else INK) + ax.text(0.53, 0.35, "5. filter: keep clusters with 25 <= carriers <= N-25: 248,254 -> 6,551", fontsize=9) + ax.text(0.53, 0.29, "6. co-occurrence: clusters found together in people more than chance -> CO_OCCURS edges (lift)", fontsize=9) + ax.text(0.53, 0.23, "7. every person carries ~928 kept clusters (2 haplotypes x 669 blocks, minus filtered ones)", fontsize=9) + ax.text(0.53, 0.17, "8. the carrier matrix is the CARRIES edge list; the SVD of it gives the starting embeddings", fontsize=9) + fig.savefig(FIG / "genome_to_graph.png", dpi=150, bbox_inches="tight"); plt.close(fig) + + +# --------------------------------------------------------------------- 4. federated topology +def federated_topology(): + fig, ax = plt.subplots(figsize=(14, 6)); ax.set_xlim(0, 1); ax.set_ylim(0, 1); ax.axis("off") + W, H, Y = 0.2, 0.34, 0.5 + sites = [("Site 1 (hospital)", "835 people\n774,753 CARRIES\n355,294 MEASURED"), ("Site 2 (hospital)", "835 people"), ("Site 3 (hospital)", "833 people")] + for i, (name, detail) in enumerate(sites): + x = 0.03 + i * 0.235 + ax.add_patch(FancyBboxPatch((x, Y), W, H, boxstyle="round,pad=0.008", fc="#ffffff", ec="#7a5af8", ls="--", lw=1.5)) + ax.text(x + W / 2, Y + H - 0.05, name, ha="center", va="center", fontsize=10.5, weight="bold") + ax.text(x + W / 2, Y + H / 2 + 0.01, "Individual nodes + labels\nCARRIES and MEASURED edges\n(never leave the site)", ha="center", va="center", fontsize=8.8) + ax.text(x + W / 2, Y + 0.02, detail, ha="center", va="bottom", fontsize=7.5, color=GREY) + ax.annotate("", xy=(x + W / 2, 0.93), xytext=(x + W / 2, Y + H + 0.01), arrowprops=dict(arrowstyle="-", lw=1.2, color=INK)) + ax.annotate("", xy=(0.86, 0.93), xytext=(0.13, 0.93), arrowprops=dict(arrowstyle="-|>", lw=1.3, color=INK)) + ax.text(0.5, 0.955, "weights only, once per round (no rows, no protein values, no embeddings of people)", ha="center", fontsize=9, color=INK) + box(ax, 0.76, Y, 0.21, H, "NVFlare server\nFedAvg\n\naverages the three\nstate dicts, keeps the best\nby validation balanced\naccuracy, sends it back", fc="#ffffff", fs=8.8, pad=0.008) + ax.annotate("", xy=(0.03 + 2 * 0.235 + W + 0.008, Y + 0.08), xytext=(0.76 - 0.008, Y + 0.08), arrowprops=dict(arrowstyle="-|>", lw=1.2, color=INK)) + ax.text(0.735, Y + 0.035, "global model", ha="center", fontsize=8, color=GREY) + ax.add_patch(FancyBboxPatch((0.03, 0.1), 0.67, 0.3, boxstyle="round,pad=0.008", fc="#d9efee", ec=TEAL, lw=1.5)) + ax.text(0.365, 0.33, "Shared and public: the Cluster, Block, Gene, Protein graph", ha="center", va="center", fontsize=10, weight="bold", color=INK) + ax.text(0.365, 0.25, "CO_OCCURS, IN_BLOCK, NEXT_BLOCK, OVERLAPS, ENCODES edges and the encoder weights;\nidentical at every site (data.haploblocks.org + UniProt)", ha="center", va="center", fontsize=8.8, color=INK) + ax.text(0.365, 0.15, "each site trains 5 local epochs on its own people per round; 30 rounds", ha="center", va="center", fontsize=8.5, color=GREY) + box(ax, 0.76, 0.1, 0.21, 0.3, "Result\n\nglobal model AUC 0.987-0.998\nvs central 0.992-0.995\non the same 376 held-out\npeople (three runs)", fc="#ffffff", fs=8.8, pad=0.008) + fig.savefig(FIG / "federated_topology.png", dpi=150, bbox_inches="tight"); plt.close(fig) + + +# --------------------------------------------------------------------- 5. person neighbourhood (real) +def person_neighbourhood(): + ctx_path = OUT / "graphrag" / "chr22" / "HG00103_context.json" + if not ctx_path.exists(): + return + ctx = json.loads(ctx_path.read_text()) + Gp = nx.Graph(); who = ctx["individual"]["id"] + Gp.add_node(who, kind="person") + for cl in ctx["most_ancestry_informative_clusters_carried"][:6]: + Gp.add_node(cl["cluster_id"], kind="cluster", anc=cl["enriched_in"]); Gp.add_edge(who, cl["cluster_id"]) + Gp.add_node(cl["block_id"], kind="block"); Gp.add_edge(cl["cluster_id"], cl["block_id"]) + for g in cl["genes_in_block"][:3]: + Gp.add_node(g, kind="gene"); Gp.add_edge(cl["block_id"], g) + for pr in cl["proteins_in_block"][:3]: + Gp.add_node(pr, kind="protein"); Gp.add_edge(cl["genes_in_block"][0] if cl["genes_in_block"] else cl["block_id"], pr) + for pr in ctx["most_extreme_protein_levels"][:5]: + Gp.add_node(pr["protein_id"], kind="protein"); Gp.add_edge(who, pr["protein_id"], measured=True, z=pr["harmonised_z"]) + for nb in ctx["nearest_individuals_in_gnn_embedding"][:3]: + Gp.add_node(nb["individual_id"], kind="neighbour", anc=nb["ancestry"]); Gp.add_edge(who, nb["individual_id"], nn=True) + pos = nx.spring_layout(Gp, seed=3, k=0.9) + fig, ax = plt.subplots(figsize=(12, 9.5)) + kinds = {"person": ("#7a5af8", 900), "neighbour": ("#c7bfff", 500), "cluster": (TEAL, 420), "block": ("#9fd3d1", 420), "gene": (AMBER, 380), "protein": ("#f2c98a", 380)} + for kind, (col, size) in kinds.items(): + ns = [n for n, d in Gp.nodes(data=True) if d["kind"] == kind] + nx.draw_networkx_nodes(Gp, pos, nodelist=ns, node_color=col, node_size=size, ax=ax, edgecolors="white") + meas = [(u, v) for u, v, d in Gp.edges(data=True) if d.get("measured")] + nnb = [(u, v) for u, v, d in Gp.edges(data=True) if d.get("nn")] + other = [(u, v) for u, v, d in Gp.edges(data=True) if not d.get("measured") and not d.get("nn")] + nx.draw_networkx_edges(Gp, pos, edgelist=other, ax=ax, edge_color="#b8c2c5", width=1.2) + nx.draw_networkx_edges(Gp, pos, edgelist=meas, ax=ax, edge_color=AMBER, width=1.6, style="dashed") + nx.draw_networkx_edges(Gp, pos, edgelist=nnb, ax=ax, edge_color="#7a5af8", width=1.2, style="dotted") + labels = {n: (n.replace("chr22_", "").replace("_cluster", "\nc") if Gp.nodes[n]["kind"] in ("cluster", "block") else n) for n in Gp} + nx.draw_networkx_labels(Gp, pos, labels=labels, font_size=6.8, ax=ax, bbox=dict(boxstyle="round,pad=0.12", fc="white", ec="none", alpha=0.75)) + from matplotlib.lines import Line2D + handles = [Line2D([0], [0], marker="o", color="w", markerfacecolor=col, markersize=11, label=lbl) for lbl, col in + [(f"this person ({who})", "#7a5af8"), ("nearest people in the GNN embedding", "#c7bfff"), ("haploblock cluster this person carries", TEAL), + ("block containing that cluster", "#9fd3d1"), ("gene overlapping the block", AMBER), ("protein (encoded by the gene, or measured in this person)", "#f2c98a")]] + handles += [Line2D([0], [0], color="#b8c2c5", lw=1.6, label="graph edge: CARRIES, IN_BLOCK, OVERLAPS, ENCODES"), + Line2D([0], [0], color=AMBER, lw=1.6, ls="--", label="MEASURED: one of this person's most extreme protein levels"), + Line2D([0], [0], color="#7a5af8", lw=1.4, ls=":", label="nearest neighbour in the 64-d GNN embedding")] + ax.legend(handles=handles, loc="upper center", bbox_to_anchor=(0.5, -0.01), ncol=3, frameon=False, fontsize=8, handletextpad=0.6, columnspacing=1.4) + ax.set_title(f"What the decoder retrieves for {who} ({ctx['individual']['ancestry']}, {ctx['individual']['population']}); GNN prediction: {ctx['gnn_phenotype_prediction']['predicted']}", fontsize=10.5) + ax.axis("off"); fig.savefig(FIG / "person_neighbourhood.png", dpi=150, bbox_inches="tight"); plt.close(fig) + + +# --------------------------------------------------------------------- 6. sites composition / batch +def sites(): + synth = OUT / "proteomics_synth" / "chr22" + meta = pd.read_csv(synth / "sample_metadata.csv"); long = pd.read_csv(synth / "measured_long.csv") + fig, axes = plt.subplots(1, 3, figsize=(13, 3.8)) + comp = meta.groupby(["site", "ancestry"]).size().unstack(fill_value=0) + comp.plot(kind="bar", stacked=True, ax=axes[0], color=[ANC[a] for a in comp.columns], width=0.7) + axes[0].set_title("people per site by ancestry (mixed by design)"); axes[0].set_ylabel("people"); axes[0].legend(frameon=False, fontsize=7, loc="upper center", bbox_to_anchor=(0.5, -0.16), ncol=5); axes[0].tick_params(axis="x", rotation=0) + prev = meta.groupby("site")["phenotype"].mean() + axes[1].bar(prev.index, prev.values, color="#7a5af8", width=0.6); axes[1].set_ylim(0, 0.6); axes[1].set_title("case prevalence per site"); axes[1].set_ylabel("fraction cases") + for i, v in enumerate(prev.values): axes[1].text(i, v + 0.01, f"{v:.2f}", ha="center", fontsize=8) + n_per = meta.groupby("site").size(); miss = 1 - long.groupby("site").size() / (n_per * long["protein_id"].nunique()) + axes[2].bar(miss.index, miss.values, color=AMBER, width=0.6); axes[2].set_title("missing protein values per site"); axes[2].set_ylabel("fraction missing"); axes[2].set_ylim(0, 0.12) + for i, v in enumerate(miss.values): axes[2].text(i, v + 0.002, f"{v:.1%}", ha="center", fontsize=8) + fig.tight_layout(); fig.savefig(FIG / "sites_composition.png", dpi=150, bbox_inches="tight"); plt.close(fig) + + harm = hp.harmonise(long) + prots = long.groupby("protein_id")["log2_intensity"].median().sort_values().index[::23][:20] + raw = long[long.protein_id.isin(prots)].groupby(["protein_id", "site"])["log2_intensity"].median().unstack().loc[prots] + z = harm[harm.protein_id.isin(prots)].groupby(["protein_id", "site"])["z"].median().unstack().loc[prots] + fig, axes = plt.subplots(1, 2, figsize=(12, 3.8), sharex=True) + for s, col in zip(raw.columns, ["#0072b2", "#d55e00", "#009e73"]): + axes[0].plot(range(len(prots)), raw[s], marker="o", ms=3, color=col, label=s); axes[1].plot(range(len(prots)), z[s], marker="o", ms=3, color=col, label=s) + axes[0].set_title("median log2 intensity per site: raw (batch offsets visible)"); axes[1].set_title("after the harmoniser (robust z per site x protein)") + axes[0].set_xlabel("20 proteins, ordered by abundance"); axes[1].set_xlabel("20 proteins, ordered by abundance"); axes[0].legend(frameon=False, fontsize=8) + fig.tight_layout(); fig.savefig(FIG / "sites_batch_effect.png", dpi=150, bbox_inches="tight"); plt.close(fig) + + truth = json.loads((synth / "ground_truth.json").read_text()) + fig, axes = plt.subplots(1, 4, figsize=(14, 3.2)) + axes[0].bar(range(20), sorted([c["beta_phenotype"] for c in truth["causal_clusters"]]), color=TEAL); axes[0].set_title("causal clusters: effect on phenotype logit"); axes[0].set_xlabel("20 clusters (sorted)") + axes[1].bar(range(20), sorted([c["beta_cis"] for c in truth["causal_clusters"]]), color=AMBER); axes[1].set_title("cis effect on the block's protein (log2)") + bp = np.array(truth["protein_effects"]["beta_phenotype"]); axes[2].hist(bp[bp != 0], bins=20, color="#7a5af8"); axes[2].set_title(f"phenotype effect on {int((bp != 0).sum())} responsive proteins"); axes[2].set_xlabel("beta (log2)") + freq = [c for c in truth["causal_clusters"]] + kg = haplokg.load_kg(OUT / "kg" / "chr22"); sup = kg["clusters"].set_index("cluster_idx")["support_frac"] + axes[3].hist([sup.loc[c["cluster_idx"]] for c in freq], bins=10, color=GREY); axes[3].set_title("carrier frequency of causal clusters"); axes[3].set_xlabel("fraction of people") + fig.suptitle("Ground truth of the synthetic proteome (saved in ground_truth.json)", fontsize=10); fig.tight_layout(); fig.savefig(FIG / "ground_truth_effects.png", dpi=150, bbox_inches="tight"); plt.close(fig) + + +# --------------------------------------------------------------------- 7. training curves, saliency, ridge, results bars +def gnn_results(): + base = BREV / "gnn_v2" / "chr22" + runs = {"genome (raw)": "phenotype_genome_raw", "proteome (MLP)": "phenotype_proteome_svd", "genome + proteome (raw)": "phenotype_both_raw", "site control": "site_both_svd"} + fig, axes = plt.subplots(1, 2, figsize=(11, 3.6)) + for name, d in runs.items(): + h = base / d / "history.csv" + if h.exists(): + hist = pd.read_csv(h); axes[0].plot(hist["epoch"], hist["loss"], label=name); axes[1].plot(hist["epoch"], hist["val_score"], label=name) + axes[0].set_title("training loss"); axes[0].set_xlabel("epoch"); axes[1].set_title("validation balanced accuracy (early-stopping criterion)"); axes[1].set_xlabel("epoch"); axes[1].legend(frameon=False, fontsize=8) + fig.tight_layout(); fig.savefig(FIG / "training_curves.png", dpi=150, bbox_inches="tight"); plt.close(fig) + + metrics = {} + for name, d in {**runs, "genome (svd)": "phenotype_genome_svd", "genome + proteome (svd)": "phenotype_both_svd"}.items(): + m = base / d / "metrics.json" + if m.exists(): metrics[name] = json.loads(m.read_text())["test"] + order = [k for k in ["genome (svd)", "genome (raw)", "proteome (MLP)", "genome + proteome (svd)", "genome + proteome (raw)"] if k in metrics] + fig, ax = plt.subplots(figsize=(9, 3.6)); xs = np.arange(len(order)) + ax.bar(xs - 0.18, [metrics[k].get("roc_auc", np.nan) for k in order], 0.36, color=TEAL, label="AUC") + ax.bar(xs + 0.18, [metrics[k]["balanced_accuracy"] for k in order], 0.36, color=AMBER, label="balanced accuracy") + for i, k in enumerate(order): + ax.text(i - 0.18, metrics[k].get("roc_auc", 0) + 0.01, f"{metrics[k].get('roc_auc', float('nan')):.2f}", ha="center", fontsize=8) + ax.text(i + 0.18, metrics[k]["balanced_accuracy"] + 0.01, f"{metrics[k]['balanced_accuracy']:.2f}", ha="center", fontsize=8) + if "site control" in metrics: + ax.axhline(metrics["site control"]["balanced_accuracy"], ls="--", color="#7a5af8", lw=1.2, label=f"site (batch) control: balanced accuracy {metrics['site control']['balanced_accuracy']:.2f}, chance 0.33") + ax.set_xticks(xs); ax.set_xticklabels(order, fontsize=8); ax.set_ylim(0.2, 1.15); ax.legend(frameon=False, fontsize=8, loc="upper center", ncol=3); ax.set_title("Synthetic phenotype on the held-out people: genome vs proteome vs both") + fig.tight_layout(); fig.savefig(FIG / "results_modalities.png", dpi=150, bbox_inches="tight"); plt.close(fig) + + sal = base / "phenotype_both_raw" / "saliency_top100.csv" + if sal.exists(): + s = pd.read_csv(sal).head(20) + fig, ax = plt.subplots(figsize=(10, 3.4)) + ax.bar(range(20), s["saliency"], color=[TEAL if c else GREY for c in s["is_causal"]]) + ax.set_xticks(range(20)); ax.set_xticklabels([c.replace("chr22_", "").replace("_cluster", "\nc") for c in s["cluster_id"]], fontsize=6.5, rotation=90) + ax.set_ylabel("saliency (gradient of case logit)"); ax.set_title(f"Top-20 clusters by saliency (teal = ground-truth causal: {int(s['is_causal'].sum())}/20, chance 1.2)") + fig.tight_layout(); fig.savefig(FIG / "saliency_top20.png", dpi=150, bbox_inches="tight"); plt.close(fig) + + r2 = OUT / "gnn_v2" / "chr22" / "proteome_ridge_baseline" / "protein_r2.csv" + if r2.exists(): + t_ = pd.read_csv(r2) + fig, ax = plt.subplots(figsize=(8, 3.4)) + ax.hist(t_.loc[~t_.is_cis, "test_r2"].dropna(), bins=40, color=GREY, alpha=0.7, label="440 other proteins") + for v in t_.loc[t_.is_cis, "test_r2"].dropna(): ax.axvline(v, color=AMBER, lw=1.2) + ax.axvline(-9, color=AMBER, lw=1.2, label="20 cis proteins (vertical lines)"); ax.set_xlim(-0.3, 0.3) + ax.set_xlabel("test R-squared of per-protein ridge from the carrier row"); ax.set_ylabel("proteins"); ax.legend(frameon=False, fontsize=8) + ax.set_title("Genome -> proteome: only strong cis effects are recoverable\n(4 of 20 cis proteins above R2 0.1, 0 of 440 others)", fontsize=10) + fig.tight_layout(); fig.savefig(FIG / "ridge_r2.png", dpi=150, bbox_inches="tight"); plt.close(fig) + + # baseline vs gnn on real labels + bl = json.loads((OUT / "baseline" / "chr22" / "metrics.json").read_text())["targets"] + gnn = {} + for tgt, d in {"ancestry": "ancestry_svd", "population": "population_raw", "sex": "sex_svd"}.items(): + m = OUT / "gnn" / "chr22" / d / "metrics.json" + if m.exists(): gnn[tgt] = json.loads(m.read_text())["test"]["balanced_accuracy"] + fig, ax = plt.subplots(figsize=(7, 3.4)); xs = np.arange(3); tg = ["ancestry", "population", "sex"] + ax.bar(xs - 0.18, [bl[t]["test"]["balanced_accuracy"] for t in tg], 0.36, color=GREY, label="logistic regression") + ax.bar(xs + 0.18, [gnn.get(t, np.nan) for t in tg], 0.36, color=TEAL, label="GNN") + ax.axhline(0.5, ls=":", color="#999"); ax.set_xticks(xs); ax.set_xticklabels(["ancestry (5)", "population (26)", "sex (control)"]); ax.set_ylim(0, 1.05); ax.set_ylabel("test balanced accuracy"); ax.legend(frameon=False, fontsize=8) + for i, t in enumerate(tg): + ax.text(i - 0.18, bl[t]["test"]["balanced_accuracy"] + 0.01, f"{bl[t]['test']['balanced_accuracy']:.2f}", ha="center", fontsize=8); ax.text(i + 0.18, gnn.get(t, 0) + 0.01, f"{gnn.get(t, float('nan')):.2f}", ha="center", fontsize=8) + ax.set_title("Real labels: baseline vs GNN (sex must stay at chance)"); fig.tight_layout(); fig.savefig(FIG / "baseline_vs_gnn.png", dpi=150, bbox_inches="tight"); plt.close(fig) + + eq = OUT / "embeddings" / "chr22" / "embedding_quality.csv" + if eq.exists(): + e = pd.read_csv(eq); e = e[e["embedding"].isin(["svd32", "gnn_ancestry_svd", "gnn_population_svd", "gnn_sex_svd"])] + fig, axes = plt.subplots(1, 2, figsize=(10, 3.4)) + axes[0].bar(e["embedding"], e["silhouette_ancestry"], color=TEAL); axes[0].set_title("silhouette by ancestry (higher = tighter groups)"); axes[0].tick_params(axis="x", labelsize=7) + axes[1].bar(e["embedding"], e["knn5_test_balanced_accuracy_ancestry"], color=TEAL, label="ancestry"); axes[1].bar(e["embedding"], e["knn5_test_balanced_accuracy_sex"], color=GREY, alpha=0.6, label="sex (control)") + axes[1].set_title("5-NN test balanced accuracy in the embedding"); axes[1].legend(frameon=False, fontsize=8, loc="upper center", bbox_to_anchor=(0.5, -0.22), ncol=2); axes[1].tick_params(axis="x", labelsize=7) + fig.tight_layout(); fig.savefig(FIG / "embedding_quality.png", dpi=150, bbox_inches="tight"); plt.close(fig) + + +# --------------------------------------------------------------------- 8. federated vs central, inference +def federated_and_inference(): + fed10 = OUT / "federated" / "chr22" / "evaluation_10rounds.json"; fed30 = OUT / "federated" / "chr22" / "evaluation.json" + if fed30.exists(): + d30 = json.loads(fed30.read_text()); d10 = json.loads(fed10.read_text()) if fed10.exists() else None + fig, ax = plt.subplots(figsize=(8, 3.4)) + names, auc, bal = ["central"], [d30["central_model_same_test_people"]["roc_auc"]], [d30["central_model_same_test_people"]["balanced_accuracy"]] + if d10: names.append("federated 10 rounds"); auc.append(d10["federated_global_model"]["auc"]); bal.append(d10["federated_global_model"]["balanced_accuracy"]) + names.append("federated 30 rounds"); auc.append(d30["federated_global_model"]["auc"]); bal.append(d30["federated_global_model"]["balanced_accuracy"]) + for s in d30["federated_per_site"]: names.append(f"fed 30, {s['site']} (own test)"); auc.append(s["auc"]); bal.append(s["balanced_accuracy"]) + xs = np.arange(len(names)); ax.bar(xs - 0.18, auc, 0.36, color=TEAL, label="AUC"); ax.bar(xs + 0.18, bal, 0.36, color=AMBER, label="balanced accuracy") + for i in range(len(names)): ax.text(i - 0.18, auc[i] + 0.005, f"{auc[i]:.3f}", ha="center", fontsize=7); ax.text(i + 0.18, bal[i] + 0.005, f"{bal[i]:.2f}", ha="center", fontsize=7) + ax.set_xticks(xs); ax.set_xticklabels(names, fontsize=7.5, rotation=15); ax.set_ylim(0.8, 1.06); ax.legend(frameon=False, fontsize=8, loc="upper center", ncol=2, bbox_to_anchor=(0.5, 1.0)); ax.set_title("Federated global model vs central model (same held-out people)", pad=14) + fig.tight_layout(); fig.savefig(FIG / "federated_vs_central.png", dpi=150, bbox_inches="tight"); plt.close(fig) + + inf = BREV / "gnn" / "chr22" / "ancestry_node2vec" / "inference" + fig, axes = plt.subplots(1, 2, figsize=(9, 3.2)) + axes[0].bar(["M2 CPU", "A100"], [0.95, 0.10], color=[GREY, TEAL]); axes[0].set_ylabel("seconds per epoch"); axes[0].set_title("GNN training: one full-graph epoch") + for i, v in enumerate([0.95, 0.10]): axes[0].text(i, v + 0.02, f"{v:.2f} s", ha="center", fontsize=8) + ms = [] + for name, f_ in [("eager", "benchmark_none_fp32.json"), ("torch.compile", "benchmark_inductor_fp32.json")]: + pth = inf / f_ + if pth.exists(): d = json.loads(pth.read_text()); ms.append((name, d.get("compiled_ms_per_full_graph", d.get("eager_ms_per_full_graph")))) + if ms: + axes[1].bar([m[0] for m in ms], [m[1] for m in ms], color=[GREY, TEAL]); axes[1].set_ylabel("ms per full-graph inference (A100)"); axes[1].set_title("inference: 2,548 people in one pass") + for i, (n, v) in enumerate(ms): axes[1].text(i, v + 0.5, f"{v:.1f} ms", ha="center", fontsize=8) + fig.tight_layout(); fig.savefig(FIG / "compute_benchmarks.png", dpi=150, bbox_inches="tight"); plt.close(fig) + + +# --------------------------------------------------------------------- 9. more EDA: blocks entropy/dominance, informative clusters heatmap, sex by population +def eda_extra(): + kg = haplokg.load_kg(OUT / "kg" / "chr22"); bl = kg["blocks"]; ind = kg["individuals"] + fig, axes = plt.subplots(1, 2, figsize=(11, 3.8)) + sc = axes[0].scatter(bl["dominance"], bl["shannon_entropy"], c=np.log10(bl["n_clusters"]), s=10, cmap="viridis"); axes[0].set_xlabel("dominance (share of the largest cluster)"); axes[0].set_ylabel("Shannon entropy of clusters"); fig.colorbar(sc, ax=axes[0], label="log10 clusters in block") + axes[0].set_title("blocks: one common haplotype vs many rare ones") + pop = ind.dropna(subset=["ancestry"]).groupby(["population", "sex"]).size().unstack(fill_value=0) + pop.plot(kind="bar", stacked=True, ax=axes[1], color=["#7b3294", "#008837"], width=0.8); axes[1].set_title("sex within each population"); axes[1].tick_params(axis="x", labelsize=6.5); axes[1].legend(frameon=False, fontsize=8) + fig.tight_layout(); fig.savefig(FIG / "eda_blocks_populations.png", dpi=150, bbox_inches="tight"); plt.close(fig) + + assoc = pd.read_csv(OUT / "cooccurrence" / "chr22" / "cluster_phenotype_association.csv").sort_values("ancestry_cramers_v", ascending=False).head(30) + cols = [c for c in assoc.columns if c.startswith("carrier_frac_")] + fig, ax = plt.subplots(figsize=(7, 7)) + im = ax.imshow(assoc[cols].to_numpy(), cmap="viridis", aspect="auto", vmin=0, vmax=1) + ax.set_xticks(range(len(cols))); ax.set_xticklabels([c.replace("carrier_frac_", "") for c in cols]); ax.set_yticks(range(30)); ax.set_yticklabels([c.replace("chr22_", "") for c in assoc["cluster_id"]], fontsize=6.5) + fig.colorbar(im, ax=ax, label="fraction of carriers within ancestry"); ax.set_title("The 30 most ancestry-informative clusters: who carries them") + fig.tight_layout(); fig.savefig(FIG / "informative_clusters_heatmap.png", dpi=150, bbox_inches="tight"); plt.close(fig) + + +# --------------------------------------------------------------------- 10. ground truth vs prediction: confusion matrices +def confusion_matrices(): + runs = [("real ancestry: GNN (SVD input)", OUT / "gnn" / "chr22" / "ancestry_svd"), + ("synthetic phenotype: genome + proteome (graph)", BREV / "gnn_v2" / "chr22" / "phenotype_both_raw"), + ("synthetic phenotype: genome only (graph)", BREV / "gnn_v2" / "chr22" / "phenotype_genome_raw"), + ("sex, negative control: GNN (SVD input)", OUT / "gnn" / "chr22" / "sex_svd")] + fig, axes = plt.subplots(2, 2, figsize=(11, 8.6)) + for ax, (title, run) in zip(axes.ravel(), runs): + p = pd.read_csv(run / "test_predictions.csv") + ct = pd.crosstab(p["true"], p["pred"]).reindex(index=sorted(p["true"].unique()), columns=sorted(p["true"].unique()), fill_value=0) + m = ct.to_numpy(); ax.imshow(m, cmap="Blues", vmin=0, vmax=m.max()) + for i in range(m.shape[0]): + for j in range(m.shape[1]): + ax.text(j, i, str(m[i, j]), ha="center", va="center", fontsize=10, color="white" if m[i, j] > m.max() / 2 else INK) + ax.set_xticks(range(m.shape[1])); ax.set_xticklabels(ct.columns); ax.set_yticks(range(m.shape[0])); ax.set_yticklabels(ct.index) + ax.set_xlabel("predicted"); ax.set_ylabel("ground truth") + acc = np.trace(m) / m.sum() + ax.set_title(f"{title}\n{np.trace(m)} / {m.sum()} correct ({acc:.1%}) on the held-out people", fontsize=9.5) + fig.tight_layout(); fig.savefig(FIG / "confusion_matrices.png", dpi=150, bbox_inches="tight"); plt.close(fig) + + +# --------------------------------------------------------------------- 11. where each data source enters and where each model product goes +def data_flow_map(): + fig, ax = plt.subplots(figsize=(20, 8.5)); ax.set_xlim(0, 1.02); ax.set_ylim(0, 1); ax.axis("off") + LAB, LABE, PAD = "#efe3f2", "#7b3294", 0.006 + hdr = dict(ha="center", fontsize=11, weight="bold", color=INK) + # column 1: sources + SX, SW, SH = 0.01, 0.165, 0.1 + src = [("HaploGraph nodes.csv.gz\nwho carries which cluster", "#d9efee", TEAL, 0.85), ("HaploGraph edges + block_stats\nco-occurrence, block statistics", "#d9efee", TEAL, 0.70), + ("phenotypes_real.csv\nancestry, population, sex", LAB, LABE, 0.55), ("uniprot_chr22.bed\ngenes, proteins, coordinates", "#f6e7cf", AMBER, 0.40), + ("proteomics matrices (3 sites)\nlog2 intensity per protein", "#f6e7cf", AMBER, 0.25), ("proteomics metadata\nsite, age, sex, case/control", LAB, LABE, 0.10)] + for txt, fc, ec, y in src: + box(ax, SX, y, SW, SH, txt, fc=fc, ec=ec, fs=9, pad=PAD) + ax.text(SX + SW / 2, 0.985, "DATA SOURCES", **hdr) + # column 2: knowledge graph + KX, KW = 0.24, 0.21 + box(ax, KX, 0.62, KW, 0.3, "GRAPH STRUCTURE\nIndividual, Cluster, Block,\nGene, Protein\nCARRIES, CO_OCCURS, IN_BLOCK,\nNEXT_BLOCK, OVERLAPS, ENCODES,\nMEASURED (harmonised z)", fs=9, pad=PAD) + box(ax, KX, 0.36, KW, 0.2, "NODE FEATURES\ncluster and block statistics;\nperson = SVD-32 of the carrier matrix\n(label-free) or raw row,\nplus protein z and observed mask", fs=9, pad=PAD) + box(ax, KX, 0.10, KW, 0.2, "LABELS on the person node\nancestry, population, sex,\nsite, phenotype\ntargets and ground truth only:\nnever an edge, never a feature", fc=LAB, ec=LABE, fs=9, pad=PAD) + ax.text(KX + KW / 2, 0.985, "KNOWLEDGE GRAPH (hetero_v2.pt)", **hdr) + def a(x1, y1, x2, y2, color=INK, text=None): + ax.annotate("", xy=(x2, y2), xytext=(x1, y1), arrowprops=dict(arrowstyle="-|>", lw=1.2, color=color)) + if text: ax.text((x1 + x2) / 2, (y1 + y2) / 2 + 0.02, text, ha="center", va="bottom", fontsize=8.5, color=GREY) + R = SX + SW + PAD; L = KX - PAD + a(R, 0.90, L, 0.86); a(R, 0.75, L, 0.78); a(R, 0.45, L, 0.70) # structure <- nodes, edges, BED + a(R, 0.88, L, 0.50); a(R, 0.30, L, 0.44) # features <- nodes (SVD / raw), matrices + a(R, 0.60, L, 0.24, LABE); a(R, 0.15, L, 0.18, LABE) # labels <- phenotypes, metadata + # column 3: encoder + EX, EW = 0.50, 0.17 + box(ax, EX, 0.34, EW, 0.36, "GNN ENCODER\nHeteroConv x 2, hidden 64\nSAGEConv + edge-weighted\nGraphConv\n\ntrained on TRAIN people's labels\nselected on VAL\nreported on TEST (376 people)", fc="#d9efee", ec=TEAL, fs=9, pad=PAD) + ax.text(EX + EW / 2, 0.985, "ENCODER", **hdr) + a(KX + KW + PAD, 0.77, EX - PAD, 0.62, text="messages"); a(KX + KW + PAD, 0.46, EX - PAD, 0.52, text="inputs"); a(KX + KW + PAD, 0.20, EX - PAD, 0.42, LABE, "loss") + # column 4: products + OX, OW, OH = 0.72, 0.13, 0.09 + outs = [("class probabilities\nper person", 0.86), ("64-d embedding per\nperson and cluster", 0.71), ("saliency per\nhaploblock cluster", 0.56), ("model weights\n(state dict)", 0.41), ("predictions for all\n2,548 people (infer.py)", 0.26)] + for txt, y in outs: + box(ax, OX, y, OW, OH, txt, fs=9, pad=PAD); a(EX + EW + PAD, 0.52, OX - PAD, y + OH / 2) + ax.text(OX + OW / 2, 0.985, "TRAINED MODEL GIVES", **hdr) + # column 5: consumers + CX, CW, CH = 0.885, 0.125, 0.12 + cons = [("evaluation against\nground truth\n(test labels, ground_truth.json)", 0.79, "#ffffff", INK), ("LLM decoder\n(NIM GraphRAG)\ncited insight per person", 0.56, "#f6e7cf", AMBER), + ("NVFlare FedAvg\nonly weights leave a site", 0.34, "#d9efee", TEAL), ("plots, Neo4j, CSVs\nfor the team", 0.12, "#ffffff", INK)] + for txt, y, fc, ec in cons: + box(ax, CX, y, CW, CH, txt, fc=fc, ec=ec, fs=8.8, pad=PAD) + ax.text(CX + CW / 2, 0.985, "CONSUMERS", **hdr) + OR = OX + OW + PAD; CL = CX - PAD + a(OR, 0.905, CL, 0.85); a(OR, 0.605, CL, 0.82) # probabilities, saliency -> evaluation + a(OR, 0.755, CL, 0.64); a(OR, 0.60, CL, 0.60) # embedding, saliency -> decoder + a(OR, 0.455, CL, 0.40) # weights -> FedAvg + a(OR, 0.305, CL, 0.18) # predictions -> plots + ax.text(0.51, 0.02, "Phenotypes are never neighbours of the person: the encoder reaches a label only through the loss on training people. Ground truth for the synthetic phenotype is the generator's own label; saliency and ridge are scored against ground_truth.json.", ha="center", fontsize=9, color=GREY) + fig.savefig(FIG / "data_flow_map.png", dpi=150, bbox_inches="tight"); plt.close(fig) + + + +# --------------------------------------------------------------------- 12. site alone vs federated, and starting-embedding comparison +def federated_site_alone(): + lo = json.loads((G / "outputs_brev" / "progenome-a100-verify" / "federated" / "chr22" / "local_only_vs_federated.json").read_text()) + sites = list(lo["local_only"]) + own = [lo["local_only"][s]["own_test"]["auc"] for s in sites] + worst = [min(v["auc"] for k, v in lo["local_only"][s].items() if k != "own_test") for s in sites] + fed = {r["site"]: r["auc"] for r in lo["federated_per_site"]} + fig, ax = plt.subplots(figsize=(8, 3.6)); xs = np.arange(len(sites)); w = 0.26 + ax.bar(xs - w, own, w, color=GREY, label="site alone, own held-out people") + ax.bar(xs, worst, w, color="#c9d3d5", label="site alone, worst transfer to another site") + ax.bar(xs + w, [fed[s] for s in sites], w, color=TEAL, label="federated global model, same people") + for i in range(len(sites)): + for off, v in ((-w, own[i]), (0, worst[i]), (w, fed[sites[i]])): ax.text(xs[i] + off, v + 0.004, f"{v:.3f}", ha="center", fontsize=7) + ax.set_xticks(xs); ax.set_xticklabels(sites); ax.set_ylim(0.9, 1.02); ax.set_ylabel("test AUC") + ax.set_title(f"With and without federation ({lo['steps_per_site']} optimizer steps per site, A100)", fontsize=10); ax.legend(frameon=False, fontsize=8, loc="upper center", bbox_to_anchor=(0.5, -0.14), ncol=3) + fig.tight_layout(); fig.savefig(FIG / "federated_site_alone.png", dpi=150, bbox_inches="tight"); plt.close(fig) + + +def init_comparison(): + V = G / "outputs_brev" / "progenome-a100-verify" / "gnn" / "chr22" + runs = [("SVD-32", OUT / "gnn/chr22/ancestry_svd", OUT / "gnn/chr22/population_svd"), ("Node2Vec-32 (50 epochs)", V / "ancestry_node2vec", V / "population_node2vec"), + ("free learned", OUT / "gnn/chr22/ancestry_learned", None), ("raw carrier row", None, OUT / "gnn/chr22/population_raw")] + base = json.loads((OUT / "baseline/chr22/metrics.json").read_text())["targets"] + fig, axes = plt.subplots(1, 2, figsize=(10, 3.6)) + for ax, (target, col) in zip(axes, (("ancestry", 1), ("population", 2))): + names, vals = [], [] + for name, a, pth in runs: + path = a if target == "ancestry" else pth + if path is not None and (path / "metrics.json").exists(): + names.append(name); vals.append(json.loads((path / "metrics.json").read_text())["test"]["balanced_accuracy"]) + ax.bar(names, vals, color=[TEAL if n.startswith("SVD") else GREY for n in names]) + ax.axhline(base[target]["test"]["balanced_accuracy"], color=AMBER, ls="--", lw=1, label="logistic regression") + for i, v in enumerate(vals): ax.text(i, v + 0.01, f"{v:.3f}", ha="center", fontsize=8) + ax.set_ylim(0, 1.2); ax.set_ylabel("test balanced accuracy"); ax.set_title(f"{target}: GNN by starting embedding"); ax.tick_params(axis="x", labelsize=7.5); ax.legend(frameon=False, fontsize=8, loc="upper right") + fig.tight_layout(); fig.savefig(FIG / "init_comparison.png", dpi=150, bbox_inches="tight"); plt.close(fig) + + +# --------------------------------------------------------------------- 13. one-slide system architecture (16:9) +def architecture_slide(): + fig, ax = plt.subplots(figsize=(16, 9)); ax.set_xlim(0, 1); ax.set_ylim(0, 1); ax.axis("off") + PUR, PURF, PAD = "#7a5af8", "#f1eefe", 0.006 + def a(x1, y1, x2, y2, color=INK, text=None, lw=1.6, dy=0.018): + ax.annotate("", xy=(x2, y2), xytext=(x1, y1), arrowprops=dict(arrowstyle="-|>", lw=lw, color=color, mutation_scale=16)) + if text: ax.text((x1 + x2) / 2, (y1 + y2) / 2 + dy, text, ha="center", va="bottom", fontsize=10, color=color) + ax.text(0.5, 0.975, "ProGenome: person-level genome graph + local proteomics, one GNN trained federated, decoded by an LLM", ha="center", va="top", fontsize=15, weight="bold", color=INK) + # ---- left: hospital sites (private) + SX, SW, SH = 0.02, 0.21, 0.19 + ax.text(SX + SW / 2, 0.905, "HOSPITAL SITES (private, never leave)", ha="center", fontsize=11.5, weight="bold", color=PUR) + for i, (name, n) in enumerate([("Site 1", "835 people"), ("Site 2", "835 people"), ("Site 3", "833 people")]): + y = 0.66 - i * 0.24 + ax.add_patch(FancyBboxPatch((SX, y), SW, SH, boxstyle=f"round,pad={PAD}", fc=PURF, ec=PUR, ls="--", lw=1.6)) + ax.text(SX + SW / 2, y + SH - 0.035, f"{name} ({n})", ha="center", va="center", fontsize=11.5, weight="bold", color=INK) + ax.text(SX + SW / 2, y + SH / 2 - 0.02, "Individual nodes + labels\nCARRIES: haplotype clusters carried\nMEASURED: local protein levels (z)", ha="center", va="center", fontsize=9.5, color=INK) + # ---- middle top: shared reference graph + GX, GW = 0.29, 0.42 + ax.add_patch(FancyBboxPatch((GX, 0.70), GW, 0.17, boxstyle=f"round,pad={PAD}", fc="#d9efee", ec=TEAL, lw=1.6)) + ax.text(GX + GW / 2, 0.845, "SHARED REFERENCE GRAPH (public, identical at every site)", ha="center", va="center", fontsize=10.5, weight="bold", color=INK) + chain = [("Cluster\n6,551", 0.315), ("Block\n669", 0.415), ("Gene\n458", 0.515), ("Protein\n460", 0.615)] + for (txt, x) in chain: + box(ax, x, 0.725, 0.07, 0.07, txt, fc="#ffffff", ec=TEAL, fs=9.5, pad=0.004) + for (_, x1), (_, x2) in zip(chain, chain[1:]): + a(x1 + 0.074, 0.76, x2 - 0.004, 0.76, TEAL, lw=1.3) + ax.text(0.35, 0.805, "CO_OCCURS (lift)", ha="center", fontsize=8.5, color=GREY); ax.text(0.465, 0.71, "IN_BLOCK · OVERLAPS · ENCODES (haploblocks.org + UniProt)", ha="center", va="top", fontsize=8.5, color=GREY) + # ---- middle: encoder + EX, EW, EY, EH = 0.33, 0.34, 0.40, 0.22 + box(ax, EX, EY, EW, EH, "GNN ENCODER (PyTorch Geometric)\nHeteroConv x 2, hidden 64\nSAGEConv + edge-weighted GraphConv\n\nperson input: SVD-32 / carrier row + protein z\ntrained on labels of training people only", fc="#ffffff", ec=TEAL, fs=10, pad=PAD) + a(0.5, 0.70 - PAD, 0.5, EY + EH + PAD, TEAL, "message passing over the graph", dy=0.012) + # ---- middle bottom: NVFlare server + box(ax, 0.37, 0.10, 0.26, 0.15, "NVFlare SERVER (FedAvg)\naverages the site weights\nreturns the global model\n30 rounds x 5 local epochs", fc="#ffffff", ec=INK, fs=10, pad=PAD) + a(0.5, 0.25 + PAD, 0.5, EY - PAD, INK, "global model", dy=0.012) + # sites -> encoder / server + for i in range(3): + y = 0.66 - i * 0.24 + SH / 2 + a(SX + SW + PAD, y, EX - PAD, EY + EH * (0.8 - 0.3 * i), PUR, lw=1.4) # private edges into the model at its own site + ax.text(0.255, 0.885, "each site trains the same model on its own people", ha="left", fontsize=9, color=PUR, style="italic") + a(SX + SW + PAD, 0.20, 0.37 - PAD, 0.175, INK, "weights only", dy=0.012) + # ---- right: outputs and decoder + OX, OW = 0.745, 0.235 + ax.text(OX + OW / 2, 0.905, "WHAT COMES OUT", ha="center", fontsize=11.5, weight="bold", color=INK) + outs = [("prediction per person\n(ancestry, population, case/control)", 0.76), ("64-d embedding per person and cluster", 0.65), ("saliency: which clusters drive the score", 0.54)] + for txt, y in outs: + box(ax, OX, y, OW, 0.085, txt, fc="#ffffff", ec=INK, fs=9.5, pad=PAD); a(EX + EW + PAD, EY + EH / 2, OX - PAD, y + 0.0425, INK, lw=1.3) + box(ax, OX, 0.30, OW, 0.17, "LLM DECODER (NVIDIA NIM)\nNemotron 3 Super\nGraphRAG: graph facts + prediction\n-> cited report per person\nevery cited id is checked", fc="#f6e7cf", ec=AMBER, fs=10, pad=PAD) + for y in (0.76, 0.65, 0.54): + a(OX + OW / 2 + (y - 0.65) * 0.9, y - PAD, OX + OW / 2 + (y - 0.65) * 0.9, 0.47 + PAD, AMBER, lw=1.3) + box(ax, OX, 0.10, OW, 0.14, "held-out results (chr22)\nAUC genome 0.60 / proteome 0.96 / both 0.99\nfederated 0.987-0.998 vs central 0.992-0.995\ncontrols (sex, site) at chance", fc="#ffffff", ec=INK, fs=9, pad=PAD) + ax.text(0.5, 0.025, "Stack: PyTorch 2.14 · PyTorch Geometric 2.8 · CUDA 12.6 · NVIDIA FLARE 2.9 · NVIDIA NIM · Neo4j 5.26 · Docker · NVIDIA Brev A100 80 GB Data: 1000 Genomes HaploGraph chr22 (haploblocks.org), UniProt, synthetic proteomics on 1000G ids", + ha="center", fontsize=9.5, color=GREY) + fig.savefig(FIG / "architecture_slide.png", dpi=200, bbox_inches="tight", facecolor="white"); plt.close(fig) + + +if __name__ == "__main__": + for fn in (sync_pipeline_figures, workflow, schema, genome_schematic, federated_topology, person_neighbourhood, sites, gnn_results, federated_and_inference, eda_extra, confusion_matrices, data_flow_map, federated_site_alone, init_comparison, architecture_slide): + try: + fn(); print("ok ", fn.__name__) + except Exception as exc: # keep going; report which figure failed + print("FAIL", fn.__name__, type(exc).__name__, str(exc)[:160]) + print(len(list(FIG.glob("*.png"))), "figures in", FIG) diff --git a/genomics/eda.py b/genomics/eda.py new file mode 100644 index 0000000..e6aa6c7 --- /dev/null +++ b/genomics/eda.py @@ -0,0 +1,276 @@ +#!/usr/bin/env python3 +"""Exploratory data analysis of everything the pipeline consumes: what the data is, where it comes from, +how it is shaped, and the statistics that justify the modelling choices. One script, one report. + + python eda.py --chrom chr22 # -> outputs/eda//EDA.md + tables/*.csv + plots/*.png + +Sections of the report: + 1. Provenance what / where / why: every input file, its source, size, build, licence notes + 2. Individuals ancestry, population, sex; who lacks labels; the train/val/test split + 3. Haploblocks block length, clusters per block, entropy/dominance/singletons along the chromosome + 4. Clusters carrier support distribution, the support filter, clusters carried per person + 5. Co-occurrence lift/weight distributions, degree, components, hub artefact, same-block vs long-range + 6. Genes / proteins how many genes per block, blocks per gene, isoforms, genes outside any block + 7. Proteomics intensity range, missingness (LOD) per protein and per site, batch effect before + and after harmonisation, correlation structure, phenotype/age/sex signal + 8. Genome <-> proteome cis correlation between carrying a cluster and the protein encoded in its block +Every number in the report is computed here; nothing is typed in. +""" +from __future__ import annotations + +import argparse +import gzip +import json +from pathlib import Path + +import numpy as np +import pandas as pd +from scipy import sparse, stats + +import haplokg +import haplokg_proteins as hp + +ANCESTRY_COLOURS = {"AFR": "#d55e00", "AMR": "#cc79a7", "EAS": "#009e73", "EUR": "#0072b2", "SAS": "#e69f00"} + + +def fmt(x, nd=3): + return f"{x:,.{nd}f}" if isinstance(x, float) else f"{x:,}" + + +def main() -> int: + here = Path(__file__).resolve().parent + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--chrom", default="chr22") + parser.add_argument("--synth-dir", type=Path, default=None, help="proteomics folder (measured_long.csv, sample_metadata.csv, ground_truth.json)") + parser.add_argument("--no-plots", action="store_true") + args = parser.parse_args() + chrom = args.chrom + data = here / "data" + kg_dir = here / "outputs" / "kg" / chrom + synth = args.synth_dir or here / "outputs" / "proteomics_synth" / chrom + out = here / "outputs" / "eda" / chrom + (out / "tables").mkdir(parents=True, exist_ok=True); (out / "plots").mkdir(exist_ok=True) + md = [] + T = lambda name, df: df.to_csv(out / "tables" / f"{name}.csv", index=False) + + if not args.no_plots: + import matplotlib; matplotlib.use("Agg") + import matplotlib.pyplot as plt + + # ------------------------------------------------------------------ 1. provenance + files = [ + ("haplograph//nodes.csv.gz", "HaploGraph cluster x individual 0/1 matrix", "data.haploblocks.org/haplograph/1000G", "built at MDxCORE Rigshospitalet Sept 2026 for this hackathon; 1000G phased haplotypes, GRCh38, MMseqs2 clusters per recombination-defined haploblock"), + ("haplograph//edges_lift_above_threshold.csv.gz", "cluster co-occurrence edges, lift >= 5", "same", "raw edges.csv is dominated by a population-frequency mega-hub; lift-filtered file is the signal"), + ("haplograph//islands.csv.gz", "dense multi-block extended haplotypes", "same", "candidate extended haplotypes"), + ("haplograph/phenotypes_real.csv", "ancestry / population / sex per 1000G individual", "1000G panel via IGSR", "the only real phenotypes 1000G has; anything else must be simulated"), + ("haploblocks/block_stats.tsv", "per-block length, n_clusters, entropy, dominance, singletons", "data.haploblocks.org/bidirectional_blast_samples", "QC statistics from the haploblock pipeline"), + ("haploblocks/_haploblock_boundaries_.tsv", "block START/END", "data.haploblocks.org/haploblock_hashes", "step 1 of the haploblocks pipeline (recombination-rate peaks)"), + ("../proteomics/uniprot_chr22.bed", "UniProt proteins (isoforms) with genomic spans", "UCSC UniProt track, Friederike", "gene BED + gene->protein map in one file"), + ("outputs/proteomics_synth//", "synthetic per-site proteomics on 1000G IDs with saved ground truth", "proteomics_synth_1000g.py", "stand-in until Wu 2013 LCL proteomics / UKB-PPP pQTL are wired in"), + ] + sizes = {} + for rel, *_ in files: + p = data / rel.replace("", chrom) + p = p if p.exists() else here / rel.replace("", chrom) + sizes[rel] = (p.stat().st_size / 1e6 if p.is_file() else sum(f.stat().st_size for f in p.glob("*") if f.is_file()) / 1e6) if p.exists() else float("nan") + prov = pd.DataFrame([{"file": r, "what": w, "where": s, "why / note": n, "MB": round(sizes[r], 1)} for r, w, s, n in files]) + T("provenance", prov) + md += [f"# EDA — {chrom}", "", "## 1. Provenance: what, where, why", "", prov.to_markdown(index=False), ""] + + # ------------------------------------------------------------------ 2. individuals + kg = haplokg.load_kg(kg_dir) + ind = kg["individuals"] + split = haplokg.load_or_make_split(kg, here / "outputs" / "splits" / chrom / "split_seed42.csv") + anc = ind["ancestry"].fillna("unlabelled").value_counts().rename_axis("ancestry").reset_index(name="n") + pop = ind.groupby(["ancestry", "population"]).size().reset_index(name="n").sort_values(["ancestry", "n"], ascending=[True, False]) + sex = ind.groupby(["ancestry", "sex"]).size().unstack(fill_value=0).reset_index() + T("individuals_ancestry", anc); T("individuals_population", pop); T("individuals_sex_by_ancestry", sex) + sp = pd.Series(split).value_counts().rename_axis("split").reset_index(name="n") + md += ["## 2. Individuals", "", + f"{len(ind):,} individuals are columns of `nodes.csv.gz`; {int((ind['ancestry_code'] >= 0).sum()):,} have labels in `phenotypes_real.csv`, " + f"{int((ind['ancestry_code'] < 0).sum())} do not (kept as unlabelled nodes). 1000G has **no** other phenotypes (no height, no disease), which is why height on the whiteboard had to become a simulated label.", "", + "Ancestry (super-population):", "", anc.to_markdown(index=False), "", + f"Populations: {pop['population'].nunique()} (smallest {pop['n'].min()}, largest {pop['n'].max()}). Sex: " + + ", ".join(f"{r.ancestry} {int(r.get('female', 0))}F/{int(r.get('male', 0))}M" for _, r in sex.iterrows()) + ".", "", + "Split (seed 42, stratified on ancestry; shared by every model): " + ", ".join(f"{r.split} {r.n}" for _, r in sp.iterrows()), ""] + if not args.no_plots: + fig, ax = plt.subplots(figsize=(8, 3.6)) + pop_sorted = pop.sort_values(["ancestry", "population"]) + ax.bar(pop_sorted["population"], pop_sorted["n"], color=[ANCESTRY_COLOURS[a] for a in pop_sorted["ancestry"]]) + ax.set_ylabel("individuals"); ax.set_title("1000G individuals per population, coloured by super-population"); ax.tick_params(axis="x", rotation=90, labelsize=8) + fig.tight_layout(); fig.savefig(out / "plots" / "02_populations.png", dpi=150); plt.close(fig) + + # ------------------------------------------------------------------ 3. haploblocks + blocks = kg["blocks"].copy() + blocks["length_kb"] = blocks["block_length"] / 1e3 + q = blocks["length_kb"].quantile([0.05, 0.25, 0.5, 0.75, 0.95]) + gaps = blocks.sort_values("start")["start"].to_numpy()[1:] - blocks.sort_values("start")["end"].to_numpy()[:-1] + bstats = pd.DataFrame({ + "metric": ["blocks", "span covered (Mb)", "first block start (Mb)", "last block end (Mb)", "length median (kb)", "length 5%/95% (kb)", + "longest block (kb)", "gaps between consecutive blocks > 1 kb", "clusters per block median", "clusters per block max", + "singleton rate median", "Shannon entropy median", "dominance (largest cluster share) median"], + "value": [len(blocks), round((blocks["end"].max() - blocks["start"].min()) / 1e6, 2), round(blocks["start"].min() / 1e6, 2), round(blocks["end"].max() / 1e6, 2), + round(q[0.5], 1), f"{q[0.05]:.1f} / {q[0.95]:.1f}", round(blocks["length_kb"].max(), 1), int((gaps > 1000).sum()), + int(blocks["n_clusters"].median()), int(blocks["n_clusters"].max()), round(blocks["singleton_rate"].median(), 3), + round(blocks["shannon_entropy"].median(), 2), round(blocks["dominance"].median(), 3)]}) + T("blocks_summary", bstats) + rho = stats.spearmanr(blocks["block_length"], blocks["n_clusters"]).correlation + md += ["## 3. Haploblocks (recombination-defined regions)", "", bstats.to_markdown(index=False), "", + f"Blocks tile the chromosome without overlap (gaps > 1 kb between consecutive blocks: {int((gaps > 1000).sum())}). " + f"Longer blocks carry more distinct haplotype clusters (Spearman ρ = {rho:.2f}); high-entropy blocks are the ones where the population is most diverse.", ""] + if not args.no_plots: + fig, axes = plt.subplots(1, 3, figsize=(12, 3.4)) + axes[0].hist(np.log10(blocks["block_length"]), bins=40, color="#1f6f8b"); axes[0].set_xlabel("log10 block length (bp)"); axes[0].set_ylabel("blocks") + axes[1].scatter(blocks["block_length"] / 1e3, blocks["n_clusters"], s=6, color="#1f6f8b"); axes[1].set_xscale("log"); axes[1].set_yscale("log"); axes[1].set_xlabel("block length (kb)"); axes[1].set_ylabel("clusters in block") + mid = (blocks["start"] + blocks["end"]) / 2e6 + axes[2].plot(mid, blocks["shannon_entropy"], lw=0.8, color="#1f6f8b"); axes[2].set_xlabel(f"{chrom} position (Mb)"); axes[2].set_ylabel("Shannon entropy of clusters") + fig.suptitle("Haploblocks: size, diversity and where diversity sits"); fig.tight_layout(); fig.savefig(out / "plots" / "03_blocks.png", dpi=150); plt.close(fig) + + # ------------------------------------------------------------------ 4. clusters + with gzip.open(data / "haplograph" / chrom / "nodes.csv.gz", "rt") as fh: + n_cols = len(fh.readline().split(",")) + all_support = None + stats_path = data / "haploblocks" / "block_stats.tsv" + bs = pd.read_csv(stats_path, sep="\t"); bs = bs[bs["chr"] == chrom] + total_clusters = int(bs["n_clusters"].sum()); singletons = int(bs["singleton_count"].sum()) + clusters = kg["clusters"] + carries = kg["carries"] + per_person = np.asarray(carries.sum(axis=1)).ravel() + csum = pd.DataFrame({ + "metric": ["clusters in nodes.csv.gz (all)", "of which singletons (1 carrier)", "kept after symmetric support >= 25", "kept fraction", + "support median (kept)", "support max (kept)", "clusters carried per person: mean", "min", "max", "blocks with >= 1 kept cluster"], + "value": [total_clusters, singletons, len(clusters), round(len(clusters) / total_clusters, 4), int(clusters["support"].median()), int(clusters["support"].max()), + round(per_person.mean(), 1), int(per_person.min()), int(per_person.max()), int(clusters["block_idx"].nunique())]}) + T("clusters_summary", csum) + md += ["## 4. Clusters (the nodes people connect to)", "", csum.to_markdown(index=False), "", + f"{singletons / total_clusters:.0%} of all clusters are singletons (one haplotype); they carry no population signal and are dropped. " + f"Each person carries ~{per_person.mean():.0f} of the kept clusters (≤ 2 per block: two haplotypes), so the individual × cluster matrix is " + f"{carries.nnz / (carries.shape[0] * carries.shape[1]):.1%} dense — sparse int8 storage is what makes chr22 fit in memory.", ""] + if not args.no_plots: + fig, axes = plt.subplots(1, 2, figsize=(9, 3.4)) + axes[0].hist(np.log10(clusters["support"]), bins=40, color="#1f6f8b"); axes[0].set_xlabel("log10 carriers per kept cluster"); axes[0].set_ylabel("clusters") + for a, col in ANCESTRY_COLOURS.items(): + m = (ind["ancestry"] == a).to_numpy() + axes[1].hist(per_person[m], bins=30, histtype="step", color=col, label=a, linewidth=1.5) + axes[1].set_xlabel("kept clusters carried per person"); axes[1].legend(frameon=False, fontsize=8) + fig.tight_layout(); fig.savefig(out / "plots" / "04_clusters.png", dpi=150); plt.close(fig) + + # ------------------------------------------------------------------ 5. co-occurrence + co = kg["co_occurs"].copy() + cl_block = clusters.set_index("cluster_idx")["block_idx"] + co["same_block"] = cl_block.loc[co["src"]].to_numpy() == cl_block.loc[co["dst"]].to_numpy() + bstart = blocks.set_index("block_idx")["start"] + co["distance_kb"] = np.abs(bstart.loc[cl_block.loc[co["src"]].to_numpy()].to_numpy() - bstart.loc[cl_block.loc[co["dst"]].to_numpy()].to_numpy()) / 1e3 + deg = np.bincount(np.concatenate([co["src"], co["dst"]]), minlength=len(clusters)) + esum = pd.DataFrame({ + "metric": ["edges (lift >= 5)", "lift median", "lift 95%", "lift max", "weight (shared carriers) median", "same-block edges", "median distance between endpoints (kb)", + "edges spanning > 10 Mb", "clusters with >= 1 edge", "degree median (connected)", "degree max", "top-1% of clusters hold this share of edges"], + "value": [len(co), round(co["lift"].median(), 2), round(co["lift"].quantile(0.95), 2), round(co["lift"].max(), 1), int(co["weight"].median()), + int(co["same_block"].sum()), round(co["distance_kb"].median(), 0), int((co["distance_kb"] > 10_000).sum()), int((deg > 0).sum()), + int(np.median(deg[deg > 0])), int(deg.max()), f"{np.sort(deg)[::-1][:max(1, len(deg)//100)].sum() / deg.sum():.0%}"]}) + T("cooccurrence_summary", esum) + md += ["## 5. Co-occurrence graph", "", esum.to_markdown(index=False), "", + "Lift = P(A∩B)/(P(A)P(B)): how much more often two clusters travel together than chance. Most edges are long-range (tens of Mb), i.e. " + "population structure rather than physical linkage — the HaploGraph README's mega-hub warning made visible. The co-occurrence analysis " + "(`cooccurrence_analysis.py`) shows edges overwhelmingly join clusters enriched in the same ancestry.", ""] + if not args.no_plots: + fig, axes = plt.subplots(1, 3, figsize=(12, 3.4)) + axes[0].hist(np.log10(co["lift"]), bins=40, color="#1f6f8b"); axes[0].set_xlabel("log10 lift"); axes[0].set_ylabel("edges") + axes[1].hist(np.log10(co["distance_kb"].clip(lower=1)), bins=40, color="#1f6f8b"); axes[1].set_xlabel("log10 distance between endpoints (kb)") + axes[2].hist(np.log10(deg[deg > 0]), bins=40, color="#1f6f8b"); axes[2].set_xlabel("log10 degree (connected clusters)") + fig.suptitle("Co-occurrence edges: strength, reach, concentration"); fig.tight_layout(); fig.savefig(out / "plots" / "05_cooccurrence.png", dpi=150); plt.close(fig) + + # ------------------------------------------------------------------ 6. genes / proteins + have_v2 = (kg_dir / "proteins.csv").exists() + if have_v2: + pt = hp.load_protein_tables(kg_dir) + genes, prot, bg = pt["genes"], pt["proteins"], pt["block_gene"] + gpb = bg.groupby("block_idx").size() + gsum = pd.DataFrame({ + "metric": ["proteins (UniProt accessions, isoforms collapsed)", "isoform rows in BED", "genes (symbols)", "proteins per gene max", "block–gene overlaps", + "genes overlapping no block", "genes spanning >= 2 blocks", "blocks with >= 1 gene", "genes per gene-bearing block median / max"], + "value": [len(prot), int(prot["n_isoforms"].sum()), len(genes), int(genes["n_proteins"].max()), len(bg), + int((~genes["gene_idx"].isin(bg["gene_idx"])).sum()), int((bg.groupby("gene_idx").size() >= 2).sum()), int(gpb.size), + f"{int(gpb.median())} / {int(gpb.max())}"]}) + T("genes_proteins_summary", gsum) + md += ["## 6. Genes and proteins", "", gsum.to_markdown(index=False), "", + "Genes outside every block sit in the chromosome ends / assembly gaps the haploblock boundaries do not cover; a gene spanning two blocks " + "links both blocks to its protein, which is what the OVERLAPS edge encodes.", ""] + + # ------------------------------------------------------------------ 7. proteomics + if have_v2 and (synth / "measured_long.csv").exists(): + long = pd.read_csv(synth / "measured_long.csv"); meta = pd.read_csv(synth / "sample_metadata.csv") + truth = json.loads((synth / "ground_truth.json").read_text()) if (synth / "ground_truth.json").exists() else None + n_people, n_prot = meta["sample_id"].nunique(), long["protein_id"].nunique() + missing = 1 - len(long) / (n_people * n_prot) + per_prot_missing = 1 - long.groupby("protein_id").size() / n_people + raw_site = long.groupby(["protein_id", "site"])["log2_intensity"].median().unstack() + site_spread_raw = raw_site.max(axis=1) - raw_site.min(axis=1) + harm = hp.harmonise(long) + harm_site = harm.groupby(["protein_id", "site"])["z"].median().unstack() + site_spread_harm = harm_site.max(axis=1) - harm_site.min(axis=1) + # phenotype / age / sex signal per protein on harmonised values + h = harm.merge(meta.rename(columns={"sample_id": "individual_id"})[["individual_id", "age", "sex", "phenotype"]], on="individual_id") + rows = [] + for pid, g in h.groupby("protein_id"): + t_ph = stats.ttest_ind(g.loc[g["phenotype"] == 1, "z"], g.loc[g["phenotype"] == 0, "z"], equal_var=False) + rows.append({"protein_id": pid, "n": len(g), "phenotype_t": t_ph.statistic, "phenotype_p": t_ph.pvalue, + "age_r": stats.pearsonr(g["age"], g["z"])[0], "sex_t": stats.ttest_ind(g.loc[g["sex"] == 1, "z"], g.loc[g["sex"] == 0, "z"], equal_var=False).statistic}) + assoc_p = pd.DataFrame(rows) + from cooccurrence_analysis import bh_fdr + assoc_p["phenotype_q"] = bh_fdr(assoc_p["phenotype_p"].to_numpy()) + T("proteomics_protein_associations", assoc_p.sort_values("phenotype_p")) + psum = pd.DataFrame({ + "metric": ["people with proteomics", "proteins", "sites", "people per site", "log2 intensity range (1%–99%)", "overall missing (below LOD)", + "missing per protein median / max", "between-site median shift per protein: raw (log2) median", "same after harmonisation (z)", + "proteins associated with phenotype (FDR 5%)", "proteins with |age r| > 0.2", "proteins with |sex t| > 3"], + "value": [n_people, n_prot, meta["site"].nunique(), ", ".join(f"{k}: {v}" for k, v in meta["site"].value_counts().sort_index().items()), + f"{long['log2_intensity'].quantile(0.01):.1f} – {long['log2_intensity'].quantile(0.99):.1f}", f"{missing:.1%}", + f"{per_prot_missing.median():.1%} / {per_prot_missing.max():.1%}", round(site_spread_raw.median(), 3), round(site_spread_harm.median(), 3), + int((assoc_p["phenotype_q"] < 0.05).sum()), int((assoc_p["age_r"].abs() > 0.2).sum()), int((assoc_p["sex_t"].abs() > 3).sum())]}) + T("proteomics_summary", psum) + note = "" + if truth: + responsive = truth.get("n_phenotype_responsive_proteins") + note = (f" Ground truth: {responsive} proteins were generated to respond to the phenotype, {len(truth['causal_clusters'])} causal clusters, " + f"site shift sd {truth['site_shift_sd']}, LOD missing rate {truth['missing_rate_at_lod']} — the harmoniser must remove the site shift " + f"(it does: {site_spread_raw.median():.2f} → {site_spread_harm.median():.2f}) and the phenotype signal must survive it " + f"({int((assoc_p['phenotype_q'] < 0.05).sum())} proteins recovered at FDR 5%).") + md += ["## 7. Proteomics", "", psum.to_markdown(index=False), "", + "Missingness is MNAR at the detection limit (low-abundance proteins go missing first), so it is kept as *absence of an edge*, never imputed as zero " + "in the graph. The harmoniser is a robust z-score per (site, protein); it is what makes `site` a pure batch label the model must not be able to predict." + note, ""] + if not args.no_plots: + fig, axes = plt.subplots(1, 3, figsize=(12, 3.4)) + axes[0].hist(long["log2_intensity"], bins=60, color="#b26a0c"); axes[0].set_xlabel("log2 intensity"); axes[0].set_ylabel("measurements") + axes[1].scatter(long.groupby("protein_id")["log2_intensity"].median(), per_prot_missing.loc[long.groupby("protein_id")["log2_intensity"].median().index], s=6, color="#b26a0c") + axes[1].set_xlabel("protein median log2 intensity"); axes[1].set_ylabel("fraction missing"); axes[1].set_title("MNAR: low abundance → missing") + axes[2].hist(site_spread_raw, bins=40, histtype="step", color="#b26a0c", label="raw log2", linewidth=1.5) + axes[2].hist(site_spread_harm, bins=40, histtype="step", color="#1f6f8b", label="harmonised z", linewidth=1.5) + axes[2].set_xlabel("max−min of site medians per protein"); axes[2].legend(frameon=False); axes[2].set_title("batch effect before / after") + fig.tight_layout(); fig.savefig(out / "plots" / "07_proteomics.png", dpi=150); plt.close(fig) + + # -------------------------------------------------------------- 8. genome <-> proteome + if truth: + ind_idx = dict(zip(ind["individual_id"], ind["individual_idx"])) + rows = [] + for c in truth["causal_clusters"]: + sub = harm[harm["protein_id"] == c["cis_protein"]] + g = np.asarray(carries[[ind_idx[i] for i in sub["individual_id"]], c["cluster_idx"]].todense()).ravel() + r = np.corrcoef(g, sub["z"])[0, 1] if g.std() > 0 else np.nan + rows.append({"cluster_id": c["cluster_id"], "cis_protein": c["cis_protein"], "beta_cis_truth": round(c["beta_cis"], 2), + "carrier_freq": round(g.mean(), 3), "pearson_r": round(r, 3), "r2": round(r * r, 3)}) + cis = pd.DataFrame(rows).sort_values("r2", ascending=False) + T("genome_proteome_cis", cis) + md += ["## 8. Genome ↔ proteome (cis signal)", "", + f"For each ground-truth causal cluster, the correlation between carrying it and the harmonised level of the protein in its block: " + f"mean r² {cis['r2'].mean():.3f}, {int((cis['r2'] > 0.1).sum())} of {len(cis)} above 0.1. This is the ceiling any genome→proteome model can reach on this data " + f"and the reason the ridge baseline (`proteome_linear_baseline.py`) recovers only the strongest cis effects.", "", cis.to_markdown(index=False), ""] + + (out / "EDA.md").write_text("\n".join(md)) + print("\n".join(md)) + print(f"\nwrote {out}/EDA.md, {len(list((out / 'tables').glob('*.csv')))} tables, {len(list((out / 'plots').glob('*.png')))} plots") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/genomics/embeddings.py b/genomics/embeddings.py new file mode 100644 index 0000000..7461cf2 --- /dev/null +++ b/genomics/embeddings.py @@ -0,0 +1,112 @@ +#!/usr/bin/env python3 +"""How much phenotype structure do the embeddings carry? + +For every individual embedding found (SVD of the carrier matrix, and the last +hidden layer of each trained GNN) this script + * projects to 2-D with PCA and colours by ancestry and by sex (the control), + * computes silhouette scores by ancestry / population / sex, + * fits a 5-nearest-neighbour classifier on train individuals and scores the + test split -> "how linearly-separable-for-free is the phenotype in this space". +Cluster embeddings are projected too, coloured by the ancestry they are enriched in. +""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np +import pandas as pd +from sklearn.decomposition import PCA +from sklearn.metrics import balanced_accuracy_score, silhouette_score +from sklearn.neighbors import KNeighborsClassifier + +import haplokg + +ANCESTRY_COLOURS = {"AFR": "#d55e00", "AMR": "#cc79a7", "EAS": "#009e73", "EUR": "#0072b2", "SAS": "#e69f00"} +SEX_COLOURS = {"female": "#7b3294", "male": "#008837"} + + +def score(emb: np.ndarray, ind: pd.DataFrame, split: np.ndarray, label_maps: dict) -> dict: + out = {} + for target in ("ancestry", "population", "sex"): + y = ind[f"{target}_code"].to_numpy() + ok = y >= 0 + out[f"silhouette_{target}"] = float(silhouette_score(emb[ok], y[ok])) if len(np.unique(y[ok])) > 1 else float("nan") + tr, te = (split == "train") & ok, (split == "test") & ok + knn = KNeighborsClassifier(n_neighbors=5).fit(emb[tr], y[tr]) + out[f"knn5_test_balanced_accuracy_{target}"] = float(balanced_accuracy_score(y[te], knn.predict(emb[te]))) + return out + + +def plot_individuals(emb: np.ndarray, ind: pd.DataFrame, title: str, out: Path) -> None: + import matplotlib; matplotlib.use("Agg") + import matplotlib.pyplot as plt + xy = PCA(n_components=2, random_state=0).fit_transform(emb) + fig, axes = plt.subplots(1, 2, figsize=(12, 5.2)) + for ax, (column, colours, label) in zip(axes, [("ancestry", ANCESTRY_COLOURS, "ancestry"), ("sex", SEX_COLOURS, "sex (control)")]): + values = ind[column].fillna("unlabelled") + for name, colour in list(colours.items()) + [("unlabelled", "#bbbbbb")]: + m = (values == name).to_numpy() + if m.any(): + ax.scatter(xy[m, 0], xy[m, 1], s=7, color=colour, label=f"{name} ({m.sum()})", alpha=0.75, linewidths=0) + ax.set_title(f"{title} - coloured by {label}"); ax.set_xlabel("PC1"); ax.set_ylabel("PC2") + ax.legend(frameon=False, markerscale=2, fontsize=8, loc="upper center", bbox_to_anchor=(0.5, -0.14), ncol=3) # below the axes, never on the points + fig.tight_layout(); fig.savefig(out, dpi=150, bbox_inches="tight"); plt.close(fig) + + +def plot_clusters(emb: np.ndarray, assoc: pd.DataFrame, title: str, out: Path) -> None: + import matplotlib; matplotlib.use("Agg") + import matplotlib.pyplot as plt + xy = PCA(n_components=2, random_state=0).fit_transform(emb) + fig, ax = plt.subplots(figsize=(7, 5.8)) + strong = assoc["ancestry_cramers_v"].to_numpy() >= 0.3 + ax.scatter(xy[~strong, 0], xy[~strong, 1], s=4, color="#cccccc", label="weakly ancestry-associated (V<0.3)", linewidths=0) + for name, colour in ANCESTRY_COLOURS.items(): + m = strong & (assoc["ancestry_dominant"].to_numpy() == name) + ax.scatter(xy[m, 0], xy[m, 1], s=6, color=colour, label=f"enriched in {name} ({m.sum()})", alpha=0.8, linewidths=0) + ax.set_title(title); ax.set_xlabel("PC1"); ax.set_ylabel("PC2"); ax.legend(frameon=False, markerscale=2, fontsize=8, loc="upper center", bbox_to_anchor=(0.5, -0.12), ncol=2) + fig.tight_layout(); fig.savefig(out, dpi=150, bbox_inches="tight"); plt.close(fig) + + +def main() -> int: + here = Path(__file__).resolve().parent + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--chrom", default="chr22") + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + kg_dir = here / "outputs" / "kg" / args.chrom + out_dir = here / "outputs" / "embeddings" / args.chrom + out_dir.mkdir(parents=True, exist_ok=True) + + kg = haplokg.load_kg(kg_dir) + ind = kg["individuals"] + split = haplokg.load_or_make_split(kg, here / "outputs" / "splits" / args.chrom / f"split_seed{args.seed}.csv", seed=args.seed) + assoc_path = here / "outputs" / "cooccurrence" / args.chrom / "cluster_phenotype_association.csv" + assoc = pd.read_csv(assoc_path) if assoc_path.exists() else None + + sources = {} + for path in sorted(out_dir.glob("svd*_individual.npy")): + sources[path.stem.replace("_individual", "")] = (path, path.with_name(path.name.replace("individual", "cluster"))) + for path in sorted((here / "outputs" / "gnn" / args.chrom).glob("*/embedding_individual.npy")): + sources[f"gnn_{path.parent.name}"] = (path, path.with_name("embedding_cluster.npy")) + + rows = [] + for name, (ind_path, cl_path) in sources.items(): + emb = np.load(ind_path) + result = {"embedding": name, "dim": emb.shape[1], **score(emb, ind, split, kg["label_maps"])} + rows.append(result) + plot_individuals(emb, ind, name, out_dir / f"individuals_{name}.png") + if cl_path.exists() and assoc is not None: + plot_clusters(np.load(cl_path), assoc, f"clusters - {name}", out_dir / f"clusters_{name}.png") + print(f"{name:32s} dim={emb.shape[1]:<4d} " + " ".join(f"{k.split('_')[-1]}: sil={result[f'silhouette_{k.split('_')[-1]}']:.3f} " + f"knn={result[f'knn5_test_balanced_accuracy_{k.split('_')[-1]}']:.3f}" for k in ("silhouette_ancestry", "silhouette_population", "silhouette_sex"))) + table = pd.DataFrame(rows) + table.to_csv(out_dir / "embedding_quality.csv", index=False) + (out_dir / "embedding_quality.json").write_text(json.dumps(rows, indent=2)) + print(f"wrote {out_dir}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/genomics/federated/client.py b/genomics/federated/client.py new file mode 100644 index 0000000..2a867bc --- /dev/null +++ b/genomics/federated/client.py @@ -0,0 +1,146 @@ +#!/usr/bin/env python3 +"""NVFlare Client API training script: one hospital site of the ProGenome graph. + +What stays at the site: its own Individual nodes (labels, CARRIES and MEASURED edges). What is shared: +the public cluster / block / gene / protein graph and the model weights. Each round the site receives the +global weights, evaluates them on its own validation and test people, trains a few local full-batch +epochs on its own training people, and sends back weights + metrics + the number of optimizer steps. + +Run only through NVFlare (job.py); not a standalone trainer. +""" +from __future__ import annotations + +import argparse +import sys +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F +from sklearn.metrics import balanced_accuracy_score, roc_auc_score + +import nvflare.client as flare +from nvflare.app_common.abstract.fl_model import MetaKey + +CO = ("cluster", "co_occurs", "cluster") +MEAS = ("individual", "measured", "protein") +RMEAS = ("protein", "rev_measured", "individual") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(allow_abbrev=False) + parser.add_argument("--genomics-dir", required=True, help="the genomics/ folder (for haplokg imports)") + parser.add_argument("--kg-dir", required=True, help="outputs/kg/ with hetero_v2.pt and the v2 tables") + parser.add_argument("--split-path", required=True) + parser.add_argument("--model-config", required=True) + parser.add_argument("--local-epochs", type=int, default=5) + parser.add_argument("--lr", type=float, default=5e-3) + parser.add_argument("--weight-decay", type=float, default=5e-4) + parser.add_argument("--seed", type=int, default=42) + return parser.parse_args() + + +def site_data(args, site_name: str, device): + """Load the public graph, keep only this site's individuals, build model inputs.""" + sys.path.insert(0, args.genomics_dir) + import haplokg + import haplokg_proteins as hp + import json + + kg = haplokg.load_kg(args.kg_dir) + pt = hp.load_protein_tables(args.kg_dir) + data = torch.load(Path(args.kg_dir) / "hetero_v2.pt", weights_only=False) + cfg = json.loads(Path(args.model_config).read_text()) + relations = [tuple(r) for r in cfg["relations"]] + + ind2 = pt["individuals"] + sites = pt["label_maps"]["site"] # e.g. ["SITE1","SITE2","SITE3"] + site_idx = int(site_name.rsplit("-", 1)[-1]) - 1 # site-1 -> SITE1 + my_site = sites[site_idx] + members = np.flatnonzero((ind2["site_code"] == site_idx).to_numpy() & (ind2["phenotype_code"] >= 0).to_numpy()) + split_frame = __import__("pandas").read_csv(args.split_path) + split = split_frame["split"].to_numpy()[members] + + # HeteroData.subgraph slices every attribute of the store; the id lists are plain Python lists, so drop them first + for key in list(data["individual"].keys()): + if key != "num_nodes" and not isinstance(data["individual"][key], torch.Tensor): + del data["individual"][key] + sub = data.subgraph({"individual": torch.as_tensor(members, dtype=torch.long)}) # other node types untouched + x_ind = np.concatenate([kg["carries"][members].toarray().astype(np.float32), + pt["abundance"][members].toarray().astype(np.float32), + pt["observed"][members].toarray().astype(np.float32)], axis=1) + x_dict = {"individual": torch.from_numpy(x_ind), "cluster": data["cluster"].x, "block": data["block"].x, + "gene": data["gene"].x, "protein": data["protein"].x} + edge_index = {rel: sub[rel].edge_index for rel in relations} + lift = data[CO].edge_attr[:, 1] + edge_weight = {CO: torch.log(lift) / torch.log(lift).max(), MEAS: sub[MEAS].edge_attr[:, 0], RMEAS: sub[RMEAS].edge_attr[:, 0]} + y = sub["individual"].y_phenotype + masks = {k: torch.from_numpy(split == k) for k in ("train", "val", "test")} + to = lambda d: {k: v.to(device) for k, v in d.items()} + return my_site, to(x_dict), to(edge_index), to(edge_weight), y.to(device), to(masks) + + +def evaluate(model, x, ei, ew, y, mask) -> dict: + was_training = model.training + model.eval() + try: + with torch.no_grad(): + logits, _ = model(x, ei, ew) + finally: + model.train(was_training) + prob = torch.softmax(logits, 1)[:, 1] + yt = y[mask].cpu().numpy() + pb = prob[mask].cpu().numpy() + pr = (pb >= 0.5).astype(int) + if len(yt) == 0: + raise RuntimeError("evaluation split is empty at this site") + out = {"balanced_accuracy": float(balanced_accuracy_score(yt, pr)), "n": int(len(yt))} + if len(np.unique(yt)) == 2: + out["auc"] = float(roc_auc_score(yt, pb)) + return out + + +def main() -> None: + args = parse_args() + torch.manual_seed(args.seed) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + flare.init() + site_name = flare.get_site_name() + my_site, x, ei, ew, y, masks = site_data(args, site_name, device) + sys.path.insert(0, str(Path(__file__).resolve().parent)) + from model import ProGenomeGNN + + model = ProGenomeGNN(args.model_config).to(device) + optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) + counts = torch.bincount(y[masks["train"]], minlength=2).float() + class_weight = (counts.sum() / counts.clamp(min=1) / 2).to(device) # local class balance, a training policy + print(f"[{site_name}] {my_site}: train {int(masks['train'].sum())} val {int(masks['val'].sum())} test {int(masks['test'].sum())} on {device}", flush=True) + + while flare.is_running(): + input_model = flare.receive() + model.load_state_dict(input_model.params) + val = evaluate(model, x, ei, ew, y, masks["val"]) + test = evaluate(model, x, ei, ew, y, masks["test"]) + metrics = {"val_balanced_accuracy": val["balanced_accuracy"], "test_balanced_accuracy": test["balanced_accuracy"], + "test_auc": test.get("auc", float("nan")), "n_val": val["n"], "n_test": test["n"]} + print(f"[{site_name}] round {input_model.current_round}: global model val bal-acc {val['balanced_accuracy']:.3f} test AUC {metrics['test_auc']:.3f}", flush=True) + if flare.is_evaluate(): + flare.send(flare.FLModel(metrics=metrics)) + continue + + steps = 0 + model.train() + for _ in range(args.local_epochs): # full-batch: one optimizer step per epoch + optimizer.zero_grad() + logits, _ = model(x, ei, ew) + loss = F.cross_entropy(logits[masks["train"]], y[masks["train"]], weight=class_weight) + loss.backward() + optimizer.step() + steps += 1 + params = {k: v.detach().cpu() for k, v in model.state_dict().items()} + flare.send(flare.FLModel(params=params, metrics=metrics, meta={MetaKey.NUM_STEPS_CURRENT_ROUND: steps})) + + +if __name__ == "__main__": + main() diff --git a/genomics/federated/evaluate_global.py b/genomics/federated/evaluate_global.py new file mode 100644 index 0000000..1c5f23d --- /dev/null +++ b/genomics/federated/evaluate_global.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 +"""Score the federated global model centrally on the same held-out people as the central model. + + python federated/evaluate_global.py # finds FL_global_model.pt in outputs/federated//workspace + +Loads the server's saved global model, rebuilds the full v2 graph exactly as train_gnn_v2.py does +(--modality both --init raw) and reports test balanced accuracy / AUC next to the central run's numbers. +""" +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +import numpy as np +import pandas as pd +import torch +from sklearn.metrics import balanced_accuracy_score, roc_auc_score + +HERE = Path(__file__).resolve().parent +GENOMICS = HERE.parent +sys.path.insert(0, str(GENOMICS)); sys.path.insert(0, str(HERE)) +import haplokg # noqa: E402 +import haplokg_proteins as hp # noqa: E402 +from model import ProGenomeGNN # noqa: E402 + +CO = ("cluster", "co_occurs", "cluster") +MEAS = ("individual", "measured", "protein") +RMEAS = ("protein", "rev_measured", "individual") + + +def load_global_state(workspace: Path, filename: str) -> dict: + hits = sorted(workspace.rglob(filename)) + if not hits: + raise SystemExit(f"no {filename} under {workspace}") + ckpt = torch.load(hits[-1], map_location="cpu", weights_only=False) + state = ckpt.get("model", ckpt) if isinstance(ckpt, dict) else ckpt + return {k: (torch.as_tensor(v) if not isinstance(v, torch.Tensor) else v) for k, v in state.items()}, hits[-1] + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter, allow_abbrev=False) + parser.add_argument("--chrom", default="chr22") + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--model-file", default="FL_global_model.pt", help="or best_FL_global_model.pt") + args = parser.parse_args() + kg_dir = GENOMICS / "outputs" / "kg" / args.chrom + fed_dir = GENOMICS / "outputs" / "federated" / args.chrom + cfg_path = fed_dir / "model_args.json" + + kg = haplokg.load_kg(kg_dir); pt = hp.load_protein_tables(kg_dir) + data = torch.load(kg_dir / "hetero_v2.pt", weights_only=False) + split = haplokg.load_or_make_split(kg, GENOMICS / "outputs" / "splits" / args.chrom / f"split_seed{args.seed}.csv", seed=args.seed) + cfg = json.loads(cfg_path.read_text()) + relations = [tuple(r) for r in cfg["relations"]] + x = {"individual": torch.from_numpy(np.concatenate([kg["carries"].toarray().astype(np.float32), pt["abundance"].toarray().astype(np.float32), + pt["observed"].toarray().astype(np.float32)], axis=1)), + "cluster": data["cluster"].x, "block": data["block"].x, "gene": data["gene"].x, "protein": data["protein"].x} + ei = {rel: data[rel].edge_index for rel in relations} + lift = data[CO].edge_attr[:, 1] + ew = {CO: torch.log(lift) / torch.log(lift).max(), MEAS: data[MEAS].edge_attr[:, 0], RMEAS: data[RMEAS].edge_attr[:, 0]} + y = data["individual"].y_phenotype.numpy() + test = (split == "test") & (y >= 0) + + state, path = load_global_state(fed_dir / "workspace", args.model_file) + model = ProGenomeGNN(str(cfg_path)); model.load_state_dict(state); model.eval() + with torch.no_grad(): + logits, _ = model(x, ei, ew) + prob = torch.softmax(logits, 1)[:, 1].numpy() + fed = {"balanced_accuracy": float(balanced_accuracy_score(y[test], (prob[test] >= 0.5).astype(int))), + "auc": float(roc_auc_score(y[test], prob[test])), "n_test": int(test.sum()), "model_file": str(path)} + + central_path = GENOMICS / "outputs" / "gnn_v2" / args.chrom / "phenotype_both_raw" / "metrics.json" + central_path = central_path if central_path.exists() else GENOMICS / "outputs_brev" / "progenome-a100" / "gnn_v2" / args.chrom / "phenotype_both_raw" / "metrics.json" + central = json.loads(central_path.read_text())["test"] if central_path.exists() else {} + # per-site view of the same test people (what each hospital would see) + ind2 = pt["individuals"]; sites = pt["label_maps"]["site"] + per_site = [] + for i, s in enumerate(sites): + m = test & (ind2["site_code"].to_numpy() == i) + per_site.append({"site": s, "n_test": int(m.sum()), "balanced_accuracy": float(balanced_accuracy_score(y[m], (prob[m] >= 0.5).astype(int))), + "auc": float(roc_auc_score(y[m], prob[m]))}) + report = {"federated_global_model": fed, "central_model_same_test_people": {k: central.get(k) for k in ("balanced_accuracy", "roc_auc", "n")}, + "federated_per_site": per_site} + (fed_dir / "evaluation.json").write_text(json.dumps(report, indent=2)) + print(json.dumps(report, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/genomics/federated/job.py b/genomics/federated/job.py new file mode 100644 index 0000000..fd4a22c --- /dev/null +++ b/genomics/federated/job.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 +"""Federated training of the ProGenome genome+proteome GNN with NVFlare FedAvg (simulation). + + python federated/job.py # 3 sites, 10 rounds x 5 local epochs, workspace under outputs/federated/ + python federated/job.py --rounds 20 --local-epochs 3 + +Sites = the 3 mixed-ancestry hospital sites of the synthetic proteomics. Each site trains on its own +Individual nodes only; the cluster/block/gene/protein graph is public and identical everywhere; only model +weights travel. Evaluate the resulting global model centrally with federated/evaluate_global.py. +""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch +from nvflare.app_opt.pt.recipes.fedavg import FedAvgRecipe +from nvflare.client.config import ExchangeFormat +from nvflare.recipe import SimEnv + +HERE = Path(__file__).resolve().parent # genomics/federated +GENOMICS = HERE.parent +CO = ("cluster", "co_occurs", "cluster") +RELATIONS = [["individual", "carries", "cluster"], ["cluster", "rev_carries", "individual"], + ["cluster", "in_block", "block"], ["block", "rev_in_block", "cluster"], list(CO), ["block", "next_block", "block"], + ["individual", "measured", "protein"], ["protein", "rev_measured", "individual"], + ["block", "overlaps", "gene"], ["gene", "rev_overlaps", "block"], ["gene", "encodes", "protein"], ["protein", "rev_encodes", "gene"]] + + +def write_model_config(kg_dir: Path, cfg_path: Path, hidden: int, layers: int, dropout: float) -> dict: + """Derive every constructor value from the public graph so server and sites build the same model.""" + data = torch.load(kg_dir / "hetero_v2.pt", weights_only=False) + n_clusters = data["cluster"].x.shape[0] + n_proteins = data["protein"].x.shape[0] + cfg = {"in_dims": {"individual": n_clusters + 2 * n_proteins, # raw carrier row + protein z + observed mask + "cluster": data["cluster"].x.shape[1], "block": data["block"].x.shape[1], + "gene": data["gene"].x.shape[1], "protein": data["protein"].x.shape[1]}, + "relations": RELATIONS, "hidden": hidden, "out_dim": 2, "layers": layers, "dropout": dropout, "aggr": "mean"} + cfg_path.write_text(json.dumps(cfg, indent=2)) + return cfg + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter, allow_abbrev=False) + parser.add_argument("--chrom", default="chr22") + parser.add_argument("--sites", type=int, default=3) + parser.add_argument("--rounds", type=int, default=10) + parser.add_argument("--local-epochs", type=int, default=5) + parser.add_argument("--hidden", type=int, default=64) + parser.add_argument("--layers", type=int, default=2) + parser.add_argument("--dropout", type=float, default=0.3) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--workspace", type=Path, default=None, help="default outputs/federated//workspace") + args = parser.parse_args() + + kg_dir = GENOMICS / "outputs" / "kg" / args.chrom + split_path = GENOMICS / "outputs" / "splits" / args.chrom / f"split_seed{args.seed}.csv" + out_dir = GENOMICS / "outputs" / "federated" / args.chrom + out_dir.mkdir(parents=True, exist_ok=True) + workspace = args.workspace or out_dir / "workspace" + cfg_path = out_dir / "model_args.json" + cfg = write_model_config(kg_dir, cfg_path, args.hidden, args.layers, args.dropout) + print("model config:", json.dumps(cfg["in_dims"]), f"hidden {cfg['hidden']} layers {cfg['layers']}") + + train_args = (f"--genomics-dir {GENOMICS} --kg-dir {kg_dir} --split-path {split_path} --model-config {cfg_path} " + f"--local-epochs {args.local_epochs} --seed {args.seed}") + recipe = FedAvgRecipe( + name="progenome_fedavg", + model={"class_path": "model.ProGenomeGNN", "args": {"config_path": str(cfg_path)}}, + min_clients=args.sites, + num_rounds=args.rounds, + train_script=str(HERE / "client.py"), + train_args=train_args, + key_metric="val_balanced_accuracy", + key_metric_mode="max", + server_expected_format=ExchangeFormat.PYTORCH, + ) + recipe.add_decomposers(["nvflare.app_opt.pt.decomposers.TensorDecomposer"]) + recipe.add_server_file(str(HERE / "model.py")) + + env = SimEnv(num_clients=args.sites, workspace_root=str(workspace)) + run = recipe.execute(env) + print("status:", run.get_status()) + print("result:", run.get_result()) + (out_dir / "run_info.json").write_text(json.dumps({"rounds": args.rounds, "local_epochs": args.local_epochs, "sites": args.sites, + "workspace": str(workspace), "status": str(run.get_status())}, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/genomics/federated/local_only.py b/genomics/federated/local_only.py new file mode 100644 index 0000000..5774e34 --- /dev/null +++ b/genomics/federated/local_only.py @@ -0,0 +1,87 @@ +#!/usr/bin/env python3 +"""What federated learning buys here: each site training ALONE on its own people, versus the federated global model. + + python federated/local_only.py # after federated/job.py + evaluate_global.py (uses their model_args.json) + +For every site the script builds the site's subgraph exactly as client.py does, trains the same model for the same +number of optimizer steps the site performed in the federated run (rounds x local epochs, default 30 x 5 = 150), +and scores it on the site's own held-out people and on the other sites' held-out people (what a lone hospital's +model would do on another hospital's patients). Compare with federated_per_site in evaluation.json, where one +global model was trained on all sites' people without any row leaving its site. +""" +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path + +import torch +import torch.nn.functional as F + +HERE = Path(__file__).resolve().parent +GENOMICS = HERE.parent +sys.path.insert(0, str(GENOMICS)); sys.path.insert(0, str(HERE)) +import client # noqa: E402 (site_data / evaluate; importing does not start NVFlare) +from job import write_model_config # noqa: E402 +from model import ProGenomeGNN # noqa: E402 + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter, allow_abbrev=False) + parser.add_argument("--chrom", default="chr22") + parser.add_argument("--sites", type=int, default=3) + parser.add_argument("--steps", type=int, default=150, help="optimizer steps per site = rounds x local epochs of the federated run") + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + kg_dir = GENOMICS / "outputs" / "kg" / args.chrom + fed_dir = GENOMICS / "outputs" / "federated" / args.chrom + fed_dir.mkdir(parents=True, exist_ok=True) + cfg_path = fed_dir / "model_args.json" + if not cfg_path.exists(): + write_model_config(kg_dir, cfg_path, hidden=64, layers=2, dropout=0.3) + ns = argparse.Namespace(genomics_dir=str(GENOMICS), kg_dir=str(kg_dir), model_config=str(cfg_path), + split_path=str(GENOMICS / "outputs" / "splits" / args.chrom / f"split_seed{args.seed}.csv")) + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + sites = {} + for i in range(1, args.sites + 1): + name, x, ei, ew, y, masks = client.site_data(ns, f"site-{i}", device) + sites[name] = (x, ei, ew, y, masks) + print(f"{name}: train {int(masks['train'].sum())} test {int(masks['test'].sum())}", flush=True) + + results = {} + for name, (x, ei, ew, y, masks) in sites.items(): + torch.manual_seed(args.seed) + model = ProGenomeGNN(str(cfg_path)).to(device) + optimizer = torch.optim.Adam(model.parameters(), lr=5e-3, weight_decay=5e-4) + counts = torch.bincount(y[masks["train"]], minlength=2).float() + class_weight = (counts.sum() / counts.clamp(min=1) / 2).to(device) + model.train() + for _ in range(args.steps): + optimizer.zero_grad() + logits, _ = model(x, ei, ew) + F.cross_entropy(logits[masks["train"]], y[masks["train"]], weight=class_weight).backward() + optimizer.step() + scores = {"own_test": client.evaluate(model, x, ei, ew, y, masks["test"])} + for other, (x2, ei2, ew2, y2, m2) in sites.items(): + if other != name: + scores[f"on_{other}_test"] = client.evaluate(model, x2, ei2, ew2, y2, m2["test"]) + results[name] = scores + print(f"{name} alone ({args.steps} steps): own test AUC {scores['own_test'].get('auc', float('nan')):.3f} " + + " ".join(f"{k} AUC {v.get('auc', float('nan')):.3f}" for k, v in scores.items() if k != "own_test"), flush=True) + + ev_path = fed_dir / "evaluation.json" + ev = json.loads(ev_path.read_text()) if ev_path.exists() else {} + report = {"steps_per_site": args.steps, "local_only": results, + "federated_per_site": ev.get("federated_per_site"), "federated_global_model": ev.get("federated_global_model"), + "central_model_same_test_people": ev.get("central_model_same_test_people")} + (fed_dir / "local_only_vs_federated.json").write_text(json.dumps(report, indent=2)) + if ev: + print("\nfederated global model on each site's own test people: " + + " ".join(f"{s['site']} AUC {s['auc']:.3f}" for s in ev["federated_per_site"])) + print(f"wrote {fed_dir / 'local_only_vs_federated.json'}") + + +if __name__ == "__main__": + main() diff --git a/genomics/federated/model.py b/genomics/federated/model.py new file mode 100644 index 0000000..dc83963 --- /dev/null +++ b/genomics/federated/model.py @@ -0,0 +1,41 @@ +"""Model shared by the NVFlare server and every site. + +Architecture is identical to `train_gnn_v2.HeteroGNNv2` (kept here as a self-contained copy so the +server app can import it without the rest of the pipeline). All constructor values come from one JSON +file written by job.py, so server and clients build byte-identical state dicts. +""" +from __future__ import annotations + +import json +from pathlib import Path + +import torch +import torch.nn.functional as F +from torch import nn +from torch_geometric.nn import GraphConv, HeteroConv, SAGEConv + +WEIGHTED = {("cluster", "co_occurs", "cluster"), ("individual", "measured", "protein"), ("protein", "rev_measured", "individual")} + + +class ProGenomeGNN(nn.Module): + def __init__(self, config_path: str): + super().__init__() + cfg = json.loads(Path(config_path).read_text()) + in_dims: dict = cfg["in_dims"] + relations = [tuple(r) for r in cfg["relations"]] + hidden, out_dim, layers, dropout, aggr = cfg["hidden"], cfg["out_dim"], cfg["layers"], cfg["dropout"], cfg["aggr"] + self.node_types = list(in_dims) + self.proj = nn.ModuleDict({t: nn.Linear(d, hidden) for t, d in in_dims.items()}) + self.convs = nn.ModuleList([HeteroConv({ + rel: (GraphConv(hidden, hidden, aggr=aggr) if rel in WEIGHTED else SAGEConv((hidden, hidden), hidden, aggr=aggr)) + for rel in relations}, aggr="sum") for _ in range(layers)]) + self.norms = nn.ModuleList([nn.ModuleDict({t: nn.LayerNorm(hidden) for t in self.node_types}) for _ in range(layers)]) + self.head = nn.Linear(hidden, out_dim) + self.dropout = dropout + + def forward(self, x_dict, edge_index_dict, edge_weight_dict): + h = {t: self.proj[t](x_dict[t]) for t in self.node_types} + for conv, norm in zip(self.convs, self.norms): + out = conv(h, edge_index_dict, edge_weight_dict=edge_weight_dict) + h = {t: F.dropout(F.relu(norm[t](out[t] + h[t])), p=self.dropout, training=self.training) if t in out else h[t] for t in h} + return self.head(h["individual"]), h diff --git a/genomics/fetch_data.sh b/genomics/fetch_data.sh new file mode 100755 index 0000000..94b06cc --- /dev/null +++ b/genomics/fetch_data.sh @@ -0,0 +1,39 @@ +#!/usr/bin/env bash +# Download the chr22 HaploGraph (pre-built 1000G haploblock-cluster graph), +# the real 1000G phenotype labels and the per-block statistics into data/. +# Source: https://data.haploblocks.org (built for this hackathon, see +# haplograph/1000G/README.txt on that server). Idempotent; verifies md5. +set -euo pipefail + +HERE="$(cd "$(dirname "$0")" && pwd)" +. "$HERE/load_env.sh" # HAPLOBLOCKS_BASE, CHROM from .env if set +CHR="${1:-${CHROM:-chr22}}" +BASE="${HAPLOBLOCKS_BASE:-https://data.haploblocks.org}" +DATA="$HERE/data" + +mkdir -p "$DATA/haplograph/$CHR" "$DATA/haploblocks" + +fetch() { # fetch + if [ -s "$2" ]; then echo "have $2"; else echo "fetch $1"; curl -sSL --fail --max-time 600 -o "$2" "$1"; fi +} + +for f in nodes.csv.gz edges_lift_above_threshold.csv.gz islands.csv.gz top_edges_by_lift.csv.gz; do + fetch "$BASE/haplograph/1000G/$CHR/$f" "$DATA/haplograph/$CHR/$f" +done +fetch "$BASE/haplograph/1000G/phenotypes_real.csv" "$DATA/haplograph/phenotypes_real.csv" +fetch "$BASE/haplograph/1000G/README.txt" "$DATA/haplograph/README.txt" +fetch "$BASE/haplograph/1000G/checksums.md5" "$DATA/haplograph/checksums.md5" +fetch "$BASE/bidirectional_blast_samples/block_stats.tsv" "$DATA/haploblocks/block_stats.tsv" +fetch "$BASE/haploblock_hashes/1000G/$CHR/${CHR}_haploblock_boundaries_${CHR}.tsv" \ + "$DATA/haploblocks/${CHR}_haploblock_boundaries_${CHR}.tsv" + +# md5 verification of the graph files (checksums.md5 uses ./chrN/file paths) +cd "$DATA/haplograph" +if command -v md5sum >/dev/null; then + grep -E "^\S+ \./$CHR/(nodes|edges_lift_above_threshold|islands|top_edges_by_lift)\.csv\.gz$" checksums.md5 | md5sum -c - +else # macOS without coreutils + grep -E "\./$CHR/(nodes|edges_lift_above_threshold|islands|top_edges_by_lift)\.csv\.gz$" checksums.md5 | while read -r sum path; do + [ "$sum" = "$(md5 -q "$path")" ] && echo "$path: OK" || { echo "$path: MISMATCH"; exit 1; } + done +fi +echo "done -> $DATA" diff --git a/genomics/graph_explore.py b/genomics/graph_explore.py new file mode 100644 index 0000000..c34b71c --- /dev/null +++ b/genomics/graph_explore.py @@ -0,0 +1,204 @@ +#!/usr/bin/env python3 +"""Look at the graph *before* embedding it. + +Builds the cluster co-occurrence graph (nodes = haploblock clusters, edges = +lift >= 5 co-occurrence) with NetworkX, prints its statistics, exports GraphML +for Gephi / Cytoscape, and draws: + + edge_positions.png every edge as (position of source, position of target) + coloured by lift -> long-range structure along the chromosome + region_-.png the subgraph of one region, nodes coloured by the + ancestry they are enriched in, sized by carrier support + degree_distribution.png + +With the nx-cugraph backend installed (GPU image) NetworkX dispatches the heavy +algorithms to cuGraph automatically when NX_CUGRAPH_AUTOCONFIG=True. +""" +from __future__ import annotations + +import argparse +import gzip +import json +import os +from pathlib import Path + +import networkx as nx +import numpy as np +import pandas as pd + +import haplokg + +ANCESTRY_COLOURS = {"AFR": "#d55e00", "AMR": "#cc79a7", "EAS": "#009e73", "EUR": "#0072b2", "SAS": "#e69f00"} + + +def backend_note() -> str: + try: + import nx_cugraph # noqa: F401 + auto = os.environ.get("NX_CUGRAPH_AUTOCONFIG", "") + return f"nx-cugraph installed (NX_CUGRAPH_AUTOCONFIG={auto or 'unset'})" + except ImportError: + return "nx-cugraph not installed -> pure NetworkX on CPU" + + +def cluster_graph(kg: dict, assoc: pd.DataFrame | None) -> nx.Graph: + cl = kg["clusters"].copy() + coords = cl["block_id"].map(lambda b: haplokg.parse_block_id(b)) + cl["start"] = [c[1] for c in coords] + cl["end"] = [c[2] for c in coords] + if assoc is not None: + cl = cl.merge(assoc[["cluster_idx", "ancestry_cramers_v", "ancestry_dominant"]], on="cluster_idx", how="left") + G = nx.Graph(name="haploblock cluster co-occurrence") + for row in cl.itertuples(index=False): + G.add_node(int(row.cluster_idx), cluster_id=row.cluster_id, block_id=row.block_id, block_idx=int(row.block_idx), + start=int(row.start), end=int(row.end), support=int(row.support), + ancestry_dominant=getattr(row, "ancestry_dominant", "") or "", + ancestry_v=float(getattr(row, "ancestry_cramers_v", float("nan")))) + co = kg["co_occurs"] + G.add_edges_from(zip(co["src"].astype(int), co["dst"].astype(int), + ({"weight": float(w), "lift": float(l)} for w, l in zip(co["weight"], co["lift"])))) + return G + + +def block_graph(G: nx.Graph) -> nx.Graph: + """Collapse clusters into their haploblocks; edge = total lift between two blocks.""" + B = nx.Graph(name="haploblock co-occurrence (collapsed)") + for _, d in G.nodes(data=True): + if d["block_idx"] not in B: + B.add_node(d["block_idx"], block_id=d["block_id"], start=d["start"], end=d["end"], n_clusters=0) + B.nodes[d["block_idx"]]["n_clusters"] += 1 + for u, v, d in G.edges(data=True): + bu, bv = G.nodes[u]["block_idx"], G.nodes[v]["block_idx"] + if bu == bv: + continue + if B.has_edge(bu, bv): + B[bu][bv]["total_lift"] += d["lift"]; B[bu][bv]["n_edges"] += 1 + else: + B.add_edge(bu, bv, total_lift=d["lift"], n_edges=1) + return B + + +def statistics(G: nx.Graph, B: nx.Graph) -> dict: + degrees = np.array([d for _, d in G.degree()]) + lift_degree = dict(G.degree(weight="lift")) + components = sorted((len(c) for c in nx.connected_components(G)), reverse=True) + hubs = sorted(G.nodes, key=lambda n: G.degree(n), reverse=True)[:10] + return { + "backend": backend_note(), + "clusters": G.number_of_nodes(), "co_occurrence_edges": G.number_of_edges(), + "density": float(nx.density(G)), + "degree_mean": float(degrees.mean()), "degree_median": float(np.median(degrees)), "degree_max": int(degrees.max()), + "isolated_clusters": int((degrees == 0).sum()), + "connected_components": len(components), "largest_component": components[0] if components else 0, + "average_clustering_coefficient": float(nx.average_clustering(G)), + "blocks_in_graph": B.number_of_nodes(), "block_block_edges": B.number_of_edges(), + "same_block_edges": int(sum(1 for u, v in G.edges if G.nodes[u]["block_idx"] == G.nodes[v]["block_idx"])), + "top_hubs": [{"cluster_id": G.nodes[n]["cluster_id"], "degree": G.degree(n), "lift_sum": round(lift_degree[n], 1), + "support": G.nodes[n]["support"], "ancestry_dominant": G.nodes[n]["ancestry_dominant"]} for n in hubs], + } + + +def default_region(data_dir: Path, chrom: str) -> tuple[int, int]: + """The densest published 'island' (extended haplotype) is a good first thing to look at.""" + islands = data_dir / "haplograph" / chrom / "islands.csv.gz" + if islands.exists(): + with gzip.open(islands, "rt") as fh: + top = pd.read_csv(fh).iloc[0] + return int(top["span_start"]), int(top["span_end"]) + return 45_000_000, 45_500_000 + + +def plot_edge_positions(G: nx.Graph, chrom: str, out: Path) -> None: + import matplotlib; matplotlib.use("Agg") + import matplotlib.pyplot as plt + xs, ys, lifts = [], [], [] + for u, v, d in G.edges(data=True): + a, b = sorted((G.nodes[u]["start"], G.nodes[v]["start"])) + xs.append(a / 1e6); ys.append(b / 1e6); lifts.append(d["lift"]) + order = np.argsort(lifts) + fig, ax = plt.subplots(figsize=(7, 6.5)) + sc = ax.scatter(np.array(xs)[order], np.array(ys)[order], c=np.array(lifts)[order], s=2, cmap="viridis", alpha=0.6, + norm=matplotlib.colors.LogNorm()) + ax.set_xlabel(f"{chrom} position of cluster A (Mb)"); ax.set_ylabel(f"{chrom} position of cluster B (Mb)") + ax.set_title(f"{G.number_of_edges():,} co-occurrence edges (lift >= 5) between {G.number_of_nodes():,} clusters") + fig.colorbar(sc, ax=ax, label="lift (log scale)") + fig.tight_layout(); fig.savefig(out, dpi=150); plt.close(fig) + + +def plot_region(G: nx.Graph, chrom: str, start: int, end: int, out: Path, seed: int = 42) -> dict: + import matplotlib; matplotlib.use("Agg") + import matplotlib.pyplot as plt + from matplotlib.lines import Line2D + nodes = [n for n, d in G.nodes(data=True) if d["end"] >= start and d["start"] <= end] + H = G.subgraph(nodes).copy() + H.remove_nodes_from([n for n in H if H.degree(n) == 0]) + if H.number_of_nodes() == 0: + return {"region": f"{chrom}:{start}-{end}", "nodes": 0, "edges": 0} + pos = nx.spring_layout(H, weight="lift", seed=seed, k=1.5 / np.sqrt(H.number_of_nodes())) + colours = [ANCESTRY_COLOURS.get(H.nodes[n]["ancestry_dominant"], "#888888") for n in H] + sizes = [20 + 180 * H.nodes[n]["support"] / 2548 for n in H] + widths = [0.2 + 1.5 * np.log10(d["lift"] / 5 + 1) for _, _, d in H.edges(data=True)] + fig, ax = plt.subplots(figsize=(9, 8)) + nx.draw_networkx_edges(H, pos, ax=ax, width=widths, edge_color="#b0b0b0", alpha=0.5) + nx.draw_networkx_nodes(H, pos, ax=ax, node_color=colours, node_size=sizes, linewidths=0.3, edgecolors="white") + handles = [Line2D([0], [0], marker="o", color="w", markerfacecolor=c, markersize=9, label=a) for a, c in ANCESTRY_COLOURS.items()] + ax.legend(handles=handles, title="enriched in", frameon=False, loc="upper left") + blocks = sorted({H.nodes[n]["block_id"] for n in H}) + ax.set_title(f"{chrom}:{start:,}-{end:,} | {H.number_of_nodes()} clusters in {len(blocks)} haploblocks, " + f"{H.number_of_edges()} co-occurrence edges\nnode size = carrier support, edge width = lift") + ax.axis("off"); fig.tight_layout(); fig.savefig(out, dpi=150); plt.close(fig) + return {"region": f"{chrom}:{start}-{end}", "nodes": H.number_of_nodes(), "edges": H.number_of_edges(), "blocks": len(blocks)} + + +def plot_degree_distribution(G: nx.Graph, out: Path) -> None: + import matplotlib; matplotlib.use("Agg") + import matplotlib.pyplot as plt + degrees = np.array([d for _, d in G.degree()]) + fig, ax = plt.subplots(figsize=(5.5, 4)) + ax.hist(degrees, bins=np.logspace(0, np.log10(degrees.max() + 1), 40), color="#1f6f8b") + ax.set_xscale("log"); ax.set_yscale("log") + ax.set_xlabel("degree (co-occurring clusters)"); ax.set_ylabel("clusters"); ax.set_title("Degree distribution") + fig.tight_layout(); fig.savefig(out, dpi=150); plt.close(fig) + + +def main() -> int: + here = Path(__file__).resolve().parent + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--chrom", default="chr22") + parser.add_argument("--kg-dir", type=Path, default=None) + parser.add_argument("--assoc", type=Path, default=None, help="cluster_phenotype_association.csv from cooccurrence_analysis.py") + parser.add_argument("--out-dir", type=Path, default=None) + parser.add_argument("--region", default=None, help="start-end in bp, e.g. 45035149-45534032 (default: densest island)") + parser.add_argument("--data-dir", type=Path, default=here / "data") + parser.add_argument("--no-graphml", action="store_true") + args = parser.parse_args() + kg_dir = args.kg_dir or here / "outputs" / "kg" / args.chrom + out_dir = args.out_dir or here / "outputs" / "graph" / args.chrom + assoc_path = args.assoc or here / "outputs" / "cooccurrence" / args.chrom / "cluster_phenotype_association.csv" + out_dir.mkdir(parents=True, exist_ok=True) + + kg = haplokg.load_kg(kg_dir) + assoc = pd.read_csv(assoc_path) if assoc_path.exists() else None + G = cluster_graph(kg, assoc) + B = block_graph(G) + stats = statistics(G, B) + + if not args.no_graphml: + nx.write_graphml(G, out_dir / f"cluster_cooccurrence_{args.chrom}.graphml") + nx.write_graphml(B, out_dir / f"block_cooccurrence_{args.chrom}.graphml") + + plot_edge_positions(G, args.chrom, out_dir / "edge_positions.png") + plot_degree_distribution(G, out_dir / "degree_distribution.png") + if args.region: + start, end = (int(x) for x in args.region.split("-")) + else: + start, end = default_region(args.data_dir, args.chrom) + stats["region_plot"] = plot_region(G, args.chrom, start, end, out_dir / f"region_{start}-{end}.png") + + (out_dir / "graph_stats.json").write_text(json.dumps(stats, indent=2)) + print(json.dumps(stats, indent=2)) + print(f"wrote {out_dir}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/genomics/graphrag_decoder.py b/genomics/graphrag_decoder.py new file mode 100644 index 0000000..ce2abab --- /dev/null +++ b/genomics/graphrag_decoder.py @@ -0,0 +1,208 @@ +#!/usr/bin/env python3 +"""GraphRAG decoder: knowledge-graph neighbourhood + GNN outputs -> LLM -> cited insight. + + python graphrag_decoder.py --individual HG00096 --dry-run # print the prompt, no LLM call + python graphrag_decoder.py --individual HG00096 # calls the NVIDIA NIM endpoint + python graphrag_decoder.py --individual HG00096 --run phenotype_both_raw + +The GNN encodes; the LLM decodes. Retrieval is deterministic and comes only from the graph: + * who the person is (ancestry / population / sex / site / age, from the graph; the true phenotype is withheld) + * the GNN's prediction for them (test_predictions.csv of the chosen run) and globally salient clusters + * the clusters they carry that are most ancestry-informative, with block, genes and proteins in that block + * their most extreme harmonised protein levels, with the gene and whether that gene sits in a block where + they carry a notable cluster (the genome<->proteome link) + * their nearest neighbours in the GNN embedding space +The LLM (OpenAI-compatible NIM API; key, model and endpoint come from config.py: NVIDIA_API_KEY, NIM_MODEL, NIM_URL in +genomics/.env or ~/.progenome.env, see .env.example) must answer as JSON and may only cite ids that appear in the +context; citations are validated before anything is written. +""" +from __future__ import annotations + +import argparse +import json +import os +import re +import time +from pathlib import Path + +import numpy as np +import pandas as pd + +import haplokg +import haplokg_proteins as hp + +import config # credentials and endpoints: .env / ~/.progenome.env / defaults (see .env.example) + +NIM_URL = config.settings.nim_url # any OpenAI-compatible chat-completions endpoint +DEFAULT_MODEL = config.settings.nim_model + +SYSTEM = """You are a careful genomics analyst writing for a clinical research team. +You receive a structured context extracted from a knowledge graph (haploblock clusters, blocks, genes, proteins, +a graph-neural-network prediction and similar individuals). Rules: +1. Use only the context. Do not invent genes, proteins, variants, diseases or numbers. +2. Every claim that refers to a graph entity must cite its id exactly as written (cluster ids like + chr22_17099658-17118145_cluster1, protein ids like P14174, gene symbols as given). +3. Ancestry in this graph is population structure, not a medical finding. The phenotype is a synthetic case/control + label used to test the pipeline; say so. +4. Answer as a single JSON object with keys: summary (2-3 sentences), ancestry_assessment, phenotype_assessment, + genome_proteome_links (list of {cluster_id, block_id, gene, protein_id, observation}), caveats (list of strings), + cited_ids (list of every id you cited). No markdown fences, no text outside the JSON.""" + + +def load_everything(here: Path, chrom: str, run: str): + kg_dir = here / "outputs" / "kg" / chrom + kg = haplokg.load_kg(kg_dir) + pt = hp.load_protein_tables(kg_dir) + assoc = pd.read_csv(here / "outputs" / "cooccurrence" / chrom / "cluster_phenotype_association.csv") + run_dirs = [here / "outputs" / "gnn_v2" / chrom / run, here / "outputs_brev" / "progenome-a100" / "gnn_v2" / chrom / run] + run_dir = next((d for d in run_dirs if (d / "metrics.json").exists()), None) + if run_dir is None: + raise SystemExit(f"no GNN run named {run!r} under outputs/gnn_v2 or outputs_brev/progenome-a100/gnn_v2") + preds = pd.read_csv(run_dir / "test_predictions.csv").set_index("individual_id") + emb = np.load(run_dir / "embedding_individual.npy") + saliency = pd.read_csv(run_dir / "saliency_top100.csv") if (run_dir / "saliency_top100.csv").exists() else None + return kg, pt, assoc, run_dir, preds, emb, saliency + + +def retrieve(individual: str, kg, pt, assoc, preds, emb, saliency, k_clusters=8, k_proteins=8, k_neighbours=5) -> dict: + ind = pt["individuals"].set_index("individual_id") + if individual not in ind.index: + raise SystemExit(f"{individual} is not in the graph") + row = ind.loc[individual] + i = int(row["individual_idx"]) + clusters, blocks, genes, prot = kg["clusters"], kg["blocks"], pt["genes"], pt["proteins"] + block_gene = pt["block_gene"].merge(genes[["gene_idx", "gene_symbol"]], on="gene_idx") + gene_prot = pt["gene_protein"].merge(prot[["protein_idx", "protein_id"]], on="protein_idx") + gene_to_prot = gene_prot.groupby("gene_idx")["protein_id"].apply(list).to_dict() + block_to_genes = block_gene.groupby("block_idx")["gene_symbol"].apply(list).to_dict() + block_to_gene_idx = block_gene.groupby("block_idx")["gene_idx"].apply(list).to_dict() + + carried = kg["carries"][i].indices + a = assoc.set_index("cluster_idx") + top = a.loc[carried].sort_values("ancestry_cramers_v", ascending=False).head(k_clusters) + cluster_ctx = [] + for cidx, r in top.iterrows(): + blk = int(r["block_idx"]); b = blocks.loc[blocks["block_idx"] == blk].iloc[0] + cluster_ctx.append({ + "cluster_id": r["cluster_id"], "block_id": b["block_id"], "carriers": int(r["support"]), + "enriched_in": r["ancestry_dominant"], "cramers_v": round(float(r["ancestry_cramers_v"]), 3), + "carrier_fraction_by_ancestry": {k.replace("carrier_frac_", ""): round(float(r[k]), 3) for k in top.columns if k.startswith("carrier_frac_")}, + "genes_in_block": block_to_genes.get(blk, []), + "proteins_in_block": sorted({p for g in block_to_gene_idx.get(blk, []) for p in gene_to_prot.get(g, [])}), + }) + notable_blocks = {c["block_id"] for c in cluster_ctx} + + m = pt["measured"] + mine = m[m["individual_idx"] == i].copy() + mine["abs_z"] = mine["z"].abs() + mine = mine.sort_values("abs_z", ascending=False).head(k_proteins) + prot_by_idx = prot.set_index("protein_idx") + protein_ctx = [] + for _, r in mine.iterrows(): + p = prot_by_idx.loc[int(r["protein_idx"])] + gblocks = block_gene.loc[block_gene["gene_idx"] == p["gene_idx"], "block_idx"].tolist() + gblock_ids = [blocks.loc[blocks["block_idx"] == b, "block_id"].iloc[0] for b in gblocks] + protein_ctx.append({"protein_id": p["protein_id"], "gene": p["gene_symbol"], "harmonised_z": round(float(r["z"]), 2), + "log2_intensity": round(float(r["log2_intensity"]), 2), "site": r["site"], + "encoded_in_blocks": gblock_ids, + "in_a_block_with_a_notable_cluster": any(b in notable_blocks for b in gblock_ids)}) + + labelled = ind[ind["ancestry_code"] >= 0] + others = labelled.index[labelled.index != individual] + d = np.linalg.norm(emb[labelled.loc[others, "individual_idx"].to_numpy()] - emb[i], axis=1) + nn = np.argsort(d)[:k_neighbours] + neighbours = [{"individual_id": others[j], "distance": round(float(d[j]), 3), "ancestry": labelled.loc[others[j], "ancestry"], + "population": labelled.loc[others[j], "population"], + "predicted_phenotype": preds.loc[others[j], "pred"] if others[j] in preds.index else "n/a"} for j in nn] + + prediction = ({"predicted": preds.loc[individual, "pred"]} if individual in preds.index + else {"predicted": "n/a (not in the held-out test split)"}) + global_saliency = saliency.head(10)[["cluster_id", "saliency"]].to_dict("records") if saliency is not None else [] + carried_salient = [s for s in global_saliency if s["cluster_id"] in set(clusters.loc[carried, "cluster_id"])] + + return { + "individual": {"id": individual, "ancestry": row["ancestry"], "population": row["population"], "sex": row["sex"], + "site": row.get("site", "n/a"), "age": int(row["age"]) if pd.notna(row.get("age", np.nan)) else "n/a", + "n_clusters_carried": int(len(carried))}, + "gnn_phenotype_prediction": prediction, + "globally_salient_clusters_this_person_carries": carried_salient, + "most_ancestry_informative_clusters_carried": cluster_ctx, + "most_extreme_protein_levels": protein_ctx, + "nearest_individuals_in_gnn_embedding": neighbours, + "note": "phenotype is a synthetic case/control label with a saved ground truth; ancestry labels are real 1000G panel data", + } + + +def call_nim(context: dict, model: str, temperature: float = 0.2, max_tokens: int = 4000, reasoning: str = "none") -> tuple[str, dict]: + import requests + key = config.require("NVIDIA_API_KEY") + # Nemotron 3 is a reasoning model: without reasoning_effort it thinks inline and can exhaust the token budget + # before the JSON; "none" answers directly, "low"/"high" keep the thinking in a separate reasoning field. + payload = {"model": model, "temperature": temperature, "max_tokens": max_tokens, "reasoning_effort": reasoning, + "messages": [{"role": "system", "content": SYSTEM}, + {"role": "user", "content": "CONTEXT (JSON):\n" + json.dumps(context, indent=1) + "\n\nWrite the JSON answer now."}]} + t0 = time.time() + r = requests.post(NIM_URL, headers={"Authorization": f"Bearer {key}", "Accept": "application/json"}, json=payload, timeout=300) + r.raise_for_status() + d = r.json() + msg = d["choices"][0]["message"] + return msg.get("content") or "", {"model": model, "reasoning_effort": reasoning, "seconds": round(time.time() - t0, 1), + "usage": d.get("usage"), "reasoning_chars": len(msg.get("reasoning_content") or "")} + + +def validate_citations(answer: dict, context: dict) -> dict: + blob = json.dumps(context) + ids = set(re.findall(r"chr\d+_\d+-\d+_cluster\d+", blob)) | set(re.findall(r"chr\d+_\d+-\d+", blob)) + ids |= {p["protein_id"] for p in context["most_extreme_protein_levels"]} + ids |= {p for c in context["most_ancestry_informative_clusters_carried"] for p in c["proteins_in_block"]} + ids |= {g for c in context["most_ancestry_informative_clusters_carried"] for g in c["genes_in_block"]} + ids |= {p["gene"] for p in context["most_extreme_protein_levels"]} + ids |= {n["individual_id"] for n in context["nearest_individuals_in_gnn_embedding"]} | {context["individual"]["id"]} + cited = set(answer.get("cited_ids", [])) + return {"cited": len(cited), "unknown_ids": sorted(cited - ids)} + + +def main() -> int: + here = Path(__file__).resolve().parent + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--chrom", default="chr22") + parser.add_argument("--individual", default="HG00096") + parser.add_argument("--run", default="phenotype_both_raw", help="GNN run folder under outputs/gnn_v2//") + parser.add_argument("--model", default=DEFAULT_MODEL) + parser.add_argument("--reasoning", default="none", choices=["none", "low", "medium", "high"], help="NIM reasoning_effort") + parser.add_argument("--dry-run", action="store_true", help="print the prompt context and stop") + args = parser.parse_args() + + kg, pt, assoc, run_dir, preds, emb, saliency = load_everything(here, args.chrom, args.run) + context = retrieve(args.individual, kg, pt, assoc, preds, emb, saliency) + out_dir = here / "outputs" / "graphrag" / args.chrom + out_dir.mkdir(parents=True, exist_ok=True) + (out_dir / f"{args.individual}_context.json").write_text(json.dumps(context, indent=2)) + if args.dry_run: + print(SYSTEM); print("\nCONTEXT:"); print(json.dumps(context, indent=1)); print(f"\nwrote {out_dir}/{args.individual}_context.json") + return 0 + + raw, meta = call_nim(context, args.model, reasoning=args.reasoning) + try: + answer = json.loads(raw[raw.find("{"): raw.rfind("}") + 1]) + except json.JSONDecodeError: + (out_dir / f"{args.individual}_raw.txt").write_text(raw) + raise SystemExit(f"model did not return JSON; raw reply saved to {out_dir}/{args.individual}_raw.txt") + check = validate_citations(answer, context) + result = {"individual": args.individual, "gnn_run": run_dir.name, "llm": meta, "citation_check": check, "answer": answer} + (out_dir / f"{args.individual}_insight.json").write_text(json.dumps(result, indent=2)) + print(f"{args.model} reasoning={meta['reasoning_effort']} {meta['seconds']}s tokens {meta['usage']}\n") + print(answer.get("summary", "")) + print("\nancestry:", answer.get("ancestry_assessment", "")) + print("phenotype:", answer.get("phenotype_assessment", "")) + for link in answer.get("genome_proteome_links", []): + print(" link:", link) + for c in answer.get("caveats", []): + print(" caveat:", c) + print(f"\ncitations: {check['cited']} cited, unknown ids: {check['unknown_ids'] or 'none'}") + print(f"wrote {out_dir}/{args.individual}_insight.json") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/genomics/haplokg.py b/genomics/haplokg.py new file mode 100644 index 0000000..62d561d --- /dev/null +++ b/genomics/haplokg.py @@ -0,0 +1,337 @@ +"""Knowledge graph from the published 1000G HaploGraph (data.haploblocks.org). + +Inputs (see fetch_data.sh): + nodes.csv.gz one row per haploblock *cluster*, one 0/1 column per + 1000G individual: does this individual carry the cluster + edges_lift_above_threshold.csv.gz cluster-cluster co-occurrence edges (weight, lift >= 5) + block_stats.tsv per-haploblock statistics (length, n_clusters, entropy, ...) + phenotypes_real.csv long-format ancestry / population / sex labels (1000G panel) + +The graph has three node types + individual --carries--> cluster --in_block--> block --next_block--> block + cluster <--co_occurs--> cluster +and the labels live on the individual nodes. Only the *cluster* rows are +filtered (by carrier support); every individual and every block is kept. +""" +from __future__ import annotations + +import gzip +import json +import re +from pathlib import Path + +import numpy as np +import pandas as pd +from scipy import sparse + +BLOCK_RE = re.compile(r"^(?Pchr[0-9XY]+)_(?P\d+)-(?P\d+)$") +CLUSTER_RE = re.compile(r"^(?Pchr[0-9XY]+_\d+-\d+)_cluster(?P\d+)$") +PHENOTYPES = ("ancestry", "population", "sex") + + +# ----------------------------------------------------------------------------- ids +def parse_block_id(block_id: str) -> tuple[str, int, int]: + m = BLOCK_RE.match(block_id) + if not m: + raise ValueError(f"not a block id: {block_id!r}") + return m["chrom"], int(m["start"]), int(m["end"]) + + +def parse_cluster_id(cluster_id: str) -> tuple[str, int]: + m = CLUSTER_RE.match(cluster_id) + if not m: + raise ValueError(f"not a cluster id: {cluster_id!r}") + return m["block"], int(m["num"]) + + +# ------------------------------------------------------------------------- loading +def read_node_matrix(path: Path, chunksize: int = 8192): + """Stream nodes.csv.gz into a CSR matrix of shape (clusters, individuals), int8. + + The file is ~1.3 GB uncompressed for chr22; reading it as int8 in chunks and + converting each chunk to sparse keeps peak memory well under 1 GB. + """ + with gzip.open(path, "rt") as fh: + header = fh.readline().rstrip("\n").split(",") + individuals = header[2:] + dtypes = {name: np.int8 for name in individuals} + dtypes[header[0]] = str + dtypes[header[1]] = str + + cluster_ids: list[str] = [] + block_ids: list[str] = [] + parts = [] + for chunk in pd.read_csv(path, dtype=dtypes, chunksize=chunksize, engine="c"): + cluster_ids.extend(chunk.iloc[:, 0].tolist()) + block_ids.extend(chunk.iloc[:, 1].tolist()) + parts.append(sparse.csr_matrix(chunk.iloc[:, 2:].to_numpy(dtype=np.int8))) + if parts: + matrix = sparse.vstack(parts, format="csr") + else: + matrix = sparse.csr_matrix((0, len(individuals)), dtype=np.int8) + return cluster_ids, block_ids, individuals, matrix + + +def load_phenotypes(path: Path) -> pd.DataFrame: + """Long (individual_id, phenotype, value) -> wide, indexed by individual_id.""" + long = pd.read_csv(path, dtype=str) + wide = long.pivot_table(index="individual_id", columns="phenotype", values="value", aggfunc="first") + for name in PHENOTYPES: + if name not in wide.columns: + wide[name] = np.nan + return wide[list(PHENOTYPES)] + + +def load_block_stats(path: Path, chrom: str) -> pd.DataFrame: + stats = pd.read_csv(path, sep="\t") + stats = stats[stats["chr"] == chrom].copy() + stats = stats.sort_values("start").reset_index(drop=True) + stats["block_idx"] = np.arange(len(stats)) + stats["singleton_rate"] = stats["singleton_count"] / stats["n_clusters"].clip(lower=1) + return stats.rename(columns={"block": "block_id"}) + + +def load_edges(path: Path) -> pd.DataFrame: + return pd.read_csv(path) + + +# ------------------------------------------------------------------------ building +def encode_labels(values: pd.Series) -> tuple[np.ndarray, list[str]]: + """Sorted class codes; missing -> -1.""" + classes = sorted(v for v in values.dropna().unique()) + lookup = {c: i for i, c in enumerate(classes)} + codes = np.array([lookup.get(v, -1) if isinstance(v, str) else -1 for v in values], dtype=np.int64) + return codes, classes + + +def build_tables( + cluster_ids: list[str], + block_ids: list[str], + individual_ids: list[str], + matrix: sparse.csr_matrix, + phenotypes: pd.DataFrame, + block_stats: pd.DataFrame, + edges: pd.DataFrame, + min_support: int = 25, + symmetric: bool = True, +) -> dict: + """Filter clusters by carrier support and assemble every table of the graph. + + symmetric=True mirrors the HaploGraph edge filter: a cluster is kept only if + min(carriers, non-carriers) >= min_support, so near-universal clusters go too. + """ + n_ind = len(individual_ids) + support = np.asarray(matrix.sum(axis=1)).ravel().astype(np.int64) + keep = support >= min_support + if symmetric: + keep &= (n_ind - support) >= min_support + + # blocks: everything in block_stats for this chromosome, plus any block that + # appears in nodes.csv but is missing from block_stats (should not happen) + blocks = block_stats.copy() + known = set(blocks["block_id"]) + extra = sorted(set(block_ids) - known, key=lambda b: parse_block_id(b)[1]) + if extra: + rows = [] + for b in extra: + chrom, start, end = parse_block_id(b) + rows.append({"chr": chrom, "block_id": b, "start": start, "end": end, "block_length": end - start}) + blocks = pd.concat([blocks, pd.DataFrame(rows)], ignore_index=True) + blocks = blocks.sort_values("start").reset_index(drop=True) + blocks["block_idx"] = np.arange(len(blocks)) + block_index = dict(zip(blocks["block_id"], blocks["block_idx"])) + + clusters = pd.DataFrame({"cluster_id": cluster_ids, "block_id": block_ids, "support": support}) + clusters["row"] = np.arange(len(clusters)) + clusters = clusters[keep].reset_index(drop=True) + clusters["cluster_idx"] = np.arange(len(clusters)) + clusters["block_idx"] = clusters["block_id"].map(block_index).astype(np.int64) + clusters["cluster_num"] = [parse_cluster_id(c)[1] for c in clusters["cluster_id"]] + clusters["support_frac"] = clusters["support"] / n_ind + + carries = matrix[clusters["row"].to_numpy()].T.tocsr() # individuals x kept clusters + carries.data = np.ones_like(carries.data, dtype=np.int8) + + individuals = pd.DataFrame({"individual_id": individual_ids}) + individuals["individual_idx"] = np.arange(n_ind) + individuals = individuals.join(phenotypes, on="individual_id") + label_maps = {} + for name in PHENOTYPES: + codes, classes = encode_labels(individuals[name]) + individuals[f"{name}_code"] = codes + label_maps[name] = classes + + cluster_index = dict(zip(clusters["cluster_id"], clusters["cluster_idx"])) + co = edges.copy() + co["src"] = co["source"].map(cluster_index) + co["dst"] = co["target"].map(cluster_index) + dropped = int(co["src"].isna().sum() + co["dst"].isna().sum() - (co["src"].isna() & co["dst"].isna()).sum()) + co = co.dropna(subset=["src", "dst"]).astype({"src": np.int64, "dst": np.int64}) + co_occurs = co[["src", "dst", "weight", "lift"]].reset_index(drop=True) + + order = blocks.sort_values("start")["block_idx"].to_numpy() + next_block = pd.DataFrame({"src": order[:-1], "dst": order[1:]}) + + return { + "individuals": individuals, + "clusters": clusters.drop(columns=["row"]), + "blocks": blocks, + "carries": carries, + "co_occurs": co_occurs, + "next_block": next_block, + "label_maps": label_maps, + "n_edges_dropped": dropped, + "min_support": min_support, + "symmetric": symmetric, + } + + +# ---------------------------------------------------------------------- features +def _zscore(frame: pd.DataFrame) -> np.ndarray: + values = frame.to_numpy(dtype=np.float64) + mean = np.nanmean(values, axis=0) + std = np.nanstd(values, axis=0) + std[std == 0] = 1.0 + z = (values - mean) / std + return np.nan_to_num(z, nan=0.0).astype(np.float32) + + +BLOCK_FEATURES = ["log_length", "shannon_entropy", "dominance", "log_n_clusters", "singleton_rate"] + + +def block_feature_frame(blocks: pd.DataFrame) -> pd.DataFrame: + out = pd.DataFrame(index=blocks.index) + out["log_length"] = np.log10(blocks["block_length"].astype(float).clip(lower=1)) + out["shannon_entropy"] = blocks.get("shannon_entropy", np.nan) + out["dominance"] = blocks.get("dominance", np.nan) + out["log_n_clusters"] = np.log1p(blocks.get("n_clusters", np.nan).astype(float)) + out["singleton_rate"] = blocks.get("singleton_rate", np.nan) + return out + + +def cluster_feature_frame(clusters: pd.DataFrame, blocks: pd.DataFrame) -> pd.DataFrame: + bf = block_feature_frame(blocks).set_index(blocks["block_idx"]) + out = pd.DataFrame(index=clusters.index) + out["log_support"] = np.log1p(clusters["support"].astype(float)) + out["support_frac"] = clusters["support_frac"] + joined = bf.reindex(clusters["block_idx"].to_numpy()) + for col in BLOCK_FEATURES: + out[f"block_{col}"] = joined[col].to_numpy() + return out + + +def to_hetero_data(tables: dict): + """Assemble a torch_geometric HeteroData (imported lazily so the tables work without torch).""" + import torch + from torch_geometric.data import HeteroData + + ind, cl, bl = tables["individuals"], tables["clusters"], tables["blocks"] + data = HeteroData() + + data["individual"].num_nodes = len(ind) + data["individual"].individual_id = list(ind["individual_id"]) + for name in PHENOTYPES: + data["individual"][f"y_{name}"] = torch.tensor(ind[f"{name}_code"].to_numpy(), dtype=torch.long) + + data["cluster"].x = torch.from_numpy(_zscore(cluster_feature_frame(cl, bl))) + data["cluster"].cluster_id = list(cl["cluster_id"]) + data["cluster"].support = torch.tensor(cl["support"].to_numpy(), dtype=torch.long) + data["block"].x = torch.from_numpy(_zscore(block_feature_frame(bl))) + data["block"].block_id = list(bl["block_id"]) + + coo = tables["carries"].tocoo() + carries = torch.from_numpy(np.vstack([coo.row, coo.col]).astype(np.int64)) + data["individual", "carries", "cluster"].edge_index = carries + data["cluster", "rev_carries", "individual"].edge_index = carries.flip(0) + + in_block = torch.from_numpy(np.vstack([cl["cluster_idx"].to_numpy(), cl["block_idx"].to_numpy()]).astype(np.int64)) + data["cluster", "in_block", "block"].edge_index = in_block + data["block", "rev_in_block", "cluster"].edge_index = in_block.flip(0) + + co = tables["co_occurs"] + src = torch.tensor(co["src"].to_numpy(), dtype=torch.long) + dst = torch.tensor(co["dst"].to_numpy(), dtype=torch.long) + attr = torch.tensor(co[["weight", "lift"]].to_numpy(dtype=np.float32)) + data["cluster", "co_occurs", "cluster"].edge_index = torch.stack([torch.cat([src, dst]), torch.cat([dst, src])]) + data["cluster", "co_occurs", "cluster"].edge_attr = torch.cat([attr, attr]) + + nb = tables["next_block"] + a = torch.tensor(nb["src"].to_numpy(), dtype=torch.long) + b = torch.tensor(nb["dst"].to_numpy(), dtype=torch.long) + data["block", "next_block", "block"].edge_index = torch.stack([torch.cat([a, b]), torch.cat([b, a])]) + + data.label_maps = tables["label_maps"] + return data + + +# -------------------------------------------------------------------------- saving +def save_tables(tables: dict, out_dir: Path) -> dict: + out_dir.mkdir(parents=True, exist_ok=True) + tables["individuals"].to_csv(out_dir / "individuals.csv", index=False) + tables["clusters"].to_csv(out_dir / "clusters.csv", index=False) + tables["blocks"].to_csv(out_dir / "blocks.csv", index=False) + tables["co_occurs"].to_csv(out_dir / "co_occurs.csv", index=False) + tables["next_block"].to_csv(out_dir / "next_block.csv", index=False) + sparse.save_npz(out_dir / "carries.npz", tables["carries"]) + (out_dir / "label_maps.json").write_text(json.dumps(tables["label_maps"], indent=2)) + ind = tables["individuals"] + summary = { + "n_individuals": int(len(ind)), + "n_individuals_labelled": int((ind["ancestry_code"] >= 0).sum()), + "n_clusters_kept": int(len(tables["clusters"])), + "n_blocks": int(len(tables["blocks"])), + "n_carries_edges": int(tables["carries"].nnz), + "n_co_occurs_edges": int(len(tables["co_occurs"])), + "n_co_occurs_dropped_by_filter": int(tables["n_edges_dropped"]), + "min_support": tables["min_support"], + "symmetric_filter": tables["symmetric"], + "clusters_per_individual_mean": float(tables["carries"].sum(axis=1).mean()), + } + (out_dir / "summary.json").write_text(json.dumps(summary, indent=2)) + return summary + + +# ------------------------------------------------------------------- reloading +def load_kg(kg_dir: Path) -> dict: + """Reload the tables written by save_tables (no torch needed).""" + kg_dir = Path(kg_dir) + return { + "individuals": pd.read_csv(kg_dir / "individuals.csv"), + "clusters": pd.read_csv(kg_dir / "clusters.csv"), + "blocks": pd.read_csv(kg_dir / "blocks.csv"), + "co_occurs": pd.read_csv(kg_dir / "co_occurs.csv"), + "next_block": pd.read_csv(kg_dir / "next_block.csv"), + "carries": sparse.load_npz(kg_dir / "carries.npz").tocsr(), + "label_maps": json.loads((kg_dir / "label_maps.json").read_text()), + } + + +def stratified_split(codes: np.ndarray, seed: int = 42, val_frac: float = 0.15, test_frac: float = 0.15) -> np.ndarray: + """Return an array of 'train' / 'val' / 'test' / 'unlabelled' per individual. + + Stratified on `codes` (use ancestry); individuals with code -1 are never used + for training or evaluation. Deterministic for a given seed so the baseline + and the GNN score the very same held-out people. + """ + from sklearn.model_selection import train_test_split + + codes = np.asarray(codes) + split = np.full(len(codes), "unlabelled", dtype=object) + labelled = np.flatnonzero(codes >= 0) + hold = val_frac + test_frac + train_idx, hold_idx = train_test_split(labelled, test_size=hold, random_state=seed, stratify=codes[labelled]) + val_idx, test_idx = train_test_split(hold_idx, test_size=test_frac / hold, random_state=seed, stratify=codes[hold_idx]) + split[train_idx], split[val_idx], split[test_idx] = "train", "val", "test" + return split + + +def load_or_make_split(kg: dict, split_path: Path, seed: int = 42) -> np.ndarray: + split_path = Path(split_path) + if split_path.exists(): + frame = pd.read_csv(split_path) + assert (frame["individual_id"].to_numpy() == kg["individuals"]["individual_id"].to_numpy()).all() + return frame["split"].to_numpy() + split = stratified_split(kg["individuals"]["ancestry_code"].to_numpy(), seed=seed) + split_path.parent.mkdir(parents=True, exist_ok=True) + pd.DataFrame({"individual_id": kg["individuals"]["individual_id"], "split": split}).to_csv(split_path, index=False) + return split diff --git a/genomics/haplokg_proteins.py b/genomics/haplokg_proteins.py new file mode 100644 index 0000000..c2af4b0 --- /dev/null +++ b/genomics/haplokg_proteins.py @@ -0,0 +1,166 @@ +"""Protein layer of the knowledge graph (schema v2). + +Adds to the genome graph built by haplokg: + Block -OVERLAPS-> Gene -ENCODES-> Protein (proteomics/uniprot_chr22.bed + gene symbols) + Individual -MEASURED{log2, z}-> Protein (a proteomics matrix keyed by 1000G IDs) +and per-individual phenotype / site / age columns from the proteomics sample metadata. + +The harmoniser: MEASURED carries the raw log2 intensity and a z-score computed +*within each site and protein* (median / MAD), which removes per-site batch offsets +before anything crosses sites - the simplest version of the whiteboard's +"NORM: HARMONIZER" box. Detection-limit missingness stays missing (no edge). +""" +from __future__ import annotations + +import json +from pathlib import Path + +import numpy as np +import pandas as pd +from scipy import sparse + + +def load_protein_bed(bed_path: Path, symbols_path: Path | None = None) -> pd.DataFrame: + """One row per protein (isoforms collapsed to the base accession) with its genomic span and gene symbol.""" + bed = pd.read_csv(bed_path, sep="\t", header=None, usecols=[0, 1, 2, 3, 5], names=["chrom", "start", "end", "uniprot", "strand"]) + bed["protein_id"] = bed["uniprot"].str.split("-").str[0] + prot = bed.groupby(["chrom", "protein_id"], as_index=False).agg( + start=("start", "min"), end=("end", "max"), strand=("strand", "first"), n_isoforms=("uniprot", "nunique")) + symbols = {} + if symbols_path and Path(symbols_path).exists(): + symbols = pd.read_csv(symbols_path).set_index("protein_id")["gene_symbol"].to_dict() + prot["gene_symbol"] = [symbols.get(p, p) for p in prot["protein_id"]] + return prot + + +def harmonise(long: pd.DataFrame) -> pd.DataFrame: + """Robust z-score per (site, protein): (x - median) / (1.4826 * MAD).""" + grp = long.groupby(["site", "protein_id"])["log2_intensity"] + med = grp.transform("median") + mad = (long["log2_intensity"] - med).abs().groupby([long["site"], long["protein_id"]]).transform("median") * 1.4826 + long = long.copy() + long["z"] = ((long["log2_intensity"] - med) / mad.replace(0, np.nan)).fillna(0.0) + return long + + +def build_protein_tables(kg: dict, bed_path: Path, measured_path: Path, metadata_path: Path, + symbols_path: Path | None = None) -> dict: + ind, blocks = kg["individuals"], kg["blocks"] + chrom = blocks["chr"].iloc[0] + prot = load_protein_bed(bed_path, symbols_path) + prot = prot[prot["chrom"] == chrom].reset_index(drop=True) + + genes = prot.groupby("gene_symbol", as_index=False).agg(start=("start", "min"), end=("end", "max"), n_proteins=("protein_id", "nunique")) + genes["gene_idx"] = np.arange(len(genes)) + gene_index = dict(zip(genes["gene_symbol"], genes["gene_idx"])) + prot["protein_idx"] = np.arange(len(prot)) + prot["gene_idx"] = prot["gene_symbol"].map(gene_index) + + # Block -OVERLAPS-> Gene by coordinate intersection (BED is 0-based half-open; blocks are 1-based inclusive) + b = blocks[["block_idx", "start", "end"]].to_numpy() + rows = [] + for g in genes.itertuples(index=False): + hit = b[(b[:, 1] <= g.end) & (b[:, 2] >= g.start + 1)] + for blk_idx, bs, be in hit: + rows.append({"block_idx": int(blk_idx), "gene_idx": int(g.gene_idx), + "overlap_bp": int(min(be, g.end) - max(bs, g.start + 1) + 1)}) + block_gene = pd.DataFrame(rows, columns=["block_idx", "gene_idx", "overlap_bp"]) + gene_protein = prot[["gene_idx", "protein_idx"]].copy() + + # Individual -MEASURED-> Protein + long = pd.read_csv(measured_path) + long = long[long["protein_id"].isin(prot["protein_id"])] + long = harmonise(long) + ind_index = dict(zip(ind["individual_id"], ind["individual_idx"])) + prot_index = dict(zip(prot["protein_id"], prot["protein_idx"])) + long["individual_idx"] = long["individual_id"].map(ind_index) + long["protein_idx"] = long["protein_id"].map(prot_index) + unmatched = int(long["individual_idx"].isna().sum()) + long = long.dropna(subset=["individual_idx"]).astype({"individual_idx": np.int64, "protein_idx": np.int64}) + measured = long[["individual_idx", "protein_idx", "site", "log2_intensity", "z"]].reset_index(drop=True) + + # dense harmonised abundance matrix (individuals x proteins) + observed mask, for node features + abundance = sparse.coo_matrix((measured["z"].to_numpy(), (measured["individual_idx"], measured["protein_idx"])), + shape=(len(ind), len(prot))).tocsr() + observed = sparse.coo_matrix((np.ones(len(measured), dtype=np.int8), (measured["individual_idx"], measured["protein_idx"])), + shape=(len(ind), len(prot))).tocsr() + + # phenotype / site / age onto individuals + meta = pd.read_csv(metadata_path).rename(columns={"sample_id": "individual_id"}) + ind2 = ind.merge(meta[["individual_id", "site", "age", "phenotype"]], on="individual_id", how="left") + ind2["phenotype_code"] = ind2["phenotype"].fillna(-1).astype(np.int64) + sites = sorted(meta["site"].dropna().unique()) + ind2["site_code"] = ind2["site"].map({s: i for i, s in enumerate(sites)}).fillna(-1).astype(np.int64) + label_maps = dict(kg["label_maps"]) + label_maps["phenotype"] = ["control", "case"] + label_maps["site"] = sites + + return {"individuals": ind2, "genes": genes, "proteins": prot, "block_gene": block_gene, "gene_protein": gene_protein, + "measured": measured, "abundance": abundance, "observed": observed, "label_maps": label_maps, + "n_unmatched_measurements": unmatched} + + +def save_protein_tables(t: dict, out_dir: Path) -> dict: + out_dir = Path(out_dir) + t["individuals"].to_csv(out_dir / "individuals_v2.csv", index=False) + t["genes"].to_csv(out_dir / "genes.csv", index=False) + t["proteins"].to_csv(out_dir / "proteins.csv", index=False) + t["block_gene"].to_csv(out_dir / "block_gene.csv", index=False) + t["gene_protein"].to_csv(out_dir / "gene_protein.csv", index=False) + t["measured"].to_csv(out_dir / "measured.csv", index=False) + sparse.save_npz(out_dir / "abundance_z.npz", t["abundance"]) + sparse.save_npz(out_dir / "abundance_observed.npz", t["observed"]) + (out_dir / "label_maps_v2.json").write_text(json.dumps(t["label_maps"], indent=2)) + summary = { + "n_genes": int(len(t["genes"])), "n_proteins": int(len(t["proteins"])), + "n_block_gene_edges": int(len(t["block_gene"])), "n_measured_edges": int(len(t["measured"])), + "n_individuals_with_proteomics": int((t["individuals"]["site_code"] >= 0).sum()), + "n_unmatched_measurements": int(t["n_unmatched_measurements"]), + "genes_without_block": int((~t["genes"]["gene_idx"].isin(t["block_gene"]["gene_idx"])).sum()), + } + (out_dir / "summary_v2.json").write_text(json.dumps(summary, indent=2)) + return summary + + +def load_protein_tables(kg_dir: Path) -> dict: + kg_dir = Path(kg_dir) + return { + "individuals": pd.read_csv(kg_dir / "individuals_v2.csv"), + "genes": pd.read_csv(kg_dir / "genes.csv"), "proteins": pd.read_csv(kg_dir / "proteins.csv"), + "block_gene": pd.read_csv(kg_dir / "block_gene.csv"), "gene_protein": pd.read_csv(kg_dir / "gene_protein.csv"), + "measured": pd.read_csv(kg_dir / "measured.csv"), + "abundance": sparse.load_npz(kg_dir / "abundance_z.npz").tocsr(), + "observed": sparse.load_npz(kg_dir / "abundance_observed.npz").tocsr(), + "label_maps": json.loads((kg_dir / "label_maps_v2.json").read_text()), + } + + +def extend_hetero_data(data, t: dict): + """Add gene/protein nodes and their edges to an existing HeteroData (in place) and return it.""" + import torch + + genes, prot = t["genes"], t["proteins"] + data["individual"].y_phenotype = torch.tensor(t["individuals"]["phenotype_code"].to_numpy(), dtype=torch.long) + data["individual"].y_site = torch.tensor(t["individuals"]["site_code"].to_numpy(), dtype=torch.long) + gx = np.stack([np.log10(genes["end"] - genes["start"] + 1), np.log1p(genes["n_proteins"])], axis=1).astype(np.float32) + px = np.stack([np.log10(prot["end"] - prot["start"] + 1), np.log1p(prot["n_isoforms"])], axis=1).astype(np.float32) + data["gene"].x = torch.from_numpy((gx - gx.mean(0)) / (gx.std(0) + 1e-6)) + data["gene"].gene_symbol = list(genes["gene_symbol"]) + data["protein"].x = torch.from_numpy((px - px.mean(0)) / (px.std(0) + 1e-6)) + data["protein"].protein_id = list(prot["protein_id"]) + + bg = torch.tensor(t["block_gene"][["block_idx", "gene_idx"]].to_numpy().T, dtype=torch.long) + data["block", "overlaps", "gene"].edge_index = bg + data["gene", "rev_overlaps", "block"].edge_index = bg.flip(0) + gp = torch.tensor(t["gene_protein"][["gene_idx", "protein_idx"]].to_numpy().T, dtype=torch.long) + data["gene", "encodes", "protein"].edge_index = gp + data["protein", "rev_encodes", "gene"].edge_index = gp.flip(0) + m = t["measured"] + mi = torch.tensor(m[["individual_idx", "protein_idx"]].to_numpy().T, dtype=torch.long) + attr = torch.tensor(m[["z", "log2_intensity"]].to_numpy(dtype=np.float32)) + data["individual", "measured", "protein"].edge_index = mi + data["individual", "measured", "protein"].edge_attr = attr + data["protein", "rev_measured", "individual"].edge_index = mi.flip(0) + data["protein", "rev_measured", "individual"].edge_attr = attr + data.label_maps = t["label_maps"] + return data diff --git a/genomics/infer.py b/genomics/infer.py new file mode 100644 index 0000000..3fc5f71 --- /dev/null +++ b/genomics/infer.py @@ -0,0 +1,152 @@ +#!/usr/bin/env python3 +"""Inference with a trained GNN: predictions + embeddings for every individual, timed. + + python infer.py --run ancestry_svd # eager PyTorch (CPU or GPU) + python infer.py --run ancestry_svd --compile tensorrt # Torch-TensorRT via torch.compile (GPU) + python infer.py --run ancestry_svd --compile inductor # torch.compile default backend + +TensorRT note: the GNN's message passing is scatter/gather over 2.4M edges, which +TensorRT does not compile natively. torch.compile(backend="torch_tensorrt") +partitions the graph, runs the dense parts (all Linear layers, norms, the head) +as TensorRT engines and falls back to PyTorch for the rest, so the speed-up is +real but bounded - the benchmark prints both numbers instead of assuming one. +""" +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import numpy as np +import pandas as pd +import torch + +import haplokg +from train_gnn import CO, HeteroGNN, pick_device, svd_embeddings + + +def build_inputs(kg: dict, data, report: dict, device: torch.device): + """Recreate exactly the inputs train_gnn.py used for this run.""" + init, k = report["init"], report["embed_dim"] + x_dict = {"cluster": data["cluster"].x, "block": data["block"].x} + if init in ("svd", "node2vec", "raw"): + emb_dir = Path(__file__).resolve().parent / "outputs" / "embeddings" / report["chrom"] + cl_path = emb_dir / f"{'node2vec' if init == 'node2vec' else 'svd'}{k}_cluster.npy" + ind_path = emb_dir / f"{'node2vec' if init == 'node2vec' else 'svd'}{k}_individual.npy" + if cl_path.exists() and ind_path.exists(): + cl_init, ind_init = np.load(cl_path), np.load(ind_path) + else: + ind_init, cl_init, _ = svd_embeddings(kg["carries"], k, 42) + x_dict["cluster"] = torch.cat([x_dict["cluster"], torch.from_numpy(cl_init)], dim=1) + x_dict["individual"] = torch.from_numpy(kg["carries"].toarray().astype(np.float32) if init == "raw" else ind_init) + lift = data[CO].edge_attr[:, 1] + edge_weight = torch.log(lift) / torch.log(lift).max() + return ({t: x.to(device) for t, x in x_dict.items()}, + {k_: v.to(device) for k_, v in data.edge_index_dict.items()}, edge_weight.to(device)) + + +class Wrapped(torch.nn.Module): + """Fixed-argument wrapper so torch.compile sees plain tensors, not dicts.""" + def __init__(self, model, edge_index_dict, edge_weight): + super().__init__() + self.model, self.edge_index_dict, self.edge_weight = model, edge_index_dict, edge_weight + + def forward(self, x_individual, x_cluster, x_block): + logits, h = self.model({"individual": x_individual, "cluster": x_cluster, "block": x_block}, + self.edge_index_dict, self.edge_weight) + return logits, h["individual"] + + +def bench(fn, *args, warmup: int = 3, iters: int = 10, device=None) -> float: + for _ in range(warmup): + fn(*args) + if device is not None and device.type == "cuda": + torch.cuda.synchronize() + t0 = time.perf_counter() + for _ in range(iters): + fn(*args) + if device is not None and device.type == "cuda": + torch.cuda.synchronize() + return (time.perf_counter() - t0) / iters + + +def main() -> int: + here = Path(__file__).resolve().parent + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--chrom", default="chr22") + parser.add_argument("--run", default="ancestry_svd", help="folder name under outputs/gnn//") + parser.add_argument("--compile", default="none", choices=["none", "inductor", "tensorrt"]) + parser.add_argument("--precision", default="fp32", choices=["fp32", "fp16"]) + parser.add_argument("--device", default="auto") + parser.add_argument("--iters", type=int, default=10) + args = parser.parse_args() + device = pick_device(args.device) + run_dir = here / "outputs" / "gnn" / args.chrom / args.run + report = json.loads((run_dir / "metrics.json").read_text()) + kg_dir = here / "outputs" / "kg" / args.chrom + kg = haplokg.load_kg(kg_dir) + data = torch.load(kg_dir / "hetero.pt", weights_only=False) + x_dict, edge_index_dict, edge_weight = build_inputs(kg, data, report, device) + + model = HeteroGNN({t: x.shape[1] for t, x in x_dict.items()}, report["hidden"], len(report["classes"]), + report["layers"], 0.0, learned_individual=data["individual"].num_nodes if report["init"] == "learned" else None, + aggr=report.get("aggr", "mean")).to(device) + model.load_state_dict(torch.load(run_dir / "model.pt", map_location=device)) + model.eval() + wrapped = Wrapped(model, edge_index_dict, edge_weight).eval() + inputs = (x_dict["individual"], x_dict["cluster"], x_dict["block"]) + + with torch.no_grad(): + eager_s = bench(wrapped, *inputs, iters=args.iters, device=device) + logits, emb = wrapped(*inputs) + result = {"run": args.run, "device": str(device), "eager_ms_per_full_graph": round(eager_s * 1000, 2), + "compile": args.compile, "precision": args.precision} + + if args.compile != "none": + if args.compile == "tensorrt": + try: + import torch_tensorrt # noqa: F401 + kwargs = {"backend": "torch_tensorrt", "options": { + "enabled_precisions": {torch.float16 if args.precision == "fp16" else torch.float32}, + "min_block_size": 1, "truncate_double": True}} + except ImportError: + print("torch_tensorrt is not installed -> falling back to the inductor backend") + kwargs = {"backend": "inductor"} + result["compile"] = "inductor (tensorrt unavailable)" + else: + kwargs = {"backend": "inductor"} + compiled = torch.compile(wrapped, **kwargs) + try: + with torch.no_grad(): + compiled_s = bench(compiled, *inputs, iters=args.iters, device=device) + c_logits, c_emb = compiled(*inputs) + result["compiled_ms_per_full_graph"] = round(compiled_s * 1000, 2) + result["speedup"] = round(eager_s / compiled_s, 2) + result["max_abs_logit_diff_vs_eager"] = float((c_logits.float() - logits.float()).abs().max()) + logits, emb = c_logits, c_emb + except Exception as exc: # compilation problems must be visible, not silent + result["compile_error"] = f"{type(exc).__name__}: {str(exc)[:300]}" + print("compile failed, results below are eager:", result["compile_error"]) + + pred = logits.float().argmax(1).cpu().numpy() + prob = torch.softmax(logits.float(), dim=1).max(1).values.cpu().numpy() + classes = report["classes"] + ind = kg["individuals"] + out = pd.DataFrame({"individual_id": ind["individual_id"], "true": ind[report["target"]].fillna("unlabelled"), + "pred": [classes[i] for i in pred], "confidence": np.round(prob, 4)}) + out_dir = run_dir / "inference" + out_dir.mkdir(exist_ok=True) + out.to_csv(out_dir / "predictions_all_individuals.csv", index=False) + np.save(out_dir / "embedding_individual.npy", emb.float().cpu().numpy()) + labelled = out["true"] != "unlabelled" + result["accuracy_all_labelled_individuals"] = float((out.loc[labelled, "true"] == out.loc[labelled, "pred"]).mean()) + result["n_individuals"] = int(len(out)) + (out_dir / f"benchmark_{args.compile}_{args.precision}.json").write_text(json.dumps(result, indent=2)) + print(json.dumps(result, indent=2)) + print(f"wrote {out_dir}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/genomics/load_env.sh b/genomics/load_env.sh new file mode 100644 index 0000000..58ffb9e --- /dev/null +++ b/genomics/load_env.sh @@ -0,0 +1,18 @@ +#!/usr/bin/env bash +# Source this from the shell scripts: . "$(dirname "$0")/load_env.sh" +# Same rules as config.py: exported variables win, then genomics/.env, then ~/.progenome.env. +# Lines are KEY=VALUE (optional 'export '), quotes stripped, full-line comments ignored. +_progenome_load() { + [ -f "$1" ] || return 0 + while IFS= read -r line || [ -n "$line" ]; do + line="${line#"${line%%[![:space:]]*}"}" + case "$line" in ''|'#'*) continue ;; esac + case "$line" in *=*) ;; *) continue ;; esac + line="${line#export }" + key="${line%%=*}"; val="${line#*=}" + val="${val%\"}"; val="${val#\"}"; val="${val%\'}"; val="${val#\'}" + if [ -z "${!key:-}" ]; then export "$key=$val"; fi + done < "$1" +} +_progenome_load "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/.env" +_progenome_load "$HOME/.progenome.env" diff --git a/genomics/neo4j_load.py b/genomics/neo4j_load.py new file mode 100644 index 0000000..1702da9 --- /dev/null +++ b/genomics/neo4j_load.py @@ -0,0 +1,142 @@ +#!/usr/bin/env python3 +"""Load the haploblock knowledge graph into Neo4j so it can be browsed at localhost:7474. + +Nodes (:Individual {id, ancestry, population, sex, split}) + (:Cluster {id, idx, block_id, support, support_frac, ancestry_dominant, ancestry_v}) + (:Block {id, idx, chrom, start, end, length, n_clusters, entropy, dominance}) +Relations (Individual)-[:CARRIES]->(Cluster) + (Cluster)-[:IN_BLOCK]->(Block) + (Block)-[:NEXT_BLOCK]->(Block) + (Cluster)-[:CO_OCCURS {weight, lift}]->(Cluster) + + docker compose up -d neo4j + python neo4j_load.py --chrom chr22 # everything (2.4M CARRIES edges, a few minutes) + python neo4j_load.py --chrom chr22 --region 45035149-45534032 # just one region, seconds + +Example Cypher once loaded: + MATCH (i:Individual {id:'HG00096'})-[:CARRIES]->(c:Cluster)-[:IN_BLOCK]->(b:Block) RETURN i,c,b LIMIT 50 + MATCH (c:Cluster)-[r:CO_OCCURS]->(d:Cluster) WHERE r.lift > 40 RETURN c,r,d + MATCH (c:Cluster) WHERE c.ancestry_v > 0.7 RETURN c.id, c.ancestry_dominant, c.support ORDER BY c.ancestry_v DESC +""" +from __future__ import annotations + +import argparse +import os +import time +from pathlib import Path + +import numpy as np +import pandas as pd +from neo4j import GraphDatabase + +import config +import haplokg + + +def run_batches(session, query: str, rows: list[dict], batch: int, label: str) -> None: + t0 = time.time() + for i in range(0, len(rows), batch): + session.run(query, rows=rows[i:i + batch]).consume() + print(f" {label:32s} {len(rows):>9,} rows {time.time() - t0:6.1f}s") + + +def main() -> int: + here = Path(__file__).resolve().parent + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--chrom", default="chr22") + parser.add_argument("--kg-dir", type=Path, default=None) + parser.add_argument("--uri", default=config.settings.neo4j_uri) # NEO4J_URI / NEO4J_USER / NEO4J_PASSWORD from .env + parser.add_argument("--user", default=config.settings.neo4j_user) + parser.add_argument("--password", default=config.settings.neo4j_password) + parser.add_argument("--region", default=None, help="start-end (bp): only clusters in these blocks and their edges") + parser.add_argument("--skip-carries", action="store_true", help="skip the 2.4M Individual->Cluster edges") + parser.add_argument("--batch", type=int, default=10000) + parser.add_argument("--wipe", action="store_true", help="delete everything in the database first") + args = parser.parse_args() + kg_dir = args.kg_dir or here / "outputs" / "kg" / args.chrom + + kg = haplokg.load_kg(kg_dir) + ind, cl, bl, co, nb = kg["individuals"], kg["clusters"], kg["blocks"], kg["co_occurs"], kg["next_block"] + assoc_path = here / "outputs" / "cooccurrence" / args.chrom / "cluster_phenotype_association.csv" + if assoc_path.exists(): + assoc = pd.read_csv(assoc_path)[["cluster_idx", "ancestry_dominant", "ancestry_cramers_v"]] + cl = cl.merge(assoc, on="cluster_idx", how="left") + split_path = here / "outputs" / "splits" / args.chrom / "split_seed42.csv" + ind["split"] = pd.read_csv(split_path)["split"].to_numpy() if split_path.exists() else "n/a" + + if args.region: + start, end = (int(x) for x in args.region.split("-")) + bl = bl[(bl["end"] >= start) & (bl["start"] <= end)] + cl = cl[cl["block_idx"].isin(bl["block_idx"])] + keep = set(cl["cluster_idx"]) + co = co[co["src"].isin(keep) & co["dst"].isin(keep)] + nb = nb[nb["src"].isin(bl["block_idx"]) & nb["dst"].isin(bl["block_idx"])] + carries = kg["carries"].tocsc()[:, cl["cluster_idx"].to_numpy()].tocoo() + + driver = GraphDatabase.driver(args.uri, auth=(args.user, args.password)) + with driver.session() as s: + if args.wipe: + while s.run("MATCH (n) WITH n LIMIT 50000 DETACH DELETE n RETURN count(n) AS c").single()["c"]: + pass + for q in ( + "CREATE CONSTRAINT individual_id IF NOT EXISTS FOR (i:Individual) REQUIRE i.id IS UNIQUE", + "CREATE CONSTRAINT cluster_idx IF NOT EXISTS FOR (c:Cluster) REQUIRE c.idx IS UNIQUE", + "CREATE CONSTRAINT block_idx IF NOT EXISTS FOR (b:Block) REQUIRE b.idx IS UNIQUE", + "CREATE INDEX cluster_id IF NOT EXISTS FOR (c:Cluster) ON (c.id)", + "CREATE INDEX individual_ancestry IF NOT EXISTS FOR (i:Individual) ON (i.ancestry)", + ): + s.run(q).consume() + + def clean(frame: pd.DataFrame) -> list[dict]: + return [{k: (None if (isinstance(v, float) and np.isnan(v)) else v) for k, v in r.items()} + for r in frame.to_dict("records")] + + print(f"loading {args.chrom} into {args.uri}") + run_batches(s, """UNWIND $rows AS r MERGE (b:Block {idx: r.idx}) + SET b.id = r.id, b.chrom = r.chrom, b.start = r.start, b.end = r.end, b.length = r.length, + b.n_clusters = r.n_clusters, b.entropy = r.entropy, b.dominance = r.dominance""", + clean(bl.rename(columns={"block_idx": "idx", "block_id": "id", "chr": "chrom", "block_length": "length", + "shannon_entropy": "entropy"}) + [["idx", "id", "chrom", "start", "end", "length", "n_clusters", "entropy", "dominance"]] + .astype({"idx": int, "start": int, "end": int, "length": int})), args.batch, "Block nodes") + + cl_rows = cl.rename(columns={"cluster_idx": "idx", "cluster_id": "id"}) + cl_rows["ancestry_dominant"] = cl_rows.get("ancestry_dominant", pd.Series([None] * len(cl_rows))) + cl_rows["ancestry_v"] = cl_rows.get("ancestry_cramers_v", pd.Series([np.nan] * len(cl_rows))) + run_batches(s, """UNWIND $rows AS r MERGE (c:Cluster {idx: r.idx}) + SET c.id = r.id, c.block_id = r.block_id, c.support = r.support, c.support_frac = r.support_frac, + c.ancestry_dominant = r.ancestry_dominant, c.ancestry_v = r.ancestry_v + WITH c, r MATCH (b:Block {idx: r.block_idx}) MERGE (c)-[:IN_BLOCK]->(b)""", + clean(cl_rows[["idx", "id", "block_id", "block_idx", "support", "support_frac", "ancestry_dominant", "ancestry_v"]] + .astype({"idx": int, "block_idx": int, "support": int})), args.batch, "Cluster nodes + IN_BLOCK") + + run_batches(s, """UNWIND $rows AS r MATCH (a:Block {idx: r.src}), (b:Block {idx: r.dst}) MERGE (a)-[:NEXT_BLOCK]->(b)""", + clean(nb.astype(int)), args.batch, "NEXT_BLOCK") + + run_batches(s, """UNWIND $rows AS r MATCH (a:Cluster {idx: r.src}), (b:Cluster {idx: r.dst}) + MERGE (a)-[e:CO_OCCURS]->(b) SET e.weight = r.weight, e.lift = r.lift""", + clean(co.astype({"src": int, "dst": int, "weight": float, "lift": float})), args.batch, "CO_OCCURS") + + run_batches(s, """UNWIND $rows AS r MERGE (i:Individual {id: r.id}) + SET i.ancestry = r.ancestry, i.population = r.population, i.sex = r.sex, i.split = r.split""", + clean(ind.rename(columns={"individual_id": "id"})[["id", "ancestry", "population", "sex", "split"]]), + args.batch, "Individual nodes") + + if not args.skip_carries: + ids = ind["individual_id"].to_numpy() + cidx = cl["cluster_idx"].to_numpy() + rows = [{"i": ids[r], "c": int(cidx[c])} for r, c in zip(carries.row, carries.col)] + run_batches(s, """UNWIND $rows AS r MATCH (i:Individual {id: r.i}), (c:Cluster {idx: r.c}) MERGE (i)-[:CARRIES]->(c)""", + rows, args.batch, "CARRIES") + + counts = s.run("""MATCH (n) WITH labels(n)[0] AS l, count(*) AS c RETURN l, c ORDER BY l""").data() + rels = s.run("""MATCH ()-[r]->() WITH type(r) AS t, count(*) AS c RETURN t, c ORDER BY t""").data() + driver.close() + print("nodes:", {r["l"]: r["c"] for r in counts}) + print("relationships:", {r["t"]: r["c"] for r in rels}) + print("browse: http://localhost:7474 (neo4j / progenome)") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/genomics/proteome_linear_baseline.py b/genomics/proteome_linear_baseline.py new file mode 100644 index 0000000..17de8ec --- /dev/null +++ b/genomics/proteome_linear_baseline.py @@ -0,0 +1,84 @@ +#!/usr/bin/env python3 +"""Per-protein ridge regression: can the genome (carrier matrix) predict each protein's harmonised level? + +The GNN compresses the genome into a 64-d embedding, which is the wrong tool for sparse +cis effects (one cluster -> one protein). This is the honest baseline for that question: +one linear model per protein from the 6,551 carrier features, scored as test R^2, split by +whether the ground truth says the protein is cis-affected by a causal cluster. +""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np +import pandas as pd +from sklearn.linear_model import Ridge + +import haplokg +import haplokg_proteins as hp + + +def main() -> int: + here = Path(__file__).resolve().parent + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--chrom", default="chr22") + parser.add_argument("--synth-dir", type=Path, default=None, help="folder with measured_long.csv, sample_metadata.csv, ground_truth.json") + parser.add_argument("--alpha", type=float, nargs="+", default=[10.0, 100.0, 1000.0]) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + kg_dir = here / "outputs" / "kg" / args.chrom + synth = args.synth_dir or here / "outputs" / "proteomics_synth" / args.chrom + out_dir = here / "outputs" / "gnn_v2" / args.chrom / "proteome_ridge_baseline" + out_dir.mkdir(parents=True, exist_ok=True) + + kg = haplokg.load_kg(kg_dir) + bed = here.parent / "proteomics" / "uniprot_chr22.bed" + symbols = here.parent / "proteomics" / "synthetic_proteomics_chr22" / "gene_symbol_cache.csv" + t = hp.build_protein_tables(kg, bed, synth / "measured_long.csv", synth / "sample_metadata.csv", symbols) + Y, M = t["abundance"].toarray(), t["observed"].toarray().astype(bool) + X = kg["carries"].astype(np.float32) + split = haplokg.load_or_make_split(kg, here / "outputs" / "splits" / args.chrom / f"split_seed{args.seed}.csv", seed=args.seed) + tr, va, te = (split == "train"), (split == "val"), (split == "test") + Yf = np.where(M, Y, 0.0) # unobserved -> 0 (= the harmonised mean); good enough for a baseline + + def r2(mask, yhat): + out = np.full(Y.shape[1], np.nan) + for j in range(Y.shape[1]): + m = mask & M[:, j] + if m.sum() >= 5: + ss_res = ((Y[m, j] - yhat[m, j]) ** 2).sum(); ss_tot = ((Y[m, j] - Y[m, j].mean()) ** 2).sum() + out[j] = 1 - ss_res / ss_tot if ss_tot > 0 else np.nan + return out + + best = None + for alpha in args.alpha: + model = Ridge(alpha=alpha).fit(X[tr], Yf[tr]) + yhat = model.predict(X) # all rows once; r2() masks by split + score = np.nanmean(r2(va, yhat)) + print(f"alpha={alpha:<7} val mean R^2 {score:.4f}") + if best is None or score > best[1]: + best = (alpha, score, yhat) + alpha, _, yhat = best + test_r2 = r2(te, yhat) + + truth = json.loads((synth / "ground_truth.json").read_text()) + prot = t["proteins"]["protein_id"].tolist() + cis = {c["cis_protein"]: c for c in truth["causal_clusters"]} + is_cis = np.array([p in cis for p in prot]) + table = pd.DataFrame({"protein_id": prot, "test_r2": test_r2, "is_cis": is_cis, + "beta_cis": [cis[p]["beta_cis"] if p in cis else np.nan for p in prot]}) + table.to_csv(out_dir / "protein_r2.csv", index=False) + report = {"alpha": alpha, "mean_r2_all": float(np.nanmean(test_r2)), "median_r2_all": float(np.nanmedian(test_r2)), + "mean_r2_cis": float(np.nanmean(test_r2[is_cis])), "mean_r2_other": float(np.nanmean(test_r2[~is_cis])), + "n_cis": int(is_cis.sum()), "cis_proteins_with_r2_over_0.1": int((test_r2[is_cis] > 0.1).sum()), + "other_proteins_with_r2_over_0.1": int((test_r2[~is_cis] > 0.1).sum())} + (out_dir / "metrics.json").write_text(json.dumps(report, indent=2)) + print(json.dumps(report, indent=2)) + print(table[is_cis].sort_values("test_r2", ascending=False).to_string(index=False, float_format=lambda v: f"{v:.3f}")) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/genomics/proteomics_synth_1000g.py b/genomics/proteomics_synth_1000g.py new file mode 100644 index 0000000..f8d7b1c --- /dev/null +++ b/genomics/proteomics_synth_1000g.py @@ -0,0 +1,150 @@ +#!/usr/bin/env python3 +"""Synthetic proteomics for the *real* 1000 Genomes individuals in the knowledge graph. + +Why: the graph's people are 1000G IDs (HG00096, ...). Proteomics keyed to the same IDs is +what makes `Individual -MEASURED-> Protein` edges exist. Until real per-individual +proteomics is in hand, this generator (same spirit as proteomics/generate_synthetic_proteomics.py) +produces log2 intensities for the chr22 proteins with a KNOWN ground truth, so the +genome+proteome model can be scored on whether it recovers it: + + * 3 hospital sites, mixed ancestry (so `site` is a pure batch label -> negative control) + * phenotype (case/control) depends on a few "causal" haploblock clusters (+ age) <- genomic signal + * each causal cluster shifts one protein encoded in the same block (cis-pQTL-like) <- genome->proteome link + * proteins also respond to phenotype / age / sex (like Nolan's generator) + * per-site batch shift + detection-limit missingness (the whiteboard's harmoniser problems) + +Outputs (outputs/proteomics_synth//): site{1,2,3}_proteomics_log2.csv (rows proteins, cols samples, +same layout as proteomics/synthetic_proteomics_chr22), sample_metadata.csv, measured_long.csv +(individual_id, protein_id, site, log2_intensity) and ground_truth.json. +""" +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np +import pandas as pd + +import haplokg + + +def load_protein_bed(path: Path) -> pd.DataFrame: + bed = pd.read_csv(path, sep="\t", header=None, usecols=[0, 1, 2, 3], names=["chrom", "start", "end", "uniprot"]) + bed["protein_id"] = bed["uniprot"].str.split("-").str[0] + # one row per protein: the union span of its isoforms + return bed.groupby(["chrom", "protein_id"], as_index=False).agg(start=("start", "min"), end=("end", "max")) + + +def main() -> int: + here = Path(__file__).resolve().parent + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--chrom", default="chr22") + parser.add_argument("--bed", type=Path, default=here.parent / "proteomics" / "uniprot_chr22.bed") + parser.add_argument("--out-dir", type=Path, default=None) + parser.add_argument("--n-sites", type=int, default=3) + parser.add_argument("--n-causal", type=int, default=20, help="causal clusters driving the phenotype") + parser.add_argument("--prevalence", type=float, default=0.35) + parser.add_argument("--causal-beta-sd", type=float, default=1.5, help="effect size sd of causal clusters on the phenotype logit") + parser.add_argument("--cis-beta-sd", type=float, default=1.0, help="effect size sd of a causal cluster on its cis protein (log2)") + parser.add_argument("--pheno-protein-frac", type=float, default=0.15, help="fraction of proteins that respond to the phenotype") + parser.add_argument("--pheno-effect-sd", type=float, default=0.5, help="effect size sd of the phenotype on responding proteins") + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + out_dir = args.out_dir or here / "outputs" / "proteomics_synth" / args.chrom + out_dir.mkdir(parents=True, exist_ok=True) + rng = np.random.default_rng(args.seed) + + kg = haplokg.load_kg(here / "outputs" / "kg" / args.chrom) + ind, clusters, blocks, carries = kg["individuals"], kg["clusters"], kg["blocks"], kg["carries"] + labelled = ind["ancestry_code"].to_numpy() >= 0 + people = ind[labelled].reset_index(drop=True) + X = carries[labelled].toarray().astype(np.float32) # people x clusters + + # ---- sites: random, stratified on ancestry so every site is mixed -------------------- + site = np.empty(len(people), dtype=object) + for anc in people["ancestry"].unique(): + idx = np.flatnonzero(people["ancestry"] == anc) + rng.shuffle(idx) + for k, chunk in enumerate(np.array_split(idx, args.n_sites)): + site[chunk] = f"SITE{k + 1}" + age = rng.integers(18, 86, size=len(people)) + sex = (people["sex"] == "male").astype(int).to_numpy() + + # ---- proteins and which block encodes them ------------------------------------------ + prot = load_protein_bed(args.bed) + prot = prot[prot["chrom"] == args.chrom].reset_index(drop=True) + b = blocks[["block_idx", "start", "end"]].to_numpy() + prot["block_idx"] = [ + int(b[(b[:, 1] <= e) & (b[:, 2] >= s)][:, 0][0]) if ((b[:, 1] <= e) & (b[:, 2] >= s)).any() else -1 + for s, e in zip(prot["start"], prot["end"]) + ] + proteins = prot["protein_id"].tolist() + P = len(proteins) + + # ---- causal clusters: in blocks that encode a protein, moderately common -------------- + eligible = clusters[(clusters["block_idx"].isin(prot["block_idx"])) & (clusters["support_frac"].between(0.05, 0.6))] + causal = eligible.sample(n=min(args.n_causal, len(eligible)), random_state=args.seed) + beta_pheno = rng.normal(0, args.causal_beta_sd, size=len(causal)) + # each causal cluster shifts one protein from its own block (cis effect) + causal_protein = [] + for blk in causal["block_idx"]: + choices = prot.index[prot["block_idx"] == blk].tolist() + causal_protein.append(int(rng.choice(choices))) + beta_cis = rng.normal(0, args.cis_beta_sd, size=len(causal)) + + # ---- phenotype from genotype (+ age), calibrated to the requested prevalence ---------- + G = X[:, causal["cluster_idx"].to_numpy()] # people x causal (0/1) + logit = G @ beta_pheno + 0.02 * (age - age.mean()) + intercept = np.quantile(logit, 1 - args.prevalence) # top `prevalence` fraction become cases (softly) + p = 1 / (1 + np.exp(-(logit - intercept))) + phenotype = (rng.random(len(people)) < p).astype(int) + + # ---- protein intensities ----------------------------------------------------------- + baseline = rng.uniform(6.0, 16.0, size=P) + b_age, b_sex = rng.normal(0, 0.5, P), rng.normal(0, 0.5, P) + responds = rng.random(P) < args.pheno_protein_frac # only some proteins track the phenotype + b_pheno = np.where(responds, rng.normal(0, args.pheno_effect_sd, P), 0.0) + site_shift = {f"SITE{k + 1}": rng.normal(0, 0.3, size=P) for k in range(args.n_sites)} # batch effect + age_z = (age - age.mean()) / age.std() + M = (baseline[None, :] + np.outer(age_z, b_age) + np.outer(sex, b_sex) + np.outer(phenotype, b_pheno) + + rng.normal(0, 0.8, size=(len(people), P)) + rng.normal(0, 0.3, size=(len(people), P))) + for j, (pi, bc) in enumerate(zip(causal_protein, beta_cis)): + M[:, pi] += bc * G[:, j] + for k in range(len(people)): + M[k] += site_shift[site[k]] + scarcity = (baseline.max() - baseline) / (baseline.max() - baseline.min() + 1e-9) + missing = rng.random(M.shape) < 0.15 * scarcity[None, :] # MNAR at the detection limit + M = np.round(M, 3); M[missing] = np.nan + + # ---- write ------------------------------------------------------------------------ + meta = pd.DataFrame({"sample_id": people["individual_id"], "site": site, "age": age, "sex": sex, + "phenotype": phenotype, "ancestry": people["ancestry"]}) + meta.to_csv(out_dir / "sample_metadata.csv", index=False) + for s in sorted(set(site)): + cols = meta.index[meta["site"] == s] + frame = pd.DataFrame(M[cols].T, index=proteins, columns=meta.loc[cols, "sample_id"]) + frame.index.name = "protein_id" + frame.to_csv(out_dir / f"{s.lower()}_proteomics_log2.csv") + long = pd.DataFrame(M, index=meta["sample_id"], columns=proteins).stack(future_stack=True).dropna().reset_index() + long.columns = ["individual_id", "protein_id", "log2_intensity"] + long = long.merge(meta[["sample_id", "site"]], left_on="individual_id", right_on="sample_id").drop(columns="sample_id") + long.to_csv(out_dir / "measured_long.csv", index=False) + truth = { + "seed": args.seed, "n_people": int(len(people)), "n_proteins": P, "prevalence_observed": float(phenotype.mean()), + "sites": {s: int((site == s).sum()) for s in sorted(set(site))}, + "causal_clusters": [{"cluster_id": c, "cluster_idx": int(i), "beta_phenotype": float(bp), + "cis_protein": proteins[pi], "beta_cis": float(bc)} + for c, i, bp, pi, bc in zip(causal["cluster_id"], causal["cluster_idx"], beta_pheno, causal_protein, beta_cis)], + "protein_effects": {"beta_age": b_age.tolist(), "beta_sex": b_sex.tolist(), "beta_phenotype": b_pheno.tolist(), "proteins": proteins}, + "site_shift_sd": 0.3, "missing_rate_at_lod": 0.15, "n_phenotype_responsive_proteins": int(responds.sum()), + "params": vars(args) | {"bed": str(args.bed), "out_dir": str(out_dir)}, + } + (out_dir / "ground_truth.json").write_text(json.dumps(truth, indent=2)) + print(f"{len(people)} people, {P} proteins, sites {truth['sites']}, cases {phenotype.mean():.2f}, " + f"{len(causal)} causal clusters, {long.shape[0]:,} measured values -> {out_dir}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/genomics/requirements.txt b/genomics/requirements.txt new file mode 100644 index 0000000..34318d5 --- /dev/null +++ b/genomics/requirements.txt @@ -0,0 +1,14 @@ +# Pinned to the versions the chr22 results were produced with (2026-09-17). +# torch itself is installed separately (CPU or CUDA index) - see setup.sh / Dockerfile. +# torch==2.14.0 +torch_geometric==2.8.0.post1 +pandas==3.0.5 +scipy==1.18.1 +scikit-learn==1.9.1 +matplotlib==3.11.2 +pyarrow==25.0.1 +networkx==3.6.1 +neo4j==6.3.1 +pytest==9.1.1 +tabulate==0.10.0 +nvflare==2.9.0 diff --git a/genomics/run_all.sh b/genomics/run_all.sh new file mode 100755 index 0000000..ab73091 --- /dev/null +++ b/genomics/run_all.sh @@ -0,0 +1,20 @@ +#!/usr/bin/env bash +# Whole pipeline for one chromosome. Env: CHROM (chr22), INIT (svd|node2vec|learned), TARGETS. +set -euo pipefail +cd "$(dirname "$0")" +. ./load_env.sh +CHROM="${CHROM:-chr22}" +INIT="${INIT:-svd}" +TARGETS="${TARGETS:-ancestry population sex}" +PY="${PYTHON:-python}" + +echo "== 1. fetch"; bash fetch_data.sh "$CHROM" +echo "== 2. knowledge graph"; $PY build_kg.py --chrom "$CHROM" +echo "== 3. co-occurrence"; $PY cooccurrence_analysis.py --chrom "$CHROM" +echo "== 4. baseline"; $PY baseline.py --chrom "$CHROM" +echo "== 5. graph display"; $PY graph_explore.py --chrom "$CHROM" +for target in $TARGETS; do + echo "== 6. GNN ($target, init=$INIT)"; $PY train_gnn.py --chrom "$CHROM" --target "$target" --init "$INIT" +done +echo "== 7. embeddings"; $PY embeddings.py --chrom "$CHROM" +echo "done -> outputs/" diff --git a/genomics/run_v2.sh b/genomics/run_v2.sh new file mode 100755 index 0000000..b3b87b9 --- /dev/null +++ b/genomics/run_v2.sh @@ -0,0 +1,28 @@ +#!/usr/bin/env bash +# Schema v2 chain: proteomics on 1000G IDs -> genes/proteins in the graph -> genome+proteome GNN ablations -> EDA. +# Run after run_all.sh (needs outputs/kg/ and the shared split). Env: CHROM (chr22), PYTHON. +set -euo pipefail +cd "$(dirname "$0")" +. ./load_env.sh # NVIDIA_API_KEY etc. from .env / ~/.progenome.env +CHROM="${CHROM:-chr22}" +PY="${PYTHON:-python}" + +echo "== v2.1 synthetic proteomics on 1000G IDs (3 mixed-ancestry sites, saved ground truth)" +$PY proteomics_synth_1000g.py --chrom "$CHROM" +echo "== v2.2 genes + proteins + MEASURED edges (harmonised) -> hetero_v2.pt" +$PY build_kg_v2.py --chrom "$CHROM" +echo "== v2.3 EDA report" +$PY eda.py --chrom "$CHROM" +echo "== v2.4 phenotype: genome vs proteome vs both" +for m in genome proteome both; do $PY train_gnn_v2.py --chrom "$CHROM" --target phenotype --modality "$m"; done +for m in genome both; do $PY train_gnn_v2.py --chrom "$CHROM" --target phenotype --modality "$m" --init raw; done # raw carrier row: saliency vs ground truth +echo "== v2.5 controls: site (batch) and ancestry on the full graph" +$PY train_gnn_v2.py --chrom "$CHROM" --target site --modality both +$PY train_gnn_v2.py --chrom "$CHROM" --target ancestry --modality both +echo "== v2.6 genome -> proteome: per-protein ridge baseline (the honest cis test)" +$PY proteome_linear_baseline.py --chrom "$CHROM" +echo "== v2.7 GraphRAG decoder (dry run; set NVIDIA_API_KEY to call the LLM)" +WHO=$(awk -F, 'NR==2{print $1}' "outputs/gnn_v2/$CHROM/phenotype_both_raw/test_predictions.csv") +$PY graphrag_decoder.py --chrom "$CHROM" --individual "$WHO" --run phenotype_both_raw --dry-run > /dev/null +if [ -n "${NVIDIA_API_KEY:-}" ]; then $PY graphrag_decoder.py --chrom "$CHROM" --individual "$WHO" --run phenotype_both_raw; fi +echo "done -> outputs/{proteomics_synth,kg,eda,gnn_v2,graphrag}/$CHROM" diff --git a/genomics/setup.sh b/genomics/setup.sh new file mode 100755 index 0000000..e2422c4 --- /dev/null +++ b/genomics/setup.sh @@ -0,0 +1,27 @@ +#!/usr/bin/env bash +# One-shot local environment: creates ./.venv with torch (CPU by default, CUDA if a GPU is present) + pinned deps. +# bash setup.sh # auto: CUDA wheels if nvidia-smi works, else CPU wheels +# TORCH_INDEX=cpu bash setup.sh +set -euo pipefail +cd "$(dirname "$0")" +TORCH_VERSION="${TORCH_VERSION:-2.14.0}" +if [ -z "${TORCH_INDEX:-}" ]; then + if command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi >/dev/null 2>&1; then TORCH_INDEX=cu126; else TORCH_INDEX=cpu; fi +fi +PY="${PYTHON:-python3}" +if command -v uv >/dev/null 2>&1; then + uv venv --python 3.13 .venv >/dev/null + uv pip install --python .venv/bin/python "torch==${TORCH_VERSION}" --index-url "https://download.pytorch.org/whl/${TORCH_INDEX}" + uv pip install --python .venv/bin/python -r requirements.txt + if [ "$TORCH_INDEX" != cpu ]; then + uv pip install --python .venv/bin/python pyg_lib -f "https://data.pyg.org/whl/torch-${TORCH_VERSION}+${TORCH_INDEX}.html" || echo "(no pyg_lib wheel -> Node2Vec falls back to SVD)" + fi +else + "$PY" -m venv .venv + .venv/bin/pip install -q --upgrade pip + .venv/bin/pip install "torch==${TORCH_VERSION}" --index-url "https://download.pytorch.org/whl/${TORCH_INDEX}" + .venv/bin/pip install -r requirements.txt +fi +.venv/bin/python -c "import torch, torch_geometric; print('torch', torch.__version__, '| pyg', torch_geometric.__version__, '| cuda', torch.cuda.is_available())" +.venv/bin/python -m pytest tests -q +echo "ready: source .venv/bin/activate (or just use: make run)" diff --git a/genomics/tests/conftest.py b/genomics/tests/conftest.py new file mode 100644 index 0000000..5801bea --- /dev/null +++ b/genomics/tests/conftest.py @@ -0,0 +1,4 @@ +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) diff --git a/genomics/tests/test_haplokg.py b/genomics/tests/test_haplokg.py new file mode 100644 index 0000000..90d7258 --- /dev/null +++ b/genomics/tests/test_haplokg.py @@ -0,0 +1,150 @@ +"""Unit tests for haplokg on a 4-individual, 5-cluster, 2-block toy graph.""" +import gzip + +import numpy as np +import pandas as pd +import pytest + +import haplokg + +INDIVIDUALS = ["I1", "I2", "I3", "I4"] +NODE_ROWS = [ + # cluster_id, block_id, I1 I2 I3 I4 support + ("chr22_100-200_cluster1", "chr22_100-200", [1, 1, 0, 0]), # 2 keep + ("chr22_100-200_cluster2", "chr22_100-200", [0, 0, 1, 1]), # 2 keep + ("chr22_100-200_cluster3", "chr22_100-200", [0, 0, 0, 1]), # 1 singleton -> drop + ("chr22_200-300_cluster1", "chr22_200-300", [1, 1, 1, 1]), # 4 universal -> drop (symmetric) + ("chr22_200-300_cluster2", "chr22_200-300", [1, 0, 1, 0]), # 2 keep +] + + +@pytest.fixture +def toy(tmp_path): + nodes = tmp_path / "nodes.csv.gz" + with gzip.open(nodes, "wt") as fh: + fh.write("id,high_dim_edge," + ",".join(INDIVIDUALS) + "\n") + for cid, bid, bits in NODE_ROWS: + fh.write(f"{cid},{bid}," + ",".join(map(str, bits)) + "\n") + + stats = tmp_path / "block_stats.tsv" + pd.DataFrame( + { + "chr": ["chr22", "chr22", "chr1"], + "block": ["chr22_200-300", "chr22_100-200", "chr1_5-9"], + "start": [200, 100, 5], + "end": [300, 200, 9], + "block_length": [100, 100, 4], + "n_haplotypes": [8, 8, 8], + "n_clusters": [2, 3, 1], + "max_cluster_size": [4, 2, 1], + "singleton_count": [0, 1, 1], + "dominance": [0.5, 0.25, 1.0], + "shannon_entropy": [1.0, 1.5, 0.0], + } + ).to_csv(stats, sep="\t", index=False) + + pheno = tmp_path / "phenotypes_real.csv" + rows = [] + for ind, anc, pop, sex in [("I1", "EUR", "GBR", "male"), ("I2", "AFR", "YRI", "female"), ("I3", "EUR", "FIN", "female")]: + rows += [(ind, "ancestry", anc), (ind, "population", pop), (ind, "sex", sex)] + pd.DataFrame(rows, columns=["individual_id", "phenotype", "value"]).assign(source="test").to_csv(pheno, index=False) + + edges = tmp_path / "edges.csv.gz" + pd.DataFrame( + { + "source": ["chr22_100-200_cluster1", "chr22_100-200_cluster3"], + "target": ["chr22_200-300_cluster2", "chr22_200-300_cluster2"], + "weight": [1, 1], + "lift": [2.0, 4.0], + } + ).to_csv(edges, index=False) + return dict(nodes=nodes, stats=stats, pheno=pheno, edges=edges) + + +def test_parse_ids(): + assert haplokg.parse_block_id("chr22_17099658-17118145") == ("chr22", 17099658, 17118145) + assert haplokg.parse_cluster_id("chr22_17099658-17118145_cluster219") == ("chr22_17099658-17118145", 219) + with pytest.raises(ValueError): + haplokg.parse_cluster_id("chr22_1-2") + + +def test_read_node_matrix_streams_to_sparse(toy): + ids, blocks, individuals, matrix = haplokg.read_node_matrix(toy["nodes"], chunksize=2) + assert individuals == INDIVIDUALS + assert ids == [r[0] for r in NODE_ROWS] + assert blocks == [r[1] for r in NODE_ROWS] + assert matrix.shape == (5, 4) and matrix.dtype == np.int8 + assert np.asarray(matrix.sum(axis=1)).ravel().tolist() == [2, 2, 1, 4, 2] + + +def _tables(toy, **kw): + ids, blocks, individuals, matrix = haplokg.read_node_matrix(toy["nodes"]) + return haplokg.build_tables( + ids, blocks, individuals, matrix, + haplokg.load_phenotypes(toy["pheno"]), + haplokg.load_block_stats(toy["stats"], "chr22"), + haplokg.load_edges(toy["edges"]), + **kw, + ) + + +def test_symmetric_support_filter_drops_singletons_and_universal(toy): + t = _tables(toy, min_support=2, symmetric=True) + assert t["clusters"]["cluster_id"].tolist() == [ + "chr22_100-200_cluster1", "chr22_100-200_cluster2", "chr22_200-300_cluster2", + ] + assert t["carries"].shape == (4, 3) # individuals x kept clusters + assert t["carries"].sum() == 6 # 2 + 2 + 2 carriers + assert t["carries"][0].toarray().ravel().tolist() == [1, 0, 1] # I1 carries b1c1 and b2c2 + + +def test_non_symmetric_filter_keeps_universal(toy): + t = _tables(toy, min_support=2, symmetric=False) + assert "chr22_200-300_cluster1" in t["clusters"]["cluster_id"].tolist() + + +def test_blocks_sorted_by_position_and_linked(toy): + t = _tables(toy, min_support=2) + assert t["blocks"]["block_id"].tolist() == ["chr22_100-200", "chr22_200-300"] # chr1 row filtered out + assert t["next_block"].to_dict("records") == [{"src": 0, "dst": 1}] + assert t["clusters"]["block_idx"].tolist() == [0, 0, 1] + + +def test_labels_encoded_with_missing_as_minus_one(toy): + t = _tables(toy, min_support=2) + ind = t["individuals"].set_index("individual_id") + assert t["label_maps"]["ancestry"] == ["AFR", "EUR"] + assert ind.loc["I1", "ancestry_code"] == 1 and ind.loc["I2", "ancestry_code"] == 0 + assert ind.loc["I4", "ancestry_code"] == -1 and ind.loc["I4", "sex_code"] == -1 + assert t["label_maps"]["sex"] == ["female", "male"] + + +def test_edges_remapped_and_filtered_endpoints_dropped(toy): + t = _tables(toy, min_support=2) + assert t["co_occurs"].to_dict("records") == [{"src": 0, "dst": 2, "weight": 1, "lift": 2.0}] + assert t["n_edges_dropped"] == 1 + + +def test_hetero_data_shapes(toy): + torch = pytest.importorskip("torch") + t = _tables(toy, min_support=2) + data = haplokg.to_hetero_data(t) + assert data["individual"].num_nodes == 4 + assert data["cluster"].x.shape == (3, 2 + len(haplokg.BLOCK_FEATURES)) + assert data["block"].x.shape == (2, len(haplokg.BLOCK_FEATURES)) + assert data["individual", "carries", "cluster"].edge_index.shape == (2, 6) + assert data["cluster", "rev_carries", "individual"].edge_index.shape == (2, 6) + assert data["cluster", "co_occurs", "cluster"].edge_index.shape == (2, 2) # both directions + assert data["cluster", "co_occurs", "cluster"].edge_attr.shape == (2, 2) + assert data["block", "next_block", "block"].edge_index.shape == (2, 2) + assert data["individual"].y_ancestry.tolist() == [1, 0, 1, -1] + assert not torch.isnan(data["cluster"].x).any() + data.validate() + + +def test_save_tables_roundtrip(toy, tmp_path): + t = _tables(toy, min_support=2) + summary = haplokg.save_tables(t, tmp_path / "kg") + assert summary["n_clusters_kept"] == 3 and summary["n_individuals_labelled"] == 3 + reloaded = pd.read_csv(tmp_path / "kg" / "clusters.csv") + assert len(reloaded) == 3 diff --git a/genomics/tests/test_haplokg_proteins.py b/genomics/tests/test_haplokg_proteins.py new file mode 100644 index 0000000..a8742c0 --- /dev/null +++ b/genomics/tests/test_haplokg_proteins.py @@ -0,0 +1,66 @@ +"""Protein layer on top of the toy graph from test_haplokg.""" +import numpy as np +import pandas as pd +import pytest + +import haplokg +import haplokg_proteins as hp +from test_haplokg import toy, _tables # noqa: F401 (fixture reuse) + + +@pytest.fixture +def protein_files(tmp_path): + bed = tmp_path / "uniprot.bed" + pd.DataFrame([ + ["chr22", 120, 180, "P00001-1", 1000, "+"], # gene GA, inside block chr22_100-200 + ["chr22", 120, 190, "P00001-2", 0, "+"], + ["chr22", 190, 260, "P00002", 1000, "-"], # gene GB, spans both blocks + ["chr22", 900, 950, "P00003", 1000, "+"], # gene GC, in no block + ]).to_csv(bed, sep="\t", header=False, index=False) + symbols = tmp_path / "symbols.csv" + pd.DataFrame({"protein_id": ["P00001", "P00002", "P00003"], "gene_symbol": ["GA", "GB", "GC"]}).to_csv(symbols, index=False) + measured = tmp_path / "measured.csv" + pd.DataFrame({ + "individual_id": ["I1", "I1", "I2", "I2", "I3", "ZZ"], + "protein_id": ["P00001", "P00002", "P00001", "P00002", "P00001", "P00001"], + "log2_intensity": [10.0, 8.0, 12.0, 8.5, 11.0, 5.0], + "site": ["S1", "S1", "S1", "S1", "S2", "S2"], + }).to_csv(measured, index=False) + meta = tmp_path / "meta.csv" + pd.DataFrame({"sample_id": ["I1", "I2", "I3"], "site": ["S1", "S1", "S2"], "age": [30, 40, 50], "sex": [1, 0, 0], + "phenotype": [1, 0, 1]}).to_csv(meta, index=False) + return dict(bed=bed, symbols=symbols, measured=measured, meta=meta) + + +def test_protein_tables(toy, protein_files): + kg = _tables(toy, min_support=2) + t = hp.build_protein_tables(kg, protein_files["bed"], protein_files["measured"], protein_files["meta"], protein_files["symbols"]) + assert t["proteins"]["protein_id"].tolist() == ["P00001", "P00002", "P00003"] + assert t["proteins"].set_index("protein_id").loc["P00001", "n_isoforms"] == 2 + assert t["genes"]["gene_symbol"].tolist() == ["GA", "GB", "GC"] + bg = t["block_gene"].sort_values(["gene_idx", "block_idx"])[["block_idx", "gene_idx"]].values.tolist() + assert bg == [[0, 0], [0, 1], [1, 1]] # GA in block0; GB overlaps block0 and block1; GC nowhere + assert t["n_unmatched_measurements"] == 1 # ZZ is not in the graph + assert len(t["measured"]) == 5 + z = t["measured"].set_index(["individual_idx", "protein_idx"])["z"] + # harmonisation runs per site over *all* the site's samples (ZZ included, even though ZZ is not in the graph): + # S2/P00001 = {11.0 (I3), 5.0 (ZZ)} -> median 8, MAD 3 -> z(I3) = 3 / (3 * 1.4826) + assert abs(z.loc[(2, 0)] - 3 / (3 * 1.4826)) < 1e-6 + ind = t["individuals"].set_index("individual_id") + assert ind.loc["I1", "phenotype_code"] == 1 and ind.loc["I4", "phenotype_code"] == -1 + assert t["label_maps"]["site"] == ["S1", "S2"] and ind.loc["I3", "site_code"] == 1 + assert t["abundance"].shape == (4, 3) and t["observed"].sum() == 5 + + +def test_extend_hetero_data(toy, protein_files): + torch = pytest.importorskip("torch") + kg = _tables(toy, min_support=2) + t = hp.build_protein_tables(kg, protein_files["bed"], protein_files["measured"], protein_files["meta"], protein_files["symbols"]) + data = hp.extend_hetero_data(haplokg.to_hetero_data(kg), t) + assert data["gene"].x.shape == (3, 2) and data["protein"].x.shape == (3, 2) + assert data["block", "overlaps", "gene"].edge_index.shape == (2, 3) + assert data["gene", "encodes", "protein"].edge_index.shape == (2, 3) + assert data["individual", "measured", "protein"].edge_index.shape == (2, 5) + assert data["individual", "measured", "protein"].edge_attr.shape == (5, 2) + assert data["individual"].y_phenotype.tolist() == [1, 0, 1, -1] + data.validate() diff --git a/genomics/train_gnn.py b/genomics/train_gnn.py new file mode 100644 index 0000000..84dfef3 --- /dev/null +++ b/genomics/train_gnn.py @@ -0,0 +1,277 @@ +#!/usr/bin/env python3 +"""Heterogeneous GNN (PyTorch Geometric) on the haploblock knowledge graph. + +Node classification on *individual* nodes. Messages flow + individual <-carries-> cluster <-co_occurs (lift-weighted)-> cluster <-in_block-> block <-next_block-> block +so an individual's representation is built from the clusters they carry, from +what those clusters co-occur with across the population, and from the block +structure along the chromosome. + +Embeddings ("--init"): + svd (default) truncated SVD of the individual x cluster carrier matrix + gives both individuals and clusters a shared k-dim starting embedding + node2vec Node2Vec pretraining on the carries + co_occurs graph (needs pyg-lib; + available in the GPU image), falls back to svd if missing + learned free nn.Embedding per individual (no prior) + raw the individual's own 0/1 carrier row (all kept clusters) through a linear + layer, plus SVD for clusters - closest to the logistic-regression baseline +Cluster and block nodes additionally get their z-scored statistics (support, block +length, entropy, dominance, ...). The final hidden layer is exported as the +learned embedding of every individual and cluster. + +Targets: ancestry (5), population (26), sex (2; negative control - autosomal +chromosome, so the honest answer is chance level). +""" +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import numpy as np +import pandas as pd +import torch +import torch.nn.functional as F +from scipy.sparse.linalg import svds +from sklearn.metrics import accuracy_score, balanced_accuracy_score, f1_score +from torch import nn +from torch_geometric.nn import GraphConv, HeteroConv, SAGEConv + +import haplokg + +CO = ("cluster", "co_occurs", "cluster") + + +def pick_device(name: str) -> torch.device: + if name != "auto": + return torch.device(name) + if torch.cuda.is_available(): + return torch.device("cuda") + return torch.device("cpu") # MPS is not enabled by default: scatter ops are still patchy there + + +def zscore(a: np.ndarray) -> np.ndarray: + a = np.asarray(a, dtype=np.float32) + std = a.std(axis=0) + std[std == 0] = 1.0 + return (a - a.mean(axis=0)) / std + + +def svd_embeddings(carries, k: int, seed: int): + """Shared k-dim embedding for individuals (rows) and clusters (columns).""" + u, s, vt = svds(carries.astype(np.float32), k=k, random_state=seed) + order = np.argsort(s)[::-1] + u, s, vt = u[:, order], s[order], vt[order] + return zscore(u * s), zscore(vt.T * s), s + + +def node2vec_embeddings(data, k: int, device: torch.device, seed: int, epochs: int = 50): + """Optional: Node2Vec on the homogeneous (individual + cluster) graph; None if pyg-lib is missing.""" + from torch_geometric.typing import WITH_PYG_LIB + if not WITH_PYG_LIB: # PyG >= 2.6: random walks live in pyg-lib (torch_cluster is deprecated) + return None + from torch_geometric.nn import Node2Vec + n_ind, n_cl = data["individual"].num_nodes, data["cluster"].num_nodes + carries = data["individual", "carries", "cluster"].edge_index + co = data[CO].edge_index + edge_index = torch.cat([ + torch.stack([carries[0], carries[1] + n_ind]), torch.stack([carries[1] + n_ind, carries[0]]), + co + n_ind, + ], dim=1).to(device) + torch.manual_seed(seed) + model = Node2Vec(edge_index, embedding_dim=k, walk_length=20, context_size=10, walks_per_node=10, + num_negative_samples=1, sparse=True, num_nodes=n_ind + n_cl).to(device) + loader = model.loader(batch_size=256, shuffle=True, num_workers=0) + optimizer = torch.optim.SparseAdam(list(model.parameters()), lr=0.01) + for epoch in range(epochs): + total = 0.0 + for pos_rw, neg_rw in loader: + optimizer.zero_grad() + loss = model.loss(pos_rw.to(device), neg_rw.to(device)) + loss.backward(); optimizer.step(); total += loss.item() + print(f" node2vec epoch {epoch + 1}/{epochs} loss {total / len(loader):.4f}") + emb = model.embedding.weight.detach().cpu().numpy() + return zscore(emb[:n_ind]), zscore(emb[n_ind:]) + + +class HeteroGNN(nn.Module): + def __init__(self, in_dims: dict, hidden: int, n_classes: int, layers: int = 2, dropout: float = 0.3, + learned_individual: int | None = None, aggr: str = "mean"): + super().__init__() + self.learned = nn.Embedding(learned_individual, hidden) if learned_individual else None + self.proj = nn.ModuleDict({t: nn.Linear(d, hidden) for t, d in in_dims.items() if d > 0}) + self.convs = nn.ModuleList() + for _ in range(layers): + self.convs.append(HeteroConv({ + ("individual", "carries", "cluster"): SAGEConv((hidden, hidden), hidden, aggr=aggr), + ("cluster", "rev_carries", "individual"): SAGEConv((hidden, hidden), hidden, aggr=aggr), + ("cluster", "in_block", "block"): SAGEConv((hidden, hidden), hidden, aggr=aggr), + ("block", "rev_in_block", "cluster"): SAGEConv((hidden, hidden), hidden, aggr=aggr), + CO: GraphConv(hidden, hidden, aggr="mean"), # takes edge_weight = normalised log-lift + ("block", "next_block", "block"): SAGEConv((hidden, hidden), hidden, aggr=aggr), + }, aggr="sum")) + self.norms = nn.ModuleList([nn.ModuleDict({t: nn.LayerNorm(hidden) for t in ("individual", "cluster", "block")}) + for _ in range(layers)]) + self.head = nn.Linear(hidden, n_classes) + self.dropout = dropout + + def encode(self, x_dict, edge_index_dict, edge_weight): + h = {t: self.proj[t](x) for t, x in x_dict.items() if t in self.proj} + if self.learned is not None: + h["individual"] = self.learned.weight if "individual" not in h else h["individual"] + self.learned.weight + for conv, norm in zip(self.convs, self.norms): + out = conv(h, edge_index_dict, edge_weight_dict={CO: edge_weight}) + h = {t: F.dropout(F.relu(norm[t](out[t] + h[t])), p=self.dropout, training=self.training) for t in out} + return h + + def forward(self, x_dict, edge_index_dict, edge_weight): + h = self.encode(x_dict, edge_index_dict, edge_weight) + return self.head(h["individual"]), h + + +def metrics(y_true, y_pred) -> dict: + return {"accuracy": float(accuracy_score(y_true, y_pred)), + "balanced_accuracy": float(balanced_accuracy_score(y_true, y_pred)), + "macro_f1": float(f1_score(y_true, y_pred, average="macro", zero_division=0)), "n": int(len(y_true))} + + +def main() -> int: + here = Path(__file__).resolve().parent + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--chrom", default="chr22") + parser.add_argument("--target", default="ancestry", choices=["ancestry", "population", "sex"]) + parser.add_argument("--init", default="svd", choices=["svd", "node2vec", "learned", "raw"]) + parser.add_argument("--aggr", default="mean", choices=["mean", "sum"], help="neighbourhood aggregation of the SAGE layers") + parser.add_argument("--embed-dim", type=int, default=32, help="k for svd / node2vec") + parser.add_argument("--node2vec-epochs", type=int, default=50, help="Node2Vec pretraining epochs (5 was clearly undertrained: loss still falling)") + parser.add_argument("--hidden", type=int, default=64) + parser.add_argument("--layers", type=int, default=2) + parser.add_argument("--dropout", type=float, default=0.3) + parser.add_argument("--lr", type=float, default=0.005) + parser.add_argument("--weight-decay", type=float, default=5e-4) + parser.add_argument("--epochs", type=int, default=300) + parser.add_argument("--patience", type=int, default=30) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--device", default="auto") + parser.add_argument("--kg-dir", type=Path, default=None) + parser.add_argument("--out-dir", type=Path, default=None) + args = parser.parse_args() + kg_dir = args.kg_dir or here / "outputs" / "kg" / args.chrom + run_name = f"{args.target}_{args.init}" + ("" if args.aggr == "mean" else f"_{args.aggr}") + out_dir = args.out_dir or here / "outputs" / "gnn" / args.chrom / run_name + out_dir.mkdir(parents=True, exist_ok=True) + torch.manual_seed(args.seed); np.random.seed(args.seed) + device = pick_device(args.device) + print(f"device: {device} target: {args.target} init: {args.init}") + + kg = haplokg.load_kg(kg_dir) + data = torch.load(kg_dir / "hetero.pt", weights_only=False) + split = haplokg.load_or_make_split(kg, here / "outputs" / "splits" / args.chrom / f"split_seed{args.seed}.csv", seed=args.seed) + y = data["individual"][f"y_{args.target}"].clone() + classes = data.label_maps[args.target] + masks = {name: torch.from_numpy((split == name) & (y.numpy() >= 0)) for name in ("train", "val", "test")} + print({k: int(v.sum()) for k, v in masks.items()}) + + # ---- input embeddings ----------------------------------------------------- + t0 = time.time() + emb_dir = here / "outputs" / "embeddings" / args.chrom + emb_dir.mkdir(parents=True, exist_ok=True) + ind_init = cl_init = None + if args.init in ("svd", "node2vec", "raw"): + ind_svd, cl_svd, sing = svd_embeddings(kg["carries"], args.embed_dim, args.seed) + np.save(emb_dir / f"svd{args.embed_dim}_individual.npy", ind_svd) + np.save(emb_dir / f"svd{args.embed_dim}_cluster.npy", cl_svd) + ind_init, cl_init = ind_svd, cl_svd + print(f"svd k={args.embed_dim} done ({time.time() - t0:.1f}s), top singular values {np.round(sing[:5], 1)}") + if args.init == "node2vec": + n2v = node2vec_embeddings(data, args.embed_dim, device, args.seed, epochs=args.node2vec_epochs) + if n2v is None: + print("pyg-lib not available -> using svd embeddings instead") + else: + ind_init, cl_init = n2v + np.save(emb_dir / f"node2vec{args.embed_dim}_individual.npy", ind_init) + np.save(emb_dir / f"node2vec{args.embed_dim}_cluster.npy", cl_init) + + x_dict = {"cluster": data["cluster"].x, "block": data["block"].x} + if cl_init is not None: + x_dict["cluster"] = torch.cat([x_dict["cluster"], torch.from_numpy(cl_init)], dim=1) + if args.init == "raw": + ind_init = kg["carries"].toarray().astype(np.float32) # 2,548 x 6,551 -> 67 MB, fine + if ind_init is not None: + x_dict["individual"] = torch.from_numpy(ind_init) + in_dims = {t: x.shape[1] for t, x in x_dict.items()} + lift = data[CO].edge_attr[:, 1] + edge_weight = torch.log(lift) / torch.log(lift).max() + + x_dict = {t: x.to(device) for t, x in x_dict.items()} + edge_index_dict = {k: v.to(device) for k, v in data.edge_index_dict.items()} + edge_weight = edge_weight.to(device) + y = y.to(device) + masks = {k: v.to(device) for k, v in masks.items()} + + model = HeteroGNN(in_dims, args.hidden, len(classes), args.layers, args.dropout, + learned_individual=data["individual"].num_nodes if args.init == "learned" else None, aggr=args.aggr).to(device) + optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) + counts = torch.bincount(y[masks["train"]], minlength=len(classes)).float() + class_weight = (counts.sum() / counts.clamp(min=1) / len(classes)).to(device) # rebalance minority classes + print(f"model parameters: {sum(p.numel() for p in model.parameters()):,}") + + def evaluate(mask): + model.eval() + with torch.no_grad(): + logits, h = model(x_dict, edge_index_dict, edge_weight) + pred = logits.argmax(1) + return metrics(y[mask].cpu().numpy(), pred[mask].cpu().numpy()), pred, h + + best, best_state, best_epoch, wait, history = -1.0, None, 0, 0, [] + t0 = time.time() + for epoch in range(1, args.epochs + 1): + model.train(); optimizer.zero_grad() + logits, _ = model(x_dict, edge_index_dict, edge_weight) + loss = F.cross_entropy(logits[masks["train"]], y[masks["train"]], weight=class_weight) + loss.backward(); optimizer.step() + val, _, _ = evaluate(masks["val"]) + history.append({"epoch": epoch, "loss": loss.item(), "val_balanced_accuracy": val["balanced_accuracy"]}) + if val["balanced_accuracy"] > best: + best, best_epoch, wait = val["balanced_accuracy"], epoch, 0 + best_state = {k: v.detach().clone() for k, v in model.state_dict().items()} + else: + wait += 1 + if epoch % 10 == 0 or epoch == 1: + print(f"epoch {epoch:4d} loss {loss:.4f} val bal-acc {val['balanced_accuracy']:.3f} (best {best:.3f} @ {best_epoch})") + if wait >= args.patience: + print(f"early stop at epoch {epoch}"); break + train_time = time.time() - t0 + + model.load_state_dict(best_state) + val, _, _ = evaluate(masks["val"]) + test, pred, h = evaluate(masks["test"]) + print(f"\n{args.target}: TEST acc={test['accuracy']:.3f} bal-acc={test['balanced_accuracy']:.3f} macro-F1={test['macro_f1']:.3f}" + f" (best epoch {best_epoch}, {train_time:.0f}s, {train_time / len(history):.2f}s/epoch)") + baseline_path = here / "outputs" / "baseline" / args.chrom / "metrics.json" + if baseline_path.exists(): + b = json.loads(baseline_path.read_text())["targets"].get(args.target, {}).get("test") + if b: + print(f"logistic-regression baseline: acc={b['accuracy']:.3f} bal-acc={b['balanced_accuracy']:.3f} macro-F1={b['macro_f1']:.3f}") + + report = {"chrom": args.chrom, "target": args.target, "init": args.init, "device": str(device), "classes": classes, + "hidden": args.hidden, "layers": args.layers, "aggr": args.aggr, "embed_dim": args.embed_dim, "best_epoch": best_epoch, + "epochs_run": len(history), "train_seconds": train_time, "val": val, "test": test, + "n_parameters": sum(p.numel() for p in model.parameters())} + (out_dir / "metrics.json").write_text(json.dumps(report, indent=2)) + pd.DataFrame(history).to_csv(out_dir / "history.csv", index=False) + ind = kg["individuals"] + te = masks["test"].cpu().numpy() + pd.DataFrame({"individual_id": ind.loc[te, "individual_id"].to_numpy(), + "true": [classes[i] for i in y[masks["test"]].cpu().numpy()], + "pred": [classes[i] for i in pred[masks["test"]].cpu().numpy()]}).to_csv(out_dir / "test_predictions.csv", index=False) + np.save(out_dir / "embedding_individual.npy", h["individual"].cpu().numpy()) + np.save(out_dir / "embedding_cluster.npy", h["cluster"].cpu().numpy()) + torch.save(best_state, out_dir / "model.pt") + print(f"wrote {out_dir}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/genomics/train_gnn_v2.py b/genomics/train_gnn_v2.py new file mode 100644 index 0000000..741b9f7 --- /dev/null +++ b/genomics/train_gnn_v2.py @@ -0,0 +1,302 @@ +#!/usr/bin/env python3 +"""Genome + proteome GNN (schema v2): does integrating the two beat either alone? + +Modalities ("--modality"): + genome the v1 graph only (individual <-> cluster <-> block, co-occurrence) + proteome no graph: an MLP on the individual's harmonised protein abundances (+ observed mask) + both the full v2 graph: genome relations + individual <-measured-> protein <-encodes- gene <-overlaps- block, + and the abundance vector as part of the individual's input + +Targets ("--target"): + phenotype the case/control label from the proteomics metadata (synthetic ground truth known) + site negative control - sites are mixed-ancestry batches, must stay at chance + ancestry / sex as in v1 + proteome regression: predict every protein's harmonised abundance from the GENOME graph alone + (genome modality) - scored as R^2 on test individuals, separately for the proteins the + ground truth says are cis-affected by a causal cluster + +With --init raw (genome/both) the individual's input is its own carrier row, which makes a +per-cluster saliency possible: gradient of the case logit w.r.t. the carrier row, averaged over +test cases, ranked -> precision@k against the ground-truth causal clusters. +""" +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import numpy as np +import pandas as pd +import torch +import torch.nn.functional as F +from sklearn.metrics import accuracy_score, balanced_accuracy_score, f1_score, roc_auc_score +from torch import nn +from torch_geometric.nn import GraphConv, HeteroConv, SAGEConv + +import haplokg +import haplokg_proteins as hp +from train_gnn import pick_device, svd_embeddings, zscore + +CO = ("cluster", "co_occurs", "cluster") +MEAS = ("individual", "measured", "protein") +RMEAS = ("protein", "rev_measured", "individual") + +GENOME_RELS = [("individual", "carries", "cluster"), ("cluster", "rev_carries", "individual"), + ("cluster", "in_block", "block"), ("block", "rev_in_block", "cluster"), CO, ("block", "next_block", "block")] +PROTEIN_RELS = [MEAS, RMEAS, ("block", "overlaps", "gene"), ("gene", "rev_overlaps", "block"), + ("gene", "encodes", "protein"), ("protein", "rev_encodes", "gene")] +WEIGHTED = {CO, MEAS, RMEAS} + + +class HeteroGNNv2(nn.Module): + def __init__(self, in_dims: dict, relations: list, hidden: int, out_dim: int, layers: int = 2, dropout: float = 0.3, aggr: str = "mean"): + super().__init__() + self.proj = nn.ModuleDict({t: nn.Linear(d, hidden) for t, d in in_dims.items()}) + self.node_types = list(in_dims) + self.convs = nn.ModuleList([HeteroConv({ + rel: (GraphConv(hidden, hidden, aggr=aggr) if rel in WEIGHTED else SAGEConv((hidden, hidden), hidden, aggr=aggr)) + for rel in relations}, aggr="sum") for _ in range(layers)]) + self.norms = nn.ModuleList([nn.ModuleDict({t: nn.LayerNorm(hidden) for t in self.node_types}) for _ in range(layers)]) + self.head = nn.Linear(hidden, out_dim) + self.dropout = dropout + + def encode(self, x_dict, edge_index_dict, edge_weight_dict): + h = {t: self.proj[t](x_dict[t]) for t in self.node_types} + for conv, norm in zip(self.convs, self.norms): + out = conv(h, edge_index_dict, edge_weight_dict=edge_weight_dict) + h = {t: F.dropout(F.relu(norm[t](out[t] + h[t])), p=self.dropout, training=self.training) if t in out else h[t] for t in h} + return h + + def forward(self, x_dict, edge_index_dict, edge_weight_dict): + h = self.encode(x_dict, edge_index_dict, edge_weight_dict) + return self.head(h["individual"]), h + + +class MLP(nn.Module): + """proteome-only baseline: same capacity, no graph.""" + def __init__(self, in_dim: int, hidden: int, out_dim: int, dropout: float = 0.3): + super().__init__() + self.net = nn.Sequential(nn.Linear(in_dim, hidden), nn.LayerNorm(hidden), nn.ReLU(), nn.Dropout(dropout), + nn.Linear(hidden, hidden), nn.LayerNorm(hidden), nn.ReLU(), nn.Dropout(dropout)) + self.head = nn.Linear(hidden, out_dim) + + def forward(self, x_dict, *_): + h = self.net(x_dict["individual"]) + return self.head(h), {"individual": h} + + +def cls_metrics(y, pred, prob=None) -> dict: + out = {"accuracy": float(accuracy_score(y, pred)), "balanced_accuracy": float(balanced_accuracy_score(y, pred)), + "macro_f1": float(f1_score(y, pred, average="macro", zero_division=0)), "n": int(len(y))} + if prob is not None and len(np.unique(y)) == 2: + out["roc_auc"] = float(roc_auc_score(y, prob)) + return out + + +def r2_per_protein(y, yhat, mask) -> np.ndarray: + """R^2 per column over observed entries; NaN where a protein has < 5 observed test values.""" + out = np.full(y.shape[1], np.nan) + for j in range(y.shape[1]): + m = mask[:, j] + if m.sum() >= 5: + ss_res = ((y[m, j] - yhat[m, j]) ** 2).sum() + ss_tot = ((y[m, j] - y[m, j].mean()) ** 2).sum() + out[j] = 1 - ss_res / ss_tot if ss_tot > 0 else np.nan + return out + + +def main() -> int: + here = Path(__file__).resolve().parent + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--chrom", default="chr22") + parser.add_argument("--target", default="phenotype", choices=["phenotype", "site", "ancestry", "sex", "proteome"]) + parser.add_argument("--modality", default="both", choices=["genome", "proteome", "both"]) + parser.add_argument("--init", default="svd", choices=["svd", "raw"]) + parser.add_argument("--embed-dim", type=int, default=32) + parser.add_argument("--hidden", type=int, default=64) + parser.add_argument("--layers", type=int, default=2) + parser.add_argument("--dropout", type=float, default=0.3) + parser.add_argument("--lr", type=float, default=0.005) + parser.add_argument("--weight-decay", type=float, default=5e-4) + parser.add_argument("--epochs", type=int, default=300) + parser.add_argument("--patience", type=int, default=30) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--device", default="auto") + parser.add_argument("--top-k", type=int, default=20, help="saliency precision@k vs ground-truth causal clusters") + args = parser.parse_args() + if args.target == "proteome" and args.modality != "genome": + parser.error("--target proteome only makes sense with --modality genome (predict the proteome from the genome graph)") + torch.manual_seed(args.seed); np.random.seed(args.seed) + device = pick_device(args.device) + kg_dir = here / "outputs" / "kg" / args.chrom + out_dir = here / "outputs" / "gnn_v2" / args.chrom / f"{args.target}_{args.modality}_{args.init}" + out_dir.mkdir(parents=True, exist_ok=True) + print(f"device: {device} target: {args.target} modality: {args.modality} init: {args.init}") + + kg = haplokg.load_kg(kg_dir) + pt = hp.load_protein_tables(kg_dir) + data = torch.load(kg_dir / "hetero_v2.pt", weights_only=False) + split = haplokg.load_or_make_split(kg, here / "outputs" / "splits" / args.chrom / f"split_seed{args.seed}.csv", seed=args.seed) + ind2 = pt["individuals"] + n_ind = len(ind2) + + # ---- individual inputs ---------------------------------------------------------- + abundance = pt["abundance"].toarray().astype(np.float32) # harmonised z, 0 where unobserved + observed = pt["observed"].toarray().astype(np.float32) + parts = [] + if args.modality in ("genome", "both"): + if args.init == "raw": + parts.append(kg["carries"].toarray().astype(np.float32)) + raw_offset, n_clusters = 0, kg["carries"].shape[1] + else: + ind_svd, cl_svd, _ = svd_embeddings(kg["carries"], args.embed_dim, args.seed) + parts.append(ind_svd) + if args.modality in ("proteome", "both"): + parts += [abundance, observed] + x_ind = np.concatenate(parts, axis=1) + + x_dict = {"individual": torch.from_numpy(x_ind)} + relations, ew = [], {} + if args.modality in ("genome", "both"): + cl_x = data["cluster"].x + if args.init == "svd": + cl_x = torch.cat([cl_x, torch.from_numpy(cl_svd)], dim=1) + x_dict["cluster"], x_dict["block"] = cl_x, data["block"].x + relations += GENOME_RELS + lift = data[CO].edge_attr[:, 1] + ew[CO] = torch.log(lift) / torch.log(lift).max() + if args.modality == "both": + x_dict["gene"], x_dict["protein"] = data["gene"].x, data["protein"].x + relations += PROTEIN_RELS + ew[MEAS] = data[MEAS].edge_attr[:, 0] # harmonised z as message weight + ew[RMEAS] = data[RMEAS].edge_attr[:, 0] + edge_index_dict = {rel: data[rel].edge_index.to(device) for rel in relations} + ew = {k: v.to(device) for k, v in ew.items()} + x_dict = {t: x.to(device) for t, x in x_dict.items()} + + # ---- targets -------------------------------------------------------------------- + regression = args.target == "proteome" + if regression: + Y = torch.from_numpy(abundance).to(device); Mobs = torch.from_numpy(observed).to(device) + has_label = observed.sum(1) > 0 + out_dim, classes = abundance.shape[1], None + else: + y_np = ind2[f"{args.target}_code"].to_numpy() + classes = pt["label_maps"][args.target] + y = torch.from_numpy(y_np).to(device) + has_label = y_np >= 0 + out_dim = len(classes) + masks = {k: torch.from_numpy((split == k) & has_label).to(device) for k in ("train", "val", "test")} + print({k: int(v.sum()) for k, v in masks.items()}) + + in_dims = {t: x.shape[1] for t, x in x_dict.items()} + if args.modality == "proteome": + model = MLP(in_dims["individual"], args.hidden, out_dim, args.dropout).to(device) + else: + model = HeteroGNNv2(in_dims, relations, args.hidden, out_dim, args.layers, args.dropout).to(device) + optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) + if not regression: + counts = torch.bincount(y[masks["train"]], minlength=out_dim).float() + class_weight = (counts.sum() / counts.clamp(min=1) / out_dim).to(device) + + def loss_fn(logits, mask): + if regression: + diff = (logits - Y) ** 2 * Mobs + return diff[mask].sum() / Mobs[mask].sum().clamp(min=1) + return F.cross_entropy(logits[mask], y[mask], weight=class_weight) + + def evaluate(mask): + model.eval() + with torch.no_grad(): + logits, h = model(x_dict, edge_index_dict, ew) + m = mask.cpu().numpy() + if regression: + yhat, yt, mo = logits.cpu().numpy()[m], abundance[m], observed[m].astype(bool) + r2 = r2_per_protein(yt, yhat, mo) + # select on the masked validation MSE (what is optimised); mean R^2 over 460 mostly + # unpredictable proteins is too noisy to stop on + score = -float(loss_fn(logits, mask)) + return {"mean_r2": float(np.nanmean(r2)), "val_masked_mse": -score, "n": int(m.sum())}, score, r2, h + prob = torch.softmax(logits, 1) + pred = prob.argmax(1) + met = cls_metrics(y[mask].cpu().numpy(), pred[mask].cpu().numpy(), prob[mask, 1].cpu().numpy() if out_dim == 2 else None) + return met, met["balanced_accuracy"], pred, h + + best, best_state, best_epoch, wait, history = -np.inf, None, 0, 0, [] + t0 = time.time() + for epoch in range(1, args.epochs + 1): + model.train(); optimizer.zero_grad() + logits, _ = model(x_dict, edge_index_dict, ew) + loss = loss_fn(logits, masks["train"]) + loss.backward(); optimizer.step() + val, score, _, _ = evaluate(masks["val"]) + history.append({"epoch": epoch, "loss": loss.item(), "val_score": score}) + if score > best: + best, best_epoch, wait = score, epoch, 0 + best_state = {k: v.detach().clone() for k, v in model.state_dict().items()} + else: + wait += 1 + if epoch % 25 == 0 or epoch == 1: + print(f"epoch {epoch:4d} loss {loss.item():.4f} val {score:.3f} (best {best:.3f} @ {best_epoch})") + if wait >= args.patience: + print(f"early stop at epoch {epoch}"); break + train_time = time.time() - t0 + model.load_state_dict(best_state) + val, _, _, _ = evaluate(masks["val"]) + test, _, extra, h = evaluate(masks["test"]) + report = {"chrom": args.chrom, "target": args.target, "modality": args.modality, "init": args.init, "device": str(device), + "hidden": args.hidden, "layers": args.layers, "best_epoch": best_epoch, "epochs_run": len(history), + "train_seconds": round(train_time, 1), "val": val, "test": test, + "n_parameters": sum(p.numel() for p in model.parameters())} + + truth_path = here / "outputs" / "proteomics_synth" / args.chrom / "ground_truth.json" + truth = json.loads(truth_path.read_text()) if truth_path.exists() else None + if regression: + r2 = extra + report["test"]["median_r2"] = float(np.nanmedian(r2)) + if truth: + prot_ids = list(pt["proteins"]["protein_id"]) + cis = {c["cis_protein"] for c in truth["causal_clusters"]} + is_cis = np.array([p in cis for p in prot_ids]) + report["test"]["mean_r2_cis_proteins"] = float(np.nanmean(r2[is_cis])) + report["test"]["mean_r2_other_proteins"] = float(np.nanmean(r2[~is_cis])) + pd.DataFrame({"protein_id": pt["proteins"]["protein_id"], "test_r2": r2}).to_csv(out_dir / "protein_r2.csv", index=False) + print(f"\nproteome from genome graph: mean test R^2 {test['mean_r2']:.3f} (median {report['test']['median_r2']:.3f})" + + (f" cis-affected proteins {report['test']['mean_r2_cis_proteins']:.3f} vs others {report['test']['mean_r2_other_proteins']:.3f}" if truth else "")) + else: + print(f"\n{args.target} [{args.modality}/{args.init}]: TEST acc={test['accuracy']:.3f} bal-acc={test['balanced_accuracy']:.3f} " + f"macro-F1={test['macro_f1']:.3f}" + (f" AUC={test['roc_auc']:.3f}" if "roc_auc" in test else "")) + te = masks["test"].cpu().numpy() + pd.DataFrame({"individual_id": ind2.loc[te, "individual_id"], "true": [classes[i] for i in y[masks["test"]].cpu().numpy()], + "pred": [classes[i] for i in extra[masks["test"]].cpu().numpy()]}).to_csv(out_dir / "test_predictions.csv", index=False) + + # saliency vs ground truth (raw carrier input, case/control target) + if args.init == "raw" and args.target == "phenotype" and truth and args.modality != "proteome": + model.eval() + x_req = {t: x.clone() for t, x in x_dict.items()} + x_req["individual"].requires_grad_(True) + logits, _ = model(x_req, edge_index_dict, ew) + cases = masks["test"] & (y == 1) + logits[cases, 1].sum().backward() + sal = x_req["individual"].grad[:, :n_clusters][cases].mean(0).cpu().numpy() + order = np.argsort(sal)[::-1] + causal_idx = {c["cluster_idx"] for c in truth["causal_clusters"]} + topk = order[:args.top_k] + hits = int(sum(i in causal_idx for i in topk)) + report["saliency"] = {"top_k": args.top_k, "hits_in_ground_truth": hits, "precision_at_k": hits / args.top_k, + "expected_by_chance": len(causal_idx) / n_clusters * args.top_k} + pd.DataFrame({"cluster_id": kg["clusters"]["cluster_id"].to_numpy()[order[:100]], "saliency": sal[order[:100]], + "is_causal": [i in causal_idx for i in order[:100]]}).to_csv(out_dir / "saliency_top100.csv", index=False) + print(f"saliency: {hits}/{args.top_k} of the top-{args.top_k} clusters are ground-truth causal (chance {report['saliency']['expected_by_chance']:.2f})") + + (out_dir / "metrics.json").write_text(json.dumps(report, indent=2)) + pd.DataFrame(history).to_csv(out_dir / "history.csv", index=False) + np.save(out_dir / "embedding_individual.npy", h["individual"].detach().cpu().numpy()) + torch.save(best_state, out_dir / "model.pt") + print(f"wrote {out_dir}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())