Refactor: Continued cleanup

Author: Linked-Liszt

configs/trainer.yaml

@@ -1,6 +1,3 @@
-# AlphaDiffract trainer configuration — ConvNeXt (paper-matching lightweight variant)
-# Use with: PYTHONPATH=src python -m trainer.train_paper configs/trainer_convnext_paper.yaml
-
 data:
   manifest_dir: "../../../ad_data/manifests"
   dataset_root: "../../../ad_data/data/dataset"
@@ -40,7 +37,8 @@ data:
     dtype: "float32"
     mmap_mode: null
     floor_at_zero: true
-    normalize_log1p: False          # paper used log1p preprocessing
+    normalize_log1p: False
+    shift_labels: true
 
   augmentation:
     noise_poisson_range: [1.0, 100.0]

src/trainer/dataset/datamodule.py

@@ -4,28 +4,6 @@
 This module wires NpyManifestDataset to PyTorch Lightning and can optionally
 auto-generate manifests from a dataset root if they are missing.
 
-Typical usage:
-    from pytorch_lightning import Trainer
-    from dataset.datamodule import NpyDataModule
-
-    dm = NpyDataModule(
-        manifest_dir="data/manifests",
-        batch_size=64,
-        num_workers=8,
-        pin_memory=True,
-        persistent_workers=True,
-        # Optional: auto-generate manifests if missing
-        dataset_root="data/dataset",
-        auto_generate_manifests=True,
-        train_ratio=0.8, val_ratio=0.1, test_ratio=0.1, seed=42,
-        # Dataset-specific kwargs
-        dataset_kwargs={"dtype": torch.float32, "mmap_mode": "r", "return_meta": True},
-    )
-
-    trainer = Trainer(max_epochs=10, accelerator="auto", devices="auto")
-    trainer.fit(model, dm)
-    trainer.test(model, dm)
-
 Notes:
 - Splitting is performed by material ID when generating manifests (never per-file).
 - Manifests avoid scanning the entire dataset during training.
@@ -80,34 +58,6 @@ def _transform(x: torch.Tensor) -> torch.Tensor:
     return _transform
 
 
