diff --git a/.gitignore b/.gitignore index de6497d..ba5e136 100644 --- a/.gitignore +++ b/.gitignore @@ -118,3 +118,8 @@ src/ops_model/vesuvius/ *.csv *.svg slurm_logs/ + +# Keep the set-classifier interpretability example configs (placeholder paths) +!src/ops_model/interpretability/classifier/configs/ +!src/ops_model/interpretability/classifier/configs/*.yaml +!src/ops_model/interpretability/classifier/configs/*.csv diff --git a/pyproject.toml b/pyproject.toml index ed3724a..44dcea8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,6 +50,15 @@ dependencies = [ # https://peps.python.org/pep-0621/#dependencies-optional-dependencies [project.optional-dependencies] models = ["transformers<4.40"] +# Set-classifier interpretability subpackage (interpretability/classifier): train / eval / +# SHAP over cell embeddings. Kept optional so it doesn't bloat the base install. +classifier = [ + "hydra-core", + "omegaconf", + "pyarrow", + "tqdm", + "matplotlib", +] test = ["pytest>=6.0", "pytest-cov"] dev = [ "ipykernel", diff --git a/src/ops_model/interpretability/__init__.py b/src/ops_model/interpretability/__init__.py new file mode 100644 index 0000000..155a3e5 --- /dev/null +++ b/src/ops_model/interpretability/__init__.py @@ -0,0 +1 @@ +"""Interpretability tools for OPS models.""" diff --git a/src/ops_model/interpretability/classifier/__init__.py b/src/ops_model/interpretability/classifier/__init__.py new file mode 100644 index 0000000..2b429fc --- /dev/null +++ b/src/ops_model/interpretability/classifier/__init__.py @@ -0,0 +1,10 @@ +"""Set-classifier interpretability subpackage. + +Train / evaluate / explain a permutation-invariant *set classifier* that predicts a +gene knockout (or a protein-complex / pathway label) from a *set* of single-cell +embeddings. Runs standalone from an embedding parquet (+ an optional gene->label +metadata CSV). Entry points: ``train.py`` (hydra ``--config-name``, see ``configs/``), +``eval.py`` (accuracy vs cells-per-set), and ``score.py`` (per-cell scores, argparse). + +Requires the optional dependencies: ``pip install ops_model[classifier]``. +""" diff --git a/src/ops_model/interpretability/classifier/configs/ebi_complexes.csv b/src/ops_model/interpretability/classifier/configs/ebi_complexes.csv new file mode 100644 index 0000000..52cdbda --- /dev/null +++ b/src/ops_model/interpretability/classifier/configs/ebi_complexes.csv @@ -0,0 +1,312 @@ +,gene_name,pathway +0,RAD21,"Nuclear mitotic cohesin complex, STAG1 variant" +1,SMC1A,"Nuclear mitotic cohesin complex, STAG1 variant" +2,SMC3,"Nuclear mitotic cohesin complex, STAG1 variant" +3,CCT3,Chaperonin-containing T-complex +4,CCT4,Chaperonin-containing T-complex +5,CCT6A,Chaperonin-containing T-complex +6,CCT7,Chaperonin-containing T-complex +7,TCP1,Chaperonin-containing T-complex +8,SNRPD1,Sm complex +9,SNRPD2,Sm complex +10,EIF3A,Eukaryotic translation initiation factor 3 complex +11,EIF3B,Eukaryotic translation initiation factor 3 complex +12,EIF3L,Eukaryotic translation initiation factor 3 complex +13,EIF3M,Eukaryotic translation initiation factor 3 complex +14,COX7B,Mitochondrial respiratory chain complex IV +15,COX7C,Mitochondrial respiratory chain complex IV +16,DNAJC19,"TIM23 mitochondrial inner membrane pre-sequence translocase complex, TIM17A variant" +17,GRPEL1,"TIM23 mitochondrial inner membrane pre-sequence translocase complex, TIM17A variant" +18,HSPA9,"TIM23 mitochondrial inner membrane pre-sequence translocase complex, TIM17A variant" +19,TIMM23,"TIM23 mitochondrial inner membrane pre-sequence translocase complex, TIM17A variant" +20,TIMM44,"TIM23 mitochondrial inner membrane pre-sequence translocase complex, TIM17A variant" +21,TSC1,TSC1-TSC2 complex +22,TSC2,TSC1-TSC2 complex +23,TELO2,TTT complex +24,TTI1,TTT complex +25,TTI2,TTT complex +26,ATP5F1B,Mitochondrial proton-transporting ATP synthase complex +27,ATP5ME,Mitochondrial proton-transporting ATP synthase complex +28,TRAPPC2L,"TRAPP II complex, TRAPPC2 variant" +29,TRAPPC4,"TRAPP II complex, TRAPPC2 variant" +30,RPL15,60S cytosolic large ribosomal subunit +31,RPL18,60S cytosolic large ribosomal subunit +32,RPL23,60S cytosolic large ribosomal subunit +33,RPL27A,60S cytosolic large ribosomal subunit +34,RPL30,60S cytosolic large ribosomal subunit +35,RPL32,60S cytosolic large ribosomal subunit +36,RPL34,60S cytosolic large ribosomal subunit +37,RPL35,60S cytosolic large ribosomal subunit +38,RPL37A,60S cytosolic large ribosomal subunit +39,RPL41,60S cytosolic large ribosomal subunit +40,RPL5,60S cytosolic large ribosomal subunit +41,RPL9,60S cytosolic large ribosomal subunit +42,RPLP1,60S cytosolic large ribosomal subunit +43,RPLP2,60S cytosolic large ribosomal subunit +44,UBA52,60S cytosolic large ribosomal subunit +45,COPA,"COPI vesicle coat complex, COPG1-COPZ1 variant" +46,COPB1,"COPI vesicle coat complex, COPG1-COPZ1 variant" +47,COPB2,"COPI vesicle coat complex, COPG1-COPZ1 variant" +48,COPZ1,"COPI vesicle coat complex, COPG1-COPZ1 variant" +49,VPS26A,"Retromer complex, VPS26A variant" +50,VPS29,"Retromer complex, VPS26A variant" +51,FAU,40S cytosolic small ribosomal subunit +52,RPS10,40S cytosolic small ribosomal subunit +53,RPS14,40S cytosolic small ribosomal subunit +54,RPS16,40S cytosolic small ribosomal subunit +55,RPS25,40S cytosolic small ribosomal subunit +56,RPS27A,40S cytosolic small ribosomal subunit +57,RPS28,40S cytosolic small ribosomal subunit +58,RPS29,40S cytosolic small ribosomal subunit +59,RPS3A,40S cytosolic small ribosomal subunit +60,RPS5,40S cytosolic small ribosomal subunit +61,RPS6,40S cytosolic small ribosomal subunit +62,RPS8,40S cytosolic small ribosomal subunit +63,RPS9,40S cytosolic small ribosomal subunit +64,RPSA,40S cytosolic small ribosomal subunit +65,ASNA1,GET complex +66,WRB,GET complex +67,POLR2B,DNA-directed RNA polymerase II complex +68,POLR2C,DNA-directed RNA polymerase II complex +69,POLR2D,DNA-directed RNA polymerase II complex +70,POLR2F,DNA-directed RNA polymerase II complex +71,POLR2H,DNA-directed RNA polymerase II complex +72,POLR2I,DNA-directed RNA polymerase II complex +73,EXOC4,"Exocyst, EXOC6 variant" +74,EXOC5,"Exocyst, EXOC6 variant" +75,EXOC7,"Exocyst, EXOC6 variant" +76,BORCS7,BORC complex +77,BORCS8,BORC complex +78,AP1M1,"Ubiquitous AP-1 Adaptor complex, sigma1a variant" +79,AP1S1,"Ubiquitous AP-1 Adaptor complex, sigma1a variant" +80,AP3M1,"Ubiquitous AP-3 Adaptor complex, sigma3a variant" +81,AP3S1,"Ubiquitous AP-3 Adaptor complex, sigma3a variant" +82,AP2A1,"AP-2 Adaptor complex, alpha1 variant" +83,AP2M1,"AP-2 Adaptor complex, alpha1 variant" +84,AP2S1,"AP-2 Adaptor complex, alpha1 variant" +85,MRPL20,39S mitochondrial large ribosomal subunit +86,MRPL34,39S mitochondrial large ribosomal subunit +87,MRPL39,39S mitochondrial large ribosomal subunit +88,MRPL40,39S mitochondrial large ribosomal subunit +89,MRPL42,39S mitochondrial large ribosomal subunit +90,MRPL52,39S mitochondrial large ribosomal subunit +91,LAMTOR1,Ragulator complex +92,LAMTOR2,Ragulator complex +93,LAMTOR3,Ragulator complex +94,LAMTOR4,Ragulator complex +95,NUP214,Nuclear pore complex +96,NUP37,Nuclear pore complex +97,NUP54,Nuclear pore complex +98,NUP98,Nuclear pore complex +99,RAE1,Nuclear pore complex +100,CHMP1A,ESCRT-III complex +101,CHMP2A,ESCRT-III complex +102,CHMP4B,ESCRT-III complex +103,CHMP4C,ESCRT-III complex +104,CHMP5,ESCRT-III complex +105,CHMP6,ESCRT-III complex +106,CHMP7,ESCRT-III complex +107,DARS,Multiaminoacyl-tRNA synthetase complex +108,LARS,Multiaminoacyl-tRNA synthetase complex +109,MARS,Multiaminoacyl-tRNA synthetase complex +110,QARS,Multiaminoacyl-tRNA synthetase complex +111,RARS,Multiaminoacyl-tRNA synthetase complex +112,ATP6V0A1,"Vacuolar proton translocating ATPase complex, ATP6V0A1 variant" +113,ATP6V1H,"Vacuolar proton translocating ATPase complex, ATP6V0A1 variant" +114,DDX39B,"TREX transcription-export complex, DX39B variant" +115,THOC3,"TREX transcription-export complex, DX39B variant" +116,THOC7,"TREX transcription-export complex, DX39B variant" +117,ARPC2,"Actin-related protein 2/3 complex, ARPC1A-ACTR3B-ARPC5 variant" +118,ARPC3,"Actin-related protein 2/3 complex, ARPC1A-ACTR3B-ARPC5 variant" +119,ARPC4,"Actin-related protein 2/3 complex, ARPC1A-ACTR3B-ARPC5 variant" +120,AATF,Small ribosomal subunit processome +121,BMS1,Small ribosomal subunit processome +122,NOL6,Small ribosomal subunit processome +123,RCL1,Small ribosomal subunit processome +124,RRP7A,Small ribosomal subunit processome +125,UTP14A,Small ribosomal subunit processome +126,CNOT1,"CCR4-NOT mRNA deadenylase complex, CNOT6L-CNOT7 variant" +127,CNOT10,"CCR4-NOT mRNA deadenylase complex, CNOT6L-CNOT7 variant" +128,CNOT11,"CCR4-NOT mRNA deadenylase complex, CNOT6L-CNOT7 variant" +129,CNOT2,"CCR4-NOT mRNA deadenylase complex, CNOT6L-CNOT7 variant" +130,CNOT4,"CCR4-NOT mRNA deadenylase complex, CNOT6L-CNOT7 variant" +131,TIMELESS,Replication fork protection complex +132,TIPIN,Replication fork protection complex +133,DDX46,U2 small nuclear ribonucleoprotein complex +134,SNRPB2,U2 small nuclear ribonucleoprotein complex +135,GINS1,GINS complex +136,GINS2,GINS complex +137,KANSL1,NSL histone acetyltransferase complex +138,MCRS1,NSL histone acetyltransferase complex +139,OGT,NSL histone acetyltransferase complex +140,MTOR,mTORC1 complex +141,RPTOR,mTORC1 complex +142,GTF2A1,General transcription factor complex TFIIA +143,GTF2A2,General transcription factor complex TFIIA +144,EXOSC3,"Nucleolar exosome complex, EXOSC10 variant" +145,EXOSC5,"Nucleolar exosome complex, EXOSC10 variant" +146,EXOSC6,"Nucleolar exosome complex, EXOSC10 variant" +147,EXOSC8,"Nucleolar exosome complex, EXOSC10 variant" +148,SRPRA,Signal recognition particle receptor complex +149,SRPRB,Signal recognition particle receptor complex +150,LIN37,Myb-MuvB transcriptional activation complex +151,LIN9,Myb-MuvB transcriptional activation complex +152,GTF3C1,General transcription factor TFIIIC complex +153,GTF3C4,General transcription factor TFIIIC complex +154,GTF3C6,General transcription factor TFIIIC complex +155,CD3EAP,DNA-directed RNA polymerase I complex +156,POLR1B,DNA-directed RNA polymerase I complex +157,POLR1C,DNA-directed RNA polymerase I complex +158,POLR1E,DNA-directed RNA polymerase I complex +159,ZNRD1,DNA-directed RNA polymerase I complex +160,SNRNP70,U1 small nuclear ribonucleoprotein complex +161,SNRPC,U1 small nuclear ribonucleoprotein complex +162,POLR3A,"DNA-directed RNA polymerase III complex, POLR3G variant" +163,POLR3B,"DNA-directed RNA polymerase III complex, POLR3G variant" +164,POLR3E,"DNA-directed RNA polymerase III complex, POLR3G variant" +165,POLR3F,"DNA-directed RNA polymerase III complex, POLR3G variant" +166,GTF2B,General transcription factor TFIIB-TBP complex +167,TBP,General transcription factor TFIIB-TBP complex +168,AURKB,"Chromosomal passenger complex, AURKB variant" +169,INCENP,"Chromosomal passenger complex, AURKB variant" +170,TERF1,Shelterin complex +171,TERF2IP,Shelterin complex +172,MED18,Core mediator complex +173,MED21,Core mediator complex +174,MED6,Core mediator complex +175,SEL1L,Luminal surveillance complex +176,SYVN1,Luminal surveillance complex +177,FTH1,"FTH1-FTL, Ferritin complex" +178,FTL,"FTH1-FTL, Ferritin complex" +179,DAD1,"Oligosaccharyltransferase complex B, MAGT1 variant" +180,DDOST,"Oligosaccharyltransferase complex B, MAGT1 variant" +181,OST4,"Oligosaccharyltransferase complex B, MAGT1 variant" +182,STT3B,"Oligosaccharyltransferase complex B, MAGT1 variant" +183,COPS2,COP9 signalosome variant 1 +184,COPS3,COP9 signalosome variant 1 +185,COPS4,COP9 signalosome variant 1 +186,COPS5,COP9 signalosome variant 1 +187,COPS6,COP9 signalosome variant 1 +188,COPS8,COP9 signalosome variant 1 +189,GPS1,COP9 signalosome variant 1 +190,BLOC1S2,BLOC-1 complex +191,BLOC1S6,BLOC-1 complex +192,SNRPD3,U7 small nuclear ribonucleoprotein complex +193,SNRPG,U7 small nuclear ribonucleoprotein complex +194,FARSA,Phenylalanyl-tRNA synthetase complex +195,FARSB,Phenylalanyl-tRNA synthetase complex +196,COMMD4,Commander complex +197,COMMD5,Commander complex +198,CUL2,LRR1-Elongin C-Elongin B E3 ubiquitin ligase complex +199,ELOC,LRR1-Elongin C-Elongin B E3 ubiquitin ligase complex +200,SF3B3,SF3B complex +201,SF3B5,SF3B complex +202,SRP19,Signal recognition particle +203,SRP54,Signal recognition particle +204,SRP68,Signal recognition particle +205,SRP72,Signal recognition particle +206,EIF4A1,"Eukaryotic translation initiation factor 4F, EIF4A1 and EIF4G1 variant" +207,EIF4G1,"Eukaryotic translation initiation factor 4F, EIF4A1 and EIF4G1 variant" +208,TBL3,UTP-B complex +209,WDR36,UTP-B complex +210,CPSF2,Histone pre-RNA core cleavage complex +211,CPSF3,Histone pre-RNA core cleavage complex +212,HAUS4,HAUS complex +213,HAUS5,HAUS complex +214,HAUS6,HAUS complex +215,SPCS2,"Signal peptidase complex, SEC11A variant" +216,SPCS3,"Signal peptidase complex, SEC11A variant" +217,POLA1,DNA polymerase alpha:primase complex +218,POLA2,DNA polymerase alpha:primase complex +219,PRIM1,DNA polymerase alpha:primase complex +220,POLE,DNA polymerase epsilon complex +221,POLE2,DNA polymerase epsilon complex +222,POLE4,DNA polymerase epsilon complex +223,MCM3,MCM complex +224,MCM5,MCM complex +225,MCM7,MCM complex +226,IDH3A,Mitochondrial isocitrate dehydrogenase complex (NAD+) +227,IDH3B,Mitochondrial isocitrate dehydrogenase complex (NAD+) +228,CCNH,Cyclin-dependent protein kinase-activating kinase complex +229,CDK7,Cyclin-dependent protein kinase-activating kinase complex +230,EMC1,"Endoplasmic reticulum membrane complex, EMC8 variant" +231,EMC2,"Endoplasmic reticulum membrane complex, EMC8 variant" +232,EMC3,"Endoplasmic reticulum membrane complex, EMC8 variant" +233,COQ4,CoQ biosynthetic complex +234,COQ7,CoQ biosynthetic complex +235,COG3,COG tethering complex +236,COG4,COG tethering complex +237,COG6,COG tethering complex +238,COG7,COG tethering complex +239,COG8,COG tethering complex +240,FANCC,Fanconi anemia ubiquitin ligase complex +241,FANCG,Fanconi anemia ubiquitin ligase complex +242,CENPC,Kinetochore CCAN complex +243,CENPH,Kinetochore CCAN complex +244,CENPW,Kinetochore CCAN complex +245,KRT18,Keratin-8 - Keratin-18 dimer complex +246,KRT8,Keratin-8 - Keratin-18 dimer complex +247,PHB,Prohibitin complex +248,PHB2,Prohibitin complex +249,ELP2,Elongator holoenzyme complex +250,ELP4,Elongator holoenzyme complex +251,INO80D,INO80 chromatin remodeling complex +252,INO80E,INO80 chromatin remodeling complex +253,CHD4,MBD2/NuRD nucleosome remodeling and deacetylase complex +254,MTA2,MBD2/NuRD nucleosome remodeling and deacetylase complex +255,CPSF6,"mRNA cleavage factor I(m) complex, CPSF6 variant" +256,NUDT21,"mRNA cleavage factor I(m) complex, CPSF6 variant" +257,BRD8,NuA4 histone acetyltransferase complex +258,MBTD1,NuA4 histone acetyltransferase complex +259,MRGBP,NuA4 histone acetyltransferase complex +260,SMARCC1,"SWI/SNF ATP-dependent chromatin remodeling complex, ACTL6A-ARID1A-SMARCA2 variant" +261,SMARCD1,"SWI/SNF ATP-dependent chromatin remodeling complex, ACTL6A-ARID1A-SMARCA2 variant" +262,PPIE,Intron-binding complex +263,XAB2,Intron-binding complex +264,SEC61A1,"SEC61 protein-conducting channel complex, SEC1A1 variant" +265,SEC61B,"SEC61 protein-conducting channel complex, SEC1A1 variant" +266,SEC61G,"SEC61 protein-conducting channel complex, SEC1A1 variant" +267,ADRM1,19S proteasome regulatory complex +268,PSMC2,19S proteasome regulatory complex +269,PSMC6,19S proteasome regulatory complex +270,PSMD11,19S proteasome regulatory complex +271,PSMD14,19S proteasome regulatory complex +272,PSMD2,19S proteasome regulatory complex +273,PSMD4,19S proteasome regulatory complex +274,DCTN4,Dynactin complex +275,DCTN5,Dynactin complex +276,DYNC1H1,"Dynein-1 complex, variant 2" +277,DYNC1I2,"Dynein-1 complex, variant 2" +278,DYNC1LI1,"Dynein-1 complex, variant 2" +279,DYNLL1,"Dynein-1 complex, variant 2" +280,NOP56,"Box C/D snoRNA-Guided RNP methyltransferase complex, FBLL1 variant" +281,NOP58,"Box C/D snoRNA-Guided RNP methyltransferase complex, FBLL1 variant" +282,DENR,"DENR-MCTS1, translation reinitiation complex" +283,MCTS1,"DENR-MCTS1, translation reinitiation complex" +284,EFTUD2,"U5 small nuclear ribonucleoprotein complex, AAR2 variant" +285,SNRNP200,"U5 small nuclear ribonucleoprotein complex, AAR2 variant" +286,LSM4,LSM2-8 complex +287,LSM5,LSM2-8 complex +288,LSM6,LSM2-8 complex +289,CABIN1,"HIR histone chaperone complex, UBN1 variant" +290,UBN1,"HIR histone chaperone complex, UBN1 variant" +291,EIF2S1,Eukaryotic translation initiation factor 2 complex +292,EIF2S2,Eukaryotic translation initiation factor 2 complex +293,DDRGK1,UFM1 ribosome E3 ligase complex +294,UFL1,UFM1 ribosome E3 ligase complex +295,EIF2B2,Eukaryotic translation initiation factor 2B complex +296,EIF2B3,Eukaryotic translation initiation factor 2B complex +297,EIF2B4,Eukaryotic translation initiation factor 2B complex +298,EIF2B5,Eukaryotic translation initiation factor 2B complex +299,INTS1,Integrator core complex +300,INTS10,Integrator core complex +301,INTS11,Integrator core complex +302,INTS2,Integrator core complex +303,INTS5,Integrator core complex +304,INTS7,Integrator core complex +305,CD2BP2,U5 small nuclear ribonucleoprotein complete complex +306,PRPF6,U5 small nuclear ribonucleoprotein complete complex +307,SNRNP40,U5 small nuclear ribonucleoprotein complete complex +308,BUD31,"Intron Lariat Spliceosome, type 1 complex" +309,RBM22,"Intron Lariat Spliceosome, type 1 complex" +310,NTC,NTC diff --git a/src/ops_model/interpretability/classifier/configs/eval_set_classifier.yaml b/src/ops_model/interpretability/classifier/configs/eval_set_classifier.yaml new file mode 100644 index 0000000..d7b783f --- /dev/null +++ b/src/ops_model/interpretability/classifier/configs/eval_set_classifier.yaml @@ -0,0 +1,100 @@ +# Eval: accuracy vs. N cells with bootstrap-style error bars (Hydra → eval_set_classifier.py) +# +# Usage (mixed-channel checkpoint + val dump from train_set_classifier dump_val_dir): +# python katamari/evals/image_verifier/eval_set_classifier.py \\ +# checkpoint_path=/path/to/best_set_classifier.pt \\ +# val_dump_dir=/path/to/val_ops \\ +# output_plot_path=/path/to/n_cells_accuracy.png \\ +# n_cells_list=[100,250,500,1000] \\ +# n_repetitions=30 \\ +# device=cuda +# +# Per-channel checkpoint (PerturbationDataset + parquet): set mixed_channels_mode=false +# and fill data.* to match training (same parquet, val_fraction, seed). + +seed: 42 +device: cuda + +# Path to the trained checkpoint to evaluate. +checkpoint_path: null + +# Required when mixed_channels_mode=true: directory with metadata.pt + one .pt per gene +# val_dump_dir: /path/to/val_ops_zstdcontrol_reprocess_update_fluorescent_no_auto2 +# val_dump_dir: /path/to/val_ops_zstdcontrol_cdino +val_dump_dir: null # set to your val dump dir (user-specific path removed) + +# Output figure path (.png / .pdf). A sidecar JSON (same basename) stores per-rep accuracies. +output_plot_path: null # set output figure path (user-specific path removed) + +plot_title: "Validation accuracy vs. N cells per set" + +# Log-scaled x-axis (N cells); set false for linear +plot_x_log: true + +# List of N_CELLS values (x-axis). +n_cells_list: +# - 1 +# - 2 +# - 5 +# - 10 +# - 20 +# - 50 +- 10 +- 20 +- 50 +- 100 +- 200 +- 500 +- 1000 +- 2000 +- 5000 + # - 100 + # - 200 + # - 500 + # - 1000 + +# Number of independent full-val passes per N (different RNG seeds). Error bars = SEM across these. +n_repetitions: 50 + +batch_size: 64 +num_workers: 4 + +# tqdm: outer bar over n_cells_list, middle bar over repetitions, inner bar over val batches +show_progress: true + +# Must match the checkpoint (mixed: MixedChannelClassifier + val dump; per-channel: SetClassifier + parquet). +mixed_channels_mode: true + +# Restrict validation to these channel names only (null = use all channels). +# only_channels: ["Phase2D"] + +# split_channels (mixed_channels_mode=true only): +# false (default): all channels in the allowset are fed to the model together per set. +# true: each channel in the allowset is evaluated independently — the model sees a +# single-channel set with `n_cells` cells of just that one channel. The per-class +# CSV gains a `channel_name` column with one row per (channel, n_cells, gene). +# The allowset is `only_channels` if set, else all checkpoint channels present in the val export. +split_channels: false + +# Label map: CSV mapping genes to class labels (e.g. pathway categories). +# If the checkpoint was trained with a label map, it is auto-detected from the checkpoint. +# Set label_map_path to override or supply a label map for a gene-level checkpoint. +# label_map_path: /path/to/label_map.csv +# label_map_gene_col: gene_name +# label_map_label_col: pathway + +# Per-channel mode only (ignored when val_dump_dir is used — val cells are fixed in the dump). +val_fraction: 0.2 +n_cells_per_set_overrides: {} + +per_class_output_file: null # set per-class output path (user-specific path removed) + +data: + parquet_entries: null # list of {path, exclude_experiments?, exclude_fluorescent_experiments?} + max_row_groups: null + max_cells_per_group: null + max_genes: null + max_channels: null + exclude_channel_names: null + min_cells_per_group: null + min_cells_drop_val: false diff --git a/src/ops_model/interpretability/classifier/configs/train_set_classifier_fluor_1K.yaml b/src/ops_model/interpretability/classifier/configs/train_set_classifier_fluor_1K.yaml new file mode 100644 index 0000000..5ad639b --- /dev/null +++ b/src/ops_model/interpretability/classifier/configs/train_set_classifier_fluor_1K.yaml @@ -0,0 +1,70 @@ +# Config for train_set_classifier.py +# Combined fluorescent set-classifier: phenotyping fluorescent markers + Cell +# Painting (ops0094) + 4i (ops0144). 1001 gene + NTC classification. +# null channel sampling (each set uses ALL available channels), 100 cells/set. +# CELL-STRATIFIED train/val split (no cross-cell leakage). Creates the shared dump. + +name: "celldino-fluor-cp-4i-paper-v2-cellstrat-null" + +seed: 42 +device: "cuda" # auto-detect + +# ---- Data ---- +data: + parquet_entries: + - path: /path/to/embeddings/celldino_paper_phenotyping_v2_160.parquet + - path: /path/to/embeddings/celldino_paper_v2_cp_160.parquet + - path: /path/to/embeddings/celldino_paper_v2_4i_160.parquet + max_row_groups: null + max_cells_per_group: null + max_genes: null + max_channels: null + # Drop phase + unannotated channels; keep all fluorescent markers (pheno + CP + 4i). + exclude_channel_names: ["Phase2D", "no label_no label"] + min_cells_per_group: null + min_cells_drop_val: false + z_standardize_control_only: true # z-standardize per (channel x experiment) using NTC (control) cell stats + # Split train/val by unique cell (experiment, well, segmentation_id) so all + # markers/modalities of a cell stay on one side (no leakage). + cell_stratify: true + +# ---- Training ---- +n_cells_per_set: 100 +channels_per_set: null # all markers per set; set [1, 2, 5, 10, 20, null] to also train on smaller random marker subsets +train_n_cell_sets_per_gene: 32 +val_n_cell_sets_per_gene: 5 +channel_drop_fraction: 0.0 +mixed_channels_mode: true +val_n_cells_per_set: [500, 200] +phase2d_val: false +eval_every: 1 +batch_size: 64 +num_epochs: 100 +learning_rate: 1e-4 +weight_decay: 0.0 +max_grad_norm: null +warmup_epochs: 20 +val_fraction: 0.2 +num_workers: 4 +save_path: best_set_classifier_paper_v2_fluor_cp_4i_null.pt +# Dump the cell-stratified processed dataset once; the sweep run reuses it. +dump_train_dir: /path/to/dumps/paper_v2_fluor_cp_4i_cellstrat/train +dump_val_dir: /path/to/dumps/paper_v2_fluor_cp_4i_cellstrat/val + +# ---- Model ---- +model: + d_model: 512 + n_heads: 4 + d_ff: null # null = 4 * d_model + dropout: 0.0 + + # Per-channel set transformer (pools N_cells → 1 vector per channel) + n_layers_cell: 2 + n_inducing_cell: 32 + + channel_conditioning: "concat" # "add" | "concat" | "adaln" | "adaln-token" | "none" + cosine_classifier: True # use cosine-similarity head with learned temperature + + # Cross-channel set transformer (full attention, pools channel vectors → 1 final vector) + +# ---- Logging ---- diff --git a/src/ops_model/interpretability/classifier/configs/train_set_classifier_fluor_ebi.yaml b/src/ops_model/interpretability/classifier/configs/train_set_classifier_fluor_ebi.yaml new file mode 100644 index 0000000..b4f9b8a --- /dev/null +++ b/src/ops_model/interpretability/classifier/configs/train_set_classifier_fluor_ebi.yaml @@ -0,0 +1,66 @@ +# Config for train_set_classifier.py +# fluor set-classifier with EBI complex labels +# (EBI only; non-EBI genes dropped). +# Reuses the pre-dumped dataset (no reprocessing). + +name: "celldino-fluor-cp-4i-ebi-only-null" + +seed: 42 +device: "cuda" # auto-detect + +# ---- Data ---- +data: + parquet_entries: + - path: /path/to/embeddings/celldino_paper_phenotyping_v2_160.parquet + - path: /path/to/embeddings/celldino_paper_v2_cp_160.parquet + - path: /path/to/embeddings/celldino_paper_v2_4i_160.parquet + max_row_groups: null + max_cells_per_group: null + max_genes: null + max_channels: null + exclude_channel_names: ["Phase2D", "no label_no label"] + min_cells_per_group: null + min_cells_drop_val: false + z_standardize_control_only: true + cell_stratify: true + # EBI complex labels: genes in the map -> their complex; NTC -> NTC. + label_map_path: src/ops_model/interpretability/classifier/configs/ebi_complexes.csv + label_map_gene_col: gene_name + label_map_label_col: pathway + label_map_fallback_to_gene: false # non-EBI genes are dropped + +# ---- Training ---- +n_cells_per_set: 100 +channels_per_set: null # all markers per set; set [1, 2, 5, 10, 20, null] to also train on smaller random marker subsets +train_n_cell_sets_per_gene: 32 +val_n_cell_sets_per_gene: 5 +channel_drop_fraction: 0.0 +mixed_channels_mode: true +val_n_cells_per_set: [500, 200] +phase2d_val: false +eval_every: 1 +batch_size: 64 +num_epochs: 100 +learning_rate: 1e-4 +weight_decay: 0.0 +max_grad_norm: null +warmup_epochs: 20 +val_fraction: 0.2 +num_workers: 4 +save_path: best_set_classifier_fluor_cp_4i_ebionly_null.pt +# Reuse the pre-dumped dataset (skip parquet processing). +load_train_dir: /path/to/dumps/paper_v2_fluor_cp_4i_cellstrat/train +load_val_dir: /path/to/dumps/paper_v2_fluor_cp_4i_cellstrat/val + +# ---- Model ---- +model: + d_model: 512 + n_heads: 4 + d_ff: null # null = 4 * d_model + dropout: 0.0 + n_layers_cell: 2 + n_inducing_cell: 32 + channel_conditioning: "concat" + cosine_classifier: True + +# ---- Logging ---- diff --git a/src/ops_model/interpretability/classifier/configs/train_set_classifier_phase_1K.yaml b/src/ops_model/interpretability/classifier/configs/train_set_classifier_phase_1K.yaml new file mode 100644 index 0000000..2ad82ea --- /dev/null +++ b/src/ops_model/interpretability/classifier/configs/train_set_classifier_phase_1K.yaml @@ -0,0 +1,63 @@ +# Config for train_set_classifier.py +# Phase-only (Phase2D) set-classifier: 1001 gene + NTC classification. +# Embeddings from extract_embeddings_paper_v2 (CellDINO ViT-L/16, 160x160 crops). + +name: "celldino-phase-paper-v2" + +seed: 42 +device: "cuda" # auto-detect + +# ---- Data ---- +data: + parquet_entries: + - path: /path/to/embeddings/celldino_paper_phenotyping_v2_160.parquet + max_row_groups: null # set to limit row groups loaded (e.g. 10 for quick debug) + max_cells_per_group: null # cap cells per (gene, channel) to limit memory + max_genes: null # keep only top-N genes by cell count (null = all) + max_channels: 1 # keep only top-N channels by cell count (Phase2D dominates: present in all wells) + # Exact channel labels as in _build_channel_label / printed "Cells per channel" (not raw parquet name if annotation overrides) + exclude_channel_names: null # e.g. [Phase2D] to drop that channel entirely + min_cells_per_group: null # drop (gene, channel) pairs with fewer cells (null = no filter) + min_cells_drop_val: false # also drop from val when min_cells_per_group filters (default: train only) + z_standardize_control_only: true # z-standardize per (channel x experiment) using NTC (control) cell stats + +# ---- Training ---- +n_cells_per_set: 100 +train_n_cell_sets_per_gene: 32 # how many cell-set draws per gene per train epoch +val_n_cell_sets_per_gene: 5 # how many cell-set draws per gene per val epoch +channel_drop_fraction: 0.0 # fraction of channels to randomly drop per training sample (0 = disabled) +mixed_channels_mode: true # pool all channels into one set (single-level model) +val_n_cells_per_set: [500, 200] # evaluate at multiple n_cells values (mixed mode only) +phase2d_val: false # enable Phase2D-only validation (false to disable) +eval_every: 1 # validate every N training epochs +batch_size: 64 +num_epochs: 100 +learning_rate: 1e-4 +weight_decay: 0.0 +max_grad_norm: null +warmup_epochs: 20 +val_fraction: 0.2 +num_workers: 4 +save_path: best_set_classifier_paper_v2_phase.pt +# Dump processed train/val datasets (one .pt per gene) for reuse. To skip +# reprocessing on a later run, switch these to load_train_dir/load_val_dir. +dump_train_dir: /path/to/dumps/paper_v2_phase/train +dump_val_dir: /path/to/dumps/paper_v2_phase/val + +# ---- Model ---- +model: + d_model: 512 + n_heads: 4 + d_ff: null # null = 4 * d_model + dropout: 0.0 + + # Per-channel set transformer (pools N_cells → 1 vector per channel) + n_layers_cell: 2 + n_inducing_cell: 32 + + channel_conditioning: "concat" # "add" | "concat" | "adaln" | "adaln-token" | "none" + cosine_classifier: True # use cosine-similarity head with learned temperature + + # Cross-channel set transformer (full attention, pools channel vectors → 1 final vector) + +# ---- Logging ---- diff --git a/src/ops_model/interpretability/classifier/configs/train_set_classifier_phase_ebi.yaml b/src/ops_model/interpretability/classifier/configs/train_set_classifier_phase_ebi.yaml new file mode 100644 index 0000000..1d27067 --- /dev/null +++ b/src/ops_model/interpretability/classifier/configs/train_set_classifier_phase_ebi.yaml @@ -0,0 +1,69 @@ +# Config for train_set_classifier.py +# Train a set-transformer classifier on per-cell embeddings. + +name: "celldino-phase-ebi" + +seed: 42 +device: "cuda" # auto-detect + +# ---- Data ---- +data: + parquet_entries: + - path: /path/to/embeddings/celldino_paper_phenotyping_v2_160.parquet + max_row_groups: null # set to limit row groups loaded (e.g. 10 for quick debug) + max_cells_per_group: null # cap cells pe (gene, channel) to limit memory + max_genes: null # keep only top-N genes by cell count (null = all) + max_channels: 1 # keep only top-N channels by cell count (null = all) + # Exact channel labels as in _build_channel_label / printed "Cells per channel" (not raw parquet name if annotation overrides) + exclude_channel_names: null # e.g. [Phase2D] to drop that channel entirely + min_cells_per_group: null # drop (gene, channel) pairs with fewer cells (null = no filter) + min_cells_drop_val: false # also drop from val when min_cells_per_group filters (default: train only) + z_standardize_control_only: true # compute z-standardization stats from NTC (control) cells only + label_map_path: src/ops_model/interpretability/classifier/configs/ebi_complexes.csv # CSV mapping gene names to class labels (e.g. pathways); trains only on genes in this file + label_map_gene_col: gene_name # column in label_map CSV containing gene names (empty/NaN → NTC) + label_map_label_col: pathway # column in label_map CSV containing class labels + +# ---- Training ---- +# n_cells_per_set: 10 # default cells sampled per channel per set +# n_cells_per_set_overrides: # per-channel overrides, e.g. {Phase2D: 200} +# Phase2D: 100 +n_cells_per_set: 100 +train_n_cell_sets_per_gene: 32 # how many cell-set draws per gene per train epoch +val_n_cell_sets_per_gene: 5 # how many cell-set draws per gene per val epoch +channel_drop_fraction: 0.0 # fraction of channels to randomly drop per training sample (0 = disabled) +mixed_channels_mode: true # pool all channels into one set (single-level model) +val_n_cells_per_set: [500, 200] # evaluate at multiple n_cells values (mixed mode only) +phase2d_val: false # enable Phase2D-only validation (false to disable) +eval_every: 1 # validate every N training epochs +batch_size: 64 +num_epochs: 150 +learning_rate: 1e-4 +weight_decay: 0.0 +max_grad_norm: null +warmup_epochs: 20 +val_fraction: 0.2 +num_workers: 4 +save_path: best_set_classifier.pt +# dump_val_dir: null # if set, save val dataset to this directory (one .pt per gene) +# dump_train_dir: /path/to/train_ops_zstdcontrol_cdino # if set, save train dataset to this directory (one .pt per gene) +# dump_val_dir: /path/to/val_ops_zstdcontrol_cdino +load_train_dir: /path/to/train_ops_zstdcontrol_cdino # if set, load pre-dumped train dataset (skip parquet processing) +load_val_dir: /path/to/val_ops_zstdcontrol_cdino # if set, load pre-dumped val dataset (skip parquet processing) + +# ---- Model ---- +model: + d_model: 512 + n_heads: 4 + d_ff: null # null = 4 * d_model + dropout: 0.0 + + # Per-channel set transformer (pools N_cells → 1 vector per channel) + n_layers_cell: 2 + n_inducing_cell: 32 + + channel_conditioning: "concat" # "add" | "concat" | "adaln" | "adaln-token" | "none" + cosine_classifier: True # use cosine-similarity head with learned temperature + + # Cross-channel set transformer (full attention, pools channel vectors → 1 final vector) + +# ---- Logging ---- diff --git a/src/ops_model/interpretability/classifier/eval.py b/src/ops_model/interpretability/classifier/eval.py new file mode 100644 index 0000000..4750cf8 --- /dev/null +++ b/src/ops_model/interpretability/classifier/eval.py @@ -0,0 +1,977 @@ +#!/usr/bin/env python +r"""Evaluate a trained set classifier vs. number of cells per set (with resampling error bars). + +Loads a checkpoint (local path or W&B artifact) and validation data the same way as +``analysis-mixed-channel-attn.ipynb``: mixed-channel mode uses :class:`MixedChannelDataset` +from a ``dump_val_dir`` export. Only ``mixed_channels_mode=true`` is supported. + +For each ``n_cells`` in ``n_cells_list``, runs ``n_repetitions`` full validation passes +with different RNG seeds (stochastic subsampling in the dataset). Plots mean **top-1** +and **top-5** accuracy (top-``k`` uses ``k = min(5, n_classes)``) with SEM error bars. + +Usage +----- +.. code-block:: bash + + cd projects/katamari + python katamari/evals/image_verifier/eval_set_classifier.py \\ + checkpoint_path=/path/to/best_set_classifier.pt \\ + val_dump_dir=/path/to/val_ops \\ + output_plot_path=/path/to/plot.png \\ + n_cells_list='[100,250,500]' \\ + n_repetitions=20 \\ + device=cuda +""" + +from __future__ import annotations + +import csv +import json +import os +import random +from pathlib import Path +from typing import cast + +import hydra +import matplotlib.pyplot as plt +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from omegaconf import DictConfig, OmegaConf +from torch.utils.data import DataLoader, Dataset +from tqdm import tqdm + +from .train import ( + MixedChannelClassifier, + MixedChannelDataset, + _load_label_map, + load_val_dataset, +) + +_CONFIG_DIR = str( + Path(os.environ.get("CONFIG_PATH", Path(__file__).resolve().parent / "configs")) +) + + +def _load_checkpoint(checkpoint_path: str | None) -> dict: + if not checkpoint_path: + raise ValueError("checkpoint_path is required") + path = Path(checkpoint_path).expanduser() + if not path.is_file(): + raise FileNotFoundError(f"Checkpoint not found: {path}") + return torch.load(path, map_location="cpu", weights_only=False) + + +def _per_class_stats( + logits: torch.Tensor, + labels: torch.Tensor, + n_classes: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return (top1_hits, top5_hits, ones) bool/long tensors on CPU for scatter accumulation.""" + k = min(5, n_classes) + preds1 = logits.argmax(dim=-1) + top1_hits = preds1 == labels + if k >= 1: + topk_idx = logits.topk(k, dim=-1).indices + top5_hits = (topk_idx == labels.unsqueeze(1)).any(dim=1) + else: + top5_hits = torch.zeros_like(labels, dtype=torch.bool) + return top1_hits, top5_hits, torch.ones_like(labels, dtype=torch.long) + + +@torch.no_grad() +def _evaluate_mixed_with_topk( + model: MixedChannelClassifier, + loader: DataLoader, + device: torch.device, + *, + n_classes: int, + n_track_genes: int, + gene_indices_in_dataset_order: torch.Tensor, + desc: str | None, + show_batches: bool, +) -> tuple[float, float, float, torch.Tensor, torch.Tensor, torch.Tensor]: + """Returns (mean_loss, top1_accuracy, top5_accuracy, per_gene_top1, per_gene_top5, per_gene_n). + + Per-gene tensors are sized ``n_track_genes`` and indexed by checkpoint gene index + (regardless of whether the model outputs label or gene logits). + """ + model.eval() + total_loss = 0.0 + correct1 = 0 + correct5 = 0 + total = 0 + pc_top1 = torch.zeros(n_track_genes, dtype=torch.long) + pc_top5 = torch.zeros(n_track_genes, dtype=torch.long) + pc_n = torch.zeros(n_track_genes, dtype=torch.long) + it = ( + tqdm(loader, desc=desc or "val", leave=False, unit="batch") + if show_batches + else loader + ) + processed = 0 + for embs, ch_ids, masks, labels in it: + embs = embs.to(device) + ch_ids = ch_ids.to(device) + masks = masks.to(device) + labels = labels.to(device) + logits = model(embs, ch_ids, masks) + loss = F.cross_entropy(logits, labels) + total_loss += loss.item() * labels.size(0) + top1_hits, top5_hits, ones = _per_class_stats(logits, labels, n_classes) + correct1 += top1_hits.sum().item() + correct5 += top5_hits.sum().item() + total += labels.size(0) + bs = labels.size(0) + # shuffle=False on the loader guarantees batches arrive in dataset-index order. + gene_ids = gene_indices_in_dataset_order[processed : processed + bs] + processed += bs + pc_n.scatter_add_(0, gene_ids, ones.cpu()) + pc_top1.scatter_add_(0, gene_ids, top1_hits.long().cpu()) + pc_top5.scatter_add_(0, gene_ids, top5_hits.long().cpu()) + + if total == 0: + raise ValueError( + "Mixed-channel eval saw zero samples (no batches). " + "Check for an empty val dataset or val DataLoader configuration." + ) + return ( + total_loss / total, + correct1 / total, + correct5 / total, + pc_top1, + pc_top5, + pc_n, + ) + + +def _build_channel_alignment( + val_channel_to_idx: dict[str, int], + ckpt_channel_to_idx: dict[str, int], +) -> tuple[torch.Tensor, frozenset[int]]: + """Val-export channel indices → checkpoint indices; return remap tensor and dropped val ids.""" + ckpt_names = set(ckpt_channel_to_idx.keys()) + val_names = set(val_channel_to_idx.keys()) + val_only = val_names - ckpt_names + shared = ckpt_names & val_names + if not shared: + raise ValueError( + "No channel names in common between val export and checkpoint." + ) + + dropped_val_channel_ids = frozenset(val_channel_to_idx[n] for n in val_only) + + ch_remap: dict[int, int] = {} + for name in shared: + ch_remap[val_channel_to_idx[name]] = ckpt_channel_to_idx[name] + + _max_val_ch = max(val_channel_to_idx.values()) + ch_remap_tensor = torch.zeros(_max_val_ch + 1, dtype=torch.long) + for old, new in ch_remap.items(): + ch_remap_tensor[old] = new + + return ch_remap_tensor, dropped_val_channel_ids + + +def _filter_mixed_val_channels( + ds: MixedChannelDataset, + ch_remap_tensor: torch.Tensor, + dropped_val_channel_ids: frozenset[int], +) -> None: + """In-place: drop cells on channels absent from ckpt and remap channel ids to ckpt space. + + Leaves val's gene_to_idx and gene indices unchanged. + """ + dropped_t = ( + torch.tensor(sorted(dropped_val_channel_ids), dtype=torch.long) + if dropped_val_channel_ids + else None + ) + + for g_idx in list(ds._gene_pools.keys()): + pool = ds._gene_pools[g_idx] + ch_ids = ds._gene_ch_ids[g_idx] + if len(pool) == 0: + continue + if dropped_t is not None and len(dropped_t) > 0: + keep = ~torch.isin(ch_ids, dropped_t) + pool = pool[keep] + ch_ids = ch_ids[keep] + if len(pool) == 0: + ds._gene_pools[g_idx] = torch.zeros(0, ds.emb_dim) + ds._gene_ch_ids[g_idx] = torch.zeros(0, dtype=torch.long) + if ds._gene_dump_meta is not None and g_idx in ds._gene_dump_meta: + del ds._gene_dump_meta[g_idx] + continue + ch_ids = ch_remap_tensor[ch_ids.long()] + ds._gene_pools[g_idx] = pool + ds._gene_ch_ids[g_idx] = ch_ids + if ds._gene_dump_meta is not None and g_idx in ds._gene_dump_meta: + del ds._gene_dump_meta[g_idx] + + +def _filter_and_remap_mixed_val( + ds: MixedChannelDataset, + ch_remap_tensor: torch.Tensor, + dropped_val_channel_ids: frozenset[int], + ckpt_gene_to_idx: dict[str, int], +) -> None: + """In-place: drop cells on channels absent from ckpt, remap channel ids, fix labels, drop empty genes.""" + _filter_mixed_val_channels(ds, ch_remap_tensor, dropped_val_channel_ids) + + # Labels must match checkpoint class indices + val_gene_to_idx = ds.gene_to_idx + name_to_ckpt_idx = {n: ckpt_gene_to_idx[n] for n in ckpt_gene_to_idx} + ds.gene_to_idx = dict(ckpt_gene_to_idx) + + kept_genes: list[str] = [] + remap_pool: dict[int, torch.Tensor] = {} + remap_ch: dict[int, torch.Tensor] = {} + for name in ds.perturbation_list: + vgi = val_gene_to_idx[name] + if vgi not in ds._gene_pools or len(ds._gene_pools[vgi]) == 0: + continue + cgi = name_to_ckpt_idx[name] + remap_pool[cgi] = ds._gene_pools[vgi] + remap_ch[cgi] = ds._gene_ch_ids[vgi] + kept_genes.append(name) + + ds.perturbation_list = kept_genes + ds._gene_pools = remap_pool + ds._gene_ch_ids = remap_ch + if ds._gene_dump_meta is not None: + new_meta: dict = {} + for name in kept_genes: + vgi = val_gene_to_idx[name] + cgi = name_to_ckpt_idx[name] + if vgi in ds._gene_dump_meta: + new_meta[cgi] = ds._gene_dump_meta[vgi] + ds._gene_dump_meta = new_meta if new_meta else None + + +def _snapshot_mixed_val( + ds: MixedChannelDataset, +) -> tuple[dict[int, torch.Tensor], dict[int, torch.Tensor], list[str]]: + """Snapshot pools/ch_ids/perturbation_list so we can restore between channel iters.""" + pools = {g: p.clone() for g, p in ds._gene_pools.items()} + ch_ids = {g: c.clone() for g, c in ds._gene_ch_ids.items()} + perts = list(ds.perturbation_list) + return pools, ch_ids, perts + + +def _restrict_mixed_val_to_channel( + ds: MixedChannelDataset, + target_ckpt_ch_id: int, + snapshot_pools: dict[int, torch.Tensor], + snapshot_ch_ids: dict[int, torch.Tensor], + snapshot_perturbation_list: list[str], + active_gene_to_idx: dict[str, int], +) -> None: + """Restore ds from snapshot then drop all cells not on ``target_ckpt_ch_id``. + + Genes with zero cells of the target channel are dropped from ``perturbation_list`` + so they don't pollute accuracy stats with fully-masked padded sets. + """ + new_pools: dict[int, torch.Tensor] = {} + new_ch_ids: dict[int, torch.Tensor] = {} + kept_gene_names: list[str] = [] + for name in snapshot_perturbation_list: + g_idx = active_gene_to_idx[name] + pool = snapshot_pools[g_idx] + ch_ids = snapshot_ch_ids[g_idx] + keep_mask = ch_ids == target_ckpt_ch_id + if not bool(keep_mask.any()): + continue + new_pools[g_idx] = pool[keep_mask] + new_ch_ids[g_idx] = ch_ids[keep_mask] + kept_gene_names.append(name) + # Genes outside perturbation_list still need entries (used elsewhere by index). + for g_idx, pool in snapshot_pools.items(): + if g_idx in new_pools: + continue + new_pools[g_idx] = torch.zeros(0, ds.emb_dim) + new_ch_ids[g_idx] = torch.zeros(0, dtype=torch.long) + ds._gene_pools = new_pools + ds._gene_ch_ids = new_ch_ids + ds.perturbation_list = kept_gene_names + + +def _run_sweep( + *, + val_ds: Dataset, + loader: DataLoader, + model: nn.Module, + device: torch.device, + mixed_mode: bool, + n_cells_list: list[int], + n_repetitions: int, + base_seed: int, + n_classes: int, + n_active_genes: int, + gene_indices_in_dataset_order: torch.Tensor, + loader_idx_to_ch: dict[int, str] | None, + n_parquet_channels: int | None, + overrides_raw, + show_progress: bool, + show_batch_pbar: bool, + desc_prefix: str = "", +) -> tuple[ + list[float], + list[float], + list[float], + list[float], + dict[str, dict[str, list[int]]], + dict[str, dict[str, list[float]]], +]: + """Run the n_cells × n_repetitions sweep and return aggregate + per-gene results.""" + results: dict[str, dict[str, list[float]]] = {} + means_top1: list[float] = [] + stderrs_top1: list[float] = [] + means_top5: list[float] = [] + stderrs_top5: list[float] = [] + per_gene_results: dict[str, dict[str, list[int]]] = {} + + outer_desc = f"{desc_prefix}n_cells" if desc_prefix else "n_cells" + n_cells_iter = ( + tqdm(n_cells_list, desc=outer_desc) if show_progress else n_cells_list + ) + for n_cells in n_cells_iter: + accs1: list[float] = [] + accs5: list[float] = [] + key = str(n_cells) + results[key] = {"top1": [], "top5": []} + cum_pc_top1 = torch.zeros(n_active_genes, dtype=torch.long) + cum_pc_top5 = torch.zeros(n_active_genes, dtype=torch.long) + cum_pc_n = torch.zeros(n_active_genes, dtype=torch.long) + rep_iter = ( + tqdm( + range(n_repetitions), + desc=f"{desc_prefix}n_cells={n_cells} reps", + leave=False, + unit="rep", + ) + if show_progress + else range(n_repetitions) + ) + for rep in rep_iter: + rseed = base_seed + rep * 1_000_003 + n_cells * 17 + torch.manual_seed(rseed) + np.random.seed(rseed % (2**32)) + random.seed(rseed) + if device.type == "cuda": + torch.cuda.manual_seed_all(rseed) + if mixed_mode: + assert isinstance(val_ds, MixedChannelDataset) + val_ds.set_n_cells(n_cells) + else: + raise ValueError("Only mixed_channels_mode=true is supported.") + + batch_desc = f"{desc_prefix}n={n_cells} rep {rep + 1}/{n_repetitions}" + if mixed_mode: + _, acc1, acc5, pc_top1, pc_top5, pc_n = _evaluate_mixed_with_topk( + cast(MixedChannelClassifier, model), + loader, + device, + n_classes=n_classes, + n_track_genes=n_active_genes, + gene_indices_in_dataset_order=gene_indices_in_dataset_order, + desc=batch_desc, + show_batches=show_batch_pbar, + ) + else: + raise ValueError("Only mixed_channels_mode=true is supported.") + accs1.append(acc1) + accs5.append(acc5) + results[key]["top1"].append(acc1) + results[key]["top5"].append(acc5) + cum_pc_top1 += pc_top1 + cum_pc_top5 += pc_top5 + cum_pc_n += pc_n + + m1 = float(np.mean(accs1)) + m5 = float(np.mean(accs5)) + if n_repetitions > 1: + se1 = float(np.std(accs1, ddof=1) / np.sqrt(len(accs1))) + se5 = float(np.std(accs5, ddof=1) / np.sqrt(len(accs5))) + else: + se1 = 0.0 + se5 = 0.0 + means_top1.append(m1) + stderrs_top1.append(se1) + means_top5.append(m5) + stderrs_top5.append(se5) + per_gene_results[key] = { + "top1_correct": cum_pc_top1.tolist(), + "top5_correct": cum_pc_top5.tolist(), + "n_samples": cum_pc_n.tolist(), + } + + return ( + means_top1, + stderrs_top1, + means_top5, + stderrs_top5, + per_gene_results, + results, + ) + + +def run(cfg: DictConfig) -> None: + ckpt = _load_checkpoint(cfg.get("checkpoint_path")) + config = ckpt["config"] + ckpt_gene_to_idx: dict[str, int] = ckpt["gene_to_idx"] + ckpt_channel_to_idx: dict[str, int] = ckpt["channel_to_idx"] + n_ckpt_channels = len(ckpt_channel_to_idx) + state_dict = ckpt["model_state_dict"] + if any(k.startswith("module.") for k in state_dict): + state_dict = {k.replace("module.", "", 1): v for k, v in state_dict.items()} + + # ---- Label map (e.g. gene → pathway) ---- + label_to_idx: dict[str, int] | None = None + label_remap: dict[int, int] | None = None + gene_to_label: dict[str, str] | None = None + + label_map_path = cfg.get("label_map_path", None) + if label_map_path is not None: + # Supersedes any gene→label mapping stored in the checkpoint, letting us + # evaluate on a different gene set. The label SPACE (label_to_idx) must + # match training though, otherwise model output indices misalign — so we + # take label_to_idx from ckpt when it's saved there. CSV labels not in + # the trained label space are dropped (model can't predict them). + gene_col = cfg.get("label_map_gene_col", "gene_name") + label_col = cfg.get("label_map_label_col", "pathway") + gene_to_label = _load_label_map(label_map_path, gene_col, label_col) + if "label_to_idx" in ckpt: + label_to_idx = cast(dict[str, int], ckpt["label_to_idx"]) + new_labels = set(gene_to_label.values()) - set(label_to_idx) + if new_labels: + print( + f" Warning: {len(new_labels)} labels in CSV not in trained " + f"label set (genes dropped): {sorted(new_labels)[:10]}" + ) + gene_to_label = { + g: lab for g, lab in gene_to_label.items() if lab in label_to_idx + } + else: + unique_labels = sorted(set(gene_to_label.values())) + label_to_idx = {label: i for i, label in enumerate(unique_labels)} + print( + f"Label map from config: {len(gene_to_label)} genes in CSV → " + f"{len(label_to_idx)} classes ({label_col})" + ) + elif "label_to_idx" in ckpt: + label_to_idx = cast(dict[str, int], ckpt["label_to_idx"]) + if "label_remap" in ckpt: + label_remap = cast(dict[int, int], ckpt["label_remap"]) + print( + f"Label map from checkpoint: {len(label_remap)} genes → " + f"{len(label_to_idx)} classes" + ) + else: + raise ValueError( + "Checkpoint has label_to_idx but no label_remap (older checkpoint). " + "Set label_map_path in the config to the same CSV used during training " + "so the gene→label mapping can be reconstructed." + ) + + n_classes = state_dict["head.1.weight"].shape[0] + if label_to_idx is not None and len(label_to_idx) != n_classes: + raise ValueError( + f"label_to_idx has {len(label_to_idx)} classes but model output dim " + f"is {n_classes}. The CSV's label set must match what the model was " + f"trained on." + ) + + mixed_mode = bool( + cfg.get("mixed_channels_mode", config.get("mixed_channels_mode", True)) + ) + + device = torch.device( + cfg.get("device", "cuda" if torch.cuda.is_available() else "cpu") + ) + base_seed = int(cfg.get("seed", 42)) + + n_cells_list = [int(x) for x in list(cfg.n_cells_list)] + if not n_cells_list: + raise ValueError("n_cells_list must be non-empty") + n_repetitions = int(cfg["n_repetitions"]) + if n_repetitions < 1: + raise ValueError("n_repetitions must be >= 1") + + batch_size = int(cfg.get("batch_size", 64)) + num_workers = int(cfg.get("num_workers", 0)) + + # Checkpoint stores Hydra config as a plain dict after torch.load — not an OmegaConf node. + mcfg = OmegaConf.to_container(OmegaConf.create(config["model"]), resolve=True) + assert isinstance(mcfg, dict) + + # ---- Model ---- + emb_dim = int(state_dict["input_proj.weight"].shape[1]) + if mixed_mode: + model = MixedChannelClassifier( + emb_dim=emb_dim, + n_classes=n_classes, + n_channels=n_ckpt_channels, + d_model=mcfg.get("d_model", 256), + n_heads=mcfg.get("n_heads", 4), + n_layers=mcfg.get("n_layers_cell", 2), + n_inducing=mcfg.get("n_inducing_cell", 32), + d_ff=mcfg.get("d_ff", None), + dropout=mcfg.get("dropout", 0.1), + cosine_classifier=mcfg.get("cosine_classifier", False), + channel_conditioning=mcfg.get("channel_conditioning", "none"), + ) + else: + raise ValueError("Only mixed_channels_mode=true is supported.") + + model.load_state_dict(state_dict) + model = model.to(device) + model.eval() + + # ---- Data ---- + if mixed_mode: + vd = cfg.get("val_dump_dir") + if vd in (None, "", "null"): + raise ValueError( + "mixed_channels_mode=true requires val_dump_dir (directory with metadata.pt from dump_val_dir)." + ) + val_dir = Path(str(vd)).expanduser() + meta_path = val_dir / "metadata.pt" + if not meta_path.is_file(): + raise FileNotFoundError(f"Expected val export at {meta_path}") + val_meta = torch.load(meta_path, map_location="cpu", weights_only=False) + val_gene_to_idx: dict[str, int] = val_meta["gene_to_idx"] + val_channel_to_idx: dict[str, int] = val_meta["channel_to_idx"] + + # Gene set check: with label_map_path, val genes only need to be in the + # CSV (not in ckpt). Otherwise, val genes must be a subset of ckpt's. + if gene_to_label is None: + val_only = set(val_gene_to_idx) - set(ckpt_gene_to_idx) + if val_only: + raise ValueError( + "Val export has genes not present in checkpoint " + f"(cannot be evaluated): {sorted(val_only)[:20]}" + ) + ckpt_only = set(ckpt_gene_to_idx) - set(val_gene_to_idx) + if ckpt_only: + print( + f" {len(ckpt_only)} genes in checkpoint absent from val export " + f"(skipped): {sorted(ckpt_only)[:10]}{'...' if len(ckpt_only) > 10 else ''}" + ) + + ch_remap_tensor, dropped_val_ch = _build_channel_alignment( + val_channel_to_idx, ckpt_channel_to_idx + ) + + only_channels_raw = cfg.get("only_channels") + if only_channels_raw is not None: + only_channels = set(only_channels_raw) + unknown = only_channels - set(val_channel_to_idx.keys()) + if unknown: + raise ValueError( + f"only_channels names not found in val export: {unknown}. " + f"Available: {sorted(val_channel_to_idx.keys())}" + ) + extra_drop = frozenset( + val_channel_to_idx[n] + for n in val_channel_to_idx + if n not in only_channels + ) + dropped_val_ch = dropped_val_ch | extra_drop + print(f"Restricting validation to channels: {sorted(only_channels)}") + + max_n = max(n_cells_list) + val_ds = load_val_dataset(str(val_dir), n_cells=max_n) + val_ds.replacement = not bool(cfg.get("sample_without_replacement", False)) + + if gene_to_label is not None: + # label_map_path mode: keep val's gene indices, build label_remap + # from val_gene_to_idx + CSV. Val genes can be entirely new wrt ckpt. + assert label_to_idx is not None + _filter_mixed_val_channels(val_ds, ch_remap_tensor, dropped_val_ch) + kept = [ + g + for g in val_ds.perturbation_list + if g in gene_to_label + and val_gene_to_idx[g] in val_ds._gene_pools + and len(val_ds._gene_pools[val_gene_to_idx[g]]) > 0 + ] + if len(kept) < len(val_ds.perturbation_list): + dropped = len(val_ds.perturbation_list) - len(kept) + print(f" Label map: dropped {dropped} genes without label mapping") + val_ds.perturbation_list = kept + label_remap = { + val_gene_to_idx[g]: label_to_idx[gene_to_label[g]] for g in kept + } + val_ds.label_remap = label_remap + active_gene_to_idx: dict[str, int] = val_gene_to_idx + else: + _filter_and_remap_mixed_val( + val_ds, ch_remap_tensor, dropped_val_ch, ckpt_gene_to_idx + ) + if label_remap is not None: + val_ds.label_remap = label_remap + kept = [ + g + for g in val_ds.perturbation_list + if ckpt_gene_to_idx[g] in label_remap + ] + if len(kept) < len(val_ds.perturbation_list): + dropped = len(val_ds.perturbation_list) - len(kept) + print(f" Label map: dropped {dropped} genes without label mapping") + val_ds.perturbation_list = kept + active_gene_to_idx = ckpt_gene_to_idx + + if len(val_ds.perturbation_list) == 0: + raise ValueError( + "No genes left in val set after channel filtering (empty pools)." + ) + + loader = DataLoader( + val_ds, + batch_size=batch_size, + shuffle=False, + num_workers=num_workers, + pin_memory=device.type == "cuda", + ) + gene_indices_in_dataset_order = torch.tensor( + [active_gene_to_idx[name] for name in val_ds.perturbation_list], + dtype=torch.long, + ) + loader_idx_to_ch: dict[int, str] | None = None + n_parquet_channels: int | None = None + else: + raise ValueError("Only mixed_channels_mode=true is supported.") + + # ---- Eval sweeps ---- + show_progress = bool(cfg.get("show_progress", True)) + show_batch_pbar = show_progress + n_active_genes = len(active_gene_to_idx) + split_channels = bool(cfg.get("split_channels", False)) + overrides_raw = cfg.get("n_cells_per_set_overrides", {}) + + # Per-channel sweep results when split_channels=true; otherwise a single "all" entry. + sweep_by_channel: dict[ + str, + tuple[ + list[float], + list[float], + list[float], + list[float], + dict[str, dict[str, list[int]]], + dict[str, dict[str, list[float]]], + ], + ] = {} + + if split_channels: + if not mixed_mode: + raise ValueError("split_channels=true requires mixed_channels_mode=true.") + assert isinstance(val_ds, MixedChannelDataset) + # Snapshot the val_ds state after the ckpt-wide channel filter+remap, so each + # per-channel iteration starts from the same clean baseline. + snap_pools, snap_ch_ids, snap_perts = _snapshot_mixed_val(val_ds) + + if only_channels_raw is not None: + eval_channel_names = list(only_channels_raw) + else: + eval_channel_names = sorted( + set(ckpt_channel_to_idx.keys()) & set(val_channel_to_idx.keys()) + ) + if not eval_channel_names: + raise ValueError( + "split_channels=true but no channels available to evaluate " + "(empty intersection of checkpoint and val export)." + ) + unknown_ch = [n for n in eval_channel_names if n not in ckpt_channel_to_idx] + if unknown_ch: + raise ValueError( + f"split_channels: channel names not in checkpoint vocab: {unknown_ch}" + ) + # Optional channel sharding for parallel eval jobs: each job handles a + # deterministic stride of the (sorted) channel list; combine CSVs after. + n_channel_shards = int(cfg.get("n_channel_shards", 1)) + channel_shard_id = int(cfg.get("channel_shard_id", 0)) + if n_channel_shards > 1: + eval_channel_names = eval_channel_names[channel_shard_id::n_channel_shards] + print( + f"channel shard {channel_shard_id}/{n_channel_shards}: " + f"{len(eval_channel_names)} channels" + ) + print(f"split_channels: evaluating channels {eval_channel_names}") + + for ch_name in eval_channel_names: + target_ch_id = ckpt_channel_to_idx[ch_name] + _restrict_mixed_val_to_channel( + val_ds, + target_ch_id, + snap_pools, + snap_ch_ids, + snap_perts, + active_gene_to_idx, + ) + if not val_ds.perturbation_list: + print( + f" Skipping channel {ch_name!r}: no genes with cells in this channel" + ) + continue + gene_indices_in_dataset_order = torch.tensor( + [active_gene_to_idx[name] for name in val_ds.perturbation_list], + dtype=torch.long, + ) + sweep_out = _run_sweep( + val_ds=val_ds, + loader=loader, + model=model, + device=device, + mixed_mode=mixed_mode, + n_cells_list=n_cells_list, + n_repetitions=n_repetitions, + base_seed=base_seed, + n_classes=n_classes, + n_active_genes=n_active_genes, + gene_indices_in_dataset_order=gene_indices_in_dataset_order, + loader_idx_to_ch=loader_idx_to_ch, + n_parquet_channels=n_parquet_channels, + overrides_raw=overrides_raw, + show_progress=show_progress, + show_batch_pbar=show_batch_pbar, + desc_prefix=f"[{ch_name}] ", + ) + sweep_by_channel[ch_name] = sweep_out + else: + sweep_out = _run_sweep( + val_ds=val_ds, + loader=loader, + model=model, + device=device, + mixed_mode=mixed_mode, + n_cells_list=n_cells_list, + n_repetitions=n_repetitions, + base_seed=base_seed, + n_classes=n_classes, + n_active_genes=n_active_genes, + gene_indices_in_dataset_order=gene_indices_in_dataset_order, + loader_idx_to_ch=loader_idx_to_ch, + n_parquet_channels=n_parquet_channels, + overrides_raw=overrides_raw, + show_progress=show_progress, + show_batch_pbar=show_batch_pbar, + ) + # Use a single synthetic channel name for the unified-eval case. + sweep_by_channel["__all__"] = sweep_out + + # ---- Plot ---- + out_raw = cfg.get("output_plot_path") + out_plot = ( + Path(str(out_raw)).expanduser() if out_raw not in (None, "", "null") else None + ) + plot_x_log = bool(cfg.get("plot_x_log", True)) + x = np.array(sorted(n_cells_list), dtype=float) + plot_title = cfg.get("plot_title", "Set classifier accuracy vs. N cells") + + if out_plot is None: + print("output_plot_path not set; skipping plot and sidecar JSON.") + else: + out_plot.parent.mkdir(parents=True, exist_ok=True) + if split_channels: + fig, (ax1, ax5) = plt.subplots(1, 2, figsize=(14, 5)) + for i, (ch_name, sweep_out) in enumerate(sweep_by_channel.items()): + m1, e1, m5, e5, *_ = sweep_out + color = f"C{i % 10}" + ax1.errorbar( + x, + np.array(m1), + yerr=np.array(e1), + fmt="o-", + capsize=3, + capthick=1.0, + linewidth=1.8, + markersize=5, + label=ch_name, + color=color, + ) + ax5.errorbar( + x, + np.array(m5), + yerr=np.array(e5), + fmt="s-", + capsize=3, + capthick=1.0, + linewidth=1.8, + markersize=5, + label=ch_name, + color=color, + ) + for ax, title in ((ax1, "Top-1"), (ax5, "Top-5")): + ax.set_xlabel("N cells per set") + ax.set_ylabel("Validation accuracy") + ax.set_title(f"{plot_title} ({title})") + if plot_x_log: + ax.set_xscale("log") + tick_vals = sorted({float(n) for n in n_cells_list}) + ax.set_xticks(tick_vals) + ax.set_xticklabels( + [ + str(int(v)) if v >= 1.0 and v == int(v) else f"{v:g}" + for v in tick_vals + ] + ) + ax.grid(True, alpha=0.3) + ax.set_ylim(0.0, 1.0) + ax.legend(loc="lower right", fontsize=8) + else: + m1, e1, m5, e5, *_ = sweep_by_channel["__all__"] + fig, ax = plt.subplots(figsize=(8, 5)) + ax.errorbar( + x, + np.array(m1), + yerr=np.array(e1), + fmt="o-", + capsize=4, + capthick=1.5, + linewidth=2, + markersize=6, + label="Top-1", + color="C0", + ) + ax.errorbar( + x, + np.array(m5), + yerr=np.array(e5), + fmt="s-", + capsize=4, + capthick=1.5, + linewidth=2, + markersize=6, + label="Top-5", + color="C1", + ) + ax.set_xlabel("N cells per set") + ax.set_ylabel("Validation accuracy") + ax.set_title(plot_title) + if plot_x_log: + ax.set_xscale("log") + tick_vals = sorted({float(n) for n in n_cells_list}) + ax.set_xticks(tick_vals) + ax.set_xticklabels( + [ + str(int(v)) if v >= 1.0 and v == int(v) else f"{v:g}" + for v in tick_vals + ] + ) + ax.grid(True, alpha=0.3) + ax.set_ylim(0.0, 1.0) + ax.legend(loc="lower right") + fig.tight_layout() + fig.savefig(out_plot, dpi=150) + plt.close(fig) + print(f"Saved plot to {out_plot}") + + # ---- Sidecar JSON ---- + sidecar = out_plot.with_suffix(".json") + payload: dict = { + "n_cells_list": n_cells_list, + "n_repetitions": n_repetitions, + "plot_x_log": plot_x_log, + "device": str(device), + "mixed_channels_mode": mixed_mode, + "split_channels": split_channels, + "n_classes": n_classes, + } + if split_channels: + payload["per_channel"] = { + ch_name: { + "mean_accuracy_top1": m1, + "stderr_top1": e1, + "mean_accuracy_top5": m5, + "stderr_top5": e5, + "repetitions": reps, + } + for ch_name, (m1, e1, m5, e5, _, reps) in sweep_by_channel.items() + } + else: + m1, e1, m5, e5, _, reps = sweep_by_channel["__all__"] + payload["mean_accuracy_top1"] = m1 + payload["stderr_top1"] = e1 + payload["mean_accuracy_top5"] = m5 + payload["stderr_top5"] = e5 + payload["repetitions"] = reps + if label_to_idx is not None: + idx_to_label = {v: k for k, v in label_to_idx.items()} + payload["label_to_idx"] = label_to_idx + payload["class_names"] = [idx_to_label[i] for i in range(len(idx_to_label))] + sidecar.write_text(json.dumps(payload, indent=2)) + print(f"Saved metrics to {sidecar}") + + # ---- Per-class CSV ---- + per_class_path = cfg.get("per_class_output_file") + if per_class_path not in (None, "", "null"): + per_class_out = Path(str(per_class_path)).expanduser() + per_class_out.parent.mkdir(parents=True, exist_ok=True) + idx_to_gene_name = {v: k for k, v in active_gene_to_idx.items()} + has_label_map = label_remap is not None + idx_to_label_name: dict[int, str] = {} + if has_label_map: + assert label_to_idx is not None and label_remap is not None + idx_to_label_name = {v: k for k, v in label_to_idx.items()} + gene_indices_to_write = sorted(label_remap.keys()) + else: + gene_indices_to_write = sorted(active_gene_to_idx.values()) + with open(per_class_out, "w", newline="") as f: + writer = csv.writer(f) + header: list[str] = [] + if split_channels: + header.append("channel_name") + header += ["n_cells", "gene_idx", "gene_name"] + if has_label_map: + header += ["label_name"] + header += [ + "n_repetitions", + "n_samples", + "top1_correct", + "top5_correct", + "top1_acc", + "top5_acc", + ] + writer.writerow(header) + for ch_name, sweep_out in sweep_by_channel.items(): + _, _, _, _, per_gene_results, _ = sweep_out + for n_cells in n_cells_list: + pc = per_gene_results[str(n_cells)] + for gene_idx in gene_indices_to_write: + n_samples = pc["n_samples"][gene_idx] + top1_c = pc["top1_correct"][gene_idx] + top5_c = pc["top5_correct"][gene_idx] + top1_acc = top1_c / n_samples if n_samples > 0 else 0.0 + top5_acc = top5_c / n_samples if n_samples > 0 else 0.0 + row: list = [] + if split_channels: + row.append(ch_name) + row += [n_cells, gene_idx, idx_to_gene_name[gene_idx]] + if has_label_map: + assert label_remap is not None + row.append(idx_to_label_name[label_remap[gene_idx]]) + row += [ + n_repetitions, + n_samples, + top1_c, + top5_c, + f"{top1_acc:.6f}", + f"{top5_acc:.6f}", + ] + writer.writerow(row) + print(f"Saved per-gene metrics to {per_class_out}") + + +@hydra.main( + version_base="1.3.0", + config_path=_CONFIG_DIR, + config_name="eval_set_classifier", +) +def main(cfg: DictConfig) -> None: + run(cfg) + + +if __name__ == "__main__": + main() # type: ignore[call-arg] diff --git a/src/ops_model/interpretability/classifier/score.py b/src/ops_model/interpretability/classifier/score.py new file mode 100644 index 0000000..7394b49 --- /dev/null +++ b/src/ops_model/interpretability/classifier/score.py @@ -0,0 +1,327 @@ +#!/usr/bin/env python +r"""Rank cells of (gene, marker) pairs by a per-cell *score* from the set classifier. + +For each cell the score is its leave-one-out marginal contribution to ``P(true class)`` — +``P(class | bag) - P(class | bag without the cell)`` — averaged over random bags and then +averaged uniformly over a grid of bag sizes. Bag size 1 is the deterministic single-cell +probability ``P(class | cell)``. + +Bags are single-channel (all cells of one marker). One CSV row per cell carries the score, +the per-bag-size marginals, and the cell's coordinates (so montages/viewers can render +straight from the ranking without a lossy re-merge to the dump). + +Example: +------- + python -m ops_model.interpretability.classifier.score \\ + --checkpoint set_classifier.pt \\ + --dump_dir /path/to/dumps/train /path/to/dumps/val \\ + --channel Phase2D --genes KIF23 HSPA5 AURKB \\ + --bag_sizes 1 2 5 10 20 50 100 200 500 --reps 50 \\ + --out_csv score_phase.csv +""" + +import argparse +import csv +from pathlib import Path + +import numpy as np +import torch + +from .train import build_model + + +def flatten_meta(cm: dict) -> dict: + """Flatten a per-cell metadata dict whose values may be lists-of-lists (one inner + list per experiment chunk) into flat per-cell lists. + """ + return { + k: ([x for seg in v for x in seg] if v and isinstance(v[0], list) else list(v)) + for k, v in cm.items() + } + + +@torch.no_grad() +def single_cell_prob(model, emb: torch.Tensor, ch: torch.Tensor, tl: int) -> np.ndarray: + """Deterministic bag-size-1 attribution: ``P(true class | cell alone)`` per cell.""" + n = emb.shape[0] + out = torch.empty(n, dtype=torch.float64) + for i in range(0, n, 4096): + e = emb[i : i + 4096][:, None, :] + c = ch[i : i + 4096][:, None] + m = torch.ones_like(c, dtype=torch.bool) + out[i : i + 4096] = torch.softmax(model(e, c, m), -1)[:, tl].double().cpu() + return out.numpy() + + +@torch.no_grad() +def marginal( + model, + emb: torch.Tensor, + ch: torch.Tensor, + tl: int, + bag: int, + n_reps: int, + seed: int, + device: torch.device, +) -> np.ndarray: + """Mean leave-one-out marginal of each cell at a given bag size. + + Over ``n_reps`` random partitions of the cells into bags of size ``bag``, each cell lands + in exactly one bag per rep. Within a bag we compute the full-bag prediction once and each + leave-one-out prediction (``bag + 1`` forward passes per bag), so a cell's marginal is + ``P(class | bag) - P(class | bag without cell)`` averaged over the ``n_reps`` bags it fell + into. The leave-one-out set is materialized by masking, and bags are processed in blocks to + bound peak activation memory (the block spans a ``(block*bag, bag, d)`` expansion). + """ + n = emb.shape[0] + marg_sum = torch.zeros(n, dtype=torch.float64, device=device) + cnt = torch.zeros(n, dtype=torch.float64, device=device) + nbpr = (n + bag - 1) // bag + eye = torch.eye(bag, dtype=torch.bool, device=device) + block = max(1, 300_000 // (bag * bag)) + for rep in range(n_reps): + g = torch.Generator().manual_seed(seed * 1_000_003 + rep) + perm = torch.randperm(n, generator=g) + pad = nbpr * bag - n + if pad: + perm = torch.cat( + [perm, perm.new_zeros(pad)] + ) # pad with a real index; masked out below + valid = torch.ones(nbpr * bag, dtype=torch.bool) + valid[n:] = False + bag_idx = perm.view(nbpr, bag).to(device) + vm = valid.view(nbpr, bag).to(device) + for b0 in range(0, nbpr, block): + sub = bag_idx[b0 : b0 + block] + vmb = vm[b0 : b0 + block] + nb = sub.shape[0] + e = emb[sub] + c = ch[sub] + p_full = torch.softmax(model(e, c, vmb), -1)[:, tl].double() + mask_loo = vmb[:, None, :] & (~eye)[None, :, :] + er = e[:, None, :, :].expand(nb, bag, bag, -1).reshape(nb * bag, bag, -1) + cr = c[:, None, :].expand(nb, bag, bag).reshape(nb * bag, bag) + mr = mask_loo.reshape(nb * bag, bag) + p_loo = torch.softmax(model(er, cr, mr), -1)[:, tl].double().view(nb, bag) + mg = p_full[:, None] - p_loo + flat = sub[vmb] + marg_sum.index_add_(0, flat, mg[vmb]) + cnt.index_add_( + 0, flat, torch.ones(int(vmb.sum()), dtype=torch.float64, device=device) + ) + if device.type == "cuda": + torch.cuda.empty_cache() + return (marg_sum / cnt.clamp(min=1)).cpu().numpy() + + +# Relative reps taper (reps_B ∝ σ_B/√B, from the measured per-sample marginal std on phase +# genes): fraction of the anchor reps to spend at each bag size. Off-grid bags use a power-law +# fallback. Bags land no lower than --min_reps. +_TAPER_FRAC = { + 2: 1.0, + 5: 0.6, + 10: 0.3, + 20: 0.1, + 50: 0.04, + 100: 0.02, + 200: 0.01, + 500: 0.002, +} + + +def reps_for_bag(bag: int, anchor: int, min_reps: int) -> int: + """Reps to use at a given bag size: the anchor scaled by the variance-based taper, floored + at min_reps. Bag 1 is deterministic (handled by the caller) so this is only for bag >= 2. + """ + frac = _TAPER_FRAC.get(bag, (2.0 / bag) ** 0.85) + return max(min_reps, round(anchor * frac)) + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--checkpoint", required=True) + ap.add_argument( + "--dump_dir", required=True, nargs="+", help="dump dirs pooled (train val)" + ) + ap.add_argument( + "--channel", required=True, help="marker/channel name, e.g. Phase2D or 5xUPRE" + ) + ap.add_argument("--genes", required=True, nargs="+", help="genes to rank") + ap.add_argument( + "--bag_sizes", + type=int, + nargs="+", + default=[1, 2, 5, 10, 20, 50], + help="bag sizes to average the marginal over (bag 1 = single-cell P(class))", + ) + ap.add_argument( + "--reps", + type=int, + default=100, + help="anchor reps for the smallest bag; larger bags are tapered down toward " + "--min_reps (the per-sample marginal variance collapses with bag size). " + "Use --flat_reps to disable the taper.", + ) + ap.add_argument( + "--min_reps", type=int, default=10, help="floor on reps at any bag size" + ) + ap.add_argument( + "--flat_reps", + action="store_true", + help="use --reps uniformly for every bag size (no taper)", + ) + ap.add_argument( + "--reps_schedule", + type=int, + nargs="+", + default=None, + help="explicit reps per bag size (must match --bag_sizes length); overrides taper", + ) + ap.add_argument( + "--max_cells", + type=int, + default=65000, + help="subsample a (gene, channel) pool larger than this (fixed seed); the " + "full leave-one-out score is infeasible for huge pools like the NTC control", + ) + ap.add_argument("--out_csv", required=True) + ap.add_argument( + "--seed", type=int, default=0, help="seed for the random bag partitions" + ) + ap.add_argument( + "--subsample_seed", + type=int, + default=0, + help="seed for the >--max_cells subsample; kept separate from --seed so the " + "same cells are scored when only the partition seed changes (e.g. reruns)", + ) + ap.add_argument("--device", default="cuda") + args = ap.parse_args() + + device = torch.device(args.device if torch.cuda.is_available() else "cpu") + ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False) + model = build_model(ckpt, device) + gene_to_idx = ckpt["gene_to_idx"] + cci = ckpt["channel_to_idx"][args.channel] + label_remap = ckpt.get("label_remap") + md = torch.load( + Path(args.dump_dir[0]) / "metadata.pt", map_location="cpu", weights_only=False + ) + dci = md["channel_to_idx"][args.channel] + bags = sorted(set(args.bag_sizes)) + # reps per bag: explicit schedule > flat > variance-based taper (default) + if args.reps_schedule is not None: + if len(args.reps_schedule) != len(args.bag_sizes): + raise ValueError("--reps_schedule length must match --bag_sizes length") + reps_map = dict(zip(args.bag_sizes, args.reps_schedule)) + elif args.flat_reps: + reps_map = dict.fromkeys(bags, args.reps) + else: + reps_map = {b: reps_for_bag(b, args.reps, args.min_reps) for b in bags} + reps_map[1] = 1 # bag 1 is deterministic (single forward pass) + print("reps per bag: " + ", ".join(f"{b}:{reps_map[b]}" for b in bags), flush=True) + + with open(args.out_csv, "w", newline="") as fo: + w = csv.writer(fo) + w.writerow( + ["gene", "channel_name", "rank", "score", "bag1"] + + [f"marg_{b}" for b in bags] + + [ + "split", + "experiment", + "well", + "y_pheno", + "x_pheno", + "segmentation_id", + "zarr_channel_index", + "n_cells", + ] + ) + for gene in args.genes: + if gene not in gene_to_idx: + print(f"skip {gene}: not in gene_to_idx", flush=True) + continue + tl = ( + label_remap[gene_to_idx[gene]] + if label_remap is not None + else gene_to_idx[gene] + ) + # pool the gene's cells for this one channel across dump dirs, tracking metadata + split + embs_l, meta_l, split_l = [], [], [] + for dd in args.dump_dir: + gp = Path(dd) / f"{gene}.pt" + if not gp.exists(): + continue + d = torch.load(gp, map_location="cpu", weights_only=False) + sel = np.flatnonzero(d["channel_ids"].long().numpy() == dci) + if sel.size == 0: + continue + embs_l.append(d["embeddings"].float().numpy()[sel]) + m = flatten_meta(d["cell_metadata"]) + meta_l.append( + { + k: [m[k][i] for i in sel] + for k in ( + "experiment", + "well", + "y_pheno", + "x_pheno", + "segmentation_id", + "index", + ) + } + ) + split_l.extend([Path(dd).name] * sel.size) + if not embs_l: + print(f"skip {gene}: no {args.channel} cells", flush=True) + continue + emb_np = np.concatenate(embs_l, 0) + meta = {k: [v for m in meta_l for v in m[k]] for k in meta_l[0]} + splits = np.array(split_l) + n_full = emb_np.shape[0] + keep = np.arange(n_full) + if n_full > args.max_cells: + rng = np.random.default_rng(args.subsample_seed) + keep = np.sort(rng.choice(n_full, size=args.max_cells, replace=False)) + emb_np = emb_np[keep] + splits = splits[keep] + + emb = torch.tensor(emb_np, device=device) + n = emb.shape[0] + ch = torch.full((n,), cci, dtype=torch.long, device=device) + margs = { + b: ( + single_cell_prob(model, emb, ch, tl) + if b == 1 + else marginal(model, emb, ch, tl, b, reps_map[b], args.seed, device) + ) + for b in bags + } + score = np.mean(np.stack([margs[b] for b in bags], 0), 0) + bag1 = margs[1] if 1 in margs else np.full(n, np.nan) + order = np.argsort(-score) + for rk, ci in enumerate(order): + gi = int(keep[ci]) + w.writerow( + [gene, args.channel, rk + 1, f"{score[ci]:.6f}", f"{bag1[ci]:.6f}"] + + [f"{margs[b][ci]:.6f}" for b in bags] + + [ + splits[ci], + meta["experiment"][gi], + meta["well"][gi], + meta["y_pheno"][gi], + meta["x_pheno"][gi], + meta["segmentation_id"][gi], + meta["index"][gi], + n, + ] + ) + print( + f"{gene} x {args.channel}: n={n} ranked (score mean={score.mean():+.5f})", + flush=True, + ) + print(f"done -> {args.out_csv}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/src/ops_model/interpretability/classifier/train.py b/src/ops_model/interpretability/classifier/train.py new file mode 100644 index 0000000..5d7784a --- /dev/null +++ b/src/ops_model/interpretability/classifier/train.py @@ -0,0 +1,2555 @@ +#!/usr/bin/env python +"""Train a set-transformer classifier on per-cell embeddings. + +Architecture +------------ +1. For each perturbation (gene), sample ``n_cells`` cell embeddings per channel. +2. Add a learned channel embedding to every cell. +3. Per-channel Set Transformer (ISAB + PMA) pools each channel's cells + into a single vector. +4. A cross-channel Set Transformer aggregates channel vectors into a + final representation. +5. A linear head predicts the gene class. + +Uses inducing-point set attention (ISAB) so that cost is O(N·m) instead +of O(N²), where m is the number of inducing points. + +Usage +----- +.. code-block:: bash + + python train_set_classifier.py + python train_set_classifier.py n_cells=500 model.d_model=256 +""" + +from __future__ import annotations + +import gc +import hashlib +import json +import math +import os +import random +import time +from pathlib import Path + +import hydra +import numpy as np +import pandas as pd +import pyarrow.parquet as pq +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint +from omegaconf import DictConfig, OmegaConf +from torch.utils.data import DataLoader, Dataset, Sampler +from tqdm import tqdm + +_CONFIG_DIR = str( + Path(os.environ.get("CONFIG_PATH", Path(__file__).resolve().parent / "configs")) +) + +# --------------------------------------------------------------------------- +# Set Transformer building blocks +# --------------------------------------------------------------------------- + + +def _cond_kw(cond: torch.Tensor | None) -> dict[str, torch.Tensor]: + """Build kwargs dict for AdaLNMAB (empty dict when no conditioning).""" + return {"cond": cond} if cond is not None else {} + + +class CrossAttnBlock(nn.Module): + """Pre-norm cross-attention block. + + Only the query stream (X) is normalized; KV (Y) passes through raw. + H = X + drop(Attn(LN(X), Y, Y)) + out = H + drop(FF(LN(H))) + """ + + def __init__( + self, d_model: int, n_heads: int, d_ff: int | None = None, dropout: float = 0.0 + ): + super().__init__() + d_ff = d_ff or 4 * d_model + self.attn = nn.MultiheadAttention( + d_model, n_heads, dropout=dropout, batch_first=True + ) + self.norm_q = nn.LayerNorm(d_model) + self.norm_ff = nn.LayerNorm(d_model) + self.ff = nn.Sequential( + nn.Linear(d_model, d_ff), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(d_ff, d_model), + nn.Dropout(dropout), + ) + + def forward( + self, + x: torch.Tensor, + y: torch.Tensor, + y_key_padding_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + x_norm = self.norm_q(x) + h = ( + x + + self.attn( + x_norm, y, y, key_padding_mask=y_key_padding_mask, need_weights=False + )[0] + ) + return h + self.ff(self.norm_ff(h)) + + +class AdaLNCrossAttnBlock(nn.Module): + """Pre-norm cross-attention block with adaptive LayerNorm conditioning. + + Only the query stream is normalized; KV passes through raw. + AdaLN modulates the two LN outputs (query norm, FF norm) via + per-sample scale/shift from a conditioning vector. + + Accepts cond as either (B, D) for global conditioning (broadcast across + all query positions) or (B, N, D) for per-token conditioning. + """ + + def __init__( + self, d_model: int, n_heads: int, d_ff: int | None = None, dropout: float = 0.0 + ): + super().__init__() + d_ff = d_ff or 4 * d_model + self.attn = nn.MultiheadAttention( + d_model, n_heads, dropout=dropout, batch_first=True + ) + self.norm_q = nn.LayerNorm(d_model, elementwise_affine=False) + self.norm_ff = nn.LayerNorm(d_model, elementwise_affine=False) + self.ff = nn.Sequential( + nn.Linear(d_model, d_ff), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(d_ff, d_model), + nn.Dropout(dropout), + ) + self.adaln_proj = nn.Linear(d_model, 4 * d_model) + + def forward( + self, + x: torch.Tensor, + y: torch.Tensor, + y_key_padding_mask: torch.Tensor | None = None, + cond: torch.Tensor | None = None, + ) -> torch.Tensor: + if cond is not None: + params = self.adaln_proj(cond) + if params.dim() == 2: + params = params.unsqueeze(1) + gamma_q, beta_q, gamma_ff, beta_ff = params.chunk(4, dim=-1) + x_norm = self.norm_q(x) * (1 + gamma_q) + beta_q + h = ( + x + + self.attn( + x_norm, + y, + y, + key_padding_mask=y_key_padding_mask, + need_weights=False, + )[0] + ) + return h + self.ff(self.norm_ff(h) * (1 + gamma_ff) + beta_ff) + else: + x_norm = self.norm_q(x) + h = ( + x + + self.attn( + x_norm, + y, + y, + key_padding_mask=y_key_padding_mask, + need_weights=False, + )[0] + ) + return h + self.ff(self.norm_ff(h)) + + +class SelfAttnBlock(nn.Module): + """Pre-norm self-attention block. + + The single input is normalized once, then used for Q, K, V. + H = X + drop(Attn(LN(X), LN(X), LN(X))) + out = H + drop(FF(LN(H))) + """ + + def __init__( + self, d_model: int, n_heads: int, d_ff: int | None = None, dropout: float = 0.0 + ): + super().__init__() + d_ff = d_ff or 4 * d_model + self.attn = nn.MultiheadAttention( + d_model, n_heads, dropout=dropout, batch_first=True + ) + self.norm_attn = nn.LayerNorm(d_model) + self.norm_ff = nn.LayerNorm(d_model) + self.ff = nn.Sequential( + nn.Linear(d_model, d_ff), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(d_ff, d_model), + nn.Dropout(dropout), + ) + + def forward( + self, + x: torch.Tensor, + key_padding_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + x_norm = self.norm_attn(x) + h = ( + x + + self.attn( + x_norm, + x_norm, + x_norm, + key_padding_mask=key_padding_mask, + need_weights=False, + )[0] + ) + return h + self.ff(self.norm_ff(h)) + + +class ISAB(nn.Module): + """Inducing-point Set Attention Block. + + Reduces O(N²) self-attention to O(N·m) by routing through *m* + learnable inducing points. Both sub-blocks are cross-attention: + mab1 attends inducing→input, mab2 attends input→inducing summary. + """ + + def __init__( + self, + d_model: int, + n_heads: int, + n_inducing: int, + d_ff: int | None = None, + adaln: bool = False, + dropout: float = 0.0, + ): + super().__init__() + self.inducing = nn.Parameter(torch.randn(1, n_inducing, d_model) * 0.02) + cross_cls = AdaLNCrossAttnBlock if adaln else CrossAttnBlock + self.cross1 = cross_cls(d_model, n_heads, d_ff, dropout=dropout) + self.cross2 = cross_cls(d_model, n_heads, d_ff, dropout=dropout) + + def forward( + self, + x: torch.Tensor, + key_padding_mask: torch.Tensor | None = None, + cond: torch.Tensor | None = None, + ) -> torch.Tensor: + # cross1: inducing points query the input tokens. When cond is per-token + # (B, N, D), average to (B, D) for the m inducing-point queries. + cond1 = cond.mean(dim=1) if cond is not None and cond.dim() == 3 else cond + h = self.cross1( + self.inducing.expand(x.size(0), -1, -1), + x, + y_key_padding_mask=key_padding_mask, + **_cond_kw(cond1), + ) + # cross2: input tokens query the inducing summaries. Per-token cond + # (B, N, D) is passed directly so each token gets its own modulation. + return self.cross2(x, h, **_cond_kw(cond)) + + +class PMA(nn.Module): + """Pooling by Multihead Attention (cross-attention from seeds to set).""" + + def __init__( + self, + d_model: int, + n_heads: int, + n_seeds: int = 1, + d_ff: int | None = None, + adaln: bool = False, + dropout: float = 0.0, + ): + super().__init__() + self.seeds = nn.Parameter(torch.randn(1, n_seeds, d_model) * 0.02) + cross_cls = AdaLNCrossAttnBlock if adaln else CrossAttnBlock + self.cross = cross_cls(d_model, n_heads, d_ff, dropout=dropout) + + def forward( + self, + x: torch.Tensor, + key_padding_mask: torch.Tensor | None = None, + cond: torch.Tensor | None = None, + ) -> torch.Tensor: + # Seeds are learnable queries, not input tokens. When cond is per-token + # (B, N, D), average to (B, D) for the seed queries. + cond_pool = cond.mean(dim=1) if cond is not None and cond.dim() == 3 else cond + return self.cross( + self.seeds.expand(x.size(0), -1, -1), + x, + y_key_padding_mask=key_padding_mask, + **_cond_kw(cond_pool), + ) + + +class MILAttentionPool(nn.Module): + """Gated-attention multiple-instance-learning pooling (Ilse et al. 2018). + + Computes a scalar attention weight per instance, + ``a_i = softmax_i(w^T (tanh(V h_i) ⊙ sigmoid(U h_i)))``, and returns the + convex combination ``sum_i a_i h_i``. + + Unlike PMA, each ``a_i`` depends only on cell ``i`` (no inter-cell + interaction) and pooling is a plain convex combination, so the weights are a + clean, faithful per-cell importance score for downstream cell selection. The + most recent (masked, softmaxed) weights are stored on ``last_attn`` (B, N) + so selection code can read them off after a forward pass. + """ + + def __init__(self, d_model: int, d_attn: int | None = None): + super().__init__() + d_attn = d_attn or d_model + self.V = nn.Linear(d_model, d_attn) + self.U = nn.Linear(d_model, d_attn) + self.w = nn.Linear(d_attn, 1) + self.last_attn = None + + def forward(self, x, key_padding_mask=None): + scores = self.w(torch.tanh(self.V(x)) * torch.sigmoid(self.U(x))).squeeze(-1) + if key_padding_mask is not None: + scores = scores.masked_fill(key_padding_mask, float("-inf")) + attn = torch.softmax(scores, dim=1) + self.last_attn = attn.detach() + return torch.bmm(attn.unsqueeze(1), x).squeeze(1) + + +class ISABBlock(nn.Module): + """Stack of ISAB layers followed by pooling. + + pool_type: + - "pma": Pooling by Multihead Attention (learnable seed tokens). + - "mean": masked mean over the set axis. + """ + + def __init__( + self, + d_model: int, + n_heads: int, + n_layers: int, + n_inducing: int, + n_pool_seeds: int = 1, + d_ff: int | None = None, + adaln: bool = False, + dropout: float = 0.0, + pool_type: str = "pma", + grad_checkpoint: bool = False, + ): + super().__init__() + assert pool_type in ("pma", "mean") + self.pool_type = pool_type + self.grad_checkpoint = grad_checkpoint + self.layers = nn.ModuleList( + [ + ISAB(d_model, n_heads, n_inducing, d_ff, adaln=adaln, dropout=dropout) + for _ in range(n_layers) + ] + ) + if pool_type == "pma": + self.pool: nn.Module | None = PMA( + d_model, n_heads, n_pool_seeds, d_ff, adaln=adaln, dropout=dropout + ) + else: + self.pool = None + self.final_norm = nn.LayerNorm(d_model) + + def forward( + self, + x: torch.Tensor, + key_padding_mask: torch.Tensor | None = None, + cond: torch.Tensor | None = None, + ) -> torch.Tensor: + for layer in self.layers: + if self.grad_checkpoint and self.training: + # Discard the ISAB layer's activations and recompute them in + # backward -- trades ~25-30% compute for much lower memory, + # enabling larger bags. Identical math; loss is unaffected. + x = torch.utils.checkpoint.checkpoint( + layer, x, key_padding_mask, cond, use_reentrant=False + ) + else: + x = layer(x, key_padding_mask=key_padding_mask, cond=cond) + if self.pool_type == "pma": + assert self.pool is not None + pooled = self.pool(x, key_padding_mask=key_padding_mask, cond=cond).squeeze( + 1 + ) + else: + # key_padding_mask convention: True = pad. + if key_padding_mask is not None: + valid = (~key_padding_mask).unsqueeze(-1).to(x.dtype) + pooled = (x * valid).sum(dim=1) / valid.sum(dim=1).clamp(min=1) + else: + pooled = x.mean(dim=1) + return self.final_norm(pooled) + + +# --------------------------------------------------------------------------- +# Full classifier +# --------------------------------------------------------------------------- + + +class CosineClassifier(nn.Module): + """Cosine-similarity head with learned temperature. + + Normalizes both the input and weight vectors so the logit for each class + is ``scale * cos(x, w_c)``. + """ + + def __init__(self, in_dim: int, num_classes: int, init_scale: float = 20.0): + super().__init__() + self.weight = nn.Parameter(torch.randn(num_classes, in_dim)) + self.log_scale = nn.Parameter(torch.tensor(math.log(init_scale))) + nn.init.normal_(self.weight, std=0.01) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = F.normalize(x, dim=-1) + w = F.normalize(self.weight, dim=-1) + return torch.exp(self.log_scale) * (x @ w.t()) + + +# --------------------------------------------------------------------------- +# Dataset +# --------------------------------------------------------------------------- + + +CellIndex = dict[tuple[int, int], torch.Tensor] +"""Mapping from (gene_idx, channel_idx) → (N, D) embedding tensor.""" + +CellDumpMetaChunk = dict[str, list] +"""One embedding-chunk's metadata: column → flat list of per-cell values (Python scalars/str).""" + +CellDumpMetaLoL = dict[str, list[list]] +"""Column → list of segments; each segment is a list aligned with one concatenated emb block.""" + +CellDumpMetaFlat = dict[str, list] +"""Column → one flat list per cell (used after ``load_val_dataset``).""" + +VAL_DUMP_META_COLS = ( + "well", + "x_pheno", + "y_pheno", + "segmentation_id", + "channel_type", + "index", +) +"""Parquet columns read only when ``load_val_dump_metadata`` is True (plus ``experiment``, already required).""" + +CELL_METADATA_DUMP_KEY = "cell_metadata" +CELL_DUMP_KEYS = ( + "experiment", + "well", + "x_pheno", + "y_pheno", + "segmentation_id", + "channel_type", + "index", +) + + +def _meta_lol_n_cells(m: CellDumpMetaLoL) -> int: + return sum(len(seg) for seg in m["experiment"]) + + +def _is_cell_meta_lol(dm: CellDumpMetaLoL | CellDumpMetaFlat) -> bool: + ex = dm["experiment"] + return bool(ex) and isinstance(ex[0], list) + + +def _pack_val_dump_meta(meta: pd.DataFrame, rows: np.ndarray) -> CellDumpMetaChunk: + """Build one chunk's metadata as Python lists (no numpy allocation for the lists).""" + sub = meta.iloc[rows] + seg_f = np.asarray( + pd.to_numeric(sub["segmentation_id"], errors="coerce"), + dtype=np.float64, + ) + seg_int = np.nan_to_num(seg_f, nan=-1.0).astype(np.int64) + chix_f = np.asarray( + pd.to_numeric(sub["index"], errors="coerce"), + dtype=np.float64, + ) + chix_int = np.nan_to_num(chix_f, nan=-1.0).astype(np.int64) + return { + "experiment": sub["experiment"].astype("string").fillna("").tolist(), + "well": sub["well"].astype("string").fillna("").tolist(), + "x_pheno": [float(x) for x in sub["x_pheno"].tolist()], + "y_pheno": [float(x) for x in sub["y_pheno"].tolist()], + "segmentation_id": [int(x) for x in seg_int.tolist()], + "channel_type": sub["channel_type"].astype("string").fillna("").tolist(), + "index": [int(x) for x in chix_int.tolist()], + } + + +def _chunks_to_lol(chunks: list[CellDumpMetaChunk]) -> CellDumpMetaLoL: + """Many parquet chunks for one (g,c,e) → column → [seg0, seg1, …] (extend-style concat).""" + cols = chunks[0].keys() + return {col: [list(c[col]) for c in chunks] for col in cols} + + +def _merge_lol_for_experiments(meta_list: list[CellDumpMetaLoL]) -> CellDumpMetaLoL: + """Concat along embedding axis: extend each column's list-of-lists with more segments.""" + cols = meta_list[0].keys() + out: CellDumpMetaLoL = {col: [] for col in cols} + for m in meta_list: + for col in cols: + out[col].extend(m[col]) + return out + + +def _permute_meta_lol(m: CellDumpMetaLoL, perm: torch.Tensor) -> CellDumpMetaLoL: + """Flatten, permute, store as a single segment (still list-of-lists: one inner list).""" + flat = {col: [x for seg in m[col] for x in seg] for col in m} + idx = perm.detach().cpu().long().tolist() + return {col: [[flat[col][i] for i in idx]] for col in m} + + +def _merge_experiments_val_with_meta( + gce: dict[tuple[int, int, int], torch.Tensor], + meta_gce: dict[tuple[int, int, int], CellDumpMetaLoL], + max_cells_per_group: int | None, + desc: str, +) -> tuple[CellIndex, dict[tuple[int, int], CellDumpMetaLoL]]: + """Like merging val (gene,ch,exp)→(gene,ch), keeping per-cell metadata aligned.""" + gc_tensor_lists: dict[tuple[int, int], list[torch.Tensor]] = {} + gc_meta_lists: dict[tuple[int, int], list[CellDumpMetaLoL]] = {} + for (g_idx, ch_idx, _exp_idx), t in gce.items(): + gc_key = (g_idx, ch_idx) + gc_tensor_lists.setdefault(gc_key, []).append(t) + gc_meta_lists.setdefault(gc_key, []).append(meta_gce[(g_idx, ch_idx, _exp_idx)]) + + index: CellIndex = {} + meta_index: dict[tuple[int, int], CellDumpMetaLoL] = {} + n_capped = 0 + for gc_key, tensors in tqdm( + gc_tensor_lists.items(), desc=desc, total=len(gc_tensor_lists), unit="group" + ): + metas = gc_meta_lists[gc_key] + combined = torch.cat(tensors) if len(tensors) > 1 else tensors[0] + combined_meta = _merge_lol_for_experiments(metas) + assert len(combined) == _meta_lol_n_cells(combined_meta), gc_key + if max_cells_per_group is not None and len(combined) > max_cells_per_group: + perm = torch.randperm(len(combined))[:max_cells_per_group] + combined = combined[perm] + combined_meta = _permute_meta_lol(combined_meta, perm) + n_capped += 1 + index[gc_key] = combined + meta_index[gc_key] = combined_meta + if max_cells_per_group is not None and n_capped > 0: + print(f" Capped {n_capped} groups to {max_cells_per_group} cells") + return index, meta_index + + +class RepeatSampler(Sampler[int]): + """Repeats each index exactly ``multiplier`` times per epoch, shuffled.""" + + def __init__(self, n_items: int, multiplier: int = 1): + self.n_items = n_items + self.multiplier = multiplier + + def __len__(self) -> int: + return self.n_items * self.multiplier + + def __iter__(self): + indices = list(range(self.n_items)) * self.multiplier + random.shuffle(indices) + yield from indices + + +# --------------------------------------------------------------------------- +# Mixed-channel mode +# --------------------------------------------------------------------------- + + +def _compute_ch_indices(ch_ids: torch.Tensor) -> dict[int, torch.Tensor]: + """Group cell positions by channel: ch_idx → LongTensor of indices into the pool.""" + if len(ch_ids) == 0: + return {} + unique = torch.unique(ch_ids).tolist() + return {int(c): (ch_ids == c).nonzero(as_tuple=False).squeeze(-1) for c in unique} + + +def _normalize_cps_choices( + channels_per_set: int | list[int | None] | None, +) -> list[int | None]: + """Normalize ``channels_per_set`` to a list of per-set choices. + + A single int or None acts as a one-element list. A list value is used as-is, + and ``None`` entries within it represent "use all channels for this set". + """ + if channels_per_set is None: + return [None] + if isinstance(channels_per_set, int): + return [channels_per_set] + return list(channels_per_set) + + +class MixedChannelDataset(Dataset): + """Each item samples n_cells across all channels proportionally. + + Tracks per-cell channel IDs so the model can apply channel conditioning. + + ``channels_per_set`` controls how many channels each set is sampled from: + + - ``None`` (default): no subsetting — sample from cells across all channels. + - ``int`` (e.g. ``3``): every set picks exactly that many random channels. + - ``list[int | None]`` (e.g. ``[1, 2, None]``): every set independently + picks one value from the list; ``None`` in the list means "all channels" + for that set. + + If a gene has fewer channels than the chosen value, all available are used. + Applies to whichever dataset (train and/or val) it is set on. + """ + + def __init__( + self, + cell_index: CellIndex, + emb_dim: int, + gene_to_idx: dict[str, int], + perturbation_list: list[str], + n_cells: int, + cell_dump_index: dict[tuple[int, int], CellDumpMetaLoL] | None = None, + label_remap: dict[int, int] | None = None, + replacement: bool = True, + channels_per_set: int | list[int | None] | None = None, + ): + self.emb_dim = emb_dim + self.gene_to_idx = gene_to_idx + self.perturbation_list = perturbation_list + self.n_cells = n_cells + self.label_remap = label_remap + self.replacement = replacement + self.channels_per_set = channels_per_set + self._cps_choices = _normalize_cps_choices(channels_per_set) + + self._gene_pools: dict[int, torch.Tensor] = {} + self._gene_ch_ids: dict[int, torch.Tensor] = {} + self._gene_ch_indices: dict[int, dict[int, torch.Tensor]] = {} + self._gene_dump_meta: dict[int, CellDumpMetaLoL | CellDumpMetaFlat] | None = ( + None + ) + if cell_dump_index is not None: + self._gene_dump_meta = {} + + for g_name in tqdm(perturbation_list, desc="Building gene pools", unit="gene"): + g_idx = gene_to_idx[g_name] + if g_idx in self._gene_pools: + continue + emb_chunks: list[torch.Tensor] = [] + ch_chunks: list[torch.Tensor] = [] + meta_chunks: list[CellDumpMetaLoL] = [] + for key, t in cell_index.items(): + if key[0] == g_idx: + emb_chunks.append(t) + ch_chunks.append(torch.full((len(t),), key[1], dtype=torch.long)) + if cell_dump_index is not None: + m = cell_dump_index[key] + assert len(t) == _meta_lol_n_cells(m), ( + key, + len(t), + _meta_lol_n_cells(m), + ) + meta_chunks.append(m) + if emb_chunks: + self._gene_pools[g_idx] = torch.cat(emb_chunks) + self._gene_ch_ids[g_idx] = torch.cat(ch_chunks) + if cell_dump_index is not None: + assert self._gene_dump_meta is not None + self._gene_dump_meta[g_idx] = _merge_lol_for_experiments( + meta_chunks + ) + else: + self._gene_pools[g_idx] = torch.zeros(0, emb_dim) + self._gene_ch_ids[g_idx] = torch.zeros(0, dtype=torch.long) + if cell_dump_index is not None: + assert self._gene_dump_meta is not None + self._gene_dump_meta[g_idx] = {k: [] for k in CELL_DUMP_KEYS} + self._gene_ch_indices[g_idx] = _compute_ch_indices(self._gene_ch_ids[g_idx]) + + self._t_randperm = 0.0 + self._t_index = 0.0 + self._t_total = 0.0 + self._n_calls = 0 + + def reset_timers(self) -> None: + self._t_randperm = 0.0 + self._t_index = 0.0 + self._t_total = 0.0 + self._n_calls = 0 + + def print_timers(self) -> None: + print( + f" [__getitem__] {self._n_calls} calls: " + f"total={self._t_total:.3f}s " + f"randperm={self._t_randperm:.3f}s " + f"index={self._t_index:.3f}s " + f"other={self._t_total - self._t_randperm - self._t_index:.3f}s" + ) + + def set_n_cells(self, n_cells: int) -> None: + self.n_cells = n_cells + + def __len__(self) -> int: + return len(self.perturbation_list) + + def __getitem__( + self, idx: int + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: + """Returns (cell_embs, ch_ids, mask, label). + + cell_embs: (n_cells, D) + ch_ids: (n_cells,) long — channel index per cell + mask: (n_cells,) bool + """ + t_start = time.perf_counter() + g_idx = self.gene_to_idx[self.perturbation_list[idx]] + label = self.label_remap[g_idx] if self.label_remap is not None else g_idx + pool = self._gene_pools[g_idx] + ch_ids = self._gene_ch_ids[g_idx] + + if len(pool) > 0: + cps_choice = random.choice(self._cps_choices) + if cps_choice is not None: + ch_indices = self._gene_ch_indices[g_idx] + avail = list(ch_indices.keys()) + n_sel = min(cps_choice, len(avail)) + if n_sel < len(avail): + sel = random.sample(avail, n_sel) + sel_idx = torch.cat([ch_indices[c] for c in sel]) + pool = pool[sel_idx] + ch_ids = ch_ids[sel_idx] + + n_available = len(pool) + + if n_available >= self.n_cells: + t0 = time.perf_counter() + if self.replacement: + selected = torch.randint(n_available, (self.n_cells,)) + else: + selected = torch.randperm(n_available)[: self.n_cells] + t1 = time.perf_counter() + embs_out = pool[selected] + ch_out = ch_ids[selected] + t2 = time.perf_counter() + self._t_randperm += t1 - t0 + self._t_index += t2 - t1 + self._t_total += t2 - t_start + self._n_calls += 1 + return ( + embs_out, + ch_out, + torch.ones(self.n_cells, dtype=torch.bool), + label, + ) + + padded = torch.zeros(self.n_cells, self.emb_dim) + padded_ch = torch.zeros(self.n_cells, dtype=torch.long) + if n_available > 0: + padded[:n_available] = pool + padded_ch[:n_available] = ch_ids + mask = torch.zeros(self.n_cells, dtype=torch.bool) + mask[:n_available] = True + self._t_total += time.perf_counter() - t_start + self._n_calls += 1 + return padded, padded_ch, mask, label + + +class MixedChannelClassifier(nn.Module): + """Single-level set transformer: ISAB pools all cells → classify. + + Channel conditioning modes: + - "add": add learned channel embedding to each cell. + - "concat": concatenate channel embedding with projected cell embedding + and linearly project back to d_model. + - "adaln": modulate LayerNorm via adaptive scale/shift from the *mean* + channel embedding (global conditioning, shape (B, D)). + - "adaln-token": modulate LayerNorm via adaptive scale/shift from + *per-token* channel embeddings (shape (B, N, D)). Each cell gets + conditioning from its own channel rather than an average. + - "none": no channel conditioning. + """ + + def __init__( + self, + emb_dim: int, + n_classes: int, + n_channels: int = 0, + d_model: int = 256, + n_heads: int = 4, + n_layers: int = 2, + n_inducing: int = 32, + d_ff: int | None = None, + dropout: float = 0.1, + cosine_classifier: bool = False, + channel_conditioning: str = "none", + pool_type: str = "pma", + grad_checkpoint: bool = False, + ): + super().__init__() + assert channel_conditioning in ("add", "concat", "adaln", "adaln-token", "none") + self.channel_conditioning = channel_conditioning + + self.input_proj = nn.Linear(emb_dim, d_model) + if channel_conditioning in ("add", "concat", "adaln", "adaln-token"): + self.channel_embeddings = nn.Embedding(n_channels, d_model) + if channel_conditioning == "concat": + self.concat_proj = nn.Linear(2 * d_model, d_model) + + use_adaln = channel_conditioning in ("adaln", "adaln-token") + self.encoder = ISABBlock( + d_model=d_model, + n_heads=n_heads, + n_layers=n_layers, + n_inducing=n_inducing, + n_pool_seeds=1, + d_ff=d_ff, + adaln=use_adaln, + dropout=dropout, + pool_type=pool_type, + grad_checkpoint=grad_checkpoint, + ) + classifier: nn.Module + if cosine_classifier: + classifier = CosineClassifier(d_model, n_classes) + else: + classifier = nn.Linear(d_model, n_classes) + self.head = nn.Sequential(nn.Dropout(dropout), classifier) + + def encode( + self, + embs: torch.Tensor, + ch_ids: torch.Tensor, + masks: torch.Tensor, + ) -> torch.Tensor: + """Pool a set of cells into the penultimate representation. + + Args: + embs: (B, N, D) cell embeddings. + ch_ids: (B, N) long — per-cell channel index. + masks: (B, N) bool — True = valid. + + Returns: + Pooled set vector of shape (B, d_model) — the input to the head. + """ + x = self.input_proj(embs) + + cond = None + if self.channel_conditioning == "add": + x = x + self.channel_embeddings(ch_ids) + elif self.channel_conditioning == "concat": + x = self.concat_proj( + torch.cat([x, self.channel_embeddings(ch_ids)], dim=-1) + ) + elif self.channel_conditioning == "adaln": + cond = self.channel_embeddings(ch_ids).mean(dim=1) + elif self.channel_conditioning == "adaln-token": + cond = self.channel_embeddings(ch_ids) + + return self.encoder(x, key_padding_mask=~masks, cond=cond) + + def forward( + self, + embs: torch.Tensor, + ch_ids: torch.Tensor, + masks: torch.Tensor, + ) -> torch.Tensor: + """Forward pass. + + Args: + embs: (B, N, D) cell embeddings. + ch_ids: (B, N) long — per-cell channel index. + masks: (B, N) bool — True = valid. + + Returns: + Logits of shape (B, n_classes). + """ + return self.head(self.encode(embs, ch_ids, masks)) + + +def train_one_epoch_mixed( + model: MixedChannelClassifier, + loader: DataLoader, + optimizer: torch.optim.Optimizer, + device: torch.device, + max_grad_norm: float | None = None, +) -> tuple[float, float, float]: + model.train() + total_loss = 0.0 + correct = 0 + total = 0 + grad_norm_sum = 0.0 + n_steps = 0 + + t_data_sum = 0.0 + t_transfer_sum = 0.0 + t_forward_sum = 0.0 + t_backward_sum = 0.0 + t_step_sum = 0.0 + + print( + f" [loader] len={len(loader)} batches, dataset={len(loader.dataset)} samples" + ) + + t_batch_start = time.perf_counter() + for embs, ch_ids, masks, labels in loader: + if n_steps == 0: + print( + f" [batch shape] embs={list(embs.shape)} ch_ids={list(ch_ids.shape)} masks={list(masks.shape)} labels={list(labels.shape)}" + ) + t_data = time.perf_counter() + + embs = embs.to(device) + ch_ids = ch_ids.to(device) + masks = masks.to(device) + labels = labels.to(device) + t_transfer = time.perf_counter() + + logits = model(embs, ch_ids, masks) + loss = F.cross_entropy(logits, labels) + t_forward = time.perf_counter() + + optimizer.zero_grad() + loss.backward() + if max_grad_norm is not None: + nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm) + grad_norm_sum += _grad_norm(model) + t_backward = time.perf_counter() + + n_steps += 1 + optimizer.step() + t_step = time.perf_counter() + + t_data_sum += t_data - t_batch_start + t_transfer_sum += t_transfer - t_data + t_forward_sum += t_forward - t_transfer + t_backward_sum += t_backward - t_forward + t_step_sum += t_step - t_backward + + total_loss += loss.item() * labels.size(0) + correct += (logits.argmax(dim=-1) == labels).sum().item() + total += labels.size(0) + + t_batch_start = time.perf_counter() + + print( + f" [timing] {n_steps} steps: " + f"data={t_data_sum:.2f}s " + f"transfer={t_transfer_sum:.2f}s " + f"forward={t_forward_sum:.2f}s " + f"backward={t_backward_sum:.2f}s " + f"optim_step={t_step_sum:.2f}s" + ) + + if total == 0: + raise ValueError( + "Mixed-channel training epoch saw zero samples (no batches). " + "Check for an empty dataset or train DataLoader drop_last dropping " + "all indices when len(dataset)*train_multiplier < batch_size." + ) + + return total_loss / total, correct / total, grad_norm_sum / max(n_steps, 1) + + +@torch.no_grad() +def evaluate_mixed( + model: MixedChannelClassifier, + loader: DataLoader, + device: torch.device, +) -> tuple[float, float]: + model.eval() + total_loss = 0.0 + correct = 0 + total = 0 + + for embs, ch_ids, masks, labels in loader: + embs = embs.to(device) + ch_ids = ch_ids.to(device) + masks = masks.to(device) + labels = labels.to(device) + logits = model(embs, ch_ids, masks) + loss = F.cross_entropy(logits, labels) + + total_loss += loss.item() * labels.size(0) + correct += (logits.argmax(dim=-1) == labels).sum().item() + total += labels.size(0) + + if total == 0: + raise ValueError( + "Mixed-channel eval saw zero samples (no batches). " + "Check for an empty val dataset or val DataLoader configuration." + ) + + return total_loss / total, correct / total + + +# --------------------------------------------------------------------------- +# Data loading +# --------------------------------------------------------------------------- + + +def _load_label_map( + label_map_path: str, + gene_col: str = "gene_name", + label_col: str = "pathway", +) -> dict[str, str]: + """Load an external CSV that maps gene names to class labels (e.g. pathways). + + Empty/NaN values in ``gene_col`` are filled as ``"NTC"`` to match the + ``fillna("NTC")`` convention used during data loading. + """ + df = pd.read_csv(label_map_path) + if gene_col not in df.columns: + raise ValueError(f"Column {gene_col!r} not found in {label_map_path}") + if label_col not in df.columns: + raise ValueError(f"Column {label_col!r} not found in {label_map_path}") + df[gene_col] = df[gene_col].fillna("NTC") + return dict(zip(df[gene_col], df[label_col])) + + +def _build_channel_label(meta: pd.DataFrame) -> pd.Series: + """Build channel label from biological annotations, falling back to ``name``. + + Rows with ``name == "Phase2D"`` are always labeled ``Phase2D`` regardless of + annotation columns (some batches store sentinel strings like ``'no label'`` + instead of NULL there). Otherwise: if either + ``biological_annotation.organelle`` or ``biological_annotation.marker`` is + non-null, the channel label is their concatenation (separated by ``_``, + skipping nulls). Otherwise falls back to the ``name`` column. + """ + name = meta["name"] + org = meta["biological_annotation.organelle"] + marker = meta["biological_annotation.marker"] + has_annotation = org.notna() | marker.notna() + + parts = org.fillna("").str.cat(marker.fillna(""), sep="_").str.strip("_") + label = parts.where(has_annotation, name) + return label.where(name != "Phase2D", "Phase2D") + + +def _resolve_label_map_genes( + unique_genes: list[str], + gene_to_label: dict[str, str], + fallback_to_gene: bool, +) -> tuple[list[str], dict[str, str]]: + """Resolve kept genes and their effective labels under an external label map. + + Parameters + ---------- + unique_genes + All gene names present in the dataset. + gene_to_label + Mapping from gene name to class label (e.g. EBI complex), from + :func:`_load_label_map`. + fallback_to_gene + If True, genes absent from ``gene_to_label`` are kept and labeled by + their own gene name (each forms its own class). If False, such genes are + dropped (train only on genes present in the map). + + Returns + ------- + tuple[list[str], dict[str, str]] + ``(kept_genes, effective_gene_to_label)`` where every gene in + ``kept_genes`` has an entry in ``effective_gene_to_label``. + """ + mapped = [g for g in unique_genes if g in gene_to_label] + if fallback_to_gene: + effective = dict(gene_to_label) + for g in unique_genes: + effective.setdefault(g, g) + return sorted(unique_genes), effective + return sorted(mapped), gene_to_label + + +def _cell_stratify_val_mask( + experiment: np.ndarray, + well: np.ndarray, + segmentation_id: np.ndarray, + val_fraction: float, + seed: int, +) -> np.ndarray: + """Assign each unique cell to the val (True) or train (False) split. + + A cell is keyed by ``(experiment, well, segmentation_id)``. The assignment is + a deterministic hash of that key, so every row belonging to a given cell + receives the same assignment regardless of row ordering, which row group it + came from, or which parquet (modality) it lives in. This prevents train/val + leakage when a single cell contributes multiple embeddings (e.g. one row per + fluorescent marker / channel, or the same cell imaged in phenotyping, Cell + Painting and 4i). + + Parameters + ---------- + experiment, well, segmentation_id + Per-row cell identifier components (any dtype; coerced to a canonical + string form — ``segmentation_id`` is normalized to an integer). + val_fraction + Target fraction of *cells* (not rows) assigned to val. + seed + Seed folded into the hash so the split is reproducible yet re-seedable. + + Returns + ------- + np.ndarray + Boolean mask, True where the row's cell is assigned to val. + """ + seg_num = pd.Series(pd.to_numeric(pd.Series(segmentation_id), errors="coerce")) + seg = seg_num.fillna(-1).astype(np.int64).astype(str) + keys = ( + pd.Series(experiment) + .astype(str) + .str.cat([pd.Series(well).astype(str), seg], sep="|") + ) + # Deterministic per-cell hash: factorize to unique cell keys, hash each once + # (few uniques per call), then broadcast back to rows. blake2b is stable + # across processes (unlike the salted builtin hash()). + codes, uniques = pd.factorize(keys, sort=False) + threshold = round(val_fraction * 1_000_000) + unique_val = np.array( + [ + int.from_bytes( + hashlib.blake2b(f"{seed}|{k}".encode(), digest_size=8).digest(), + "big", + ) + % 1_000_000 + < threshold + for k in uniques + ], + dtype=bool, + ) + return unique_val[codes] + + +def load_data( + parquet_entries: list[dict], + max_row_groups: int | None, + max_cells_per_group: int | None, + val_fraction: float = 0.2, + seed: int = 42, + cell_stratify: bool = False, + max_genes: int | None = None, + max_channels: int | None = None, + min_cells_per_group: int | None = None, + min_cells_drop_val: bool = False, + load_val_dump_metadata: bool = False, + load_train_dump_metadata: bool = False, + exclude_channel_names: list[str] | None = None, + include_channel_names: list[str] | None = None, + z_standardize: bool = True, + z_standardize_control_only: bool = False, +) -> tuple[ + CellIndex, + CellIndex, + int, + dict[str, int], + dict[str, int], + dict[tuple[int, int], CellDumpMetaLoL] | None, + dict[tuple[int, int], CellDumpMetaLoL] | None, +]: + """Load embeddings from parquet files into train/val :data:`CellIndex` dicts. + + Each row is randomly assigned to train or val as it is read, so no + post-hoc pass over the full data is needed. Null ``gene_name`` + values are filled as ``"NTC"`` (non-targeting control). + + Args: + parquet_entries: List of dicts, each with ``"path"`` (str) and optional + per-file filters ``"exclude_experiments"`` (list[str] | None) and + ``"exclude_fluorescent_experiments"`` (list[str] | None). + exclude_channel_names: If set, drop any row whose channel label (same string + as :func:`_build_channel_label` / keys in ``channel_to_idx``) is in this + list. Matching is exact and case-sensitive. Applied globally to all files. + include_channel_names: If set, keep ONLY rows whose channel label is in this + list (same string semantics as ``exclude_channel_names``). Applied before + the exclude filter, so a channel in both lists is dropped. Use e.g. + ``["Phase2D"]`` to train a single-channel phase-only classifier. + z_standardize: If True (default), z-standardize embeddings per + (channel, experiment). If False, skip standardization entirely + and return raw embeddings. + z_standardize_control_only: If True, compute z-standardization statistics + (mean/std per channel×experiment) only from NTC (non-targeting control) + cells. All cells are still standardized using these control-derived stats. + Ignored when ``z_standardize`` is False. + + Returns: + (train_index, val_index, emb_dim, gene_to_idx, channel_to_idx, + val_cell_dump_meta, train_cell_dump_meta). + Each dump meta is None unless the corresponding ``load_*_dump_metadata`` flag is True. + """ + base_meta_cols = [ + "gene_name", + "experiment", + "name", + "biological_annotation.organelle", + "biological_annotation.marker", + ] + if load_val_dump_metadata or load_train_dump_metadata: + base_meta_cols = base_meta_cols + list(VAL_DUMP_META_COLS) + + any_fluor_exclude = any( + entry.get("exclude_fluorescent_experiments") for entry in parquet_entries + ) + if any_fluor_exclude and "channel_type" not in base_meta_cols: + base_meta_cols.append("channel_type") + + exclude_set: frozenset[str] | None = ( + frozenset(exclude_channel_names) if exclude_channel_names else None + ) + if exclude_set is not None: + print(f"Excluding channel labels (exact match): {sorted(exclude_set)}") + include_set: frozenset[str] | None = ( + frozenset(include_channel_names) if include_channel_names else None + ) + if include_set is not None: + print(f"Including ONLY channel labels (exact match): {sorted(include_set)}") + if not z_standardize: + print("Z-standardization: DISABLED (returning raw embeddings)") + elif z_standardize_control_only: + print("Z-standardization: computing stats from NTC (control) cells only") + + rng = np.random.RandomState(seed) + gene_to_idx: dict[str, int] = {} + channel_to_idx: dict[str, int] = {} + experiment_to_idx: dict[str, int] = {} + train_buf: dict[tuple[int, int, int], list[torch.Tensor]] = {} + val_buf: dict[tuple[int, int, int], list[torch.Tensor]] = {} + val_dump_buf: dict[tuple[int, int, int], list[CellDumpMetaChunk]] = {} + train_dump_buf: dict[tuple[int, int, int], list[CellDumpMetaChunk]] = {} + ce_sums: dict[tuple[int, int], torch.Tensor] = {} + ce_sq_sums: dict[tuple[int, int], torch.Tensor] = {} + ce_counts: dict[tuple[int, int], int] = {} + emb_dim: int | None = None + total_rows = 0 + t0 = time.time() + + rg_tasks: list[tuple] = [] + for entry in parquet_entries: + parquet_path = entry["path"] + file_exclude_exp: frozenset[str] | None = ( + frozenset(entry["exclude_experiments"]) + if entry.get("exclude_experiments") + else None + ) + file_exclude_fluor: frozenset[str] | None = ( + frozenset(entry["exclude_fluorescent_experiments"]) + if entry.get("exclude_fluorescent_experiments") + else None + ) + file_col_remap: dict[str, str] = dict(entry.get("column_remap", {})) + pf = pq.ParquetFile(parquet_path) + n_rg_file = pf.metadata.num_row_groups + if max_row_groups is not None: + remaining = max_row_groups - len(rg_tasks) + if remaining <= 0: + break + n_rg_file = min(n_rg_file, remaining) + print( + f"Loading {n_rg_file}/{pf.metadata.num_row_groups} " + f"row groups from {parquet_path}" + ) + if file_exclude_exp is not None: + print(f" Excluding experiments: {sorted(file_exclude_exp)}") + if file_exclude_fluor is not None: + print( + f" Excluding fluorescent channels from: {sorted(file_exclude_fluor)}" + ) + schema_names = {f.name for f in pf.schema_arrow} + _legacy = "index" not in schema_names and "channel_index" in schema_names + fmc = list(base_meta_cols) + if _legacy and "index" in fmc: + fmc = ["channel_index" if c == "index" else c for c in fmc] + remap_reverse: dict[str, str] = {} + for canonical, source in file_col_remap.items(): + if source in schema_names: + fmc = [source if c == canonical else c for c in fmc] + remap_reverse[source] = canonical + missing = frozenset(c for c in fmc if c not in schema_names) + _optional_cols = frozenset(VAL_DUMP_META_COLS) + required_missing = missing - _optional_cols + if cell_stratify: + # These are otherwise optional (part of VAL_DUMP_META_COLS) but are + # mandatory for the cell-level split — fail loud rather than silently + # collapsing every cell into one hash bucket via the -1 default. + required_missing = required_missing | ( + missing & {"well", "segmentation_id"} + ) + if required_missing: + raise KeyError( + f"Required columns missing from {parquet_path}: {sorted(required_missing)}. " + f"Available: {sorted(schema_names)}" + ) + if missing: + fmc = [c for c in fmc if c in schema_names] + print( + f" Columns missing from schema (will fill defaults): {sorted(missing)}" + ) + for rg_idx in range(n_rg_file): + rg_tasks.append( + ( + pf, + rg_idx, + fmc, + _legacy, + file_exclude_exp, + file_exclude_fluor, + missing, + remap_reverse, + ) + ) + + _MISSING_COL_DEFAULTS: dict[str, object] = { + "x_pheno": float("nan"), + "y_pheno": float("nan"), + "segmentation_id": -1, + "index": -1, + "well": "", + "channel_type": "", + } + + n_total_rg = len(rg_tasks) + for task_idx, ( + pf, + rg_idx, + file_meta_cols, + _legacy_channel_index, + exclude_exp_set, + exclude_fluor_exp_set, + missing_cols, + col_remap, + ) in enumerate(rg_tasks): + table = pf.read_row_group(rg_idx) + meta = table.select(file_meta_cols).to_pandas() + if _legacy_channel_index and "channel_index" in meta.columns: + meta = meta.rename(columns={"channel_index": "index"}) + if col_remap: + meta = meta.rename(columns=col_remap) + for col in missing_cols: + meta[col] = _MISSING_COL_DEFAULTS.get(col, "") + meta["gene_name"] = meta["gene_name"].astype("string").fillna("NTC") + n_rows = len(meta) + + channel_label = _build_channel_label(meta) + + emb_col = table.column("embeddings") + flat = emb_col.combine_chunks().values.to_numpy(zero_copy_only=False) + dim = len(flat) // n_rows + if emb_dim is None: + emb_dim = dim + emb_2d = flat.reshape(n_rows, dim) + + if include_set is not None: + keep = channel_label.isin(include_set) + keep_arr = keep.to_numpy() + if not keep_arr.any(): + print( + f" Row group {task_idx + 1}/{n_total_rg}: 0 rows after channel include, skipping" + ) + del table, emb_col, flat, emb_2d + continue + if not keep_arr.all(): + meta = meta.loc[keep].reset_index(drop=True) + channel_label = channel_label.loc[keep].reset_index(drop=True) + emb_2d = emb_2d[keep_arr] + n_rows = len(meta) + + if exclude_set is not None: + keep = ~channel_label.isin(exclude_set) + keep_arr = keep.to_numpy() + if not keep_arr.any(): + print( + f" Row group {task_idx + 1}/{n_total_rg}: 0 rows after channel exclude, skipping" + ) + del table, emb_col, flat, emb_2d + continue + if not keep_arr.all(): + meta = meta.loc[keep].reset_index(drop=True) + channel_label = channel_label.loc[keep].reset_index(drop=True) + emb_2d = emb_2d[keep_arr] + n_rows = len(meta) + + if exclude_exp_set is not None: + keep_exp = ~meta["experiment"].isin(exclude_exp_set) + keep_exp_arr = keep_exp.to_numpy() + if not keep_exp_arr.any(): + print( + f" Row group {task_idx + 1}/{n_total_rg}: 0 rows after experiment exclude, skipping" + ) + del table, emb_col, flat, emb_2d + continue + if not keep_exp_arr.all(): + meta = meta.loc[keep_exp].reset_index(drop=True) + channel_label = channel_label.loc[keep_exp].reset_index(drop=True) + emb_2d = emb_2d[keep_exp_arr] + n_rows = len(meta) + + if exclude_fluor_exp_set is not None: + is_target_exp = meta["experiment"].isin(exclude_fluor_exp_set) + is_fluorescent = meta["channel_type"] == "fluorescent" + keep_fluor = ~(is_target_exp & is_fluorescent) + keep_fluor_arr = keep_fluor.to_numpy() + if not keep_fluor_arr.any(): + print( + f" Row group {task_idx + 1}/{n_total_rg}: " + f"0 rows after fluorescent exclude, skipping" + ) + del table, emb_col, flat, emb_2d + continue + if not keep_fluor_arr.all(): + meta = meta.loc[keep_fluor].reset_index(drop=True) + channel_label = channel_label.loc[keep_fluor].reset_index(drop=True) + emb_2d = emb_2d[keep_fluor_arr] + n_rows = len(meta) + + genes = meta["gene_name"].values + ch_labels = channel_label.values + experiments = meta["experiment"].values + for g in set(genes): + if g not in gene_to_idx: + gene_to_idx[g] = len(gene_to_idx) + for c in set(ch_labels): + if c not in channel_to_idx: + channel_to_idx[c] = len(channel_to_idx) + for e in set(experiments): + if e not in experiment_to_idx: + experiment_to_idx[e] = len(experiment_to_idx) + + gene_ids = np.array([gene_to_idx[g] for g in genes], dtype=np.int64) + ch_ids = np.array([channel_to_idx[c] for c in ch_labels], dtype=np.int64) + exp_ids = np.array([experiment_to_idx[e] for e in experiments], dtype=np.int64) + n_ch = len(channel_to_idx) + n_exp = len(experiment_to_idx) + group_keys = gene_ids * (n_ch * n_exp) + ch_ids * n_exp + exp_ids + + if cell_stratify: + # Assign whole cells (experiment, well, segmentation_id) to train/val + # so all rows of a cell — across channels/markers and modalities — + # stay on one side. No train/val leakage. + is_val = _cell_stratify_val_mask( + meta["experiment"].to_numpy(), + meta["well"].to_numpy(), + meta["segmentation_id"].to_numpy(), + val_fraction, + seed, + ) + else: + is_val = rng.random(n_rows) < val_fraction + + order = np.argsort(group_keys, kind="mergesort") + sorted_keys = group_keys[order] + split_points = np.flatnonzero(np.diff(sorted_keys)) + 1 + for chunk_indices in np.split(order, split_points): + first = chunk_indices[0] + key = (int(gene_ids[first]), int(ch_ids[first]), int(exp_ids[first])) + ce_key = (key[1], key[2]) + chunk_val = is_val[chunk_indices] + val_rows = chunk_indices[chunk_val] + train_rows = chunk_indices[~chunk_val] + if len(train_rows) > 0: + train_chunk = torch.from_numpy(emb_2d[train_rows]) + train_buf.setdefault(key, []).append(train_chunk) + use_for_stats = z_standardize and ( + not z_standardize_control_only or genes[first] == "NTC" + ) + if use_for_stats: + ce_sums[ce_key] = ce_sums.get( + ce_key, torch.zeros(dim) + ) + train_chunk.sum(0) + ce_sq_sums[ce_key] = ce_sq_sums.get( + ce_key, torch.zeros(dim) + ) + train_chunk.pow(2).sum(0) + ce_counts[ce_key] = ce_counts.get(ce_key, 0) + len(train_rows) + if load_train_dump_metadata: + train_dump_buf.setdefault(key, []).append( + _pack_val_dump_meta(meta, train_rows) + ) + if len(val_rows) > 0: + val_buf.setdefault(key, []).append(torch.from_numpy(emb_2d[val_rows])) + if load_val_dump_metadata: + val_dump_buf.setdefault(key, []).append( + _pack_val_dump_meta(meta, val_rows) + ) + + total_rows += n_rows + elapsed = time.time() - t0 + print( + f" Row group {task_idx + 1}/{n_total_rg}: {n_rows:,} rows ({elapsed:.1f}s)" + ) + + del table, emb_col, flat, emb_2d + + if emb_dim is None: + raise ValueError("No valid rows found in any parquet file.") + + if max_genes is not None and len(gene_to_idx) > max_genes: + idx_to_gene = {v: k for k, v in gene_to_idx.items()} + gene_cells: dict[int, int] = {} + for (g_idx, _ch, _exp), chunks in [*train_buf.items(), *val_buf.items()]: + gene_cells[g_idx] = gene_cells.get(g_idx, 0) + sum( + c.shape[0] for c in chunks + ) + top_genes = sorted(gene_cells, key=lambda g: gene_cells[g], reverse=True)[ + :max_genes + ] + keep = set(top_genes) + train_buf = {k: v for k, v in train_buf.items() if k[0] in keep} + val_buf = {k: v for k, v in val_buf.items() if k[0] in keep} + old_to_new = {old: new for new, old in enumerate(sorted(keep))} + gene_to_idx = {idx_to_gene[old]: new for old, new in old_to_new.items()} + train_buf = {(old_to_new[k[0]], k[1], k[2]): v for k, v in train_buf.items()} + val_buf = {(old_to_new[k[0]], k[1], k[2]): v for k, v in val_buf.items()} + if load_val_dump_metadata: + val_dump_buf = {k: v for k, v in val_dump_buf.items() if k[0] in keep} + val_dump_buf = { + (old_to_new[k[0]], k[1], k[2]): v for k, v in val_dump_buf.items() + } + if load_train_dump_metadata: + train_dump_buf = {k: v for k, v in train_dump_buf.items() if k[0] in keep} + train_dump_buf = { + (old_to_new[k[0]], k[1], k[2]): v for k, v in train_dump_buf.items() + } + print(f" Filtered to top {max_genes} genes by cell count") + + if max_channels is not None and len(channel_to_idx) > max_channels: + idx_to_channel = {v: k for k, v in channel_to_idx.items()} + ch_cells: dict[int, int] = {} + for (_g, ch_idx, _exp), chunks in [*train_buf.items(), *val_buf.items()]: + ch_cells[ch_idx] = ch_cells.get(ch_idx, 0) + sum(c.shape[0] for c in chunks) + top_channels = sorted(ch_cells, key=lambda c: ch_cells[c], reverse=True)[ + :max_channels + ] + keep_ch = set(top_channels) + train_buf = {k: v for k, v in train_buf.items() if k[1] in keep_ch} + val_buf = {k: v for k, v in val_buf.items() if k[1] in keep_ch} + old_to_new_ch = {old: new for new, old in enumerate(sorted(keep_ch))} + channel_to_idx = { + idx_to_channel[old]: new for old, new in old_to_new_ch.items() + } + train_buf = {(k[0], old_to_new_ch[k[1]], k[2]): v for k, v in train_buf.items()} + val_buf = {(k[0], old_to_new_ch[k[1]], k[2]): v for k, v in val_buf.items()} + if load_val_dump_metadata: + val_dump_buf = {k: v for k, v in val_dump_buf.items() if k[1] in keep_ch} + val_dump_buf = { + (k[0], old_to_new_ch[k[1]], k[2]): v for k, v in val_dump_buf.items() + } + if load_train_dump_metadata: + train_dump_buf = { + k: v for k, v in train_dump_buf.items() if k[1] in keep_ch + } + train_dump_buf = { + (k[0], old_to_new_ch[k[1]], k[2]): v for k, v in train_dump_buf.items() + } + ce_sums = { + (old_to_new_ch[k[0]], k[1]): v + for k, v in ce_sums.items() + if k[0] in keep_ch + } + ce_sq_sums = { + (old_to_new_ch[k[0]], k[1]): v + for k, v in ce_sq_sums.items() + if k[0] in keep_ch + } + ce_counts = { + (old_to_new_ch[k[0]], k[1]): v + for k, v in ce_counts.items() + if k[0] in keep_ch + } + print(f" Filtered to top {max_channels} channels by cell count") + + if min_cells_per_group is not None: + gc_cells: dict[tuple[int, int], int] = {} + for (g_idx, ch_idx, _exp), chunks in [*train_buf.items(), *val_buf.items()]: + gc_key = (g_idx, ch_idx) + gc_cells[gc_key] = gc_cells.get(gc_key, 0) + sum(c.shape[0] for c in chunks) + drop_gc: set[tuple[int, int]] = { + k for k, n in gc_cells.items() if n < min_cells_per_group + } + if drop_gc: + train_buf = { + k: v for k, v in train_buf.items() if (k[0], k[1]) not in drop_gc + } + if load_train_dump_metadata: + train_dump_buf = { + k: v + for k, v in train_dump_buf.items() + if (k[0], k[1]) not in drop_gc + } + if min_cells_drop_val: + val_buf = { + k: v for k, v in val_buf.items() if (k[0], k[1]) not in drop_gc + } + if load_val_dump_metadata: + val_dump_buf = { + k: v + for k, v in val_dump_buf.items() + if (k[0], k[1]) not in drop_gc + } + print( + f" Dropped {len(drop_gc)} (gene, channel) groups" + f" with < {min_cells_per_group} cells" + f" (train{' + val' if min_cells_drop_val else ', val kept'})" + ) + + print(f"Total: {total_rows:,} rows, dim={emb_dim}") + + # Compute mean/std per (channel, experiment) from train moments + ce_mean: dict[tuple[int, int], torch.Tensor] = {} + ce_std: dict[tuple[int, int], torch.Tensor] = {} + ch_fallback_mean: dict[int, torch.Tensor] = {} + ch_fallback_std: dict[int, torch.Tensor] = {} + if z_standardize: + for ce_key in ce_sums: + n = ce_counts[ce_key] + mean = ce_sums[ce_key] / n + ce_mean[ce_key] = mean + ce_std[ce_key] = ( + (ce_sq_sums[ce_key] / n - mean**2).clamp(min=0).sqrt().clamp(min=1e-6) + ) + for ch_idx in channel_to_idx.values(): + ch_means = [ce_mean[k] for k in ce_mean if k[0] == ch_idx] + if ch_means: + ch_fallback_mean[ch_idx] = torch.stack(ch_means).mean(0) + ch_fallback_std[ch_idx] = torch.stack( + [ce_std[k] for k in ce_std if k[0] == ch_idx] + ).mean(0) + del ce_sums, ce_sq_sums, ce_counts + + # Concatenate chunks per (gene, channel, experiment) + print("Concatenating per (gene, channel, experiment)...") + + def _concat_gce( + buf: dict[tuple[int, int, int], list[torch.Tensor]], + desc: str = "Concat", + ) -> dict[tuple[int, int, int], torch.Tensor]: + return { + key: (torch.cat(chunks) if len(chunks) > 1 else chunks[0]) + for key, chunks in tqdm( + buf.items(), desc=desc, total=len(buf), unit="group" + ) + } + + train_gce = _concat_gce(train_buf, desc="Concat train") + del train_buf + val_gce = _concat_gce(val_buf, desc="Concat val") + del val_buf + + val_dump_gce: dict[tuple[int, int, int], CellDumpMetaLoL] | None = None + if load_val_dump_metadata: + val_dump_gce = {} + for key, chunks in tqdm( + val_dump_buf.items(), + desc="Concat val meta", + total=len(val_dump_buf), + unit="group", + ): + val_dump_gce[key] = _chunks_to_lol(chunks) + del val_dump_buf + + train_dump_gce: dict[tuple[int, int, int], CellDumpMetaLoL] | None = None + if load_train_dump_metadata: + train_dump_gce = {} + for key, chunks in tqdm( + train_dump_buf.items(), + desc="Concat train meta", + total=len(train_dump_buf), + unit="group", + ): + train_dump_gce[key] = _chunks_to_lol(chunks) + del train_dump_buf + + # Build (ch, exp) → list of gce keys for efficient lookup + train_by_ce: dict[tuple[int, int], list[tuple[int, int, int]]] = {} + for key in train_gce: + train_by_ce.setdefault((key[1], key[2]), []).append(key) + val_by_ce: dict[tuple[int, int], list[tuple[int, int, int]]] = {} + for key in val_gce: + val_by_ce.setdefault((key[1], key[2]), []).append(key) + + if z_standardize: + # Z-standardize by looping over (channel, experiment) pairs + ce_pairs = sorted(set(train_by_ce.keys()) | set(val_by_ce.keys())) + for ch_idx, exp_idx in tqdm(ce_pairs, desc="Z-standardize", unit="(ch,exp)"): + ce_key = (ch_idx, exp_idx) + mean = ce_mean.get( + ce_key, ch_fallback_mean.get(ch_idx, torch.zeros(emb_dim)) + ) + std = ce_std.get(ce_key, ch_fallback_std.get(ch_idx, torch.ones(emb_dim))) + for gce_key in train_by_ce.get(ce_key, []): + train_gce[gce_key].sub_(mean).div_(std) + for gce_key in val_by_ce.get(ce_key, []): + val_gce[gce_key].sub_(mean).div_(std) + + n_ce_groups = len(ce_mean) + stats_source = "NTC-only" if z_standardize_control_only else "all-train" + print( + f" Z-standardized per (channel, experiment): " + f"{n_ce_groups} groups (stats from {stats_source})" + ) + else: + print(" Skipping z-standardization (z_standardize=False)") + + # Merge across experiments: (gene, channel, experiment) → (gene, channel) + def _merge_experiments( + gce: dict[tuple[int, int, int], torch.Tensor], + desc: str = "Merge", + ) -> CellIndex: + gc_groups: dict[tuple[int, int], list[torch.Tensor]] = {} + for (g_idx, ch_idx, _exp_idx), t in gce.items(): + gc_groups.setdefault((g_idx, ch_idx), []).append(t) + index: CellIndex = {} + n_capped = 0 + for gc_key, tensors in tqdm( + gc_groups.items(), desc=desc, total=len(gc_groups), unit="group" + ): + combined = torch.cat(tensors) if len(tensors) > 1 else tensors[0] + if max_cells_per_group is not None and len(combined) > max_cells_per_group: + combined = combined[torch.randperm(len(combined))[:max_cells_per_group]] + n_capped += 1 + index[gc_key] = combined + if max_cells_per_group is not None and n_capped > 0: + print(f" Capped {n_capped} groups to {max_cells_per_group} cells") + return index + + if load_train_dump_metadata: + assert train_dump_gce is not None + train_index, train_cell_dump_meta = _merge_experiments_val_with_meta( + train_gce, + train_dump_gce, + max_cells_per_group, + desc="Merge train", + ) + del train_gce, train_dump_gce + else: + train_index = _merge_experiments(train_gce, desc="Merge train") + train_cell_dump_meta = None + del train_gce + + if load_val_dump_metadata: + assert val_dump_gce is not None + val_index, val_cell_dump_meta = _merge_experiments_val_with_meta( + val_gce, + val_dump_gce, + max_cells_per_group, + desc="Merge val", + ) + del val_gce, val_dump_gce + else: + val_index = _merge_experiments(val_gce, desc="Merge val") + val_cell_dump_meta = None + del val_gce + + gene_to_idx = dict(sorted(gene_to_idx.items(), key=lambda kv: kv[1])) + channel_to_idx = dict(sorted(channel_to_idx.items(), key=lambda kv: kv[1])) + total_train = sum(t.shape[0] for t in train_index.values()) + total_val = sum(t.shape[0] for t in val_index.values()) + print( + f" {len(gene_to_idx)} genes, {len(channel_to_idx)} channels" + f" | train: {total_train:,} cells, val: {total_val:,} cells" + ) + print(f" Total load time: {time.time() - t0:.1f}s") + + return ( + train_index, + val_index, + emb_dim, + gene_to_idx, + channel_to_idx, + val_cell_dump_meta, + train_cell_dump_meta, + ) + + +# --------------------------------------------------------------------------- +# Training loop +# --------------------------------------------------------------------------- + + +def _grad_norm(model: nn.Module) -> float: + total_sq = 0.0 + for p in model.parameters(): + if p.grad is not None: + total_sq += p.grad.data.square().sum().item() + return total_sq**0.5 + + +def _dump_dataset( + ds: Dataset, + gene_to_idx: dict[str, int], + channel_to_idx: dict[str, int], + emb_dim: int, + dump_dir: str, + label: str = "dataset", +) -> None: + """Save a MixedChannelDataset's gene pools to disk: one .pt per gene + metadata.pt.""" + assert isinstance(ds, MixedChannelDataset) + out = Path(dump_dir) + out.mkdir(parents=True, exist_ok=True) + + idx_to_gene = {v: k for k, v in gene_to_idx.items()} + + for g_idx, pool in tqdm( + ds._gene_pools.items(), desc=f"Dumping {label} genes", unit="gene" + ): + gene_name = idx_to_gene[g_idx] + payload: dict = { + "embeddings": pool, + "channel_ids": ds._gene_ch_ids[g_idx], + } + if ds._gene_dump_meta is not None: + dm = ds._gene_dump_meta[g_idx] + if _is_cell_meta_lol(dm): + payload[CELL_METADATA_DUMP_KEY] = dm + elif not dm["experiment"]: + payload[CELL_METADATA_DUMP_KEY] = {k: [] for k in CELL_DUMP_KEYS} + else: + flat: CellDumpMetaFlat = dm # type: ignore[assignment] + payload[CELL_METADATA_DUMP_KEY] = {k: [v] for k, v in flat.items()} + torch.save(payload, out / f"{gene_name}.pt") + + torch.save( + { + "gene_to_idx": gene_to_idx, + "channel_to_idx": channel_to_idx, + "emb_dim": emb_dim, + "perturbation_list": ds.perturbation_list, + "n_cells": ds.n_cells, + }, + out / "metadata.pt", + ) + print(f"Dumped {label} dataset ({len(ds._gene_pools)} genes) to {out}") + + ds._gene_dump_meta = None + + +def load_dataset( + dump_dir: str, + n_cells: int | None = None, + channels_per_set: int | list[int | None] | None = None, +) -> MixedChannelDataset: + """Recreate a MixedChannelDataset from files written by :func:`_dump_dataset`. + + Per-gene ``.pt`` files may include ``cell_metadata`` (column -> list of segment lists). + When present, it is flattened into ``_gene_dump_meta`` as one list per column. + """ + root = Path(dump_dir) + meta = torch.load(root / "metadata.pt", map_location="cpu", weights_only=False) + + gene_to_idx: dict[str, int] = meta["gene_to_idx"] + idx_to_gene = {v: k for k, v in gene_to_idx.items()} + emb_dim: int = meta["emb_dim"] + perturbation_list: list[str] = meta["perturbation_list"] + if n_cells is None: + n_cells = meta["n_cells"] + + ds = MixedChannelDataset.__new__(MixedChannelDataset) + ds.emb_dim = emb_dim + ds.gene_to_idx = gene_to_idx + ds.perturbation_list = perturbation_list + ds.n_cells = n_cells + ds.replacement = True + ds.channels_per_set = channels_per_set + ds._cps_choices = _normalize_cps_choices(channels_per_set) + ds._gene_pools = {} + ds._gene_ch_ids = {} + ds._gene_ch_indices = {} + ds._gene_dump_meta = None + ds.label_remap = None + ds._t_randperm = 0.0 + ds._t_index = 0.0 + ds._t_total = 0.0 + ds._n_calls = 0 + + for g_idx, gene_name in tqdm( + idx_to_gene.items(), desc="Loading genes", unit="gene" + ): + pt_path = root / f"{gene_name}.pt" + if not pt_path.exists(): + ds._gene_pools[g_idx] = torch.zeros(0, emb_dim) + ds._gene_ch_ids[g_idx] = torch.zeros(0, dtype=torch.long) + ds._gene_ch_indices[g_idx] = {} + continue + data = torch.load(pt_path, map_location="cpu", weights_only=False) + ds._gene_pools[g_idx] = data["embeddings"] + ds._gene_ch_ids[g_idx] = data["channel_ids"] + ds._gene_ch_indices[g_idx] = _compute_ch_indices(data["channel_ids"]) + if CELL_METADATA_DUMP_KEY in data: + if ds._gene_dump_meta is None: + ds._gene_dump_meta = {} + lol = data[CELL_METADATA_DUMP_KEY] + # Backwards compat: old dumps use "channel_index" instead of "index" + if "channel_index" in lol and "index" not in lol: + lol["index"] = lol.pop("channel_index") + col_order = [c for c in CELL_DUMP_KEYS if c in lol] + col_order += [c for c in lol if c not in col_order] + ds._gene_dump_meta[g_idx] = { + col: [x for seg in lol[col] for x in seg] for col in col_order + } + + print( + f"Loaded dataset from {root}: " + f"{len(ds._gene_pools)} genes, {len(ds.perturbation_list)} perturbations, " + f"n_cells={ds.n_cells}" + ) + return ds + + +load_val_dataset = load_dataset + + +def _subset_train_cells( + ds: MixedChannelDataset, + max_train_cells: int, + seed: int, +) -> tuple[int, int]: + """Uniformly subsample a mixed-channel train dataset to ``max_train_cells`` cells. + + Cells are drawn at random across all genes (and channels), so the natural + per-gene / per-channel distribution is preserved in expectation. Pass only + the *training* dataset -- validation is always kept intact. + + Parameters + ---------- + ds + The mixed-channel training dataset to subsample in place. + max_train_cells + Target total number of training cells. If the dataset already has fewer + cells, it is left unchanged. + seed + Seed for the global permutation, so the subset is reproducible. + + Returns + ------- + tuple[int, int] + ``(original_total, new_total)`` cell counts. + """ + keys = list(ds._gene_pools.keys()) + sizes = [len(ds._gene_pools[k]) for k in keys] + total = sum(sizes) + if total <= max_train_cells: + return total, total + + g = torch.Generator().manual_seed(seed) + keep = torch.zeros(total, dtype=torch.bool) + keep[torch.randperm(total, generator=g)[:max_train_cells]] = True + + offset = 0 + for k, sz in zip(keys, sizes): + mask = keep[offset : offset + sz] + offset += sz + ds._gene_pools[k] = ds._gene_pools[k][mask] + ds._gene_ch_ids[k] = ds._gene_ch_ids[k][mask] + ds._gene_ch_indices[k] = _compute_ch_indices(ds._gene_ch_ids[k]) + # Per-cell dump metadata (if present) is only consumed for validation + # analysis, never during training, so drop it rather than re-aligning. + ds._gene_dump_meta = None + return total, max_train_cells + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def run(cfg: DictConfig) -> None: + resolved_cfg = OmegaConf.to_container(cfg, resolve=True) + run_name = resolved_cfg.get("name") if isinstance(resolved_cfg, dict) else None + print(f"Resolved run name: {run_name!r}") + + device = torch.device( + cfg.get("device", "cuda" if torch.cuda.is_available() else "cpu") + ) + print(f"Device: {device}") + + seed = cfg.get("seed", 42) + torch.manual_seed(seed) + np.random.seed(seed) + + mixed_mode: bool = cfg.get("mixed_channels_mode", False) + if not mixed_mode: + raise ValueError("Only mixed_channels_mode=true is supported.") + dump_val_dir = cfg.get("dump_val_dir", None) + dump_train_dir = cfg.get("dump_train_dir", None) + load_train_dir = cfg.get("load_train_dir", None) + load_val_dir = cfg.get("load_val_dir", None) + + for _dname, _dval in [ + ("dump_val_dir", dump_val_dir), + ("dump_train_dir", dump_train_dir), + ]: + if _dval is not None: + _dp = Path(_dval) + if _dp.exists() and any(_dp.iterdir()): + raise RuntimeError( + f"{_dname} {_dval!r} already exists and is not empty. " + f"Remove it or choose a different directory." + ) + _dp.mkdir(parents=True, exist_ok=True) + + use_preloaded = load_train_dir is not None and load_val_dir is not None + if use_preloaded and not mixed_mode: + raise ValueError( + "load_train_dir / load_val_dir require mixed_channels_mode=true" + ) + + load_val_dump_metadata = dump_val_dir is not None and mixed_mode + load_train_dump_metadata = dump_train_dir is not None and mixed_mode + + # ---- Load data ---- + if use_preloaded: + # Fast path: load pre-dumped datasets directly + n_cells_cfg: int = cfg.get("n_cells_per_set", 500) + print(f"Loading pre-dumped train dataset from {load_train_dir}") + preloaded_train_ds = load_dataset(load_train_dir, n_cells=n_cells_cfg) + print(f"Loading pre-dumped val dataset from {load_val_dir}") + preloaded_val_ds = load_dataset(load_val_dir, n_cells=n_cells_cfg) + + train_meta = torch.load( + Path(load_train_dir) / "metadata.pt", map_location="cpu", weights_only=False + ) + gene_to_idx = train_meta["gene_to_idx"] + channel_to_idx = train_meta["channel_to_idx"] + emb_dim = train_meta["emb_dim"] + + train_index = None + val_index = None + val_cell_dump_meta = None + train_cell_dump_meta = None + else: + preloaded_train_ds = None + preloaded_val_ds = None + + val_frac = cfg.get("val_fraction", 0.2) + exclude_ch_cfg = cfg.data.get("exclude_channel_names") + exclude_channel_names = ( + list(exclude_ch_cfg) if exclude_ch_cfg is not None else None + ) + include_ch_cfg = cfg.data.get("include_channel_names") + include_channel_names = ( + list(include_ch_cfg) if include_ch_cfg is not None else None + ) + + parquet_entries: list[dict] = [] + for entry in cfg.data.parquet_entries: + pe: dict = {"path": entry.path} + exc_exp = entry.get("exclude_experiments") + if exc_exp is not None: + pe["exclude_experiments"] = list(exc_exp) + exc_fluor = entry.get("exclude_fluorescent_experiments") + if exc_fluor is not None: + pe["exclude_fluorescent_experiments"] = list(exc_fluor) + col_remap = entry.get("column_remap") + if col_remap is not None: + pe["column_remap"] = dict(col_remap) + parquet_entries.append(pe) + + ( + train_index, + val_index, + emb_dim, + gene_to_idx, + channel_to_idx, + val_cell_dump_meta, + train_cell_dump_meta, + ) = load_data( + parquet_entries=parquet_entries, + max_row_groups=cfg.data.get("max_row_groups", None), + max_cells_per_group=cfg.data.get("max_cells_per_group", None), + val_fraction=val_frac, + seed=seed, + cell_stratify=cfg.data.get("cell_stratify", False), + max_genes=cfg.data.get("max_genes", None), + max_channels=cfg.data.get("max_channels", None), + min_cells_per_group=cfg.data.get("min_cells_per_group", None), + min_cells_drop_val=cfg.data.get("min_cells_drop_val", False), + load_val_dump_metadata=load_val_dump_metadata, + load_train_dump_metadata=load_train_dump_metadata, + exclude_channel_names=exclude_channel_names, + include_channel_names=include_channel_names, + z_standardize=cfg.data.get("z_standardize", True), + z_standardize_control_only=cfg.data.get( + "z_standardize_control_only", False + ), + ) + + unique_genes = sorted(gene_to_idx.keys()) + n_channels = len(channel_to_idx) + + # ---- Optional label map (e.g. gene → pathway) ---- + label_map_path = cfg.data.get("label_map_path", None) + label_to_idx: dict[str, int] | None = None + label_remap: dict[int, int] | None = None + if label_map_path is not None: + gene_col = cfg.data.get("label_map_gene_col", "gene_name") + label_col = cfg.data.get("label_map_label_col", "pathway") + fallback_to_gene = cfg.data.get("label_map_fallback_to_gene", False) + gene_to_label = _load_label_map(label_map_path, gene_col, label_col) + n_before = len(unique_genes) + unique_genes, gene_to_label = _resolve_label_map_genes( + unique_genes, gene_to_label, fallback_to_gene + ) + n_dropped = n_before - len(unique_genes) + if n_dropped: + print(f" Label map: dropping {n_dropped} genes not in {label_map_path}") + unique_labels = sorted(set(gene_to_label[g] for g in unique_genes)) + label_to_idx = {label: i for i, label in enumerate(unique_labels)} + label_remap = { + gene_to_idx[g]: label_to_idx[gene_to_label[g]] for g in unique_genes + } + n_classes = len(label_to_idx) + print( + f" Label map: {len(unique_genes)} genes → {n_classes} classes " + f"from {label_map_path}" + ) + print(f" Classes: {unique_labels}") + else: + n_classes = len(unique_genes) + + # ---- Print per-channel cell counts ---- + idx_to_channel = {v: k for k, v in channel_to_idx.items()} + gene_indices = [gene_to_idx[g] for g in unique_genes] + if train_index is not None and val_index is not None: + cells_per_gene_channel: dict[int, list[int]] = { + i: [] for i in range(n_channels) + } + for g_idx in gene_indices: + for ch_idx in range(n_channels): + n = 0 + if (g_idx, ch_idx) in train_index: + n += len(train_index[(g_idx, ch_idx)]) + if (g_idx, ch_idx) in val_index: + n += len(val_index[(g_idx, ch_idx)]) + cells_per_gene_channel[ch_idx].append(n) + + print(f"\nCells per channel ({n_channels} channels):") + for ch_idx in range(n_channels): + counts = cells_per_gene_channel[ch_idx] + total = sum(counts) + sorted_counts = sorted(counts) + mid = len(sorted_counts) // 2 + median = ( + sorted_counts[mid] + if len(sorted_counts) % 2 == 1 + else (sorted_counts[mid - 1] + sorted_counts[mid]) // 2 + ) + print( + f" {idx_to_channel[ch_idx]}: {total:,} cells " + f"(per gene: min={min(counts):,}, median={median:,}, max={max(counts):,})" + ) + + train_multiplier = cfg.get("train_n_cell_sets_per_gene", 1) + val_multiplier = cfg.get("val_n_cell_sets_per_gene", 1) + batch_size = cfg.get("batch_size", 32) + mcfg = cfg.model + + phase2d_val: bool = cfg.get("phase2d_val", True) + phase2d_ch_idx: int | None = None + if phase2d_val: + for ch_name, ch_idx in channel_to_idx.items(): + if ch_name == "Phase2D": + phase2d_ch_idx = ch_idx + break + + val_phase2d_loader: DataLoader | None = None + val_n_cells_list: list[int] = [] + mixed_val_ds: MixedChannelDataset | None = None + mixed_val_p2d_ds: MixedChannelDataset | None = None + + if mixed_mode: + # -- Mixed-channel mode: pool all channels into one set -- + n_cells: int = cfg.get("n_cells_per_set", 500) + _cps_raw = cfg.get("channels_per_set", None) + channels_per_set: int | list[int | None] | None + if _cps_raw is None or isinstance(_cps_raw, int): + channels_per_set = _cps_raw + else: + channels_per_set = list(_cps_raw) + if channels_per_set is not None: + print( + f" channels_per_set={channels_per_set}: each set picks a value " + f"(None=all channels) and samples cells only from that many " + f"randomly-chosen channels per gene (applied to both train and val)" + ) + + if preloaded_train_ds is not None and preloaded_val_ds is not None: + # Fast path: datasets already loaded from dump dirs + train_ds: Dataset = preloaded_train_ds + val_ds: Dataset = preloaded_val_ds + preloaded_train_ds.channels_per_set = channels_per_set + preloaded_train_ds._cps_choices = _normalize_cps_choices(channels_per_set) + preloaded_val_ds.channels_per_set = channels_per_set + preloaded_val_ds._cps_choices = _normalize_cps_choices(channels_per_set) + if label_remap is not None: + preloaded_train_ds.label_remap = label_remap + preloaded_val_ds.label_remap = label_remap + # The dump's perturbation_list covers every dumped gene, but + # label_remap only has entries for genes the new label map + # covers. Drop the rest so __getitem__ never hits a KeyError. + mapped_set = set(unique_genes) + preloaded_train_ds.perturbation_list = [ + g for g in preloaded_train_ds.perturbation_list if g in mapped_set + ] + preloaded_val_ds.perturbation_list = [ + g for g in preloaded_val_ds.perturbation_list if g in mapped_set + ] + else: + assert train_index is not None and val_index is not None + + def _median(vals: list[int]) -> int: + s = sorted(vals) + mid = len(s) // 2 + return s[mid] if len(s) % 2 == 1 else (s[mid - 1] + s[mid]) // 2 + + # Stats: total cells per gene across all channels + total_per_gene: list[int] = [] + phase2d_per_gene: list[int] = [] + for g_idx in gene_indices: + total = 0 + p2d = 0 + for ch_idx in range(n_channels): + n = 0 + if (g_idx, ch_idx) in train_index: + n += len(train_index[(g_idx, ch_idx)]) + if (g_idx, ch_idx) in val_index: + n += len(val_index[(g_idx, ch_idx)]) + total += n + if ch_idx == phase2d_ch_idx: + p2d = n + total_per_gene.append(total) + phase2d_per_gene.append(p2d) + + print( + f"\nMixed-channel mode: {n_cells} cells per set" + f"\n Total cells per gene: " + f"min={min(total_per_gene):,}, " + f"median={_median(total_per_gene):,}, " + f"max={max(total_per_gene):,}" + ) + if phase2d_ch_idx is not None: + print( + f" Phase2D cells per gene: " + f"min={min(phase2d_per_gene):,}, " + f"median={_median(phase2d_per_gene):,}, " + f"max={max(phase2d_per_gene):,}" + ) + + train_ds = MixedChannelDataset( + train_index, + emb_dim, + gene_to_idx, + unique_genes, + n_cells, + cell_dump_index=train_cell_dump_meta, + label_remap=label_remap, + channels_per_set=channels_per_set, + ) + val_ds = MixedChannelDataset( + val_index, + emb_dim, + gene_to_idx, + unique_genes, + n_cells, + cell_dump_index=val_cell_dump_meta, + label_remap=label_remap, + channels_per_set=channels_per_set, + ) + + if dump_train_dir is not None: + _dump_dataset( + train_ds, + gene_to_idx, + channel_to_idx, + emb_dim, + dump_train_dir, + label="train", + ) + train_cell_dump_meta = None + if dump_val_dir is not None: + _dump_dataset( + val_ds, + gene_to_idx, + channel_to_idx, + emb_dim, + dump_val_dir, + label="val", + ) + val_cell_dump_meta = None + + del train_index, val_index + gc.collect() + + # Optionally cap the TOTAL number of training cells (val kept full). + assert isinstance(train_ds, MixedChannelDataset) + max_train_cells = cfg.data.get("max_train_cells", None) + if max_train_cells is not None: + orig, new = _subset_train_cells(train_ds, int(max_train_cells), seed) + print( + f"Train-cell subset: {orig:,} -> {new:,} " + f"(max_train_cells={int(max_train_cells):,})" + ) + else: + n_train_cells = sum(len(p) for p in train_ds._gene_pools.values()) + print(f"Using full train set: {n_train_cells:,} cells") + + # If len(dataset)*multiplier < batch_size, drop_last=True yields zero batches + # (e.g. max_genes=1 with batch_size>1). + _mixed_train_indices = len(train_ds) * train_multiplier + train_loader = DataLoader( + train_ds, + batch_size=batch_size, + sampler=RepeatSampler(len(train_ds), train_multiplier), + num_workers=cfg.get("num_workers", 0), + drop_last=_mixed_train_indices >= batch_size, + pin_memory=True, + ) + val_loader = DataLoader( + val_ds, + batch_size=batch_size, + sampler=RepeatSampler(len(val_ds), val_multiplier), + num_workers=0, + pin_memory=True, + ) + + channel_conditioning = mcfg.get("channel_conditioning", "none") + model: nn.Module = MixedChannelClassifier( + emb_dim=emb_dim, + n_classes=n_classes, + n_channels=n_channels, + d_model=mcfg.get("d_model", 256), + n_heads=mcfg.get("n_heads", 4), + n_layers=mcfg.get("n_layers_cell", 2), + n_inducing=mcfg.get("n_inducing_cell", 32), + d_ff=mcfg.get("d_ff", None), + dropout=mcfg.get("dropout", 0.1), + cosine_classifier=mcfg.get("cosine_classifier", False), + channel_conditioning=channel_conditioning, + pool_type=mcfg.get("pool_type", "pma"), + ).to(device) + + _train_fn = train_one_epoch_mixed + _eval_fn = evaluate_mixed + + val_p2d_ds: MixedChannelDataset | None = None + if phase2d_ch_idx is not None and val_index is not None: + val_p2d_index: CellIndex = { + k: v for k, v in val_index.items() if k[1] == phase2d_ch_idx + } + if val_p2d_index: + val_p2d_ds = MixedChannelDataset( + val_p2d_index, + emb_dim, + gene_to_idx, + unique_genes, + n_cells, + label_remap=label_remap, + ) + val_phase2d_loader = DataLoader( + val_p2d_ds, + batch_size=batch_size, + sampler=RepeatSampler(len(val_p2d_ds), val_multiplier), + num_workers=0, + pin_memory=True, + ) + + val_n_cells_raw = cfg.get("val_n_cells_per_set", None) + val_n_cells_list = list(val_n_cells_raw) if val_n_cells_raw else [] + mixed_val_ds = val_ds # type: ignore[assignment] + mixed_val_p2d_ds = val_p2d_ds + + n_gpus = torch.cuda.device_count() + if n_gpus > 1: + print(f"Using DataParallel across {n_gpus} GPUs") + model = nn.DataParallel(model) + + n_params = sum(p.numel() for p in model.parameters() if p.requires_grad) + print(f"Model parameters: {n_params:,}") + + # ---- Optimizer & scheduler ---- + lr = cfg.get("learning_rate", 1e-3) + wd = cfg.get("weight_decay", 1e-4) + optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=wd) + + n_epochs = cfg.get("num_epochs", 100) + warmup_epochs = cfg.get("warmup_epochs", 5) + + def lr_schedule(epoch: int) -> float: + if epoch < warmup_epochs: + return (epoch + 1) / warmup_epochs + progress = (epoch - warmup_epochs) / max(1, n_epochs - warmup_epochs) + return 0.5 * (1.0 + math.cos(math.pi * progress)) + + scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_schedule) + max_grad_norm: float | None = cfg.get("max_grad_norm", None) + + eval_every = cfg.get("eval_every", 1) + + # ---- Training ---- + best_val_acc = 0.0 + # Per-N validation accuracies of the SELECTED (best primary-N) model, captured + # at the epoch the best checkpoint is saved. Written to `metrics_out` if set. + best_metrics: dict[str, float] = {} + for epoch in range(n_epochs): + t0 = time.time() + + if isinstance(train_loader.dataset, MixedChannelDataset): + train_loader.dataset.reset_timers() + + train_loss, train_acc, grad_norm = _train_fn( + model, train_loader, optimizer, device, max_grad_norm=max_grad_norm + ) + + if isinstance(train_loader.dataset, MixedChannelDataset): + train_loader.dataset.print_timers() + + scheduler.step() + elapsed = time.time() - t0 + + log_dict: dict[str, float] = { + "epoch": epoch + 1, + "train/loss": train_loss, + "train/accuracy": train_acc, + "train/grad_norm": grad_norm, + "lr": optimizer.param_groups[0]["lr"], + "epoch_time_s": elapsed, + } + + do_eval = (epoch + 1) % eval_every == 0 or (epoch + 1) == n_epochs + val_elapsed = 0.0 + if do_eval: + t_val = time.time() + val_loss, val_acc = _eval_fn(model, val_loader, device) + log_dict["val/loss"] = val_loss + log_dict["val/accuracy"] = val_acc + + if val_phase2d_loader is not None: + p2d_loss, p2d_acc = _eval_fn(model, val_phase2d_loader, device) + log_dict["val_phase2d/loss"] = p2d_loss + log_dict["val_phase2d/accuracy"] = p2d_acc + + for nc in val_n_cells_list: + assert mixed_val_ds is not None + mixed_val_ds.set_n_cells(nc) + nc_loss, nc_acc = _eval_fn(model, val_loader, device) + log_dict[f"val_n{nc}/loss"] = nc_loss + log_dict[f"val_n{nc}/accuracy"] = nc_acc + if mixed_val_p2d_ds is not None and val_phase2d_loader is not None: + mixed_val_p2d_ds.set_n_cells(nc) + nc_p2d_loss, nc_p2d_acc = _eval_fn( + model, val_phase2d_loader, device + ) + log_dict[f"val_phase2d_n{nc}/loss"] = nc_p2d_loss + log_dict[f"val_phase2d_n{nc}/accuracy"] = nc_p2d_acc + + if val_n_cells_list and mixed_val_ds is not None: + mixed_val_ds.set_n_cells(n_cells) + if mixed_val_p2d_ds is not None: + mixed_val_p2d_ds.set_n_cells(n_cells) + + val_elapsed = time.time() - t_val + + extra_parts = "" + if do_eval: + extra_parts += f"| val_loss={val_loss:.4f} val_acc={val_acc:.4f} " + if val_phase2d_loader is not None: + extra_parts += f"| phase2d_acc={p2d_acc:.4f} " + for nc in val_n_cells_list: + extra_parts += f"| val_n{nc}_acc={log_dict[f'val_n{nc}/accuracy']:.4f} " + + print( + f"Epoch {epoch + 1:3d}/{n_epochs} " + f"| train_loss={train_loss:.4f} train_acc={train_acc:.4f} " + + extra_parts + + f"| lr={optimizer.param_groups[0]['lr']:.2e} " + f"| train={elapsed:.1f}s val={val_elapsed:.1f}s" + ) + + if do_eval and val_acc > best_val_acc: + best_val_acc = val_acc + best_metrics = { + "val_accuracy": val_acc, + **{ + f"val_n{nc}_accuracy": log_dict[f"val_n{nc}/accuracy"] + for nc in val_n_cells_list + }, + } + save_path = cfg.get("save_path", "best_set_classifier.pt") + state_dict = ( + model.module.state_dict() + if isinstance(model, nn.DataParallel) + else model.state_dict() + ) + ckpt: dict = { + "model_state_dict": state_dict, + "gene_to_idx": gene_to_idx, + "channel_to_idx": channel_to_idx, + "config": resolved_cfg, + "epoch": epoch + 1, + "val_acc": val_acc, + } + if label_to_idx is not None: + ckpt["label_to_idx"] = label_to_idx + ckpt["label_remap"] = label_remap + torch.save(ckpt, save_path) + print(f" Saved best model (val_acc={val_acc:.4f}) to {save_path}") + + print(f"\nBest val accuracy: {best_val_acc:.4f}") + + # Optionally persist the selected model's per-N validation accuracies so a + # wrapping eval can read them back without scraping stdout. + metrics_out = cfg.get("metrics_out", None) + if metrics_out is not None: + out_path = Path(metrics_out) + out_path.parent.mkdir(parents=True, exist_ok=True) + payload = {"n_cells_per_set": cfg.get("n_cells_per_set", None), **best_metrics} + out_path.write_text(json.dumps(payload, indent=2)) + print(f"Wrote metrics to {out_path}") + + +def build_model(ckpt: dict, device: torch.device) -> MixedChannelClassifier: + """Reconstruct a trained ``MixedChannelClassifier`` from a checkpoint dict. + + Used by ``eval.py`` and ``score.py``. Infers ``emb_dim`` / ``n_classes`` / + ``n_channels`` from the saved state dict so no separate metadata is required. + """ + config = ckpt["config"] + mcfg = OmegaConf.to_container(OmegaConf.create(config["model"]), resolve=True) + assert isinstance(mcfg, dict) + state_dict = ckpt["model_state_dict"] + if any(k.startswith("module.") for k in state_dict): + state_dict = {k.replace("module.", "", 1): v for k, v in state_dict.items()} + emb_dim = int(state_dict["input_proj.weight"].shape[1]) + n_classes = int(state_dict["head.1.weight"].shape[0]) + n_channels = len(ckpt["channel_to_idx"]) + model = MixedChannelClassifier( + emb_dim=emb_dim, + n_classes=n_classes, + n_channels=n_channels, + d_model=mcfg.get("d_model", 256), + n_heads=mcfg.get("n_heads", 4), + n_layers=mcfg.get("n_layers_cell", 2), + n_inducing=mcfg.get("n_inducing_cell", 32), + d_ff=mcfg.get("d_ff", None), + dropout=mcfg.get("dropout", 0.1), + cosine_classifier=mcfg.get("cosine_classifier", False), + channel_conditioning=mcfg.get("channel_conditioning", "none"), + ) + model.load_state_dict(state_dict) + return model.to(device).eval() + + +@hydra.main( + version_base="1.3.0", + config_path=_CONFIG_DIR, + config_name="train_set_classifier_phase_1K", +) +def main(cfg: DictConfig) -> None: + run(cfg) + + +if __name__ == "__main__": + main() # type: ignore[call-arg]