bsvae 0.1.1__py3-none-any.whl

bsvae/hyperparam.ini

@@ -0,0 +1,82 @@
+[Custom]
+    # General options
+    log_level = "info"
+    no_progress_bar = False
+    no_cuda = False
+    seed = 1234
+
+    # Training options
+    epochs = 100
+    batch_size = 64
+    lr = 5e-4
+    checkpoint_every = 30
+    dataset = "genenet"
+    experiment = "structured"
+
+    # Model options
+    model = "StructuredFactorVAE"
+    latent_dim = 10
+    hidden_dims = [512, 256]   # encoder hidden layers
+    dropout = 0.1
+    init_sd = 0.02
+    learn_var = False          # decoder per-gene variance
+
+    # Loss options
+    loss = "beta"
+    beta = 1.0                 # KL weight
+    l1_strength = 1e-3         # sparsity regularizer
+    lap_strength = 1e-4        # Laplacian smoothness
+
+    # Evaluations
+    is_metrics = True
+    no_test = False
+    is_eval_only = False
+    eval_batchsize = 512
+
+
+# ### DATASET COMMON ###
+[Common_genenet]
+    dataset = "genenet"
+    checkpoint_every = 100
+    epochs = 400
+
+
+# ### LOSS COMMON ###
+[Common_VAE]
+    loss = "VAE"
+    lr = 5e-4
+[Common_beta]
+    loss = "beta"
+    beta = 4.0
+    lr = 5e-4
+
+
+# ### EXPERIMENT SPECIFIC ###
+
+[VAE_genenet]
+    dataset = "genenet"
+    model = "StructuredFactorVAE"
+    loss = "VAE"
+    latent_dim = 8
+    epochs = 400
+    checkpoint_every = 50
+    lr = 1e-4
+    l1_strength = 1e-3
+    lap_strength = 1e-4
+
+[beta_genenet]
+    dataset = "genenet"
+    model = "StructuredFactorVAE"
+    loss = "beta"
+    beta = 4.0
+    latent_dim = 8
+    epochs = 400
+    checkpoint_every = 50
+    lr = 1e-4
+    l1_strength = 5e-4
+    lap_strength = 5e-5
+
+[debug]
+    epochs = 1
+    log_level = "debug"
+    no_test = True

bsvae/main.py