-
-def _shift_one_based_collate(batch):
-    """
-    Collate function that uses PyTorch's default_collate, then unconditionally shifts
-    cs and sg labels by -1 (assumes 1-based input). Performed under torch.no_grad to avoid
-    constructing any graphs.
-    """
-    collated = default_collate(batch)
-    with torch.no_grad():
-        def _shift(t):
-            return t - 1 if torch.is_tensor(t) else t
-
-        if isinstance(collated, dict):
-            if "cs" in collated:
-                collated["cs"] = _shift(collated["cs"])
-            if "sg" in collated:
-                collated["sg"] = _shift(collated["sg"])
-        elif isinstance(collated, (list, tuple)):
-            # Tuple-based batches: (x, cs, sg, [lp])
-            lst = list(collated)
-            if len(lst) >= 2:
-                lst[1] = _shift(lst[1])
-            if len(lst) >= 3:
-                lst[2] = _shift(lst[2])
-            collated = type(collated)(lst)
-    return collated
-
-
 class NpyDataModule(pl.LightningDataModule):
     """
     LightningDataModule that reads train/val/test JSONL manifests and constructs DataLoaders.
@@ -165,9 +115,7 @@ def __init__(
         self.persistent_workers = persistent_workers
         self.prefetch_factor = prefetch_factor
         self.collate_fn = collate_fn
-        if self.collate_fn is None:
-            self.collate_fn = _shift_one_based_collate
-
+        
         self.dataset_cls = dataset_cls
         self.dataset_kwargs = dataset_kwargs or {}
 

src/trainer/dataset/dataset.py

@@ -157,6 +157,7 @@ def __init__(
         allow_pickle: bool,
         floor_at_zero: bool,
         normalize_log1p: bool,
+        shift_labels: bool,
         transform: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
         labels_key_map: Optional[Dict[str, List[str]]] = None,
     ) -> None:
@@ -170,6 +171,7 @@ def __init__(
         self.allow_pickle = allow_pickle
         self.floor_at_zero = floor_at_zero
         self.normalize_log1p = normalize_log1p
+        self.shift_labels = shift_labels
         # Default key mapping for extracting fields from embedded containers
         # Simplified: single string keys, no search lists
         self.labels_key_map = labels_key_map or {
@@ -311,8 +313,12 @@ def _get_exact(container, key: str):
         # Attach labels if present/extracted
         if self.extract_labels:
             if y_cs_t is not None:
+                if self.shift_labels:
+                    y_cs_t = y_cs_t - 1
                 sample["cs"] = y_cs_t
             if y_sg_t is not None:
+                if self.shift_labels:
+                    y_sg_t = y_sg_t - 1
                 sample["sg"] = y_sg_t
             if y_lp_t is not None:
                 sample["lattice_params"] = y_lp_t

src/trainer/model/model.py

@@ -3,7 +3,6 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import numpy as np
 
 import pytorch_lightning as pl
 
@@ -183,7 +182,6 @@ def make_mlp(
 
 
 # -----------------------------
-# OG-style Multiscale CNN Backbone (1D) with ConvNeXt-like blocks
 # Mirrors alphadiffract.model.MultiscaleCNNBackbone behavior:
 # - sequential conv stages with specified kernel_sizes and strides
 # - optional average/max pooling between stages and at the end
@@ -212,13 +210,11 @@ def __init__(
         self.dim_in = dim_in
         self.output_type = output_type
 
-        # Build per-stage dropout schedule
         if ramped_dropout_rate:
             dropout_per_stage = torch.linspace(0.0, dropout_rate, steps=len(channels)).tolist()
         else:
             dropout_per_stage = [dropout_rate] * len(channels)
 
-        # Select pooling module
         if pooling_type == "average":
             pool_cls = nn.AvgPool1d
             pool_kwargs = {"kernel_size": 3, "stride": 2}
@@ -231,7 +227,6 @@ def __init__(
         layers: List[nn.Module] = []
         in_ch = 1
         for i, (out_ch, k, s) in enumerate(zip(channels, kernel_sizes, strides)):
-            # Build stage block matching OG ConvNextBlock1DAdaptorForMultiscaleCNN
             stage_block = ConvNextBlock1DAdaptor(
                 in_channels=in_ch,
                 out_channels=out_ch,
@@ -284,7 +279,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 class AlphaDiffractMultiscaleLightning(pl.LightningModule):
     def __init__(
         self,
-        # Backbone params (OG-style)
         dim_in: int,
         channels: Tuple[int, ...],
         kernel_sizes: Tuple[int, ...],
@@ -300,36 +294,31 @@ def __init__(
         layer_scale_init_value: float,
         drop_path_rate: float,
 
-        # Heads
         head_dropout: float,
         cs_hidden: Optional[Tuple[int, ...]],
         sg_hidden: Optional[Tuple[int, ...]],
         lp_hidden: Optional[Tuple[int, ...]],
 
-        # Task sizes
         num_cs_classes: int,
         num_sg_classes: int,
         num_lp_outputs: int,
 
-        # LP bounding
         lp_bounds_min: Tuple[float, float, float, float, float, float],
         lp_bounds_max: Tuple[float, float, float, float, float, float],
         bound_lp_with_sigmoid: bool,
 
-        # Loss weights
         lambda_cs: float,
         lambda_sg: float,
         lambda_lp: float,
 
-        # Optimizer
         lr: float,
         weight_decay: float,
         use_adamw: bool,
+        steps_per_epoch: int,
     ):
         super().__init__()
         self.save_hyperparameters()
 
-        # Backbone
         self.backbone = MultiscaleCNNBackbone1D(
             dim_in=dim_in,
             channels=channels,
@@ -348,7 +337,6 @@ def __init__(
         )
         feat_dim = self.backbone.dim_output
 
-        # Heads
         self.cs_head = make_mlp(
             input_dim=feat_dim,
             hidden_dims=cs_hidden,
@@ -371,22 +359,20 @@ def __init__(
             output_activation=None,
         )
 
-        # Losses and bounds
         self.ce = nn.CrossEntropyLoss()
         self.mse = nn.MSELoss()
         self.register_buffer("lp_min", torch.tensor(lp_bounds_min, dtype=torch.float32))
         self.register_buffer("lp_max", torch.tensor(lp_bounds_max, dtype=torch.float32))
         self.bound_lp_with_sigmoid = bound_lp_with_sigmoid
 
-        # weights and optim config
         self.lambda_cs = lambda_cs
         self.lambda_sg = lambda_sg
         self.lambda_lp = lambda_lp
         self.lr = lr
         self.weight_decay = weight_decay
         self.use_adamw = use_adamw
+        self.steps_per_epoch = steps_per_epoch
 
-        # Task sizes
         self.num_cs_classes = num_cs_classes
         self.num_sg_classes = num_sg_classes
         self.num_lp_outputs = num_lp_outputs
@@ -523,23 +509,7 @@ def configure_optimizers(self):
         else:
             optimizer = torch.optim.Adam(params, lr=self.lr, weight_decay=self.weight_decay)
 
-        # Compute steps per epoch to match OG scheduler semantics:
-        # step_size_up = 6 * iterations_per_epoch
-        steps_per_epoch = None
-        try:
-            if hasattr(self, "trainer") and self.trainer is not None:
-                total_steps = getattr(self.trainer, "estimated_stepping_batches", None)
-                max_epochs = getattr(self.trainer, "max_epochs", None)
-                if total_steps is not None and max_epochs is not None and max_epochs > 0:
-                    steps_per_epoch = max(1, total_steps // max_epochs)
-        except Exception:
-            pass
-
-        if steps_per_epoch is None:
-            # Fallback if trainer hooks are unavailable; use a conservative default
-            steps_per_epoch = 100
-
-        step_size_up = int(6 * steps_per_epoch)
+        step_size_up = int(6 * self.steps_per_epoch)
 
         scheduler = torch.optim.lr_scheduler.CyclicLR(
             optimizer,

src/trainer/train.py

@@ -6,21 +6,18 @@
 import yaml
 from pytorch_lightning import Trainer, seed_everything
 from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor, Callback
-import signal
 from pytorch_lightning.loggers import CSVLogger
 try:
     from pytorch_lightning.loggers import MLFlowLogger
 except Exception:
     MLFlowLogger = None
 
-# Project imports (expect PYTHONPATH=src or run via `python -m trainer.train_paper`)
 from dataset import NpyDataModule
 from model.model import AlphaDiffractMultiscaleLightning
 
 
 def parse_args() -> argparse.Namespace:
     p = argparse.ArgumentParser(description="Train AlphaDiffract paper model (config-required)")
-    # Require a config file path with no script-side defaults
     p.add_argument("config", type=str, help="Path to trainer config YAML (e.g., configs/trainer.yaml)")
     return p.parse_args()
 
@@ -63,6 +60,7 @@ def build_datamodule_from_cfg(cfg: Dict[str, Any]) -> NpyDataModule:
         "allow_pickle": prep_cfg["allow_pickle"],
         "floor_at_zero": prep_cfg["floor_at_zero"],
         "normalize_log1p": prep_cfg["normalize_log1p"],
+        "shift_labels": prep_cfg["shift_labels"],
     }
     labels_key_map = prep_cfg["labels_key_map"]
     if labels_key_map is not None:
@@ -95,7 +93,7 @@ def build_datamodule_from_cfg(cfg: Dict[str, Any]) -> NpyDataModule:
     return dm
 
 
-def build_model_from_cfg(cfg: Dict[str, Any]):
+def build_model_from_cfg(cfg: Dict[str, Any], steps_per_epoch: int):
     model_cfg = cfg["model"]
     backbone_cfg = model_cfg["backbone"]
     heads_cfg = model_cfg["heads"]
@@ -149,10 +147,11 @@ def build_model_from_cfg(cfg: Dict[str, Any]):
             lambda_cs=loss_cfg["lambda_cs"],
             lambda_sg=loss_cfg["lambda_sg"],
             lambda_lp=loss_cfg["lambda_lp"],
-            
+
             lr=optim_cfg["lr"],
             weight_decay=optim_cfg["weight_decay"],
             use_adamw=optim_cfg["use_adamw"],
+            steps_per_epoch=steps_per_epoch,
         )
     else:
         raise ValueError(f"Unsupported model_type '{model_type}'. Expected 'multiscale'.")
@@ -292,7 +291,11 @@ def main():
     torch.set_float32_matmul_precision('high')
 
     dm = build_datamodule_from_cfg(cfg)
-    model = build_model_from_cfg(cfg)
+    # Explicitly setup datamodule to calculate steps_per_epoch
+    dm.setup("fit")
+    steps_per_epoch = len(dm.train_dataloader())
+    
+    model = build_model_from_cfg(cfg, steps_per_epoch=steps_per_epoch)
     trainer = build_trainer_from_cfg(cfg, raw_config_path=args.config)
 
     # Train