refactor: Move to nested config structure for organization

Author: Linked-Liszt

configs/trainer.yaml

@@ -1,127 +1,124 @@
 # AlphaDiffract trainer configuration — ConvNeXt (paper-matching lightweight variant)
 # Use with: PYTHONPATH=src python -m trainer.train_paper configs/trainer_convnext_paper.yaml
 
-# --- Data / Manifests ---
-manifest_dir: "../../../ad_data/manifests"
-dataset_root: "../../../ad_data/data/dataset"
-extra_val_file: "rruff.jsonl"
-auto_generate_manifests: true
-train_ratio: 0.8
-val_ratio: 0.1
-test_ratio: 0.1
-seed: 42
+data:
+  manifest_dir: "../../../ad_data/manifests"
+  dataset_root: "../../../ad_data/data/dataset"
+  extra_val_file: "rruff.jsonl"
+  auto_generate_manifests: true
+  train_ratio: 0.8
+  val_ratio: 0.1
+  test_ratio: 0.1
+  seed: 42
 
-# --- DataLoader ---
-batch_size: 64                  # match OG run (64 per process)
-num_workers: 8
-pin_memory: true
-persistent_workers: true
+  loader:
+    # --- DataLoader ---
+    batch_size: 64                  # match OG run (64 per process)
+    num_workers: 8
+    pin_memory: true
+    persistent_workers: true
+    prefetch_factor: 2
+    train_file: "train.jsonl"
+    val_file: "val.jsonl"
+    test_file: "test.jsonl"
 
-# --- Dataset label extraction (embedded in .npy/.npz) ---
-validate_paths: false
-extract_labels: true
-allow_pickle: true
-labels_key_map:
-  x: "dp"
-  cs: "cs"
-  sg: "sg"
-  lattice_params: null
-  lp_a: "_cell_length_a"
-  lp_b: "_cell_length_b"
-  lp_c: "_cell_length_c"
-  lp_alpha: "_cell_angle_alpha"
-  lp_beta: "_cell_angle_beta"
-  lp_gamma: "_cell_angle_gamma"
-dtype: "float32"
-mmap_mode: null
-floor_at_zero: true
-normalize_log1p: False          # paper used log1p preprocessing
-model_type: "multiscale"
+  preprocessing:
+    validate_paths: false
+    extract_labels: true
+    allow_pickle: true
+    labels_key_map:
+      x: "dp"
+      cs: "cs"
+      sg: "sg"
+      lattice_params: null
+      lp_a: "_cell_length_a"
+      lp_b: "_cell_length_b"
+      lp_c: "_cell_length_c"
+      lp_alpha: "_cell_angle_alpha"
+      lp_beta: "_cell_angle_beta"
+      lp_gamma: "_cell_angle_gamma"
+    dtype: "float32"
+    mmap_mode: null
+    floor_at_zero: true
+    normalize_log1p: False          # paper used log1p preprocessing
 
-# --- ConvNeXt (OG-equivalent configuration) ---
-# 3 stages; one block per stage; large kernels; stride-5 downsampling
-# Matches OG multiscale_cnn_cls_regr_convnextBlock_LeakyReLU.json exactly
-depths: [1, 1, 1]
-dims: [80, 80, 80]
-kernel_sizes: [100, 50, 25]
-strides: [5, 5, 5]
-dropout_rate: 0.3
-# OG uses layer_scale_init_value=0 (disabled)
-layer_scale_init_value: 0.0
-# OG uses constant drop_path_rate=0.3 (not ramped)
-drop_path_rate: 0.3
-ramped_dropout_rate: false
-block_type: "convnext"
-pooling_type: "average"
-final_pool: true
-use_batchnorm: false
-output_type: "flatten"
+  augmentation:
+    noise_poisson_range: [1.0, 100.0]
+    noise_gaussian_range: [0.001, 0.1]
+    standardize_to: [0.0, 100.0]
 
-# Heads
-head_dropout: 0.5
-cs_hidden: [2300, 1150]
-sg_hidden: [2300, 1150]
-lp_hidden: [512, 256]
+model:
+  type: "multiscale"
+  
+  backbone:
+    dim_in: 8192
+    dims: [80, 80, 80]
+    kernel_sizes: [100, 50, 25]
+    strides: [5, 5, 5]
+    dropout_rate: 0.3
+    layer_scale_init_value: 0.0
+    drop_path_rate: 0.3
+    ramped_dropout_rate: false
+    block_type: "convnext"
+    pooling_type: "average"
+    final_pool: true
+    use_batchnorm: false
+    activation: "leaky_relu"
+    output_type: "flatten"
 
-# Task sizes
-num_cs_classes: 7
-num_sg_classes: 230
-num_lp_outputs: 6
+  heads:
+    head_dropout: 0.5
+    cs_hidden: [2300, 1150]
+    sg_hidden: [2300, 1150]
+    lp_hidden: [512, 256]
 
