From e6efb9f9438e16105d990fe4d23d80af03e852d8 Mon Sep 17 00:00:00 2001 From: Youssef Date: Thu, 22 Jan 2026 15:52:26 -0500 Subject: [PATCH 1/2] [WIP] wandb + hydra --- README.md | 144 ++++++--- configs/augmentation/none.yaml | 6 + configs/augmentation/standard.yaml | 35 +++ configs/augmentation/strong.yaml | 42 +++ configs/config.yaml | 51 +++ configs/dataset/mc_maze.yaml | 36 +++ configs/dataset/mc_rtt.yaml | 37 +++ configs/experiment/exp22c_teacher.yaml | 51 +++ configs/experiment/exp23_validation.yaml | 53 ++++ configs/experiment/exp25_mamba.yaml | 51 +++ configs/experiment/lstm_baseline.yaml | 36 +++ configs/model/lstm.yaml | 24 ++ configs/model/mamba.yaml | 31 ++ configs/model/transformer.yaml | 30 ++ configs/model/vqvae.yaml | 37 +++ configs/trainer/default.yaml | 39 +++ configs/trainer/progressive.yaml | 43 +++ docs/MIGRATION_GUIDE.md | 165 ++++++++++ evaluate.py | 130 ++++++++ pyproject.toml | 39 ++- requirements.txt | 26 +- src/__init__.py | 13 + src/datamodules/__init__.py | 79 +++++ src/datamodules/base.py | 253 +++++++++++++++ src/datamodules/mc_maze.py | 99 ++++++ src/datamodules/mc_rtt.py | 111 +++++++ src/trainer.py | 377 +++++++++++++++++++++++ src/utils/__init__.py | 21 ++ src/utils/augmentation.py | 232 ++++++++++++++ src/utils/logging.py | 183 +++++++++++ src/utils/metrics.py | 169 ++++++++++ src/utils/seeding.py | 152 +++++++++ train.py | 167 ++++++++++ 33 files changed, 2907 insertions(+), 55 deletions(-) create mode 100644 configs/augmentation/none.yaml create mode 100644 configs/augmentation/standard.yaml create mode 100644 configs/augmentation/strong.yaml create mode 100644 configs/config.yaml create mode 100644 configs/dataset/mc_maze.yaml create mode 100644 configs/dataset/mc_rtt.yaml create mode 100644 configs/experiment/exp22c_teacher.yaml create mode 100644 configs/experiment/exp23_validation.yaml create mode 100644 configs/experiment/exp25_mamba.yaml create mode 100644 configs/experiment/lstm_baseline.yaml create mode 100644 configs/model/lstm.yaml create mode 100644 configs/model/mamba.yaml create mode 100644 configs/model/transformer.yaml create mode 100644 configs/model/vqvae.yaml create mode 100644 configs/trainer/default.yaml create mode 100644 configs/trainer/progressive.yaml create mode 100644 docs/MIGRATION_GUIDE.md create mode 100644 evaluate.py create mode 100644 src/__init__.py create mode 100644 src/datamodules/__init__.py create mode 100644 src/datamodules/base.py create mode 100644 src/datamodules/mc_maze.py create mode 100644 src/datamodules/mc_rtt.py create mode 100644 src/trainer.py create mode 100644 src/utils/__init__.py create mode 100644 src/utils/augmentation.py create mode 100644 src/utils/logging.py create mode 100644 src/utils/metrics.py create mode 100644 src/utils/seeding.py create mode 100644 train.py diff --git a/README.md b/README.md index 23a5be5..20da853 100644 --- a/README.md +++ b/README.md @@ -103,53 +103,120 @@ An experimental project exploring: 📓 **[RESEARCH_LOG.md](RESEARCH_LOG.md)** - Detailed experiment notes, results, and analysis -## Project Structure +## Project Structure (v0.2 - Configuration-Driven) ``` -python/phantomx/ - model.py # ProgressiveVQVAE (MLP-based) - models_extended.py # CausalTransformerVQVAE, GumbelVQVAE (best performers) - trainer.py # ProgressiveTrainer (3-phase training) - tta.py # Test-Time Adaptation (TTAWrapper, OnlineTTA) - tokenizer/ # Spike tokenization - data/ # MC_Maze data loading -python/ - exp10_beat_lstm.py # Latest: CausalTransformer + Gumbel experiments - compare_models.py # Model comparisons -models/ - exp9_progressive_vqvae.pt # Progressive VQ-VAE (R²=0.71) - comparison_results.json # All experiment results +PhantomX/ +├── configs/ # 🆕 Hydra YAML configurations +│ ├── config.yaml # Main config +│ ├── model/ # Model configs (vqvae, mamba, lstm, transformer) +│ ├── dataset/ # Dataset configs (mc_maze, mc_rtt) +│ ├── trainer/ # Training configs (default, progressive) +│ ├── augmentation/ # Augmentation configs (none, standard, strong) +│ └── experiment/ # Experiment presets (exp25_mamba, etc.) +├── src/ # 🆕 Unified source code +│ ├── models/ # Model implementations +│ ├── datamodules/ # Dataset loading +│ ├── trainer.py # Unified trainer +│ └── utils/ # Logging, seeding, metrics +├── train.py # 🆕 SINGLE entry point for all experiments +├── python/ # Legacy experiment scripts (deprecated) +├── data/ # NWB data files +└── logs/ # Experiment outputs (auto-generated) ``` -## Quick Start +## Quick Start (v0.2) -```python -from phantomx.model import ProgressiveVQVAE -from phantomx.trainer import ProgressiveTrainer -from phantomx.data import MCMazeDataset - -# Load data -dataset = MCMazeDataset("path/to/mc_maze.nwb") -train_loader, val_loader, test_loader = create_dataloaders(dataset) - -# Create and train model -model = ProgressiveVQVAE(n_channels=142, window_size=10) -trainer = ProgressiveTrainer(model, train_loader, val_loader) -result = trainer.train() -print(f"Best R²: {result['best_r2']:.4f}") - -# Test-Time Adaptation for new sessions -from phantomx.tta import OnlineTTA -tta = OnlineTTA(model) -predictions = tta.predict(new_data) -``` - -## Setup +### Installation ```bash +# Create environment python -m venv venv venv\Scripts\activate # Windows -pip install -r requirements.txt +# source venv/bin/activate # Linux/Mac + +# Install with new dependencies +pip install -e ".[all]" + +# Or just core dependencies +pip install -e . +``` + +### Basic Training + +```bash +# Default: VQ-VAE on MC_Maze +python train.py + +# Mamba on MC_RTT (Exp 25) +python train.py experiment=exp25_mamba + +# LSTM baseline +python train.py model=lstm dataset=mc_maze + +# Override hyperparameters +python train.py model=mamba model.n_layers=6 trainer.learning_rate=1e-4 + +# Disable WandB for quick tests +python train.py trainer.max_epochs=5 logging.use_wandb=false +``` + +### Multi-Seed Validation (Exp 23 style) + +```bash +# Run with multiple seeds +python train.py experiment=exp23_validation --multirun seed=42,123,456,789,1337 +``` + +### Hyperparameter Sweeps + +```bash +# Model comparison +python train.py --multirun model=vqvae,lstm,mamba dataset=mc_maze + +# Full grid search +python train.py --multirun \ + model=mamba \ + model.d_model=64,128,256 \ + model.n_layers=2,4,6 +``` + +### Python API + +```python +import hydra +from omegaconf import OmegaConf + +from src.models import build_model +from src.datamodules import build_datamodule +from src.trainer import Trainer +from src.utils import seed_everything + +# Load config +cfg = OmegaConf.load("configs/config.yaml") +cfg = OmegaConf.merge(cfg, OmegaConf.load("configs/model/mamba.yaml")) +cfg = OmegaConf.merge(cfg, OmegaConf.load("configs/dataset/mc_rtt.yaml")) + +seed_everything(42) + +# Build components +datamodule = build_datamodule(cfg) +datamodule.setup() + +model = build_model(cfg, n_channels=datamodule.n_channels) +trainer = Trainer(model, datamodule.train_dataloader(), datamodule.val_dataloader(), cfg) + +# Train +results = trainer.train() +print(f"Best R²: {results['best_r2']:.4f}") +``` + +## Legacy Quick Start (v0.1) + +```python +# Still works for backward compatibility +from python.phantomx.model import ProgressiveVQVAE +from python.phantomx.trainer import ProgressiveTrainer ``` ## Current Status @@ -192,6 +259,7 @@ This project was developed with assistance from AI coding assistants and workflo - Claude Sonnet 4.5 (Anthropic) - Gemini 3.0 Pro (Google) - GPT 5.2 (OpenAI) +- Grok Code Fast 1 (xAi) All code was tested, and validated by the author. diff --git a/configs/augmentation/none.yaml b/configs/augmentation/none.yaml new file mode 100644 index 0000000..4b0da97 --- /dev/null +++ b/configs/augmentation/none.yaml @@ -0,0 +1,6 @@ +# No Augmentation Configuration +# ============================== +# Baseline: no data augmentation + +name: "none" +enabled: false diff --git a/configs/augmentation/standard.yaml b/configs/augmentation/standard.yaml new file mode 100644 index 0000000..6e5a2e5 --- /dev/null +++ b/configs/augmentation/standard.yaml @@ -0,0 +1,35 @@ +# Standard Augmentation Configuration +# ==================================== +# Data augmentation techniques from successful experiments + +name: "standard" +enabled: true + +# Electrode dropout (key for robustness) +electrode_dropout: + enabled: true + p: 0.1 # Probability of dropping each electrode + structured: false # Random vs structured dropout + +# Temporal jitter +temporal_jitter: + enabled: true + max_shift_bins: 2 # Max bins to shift (50ms at 25ms bins) + wrap: false # Wrap around or zero-pad + +# Noise injection +noise: + enabled: true + type: "gaussian" # gaussian, poisson + scale: 0.05 # Noise scale + +# Spike count scaling +spike_scaling: + enabled: true + min_scale: 0.9 + max_scale: 1.1 + +# Mixup (disabled by default) +mixup: + enabled: false + alpha: 0.2 diff --git a/configs/augmentation/strong.yaml b/configs/augmentation/strong.yaml new file mode 100644 index 0000000..f8cacd3 --- /dev/null +++ b/configs/augmentation/strong.yaml @@ -0,0 +1,42 @@ +# Strong Augmentation Configuration +# ================================== +# Aggressive augmentation for maximum robustness + +name: "strong" +enabled: true + +# Electrode dropout (aggressive) +electrode_dropout: + enabled: true + p: 0.2 # 20% electrode dropout + structured: true # Structured (realistic) dropout patterns + +# Temporal jitter (aggressive) +temporal_jitter: + enabled: true + max_shift_bins: 4 # 100ms jitter + wrap: false + +# Noise injection (stronger) +noise: + enabled: true + type: "gaussian" + scale: 0.1 + +# Spike count scaling (wider range) +spike_scaling: + enabled: true + min_scale: 0.8 + max_scale: 1.2 + +# Mixup (enabled) +mixup: + enabled: true + alpha: 0.3 + +# Cutout (neural version) +cutout: + enabled: true + n_holes: 2 + hole_size_channels: 10 + hole_size_time: 3 diff --git a/configs/config.yaml b/configs/config.yaml new file mode 100644 index 0000000..a43c6a5 --- /dev/null +++ b/configs/config.yaml @@ -0,0 +1,51 @@ +# PhantomX Main Configuration +# ============================ +# Default configuration for all experiments. +# Override via command line: python train.py model=mamba dataset=mc_rtt + +defaults: + - _self_ + - model: vqvae # Default model + - dataset: mc_maze # Default dataset + - trainer: default # Default trainer settings + - augmentation: none # No augmentation by default + - optional experiment: null # Override everything for specific experiments + +# ===== General Settings ===== +seed: 42 +output_dir: ${hydra:run.dir} +project_name: "PhantomX" +experiment_name: "${model.name}_${dataset.name}" + +# ===== Logging ===== +logging: + use_wandb: true + wandb_project: "PhantomX" + wandb_entity: null # Your wandb username/team + log_every_n_steps: 10 + save_checkpoints: true + checkpoint_every_n_epochs: 10 + +# ===== Reproducibility ===== +reproducibility: + deterministic: true + benchmark: false # Set true for faster training if input sizes don't change + save_git_hash: true + save_config: true # Hydra does this automatically + +# ===== Hardware ===== +hardware: + device: "auto" # auto, cuda, cpu, mps + num_workers: 4 + pin_memory: true + precision: 32 # 16, 32, or "bf16" + +# ===== Hydra Settings ===== +hydra: + run: + dir: logs/${experiment_name}/${now:%Y-%m-%d_%H-%M-%S} + sweep: + dir: logs/sweeps/${experiment_name} + subdir: ${hydra.job.num} + job: + chdir: false # Don't change working directory diff --git a/configs/dataset/mc_maze.yaml b/configs/dataset/mc_maze.yaml new file mode 100644 index 0000000..dfbab6e --- /dev/null +++ b/configs/dataset/mc_maze.yaml @@ -0,0 +1,36 @@ +# MC_Maze Dataset Configuration +# ============================== +# Motor cortex recordings during maze reaching task +# From Neural Latents Benchmark (NLB 2021) + +name: "mc_maze" + +# Dataset class for Hydra instantiation +_target_: src.datamodules.mc_maze.MCMazeDataModule + +# Data paths +data_dir: ${oc.env:PHANTOMX_DATA_DIR,data} +filename: "mc_maze.nwb" + +# Preprocessing +bin_size_ms: 25.0 # 40 Hz sampling +normalize: true # Z-score normalization +normalize_per_channel: true + +# Splits +train_ratio: 0.7 +val_ratio: 0.1 +test_ratio: 0.2 +shuffle_before_split: true # Shuffle trials, not time bins + +# DataLoader settings +batch_size: 64 +shuffle: true # Shuffle batches during training + +# Dataset specifics +n_channels: 137 # Neural units +task: "velocity_decoding" +target_type: "velocity" # velocity, position, acceleration + +# Trial structure +trial_based: true # MC_Maze has discrete trials diff --git a/configs/dataset/mc_rtt.yaml b/configs/dataset/mc_rtt.yaml new file mode 100644 index 0000000..210aadc --- /dev/null +++ b/configs/dataset/mc_rtt.yaml @@ -0,0 +1,37 @@ +# MC_RTT Dataset Configuration +# ============================= +# Motor cortex Random Target Tracking task +# Continuous tracking - fundamentally different from MC_Maze + +name: "mc_rtt" + +# Dataset class for Hydra instantiation +_target_: src.datamodules.mc_rtt.MCRTTDataModule + +# Data paths +data_dir: ${oc.env:PHANTOMX_DATA_DIR,data} +filename: "mc_rtt.nwb" + +# Preprocessing +bin_size_ms: 25.0 # 40 Hz sampling +normalize: true +normalize_per_channel: true + +# Splits (CRITICAL: sequential for continuous data) +train_ratio: 0.7 +val_ratio: 0.1 +test_ratio: 0.2 +shuffle_before_split: false # Keep temporal order! + +# DataLoader settings +batch_size: 32 +shuffle: false # Sequential batches for stateful models + +# Dataset specifics +n_channels: 130 # Neural units (different from MC_Maze) +task: "velocity_decoding" +target_type: "finger_velocity" # Uses finger_vel, not hand_vel + +# Continuous structure +trial_based: false # Continuous tracking, no discrete trials +continuous: true # Enable overlapping windows diff --git a/configs/experiment/exp22c_teacher.yaml b/configs/experiment/exp22c_teacher.yaml new file mode 100644 index 0000000..5e9cd44 --- /dev/null +++ b/configs/experiment/exp22c_teacher.yaml @@ -0,0 +1,51 @@ +# @package _global_ +# ================================================== +# Experiment 22c: Multi-Seed Teacher Distillation +# ================================================== +# Ensemble of transformers for robust teaching signal +# +# From RESEARCH_LOG: Architecture alone is insufficient. +# Need identical augmentation, dropout, lr settings. +# +# Best result: R² = 0.7159 + +defaults: + - override /model: transformer + - override /dataset: mc_maze + - override /trainer: default + - override /augmentation: standard + +experiment_name: "exp22c_multiseed_teacher" +notes: | + Experiment 22c: Multi-Seed Causal Transformer + - Ensemble training across multiple seeds + - Knowledge distillation target + - Standard augmentation for robustness + +# Model overrides (from exp22c) +model: + d_model: 384 + n_heads: 6 + n_layers: 6 + dropout: 0.1 + window_size: 10 + +# Dataset +dataset: + batch_size: 64 + shuffle: true + +# Training +trainer: + max_epochs: 100 + learning_rate: 3e-4 + lr_scheduler: "cosine" + patience: 20 + +# Multi-seed ensemble +ensemble: + enabled: true + seeds: [42, 123, 456, 789, 1337] + aggregate: "mean" # mean, vote, stack + +seed: 42 diff --git a/configs/experiment/exp23_validation.yaml b/configs/experiment/exp23_validation.yaml new file mode 100644 index 0000000..baa539e --- /dev/null +++ b/configs/experiment/exp23_validation.yaml @@ -0,0 +1,53 @@ +# @package _global_ +# ================================================== +# Experiment 23: Statistical Validation +# ================================================== +# Multi-seed validation to ensure reproducibility +# From README: "Validated across 10 seeds" +# +# Best R² = 0.7159 ± 0.008 + +defaults: + - override /model: vqvae + - override /dataset: mc_maze + - override /trainer: progressive + - override /augmentation: standard + +experiment_name: "exp23_statistical_validation" +notes: | + Experiment 23: Statistical Validation + - 10 seeds for reproducibility + - Progressive VQ-VAE training + - Standard augmentation + +# Model (Progressive VQ-VAE) +model: + embedding_dim: 128 + quantizer: + num_codes: 256 + type: "ema" + window_size: 10 + +# Dataset +dataset: + batch_size: 64 + shuffle: true + +# Training (Progressive) +trainer: + pretrain: + epochs: 30 + learning_rate: 1e-3 + finetune: + epochs: 50 + learning_rate: 3e-4 + +# Statistical validation +validation: + n_seeds: 10 + seeds: [42, 123, 456, 789, 1337, 2024, 3141, 4242, 5555, 6789] + report_mean: true + report_std: true + confidence_interval: 0.95 + +seed: 42 # Default seed, overridden in sweep diff --git a/configs/experiment/exp25_mamba.yaml b/configs/experiment/exp25_mamba.yaml new file mode 100644 index 0000000..b638799 --- /dev/null +++ b/configs/experiment/exp25_mamba.yaml @@ -0,0 +1,51 @@ +# @package _global_ +# ================================================== +# Experiment 25: Mamba on MC_RTT +# ================================================== +# The Navigation Filter - Mamba as a Neural Kalman Filter +# +# HYPOTHESIS: +# MC_RTT is continuous random target tracking where context IS the trajectory. +# Mamba's stateful hidden state acts as trajectory memory. +# +# Key differences from MC_Maze: +# - Window: 2 seconds (80 bins) - long context matters +# - No shuffle: sequential batches for stateful training +# - Stateful: model maintains h_t across the session +# +# TARGET: R² > 0.70 + +defaults: + - override /model: mamba + - override /dataset: mc_rtt + - override /trainer: default + - override /augmentation: none + +# Experiment metadata +experiment_name: "exp25_mamba_mc_rtt" +notes: | + Experiment 25: Mamba on Continuous Tracking (MC_RTT) + - Stateful Mamba acting as Neural Kalman Filter + - 2-second context window for trajectory integration + - Sequential training (no shuffle) + +# Model overrides (from successful exp25 results) +model: + d_model: 128 + n_layers: 4 + d_state: 16 + window_size: 80 # 2 seconds at 40Hz + stateful: true + +# Dataset overrides +dataset: + shuffle: false # CRITICAL for stateful training + batch_size: 32 + +# Training overrides +trainer: + max_epochs: 100 + learning_rate: 1e-3 + patience: 30 # Longer patience for stateful training + +seed: 42 diff --git a/configs/experiment/lstm_baseline.yaml b/configs/experiment/lstm_baseline.yaml new file mode 100644 index 0000000..6ea02f2 --- /dev/null +++ b/configs/experiment/lstm_baseline.yaml @@ -0,0 +1,36 @@ +# @package _global_ +# ================================================== +# LSTM Baseline Experiment +# ================================================== +# Standard LSTM for comparison +# Establishes baseline performance + +defaults: + - override /model: lstm + - override /dataset: mc_maze + - override /trainer: default + - override /augmentation: none + +experiment_name: "lstm_baseline" +notes: | + LSTM Baseline + - Standard 2-layer LSTM + - No augmentation (fair comparison) + - Baseline for other experiments + +model: + hidden_dim: 256 + num_layers: 2 + dropout: 0.2 + window_size: 10 + +dataset: + batch_size: 64 + shuffle: true + +trainer: + max_epochs: 100 + learning_rate: 1e-3 + patience: 20 + +seed: 42 diff --git a/configs/model/lstm.yaml b/configs/model/lstm.yaml new file mode 100644 index 0000000..52d56aa --- /dev/null +++ b/configs/model/lstm.yaml @@ -0,0 +1,24 @@ +# LSTM Baseline Configuration +# ============================ +# Standard LSTM architecture for comparison + +name: "lstm" + +# Model class for Hydra instantiation +_target_: src.models.lstm.LSTMDecoder + +# Architecture +input_dim: null # Set automatically from dataset +hidden_dim: 256 +num_layers: 2 +bidirectional: false +dropout: 0.2 + +# Output +output_dim: 2 # Velocity (vx, vy) + +# Temporal settings +window_size: 10 # 10 bins * 25ms = 250ms + +# Training +stateful: false diff --git a/configs/model/mamba.yaml b/configs/model/mamba.yaml new file mode 100644 index 0000000..58c4c83 --- /dev/null +++ b/configs/model/mamba.yaml @@ -0,0 +1,31 @@ +# Mamba (State Space Model) Configuration +# ======================================== +# Stateful Mamba architecture for continuous tracking tasks +# Optimized for MC_RTT based on Exp 25 results + +name: "mamba" + +# Model class for Hydra instantiation +_target_: src.models.mamba.StatefulMamba + +# Core S6 (Selective State Space) parameters +d_model: 128 # Model dimension +d_state: 16 # SSM state dimension +d_conv: 4 # Convolution kernel size +expand: 2 # FFN expansion factor +n_layers: 4 # Number of Mamba blocks + +# Discretization +dt_min: 0.001 +dt_max: 0.1 + +# Input/Output +output_dim: 2 # Velocity (vx, vy) +dropout: 0.1 + +# Temporal settings (CRITICAL for MC_RTT) +window_size: 80 # 80 bins * 25ms = 2 seconds context +stateful: true # Maintain hidden state across batches + +# Training specifics +no_shuffle: true # Sequential batches for stateful training diff --git a/configs/model/transformer.yaml b/configs/model/transformer.yaml new file mode 100644 index 0000000..2c7878f --- /dev/null +++ b/configs/model/transformer.yaml @@ -0,0 +1,30 @@ +# Causal Transformer Configuration +# ================================= +# Transformer with causal attention for neural decoding +# Based on Exp 22c multi-seed teacher results + +name: "transformer" + +# Model class for Hydra instantiation +_target_: src.models.transformer.CausalTransformerDecoder + +# Architecture +d_model: 384 # Model dimension (from Exp 22c) +n_heads: 6 # Attention heads +n_layers: 6 # Transformer blocks +dim_feedforward: 1536 # 4x d_model +dropout: 0.1 + +# Positional encoding +pos_encoding: "sinusoidal" # sinusoidal, learned, rotary +max_seq_len: 512 + +# Attention +attention_dropout: 0.1 +causal: true # Causal masking + +# Output +output_dim: 2 # Velocity (vx, vy) + +# Temporal settings +window_size: 10 # 10 bins * 25ms = 250ms diff --git a/configs/model/vqvae.yaml b/configs/model/vqvae.yaml new file mode 100644 index 0000000..1f7b66b --- /dev/null +++ b/configs/model/vqvae.yaml @@ -0,0 +1,37 @@ +# Progressive VQ-VAE Model Configuration +# ======================================= +# Default VQ-VAE architecture from successful experiments + +name: "vqvae" + +# Model class for Hydra instantiation +_target_: src.models.vqvae.ProgressiveVQVAE + +# Architecture +encoder: + hidden_dims: [1024, 512] + activation: "gelu" + use_layer_norm: true + dropout: 0.1 + +decoder: + hidden_dims: [512, 1024] + activation: "gelu" + use_layer_norm: true + dropout: 0.1 + +# Vector Quantization +quantizer: + type: "ema" # ema, gumbel, fsq + num_codes: 256 + embedding_dim: 128 + decay: 0.99 + commitment_cost: 0.1 + epsilon: 1e-5 + +# Input/Output +embedding_dim: 128 +output_dim: 2 # Velocity (vx, vy) + +# Temporal settings +window_size: 10 # 10 bins * 25ms = 250ms context diff --git a/configs/trainer/default.yaml b/configs/trainer/default.yaml new file mode 100644 index 0000000..8bc75f6 --- /dev/null +++ b/configs/trainer/default.yaml @@ -0,0 +1,39 @@ +# Default Trainer Configuration +# ============================== +# Standard training settings + +name: "default" + +# Training epochs +max_epochs: 100 +min_epochs: 10 + +# Learning rate +learning_rate: 1e-3 +lr_scheduler: "cosine" # cosine, step, plateau, none +lr_warmup_epochs: 5 +lr_min: 1e-6 + +# Optimizer +optimizer: "adamw" +weight_decay: 1e-4 +betas: [0.9, 0.999] +eps: 1e-8 + +# Gradient handling +gradient_clip_val: 1.0 +gradient_clip_algorithm: "norm" # norm, value +accumulate_grad_batches: 1 + +# Early stopping +early_stopping: true +patience: 20 +monitor: "val/r2" +mode: "max" + +# Validation +val_check_interval: 1.0 # Check every epoch +check_val_every_n_epoch: 1 + +# Logging +log_every_n_steps: 10 diff --git a/configs/trainer/progressive.yaml b/configs/trainer/progressive.yaml new file mode 100644 index 0000000..c91587e --- /dev/null +++ b/configs/trainer/progressive.yaml @@ -0,0 +1,43 @@ +# Progressive Training Configuration +# =================================== +# Three-phase training for VQ-VAE models + +name: "progressive" + +# Phase 1: Pre-training (encoder only) +pretrain: + epochs: 30 + learning_rate: 1e-3 + freeze_quantizer: true + description: "Pre-train encoder without VQ" + +# Phase 2: Codebook initialization +init_codebook: + method: "kmeans" + n_init: 10 + max_iter: 300 + description: "K-means initialization from encoder outputs" + +# Phase 3: Fine-tuning (full model) +finetune: + epochs: 50 + learning_rate: 3e-4 + freeze_quantizer: false + description: "Fine-tune with VQ enabled" + +# Optimizer (shared) +optimizer: "adamw" +weight_decay: 1e-4 + +# Learning rate schedule +lr_scheduler: "cosine" +lr_warmup_epochs: 3 + +# Early stopping (per phase) +early_stopping: true +patience: 15 +monitor: "val/r2" +mode: "max" + +# Gradient handling +gradient_clip_val: 1.0 diff --git a/docs/MIGRATION_GUIDE.md b/docs/MIGRATION_GUIDE.md new file mode 100644 index 0000000..198cad3 --- /dev/null +++ b/docs/MIGRATION_GUIDE.md @@ -0,0 +1,165 @@ +# Migration Guide: v0.1 → v0.2 + +This guide helps you migrate from the old "one script per experiment" approach to the new configuration-driven system. + +## Key Changes + +### 1. Single Entry Point + +**Before (v0.1):** +```bash +python python/exp25_mamba_mcrtt.py +python python/exp23_statistical_validation.py +python python/exp22c_multiseed_teacher.py +``` + +**After (v0.2):** +```bash +python train.py experiment=exp25_mamba +python train.py experiment=exp23_validation +python train.py experiment=exp22c_teacher +``` + +### 2. Configuration Files Replace Hardcoded Params + +**Before:** Hyperparameters scattered in `exp25_mamba_mcrtt.py` (lines 50-100) +```python +d_model = 128 +n_layers = 4 +window_size = 80 # 2 seconds +batch_size = 32 +learning_rate = 1e-3 +``` + +**After:** Centralized in `configs/experiment/exp25_mamba.yaml` +```yaml +model: + d_model: 128 + n_layers: 4 + window_size: 80 + +dataset: + batch_size: 32 + +trainer: + learning_rate: 1e-3 +``` + +### 3. Import Paths Changed + +**Before:** +```python +from python.phantomx.model import ProgressiveVQVAE +from python.phantomx.trainer import ProgressiveTrainer +``` + +**After:** +```python +from src.models import build_model, ProgressiveVQVAE +from src.trainer import Trainer, ProgressiveTrainer +from src.datamodules import build_datamodule +``` + +### 4. Experiment Tracking + +**Before:** Manual README.md tables + +**After:** Automatic WandB logging +```bash +# Disable if needed +python train.py logging.use_wandb=false +``` + +## Mapping Old Experiments to New Configs + +| Old Script | New Command | +|------------|-------------| +| `exp25_mamba_mcrtt.py` | `python train.py experiment=exp25_mamba` | +| `exp22c_multiseed_teacher.py` | `python train.py experiment=exp22c_teacher` | +| `exp23_statistical_validation.py` | `python train.py experiment=exp23_validation` | +| `exp10_beat_lstm.py` | `python train.py experiment=lstm_baseline` | + +## Creating New Experiments + +### 1. Create an experiment config + +```yaml +# configs/experiment/my_experiment.yaml +# @package _global_ +defaults: + - override /model: mamba + - override /dataset: mc_rtt + - override /augmentation: strong + +experiment_name: "my_experiment" +notes: "Testing new idea..." + +model: + n_layers: 6 + d_model: 256 + +seed: 42 +``` + +### 2. Run it + +```bash +python train.py experiment=my_experiment +``` + +### 3. Compare variations + +```bash +python train.py experiment=my_experiment --multirun model.n_layers=4,6,8 +``` + +## Common Migration Issues + +### Issue: ModuleNotFoundError + +**Solution:** The new structure uses `src/` instead of `python/phantomx/` +```bash +# Make sure you're in the project root +cd PhantomX +pip install -e . +``` + +### Issue: Config not found + +**Solution:** Run from project root where `configs/` is located +```bash +# Wrong (inside configs/) +cd configs && python ../train.py # ❌ + +# Right (project root) +python train.py # ✓ +``` + +### Issue: Different results than before + +**Checklist:** +1. Same seed? (`seed: 42`) +2. Same augmentation? (`augmentation: standard` vs `none`) +3. Same window size? (`model.window_size`) +4. Same batch size? (`dataset.batch_size`) + +Use the config to ensure reproducibility: +```bash +python train.py experiment=exp25_mamba --cfg job +``` + +## Backward Compatibility + +The old `python/phantomx/` code still works. You can gradually migrate: + +```python +# This still works +from python.phantomx.model import ProgressiveVQVAE +``` + +But we recommend migrating to enjoy: +- ✅ Automatic config saving +- ✅ WandB integration +- ✅ Reproducible experiments +- ✅ Easy hyperparameter sweeps +- ✅ No more code duplication diff --git a/evaluate.py b/evaluate.py new file mode 100644 index 0000000..095014a --- /dev/null +++ b/evaluate.py @@ -0,0 +1,130 @@ +#!/usr/bin/env python +""" +Evaluate a trained PhantomX model. + +Usage: + python evaluate.py checkpoint_path [--dataset mc_maze|mc_rtt] + +Example: + python evaluate.py logs/exp25_mamba/best_model.pt --dataset mc_rtt +""" + +import argparse +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent)) + +import torch +from omegaconf import OmegaConf + +from src.models import build_model +from src.datamodules import build_datamodule +from src.utils.metrics import compute_metrics +from src.utils.seeding import seed_everything + + +def main(): + parser = argparse.ArgumentParser(description="Evaluate PhantomX model") + parser.add_argument("checkpoint", type=str, help="Path to model checkpoint") + parser.add_argument("--dataset", type=str, default="mc_maze", choices=["mc_maze", "mc_rtt"]) + parser.add_argument("--split", type=str, default="test", choices=["val", "test"]) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--device", type=str, default="auto") + args = parser.parse_args() + + # Load checkpoint + checkpoint_path = Path(args.checkpoint) + if not checkpoint_path.exists(): + print(f"❌ Checkpoint not found: {checkpoint_path}") + sys.exit(1) + + print(f"📂 Loading checkpoint: {checkpoint_path}") + checkpoint = torch.load(checkpoint_path, map_location="cpu") + + # Get config from checkpoint or use defaults + if "config" in checkpoint: + cfg = OmegaConf.create(checkpoint["config"]) + else: + # Load default configs + cfg = OmegaConf.load("configs/config.yaml") + model_cfg = OmegaConf.load(f"configs/model/vqvae.yaml") + dataset_cfg = OmegaConf.load(f"configs/dataset/{args.dataset}.yaml") + cfg = OmegaConf.merge(cfg, {"model": model_cfg, "dataset": dataset_cfg}) + + # Override dataset if specified + if args.dataset: + dataset_cfg = OmegaConf.load(f"configs/dataset/{args.dataset}.yaml") + cfg.dataset = OmegaConf.merge(cfg.dataset, dataset_cfg) + + seed_everything(args.seed) + + # Device + if args.device == "auto": + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + else: + device = torch.device(args.device) + + print(f"🖥️ Device: {device}") + + # Load data + print(f"\n📊 Loading {args.dataset} dataset...") + datamodule = build_datamodule(cfg) + datamodule.setup() + + if args.split == "test": + dataloader = datamodule.test_dataloader() + else: + dataloader = datamodule.val_dataloader() + + # Build model + print(f"\n🔧 Building model...") + model = build_model(cfg, n_channels=datamodule.n_channels) + + # Load weights + if "model_state_dict" in checkpoint: + model.load_state_dict(checkpoint["model_state_dict"]) + else: + model.load_state_dict(checkpoint) + + model = model.to(device) + model.eval() + + n_params = sum(p.numel() for p in model.parameters()) + print(f" Parameters: {n_params:,}") + + # Evaluate + print(f"\n🎯 Evaluating on {args.split} set...") + + all_preds = [] + all_targets = [] + + with torch.no_grad(): + for neural, target in dataloader: + neural = neural.to(device) + output, _ = model(neural) + + all_preds.append(output.cpu()) + all_targets.append(target) + + preds = torch.cat(all_preds, dim=0) + targets = torch.cat(all_targets, dim=0) + + metrics = compute_metrics(preds, targets, prefix="") + + # Print results + print("\n" + "=" * 40) + print("Results") + print("=" * 40) + print(f" R²: {metrics['r2']:.4f}") + print(f" R² (x): {metrics.get('r2_x', 'N/A'):.4f}" if 'r2_x' in metrics else "") + print(f" R² (y): {metrics.get('r2_y', 'N/A'):.4f}" if 'r2_y' in metrics else "") + print(f" MSE: {metrics['mse']:.6f}") + print(f" RMSE: {metrics['rmse']:.6f}") + print("=" * 40) + + return metrics["r2"] + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 29de460..1f78159 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,8 +4,8 @@ build-backend = "setuptools.build_meta" [project] name = "phantomx" -version = "0.1.0" -description = "LaBraM-POYO Neural Foundation Model: Population-geometry BCI decoding with electrode-dropout robustness" +version = "0.2.0" +description = "PhantomX: Configuration-Driven Neural Foundation Model for BCI Decoding" authors = [ {name = "Youssef El abbassi", email = "youssef@elabbassi.com"} ] @@ -13,7 +13,7 @@ readme = "README.md" license = {text = "MIT"} requires-python = ">=3.10" classifiers = [ - "Development Status :: 3 - Alpha", + "Development Status :: 4 - Beta", "Intended Audience :: Science/Research", "Topic :: Scientific/Engineering :: Artificial Intelligence", "License :: OSI Approved :: MIT License", @@ -22,17 +22,32 @@ classifiers = [ "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", ] -keywords = ["bci", "neural-decoding", "vq-vae", "brain-computer-interface", "deep-learning"] +keywords = ["bci", "neural-decoding", "vq-vae", "mamba", "brain-computer-interface", "deep-learning", "hydra"] dependencies = [ + # Core ML "torch>=2.1.0", "numpy>=1.24.0", "scipy>=1.11.0", + "scikit-learn>=1.3.0", + "einops>=0.7.0", + + # Data "pynwb>=2.5.0", "h5py>=3.9.0", - "einops>=0.7.0", + + # Configuration (Hydra) + "hydra-core>=1.3.0", + "hydra-colorlog>=1.2.0", + "omegaconf>=2.3.0", + + # Experiment Tracking + "wandb>=0.16.0", + + # Utilities "matplotlib>=3.7.0", "tqdm>=4.66.0", + "rich>=13.0.0", ] [project.optional-dependencies] @@ -59,8 +74,11 @@ notebooks = [ "seaborn>=0.12.0", "plotly>=5.14.0", ] +lightning = [ + "lightning>=2.1.0", +] all = [ - "phantomx[dev,onnx,integration,notebooks]", + "phantomx[dev,onnx,integration,notebooks,lightning]", ] [project.urls] @@ -69,14 +87,15 @@ all = [ "Documentation" = "https://github.com/yelabb/PhantomX#readme" [project.scripts] -phantomx-train = "phantomx.cli:train_cli" -phantomx-test = "phantomx.cli:test_cli" +phantomx-train = "src.cli:train_cli" +phantomx-test = "src.cli:test_cli" [tool.setuptools.packages.find] -where = ["python"] +where = ["."] +include = ["src*", "python*"] [tool.setuptools.package-dir] -"" = "python" +"" = "." [tool.pytest.ini_options] testpaths = ["tests"] diff --git a/requirements.txt b/requirements.txt index cadaf7d..61c852b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,31 +1,43 @@ -# PhantomX: LaBraM-POYO Stack Dependencies +# PhantomX: Configuration-Driven Neural Foundation Model +# v0.2.0 -# Core ML Framework +# ===== Core ML Framework ===== torch>=2.1.0 torchvision>=0.16.0 -# Neural Data Processing +# ===== Neural Data Processing ===== numpy>=1.24.0 scipy>=1.11.0 +scikit-learn>=1.3.0 pynwb>=2.5.0 h5py>=3.9.0 -# Training & Optimization +# ===== Model Components ===== einops>=0.7.0 timm>=0.9.0 # For modern architectures if needed -# Visualization & Analysis +# ===== Configuration (Hydra) ===== +hydra-core>=1.3.0 +hydra-colorlog>=1.2.0 +omegaconf>=2.3.0 + +# ===== Experiment Tracking ===== +wandb>=0.16.0 + +# ===== Visualization & Analysis ===== matplotlib>=3.7.0 seaborn>=0.12.0 plotly>=5.14.0 +rich>=13.0.0 +tqdm>=4.66.0 -# Development & Testing +# ===== Development & Testing ===== pytest>=7.4.0 pytest-cov>=4.1.0 black>=23.7.0 ruff>=0.0.280 -# Integration with PhantomLink +# ===== Integration with PhantomLink ===== msgpack>=1.0.5 websockets>=11.0.3 fastapi>=0.100.0 diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000..73351ba --- /dev/null +++ b/src/__init__.py @@ -0,0 +1,13 @@ +""" +PhantomX: Neural Foundation Model for BCI Decoding + +A configuration-driven research framework for neural decoding experiments. +""" + +__version__ = "0.2.0" + +from .models import build_model +from .datamodules import build_datamodule +from .trainer import Trainer + +__all__ = ["build_model", "build_datamodule", "Trainer", "__version__"] diff --git a/src/datamodules/__init__.py b/src/datamodules/__init__.py new file mode 100644 index 0000000..c8323fe --- /dev/null +++ b/src/datamodules/__init__.py @@ -0,0 +1,79 @@ +""" +PhantomX DataModules + +PyTorch Lightning-style data wrappers for neural datasets. +""" + +from typing import Dict, Type +from omegaconf import DictConfig + +import hydra + +# Registry +DATAMODULE_REGISTRY: Dict[str, Type] = {} + + +def register_datamodule(name: str): + """Decorator to register a datamodule.""" + def decorator(cls): + DATAMODULE_REGISTRY[name] = cls + return cls + return decorator + + +def build_datamodule(cfg: DictConfig): + """ + Build a datamodule from configuration. + + Args: + cfg: Full configuration (must contain 'dataset' key) + + Returns: + Instantiated datamodule + """ + dataset_cfg = cfg.dataset + + # Try Hydra instantiate first + if '_target_' in dataset_cfg: + return hydra.utils.instantiate(dataset_cfg) + + # Fall back to registry + name = dataset_cfg.get('name', 'mc_maze') + _import_all_datamodules() + + if name not in DATAMODULE_REGISTRY: + available = ", ".join(sorted(DATAMODULE_REGISTRY.keys())) + raise KeyError(f"DataModule '{name}' not found. Available: {available}") + + return DATAMODULE_REGISTRY[name](cfg=dataset_cfg) + + +def list_datamodules() -> list: + """List all registered datamodules.""" + _import_all_datamodules() + return sorted(DATAMODULE_REGISTRY.keys()) + + +def _import_all_datamodules(): + """Import all datamodule modules.""" + try: + from . import mc_maze + except ImportError: + pass + try: + from . import mc_rtt + except ImportError: + pass + + +from .base import BaseDataModule +from .mc_maze import MCMazeDataModule +from .mc_rtt import MCRTTDataModule + +__all__ = [ + "build_datamodule", + "list_datamodules", + "BaseDataModule", + "MCMazeDataModule", + "MCRTTDataModule", +] diff --git a/src/datamodules/base.py b/src/datamodules/base.py new file mode 100644 index 0000000..0e5d2b1 --- /dev/null +++ b/src/datamodules/base.py @@ -0,0 +1,253 @@ +""" +Base DataModule + +Abstract base class for all neural data modules. +""" + +from abc import ABC, abstractmethod +from pathlib import Path +from typing import Optional, Tuple, Dict, Any + +import numpy as np +import torch +from torch.utils.data import Dataset, DataLoader, random_split +from omegaconf import DictConfig + + +class WindowedNeuralDataset(Dataset): + """ + Windowed neural dataset for velocity decoding. + + Extracts overlapping windows of neural data and corresponding targets. + """ + + def __init__( + self, + neural_data: np.ndarray, + targets: np.ndarray, + window_size: int = 10, + stride: int = 1, + normalize: bool = True, + mean: Optional[np.ndarray] = None, + std: Optional[np.ndarray] = None, + ): + """ + Args: + neural_data: [T, n_channels] spike counts + targets: [T, output_dim] velocity/position + window_size: Number of time bins per window + stride: Stride between windows + normalize: Whether to z-score normalize + mean, std: Pre-computed normalization stats + """ + self.window_size = window_size + self.stride = stride + + # Normalize if requested + if normalize: + if mean is None: + mean = neural_data.mean(axis=0) + if std is None: + std = neural_data.std(axis=0) + 1e-8 + neural_data = (neural_data - mean) / std + + self.mean = mean + self.std = std + + # Convert to tensors + self.neural = torch.from_numpy(neural_data).float() + self.targets = torch.from_numpy(targets).float() + + # Compute valid indices + n_samples = len(neural_data) - window_size + 1 + self.indices = list(range(0, n_samples, stride)) + + def __len__(self) -> int: + return len(self.indices) + + def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]: + start = self.indices[idx] + end = start + self.window_size + + # Neural window + neural_window = self.neural[start:end] # [window_size, n_channels] + + # Target at end of window (or middle, depending on task) + target = self.targets[end - 1] # [output_dim] + + return neural_window, target + + +class BaseDataModule(ABC): + """ + Abstract base class for data modules. + + Provides: + - Data loading and preprocessing + - Train/val/test splits + - DataLoader creation + - Augmentation integration + """ + + def __init__( + self, + cfg: Optional[DictConfig] = None, + data_dir: str = "data", + batch_size: int = 64, + num_workers: int = 4, + pin_memory: bool = True, + window_size: int = 10, + stride: int = 1, + train_ratio: float = 0.7, + val_ratio: float = 0.1, + test_ratio: float = 0.2, + shuffle: bool = True, + normalize: bool = True, + **kwargs + ): + # Use config if provided + if cfg is not None: + data_dir = cfg.get('data_dir', data_dir) + batch_size = cfg.get('batch_size', batch_size) + window_size = cfg.get('window_size', window_size) if 'window_size' in cfg else kwargs.get('window_size', window_size) + train_ratio = cfg.get('train_ratio', train_ratio) + val_ratio = cfg.get('val_ratio', val_ratio) + test_ratio = cfg.get('test_ratio', test_ratio) + shuffle = cfg.get('shuffle', shuffle) + normalize = cfg.get('normalize', normalize) + + self.data_dir = Path(data_dir) + self.batch_size = batch_size + self.num_workers = num_workers + self.pin_memory = pin_memory + self.window_size = window_size + self.stride = stride + self.train_ratio = train_ratio + self.val_ratio = val_ratio + self.test_ratio = test_ratio + self.shuffle = shuffle + self.normalize = normalize + + # Data storage + self.neural_data: Optional[np.ndarray] = None + self.targets: Optional[np.ndarray] = None + self.train_dataset: Optional[Dataset] = None + self.val_dataset: Optional[Dataset] = None + self.test_dataset: Optional[Dataset] = None + + # Normalization stats + self.mean: Optional[np.ndarray] = None + self.std: Optional[np.ndarray] = None + + # Dynamically determined from data + self._n_channels: Optional[int] = None + + @property + @abstractmethod + def name(self) -> str: + """Dataset name.""" + pass + + @property + def n_channels(self) -> int: + """Number of neural channels (determined from loaded data).""" + if self._n_channels is not None: + return self._n_channels + if self.neural_data is not None: + return self.neural_data.shape[1] + # Return a placeholder - will be updated after loading + return self._get_expected_n_channels() + + def _get_expected_n_channels(self) -> int: + """Override to provide expected channel count before data is loaded.""" + return 137 # Default + + @property + @abstractmethod + def output_dim(self) -> int: + """Output dimension (e.g., 2 for velocity).""" + pass + + @abstractmethod + def _load_raw_data(self) -> Tuple[np.ndarray, np.ndarray]: + """Load raw neural data and targets.""" + pass + + def setup(self, stage: Optional[str] = None): + """ + Setup data splits. + + Args: + stage: 'fit', 'test', or None (all) + """ + if self.neural_data is None: + self.neural_data, self.targets = self._load_raw_data() + # Update n_channels from loaded data + self._n_channels = self.neural_data.shape[1] + + # Compute normalization stats from all data + if self.normalize: + self.mean = self.neural_data.mean(axis=0) + self.std = self.neural_data.std(axis=0) + 1e-8 + + # Create datasets with shared normalization + full_dataset = WindowedNeuralDataset( + self.neural_data, + self.targets, + window_size=self.window_size, + stride=self.stride, + normalize=self.normalize, + mean=self.mean, + std=self.std, + ) + + # Split based on WINDOWED dataset length (not raw data length) + n_samples = len(full_dataset) + n_train = int(n_samples * self.train_ratio) + n_val = int(n_samples * self.val_ratio) + n_test = n_samples - n_train - n_val + + # Split + self.train_dataset, self.val_dataset, self.test_dataset = random_split( + full_dataset, + [n_train, n_val, n_test], + generator=torch.Generator().manual_seed(42) + ) + + print(f"📊 {self.name}: train={len(self.train_dataset)}, val={len(self.val_dataset)}, test={len(self.test_dataset)}") + + def _should_pin_memory(self) -> bool: + """Only pin memory if CUDA is available.""" + import torch + return self.pin_memory and torch.cuda.is_available() + + def train_dataloader(self) -> DataLoader: + """Training dataloader.""" + return DataLoader( + self.train_dataset, + batch_size=self.batch_size, + shuffle=self.shuffle, + num_workers=self.num_workers, + pin_memory=self._should_pin_memory(), + drop_last=True, + ) + + def val_dataloader(self) -> DataLoader: + """Validation dataloader.""" + return DataLoader( + self.val_dataset, + batch_size=self.batch_size, + shuffle=False, + num_workers=self.num_workers, + pin_memory=self._should_pin_memory(), + ) + + def test_dataloader(self) -> DataLoader: + """Test dataloader.""" + return DataLoader( + self.test_dataset, + batch_size=self.batch_size, + shuffle=False, + num_workers=self.num_workers, + pin_memory=self._should_pin_memory(), + ) diff --git a/src/datamodules/mc_maze.py b/src/datamodules/mc_maze.py new file mode 100644 index 0000000..49f3f2b --- /dev/null +++ b/src/datamodules/mc_maze.py @@ -0,0 +1,99 @@ +""" +MC_Maze DataModule + +Motor cortex recordings during maze reaching task. +Discrete reaching movements with pauses. +""" + +from pathlib import Path +from typing import Tuple, Optional +import numpy as np + +from .base import BaseDataModule +from . import register_datamodule + + +@register_datamodule('mc_maze') +class MCMazeDataModule(BaseDataModule): + """ + MC_Maze DataModule. + + Motor cortex recordings during a delayed center-out reaching task. + 137 sorted units, discrete trials. + """ + + @property + def name(self) -> str: + return "mc_maze" + + def _get_expected_n_channels(self) -> int: + """Expected channel count (actual may differ based on data file).""" + return 142 # Can vary; actual determined from data + + @property + def output_dim(self) -> int: + return 2 # Velocity (vx, vy) + + def _load_raw_data(self) -> Tuple[np.ndarray, np.ndarray]: + """Load MC_Maze data from NWB file.""" + filepath = self.data_dir / "mc_maze.nwb" + + if not filepath.exists(): + raise FileNotFoundError( + f"MC_Maze data not found at {filepath}. " + "Download from DANDI Archive or run `python upload_data.py`" + ) + + try: + from pynwb import NWBHDF5IO + except ImportError: + raise ImportError("PyNWB required. Install with: pip install pynwb") + + print(f"Loading MC_Maze from: {filepath}") + + bin_size_ms = 25.0 # 40 Hz + + with NWBHDF5IO(str(filepath), mode='r', load_namespaces=True) as io: + nwb = io.read() + + # Get neural units + units = nwb.units + n_units = len(units.id[:]) + + # Get behavior + behavior = nwb.processing.get('behavior') + hand_vel = behavior.data_interfaces.get('hand_vel') + cursor_pos = behavior.data_interfaces.get('cursor_pos') + + timestamps = cursor_pos.timestamps[:] + velocity = hand_vel.data[:].astype(np.float32) + + # Bin spikes + duration = timestamps[-1] - timestamps[0] + n_bins = int(duration * 1000 / bin_size_ms) + + spike_counts = np.zeros((n_bins, n_units), dtype=np.float32) + + for unit_idx in range(n_units): + spike_times = units.get_unit_spike_times(unit_idx) + if spike_times is not None and len(spike_times) > 0: + bin_indices = ((spike_times - timestamps[0]) * 1000 / bin_size_ms).astype(np.int32) + bin_indices = np.clip(bin_indices, 0, n_bins - 1) + for idx in bin_indices: + spike_counts[idx, unit_idx] += 1 + + # Bin velocities + bin_samples = int(bin_size_ms) + rate = 1000 # 1000 Hz behavior data + binned_vel = np.zeros((n_bins, 2), dtype=np.float32) + + for i in range(n_bins): + start_time = timestamps[0] + i * bin_size_ms / 1000 + end_time = start_time + bin_size_ms / 1000 + mask = (timestamps >= start_time) & (timestamps < end_time) + if mask.any(): + binned_vel[i] = velocity[mask].mean(axis=0) + + print(f" Loaded: {spike_counts.shape[0]} bins, {spike_counts.shape[1]} channels") + + return spike_counts, binned_vel diff --git a/src/datamodules/mc_rtt.py b/src/datamodules/mc_rtt.py new file mode 100644 index 0000000..15afb8d --- /dev/null +++ b/src/datamodules/mc_rtt.py @@ -0,0 +1,111 @@ +""" +MC_RTT DataModule + +Motor cortex Random Target Tracking task. +Continuous tracking - fundamentally different from MC_Maze. +""" + +from pathlib import Path +from typing import Tuple +import numpy as np + +from .base import BaseDataModule +from . import register_datamodule + + +@register_datamodule('mc_rtt') +class MCRTTDataModule(BaseDataModule): + """ + MC_RTT DataModule. + + Continuous random target tracking task. + 130 neural units, continuous data (no discrete trials). + + Key differences from MC_Maze: + - Uses finger_vel instead of hand_vel + - Continuous tracking, no discrete reaches + - Longer context important (trajectory integration) + """ + + def __init__(self, **kwargs): + super().__init__(**kwargs) + # Override defaults for continuous data + if 'shuffle' not in kwargs and kwargs.get('cfg') is None: + self.shuffle = False # Sequential for stateful models + + @property + def name(self) -> str: + return "mc_rtt" + + def _get_expected_n_channels(self) -> int: + """Expected channel count (actual may differ based on data file).""" + return 130 # Can vary; actual determined from data + + @property + def output_dim(self) -> int: + return 2 # Finger velocity (vx, vy) + + def _load_raw_data(self) -> Tuple[np.ndarray, np.ndarray]: + """Load MC_RTT data from NWB file.""" + filepath = self.data_dir / "mc_rtt.nwb" + + if not filepath.exists(): + raise FileNotFoundError( + f"MC_RTT data not found at {filepath}. " + "Download from DANDI Archive or run `python upload_data.py`" + ) + + try: + from pynwb import NWBHDF5IO + except ImportError: + raise ImportError("PyNWB required. Install with: pip install pynwb") + + print(f"Loading MC_RTT from: {filepath}") + + bin_size_ms = 25.0 # 40 Hz + + with NWBHDF5IO(str(filepath), mode='r', load_namespaces=True) as io: + nwb = io.read() + + # Get neural units + units = nwb.units + n_units = len(units.id[:]) + + # Get behavior - MC_RTT uses finger_vel + behavior = nwb.processing.get('behavior') + finger_vel = behavior.data_interfaces.get('finger_vel') + + velocity = finger_vel.data[:] # [T, 2] + rate = finger_vel.rate # 1000 Hz + n_samples = len(velocity) + + # Bin size in samples + bin_samples = int(bin_size_ms * rate / 1000) + n_bins = n_samples // bin_samples + + print(f" Raw: {n_samples} samples at {rate}Hz = {n_samples/rate:.1f}s") + print(f" Binning: {bin_samples} samples/bin → {n_bins} bins at {1000/bin_size_ms:.0f}Hz") + + # Bin spikes + spike_counts = np.zeros((n_bins, n_units), dtype=np.float32) + + for unit_idx in range(n_units): + spike_times = units.get_unit_spike_times(unit_idx) + if spike_times is not None and len(spike_times) > 0: + bin_indices = (spike_times * 1000 / bin_size_ms).astype(np.int32) + bin_indices = np.clip(bin_indices, 0, n_bins - 1) + for idx in bin_indices: + spike_counts[idx, unit_idx] += 1 + + # Bin velocities + binned_vel = np.zeros((n_bins, 2), dtype=np.float32) + for i in range(n_bins): + start_idx = i * bin_samples + end_idx = start_idx + bin_samples + binned_vel[i] = velocity[start_idx:end_idx].mean(axis=0) + + print(f" Loaded: {spike_counts.shape[0]} bins, {spike_counts.shape[1]} channels") + mean_rate = spike_counts.mean() * (1000 / bin_size_ms) + print(f" Mean firing rate: {mean_rate:.2f} spikes/s") + + return spike_counts, binned_vel diff --git a/src/trainer.py b/src/trainer.py new file mode 100644 index 0000000..f4515a3 --- /dev/null +++ b/src/trainer.py @@ -0,0 +1,377 @@ +""" +PhantomX Unified Trainer + +Generic trainer supporting all model types and training modes. +Replaces the scattered training logic in exp*.py files. +""" + +from pathlib import Path +from typing import Optional, Dict, Any, Callable +from dataclasses import dataclass +import time + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.utils.data import DataLoader +from torch.optim import AdamW +from torch.optim.lr_scheduler import CosineAnnealingLR, ReduceLROnPlateau +from omegaconf import DictConfig +from tqdm import tqdm + +from .utils.metrics import compute_metrics, MetricTracker +from .utils.logging import log_metrics + + +@dataclass +class TrainerState: + """Trainer state for checkpointing.""" + epoch: int + global_step: int + best_metric: float + best_epoch: int + + +class Trainer: + """ + Unified trainer for PhantomX models. + + Features: + - Standard training loop with validation + - Progressive training for VQ-VAE + - Early stopping and checkpointing + - Metric tracking and logging + """ + + def __init__( + self, + model: nn.Module, + train_loader: DataLoader, + val_loader: DataLoader, + cfg: DictConfig, + test_loader: Optional[DataLoader] = None, + device: str = "auto", + ): + self.model = model + self.train_loader = train_loader + self.val_loader = val_loader + self.test_loader = test_loader + self.cfg = cfg + + # Device setup + if device == "auto": + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + else: + self.device = torch.device(device) + + self.model = self.model.to(self.device) + + # Get trainer config + trainer_cfg = cfg.get('trainer', {}) + + # Optimizer + self.optimizer = AdamW( + model.parameters(), + lr=trainer_cfg.get('learning_rate', 1e-3), + weight_decay=trainer_cfg.get('weight_decay', 1e-4), + betas=tuple(trainer_cfg.get('betas', [0.9, 0.999])), + ) + + # LR Scheduler + max_epochs = trainer_cfg.get('max_epochs', 100) + scheduler_type = trainer_cfg.get('lr_scheduler', 'cosine') + + if scheduler_type == 'cosine': + self.scheduler = CosineAnnealingLR( + self.optimizer, + T_max=max_epochs, + eta_min=trainer_cfg.get('lr_min', 1e-6) + ) + elif scheduler_type == 'plateau': + self.scheduler = ReduceLROnPlateau( + self.optimizer, + mode='max', + factor=0.5, + patience=10, + ) + else: + self.scheduler = None + + # Training config + self.max_epochs = max_epochs + self.gradient_clip_val = trainer_cfg.get('gradient_clip_val', 1.0) + self.log_every_n_steps = cfg.get('logging', {}).get('log_every_n_steps', 10) + + # Early stopping + self.early_stopping = trainer_cfg.get('early_stopping', True) + self.patience = trainer_cfg.get('patience', 20) + self.monitor = trainer_cfg.get('monitor', 'val/r2') + + # Metric tracker + self.metric_tracker = MetricTracker( + metric_name='r2', + mode='max' + ) + + # State + self.state = TrainerState( + epoch=0, + global_step=0, + best_metric=-float('inf'), + best_epoch=0, + ) + self.best_model_state = None + + # Output directory + self.output_dir = Path(cfg.get('output_dir', 'logs')) + self.output_dir.mkdir(parents=True, exist_ok=True) + + def train(self) -> Dict[str, Any]: + """ + Run full training loop. + + Returns: + Dict with training results + """ + print(f"\n🚀 Starting training on {self.device}") + print(f" Model: {type(self.model).__name__}") + print(f" Epochs: {self.max_epochs}") + print(f" Train batches: {len(self.train_loader)}") + print(f" Val batches: {len(self.val_loader)}") + + start_time = time.time() + + for epoch in range(self.max_epochs): + self.state.epoch = epoch + + # Training + train_metrics = self._train_epoch() + + # Validation + val_metrics = self._validate() + + # Log metrics + all_metrics = {**train_metrics, **val_metrics, 'epoch': epoch} + log_metrics(all_metrics, step=epoch) + + # Update LR scheduler + if self.scheduler is not None: + if isinstance(self.scheduler, ReduceLROnPlateau): + self.scheduler.step(val_metrics['val/r2']) + else: + self.scheduler.step() + + # Track best model + r2 = val_metrics['val/r2'] + is_best = self.metric_tracker.update(r2, epoch) + + if is_best: + self.state.best_metric = r2 + self.state.best_epoch = epoch + self.best_model_state = {k: v.cpu().clone() for k, v in self.model.state_dict().items()} + self._save_checkpoint('best_model.pt') + + # Log progress + lr = self.optimizer.param_groups[0]['lr'] + print(f"Epoch {epoch:3d} | Train Loss: {train_metrics['train/loss']:.4f} | " + f"Val R²: {r2:.4f} | LR: {lr:.2e}" + (" ⭐" if is_best else "")) + + # Early stopping + if self.early_stopping and self.metric_tracker.should_stop(self.patience): + print(f"\n⚠️ Early stopping at epoch {epoch} (no improvement for {self.patience} epochs)") + break + + # Load best model + if self.best_model_state is not None: + self.model.load_state_dict(self.best_model_state) + + training_time = time.time() - start_time + + # Final test + test_metrics = {} + if self.test_loader is not None: + test_metrics = self._test() + print(f"\n🎯 Test R²: {test_metrics['test/r2']:.4f}") + + results = { + 'best_r2': self.state.best_metric, + 'best_epoch': self.state.best_epoch, + 'training_time_s': training_time, + **test_metrics, + } + + print(f"\n✅ Training complete!") + print(f" Best R²: {self.state.best_metric:.4f} (epoch {self.state.best_epoch})") + print(f" Time: {training_time/60:.1f} minutes") + + return results + + def _train_epoch(self) -> Dict[str, float]: + """Run one training epoch.""" + self.model.train() + + total_loss = 0.0 + all_preds = [] + all_targets = [] + + pbar = tqdm(self.train_loader, desc=f"Epoch {self.state.epoch}", leave=False) + + for batch_idx, (neural, target) in enumerate(pbar): + neural = neural.to(self.device) + target = target.to(self.device) + + # Forward + self.optimizer.zero_grad() + output, info = self.model(neural) + + # Loss + loss = F.mse_loss(output, target) + + # Add VQ loss if present + if 'commitment_loss' in info: + loss = loss + info['commitment_loss'] + + # Backward + loss.backward() + + # Gradient clipping + if self.gradient_clip_val > 0: + torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.gradient_clip_val) + + self.optimizer.step() + + # Track + total_loss += loss.item() + all_preds.append(output.detach()) + all_targets.append(target.detach()) + + self.state.global_step += 1 + + # Update progress bar + pbar.set_postfix({'loss': loss.item()}) + + # Compute metrics + preds = torch.cat(all_preds, dim=0) + targets = torch.cat(all_targets, dim=0) + metrics = compute_metrics(preds, targets, prefix='train/') + metrics['train/loss'] = total_loss / len(self.train_loader) + + return metrics + + @torch.no_grad() + def _validate(self) -> Dict[str, float]: + """Run validation.""" + self.model.eval() + + all_preds = [] + all_targets = [] + total_loss = 0.0 + + for neural, target in self.val_loader: + neural = neural.to(self.device) + target = target.to(self.device) + + output, _ = self.model(neural) + loss = F.mse_loss(output, target) + + total_loss += loss.item() + all_preds.append(output) + all_targets.append(target) + + preds = torch.cat(all_preds, dim=0) + targets = torch.cat(all_targets, dim=0) + + metrics = compute_metrics(preds, targets, prefix='val/') + metrics['val/loss'] = total_loss / len(self.val_loader) + + return metrics + + @torch.no_grad() + def _test(self) -> Dict[str, float]: + """Run test evaluation.""" + self.model.eval() + + all_preds = [] + all_targets = [] + + for neural, target in self.test_loader: + neural = neural.to(self.device) + target = target.to(self.device) + + output, _ = self.model(neural) + + all_preds.append(output) + all_targets.append(target) + + preds = torch.cat(all_preds, dim=0) + targets = torch.cat(all_targets, dim=0) + + return compute_metrics(preds, targets, prefix='test/') + + def _save_checkpoint(self, filename: str): + """Save model checkpoint.""" + path = self.output_dir / filename + torch.save({ + 'epoch': self.state.epoch, + 'model_state_dict': self.model.state_dict(), + 'optimizer_state_dict': self.optimizer.state_dict(), + 'best_metric': self.state.best_metric, + }, path) + + +class ProgressiveTrainer(Trainer): + """ + Progressive trainer for VQ-VAE models. + + Three phases: + 1. Pre-train encoder (VQ disabled) + 2. K-means codebook initialization + 3. Fine-tune with VQ enabled + """ + + def train(self) -> Dict[str, Any]: + """Run progressive training.""" + trainer_cfg = self.cfg.get('trainer', {}) + + # Phase 1: Pre-train + print("\n[Phase 1/3] Pre-training encoder...") + if hasattr(self.model, 'use_vq'): + self.model.use_vq = False + + pretrain_epochs = trainer_cfg.get('pretrain', {}).get('epochs', 30) + self.max_epochs = pretrain_epochs + pretrain_results = super().train() + + # Phase 2: Init codebook + print("\n[Phase 2/3] Initializing codebook with k-means...") + if hasattr(self.model, 'init_codebook'): + self.model.init_codebook(self.train_loader) + + # Phase 3: Fine-tune + print("\n[Phase 3/3] Fine-tuning with VQ...") + if hasattr(self.model, 'use_vq'): + self.model.use_vq = True + + # Reset optimizer with lower LR + finetune_lr = trainer_cfg.get('finetune', {}).get('learning_rate', 3e-4) + finetune_epochs = trainer_cfg.get('finetune', {}).get('epochs', 50) + + self.optimizer = AdamW( + self.model.parameters(), + lr=finetune_lr, + weight_decay=trainer_cfg.get('weight_decay', 1e-4), + ) + self.scheduler = CosineAnnealingLR(self.optimizer, T_max=finetune_epochs) + + # Reset state for fine-tuning + self.max_epochs = finetune_epochs + self.state = TrainerState(epoch=0, global_step=0, best_metric=-float('inf'), best_epoch=0) + self.metric_tracker = MetricTracker('r2', 'max') + + finetune_results = super().train() + + return { + 'pretrain_r2': pretrain_results.get('best_r2', 0), + 'finetune_r2': finetune_results.get('best_r2', 0), + **finetune_results, + } diff --git a/src/utils/__init__.py b/src/utils/__init__.py new file mode 100644 index 0000000..bc6aaba --- /dev/null +++ b/src/utils/__init__.py @@ -0,0 +1,21 @@ +""" +Utility modules for PhantomX. + +- logging: WandB integration and experiment tracking +- seeding: Reproducibility utilities +- metrics: Evaluation metrics (R², MSE, etc.) +- augmentation: Data augmentation transforms +""" + +from .logging import init_logger, log_metrics, finish_logger +from .seeding import seed_everything, get_git_hash +from .metrics import compute_metrics + +__all__ = [ + "init_logger", + "log_metrics", + "finish_logger", + "seed_everything", + "get_git_hash", + "compute_metrics", +] diff --git a/src/utils/augmentation.py b/src/utils/augmentation.py new file mode 100644 index 0000000..bd9ee64 --- /dev/null +++ b/src/utils/augmentation.py @@ -0,0 +1,232 @@ +""" +Data Augmentation Transforms + +Neural data augmentation techniques for robust BCI decoding. +Addresses the "Exp 22 forgot augmentation" issue by centralizing transforms. +""" + +from typing import Tuple, Optional, Dict, Any +from dataclasses import dataclass +import numpy as np +import torch +from omegaconf import DictConfig + + +@dataclass +class AugmentedSample: + """Container for augmented neural data.""" + neural: torch.Tensor + target: torch.Tensor + mask: Optional[torch.Tensor] = None + + +class Compose: + """Compose multiple transforms.""" + + def __init__(self, transforms: list): + self.transforms = [t for t in transforms if t is not None] + + def __call__(self, neural: torch.Tensor, target: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + for t in self.transforms: + neural, target = t(neural, target) + return neural, target + + +class ElectrodeDropout: + """ + Randomly drop electrodes during training. + + Key for robustness to electrode failure in real BCIs. + From Exp 22b: Critical augmentation for generalization. + + Args: + p: Probability of dropping each electrode + structured: If True, drop contiguous electrode groups + """ + + def __init__(self, p: float = 0.1, structured: bool = False): + self.p = p + self.structured = structured + + def __call__(self, neural: torch.Tensor, target: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + if not self.training: + return neural, target + + if self.structured: + return self._structured_dropout(neural, target) + else: + return self._random_dropout(neural, target) + + def _random_dropout(self, neural: torch.Tensor, target: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + # neural: [batch, time, channels] or [batch, channels] + if neural.dim() == 3: + n_channels = neural.shape[-1] + mask = torch.rand(n_channels, device=neural.device) > self.p + neural = neural * mask.unsqueeze(0).unsqueeze(0) + else: + n_channels = neural.shape[-1] + mask = torch.rand(n_channels, device=neural.device) > self.p + neural = neural * mask.unsqueeze(0) + return neural, target + + def _structured_dropout(self, neural: torch.Tensor, target: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + # Drop contiguous groups of electrodes + n_channels = neural.shape[-1] + if torch.rand(1).item() < self.p: + # Drop a contiguous block + block_size = int(n_channels * self.p * 2) + start = torch.randint(0, n_channels - block_size, (1,)).item() + mask = torch.ones(n_channels, device=neural.device) + mask[start:start + block_size] = 0 + if neural.dim() == 3: + neural = neural * mask.unsqueeze(0).unsqueeze(0) + else: + neural = neural * mask.unsqueeze(0) + return neural, target + + @property + def training(self) -> bool: + return True # Override in wrapper + + +class TemporalJitter: + """ + Randomly shift neural data in time. + + Helps with temporal alignment robustness. + + Args: + max_shift_bins: Maximum bins to shift (positive or negative) + wrap: If True, wrap around; else zero-pad + """ + + def __init__(self, max_shift_bins: int = 2, wrap: bool = False): + self.max_shift = max_shift_bins + self.wrap = wrap + + def __call__(self, neural: torch.Tensor, target: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + if self.max_shift == 0: + return neural, target + + shift = torch.randint(-self.max_shift, self.max_shift + 1, (1,)).item() + if shift == 0: + return neural, target + + if neural.dim() == 3: + # [batch, time, channels] + neural = torch.roll(neural, shifts=shift, dims=1) + target = torch.roll(target, shifts=shift, dims=0 if target.dim() == 1 else 1) + + if not self.wrap: + # Zero out wrapped values + if shift > 0: + neural[:, :shift, :] = 0 + else: + neural[:, shift:, :] = 0 + + return neural, target + + +class GaussianNoise: + """ + Add Gaussian noise to neural data. + + Args: + scale: Noise standard deviation (relative to data std) + """ + + def __init__(self, scale: float = 0.05): + self.scale = scale + + def __call__(self, neural: torch.Tensor, target: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + noise = torch.randn_like(neural) * self.scale + neural = neural + noise + return neural, target + + +class SpikeScaling: + """ + Randomly scale spike counts. + + Helps with gain variability across sessions. + + Args: + min_scale: Minimum scale factor + max_scale: Maximum scale factor + """ + + def __init__(self, min_scale: float = 0.9, max_scale: float = 1.1): + self.min_scale = min_scale + self.max_scale = max_scale + + def __call__(self, neural: torch.Tensor, target: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + scale = torch.empty(1).uniform_(self.min_scale, self.max_scale).item() + neural = neural * scale + return neural, target + + +class Mixup: + """ + Mixup augmentation for neural data. + + From Zhang et al. 2018 - interpolate between samples. + + Args: + alpha: Beta distribution parameter + """ + + def __init__(self, alpha: float = 0.2): + self.alpha = alpha + + def __call__(self, neural: torch.Tensor, target: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + # This needs batch-level mixing - implemented in dataloader + return neural, target + + +def build_augmentation(cfg: DictConfig) -> Optional[Compose]: + """ + Build augmentation pipeline from config. + + Args: + cfg: Augmentation configuration + + Returns: + Composed transforms or None if disabled + """ + if not cfg.get('enabled', False): + return None + + transforms = [] + + # Electrode dropout + if cfg.get('electrode_dropout', {}).get('enabled', False): + transforms.append(ElectrodeDropout( + p=cfg.electrode_dropout.get('p', 0.1), + structured=cfg.electrode_dropout.get('structured', False) + )) + + # Temporal jitter + if cfg.get('temporal_jitter', {}).get('enabled', False): + transforms.append(TemporalJitter( + max_shift_bins=cfg.temporal_jitter.get('max_shift_bins', 2), + wrap=cfg.temporal_jitter.get('wrap', False) + )) + + # Noise + if cfg.get('noise', {}).get('enabled', False): + if cfg.noise.get('type', 'gaussian') == 'gaussian': + transforms.append(GaussianNoise( + scale=cfg.noise.get('scale', 0.05) + )) + + # Spike scaling + if cfg.get('spike_scaling', {}).get('enabled', False): + transforms.append(SpikeScaling( + min_scale=cfg.spike_scaling.get('min_scale', 0.9), + max_scale=cfg.spike_scaling.get('max_scale', 1.1) + )) + + if not transforms: + return None + + return Compose(transforms) diff --git a/src/utils/logging.py b/src/utils/logging.py new file mode 100644 index 0000000..1f9b0b3 --- /dev/null +++ b/src/utils/logging.py @@ -0,0 +1,183 @@ +""" +Experiment Logging Utilities + +Integrates with Weights & Biases (WandB) for automatic experiment tracking. +Replaces the manual README tables with auto-generated dashboards. +""" + +import os +import json +from pathlib import Path +from typing import Dict, Any, Optional +from datetime import datetime + +from omegaconf import DictConfig, OmegaConf + +# Optional WandB import +try: + import wandb + WANDB_AVAILABLE = True +except ImportError: + WANDB_AVAILABLE = False + wandb = None + + +class ExperimentLogger: + """ + Unified experiment logger supporting WandB and local logging. + + Features: + - Automatic config serialization + - Git hash tracking for reproducibility + - Metric logging with step tracking + - Artifact management (models, configs) + """ + + def __init__( + self, + cfg: DictConfig, + use_wandb: bool = True, + project: str = "PhantomX", + entity: Optional[str] = None, + tags: Optional[list] = None, + notes: Optional[str] = None, + output_dir: Optional[Path] = None, + ): + self.cfg = cfg + self.use_wandb = use_wandb and WANDB_AVAILABLE + self.project = project + self.entity = entity + self.output_dir = Path(output_dir) if output_dir else Path("logs") + self.run = None + self.step = 0 + + # Prepare tags + self.tags = tags or [] + if hasattr(cfg, 'model') and hasattr(cfg.model, 'name'): + self.tags.append(cfg.model.name) + if hasattr(cfg, 'dataset') and hasattr(cfg.dataset, 'name'): + self.tags.append(cfg.dataset.name) + + self.notes = notes + if hasattr(cfg, 'notes'): + self.notes = cfg.notes + + def init(self) -> "ExperimentLogger": + """Initialize the logger.""" + if self.use_wandb: + self._init_wandb() + else: + self._init_local() + return self + + def _init_wandb(self): + """Initialize Weights & Biases.""" + config_dict = OmegaConf.to_container(self.cfg, resolve=True) + + self.run = wandb.init( + project=self.project, + entity=self.entity, + name=getattr(self.cfg, 'experiment_name', None), + config=config_dict, + tags=self.tags, + notes=self.notes, + dir=str(self.output_dir), + reinit=True, + ) + + # Log git hash if available + from .seeding import get_git_hash + git_hash = get_git_hash() + if git_hash: + wandb.config.update({"git_hash": git_hash}, allow_val_change=True) + + print(f"📊 WandB run initialized: {self.run.url}") + + def _init_local(self): + """Initialize local file logging.""" + self.output_dir.mkdir(parents=True, exist_ok=True) + self.log_file = self.output_dir / "metrics.jsonl" + + # Save config + config_file = self.output_dir / "config.yaml" + OmegaConf.save(self.cfg, config_file) + + print(f"📊 Local logging to: {self.output_dir}") + + def log(self, metrics: Dict[str, Any], step: Optional[int] = None): + """Log metrics.""" + if step is not None: + self.step = step + else: + self.step += 1 + + if self.use_wandb and self.run: + wandb.log(metrics, step=self.step) + else: + # Local logging + log_entry = {"step": self.step, **metrics, "timestamp": datetime.now().isoformat()} + with open(self.log_file, "a") as f: + f.write(json.dumps(log_entry) + "\n") + + def log_summary(self, summary: Dict[str, Any]): + """Log final summary metrics.""" + if self.use_wandb and self.run: + for key, value in summary.items(): + wandb.run.summary[key] = value + else: + summary_file = self.output_dir / "summary.json" + with open(summary_file, "w") as f: + json.dump(summary, f, indent=2) + + def save_model(self, model_path: Path, name: str = "model"): + """Save model as artifact.""" + if self.use_wandb and self.run: + artifact = wandb.Artifact(name, type="model") + artifact.add_file(str(model_path)) + self.run.log_artifact(artifact) + # Local: model already saved by trainer + + def finish(self): + """Finish logging.""" + if self.use_wandb and self.run: + wandb.finish() + print("📊 WandB run finished") + + +# Module-level convenience functions +_logger: Optional[ExperimentLogger] = None + + +def init_logger( + cfg: DictConfig, + use_wandb: bool = True, + **kwargs +) -> ExperimentLogger: + """Initialize the global experiment logger.""" + global _logger + + # Get settings from config if available + if hasattr(cfg, 'logging'): + use_wandb = cfg.logging.get('use_wandb', use_wandb) + kwargs.setdefault('project', cfg.logging.get('wandb_project', 'PhantomX')) + kwargs.setdefault('entity', cfg.logging.get('wandb_entity', None)) + + if hasattr(cfg, 'output_dir'): + kwargs.setdefault('output_dir', cfg.output_dir) + + _logger = ExperimentLogger(cfg, use_wandb=use_wandb, **kwargs) + return _logger.init() + + +def log_metrics(metrics: Dict[str, Any], step: Optional[int] = None): + """Log metrics to the global logger.""" + if _logger: + _logger.log(metrics, step) + + +def finish_logger(): + """Finish the global logger.""" + global _logger + if _logger: + _logger.finish() + _logger = None diff --git a/src/utils/metrics.py b/src/utils/metrics.py new file mode 100644 index 0000000..a706a4f --- /dev/null +++ b/src/utils/metrics.py @@ -0,0 +1,169 @@ +""" +Evaluation Metrics + +Standardized metrics for neural decoding evaluation. +Primary metric: R² (coefficient of determination) +""" + +from typing import Dict, Any, Optional, Union +import numpy as np +import torch +from sklearn.metrics import r2_score, mean_squared_error, mean_absolute_error + + +def compute_metrics( + predictions: Union[np.ndarray, torch.Tensor], + targets: Union[np.ndarray, torch.Tensor], + prefix: str = "", + per_dim: bool = True, +) -> Dict[str, float]: + """ + Compute standard evaluation metrics. + + Args: + predictions: Model predictions [N, D] or [N,] + targets: Ground truth [N, D] or [N,] + prefix: Prefix for metric names (e.g., "val/", "test/") + per_dim: Compute per-dimension metrics for multi-dim outputs + + Returns: + Dictionary of metrics + + Example: + >>> preds = model(neural_data) + >>> metrics = compute_metrics(preds, velocities, prefix="val/") + >>> print(metrics) + {'val/r2': 0.71, 'val/r2_x': 0.72, 'val/r2_y': 0.70, ...} + """ + # Convert to numpy + if isinstance(predictions, torch.Tensor): + predictions = predictions.detach().cpu().numpy() + if isinstance(targets, torch.Tensor): + targets = targets.detach().cpu().numpy() + + # Flatten if 1D + if predictions.ndim == 1: + predictions = predictions.reshape(-1, 1) + if targets.ndim == 1: + targets = targets.reshape(-1, 1) + + metrics = {} + + # Overall R² (macro average) + r2_overall = r2_score(targets, predictions, multioutput='uniform_average') + metrics[f"{prefix}r2"] = float(r2_overall) + + # Overall MSE and MAE + mse = mean_squared_error(targets, predictions) + mae = mean_absolute_error(targets, predictions) + metrics[f"{prefix}mse"] = float(mse) + metrics[f"{prefix}mae"] = float(mae) + metrics[f"{prefix}rmse"] = float(np.sqrt(mse)) + + # Per-dimension metrics (for velocity: x, y) + if per_dim and predictions.shape[1] > 1: + dim_names = ['x', 'y', 'z', 'w'][:predictions.shape[1]] + + for i, name in enumerate(dim_names): + r2_dim = r2_score(targets[:, i], predictions[:, i]) + mse_dim = mean_squared_error(targets[:, i], predictions[:, i]) + + metrics[f"{prefix}r2_{name}"] = float(r2_dim) + metrics[f"{prefix}mse_{name}"] = float(mse_dim) + + # Variance explained (alternative to R²) + var_explained = 1 - np.var(targets - predictions) / np.var(targets) + metrics[f"{prefix}var_explained"] = float(var_explained) + + return metrics + + +def compute_r2( + predictions: Union[np.ndarray, torch.Tensor], + targets: Union[np.ndarray, torch.Tensor], +) -> float: + """ + Simple R² computation (convenience function). + + Args: + predictions: Model predictions + targets: Ground truth + + Returns: + R² score (float) + """ + if isinstance(predictions, torch.Tensor): + predictions = predictions.detach().cpu().numpy() + if isinstance(targets, torch.Tensor): + targets = targets.detach().cpu().numpy() + + return float(r2_score(targets.flatten(), predictions.flatten())) + + +def compute_bits_per_spike( + predictions: Union[np.ndarray, torch.Tensor], + targets: Union[np.ndarray, torch.Tensor], + spike_counts: Union[np.ndarray, torch.Tensor], +) -> float: + """ + Compute bits per spike metric (NLB standard). + + Args: + predictions: Predicted spike rates (log-rates) + targets: Actual spike counts + spike_counts: Total spike counts for normalization + + Returns: + Bits per spike + """ + # This is a placeholder - implement based on NLB evaluation code + raise NotImplementedError("Bits per spike not yet implemented") + + +class MetricTracker: + """ + Track metrics over training epochs. + + Useful for early stopping and best model selection. + """ + + def __init__(self, metric_name: str = "r2", mode: str = "max"): + """ + Args: + metric_name: Name of metric to track + mode: "max" for metrics like R², "min" for loss + """ + self.metric_name = metric_name + self.mode = mode + self.best_value = -float('inf') if mode == "max" else float('inf') + self.best_epoch = 0 + self.history = [] + + def update(self, value: float, epoch: int) -> bool: + """ + Update tracker with new value. + + Returns: + True if this is a new best value + """ + self.history.append({"epoch": epoch, "value": value}) + + is_best = False + if self.mode == "max": + if value > self.best_value: + self.best_value = value + self.best_epoch = epoch + is_best = True + else: + if value < self.best_value: + self.best_value = value + self.best_epoch = epoch + is_best = True + + return is_best + + def should_stop(self, patience: int) -> bool: + """Check if training should stop based on patience.""" + if len(self.history) < patience: + return False + return self.history[-1]["epoch"] - self.best_epoch >= patience diff --git a/src/utils/seeding.py b/src/utils/seeding.py new file mode 100644 index 0000000..15a9978 --- /dev/null +++ b/src/utils/seeding.py @@ -0,0 +1,152 @@ +""" +Seeding and Reproducibility Utilities + +Ensures experiments are reproducible across runs. +Critical for statistical validation (Exp 23). +""" + +import os +import random +import subprocess +from typing import Optional + +import numpy as np +import torch + + +def seed_everything(seed: int = 42, deterministic: bool = True) -> int: + """ + Seed all random number generators for reproducibility. + + Args: + seed: Random seed + deterministic: If True, use deterministic algorithms (slower but reproducible) + + Returns: + The seed used + + Example: + >>> seed_everything(42) + >>> # All subsequent random operations are reproducible + """ + # Python + random.seed(seed) + os.environ['PYTHONHASHSEED'] = str(seed) + + # NumPy + np.random.seed(seed) + + # PyTorch + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + torch.cuda.manual_seed_all(seed) # For multi-GPU + + if deterministic: + # Deterministic algorithms (may be slower) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + + # PyTorch 1.8+ + if hasattr(torch, 'use_deterministic_algorithms'): + try: + torch.use_deterministic_algorithms(True) + except RuntimeError: + # Some operations don't have deterministic implementations + pass + else: + # Faster but non-deterministic + torch.backends.cudnn.deterministic = False + torch.backends.cudnn.benchmark = True + + print(f"🌱 Seeded everything with seed={seed}, deterministic={deterministic}") + return seed + + +def get_git_hash() -> Optional[str]: + """ + Get the current git commit hash for reproducibility tracking. + + Returns: + Short git hash or None if not in a git repo + """ + try: + result = subprocess.run( + ['git', 'rev-parse', '--short', 'HEAD'], + capture_output=True, + text=True, + timeout=5 + ) + if result.returncode == 0: + return result.stdout.strip() + except (subprocess.SubprocessError, FileNotFoundError): + pass + return None + + +def get_git_diff() -> Optional[str]: + """ + Get uncommitted changes for debugging. + + Returns: + Git diff output or None + """ + try: + result = subprocess.run( + ['git', 'diff', '--stat'], + capture_output=True, + text=True, + timeout=10 + ) + if result.returncode == 0 and result.stdout.strip(): + return result.stdout.strip() + except (subprocess.SubprocessError, FileNotFoundError): + pass + return None + + +def is_deterministic() -> bool: + """Check if PyTorch is in deterministic mode.""" + return ( + torch.backends.cudnn.deterministic and + not torch.backends.cudnn.benchmark + ) + + +class SeedContext: + """ + Context manager for temporary seeding. + + Useful for reproducible data augmentation while keeping + training stochastic. + + Example: + >>> with SeedContext(42): + ... # Reproducible operations + ... data = augment(data) + >>> # Back to previous random state + """ + + def __init__(self, seed: int): + self.seed = seed + self.numpy_state = None + self.torch_state = None + self.random_state = None + + def __enter__(self): + # Save current states + self.numpy_state = np.random.get_state() + self.torch_state = torch.get_rng_state() + self.random_state = random.getstate() + + # Set temporary seed + np.random.seed(self.seed) + torch.manual_seed(self.seed) + random.seed(self.seed) + + return self + + def __exit__(self, *args): + # Restore previous states + np.random.set_state(self.numpy_state) + torch.set_rng_state(self.torch_state) + random.setstate(self.random_state) diff --git a/train.py b/train.py new file mode 100644 index 0000000..5e72726 --- /dev/null +++ b/train.py @@ -0,0 +1,167 @@ +#!/usr/bin/env python +""" +PhantomX Unified Training Entry Point + +This is the SINGLE entry point for all experiments. +Replaces the 25+ exp*.py scripts with configuration-driven training. + +Usage: + # Default training (VQ-VAE on MC_Maze) + python train.py + + # Specific experiment + python train.py experiment=exp25_mamba + + # Override model/dataset + python train.py model=mamba dataset=mc_rtt + + # Override hyperparameters + python train.py model=lstm model.hidden_dim=512 trainer.learning_rate=1e-4 + + # Multi-seed validation + python train.py experiment=exp23_validation --multirun seed=42,123,456 + + # Sweep + python train.py --multirun model=vqvae,lstm,mamba dataset=mc_maze + +Examples: + # Replicate Exp 25: Mamba on MC_RTT + python train.py experiment=exp25_mamba + + # Replicate Exp 22c: Multi-seed teacher + python train.py experiment=exp22c_teacher + + # Quick test + python train.py trainer.max_epochs=5 logging.use_wandb=false +""" + +import sys +from pathlib import Path + +# Add src to path +sys.path.insert(0, str(Path(__file__).parent)) + +import hydra +from omegaconf import DictConfig, OmegaConf + +# Import after hydra to avoid issues +import torch + + +@hydra.main(version_base=None, config_path="configs", config_name="config") +def main(cfg: DictConfig) -> float: + """ + Main training function. + + Args: + cfg: Hydra configuration + + Returns: + Best validation R² (for Hydra sweeps) + """ + # Import here to avoid circular imports + from src.utils.seeding import seed_everything, get_git_hash + from src.utils.logging import init_logger, finish_logger + from src.models import build_model + from src.datamodules import build_datamodule + from src.trainer import Trainer, ProgressiveTrainer + + # Print config + print("=" * 60) + print("PhantomX Training") + print("=" * 60) + print(f"Experiment: {cfg.get('experiment_name', 'default')}") + print(f"Model: {cfg.model.name}") + print(f"Dataset: {cfg.dataset.name}") + print(f"Seed: {cfg.seed}") + print("=" * 60) + + # Seed everything for reproducibility + seed_everything( + seed=cfg.seed, + deterministic=cfg.reproducibility.get('deterministic', True) + ) + + # Log git hash + git_hash = get_git_hash() + if git_hash: + print(f"Git commit: {git_hash}") + + # Initialize logger + logger = init_logger(cfg) + + try: + # Build datamodule + print("\n📊 Loading data...") + datamodule = build_datamodule(cfg) + datamodule.setup() + + # Get window size from model config + window_size = cfg.model.get('window_size', 10) + + # Update datamodule window size if needed + if hasattr(datamodule, 'window_size') and datamodule.window_size != window_size: + print(f" Updating window size: {datamodule.window_size} → {window_size}") + datamodule.window_size = window_size + datamodule.setup() # Recreate datasets + + # Build model + print("\n🔧 Building model...") + model = build_model(cfg, n_channels=datamodule.n_channels) + + # Count parameters + n_params = sum(p.numel() for p in model.parameters() if p.requires_grad) + print(f" Parameters: {n_params:,}") + + # Create trainer + print("\n🏋️ Setting up trainer...") + trainer_name = cfg.trainer.get('name', 'default') + + if trainer_name == 'progressive' and hasattr(model, 'use_vq'): + trainer = ProgressiveTrainer( + model=model, + train_loader=datamodule.train_dataloader(), + val_loader=datamodule.val_dataloader(), + test_loader=datamodule.test_dataloader(), + cfg=cfg, + ) + else: + trainer = Trainer( + model=model, + train_loader=datamodule.train_dataloader(), + val_loader=datamodule.val_dataloader(), + test_loader=datamodule.test_dataloader(), + cfg=cfg, + ) + + # Train + results = trainer.train() + + # Log summary + logger.log_summary({ + 'best_r2': results['best_r2'], + 'best_epoch': results['best_epoch'], + 'training_time_s': results['training_time_s'], + 'n_parameters': n_params, + 'git_hash': git_hash, + }) + + # Save final model + model_path = Path(cfg.output_dir) / 'final_model.pt' + torch.save({ + 'model_state_dict': model.state_dict(), + 'config': OmegaConf.to_container(cfg), + 'results': results, + }, model_path) + logger.save_model(model_path) + + print(f"\n📁 Results saved to: {cfg.output_dir}") + + return results['best_r2'] + + finally: + finish_logger() + + +if __name__ == "__main__": + main() From 2f7abb50b062dbb27865257dd0df8bcbe7a114b0 Mon Sep 17 00:00:00 2001 From: Youssef Date: Fri, 23 Jan 2026 10:48:43 -0500 Subject: [PATCH 2/2] update log, introduce PhantomMusic --- RESEARCH_LOG.md | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/RESEARCH_LOG.md b/RESEARCH_LOG.md index 602cbf9..3d602eb 100644 --- a/RESEARCH_LOG.md +++ b/RESEARCH_LOG.md @@ -2443,3 +2443,22 @@ split = 70% train / 15% val / 15% test (sequential) --- +## 🎰 Architecture Roulette — Time to Build Intuition +**Date**: 2026-01-23 + +We've been spinning the wheel: LSTM → Transformer → Mamba → VQ-VAE → FSQ → distillation... Each architecture wins on one dataset, fails on another. We're pattern-matching configurations without understanding **why** they work. + +**The problem**: We lack intuition for neural signal structure. + +**The plan**: Step back from BCI decoding and build fundamental intuition on a simpler domain — **music**. Audio has: +- Clear temporal structure (rhythm, melody) +- Multi-scale patterns (beats → bars → phrases) +- Well-understood representations (spectrograms, MFCCs) +- Easy human evaluation (does it sound right?) + +If we can develop intuition for temporal tokenization on music, we can transfer those insights back to neural data. + +**New repo**: https://github.com/yelabb/PhantomMusic + +--- +