KGIA end-to-end workflow — patient pickle → fine-tune → ranking

Tracking: #124. Written after the first real end-to-end run: smoke run #1 (2026-09-28/29, #439, #114). Every command below is one that run actually used, or the documented successor of one. The run’s results are in `mods_smoke_run_1_writeup.md`; its step-by-step runbook, with every trap it hit, is `mods_smoke_run_1_playbook.md` §5. This document is the reusable skeleton of that runbook for a new patient or a new run.

Read this first. The released KGIA criterion works mechanically but does not help as released. On smoke run #1 its weights stalled in the first epoch, λ had no effect on what it learned, the adjustment could only demote a drug, and 73 % of it flowed through contraindicated_for edges (writeup §4.3). Running this workflow produces a traceable ranking, not a personalized-and-better one. The fix under way is the signed per-relation criterion (#420). When it lands, step 5 changes and this document must be updated with it. (2026-10-05: #420 is planned, not under way. The Q4 2026 plan schedules it for late November at the earliest, and only if a conditioned baseline first beats drug-mean on a cell-line proxy benchmark. Before that, the foundation has to beat the popularity count it currently loses to, 0.031 to 0.021 MRR (#523).)

---

1. What KGIA is

KGIA is the PDR2 personalization mechanism (Oğuztüzün et al., JBI 2026; doi:10.1016/j.jbi.2026.105039). A foundation model (ULTRA/NBFNet) is trained on the knowledge graph as usual. Then the foundation is frozen and a small criterion, PersonalizedRepurposingLoss (losses.py), is fine-tuned. Its only trainable parameters are one weight per disease on the patient’s PRS axis (disease_weights, w) and one per gene on her expression axis (protein_weights, v). At scoring time those weights, multiplied by her own PRS and expression values, add a bias to each candidate’s foundation score.

The architectural shift from PDR1 is where the patient enters:

	PDR1 (KG-Predict lineage)	PDR2 / KGIA
Patient enters via	the graph and embedding table: a patient node whose embedding is learned	the loss: learnable weights times her feature vectors
What training changes	entity embeddings (the model itself)	only w and v; the foundation is frozen
Cost per patient	a retrain	a fine-tune of ~200 scalars
Consequence	personalization is baked into every score	raw score = the foundation’s; personalized = raw + adj

Two facts about the released loss matter for everything downstream, and both are established on our data rather than inferred from the paper:

• The code, not the paper text, is the specification. The released loss L2-normalizes the weights and tanh-squashes the term, and has the opposite sign to the paper’s formula. Implementing the paper text literally is destructive (−92 % MRR). See `pdr2_loss_text_vs_code_note.md` and `pdr2_alt_loss_ab_writeup.md`. --loss_variant selects paper_code (default, the release), paper_text or minimal_patch.
• The criterion is closed-form. Its terms depend only on (weights, fixed patient values), not on the batch, so the optimum is ŵ ≈ −PRS/‖PRS‖ and v̂ ≈ −expr/‖expr‖. Fine-tuning finds that direction; it learns nothing from the graph. Background: `pdr_personalization_writeup.md`.

The patient also exists as a graph node (MAPLE_PATIENT:…, “Mechanism A”, mission spec §5.4). That node is trained into the foundation with her biology edges; it is separate from KGIA and is not what this workflow fine-tunes.

---

2. The patient pickle

The fine-tune reads one pickle, validated on load by utils.load_patient_data:

{
    "prs_values":    [float, ...],   # one per disease on the axis, same order as disease_names
    "expr_values":   [float, ...],   # one per gene on the axis, same order as protein_names
    "disease_names": [str, ...],     # axis v2: MONDO anchor CURIEs, e.g. "MONDO:0005575"
    "protein_names": [str, ...],     # HGNC symbols
}