-# LP output bounds
-lp_bounds_min: [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
-lp_bounds_max: [300.0, 300.0, 300.0, 180.0, 180.0, 180.0]
-bound_lp_with_sigmoid: true
+  tasks:
+    num_cs_classes: 7
+    num_sg_classes: 230
+    num_lp_outputs: 6
 
-# Loss weights
-lambda_cs: 1.0
-lambda_sg: 1.0
-lambda_lp: 1.0
+    lp_bounds_min: [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
+    lp_bounds_max: [300.0, 300.0, 300.0, 180.0, 180.0, 180.0]
+    bound_lp_with_sigmoid: true
 
-# Optional GEMD term on SG
-gemd_mu: 0.0
-gemd_distance_matrix_path: null
+  loss:
+    lambda_cs: 1.0
+    lambda_sg: 1.0
+    lambda_lp: 1.0
 
-# Optimizer (paper): AdamW, lr=2e-4, wd=0.01
-lr: 0.0002
-weight_decay: 0.01
-use_adamw: true
-gradient_clip_val: 1.0
-gradient_clip_algorithm: "norm"
+    gemd_mu: 0.0
+    gemd_distance_matrix_path: null
 
-# --- Noise augmentation (training split only; matches paper) ---
-# If provided, noise is applied dynamically per-sample in the DataModule using the same
-# sequencing as the paper: Poisson -> normalize -> add Gaussian -> renormalize -> rescale.
-# Set ranges to None to disable.
-noise_poisson_range: [1.0, 100.0]
-noise_gaussian_range: [0.001, 0.1]
+optimizer:
+  lr: 0.0002
+  weight_decay: 0.01
+  use_adamw: true
+  gradient_clip_val: 1.0
+  gradient_clip_algorithm: "norm"
 
-# Standardize after noise to match OG CLI (--standardize-to 0 100)
-standardize_to: [0.0, 100.0]
-# --- Logging ---
-logger: "mlflow"
-csv_logger_name: "model_logs_convnext_paper"
-mlflow_experiment_name: "AlphaDiffract_Paper_ConvNeXt"
-mlflow_tracking_uri: null
-mlflow_run_name: "ConvNeXt_Paper_Run"
+trainer:
+  default_root_dir: "outputs/convnext_paper"
+  max_epochs: 100
+  accumulate_grad_batches: 1
+  precision: "32"                # match OG (AMP disabled)
+  accelerator: "gpu"
+  devices: 1
+  log_every_n_steps: 200
+  deterministic: false
+  benchmark: true
 
-# --- Trainer settings ---
-default_root_dir: "outputs/convnext_paper"
-max_epochs: 100
-accumulate_grad_batches: 1
-precision: "32"                # match OG (AMP disabled)
-accelerator: "gpu"
-devices: 1
-log_every_n_steps: 200
-deterministic: false
-benchmark: true
+logging:
+  logger: "mlflow"
+  csv_logger_name: "model_logs_convnext_paper"
+  mlflow_experiment_name: "AlphaDiffract_Paper_ConvNeXt"
+  mlflow_tracking_uri: null
+  mlflow_run_name: "ConvNeXt_Paper_Run"
 
-# --- Checkpointing ---
-monitor: "val/loss"
-mode: "min"
-save_top_k: 1
-every_n_epochs: 1
-
-# --- Evaluation ---
-resume_from:
-test_after_train: true
+checkpointing:
+  monitor: "val/loss"
+  mode: "min"
+  save_top_k: 1
+  every_n_epochs: 1
+  
+  resume_from: null
+  test_after_train: true

src/trainer/dataset/datamodule.py

@@ -130,12 +130,15 @@ class NpyDataModule(pl.LightningDataModule):
 
     def __init__(
         self,
-        manifest_dir: str = "data/manifests",
-        batch_size: int = 32,
-        num_workers: int = 4,
-        pin_memory: bool = True,
-        persistent_workers: bool = True,
-        prefetch_factor: Optional[int] = None,
+        manifest_dir: str,
+        batch_size: int,
+        num_workers: int,
+        pin_memory: bool,
+        persistent_workers: bool,
+        prefetch_factor: Optional[int],
+        train_file: str,
+        val_file: str,
+        test_file: str,
         collate_fn: Optional[Callable] = None,
         dataset_cls: type = NpyManifestDataset,
         dataset_kwargs: Optional[Dict[str, Any]] = None,
@@ -146,10 +149,6 @@ def __init__(
         val_ratio: float = 0.1,
         test_ratio: float = 0.1,
         seed: int = 42,
-        # Custom manifest filenames (within manifest_dir)
-        train_file: str = "train.jsonl",
-        val_file: str = "val.jsonl",
-        test_file: str = "test.jsonl",
         # Optional: add a second validation manifest file (e.g., "rruff.jsonl")
         extra_val_file: Optional[str] = None,
         # Optional noise augmentation for training split only

src/trainer/dataset/dataset.py

@@ -149,16 +149,16 @@ class NpyManifestDataset(Dataset):
     def __init__(
         self,
         manifest_path: str,
+        dtype: torch.dtype,
+        mmap_mode: Optional[str],
+        return_meta: bool,
+        validate_paths: bool,
+        extract_labels: bool,
+        allow_pickle: bool,
+        floor_at_zero: bool,
+        normalize_log1p: bool,
         transform: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
-        dtype: torch.dtype = torch.float32,
-        mmap_mode: Optional[str] = "r",
-        return_meta: bool = True,
-        validate_paths: bool = False,
-        extract_labels: bool = False,
         labels_key_map: Optional[Dict[str, List[str]]] = None,
-        allow_pickle: bool = True,
-        floor_at_zero: bool = True,
-        normalize_log1p: bool = False,
     ) -> None:
         super().__init__()
         self.manifest_path = manifest_path