Implement Adversarial Autoencoder Components

- Added `GradReverseLayer` for gradient reversal functionality.
- Introduced `ResidualBlock` for building residual networks.
- Created `make_mlp` function to construct multi-layer perceptrons with optional residual connections.
- Developed `AEBatchCorrector` class for adversarial batch correction, integrating encoder, decoder, and adversarial components.
- Established a model factory in `factory.py` to construct models based on command-line arguments, supporting both AE and VAE+attention variants.
- Modularized code for improved reusability and maintainability.

Author: Stefan Ritchie

NN_batch_correct.py

@@ -37,6 +37,8 @@
 
 import torch
 import torch.nn as nn
+from models import AEBatchCorrector, GradReverseLayer, ResidualBlock, make_mlp
+from models.factory import build_model_from_args
 # AMP GradScaler compatibility (torch>=2 provides torch.amp.GradScaler; older used torch.cuda.amp.GradScaler)
 try:  # prefer new API to silence deprecation warnings when available
     from torch.amp import GradScaler  # type: ignore
@@ -294,161 +296,13 @@ def _eta2_by_group(Z: np.ndarray, groups: np.ndarray):
 # ----------------------------
 # Gradient Reversal
 # ----------------------------
-
-class GradReverse(torch.autograd.Function):
-    @staticmethod
-    def forward(ctx, x, lambda_):
-        ctx.lambda_ = lambda_
-        return x.view_as(x)
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        return grad_output.neg() * ctx.lambda_, None
-
-
-class GradReverseLayer(nn.Module):
-    def __init__(self, lambda_: float = 1.0):
-        super().__init__()
-        self.lambda_ = lambda_
-
-    def forward(self, x):
-        return GradReverse.apply(x, self.lambda_)
-
-    def set_lambda(self, lambda_: float):
-        self.lambda_ = lambda_
+# Moved to models.ae (imported above)
 
 
 # ----------------------------
 # Model
 # ----------------------------
