Refactor

.gitignore

@@ -10,6 +10,5 @@ __pycache__
 *.pt
 /data/*
 /logs/*
-**/figures/*
 **/runs/*
 .DS_Store

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "TOSICA"]
+	path = TOSICA
+	url = https://github.com/JackieHanLab/TOSICA.git

1_download.py

@@ -1,7 +1,26 @@
+"""Download all datasets and process into train/val/test .npy splits."""
+
+from allen_brain.TOSICA.train import set_seed
 from allen_brain.cell_data.cell_download import download_data
+from allen_brain.cell_data.cell_load import load_10x, load_smartseq
+from allen_brain.data_sets import TOSICA_DATASETS
+
 
 def main():
+    # Download raw files (CSVs + h5ad)
     download_data()
-
+
+    load_10x()
+    load_smartseq()
+
+    # TOSICA benchmark datasets (hPancreas, mPancreas, mAtlas)
+    set_seed(1)
+    for name, mod in TOSICA_DATASETS.items():
+        try:
+            mod.setup()
+        except Exception as e:
+            print(f'[WARN] {name}: {e}')
+
+
 if __name__ == '__main__':
-    main()
+    main()

2_visualize.py

@@ -0,0 +1,91 @@
+"""
+3_visualize.py
+Visualize both 10x and SmartSeq datasets.
+Generates: class distribution, PCA, UMAP, heatmap, violin, and CV^2 plots.
+All figures are saved under figures/*.png.
+"""
+import os
+
+from rich.console import Console
+from rich.panel import Panel
+
+from allen_brain.TOSICA.train import set_seed
+from allen_brain.cell_data.cell_dataset import make_dataset
+from allen_brain.cell_data.cell_load import ALL_DATASETS
+from allen_brain.cell_data.cell_vis import DatasetVisualizer
+
+SEED = 1
+console = Console()
+
+DATASETS = [(info['dir'], name) for name, info in ALL_DATASETS.items()]
+
+
+def run_visualizations(data_dir: str, tag: str):
+    set_seed(SEED)
+    """Run the full visualization suite for one dataset."""
+    if not (os.path.exists(os.path.join(data_dir, 'X_train.npy'))
+            or os.path.exists(os.path.join(data_dir, 'X_train.npz'))):
+        console.print(f"  [yellow]No data for {tag}, skipping[/yellow]")
+        return
+
+    console.print(Panel(
+        f"[bold]{tag.upper()}[/bold]  ·  {data_dir}",
+        title="Dataset", border_style="cyan", expand=False,
+    ))
+
+    train_ds = make_dataset(data_dir, split="train")
+    if train_ds is None:
+        console.print(f"  [yellow]No train split found for {tag}[/yellow], loading full matrix ...")
+        train_ds = make_dataset(data_dir)
+    if train_ds is None:
+        console.print(f"  [bold red]\\[SKIP][/bold red] No data found in {data_dir}")
+        return
+
+    gene_names = train_ds.gene_names
+    fig_dir = "figures"
+    vis = DatasetVisualizer(train_ds, fig_dir=fig_dir, seed=SEED)
+
+    vis.plot_class_distribution(
+        save_path=os.path.join(fig_dir, f"{tag}_class_distribution.png"),
+    )
+
+    pca, X_pca = vis.plot_pca(
+        n_components=20,
+        save_path=fig_dir,
+        file_name=f"{tag}_pca.png",
+    )
+
+    vis.plot_umap(
+        X_pca,
+        max_cells=6000,
+        save_path=os.path.join(fig_dir, f"{tag}_umap.png"),
+    )
+
+    vis.plot_heatmap(
+        gene_names=gene_names,
+        n_genes=20,
+        save_path=os.path.join(fig_dir, f"{tag}_heatmap.png"),
+    )
+
+    vis.plot_violin(
+        gene_names=gene_names,
+        top_n=6,
+        save_path=os.path.join(fig_dir, f"{tag}_violin.png"),
+    )
+
+    vis.plot_cv2(
+        gene_names=gene_names,
+        n_top=1000,
+        save_path=os.path.join(fig_dir, f"{tag}_cv2.png"),
+    )
+
+
+def main():
+    set_seed(SEED)
+    for data_dir, tag in DATASETS:
+        run_visualizations(data_dir, tag)
+    console.print("\nAll visualizations complete.")
+
+
+if __name__ == "__main__":
+    main()

3_CNN.py

@@ -0,0 +1,68 @@
+import math
+import os
+import pickle
+
+import numpy as np
+import torch
+import torch.optim as optim
+import torch.optim.lr_scheduler as lr_scheduler
+
+from allen_brain.TOSICA.train import set_seed
+from allen_brain.models import train as T
+from allen_brain.models.config import ExperimentConfig
+
+DATA_DIR = 'data/mPancreas'
+SEED = 1
+BATCH_SIZE = 16384
+N_HVG = 10000
+EPOCHS = 20
+LR = 0.001
+LRF = 0.01
+NORMALIZE = 'None'
+
+cfg = ExperimentConfig(
+    model='CellTypeCNN',
+    seed=SEED,
+    batch_size=BATCH_SIZE,
+    n_hvg=N_HVG,
+    epochs=EPOCHS,
+    normalize=NORMALIZE,
+)
+
+
+def main() -> None:
+    set_seed(SEED)
+    trainer = T.Trainer(cfg)
+    train_loader, val_loader, hvg_idx, scaler = trainer.make_dataloaders(
+        DATA_DIR, n_hvg=N_HVG, normalize=NORMALIZE)
+    ds = train_loader.dataset
+
+    model = T.build_model('CellTypeCNN', len(ds.gene_names), ds.n_classes)
+
+    criterion = torch.nn.CrossEntropyLoss()
+    pg = [p for p in model.parameters() if p.requires_grad]
+    optimizer = optim.SGD(pg, lr=LR, momentum=0.9, weight_decay=5E-5)
+    lf = lambda x: ((1 + math.cos(x * math.pi / EPOCHS)) / 2) * (1 - LRF) + LRF
+    scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
+
+    writer, ckpt = T.make_writer_and_ckpt(cfg, len(ds.gene_names))
+    ckpt_dir = os.path.dirname(ckpt)
+    T._save_model_kwargs(ckpt_dir, {})
+    if hvg_idx is not None:
+        np.save(os.path.join(ckpt_dir, 'hvg_indices.npy'), hvg_idx)
+    if scaler is not None:
+        with open(os.path.join(ckpt_dir, 'scaler.pkl'), 'wb') as f:
+            pickle.dump(scaler, f)
+    with open(os.path.join(ckpt_dir, 'normalize.txt'), 'w') as f:
+        f.write(NORMALIZE)
+
+    T.print_header()
+    T.train(model, (train_loader, val_loader), criterion, optimizer, scheduler,
+            EPOCHS, writer, ckpt, squeeze_channel=False)
+
+    metrics = trainer.evaluate(DATA_DIR, ckpt, squeeze_channel=False)
+    T.append_results_csv('CNN', metrics)
+
+
+if __name__ == '__main__':
+    main()

3_GNN.py

@@ -0,0 +1,66 @@
+import math
+import os
+
+import numpy as np
+import torch
+import torch.optim as optim
+import torch.optim.lr_scheduler as lr_scheduler
+
+from allen_brain.TOSICA.train import set_seed
+from allen_brain.models import train as T
+from allen_brain.models.config import ExperimentConfig
+from allen_brain.models.gnn_train import GraphTrainer, train_graph
+from allen_brain.models.CellTypeGNN import GraphBuilder
+
+DATA_DIR = 'data/mPancreas'
+SEED = 42
+BATCH_SIZE = 256
+K_NEIGHBORS = 10
+EPOCHS = 20
+LR = 0.001
+LRF = 0.01
+N_HVG = 10000
+NORMALIZE = 'None'
+
+cfg = ExperimentConfig(
+    model='CellTypeGNN',
+    seed=SEED,
+    batch_size=BATCH_SIZE,
+    n_hvg=0,
+    epochs=EPOCHS,
+    k_neighbors=K_NEIGHBORS,
+    normalize=NORMALIZE,
+)
+
+
+def main() -> None:
+    set_seed(SEED)
+    gb = GraphBuilder(k_neighbors=K_NEIGHBORS, normalize=NORMALIZE)
+    data = gb.build_graph_data(DATA_DIR, n_hvg=N_HVG).to(T.DEVICE)
+    n_classes: int = int(data.y.max().item()) + 1
+    class_names: list[str] = list(np.load(f'{DATA_DIR}/class_names.npy', allow_pickle=True))
+    n_features: int = data.x.shape[1]
+
+    model = T.build_model('CellTypeGNN', n_features, n_classes)
+
+    criterion = torch.nn.CrossEntropyLoss()
+    pg = [p for p in model.parameters() if p.requires_grad]
+    optimizer = optim.SGD(pg, lr=LR, momentum=0.9, weight_decay=5E-5)
+    lf = lambda x: ((1 + math.cos(x * math.pi / EPOCHS)) / 2) * (1 - LRF) + LRF
+    scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
+
+    writer, ckpt = T.make_writer_and_ckpt(cfg, n_features)
+    T._save_model_kwargs(os.path.dirname(ckpt), {})
+
+    T.print_header()
+    train_graph(model, data, criterion, optimizer, scheduler,
+                EPOCHS, writer, ckpt)
+
+    trainer = GraphTrainer(cfg)
+    metrics = trainer.evaluate(data, ckpt, n_features, n_classes,
+                               class_names=class_names)
+    T.append_results_csv('GNN', metrics)
+
+
+if __name__ == '__main__':
+    main()

3_MLP.py

@@ -0,0 +1,69 @@
+import math
+import os
+import pickle
+
+import numpy as np
+import torch
+import torch.optim as optim
+import torch.optim.lr_scheduler as lr_scheduler
+
+from allen_brain.TOSICA.train import set_seed
+from allen_brain.models import train as T
+from allen_brain.models.config import ExperimentConfig
+
+DATA_DIR = 'data/mPancreas'
+SEED = 1
+BATCH_SIZE = 8192
+EPOCHS = 20
+LR = 0.01
+LRF = 0.01
+N_HVG = 10000
+NORMALIZE = 'None'
+
+cfg = ExperimentConfig(
+    model='CellTypeMLP',
+    seed=SEED,
+    batch_size=BATCH_SIZE,
+    epochs=EPOCHS,
+    normalize=NORMALIZE,
+    lr=LR,
+
+)
+
+
+def main() -> None:
+    set_seed(SEED)
+    trainer = T.Trainer(cfg)
+    train_loader, val_loader, hvg_idx, scaler = trainer.make_dataloaders(
+        DATA_DIR, n_hvg=N_HVG, normalize=NORMALIZE)
+    ds = train_loader.dataset
+
+    model = T.build_model('CellTypeMLP', len(ds.gene_names), ds.n_classes)
+
+    criterion = torch.nn.CrossEntropyLoss()
+    pg = [p for p in model.parameters() if p.requires_grad]
+    optimizer = optim.SGD(pg, lr=LR, momentum=0.9, weight_decay=5E-5)
+    lf = lambda x: ((1 + math.cos(x * math.pi / EPOCHS)) / 2) * (1 - LRF) + LRF
+    scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
+
+    writer, ckpt = T.make_writer_and_ckpt(cfg, len(ds.gene_names))
+    ckpt_dir = os.path.dirname(ckpt)
+    T._save_model_kwargs(ckpt_dir, {})
+    if hvg_idx is not None:
+        np.save(os.path.join(ckpt_dir, 'hvg_indices.npy'), hvg_idx)
+    if scaler is not None:
+        with open(os.path.join(ckpt_dir, 'scaler.pkl'), 'wb') as f:
+            pickle.dump(scaler, f)
+    with open(os.path.join(ckpt_dir, 'normalize.txt'), 'w') as f:
+        f.write(NORMALIZE)
+
+    T.print_header()
+    T.train(model, (train_loader, val_loader), criterion, optimizer, scheduler,
+            EPOCHS, writer, ckpt, squeeze_channel=True)
+
+    metrics = trainer.evaluate(DATA_DIR, ckpt, squeeze_channel=True)
+    T.append_results_csv('MLP', metrics)
+
+
+if __name__ == '__main__':
+    main()