Feature attribution¶

Once m3 can predict disease, attribution asks which genes, cells and cell types drove that prediction. We train on the full reference, then trace each prediction back to the genes, cells, and cell types behind it.

In [1]:

import os
import pandas as pd
import m3

OUT = "./tutorial_out_demo/03_attribution"
os.makedirs(OUT, exist_ok=True)

import os import pandas as pd import m3 OUT = "./tutorial_out_demo/03_attribution" os.makedirs(OUT, exist_ok=True)

1. Load the demo data¶

In [2]:

data = m3.datasets.liu_demo()
print(data)

data = m3.datasets.liu_demo() print(data)

Dataset(n_cells=30534, batches=['B1', 'B2', 'B3'], modalities=[rna:1000, adt:192])

/opt/miniforge3/lib/python3.10/site-packages/anndata/_core/anndata.py:1820: UserWarning: Variable names are not unique. To make them unique, call `.var_names_make_unique`.
  utils.warn_names_duplicates("var")

2. Build and train the model¶

Same setup as the patient-prediction tutorial, minus the held-out batch (we attribute on the full reference). .train() fits both the integration model and the disease predictor; attribution explains that predictor.

In [3]:

model = m3.M3(
    data,
    condition_keys=["cond_group", "Age_interval"],
    target_condition="cond_group",
    celltype_key="mergedcelltype",
    batch_key="batch",
    donor_key="sample_id",
    embedding_dim=30,
)
model.train(max_epochs=500)
print("capabilities:", model.capabilities)

model = m3.M3( data, condition_keys=["cond_group", "Age_interval"], target_condition="cond_group", celltype_key="mergedcelltype", batch_key="batch", donor_key="sample_id", embedding_dim=30, ) model.train(max_epochs=500) print("capabilities:", model.capabilities)

Using 'matrix/data'
torch.Size([1000, 7163])
Using 'matrix/data'

torch.Size([1000, 10789])
Using 'matrix/data'
torch.Size([1000, 12582])
Using 'matrix/data'
torch.Size([192, 7163])
Using 'matrix/data'

torch.Size([192, 10789])
Using 'matrix/data'
torch.Size([192, 12582])
Batch counts: {0: 7163, 1: 10789, 2: 12582}
Minimum batch size: 7163

Epoch 1, Validation Loss: 5.1807

Epoch 2, Validation Loss: 5.1167

Epoch 3, Validation Loss: 5.0486

...  (493 epochs omitted)  ...

Epoch 497, Validation Loss: 0.5033

Epoch 498, Validation Loss: 0.5032

Epoch 499, Validation Loss: 0.5039

Epoch 500, Validation Loss: 0.5011
Using 'matrix/data'
torch.Size([1000, 7163])
Using 'matrix/data'

torch.Size([1000, 10789])
Using 'matrix/data'
torch.Size([1000, 12582])
Using 'matrix/data'
torch.Size([192, 7163])
Using 'matrix/data'

torch.Size([192, 10789])
Using 'matrix/data'
torch.Size([192, 12582])

capabilities: {'embedding': True, 'reconstruct': True, 'predict_donors': True}

3. Attribute and rank the cell types¶

model.attribute(reference_labels=["HC"]) scores how much each gene, cell, and cell type pushed the prediction away from the reference label (HC, i.e. healthy). Here we look at the top cell types: attr.top_celltypes(...) ranks them, keeping only cell types with enough cells in each condition so a small group can't mislead the ranking.

In [4]:

attr = model.attribute(reference_labels=["HC"])
print("top cell types (filtered: >= 200 cells per condition):\n",
      attr.top_celltypes(min_cells_per_condition=200).to_string(index=False))

attr = model.attribute(reference_labels=["HC"]) print("top cell types (filtered: >= 200 cells per condition):\n", attr.top_celltypes(min_cells_per_condition=200).to_string(index=False))

Computing patient-level predictions...

Running Input-level Captum IG with 50 steps...

Done.
Running Patient vector-level Captum IG with 50 steps...
Done.

top cell types (filtered: >= 200 cells per condition):
          celltype  importance
   Mono_Classical    7.129031