-
-class ResidualBlock(nn.Module):
-    """
-    Residual block with skip connection: x + FFN(x) + LayerNorm
-    """
-    def __init__(self, hidden_size: int, dropout: float = 0.0):
-        super().__init__()
-        self.ffn = nn.Sequential(
-            nn.Linear(hidden_size, hidden_size),
-            nn.SiLU(),
-            nn.Dropout(dropout),
-            nn.Linear(hidden_size, hidden_size),
-            nn.Dropout(dropout)
-        )
-        self.layer_norm = nn.LayerNorm(hidden_size)
-    
-    def forward(self, x):
-        return self.layer_norm(x + self.ffn(x))
-
-
-def make_mlp(sizes, dropout=0.0, last_activation=None, use_residual=False):
-    """
-    Hidden layers: Linear -> LayerNorm -> SiLU -> Dropout
-    Output layer: optional activation per arg
-    
-    If use_residual=True, uses ResidualBlock layers instead.
-    Note: For residual connections to work, all hidden layer sizes must be the same.
-    """
-    if use_residual:
-        if len(sizes) < 3:
-            raise ValueError("Residual networks need at least input, hidden, and output layers")
-        
-        # Check that all hidden layers have the same size for residual connections
-        hidden_sizes = sizes[1:-1]
-        if len(set(hidden_sizes)) > 1:
-            raise ValueError(f"For residual connections, all hidden layer sizes must be the same. Got: {hidden_sizes}")
-        
-        hidden_size = hidden_sizes[0]
-        num_hidden_layers = len(hidden_sizes)
-        
-        layers = []
-        
-        # Input projection to hidden size
-        layers.append(nn.Linear(sizes[0], hidden_size))
-        layers.extend([nn.LayerNorm(hidden_size), nn.SiLU(), nn.Dropout(dropout)])
-        
-        # Residual blocks
-        for _ in range(num_hidden_layers):
-            layers.append(ResidualBlock(hidden_size, dropout))
-        
-        # Output layer
-        layers.append(nn.Linear(hidden_size, sizes[-1]))
-        if last_activation == "relu":
-            layers.append(nn.ReLU())
-        elif last_activation == "tanh":
-            layers.append(nn.Tanh())
-        elif last_activation == "sigmoid":
-            layers.append(nn.Sigmoid())
-        
-        return nn.Sequential(*layers)
-    
-    else:
-        # Original implementation
-        layers = []
-        for i in range(len(sizes) - 1):
-            in_f, out_f = sizes[i], sizes[i + 1]
-            layers.append(nn.Linear(in_f, out_f))
-            if i < len(sizes) - 2:
-                layers += [nn.LayerNorm(out_f), nn.SiLU(), nn.Dropout(dropout)]
-            else:
-                if last_activation == "relu":
-                    layers += [nn.ReLU()]
-                elif last_activation == "tanh":
-                    layers += [nn.Tanh()]
-                elif last_activation == "sigmoid":
-                    layers += [nn.Sigmoid()]
-        return nn.Sequential(*layers)
-
-
-class AEBatchCorrector(nn.Module):
-    def __init__(
-        self,
-        n_genes: int,
-        latent_dim: int = 32,
-        enc_hidden=(1024, 256),
-        dec_hidden=(256, 1024),
-        adv_hidden=(128,),
-        sup_hidden=(64,),
-        n_batches: int = 2,
-        n_labels: Optional[int] = None,
-        dropout: float = 0.1,
-        adv_lambda: float = 1.0,
-        use_residual: bool = False,
-    ):
-        super().__init__()
-        self.n_labels = n_labels
-        self.grl = GradReverseLayer(lambda_=adv_lambda)
-
-        enc_sizes = [n_genes] + list(enc_hidden) + [latent_dim]
-        dec_sizes = [latent_dim] + list(dec_hidden) + [n_genes]
-        self.encoder = make_mlp(enc_sizes, dropout=dropout, last_activation=None, use_residual=use_residual)
-        self.decoder = make_mlp(dec_sizes, dropout=dropout, last_activation=None, use_residual=use_residual)
-
-        adv_sizes = [latent_dim] + list(adv_hidden) + [n_batches]
-        self.adv = make_mlp(adv_sizes, dropout=dropout, last_activation=None, use_residual=False)  # Keep simple for adversarial head
-
-        if n_labels is not None:
-            sup_sizes = [latent_dim] + list(sup_hidden) + [n_labels]
-            self.sup = make_mlp(sup_sizes, dropout=dropout, last_activation=None, use_residual=False)  # Keep simple for supervised head
-        else:
-            self.sup = None
-
-    def forward(self, x, adv_lambda: Optional[float] = None):
-        z = self.encoder(x)
-        x_hat = self.decoder(z)
-        if adv_lambda is not None:
-            self.grl.set_lambda(adv_lambda)
-        z_rev = self.grl(z)
-        batch_logits = self.adv(z_rev)
-        label_logits = self.sup(z) if self.sup is not None else None
-        return x_hat, batch_logits, label_logits, z
-
-    @torch.no_grad()
-    def reconstruct(self, x):
-        """Fast path at inference time (skips adversary & GRL)."""
-        z = self.encoder(x)
-        x_hat = self.decoder(z)
-        return x_hat, z
+# Moved to models.ae (imported above)
 
 
 # ----------------------------
@@ -1192,7 +1046,8 @@ def main():
     ap.add_argument("--viz_hvg_top", default=2000, type=int, help="Top-N most variable genes to use for PCA visualisations (0=use all)")
     ap.add_argument("--viz_pca_before", default="pca_before.png", type=str, help="Output path for PCA before correction")
     ap.add_argument("--viz_pca_after", default="pca_after.png", type=str, help="Output path for PCA after correction")
-    ap.add_argument("--viz_boxplot", default="logCPM_boxplots.png", type=str, help="Output path for logCPM boxplots")
+    # Accept optional value; if provided without a path, use the default filename
+    ap.add_argument("--viz_boxplot", nargs="?", const="logCPM_boxplots.png", default="logCPM_boxplots.png", type=str, help="Output path for logCPM boxplots (optional value)")
     # Differentiable biology preservation loss & adaptive adversary based on cond_sil
     ap.add_argument("--bio_weight", default=0.0, type=float, help="Weight for supervised center loss on latent (preserve biology)")
     ap.add_argument("--bio_gamma", default=0.5, type=float, help="Relative weight for between-class separation in center loss")