@@ -0,0 +1,237 @@
+#!/usr/bin/env python
+
+import sys
+import logging
+import argparse
+from os.path import join, dirname
+from configparser import ConfigParser
+
+import torch
+from torch import optim
+
+from bsvae.models import StructuredFactorVAE
+from bsvae.models.losses import BaseLoss
+from bsvae.utils.datasets import get_dataloaders
+from bsvae.utils.helpers import (
+    set_seed,
+    get_device,
+    get_n_params,
+    create_safe_directory,
+    FormatterNoDuplicate,
+    get_config_section,
+    update_namespace_,
+)
+from bsvae.utils import Trainer, Evaluator
+from bsvae.utils.ppi import load_ppi_laplacian
+from bsvae.utils.modelIO import save_model, load_model, load_metadata
+
+RES_DIR = "results"
+
+def load_config(config_path: str, section: str = "Custom") -> dict:
+    """Load hyperparameters from .ini file."""
+    parser = ConfigParser()
+    parser.read(config_path)
+    if section not in parser:
+        raise ValueError(f"Config section [{section}] not found in {config_path}")
+    return {k: ast_literal_eval(v) for k, v in parser[section].items()}
+
+
+def ast_literal_eval(val):
+    """Helper to parse lists, ints, floats from strings safely."""
+    import ast
+    try:
+        return ast.literal_eval(val)
+    except Exception:
+        return val
+
+
+def parse_arguments(cli_args):
+    """Parse CLI args, with defaults pulled from hyperparam.ini."""
+    pre_parser = argparse.ArgumentParser(add_help=False)
+    pre_parser.add_argument("--config", "-c", type=str,
+                            default=join(dirname(__file__), "hyperparam.ini"),
+                            help="Path to hyperparam.ini")
+    pre_parser.add_argument("--section", type=str, default="Custom",
+                            help="Section of .ini to load")
+    config_args, _ = pre_parser.parse_known_args(cli_args)
+
+    # Load defaults from config
+    config = load_config(config_args.config, config_args.section)
+
+    parser = argparse.ArgumentParser(
+        description="Training and evaluation for StructuredFactorVAE.",
+        formatter_class=FormatterNoDuplicate,
+        parents=[pre_parser],
+    )
+
+    # General
+    parser.add_argument("name", type=str, help="Experiment name.")
+    parser.add_argument("--seed", type=int, default=config.get("seed", 13))
+    parser.add_argument("--no-cuda", action="store_true",
+                        default=config.get("no_cuda", False))
+
+    # Training
+    parser.add_argument("--epochs", type=int, default=config.get("epochs", 100))
+    parser.add_argument("--batch-size", type=int,
+                        default=config.get("batch_size", 64))
+    parser.add_argument("--lr", type=float, default=config.get("lr", 5e-4))
+    parser.add_argument("--checkpoint-every", type=int,
+                        default=config.get("checkpoint_every", 10))
+
+    # Model
+    parser.add_argument("--latent-dim", "-z", type=int,
+                        default=config.get("latent_dim", 10))
+    parser.add_argument("--hidden-dims", "-Z", type=ast_literal_eval,
+                        default=config.get("hidden_dims", [128, 64]))
+    parser.add_argument("--dropout", type=float,
+                        default=config.get("dropout", 0.1))
+    parser.add_argument("--learn-var", action="store_true",
+                        default=config.get("learn_var", False))
+    parser.add_argument("--init-sd", type=float,
+                        default=config.get("init_sd", 0.02))
+
+    # Loss
+    parser.add_argument("--loss", type=str, default=config.get("loss", "VAE"),
+                        choices=["VAE", "beta"])
+    parser.add_argument("--beta", type=float, default=config.get("beta", 1.0))
+    parser.add_argument("--l1-strength", type=float,
+                        default=config.get("l1_strength", 1e-3))
+    parser.add_argument("--lap-strength", type=float,
+                        default=config.get("lap_strength", 1e-4))
+
+    # Evaluation
+    parser.add_argument("--is-eval-only", action="store_true",
+                        default=config.get("is_eval_only"))
+    parser.add_argument("--no-test", action="store_true",
+                        default=config.get("no_test"))
+    parser.add_argument("--eval-batchsize", type=int,
+                        default=config.get("eval_batchsize"))
+
+    # Dataset
+    parser.add_argument("--dataset", type=str,
+                        default=config.get("dataset", "genenet"))
+    parser.add_argument("--gene-expression-filename", type=str,
+                        default=config.get("gene_expression_filename", None),
+                        help="CSV with gene expression (genes x samples).")
+    parser.add_argument("--gene-expression-dir", type=str,
+                        default=config.get("gene_expression_dir", None),
+                        help="Directory with train/test splits (X_train.csv, X_test.csv).")
+
+    # PPI Priors
+    parser.add_argument("--ppi-taxid", type=str, default=config.get("ppi_taxid", "9606"))
+    parser.add_argument("--ppi-cache", type=str, default=config.get("ppi_cache", None))
+
+    args = parser.parse_args(cli_args)
+
+    # Validate dataset input
+    if bool(args.gene_expression_filename) == bool(args.gene_expression_dir):
+        parser.error("Specify exactly one of --gene-expression-filename or --gene-expression-dir.")
+
+    parser.set_defaults(**vars(args))
+    return parser.parse_args(cli_args)
+
+
+def setup_logging():
+    log_fmt = "%(asctime)s %(levelname)s: %(message)s"
+    logging.basicConfig(level=logging.INFO, format=log_fmt)
+    return logging.getLogger(__name__)
+
+
+def main(args):
+    logger = setup_logging()
+    set_seed(args.seed)
+    device = get_device(use_gpu=not args.no_cuda)
+    exp_dir = join(RES_DIR, args.name)
+    logger.info(f"Experiment directory: {exp_dir}")
+
+    # Training
+    if not args.is_eval_only:
+        create_safe_directory(exp_dir, logger=logger)
+
+        # Data
+        train_loader = get_dataloaders(
+            dataset=args.dataset,
+            batch_size=args.batch_size,
+            logger=logger,
+            train=True,
+            drop_last=True,
+            gene_expression_filename=args.gene_expression_filename,
+            gene_expression_dir=args.gene_expression_dir,
+        )
+        n_genes = train_loader.dataset[0][0].shape[-1]
+        logger.info(f"Training dataset size: {len(train_loader.dataset)}")
+
+        # PPI Laplacian
+        try:
+            gene_list = getattr(train_loader.dataset, "genes", None)
+            if gene_list is not None:
+                L, G = load_ppi_laplacian(
+                    gene_list,
+                    taxid=args.ppi_taxid,
+                    min_score=700,
+                    cache_dir=args.ppi_cache or "~/.bsvae/ppi",
+                )
+                logger.info(f"PPI Graph loaded: {G.number_of_nodes()} nodes, {G.number_of_edges()} edges")
+            else:
+                L = None
+        except Exception as e:
+            logger.warning(f"Could not load PPI Laplacian: {e}")
+            L = None
+
+        # Model
+        model = StructuredFactorVAE(
+            n_genes=n_genes,
+            n_latent=args.latent_dim,
+            hidden_dims=args.hidden_dims,
+            dropout=args.dropout,
+            init_sd=args.init_sd,
+            learn_var=args.learn_var,
+            L=L
+        ).to(device)
+        logger.info(f"Model initialized with {n_genes} genes, {args.latent_dim} latent dims")
+
+        # Loss + optimizer
+        optimizer = optim.Adam(model.parameters(), lr=args.lr)
+        loss_f = BaseLoss(beta=args.beta,
+                          l1_strength=args.l1_strength,
+                          lap_strength=args.lap_strength)
+
+        trainer = Trainer(
+            model, optimizer, loss_f, device=device, logger=logger,
+            save_dir=exp_dir, is_progress_bar=not args.no_cuda,
+        )
+        trainer(train_loader, epochs=int(args.epochs),
+                checkpoint_every=int(args.checkpoint_every))
+        save_model(trainer.model, exp_dir, metadata=vars(args))
+
+    # Evaluation
+    if not args.no_test:
+        model = load_model(exp_dir, is_gpu=not args.no_cuda)
+        metadata = load_metadata(exp_dir)
+        test_loader = get_dataloaders(
+            dataset="genenet",
+            batch_size=args.eval_batchsize,
+            shuffle=False,
+            logger=logger,
+            train=False,
+            drop_last=True,
+            gene_expression_filename=args.gene_expression_filename,
+            gene_expression_dir=args.gene_expression_dir,
+        )
+        loss_f = BaseLoss(beta=args.beta,
+                          l1_strength=args.l1_strength,
+                          lap_strength=args.lap_strength)
+        evaluator = Evaluator(
+            model, loss_f, device=device, logger=logger, save_dir=exp_dir,
+            is_progress_bar=not args.no_cuda,
+        )
+        evaluator(test_loader)
+
+
+def cli():
+    args = parse_arguments(sys.argv[1:])
+    main(args)
+
+
+if __name__ == "__main__":
+    cli()

