bsvae 0.3.1__tar.gz → 0.3.2__tar.gz

bsvae-0.3.2/PKG-INFO

@@ -0,0 +1,224 @@
+Metadata-Version: 2.4
+Name: bsvae
+Version: 0.3.2
+Summary: Gaussian Mixture VAE for multi-modal biological module discovery in omics data.
+License: GPL-3.0-only
+License-File: LICENSE
+Keywords: variational-autoencoder,gaussian-mixture-model,pytorch,bioinformatics,gene-expression,module-discovery,omics
+Author: Kynon J Benjamin
+Author-email: kj.benjamin90@gmail.com
+Requires-Python: >=3.11,<3.14
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: GNU General Public License v3 (GPLv3)
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Requires-Dist: faiss-cpu (>=1.7)
+Requires-Dist: h5py (>=3.0)
+Requires-Dist: numpy (<2.3)
+Requires-Dist: pandas (>=2.3.2,<3.0.0)
+Requires-Dist: scikit-learn (>=1.7.2,<2.0.0)
+Requires-Dist: scipy (>=1.16.2,<2.0.0)
+Requires-Dist: torch (>=2.8.0,<3.0.0)
+Requires-Dist: tqdm (>=4.67.1,<5.0.0)
+Project-URL: Bug Tracker, https://github.com/heart-gen/BSVAE/issues
+Project-URL: homepage, https://bsvae.readthedocs.io/en/latest/
+Project-URL: repository, https://github.com/heart-gen/BSVAE
+Description-Content-Type: text/markdown
+
+# BSVAE: Gaussian Mixture VAE for Module Discovery
+
+[![PyTorch](https://img.shields.io/badge/PyTorch-%3E%3D2.8-orange)](https://pytorch.org/)
+[![Python](https://img.shields.io/badge/Python-%3E%3D3.11-blue)](https://python.org/)
+[![License](https://img.shields.io/badge/license-GPLv3-blue.svg)](LICENSE)
+[![Documentation](https://readthedocs.org/projects/bsvae/badge/?version=latest)](https://bsvae.readthedocs.io/)
+
+BSVAE is a PyTorch package for feature-level module discovery in omics data. The main model, `GMMModuleVAE`, learns latent structure across features such as genes, transcripts, or proteins, then supports downstream network extraction, module assignment, latent analysis, and simulation-based benchmarking.
+
+## What The Package Provides
+
+- `bsvae-train` for training a two-phase GMM-VAE
+- `bsvae-sweep-k` for selecting the number of modules (`K`) with held-out validation and optional stability replicates
+- `bsvae-networks` for post-training network extraction, module extraction, latent export, and latent analysis
+- `bsvae-simulate` for synthetic data generation, scenario-grid simulation, and module-recovery benchmarking
+
+## Installation
+
+Install from PyPI:
+
+```bash
+pip install bsvae
+```
+
+Install from source:
+
+```bash
+git clone https://github.com/heart-gen/BSVAE.git
+cd BSVAE
+pip install -e .
+```
+
+Optional dependency:
+
+- `anndata` is only needed for `.h5ad` input or `.h5ad` latent export helpers
+
+## Input Data Contract
+
+Training and most downstream commands expect an expression matrix in `features x samples` orientation:
+
+- rows are features such as genes, transcripts, or proteins
+- columns are sample IDs
+- the first column is the feature index when using CSV/TSV
+
+Supported loaders:
+
+- `.csv` / `.csv.gz`
+- `.tsv` / `.tsv.gz`
+- `.h5` / `.hdf5`
+- `.h5ad` with optional `anndata`
+
+## CLI Entry Points
+
+- `bsvae-train`
+- `bsvae-sweep-k`
+- `bsvae-networks`
+- `bsvae-simulate`
+
+## Quickstart
+
+Train a small model:
+
+```bash
+bsvae-train pilot_run \
+  --dataset data/expression.csv \
+  --epochs 50 \
+  --n-modules 12 \
+  --latent-dim 16
+```
+
+Recommended production flow: select `K` first, then use the retrained final model.
+
+```bash
+bsvae-sweep-k sweep1 \
+  --dataset data/expression.csv \
+  --k-grid 8,12,16,24,32 \
+  --sweep-epochs 60 \
+  --stability-reps 5 \
+  --val-frac 0.1
+```
+
+This writes sweep outputs under `results/sweep1/sweep_k/` and retrains the selected model under `results/sweep1/final_k<K>/`.
+
+Extract module assignments from the final model:
+
+```bash
+bsvae-networks extract-modules \
+  --model-path results/sweep1/final_k16 \
+  --dataset data/expression.csv \
+  --output-dir results/sweep1/final_k16/modules \
+  --expr data/expression.csv \
+  --soft-eigengenes
+```
+
+Export latents:
+
+```bash
+bsvae-networks export-latents \
+  --model-path results/sweep1/final_k16 \
+  --dataset data/expression.csv \
+  --output results/sweep1/final_k16/latents
+```
+
+This command writes `results/sweep1/final_k16/latents.npz` containing `mu`, `logvar`, `gamma`, and `feature_ids`.
+
+Extract feature-feature networks:
+
+```bash
+bsvae-networks extract-networks \
+  --model-path results/sweep1/final_k16 \
+  --dataset data/expression.csv \
+  --output-dir results/sweep1/final_k16/networks \
+  --methods mu_cosine gamma_knn \
+  --top-k 50
+```
+
+Run latent-space analysis:
+
+```bash
+bsvae-networks latent-analysis \
+  --model-path results/sweep1/final_k16 \
+  --dataset data/expression.csv \
+  --output-dir results/sweep1/final_k16/latent_analysis \
+  --kmeans-k 16 \
+  --umap
+```
+
+## Training Outputs
+
+`bsvae-train` writes to `results/<experiment>/` by default:
+
+- `model.pt`
+- `specs.json`
+- `train_losses.csv`
+- `model-<epoch>.pt` checkpoints when `--checkpoint-every` is greater than zero
+
+`bsvae-sweep-k` writes:
+
+- `results/<name>/sweep_k/sweep_results.csv`
+- `results/<name>/sweep_k/sweep_summary.json`
+- `results/<name>/sweep_k/k<K>/rep_<rep>/...` per-sweep run outputs
+- `results/<name>/final_k<K>/...` for the final retrained model when `--train-final` is enabled
+
+## Simulation and Benchmarking
+
+Generate one synthetic dataset:
+
+```bash
+bsvae-simulate generate \
+  --output data/sim_expr.csv \
+  --save-ground-truth data/sim_truth.csv
+```
+
+Generate a scenario grid:
+
+```bash
+bsvae-simulate init-config --output sim.yaml
+
+bsvae-simulate generate-grid \
+  --config sim.yaml \
+  --outdir results/sim_pub_v1 \
+  --reps 30 \
+  --base-seed 13
+```
+
+Validate the generated grid:
+
+```bash
+bsvae-simulate validate-grid --grid-dir results/sim_pub_v1
+```
+
+Each scenario replicate writes method-ready files under `results/sim_pub_v1/scenarios/<scenario_id>/rep_<rep>/`, including:
+
+- `expr/features_x_samples.tsv.gz`
+- `expr/samples_x_features.tsv.gz`
+- `truth/modules_hard.csv`
+- `method_inputs.json`
+
+## Documentation
+
+See the docs for the full workflow and command reference:
+
+- `docs/index.md`
+- `docs/quickstart.md`
+- `docs/tutorial.md`
+- `docs/cli.md`
+- `docs/networks.md`
+- `docs/hyperparameters.md`
+
+## License
+
+This project is licensed under the [GNU General Public License v3.0](LICENSE).
+

bsvae-0.3.2/README.md

@@ -0,0 +1,192 @@
+# BSVAE: Gaussian Mixture VAE for Module Discovery
+
+[![PyTorch](https://img.shields.io/badge/PyTorch-%3E%3D2.8-orange)](https://pytorch.org/)
+[![Python](https://img.shields.io/badge/Python-%3E%3D3.11-blue)](https://python.org/)
+[![License](https://img.shields.io/badge/license-GPLv3-blue.svg)](LICENSE)
+[![Documentation](https://readthedocs.org/projects/bsvae/badge/?version=latest)](https://bsvae.readthedocs.io/)
+
+BSVAE is a PyTorch package for feature-level module discovery in omics data. The main model, `GMMModuleVAE`, learns latent structure across features such as genes, transcripts, or proteins, then supports downstream network extraction, module assignment, latent analysis, and simulation-based benchmarking.
+
+## What The Package Provides
+
+- `bsvae-train` for training a two-phase GMM-VAE
+- `bsvae-sweep-k` for selecting the number of modules (`K`) with held-out validation and optional stability replicates
+- `bsvae-networks` for post-training network extraction, module extraction, latent export, and latent analysis
+- `bsvae-simulate` for synthetic data generation, scenario-grid simulation, and module-recovery benchmarking
+
+## Installation
+
+Install from PyPI:
+
+```bash
+pip install bsvae
+```
+
+Install from source:
+
+```bash
+git clone https://github.com/heart-gen/BSVAE.git
+cd BSVAE
+pip install -e .
+```
+
+Optional dependency:
+
+- `anndata` is only needed for `.h5ad` input or `.h5ad` latent export helpers
+
+## Input Data Contract
+
+Training and most downstream commands expect an expression matrix in `features x samples` orientation:
+
+- rows are features such as genes, transcripts, or proteins
+- columns are sample IDs
+- the first column is the feature index when using CSV/TSV
+
+Supported loaders:
+
+- `.csv` / `.csv.gz`
+- `.tsv` / `.tsv.gz`
+- `.h5` / `.hdf5`
+- `.h5ad` with optional `anndata`
+
+## CLI Entry Points
+
+- `bsvae-train`
+- `bsvae-sweep-k`
+- `bsvae-networks`
+- `bsvae-simulate`
+
+## Quickstart
+
+Train a small model:
+
+```bash
+bsvae-train pilot_run \
+  --dataset data/expression.csv \
+  --epochs 50 \
+  --n-modules 12 \
+  --latent-dim 16
+```
+
+Recommended production flow: select `K` first, then use the retrained final model.
+
+```bash
+bsvae-sweep-k sweep1 \
+  --dataset data/expression.csv \
+  --k-grid 8,12,16,24,32 \
+  --sweep-epochs 60 \
+  --stability-reps 5 \
+  --val-frac 0.1
+```
+
+This writes sweep outputs under `results/sweep1/sweep_k/` and retrains the selected model under `results/sweep1/final_k<K>/`.
+
+Extract module assignments from the final model:
+
+```bash
+bsvae-networks extract-modules \
+  --model-path results/sweep1/final_k16 \
+  --dataset data/expression.csv \
+  --output-dir results/sweep1/final_k16/modules \
+  --expr data/expression.csv \
+  --soft-eigengenes
+```
+
+Export latents:
+
+```bash
+bsvae-networks export-latents \
+  --model-path results/sweep1/final_k16 \
+  --dataset data/expression.csv \
+  --output results/sweep1/final_k16/latents
+```
+
+This command writes `results/sweep1/final_k16/latents.npz` containing `mu`, `logvar`, `gamma`, and `feature_ids`.
+
+Extract feature-feature networks:
+
+```bash
+bsvae-networks extract-networks \
+  --model-path results/sweep1/final_k16 \
+  --dataset data/expression.csv \
+  --output-dir results/sweep1/final_k16/networks \
+  --methods mu_cosine gamma_knn \
+  --top-k 50
+```
+
+Run latent-space analysis:
+
+```bash
+bsvae-networks latent-analysis \
+  --model-path results/sweep1/final_k16 \
+  --dataset data/expression.csv \
+  --output-dir results/sweep1/final_k16/latent_analysis \
+  --kmeans-k 16 \
+  --umap
+```
+
+## Training Outputs
+
+`bsvae-train` writes to `results/<experiment>/` by default:
+
+- `model.pt`
+- `specs.json`
+- `train_losses.csv`
+- `model-<epoch>.pt` checkpoints when `--checkpoint-every` is greater than zero
+
+`bsvae-sweep-k` writes:
+
+- `results/<name>/sweep_k/sweep_results.csv`
+- `results/<name>/sweep_k/sweep_summary.json`
+- `results/<name>/sweep_k/k<K>/rep_<rep>/...` per-sweep run outputs
+- `results/<name>/final_k<K>/...` for the final retrained model when `--train-final` is enabled
+
+## Simulation and Benchmarking
+
+Generate one synthetic dataset:
+
+```bash
+bsvae-simulate generate \
+  --output data/sim_expr.csv \
+  --save-ground-truth data/sim_truth.csv
+```
+
+Generate a scenario grid:
+
+```bash
+bsvae-simulate init-config --output sim.yaml
+
+bsvae-simulate generate-grid \
+  --config sim.yaml \
+  --outdir results/sim_pub_v1 \
+  --reps 30 \
+  --base-seed 13
+```
+
+Validate the generated grid:
+
+```bash
+bsvae-simulate validate-grid --grid-dir results/sim_pub_v1
+```
+
+Each scenario replicate writes method-ready files under `results/sim_pub_v1/scenarios/<scenario_id>/rep_<rep>/`, including:
+
+- `expr/features_x_samples.tsv.gz`
+- `expr/samples_x_features.tsv.gz`
+- `truth/modules_hard.csv`
+- `method_inputs.json`
+
+## Documentation
+
+See the docs for the full workflow and command reference:
+
+- `docs/index.md`
+- `docs/quickstart.md`
+- `docs/tutorial.md`
+- `docs/cli.md`
+- `docs/networks.md`
+- `docs/hyperparameters.md`
+
+## License
+
+This project is licensed under the [GNU General Public License v3.0](LICENSE).

{bsvae-0.3.1 → bsvae-0.3.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "bsvae"
-version = "0.3.1"  # dynamically resolved by poetry-dynamic-versioning
+version = "0.3.2"  # dynamically resolved by poetry-dynamic-versioning
 description = "Gaussian Mixture VAE for multi-modal biological module discovery in omics data."
 readme = "README.md"
 authors = ["Kynon J Benjamin <kj.benjamin90@gmail.com>"]

{bsvae-0.3.1 → bsvae-0.3.2}/src/bsvae/cli/networks.py

@@ -64,6 +64,12 @@ def build_parser() -> argparse.ArgumentParser:
     mp.add_argument("--use-leiden", action="store_true",
                     help="Also run Leiden clustering as fallback comparison.")
     mp.add_argument("--leiden-resolution", type=float, default=1.0)
+    mp.add_argument("--tx2gene", type=str, default=None,
+                    help="TSV (transcript_id, gene_id) for isoform-to-gene aggregation.")
+    mp.add_argument("--aggregate-to-gene", action="store_true", default=False,
+                    help="Aggregate isoform-level γ to gene level by averaging across "
+                         "isoforms per gene. Saves gamma_gene.npz and "
+                         "hard_assignments_gene.npz. Requires --tx2gene.")
 
     # ------------------------------------------------------------------
     # export-latents
@@ -120,6 +126,35 @@ def handle_extract_networks(args, logger: logging.Logger) -> None:
     logger.info("Network extraction complete → %s", args.output_dir)
 
 
+def _aggregate_gamma_to_gene(
+    gamma: "np.ndarray",
+    feature_ids: List[str],
+    tx2gene_path: str,
+) -> "tuple[np.ndarray, list[str]]":
+    """Average isoform-level γ (F_tx, K) to gene level, renormalise rows to sum 1."""
+    import numpy as np
+    from collections import defaultdict
+    from bsvae.utils.hierarchy import load_tx2gene
+
+    tx2gene = load_tx2gene(path=tx2gene_path, feature_ids=feature_ids)
+    # Map every feature to its gene_id; fall back to itself for unmapped features
+    fid_to_gene = {fid: tx2gene.get(fid, fid) for fid in feature_ids}
+
+    gene_to_idxs: dict = defaultdict(list)
+    for i, fid in enumerate(feature_ids):
+        gene_to_idxs[fid_to_gene[fid]].append(i)
+
+    gene_ids = list(gene_to_idxs.keys())
+    gamma_gene = np.zeros((len(gene_ids), gamma.shape[1]), dtype=np.float32)
+    for j, gid in enumerate(gene_ids):
+        gamma_gene[j] = gamma[gene_to_idxs[gid]].mean(axis=0)
+
+    row_sums = gamma_gene.sum(axis=1, keepdims=True)
+    row_sums = np.where(row_sums > 0, row_sums, 1.0)
+    gamma_gene /= row_sums
+    return gamma_gene, gene_ids
+
+
 def handle_extract_modules(args, logger: logging.Logger) -> None:
     import numpy as np
     from bsvae.utils.modelIO import load_model
@@ -144,6 +179,30 @@ def handle_extract_modules(args, logger: logging.Logger) -> None:
         model.n_modules, len(feature_ids),
     )
 
+    if args.aggregate_to_gene:
+        if not args.tx2gene:
+            logger.warning("--aggregate-to-gene requires --tx2gene; skipping.")
+        else:
+            gamma_gene, gene_ids = _aggregate_gamma_to_gene(
+                result["gamma"], result["feature_ids"], args.tx2gene
+            )
+            hard_gene = gamma_gene.argmax(axis=1).astype(np.int32)
+            out = Path(args.output_dir)
+            np.savez_compressed(
+                out / "gamma_gene.npz",
+                gamma=gamma_gene,
+                feature_ids=np.array(gene_ids, dtype=object),
+            )
+            np.savez_compressed(
+                out / "hard_assignments_gene.npz",
+                hard_assignments=hard_gene,
+                feature_ids=np.array(gene_ids, dtype=object),
+            )
+            logger.info(
+                "Gene-level aggregation: %d transcripts → %d genes",
+                len(result["feature_ids"]), len(gene_ids),
+            )
+
     if args.soft_eigengenes and args.expr:
         import pandas as pd
         expr = pd.read_csv(args.expr, index_col=0, sep="\t" if args.expr.endswith(".tsv") else ",")

{bsvae-0.3.1 → bsvae-0.3.2}/src/bsvae/cli/sweep_k.py

@@ -320,7 +320,7 @@ def parse_args(cli_args=None):
     p.add_argument("--kl-cycle-length", type=int, default=50)
     p.add_argument("--sep-strength", type=float, default=0.1)
     p.add_argument("--sep-alpha", type=float, default=2.0)
-    p.add_argument("--bal-strength", type=float, default=0.01)
+    p.add_argument("--bal-strength", type=float, default=0.1)
     p.add_argument("--bal-ema-blend", type=float, default=0.5)
     p.add_argument("--pi-entropy-strength", type=float, default=0.0)
     p.add_argument("--hier-strength", type=float, default=0.0)

{bsvae-0.3.1 → bsvae-0.3.2}/src/bsvae/cli/train.py

@@ -85,11 +85,16 @@ def parse_args(cli_args=None):
     p.add_argument("--use-batch-norm", action="store_true", dest="use_batch_norm", default=True)
     p.add_argument("--sigma-min", type=float, default=0.3,
                    help="GMM component σ floor (default: 0.3).")
+    p.add_argument("--normalize-input", action="store_true", default=False,
+                   help="Z-score each feature profile across samples before encoding. "
+                        "Recommended for count-derived data (log2-CPM, etc.).")
 
     # --- Loss hyperparameters ---
     p.add_argument("--beta", type=float, default=1.0, help="KL weight (default: 1.0).")
-    p.add_argument("--free-bits", type=float, default=0.5,
-                   help="Per-dim KL lower bound in nats (default: 0.5).")
+    p.add_argument("--free-bits", type=float, default=0.0,
+                   help="Per-dim KL lower bound in nats (default: 0.0). "
+                        "Non-zero values hurt GMM cluster recovery by flooring γ to uniform; "
+                        "leave at 0.0 for normal use.")
     p.add_argument("--kl-warmup-epochs", type=int, default=0,
                    help="Epochs for β linear ramp within Phase 2 (default: 0).")
     p.add_argument("--kl-anneal-mode", choices=["linear", "cyclical"], default="linear")
@@ -98,16 +103,36 @@ def parse_args(cli_args=None):
                    help="λ_sep: σ-scaled separation loss weight (default: 0.1).")
     p.add_argument("--sep-alpha", type=float, default=2.0,
                    help="α: margin multiplier for separation loss (default: 2.0).")
-    p.add_argument("--bal-strength", type=float, default=0.01,
-                   help="λ_bal: γ-usage balance loss weight (default: 0.01).")
+    p.add_argument("--bal-strength", type=float, default=0.1,
+                   help="λ_bal: γ-usage balance loss weight (default: 0.1). "
+                        "Higher values force uniform usage; lower values allow natural "
+                        "module-size variation but risk dead components.")
     p.add_argument("--bal-ema-blend", type=float, default=0.5,
                    help="Balance loss blend α between EMA and batch (default: 0.5).")
+    p.add_argument("--masked-recon", action="store_true", default=False,
+                   help="Compute reconstruction loss only on non-zero input entries. "
+                        "Recommended for zero-inflated count data.")
     p.add_argument("--pi-entropy-strength", type=float, default=0.0,
                    help="λ_pi: KL(uniform || π) penalty weight (default: 0.0).")
     p.add_argument("--hier-strength", type=float, default=0.0,
                    help="λ_hier: hierarchical isoform loss weight (default: 0.0=disabled).")
+    p.add_argument("--corr-strength", type=float, default=0.0,
+                   help="λ_corr: reconstruction-space pairwise correlation loss weight (default: 0.0=disabled). "
+                        "Penalises MSE between the batch-level Pearson correlation matrices of inputs "
+                        "and reconstructions. Recommended: 0.05–0.2 for NB count data with normalize-input.")
+    p.add_argument("--latent-corr-strength", type=float, default=0.0,
+                   help="λ_lcorr: latent-space correlation alignment loss weight (default: 0.0=disabled). "
+                        "Aligns cosine similarities of encoder μ with input Pearson correlations. "
+                        "More direct than --corr-strength; recommended: 0.1–1.0 for NB count data.")
     p.add_argument("--tx2gene", type=str, default=None,
                    help="TSV with (transcript_id, gene_id) for hierarchical loss.")
+    p.add_argument("--isoform-stratified", action="store_true", default=False,
+                   help="Use IsoformStratifiedSampler so isoforms of the same gene "
+                        "co-occur in batches, making --hier-strength effective. "
+                        "Requires --tx2gene.")
+    p.add_argument("--p-multi", type=float, default=0.5,
+                   help="Probability of drawing a full batch from multi-isoform features "
+                        "(default: 0.5). Only active with --isoform-stratified.")
 
     # --- Evaluation ---
     p.add_argument("--no-eval", action="store_true", default=False,
@@ -145,12 +170,53 @@ def main(args):
 
     # --- Hierarchy ---
     gene_groups = {}
-    if args.hier_strength > 0 and args.tx2gene:
+    if args.tx2gene and (args.hier_strength > 0 or args.isoform_stratified):
         from bsvae.utils.hierarchy import load_tx2gene, group_isoforms_by_gene
         tx2gene = load_tx2gene(path=args.tx2gene, feature_ids=feature_ids)
         gene_groups = group_isoforms_by_gene(tx2gene, feature_ids)
         logger.info("Hierarchical groups: %d multi-isoform genes", len(gene_groups))
 
+    # Validate gene_groups indices are in-bounds
+    if gene_groups:
+        n_features_total = len(feature_ids)
+        for gene_id, idxs in gene_groups.items():
+            bad = [i for i in idxs if i < 0 or i >= n_features_total]
+            if bad:
+                raise ValueError(
+                    f"gene_groups['{gene_id}'] contains out-of-bounds indices {bad} "
+                    f"(dataset has {n_features_total} features)"
+                )
+
+    # Replace DataLoader with IsoformStratifiedSampler if requested
+    if args.isoform_stratified and gene_groups:
+        import torch.utils.data as tud
+        from bsvae.utils.hierarchy import IsoformStratifiedSampler
+        multi_indices = sorted({idx for idxs in gene_groups.values() for idx in idxs})
+        sampler = IsoformStratifiedSampler(
+            n_features=len(feature_ids),
+            multi_isoform_indices=multi_indices,
+            batch_size=args.batch_size,
+            p_multi=args.p_multi,
+            seed=args.seed,
+        )
+        train_loader = tud.DataLoader(
+            train_loader.dataset,
+            batch_size=args.batch_size,
+            sampler=sampler,
+            drop_last=True,
+            pin_memory=train_loader.pin_memory,
+            num_workers=train_loader.num_workers,
+        )
+        logger.info(
+            "IsoformStratifiedSampler: %d multi-isoform features, p_multi=%.2f",
+            len(multi_indices), args.p_multi,
+        )
+    elif args.isoform_stratified and not gene_groups:
+        logger.warning(
+            "--isoform-stratified requested but no multi-isoform gene groups found; "
+            "using standard random sampler."
+        )
+
     # --- Model ---
     model = GMMModuleVAE(
         n_features=n_features,
@@ -160,6 +226,7 @@ def main(args):
         dropout=args.dropout,
         use_batch_norm=args.use_batch_norm,
         sigma_min=args.sigma_min,
+        normalize_input=args.normalize_input,
     ).to(device)
     logger.info(
         "Model: %d features, %d latent dims, %d modules | params=%d",
@@ -179,6 +246,10 @@ def main(args):
         pi_entropy_strength=args.pi_entropy_strength,
         bal_ema_blend=args.bal_ema_blend,
         hier_strength=args.hier_strength,
+        normalize_input=args.normalize_input,
+        masked_recon=args.masked_recon,
+        corr_strength=args.corr_strength,
+        latent_corr_strength=args.latent_corr_strength,
     )
     warmup_loss = WarmupLoss(
         beta=args.beta,

{bsvae-0.3.1 → bsvae-0.3.2}/src/bsvae/models/gmm_prior.py

@@ -272,6 +272,11 @@ class GaussianMixturePrior(nn.Module):
         K = self.K
         X = mu_samples.cpu().numpy()
 
+        if len(X) < K:
+            raise ValueError(
+                f"kmeans_init_: need at least K={K} samples but got {len(X)}. "
+                "Increase the dataset size or reduce n_modules."
+            )
         km = KMeans(n_clusters=K, random_state=0, n_init=10)
         labels = km.fit_predict(X)
         centers = torch.tensor(km.cluster_centers_, dtype=torch.float32)