Mono_Nonclassical    6.961583
          CD8_Mem    6.789464
               NK    6.645633
        CD8_Va7.2    6.418572
          CD4_Mem    6.375149
            B_Mem    6.215689
            T_Vd2    6.005471
        CD4_Naive    3.659189
        CD8_Naive    3.085082
          B_Naive    1.832521

4. Top genes¶

attr.top_genes(...) ranks the genes by how strongly they drove the prediction. min_cells_per_condition keeps only cell types with enough cells in each condition (so a tiny group can't dominate the ranking), and housekeeping / ribosomal genes are dropped.

In [5]:

top_rna = attr.top_genes(n=100, min_cells_per_condition=200, modality="rna")
print(f"top-100 RNA genes (computed over {int(top_rna['n_celltypes_used'].iloc[0])} balanced cell types):")
print(top_rna.head(15).to_string(index=False))
top_rna.to_csv(os.path.join(OUT, "gene_importance_top100_rna.csv"), index=False)

top_rna = attr.top_genes(n=100, min_cells_per_condition=200, modality="rna") print(f"top-100 RNA genes (computed over {int(top_rna['n_celltypes_used'].iloc[0])} balanced cell types):") print(top_rna.head(15).to_string(index=False)) top_rna.to_csv(os.path.join(OUT, "gene_importance_top100_rna.csv"), index=False)

top-100 RNA genes (computed over 11 balanced cell types):
feature modality    score  n_celltypes_used
   OAZ1      rna 0.002577                11
  HSPA8      rna 0.002045                11
PIK3IP1      rna 0.002016                11
   CD69      rna 0.001964                11
   IL32      rna 0.001843                11
   IL7R      rna 0.001762                11
TSC22D3      rna 0.001632                11
  CXCR4      rna 0.001501                11
   UCP2      rna 0.001480                11
   JUNB      rna 0.001447                11
    LTB      rna 0.001406                11
 IFITM1      rna 0.001405                11
  DDIT4      rna 0.001379                11
  ANXA6      rna 0.001316                11
 FAM65B      rna 0.001314                11

5. Save the attribution¶

The full cells-by-features attribution matrix, as an HDF5 array and an AnnData.

In [6]:

import h5py
import anndata as ad

attr_mat = attr.attribution
present = [m for m in ("rna", "adt", "atac") if m in data.modality_names]
feat_names = [f for mm in present for f in list(data.var[mm])][: attr_mat.shape[1]]

with h5py.File(os.path.join(OUT, "attribution.h5"), "w") as f:
    f["data"] = attr_mat

_obs = model.cell_metadata.astype(str).reset_index(drop=True)
ad.AnnData(X=attr_mat, obs=_obs, var=pd.DataFrame(index=feat_names)).write_h5ad(
    os.path.join(OUT, "attribution.h5ad"))
print("saved attribution -> .h5 (['data']) / .h5ad in", OUT)

import h5py import anndata as ad attr_mat = attr.attribution present = [m for m in ("rna", "adt", "atac") if m in data.modality_names] feat_names = [f for mm in present for f in list(data.var[mm])][: attr_mat.shape[1]] with h5py.File(os.path.join(OUT, "attribution.h5"), "w") as f: f["data"] = attr_mat _obs = model.cell_metadata.astype(str).reset_index(drop=True) ad.AnnData(X=attr_mat, obs=_obs, var=pd.DataFrame(index=feat_names)).write_h5ad( os.path.join(OUT, "attribution.h5ad")) print("saved attribution -> .h5 (['data']) / .h5ad in", OUT)

/opt/miniforge3/lib/python3.10/site-packages/anndata/_core/aligned_df.py:67: ImplicitModificationWarning: Transforming to str index.
  warnings.warn("Transforming to str index.", ImplicitModificationWarning)
/opt/miniforge3/lib/python3.10/site-packages/anndata/_core/anndata.py:1820: UserWarning: Variable names are not unique. To make them unique, call `.var_names_make_unique`.
  utils.warn_names_duplicates("var")
... storing 'sample_id' as categorical

... storing 'Donor' as categorical

... storing 'cond_group' as categorical

... storing 'Age_interval' as categorical

... storing 'mergedcelltype' as categorical

... storing 'batch' as categorical

saved attribution -> .h5 (['data']) / .h5ad in ./tutorial_out_demo/03_attribution

Done. The genes, cells, and cell types behind m3's Severe-vs-HC prediction.