bsvae/models/__init__.py

@@ -0,0 +1,4 @@
+from .vae import BaseVAE
+from .encoder import StructuredEncoder
+from .decoder import StructuredDecoder
+from .structured import StructuredFactorVAE

bsvae/models/decoder.py

@@ -0,0 +1,90 @@
+"""
+Decoder for StructuredFactorVAE
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class StructuredDecoder(nn.Module):
+    """
+    Structured decoder for gene expression:
+      - Each latent factor corresponds to a module.
+      - Decoder weights are optionally masked by prior gene×module assignments.
+      - Supports sparsity and Laplacian penalties.
+      - Outputs mean and dispersion for Negative Binomial likelihood.
+
+    Parameters
+    ----------
+    n_genes : int
+        Number of genes (G).
+    n_latent : int
+        Number of latent factors/modules (K).
+    mask : torch.Tensor or None
+        Optional binary mask of shape (G, K) where 1 means gene↔factor allowed.
+    init_sd : float
+        Std for normal init of decoder weights.
+    """
+    def __init__(self, n_genes: int, n_latent: int,
+                 mask: torch.Tensor = None,
+                 init_sd: float = 0.02,
+                 learn_var: bool = False):
+        super().__init__()
+        self.n_genes = n_genes
+        self.n_latent = n_latent
+        self.learn_var = learn_var
+
+        # Weight matrix W: (G, K)
+        self.W = nn.Parameter(torch.randn(n_genes, n_latent) * init_sd)
+        self.bias = nn.Parameter(torch.zeros(n_genes))
+
+        # Mask (non-trainable)
+        if mask is not None:
+            assert mask.shape == (n_genes, n_latent)
+            self.register_buffer("mask", mask.float())
+        else:
+            self.mask = None
+
+        # Gene-specific variance (optional)
+        if self.learn_var:
+            self.log_var = nn.Parameter(torch.zeros(n_genes))
+        else:
+            self.log_var = None
+
+    def forward(self, z: torch.Tensor):
+        """
+        Parameters
+        ----------
+        z : torch.Tensor
+            Latent codes, shape (batch, K).
+
+        Returns
+        -------
+        recon_x : torch.Tensor
+            Reconstructed gene expression (batch, G).
+        log_var : torch.Tensor
+            Gene-specific log-variance (G,).
+        """
+        # Apply mask if provided
+        W_eff = self.W
+        if self.mask is not None:
+            W_eff = W_eff * self.mask  # elementwise mask
+
+        recon_x = z @ W_eff.T + self.bias
+        return recon_x, self.log_var
+
+    # Regularizers
+    def group_sparsity_penalty(self, l1_strength: float = 1e-3):
+        """
+        Group sparsity across genes per latent factor.
+        """
+        W_eff = self.W * self.mask if self.mask is not None else self.W
+        return l1_strength * torch.sum(torch.abs(W_eff))
+
+    def laplacian_penalty(self, L: torch.Tensor, lap_strength: float = 1e-3):
+        """
+        Laplacian smoothness: tr(W^T L W).
+        """
+        W_eff = self.W * self.mask if self.mask is not None else self.W
+        LW = torch.sparse.mm(L, W_eff) if L.is_sparse else torch.matmul(L, W_eff)
+        penalty = torch.sum(W_eff * LW)  # trace approx
+        return lap_strength * penalty

bsvae/models/encoder.py

@@ -0,0 +1,71 @@
+"""
+Encoder for StructuredFactorVAE
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class StructuredEncoder(nn.Module):
+    """
+    Encoder for StructuredFactorVAE.
+    Maps normalized gene expression (log-CPM/TPM) to latent Gaussian parameters.
+
+    Parameters
+    ----------
+    n_genes : int
+        Number of input genes (features).
+    n_latent : int
+        Dimensionality of latent space (number of biological modules).
+    hidden_dims : list of int, optional
+        Sizes of hidden layers. Default: [512, 256, 128].
+    dropout : float
+        Dropout probability applied after hidden layers.
+    """
+    def __init__(self, n_genes: int, n_latent: int,
+                 hidden_dims: list[int] = None,
+                 dropout: float = 0.1):
+        super().__init__()
+        self.n_genes = n_genes
+        self.n_latent = n_latent
+        self.hidden_dims = hidden_dims if hidden_dims is not None else [512, 256, 128]
+        self.dropout = dropout
+
+        # Build feedforward encoder network
+        modules = []
+        input_dim = n_genes
+        for h_dim in self.hidden_dims:
+            modules.append(nn.Linear(input_dim, h_dim))
+            modules.append(nn.ReLU())
+            if dropout > 0:
+                modules.append(nn.Dropout(dropout))
+            input_dim = h_dim
+        self.encoder = nn.Sequential(*modules)
+
+        # Outputs: mean and log-variance
+        self.fc_mu = nn.Linear(input_dim, n_latent)
+        self.fc_logvar = nn.Linear(input_dim, n_latent)
+
+    def forward(self, x: torch.Tensor):
+        """
+        Forward pass through encoder.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input expression data (batch, G).
+
+        Returns
+        -------
+        mu : torch.Tensor
+            Latent mean (batch, K).
+        logvar : torch.Tensor
+            Latent log variance (batch, K).
+        z : torch.Tensor
+            Sampled latent code (batch, K).
+        """
+        h = self.encoder(x)
+        mu = self.fc_mu(h)
+        logvar = self.fc_logvar(h)
+
+        return mu, logvar