@@ -1339,38 +1194,12 @@ def main():
     dl_val = DataLoader(ds_val, batch_size=args.batch_size, shuffle=False, **loader_kwargs)
 
     # Model
-    if args.model_type == "vae_attention":
-        if VaeAttentionBatchCorrector is None:
-            raise RuntimeError("vae_attention model requested but module not available.")
-        model = VaeAttentionBatchCorrector(
-            num_genes=logcpm.shape[1],
-            num_batches=len(batch_classes),
-            latent_dim=args.latent_dim,
-            hidden_dim=args.vae_hidden_dim,
-            attention_dim=args.vae_attention_dim,
-            n_heads=args.vae_attn_heads,
-            dropout=args.dropout,
-            dispersion=args.vae_dispersion,
-            attn_max_tokens=args.attn_max_tokens,
-        )
-    else:
-        enc_hidden = tuple(int(x) for x in args.enc_hidden.split(",") if x.strip())
-        dec_hidden = tuple(int(x) for x in args.dec_hidden.split(",") if x.strip())
-        adv_hidden = tuple(int(x) for x in args.adv_hidden.split(",") if x.strip())
-        sup_hidden = tuple(int(x) for x in args.sup_hidden.split(",") if x.strip())
-        model = AEBatchCorrector(
-            n_genes=logcpm.shape[1],
-            latent_dim=args.latent_dim,
-            enc_hidden=enc_hidden,
-            dec_hidden=dec_hidden,
-            adv_hidden=adv_hidden,
-            sup_hidden=sup_hidden,
-            n_batches=len(batch_classes),
-            n_labels=(len(label_classes) if label_classes is not None else None),
-            dropout=args.dropout,
-            adv_lambda=args.adv_weight,
-            use_residual=args.use_residual,
-        )
+    model = build_model_from_args(
+        args,
+        n_genes=logcpm.shape[1],
+        n_batches=len(batch_classes),
+        n_labels=(len(label_classes) if label_classes is not None else None),
+    )
 
     # Optional torch.compile (PyTorch 2+)
     if args.compile:

README.md

@@ -9,7 +9,7 @@
 
 <p align=center>
 <img src="logCPM_boxplots.png" alt="LogCPM Boxplots" height="200" align=center />
-<img src="VAEModel\sweep_pca_panel.png" alt="Before & After PCA Plots" height="200" align=center />
+<img src="assets\pca_panel.png" alt="Before & After PCA Plots" height="200" align=center />
 </p>
 
 An adversarial autoencoder for bulk RNA-seq batch effect correction. It learns a latent representation that preserves biological signal (optional supervised head) while discouraging batch-specific variation via a gradient reversal adversary. Outputs a batch-corrected expression matrix (logCPM scale) and optional latent embedding plus visual diagnostics.
@@ -34,6 +34,9 @@ An adversarial autoencoder for bulk RNA-seq batch effect correction. It learns a
 ---
 ## Repository Structure (selected)
 - `NN_batch_correct.py`  Main training + correction script
+- `models/`              Model components (adversarial AE) and factory
+  - `models/ae.py`      AEBatchCorrector, GradReverseLayer, ResidualBlock, make_mlp
+  - `models/factory.py` Helper to build AE or VAE+Attention models from CLI args
 - `visualise.py`         Standalone PCA / boxplot + architecture diagram utilities
 - `bulk_counts.csv`      Example counts matrix (shape: genes × samples or samples × genes)
 - `sample_meta.csv`      Example metadata (columns include sample,batch[,condition])
@@ -70,6 +73,11 @@ pip install -r requirements.txt
 Optional extras:
 - Set `WANDB_MODE=offline` to avoid network usage.
 
+---
+## Refactor Notes
+- Model code has been extracted from `NN_batch_correct.py` into `models/ae.py` and is re-exported via `models/__init__.py`. The CLI behavior is unchanged.
+- A lightweight factory (`models/factory.py`) centralizes model construction for AE vs. VAE+Attention to keep the training script focused on data and orchestration.
+
 ---
 ## Quick Start
 Minimal unsupervised run (only batches):