The loss is positional. Position i of prs_values is multiplied by disease_weights[i], so every pickle that trains or is scored against one criterion must be projected onto the same registered axis (#108). The axis is a pdr2_feature_schema row plus a git-tracked descriptor, backend/pdr2_schemas/kgia_schema_<version>.json; its schema_hash is recorded on each pickle. An axis change is a new schema version, never an in-place edit.

Axis	Diseases	Genes	Status
v1	6, keyed by UMLS CUI, from the 13-entry smoke-test PGS set; no colorectal cancer	84	legacy; cannot hold a released PRS value
v2	109, one per released Gate F axis pick, keyed by MONDO anchor; CRC MONDO:0005575 at position 8	84 (copied from v1)	used by smoke run #1; registered inactive

Build and register an axis (only after a Gate F release exists):

cd backend
python scripts/build_kgia_feature_schema.py --version v3 --dry-run   # print, write nothing
python scripts/build_kgia_feature_schema.py --version v3             # write + register (inactive)

--version is required; an existing version is refused without --overwrite <version>. The build refuses unless the required index disease (MONDO:0005575) is on the axis.

Build the pickle with build_pdr2_input_pickle.py:

cd backend
# Track B: real PRS + real expression
python scripts/build_pdr2_input_pickle.py --patient-id <patient uuid> --source maureen `
    --schema-version v2 --out ../runs/<date>_<run>/<patient>_trackB_v2.pkl
# Track A control: same expression, PRS all zero
python scripts/build_pdr2_input_pickle.py --patient-id <patient uuid> --source synthetic_zeros `
    --schema-version v2 --out ../runs/<date>_<run>/<patient>_trackA_v2.pkl

• Always pass --patient-id. With two patients in mods_patient the default lookup resolves nothing useful.
• Never pass --allow-raw-prs-fallback for a real run. Raw PRS sums span twelve orders of magnitude and saturate the tanh (PRS primer §1.4).
• Each build writes three files: the .pkl, a .meta.json (schema version and hash, value_field, expression transform, reference-filled count) and a .resolution.json (each axis name → its KG CURIEs, §3). Keep all three together; scoring needs the sidecar.
• Arrays are written as Python lists so a NumPy 2 host pickle loads in the container’s NumPy 1.
• Exit codes: 2 no patient / unregistered schema / index disease missing, 3 PRS rows refused, 4 index cancer unresolved, 5 expression not a percentile, 6 --source maureen with no measured PRS position.

Check before using it (smoke check M3): Track B has measured PRS positions > 0, every value is finite, .meta.json says value_field: z_score, and .resolution.json exists. Record each file’s sha256. Then upload all three to gs://BUCKET/patient_data/<date>_<run>/.

Inspect a patient’s vectors in the UI: Research dashboard → PDR Input tab, or GET /api/mods/research/patient-profile/kgia-tensors.

---

3. From clinical reports to KG entities

Two separate paths put clinical data where the model can use it.

(a) Her node in the graph (Mechanism A). Parsed reports (Caris, Guardant, Invitae, RGCC, EHR) land in the MODS tables (mods_list_alterations, mods_gene_expression, …), then mods_patient_kg_service.materialize_patient_kg writes them as maple_patient edges from MAPLE_PATIENT:<slug> to canonical KG entities: alterations to genes, expression, phenotypes, medications and RGCC sensitivities. A foundation trained with include_source: ["maple_patient"] learns that node. See #110 for the MODS pipeline.

• Exclude her drug-response and medication predicates from any foundation you intend to validate against RGCC (#214, #432). Smoke run #1 excluded six: MAPLE:sensitive_to_drug, MAPLE:partially_sensitive_to_drug, MAPLE:resistant_to_drug, MAPLE:tried_drug, MAPLE:current_medication, MAPLE:past_medication. Her biology edges stay in.
• Her variant edges are still gene-level (Phase 8.5): BRAF V600E is indistinguishable from wild-type BRAF in the graph.

(b) Her feature vectors (KGIA). The pickle builder projects her data onto the axis (mods_kgia_schema_service.project_patient_vectors):

• Expression: the Caris WTS TPM percentile for each axis gene from her most recent matching specimen, as a probit z-score, z = Φ⁻¹(p/100) clipped to ±4 (#435). Raw 0–100 percentiles saturated the tanh at step 0 and the gene weights never trained.
• PRS: each disease position reads the one released score that Gate F marked as that anchor’s axis pick (metadata.gate_f.axis_pick), using its normalized z_score (§4).
• An axis position with no value takes the reference fill 0.0 (average risk on the z scale); the count is recorded as prs_reference_filled_count.

The resolution sidecar then maps each axis name to KG entity CURIEs (resolve_disease_kg_curies, resolve_gene_kg_curies). Scoring uses it to find which KG entities are anchors (§6). Read its summary.*_unresolved lists: an unresolved axis name contributes nothing at scoring, silently.

Mapping a new report type or a new disease anchor is an ontology decision. Per CLAUDE.md, propose it as a table with the judgement calls flagged and get approval before committing it (the v2 disease anchors were approved as decision D7 in #439).

---

4. Building PRS values

The PRS feeding KGIA is the released output of the gated PRS pipeline; the arc and every gate are in `prs_primer.md` (status: #100). In outline:

Gate	Produces	Driver
A	dense callset from her gVCF (hom-ref explicit)	worker genotype_dense_callset
B	which catalog scores are computable	worker mods_prs_manifest_build
C / E	raw scores, written quarantined	scripts/prs/run_gate_e.py
D1–D3	panel frequencies, panel scores, z-score and percentile	scripts/prs/run_gate_d_panel.py, run_gate_d.py
F	release: validity='released', disease anchor, axis pick	scripts/prs/release_gate_f.py

Run the Gate D/E drivers inside a worker container (docker exec maple-worker-1 python /app/backend/scripts/prs/<driver>.py …); the exact smoke-run invocations are in playbook §5.1.

Rules that carry to a new patient:

• Only released rows feed a pickle. Every reader filters on validity = 'released'; quarantined rows never reach the model.
• Only normalized values. z_score, else percentile, and one field for the whole vector; a raw sum is refused.
• The panel reference must match the patient’s ancestry. Smoke run #1 used the 667 unrelated EUR samples of pgsc_HGDP+1kGP_v1. A non-European patient needs a different reference subset and different portability caveats.
• No genotype data → Track A. Build with --source synthetic_zeros. The PRS half of the criterion then gets no gradient (its weights move only by weight decay) and only expression personalizes. There is no sanctioned PRS proxy: the 13-entry SMOKE_TEST_PGS_SET is a test fixture, three of its entries are the wrong disease, and it must never size an axis or stand in for a patient’s scores (#244).

---

5. Fine-tune

The upstream run.py --use_ploss -l1 -l2 entry point is not what we run: it has author-local paths. Our entry point is fine_tune_maple.py, which calls training.fine_tune_patient_specific (training.py), reached in the container as MODE=fine_tune. Its flags keep the upstream names: --patient_pickle, -l1/--lambda1_prs, -l2/--lambda2_biomarker, --loss_variant.

Submit through the API, never raw gcloud (mission spec §7.3):

POST /api/mods/research/runs/{foundation_run_id}/submit_finetune
{
  "finetune_run_id":    "<date>_<run>_ft_trackB_l0.01",
  "ckpt_uri":           "gs://BUCKET/runs/<foundation>/train/<checkpoint>.pt",
  "patient_pickle_uri": "gs://BUCKET/patient_data/<date>_<run>/<patient>_trackB_v2.pkl",
  "patient_id":         "<filename label, e.g. maureen>",
  "patient_uuid":       "<patient uuid>",
  "lambda1_prs": 0.01, "lambda2_biomarker": 0.01,
  "loss_variant":       "paper_code",
  "epochs": 20, "bpe": 100,
  "gpu_type":           "A100_80GB",
  "fine_tune_max_batch": 128
}

• gpu_type: "A100_80GB" and fine_tune_max_batch: 128 at PrimeKG scale. The handler defaults to 40 GB, and the batch pre-flight probe runs under no_grad and undersizes memory ~6× (#501).
• patient_id is a filename label; patient_uuid is the id. The rank endpoint reads the UUID (#518).
• The child reuses the foundation’s recorded split and overlay (#433); nothing is re-exported.
• Output: the fine-tune bundle (p<patient_id>_pc_full_state_dict.pth, foundation + criterion) and metrics.json with disease_weights_final, protein_weights_final, the input vectors, the foundation checkpoint’s sha256 and criterion_epoch.

Verify (smoke check M6):

• cos(normalize(w), −PRS) > 0.5 with sign agreement ≥ 0.90, and the same for v against expression. This proves the values reached the right positions. (Smoke run #1: disease cos +0.649 with sign agreement 0.881, accepted; protein +0.652 / 0.952.)
• A Track A control’s disease weights move only by weight decay.
• The run log shows Successfully loaded checkpoint for the foundation.

What not to do, learned on smoke run #1:

• Don’t sweep λ at fine-tune time. Adam cancels the gradient’s λ scale, so λ 0.01, 0.1 and 1.0 learned identical weights (cos 1.000000). λ matters only as a scale at scoring (§6). One fine-tune per axis and track is enough; sweep λ in the scoring step.
• Don’t read much into criterion_epoch. Early stopping keys on the epoch’s training loss, which the frozen foundation dominates; every smoke fine-tune stopped at 12 and restored 6.

---

6. Inference: ranking candidates for a disease

Scoring runs the container’s MODE=score_edges (score_edges.py, logic in edge_scoring.py). The KGIA path is one call:

POST /api/mods/research/kgia/finetune/{finetune_run_id}/rank
{
  "disease_curie":            "MONDO:0005575",
  "patient_pickle_uri":       "gs://…/<patient>_trackB_v2.pkl",
  "resolution_sidecar_uri":   "gs://…/<patient>_trackB_v2.resolution.json",
  "adjustment_target_domain": "drug",
  "lambda1_prs": 0.1, "lambda2_biomarker": 0.1,
  "gpu_type": "A100_40GB"
}

It stages every candidate drug as (drug, biolink:treats, <disease>) (decision D1), scores them in one forward pass per query, stores the ranked list on the run, and exposes it at GET /api/mods/research/kgia/finetune/{run}/ranking (ranking: null means not computed yet, which is different from an empty list). For an arbitrary edge list, use POST /runs/{run}/submit_score_edges with the same personalization fields.

• Use the disease query, not the patient query, on a foundation trained without her MAPLE:*_to_drug edges: it has no such relation, and patient-query rows come back labelled missing_predicate.
• Personalization is all or nothing and never defaulted. Pickle, sidecar and adjustment_target_domain must be given together; λ defaults to the fine-tune’s recorded value. Leave them out and you get raw scores only.

For each candidate e the output carries raw_score, adj, pers_score = raw_score + adj, raw_rank and pers_rank (plus filtered ranks):

adj[e] =   Σ over disease anchors a adjacent to e:  λ1 · ŵ[c(a)] · prs[c(a)]
         + Σ over gene anchors a adjacent to e:     λ2 · v̂[c(a)] · expr[c(a)]

where the anchors are e’s neighbours in the training graph that the sidecar resolves to an axis position, ignoring edge direction and relation type, each adjacency counted once, and ŵ, v̂ are the weights exactly as the loss variant applies them (normalized for paper_code) (personalization_adjustment.py). The adjustment is added after message passing, so the foundation’s graph paths explain raw_score and nothing explains adj except its anchors.

Verify (smoke check M7): at λ 0, pers_score == raw_score on every row; the scores’ .summary.json reports the anchor counts, the target-domain count (if it is 0, a drug CURIE prefix is unknown to classify_entity) and which relation types carried the adjustment.

Validate against RGCC (Maureen only; smoke check M8) with the per-draw grid. Never pool the two draws.

cd backend
python scripts/score_pdr2_edges.py build-edges --disease-curie MONDO:0005575 `
    --out ../runs/<date>_<run>/rgcc_disease_edges.csv
python scripts/compute_metrics_rgcc_pdr2.py --query disease --disease-curie MONDO:0005575 `
    --scores <scores_csv_uri from the ranking row> `
    --edges ../runs/<date>_<run>/rgcc_disease_edges.csv `
    --run <fine-tune run id> --out ../runs/<date>_<run>/<fine-tune>/rgcc_metrics.json --store

--store writes mods_research_validation_result rows (validation_type='rgcc_per_draw').

---

7. Reading the output

What the ranked list is. raw_rank orders drugs by how strongly the foundation’s learned graph geometry supports (drug, treats, disease). pers_rank is that order after her PRS-and-expression bias. Neither is evidence that a drug works for her, is safe with her current regimen, or is tolerable given her labs (mission spec §1.4). Every list is a set of hypotheses for literature review and, only after that, for her clinicians.

Questions to ask of any list before anyone reads it clinically:

1. Is the top recall or repurposing? On smoke run #1, ranks 1–2 were the only two drugs whose treats-CRC edge sat in the training split. A drug with a known indication in the training graph ranking first is the model remembering. Mark each candidate’s indication status and whether its edge was in training.
2. What moved, and why? For each drug whose pers_rank differs from raw_rank, name the anchors (diseases and genes) that carried its adj. Under the released criterion a move is almost always a demotion, and often through contraindicated_for; that is the known defect, not a finding.
3. Is it in the target domain? Therapy and regimen entities from CIViC and OncoKB can rank high (her own regimen did) but sit outside the drug domain, so personalization cannot touch them. Report them, and say they are unadjusted.
4. Which foundation, which λ, which axis? A number without its run id, λ, schema hash and pickle sha256 cannot be traced (goal G4). Record them where the list is shown.

Graph paths. Explaining why a drug scored high (direct, mechanistic, analogy, personal or hub paths) uses BaseNBFNet.visualize(); the method and its positive control are in the playbook §5.4 step 4. There is no container mode for it yet.

Where it surfaces in the MODS UI. Research dashboard → KGIA Ranking tab (KgiaRankingPanel) lists fine-tune runs with their λ and shows a run’s stored ranking: rank, held-out rank, drug, score and whether the drug was seen in training, with a held-out-only filter. As of 2026-09-30 it has two limits worth knowing (#560):

• Its Rank button posts an empty body, so it runs the default patient query with no personalization. On a foundation trained without her drug edges that query has nothing to score. Submit the disease query with the personalization fields through the API (§6); the panel then displays the stored result.
• It renders the raw score and rank only. The stored ranking rows also carry adj, pers_score, pers_candidate_rank and in_target_domain (mods_kgia_finetune_service), and the full per-edge output is the run’s scores.csv (scores_csv_uri on the ranking row); read the personalized order from either.

The PDR Input tab shows the vectors that went in. Label any page that shows a smoke or test ranking “machinery validation, not a recommendation”.

---

Checklist for a new run

1. Released PRS exists for the patient (or you have chosen Track A deliberately).
2. Axis registered; its schema_hash noted.
3. Track B and Track A pickles built, M3 checked, sha256 recorded, uploaded with sidecars.
4. Foundation chosen, its export excluding the predicates you will validate against.
5. One fine-tune per track, on A100-80GB, batch 128; M6 checked.
6. Ranking per fine-tune at the λ values you want; M7 checked.
7. Per-draw validation stored, if the patient has labels.
8. Top candidates annotated per §7, with run ids, image digest, schema hash and pickle sha256 recorded on the tracking issue.