bsvae/models/losses.py

@@ -0,0 +1,116 @@
+"""
+Losses for StructuredFactorVAE.
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Dict
+
+def gaussian_nll(x, recon_x, log_var=None, reduction="mean"):
+    """
+    Gaussian negative log-likelihood per gene.
+    """
+    if log_var is None:
+        return F.mse_loss(recon_x, x, reduction=reduction)
+
+    var = torch.exp(log_var)
+    nll = 0.5 * (torch.log(2 * torch.pi * var) + (x - recon_x) ** 2 / var)
+    if reduction == "mean":
+        return nll.mean()
+    elif reduction == "sum":
+        return nll.sum()
+    else:
+        return nll
+
+
+def kl_normal_loss(mu, logvar, reduction="sum"):
+    """
+    KL divergence between q(z|x)=N(mu, sigma^2) and prior p(z)=N(0,I).
+    """
+    kl = -0.5 * (1 + logvar - mu.pow(2) - logvar.exp())
+    if reduction == "sum":
+        return kl.sum()
+    elif reduction == "mean":
+        return kl.mean()
+    else:
+        return kl
+
+
+class BaseLoss(nn.Module):
+    """
+    Base loss for StructuredFactorVAE with biological regularizers.
+    """
+    def __init__(self, beta: float = 1.0,
+                 l1_strength: float = 1e-3,
+                 lap_strength: float = 1e-4,
+                 record_loss_every: int = 50):
+        super().__init__()
+        self.beta = beta
+        self.l1_strength = l1_strength
+        self.lap_strength = lap_strength
+        self.n_train_steps = 0
+        self.record_loss_every = record_loss_every
+
+    def forward(self,
+                x: torch.Tensor, recon_x: torch.Tensor,
+                mu: torch.Tensor, logvar: torch.Tensor,
+                model: nn.Module,
+                L: Optional[torch.Tensor] = None,
+                storer: Optional[Dict[str, list]] = None,
+                is_train: bool = True) -> torch.Tensor:
+        """
+        Compute StructuredFactorVAE loss.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input data (batch, G).
+        recon_x : torch.Tensor
+            Reconstruction from decoder (batch, G).
+        mu : torch.Tensor
+            Latent mean (batch, K).
+        logvar : torch.Tensor
+            Latent log variance (batch, K).
+        model : nn.Module
+            Structured decoder (with sparsity/laplacian methods).
+        storer : dict
+            Dictionary for logging intermediate values.
+        is_train : bool
+            Whether in training mode.
+
+        Returns
+        -------
+        loss : torch.Tensor
+            Total loss for this batch.
+        """
+        # Reconstruction loss: Gaussian NLL if available, else MSE
+        if hasattr(model.decoder, "log_var") and model.decoder.log_var is not None:
+            recon_loss = gaussian_nll(x, recon_x, log_var=model.decoder.log_var,
+                                      reduction="mean")
+        else:
+            recon_loss = F.mse_loss(recon_x, x, reduction="mean")
+
+        # KL divergence
+        kl_loss = kl_normal_loss(mu, logvar, reduction="mean")
+
+        # Biological regularizers
+        sparsity_loss = model.decoder.group_sparsity_penalty(self.l1_strength)
+        laplacian_loss = 0.0
+        if model.laplacian_matrix is not None:
+            laplacian_loss = model.decoder.laplacian_penalty(
+                model.laplacian_matrix, self.lap_strength
+            )
+
+        # Total loss
+        loss = recon_loss + self.beta * kl_loss + sparsity_loss + laplacian_loss
+
+        # Logging
+        if storer is not None:
+            storer.setdefault("recon_loss", []).append(recon_loss.item())
+            storer.setdefault("kl_loss", []).append(kl_loss.item())
+            storer.setdefault("sparsity_loss", []).append(sparsity_loss.detach().item())
+            if model.laplacian_matrix is not None:
+                storer.setdefault("laplacian_loss", []).append(laplacian_loss.detach().item())
+            storer.setdefault("loss", []).append(loss.item())
+
+        return loss

bsvae/models/structured.py

@@ -0,0 +1,75 @@
+"""
+Structured Factor VAE: wraps encoder + decoder for end-to-end training.
+"""
+import torch
+import torch.nn as nn
+from .vae import BaseVAE
+from .encoder import StructuredEncoder
+from .decoder import StructuredDecoder
+from ..utils.initialization import weights_init
+
+class StructuredFactorVAE(BaseVAE):
+    """
+    Structured Factor VAE for gene expression.
+
+    Encoder: maps log-CPM input (batch, G) → latent mean + logvar.
+    Decoder: maps latent z → reconstructed expression (batch, G).
+             Supports optional gene-specific variance (log_var).
+
+    Parameters
+    ----------
+    n_genes : int
+        Number of input genes.
+    n_latent : int
+        Number of latent dimensions (modules).
+    hidden_dims : list of int
+        Hidden layer sizes for encoder.
+    dropout : float
+        Dropout probability in encoder.
+    mask : torch.Tensor or None
+        Optional binary gene×module mask for decoder (G, K).
+    init_sd : float
+        Std for decoder weight initialization.
+    learn_var : bool
+        If True, decoder learns per-gene log variance.
+    """
+    def __init__(self, n_genes: int, n_latent: int,
+                 hidden_dims=None, dropout: float = 0.1,
+                 mask: torch.Tensor = None,
+                 init_sd: float = 0.02,
+                 learn_var: bool = False,
+                 L: torch.Tensor = None):
+        super().__init__(n_genes, n_latent)
+        self.hidden_dims = hidden_dims or [512, 256, 128]
+        self.dropout = dropout
+        self.encoder = StructuredEncoder(
+            n_genes=n_genes,
+            n_latent=n_latent,
+            hidden_dims=hidden_dims,
+            dropout=dropout
+        )
+        self.decoder = StructuredDecoder(
+            n_genes=n_genes,
+            n_latent=n_latent,
+            mask=mask,
+            init_sd=init_sd,
+            learn_var=learn_var
+        )
+        self.laplacian_matrix = L
+
+    def forward(self, x: torch.Tensor):
+        mu, logvar = self.encoder(x)
+        z = self.reparameterize(mu, logvar)
+        recon_x, log_var = self.decoder(z)
+        return recon_x, mu, logvar, z, log_var
+
+    def group_sparsity_penalty(self, l1_strength: float = 1e-3):
+        return self.decoder.group_sparsity_penalty(l1_strength)
+
+    def laplacian_penalty(self, L: torch.Tensor, lap_strength: float = 1e-3):
+        return self.decoder.laplacian_penalty(L, lap_strength)
+
+    def reset_parameters(self, activation: str = "relu"):
+        """Reset all learnable parameters with custom init."""
+        self.activation = activation
+        self.apply(lambda m: weights_init(m, activation=self.activation))