code/training/config.py

"""
Hyperparameters, feature definitions, and normalization config.

All magic numbers live here so experiments are easy to re-run with
different settings.
"""

from dataclasses import dataclass, field

# ── Feature definitions ──────────────────────────────────────────────────────
# These MUST match the keys produced by batch_generate.py (circuit JSON).

# Input features for Model B (HH + ACh).  Model A drops "ach_level".
INPUT_FEATURES_B = [
    "n_exc",
    "n_inh",
    "conn_prob",
    "n_synapses",
    "mean_in_degree",
    "gS_exc_effective",
    "ou_mu_effective",
    "ou_sigma_effective",
    "ou_tau",
    "sim_duration_ms",
    "ach_level",            # <-- Model B only
]

INPUT_FEATURES_A = [f for f in INPUT_FEATURES_B if f != "ach_level"]

# The 11 output statistics (targets).  Order matters — keep consistent.
OUTPUT_STATS = [
    "mean_firing_rate",
    "mean_exc_rate",
    "mean_inh_rate",
    "mean_cv_isi",
    "mean_fano_factor",
    "synchrony_index",
    "mean_pairwise_corr",
    "peak_frequency_hz",
    "total_spectral_power",
    "n_active_neurons",
    "total_spikes",
]

# Statistics that should be log-transformed before normalization
# (they span orders of magnitude or are strictly positive counts).
LOG_TRANSFORM_STATS = {
    "total_spectral_power",
    "total_spikes",
    "n_active_neurons",
}

# Input features that should be log-transformed
LOG_TRANSFORM_INPUTS = {
    "n_synapses",
    "gS_exc_effective",
}


# ── Hyperparameters ──────────────────────────────────────────────────────────

@dataclass
class TrainConfig:
    """Training hyperparameters — small model, fast iteration."""

    # --- Model architecture ---
    d_model: int = 64           # Embedding dimension
    n_heads: int = 4            # Attention heads
    n_layers: int = 4           # Transformer encoder layers
    d_ff: int = 256             # Feed-forward hidden dim
    dropout: float = 0.1        # Dropout rate

    # --- Training ---
    batch_size: int = 2048      # Large batch — GPU is underutilized at 256
    lr: float = 1e-3            # Scaled up with batch size (linear scaling rule)
    weight_decay: float = 1e-2  # AdamW weight decay
    warmup_steps: int = 500     # Linear warmup steps
    max_epochs: int = 200       # Maximum training epochs
    patience: int = 20          # Early stopping patience (epochs)
    grad_clip: float = 1.0      # Gradient clipping norm

    # --- Data ---
    val_frac: float = 0.1       # Fraction of circuits for validation
    seed: int = 42              # Random seed for splits + init

    # --- Paths (on Modal volume) ---
    sim_dir: str = "/data/sims/prod_5k_v13/circuits"
    extra_ach0_dir: str = "/data/sims/ach0_extra/circuits"  # Additional ACh=0 data
    allen_dir: str = "/data/allen/epochs"
    checkpoint_dir: str = "/data/training/checkpoints"
    log_dir: str = "/data/training/logs"

    # --- Data augmentation ---
    aug_noise_scale: float = 0.02     # Gaussian noise std on normalized inputs
    aug_mixup_alpha: float = 0.2      # Beta distribution param for mixup

    # --- MLP architecture ---
    mlp_hidden: list = field(default_factory=lambda: [64, 64])
    mlp_dropout: float = 0.1

    # --- Derived ---
    n_input_features_a: int = field(init=False)
    n_input_features_b: int = field(init=False)
    n_output_stats: int = field(init=False)

    def __post_init__(self):
        self.n_input_features_a = len(INPUT_FEATURES_A)  # 10
        self.n_input_features_b = len(INPUT_FEATURES_B)  # 11
        self.n_output_stats = len(OUTPUT_STATS)           # 11