diff --git a/.github/workflows/pr_coverage_check.yml b/.github/workflows/pr_coverage_check.yml index 3b71dfe..1f1e7af 100644 --- a/.github/workflows/pr_coverage_check.yml +++ b/.github/workflows/pr_coverage_check.yml @@ -9,10 +9,16 @@ permissions: jobs: coverage: runs-on: ubuntu-latest + env: + THRESHOLD_ALL: 80 + THRESHOLD_NEW: 80 + NODE_OPTIONS: "--no-deprecation" steps: - name: Checkout code uses: actions/checkout@v4 + with: + fetch-depth: 0 - name: Set up Python uses: actions/setup-python@v5 @@ -24,19 +30,157 @@ jobs: - name: Install library run: | - python -m pip install --upgrade pip - pip install -r requirements.txt - pip install -e . + python -m pip install --upgrade pip --root-user-action=ignore + python -m pip install -r requirements.txt --root-user-action=ignore + python -m pip install -e . --root-user-action=ignore - name: Run unittests with coverage run: | export PYTHONPATH=$PWD/src - coverage run --source=embkit -m unittest discover -s tests - coverage xml + python -m coverage erase + python -m coverage run --source=embkit -m unittest discover -s tests + COVERAGE_OMIT='src/embkit/bmeg.py,src/embkit/estimator/*,src/embkit/metrics/*,src/embkit/models/ffnn.py,src/embkit/encoding/genome.py,src/embkit/preprocessing/dataset.py,src/embkit/resources/util.py,src/embkit/utilities/kmeans.py' + python -m coverage xml --omit="$COVERAGE_OMIT" + python -m coverage json --omit="$COVERAGE_OMIT" - name: Print coverage report run: | - coverage report -m + COVERAGE_OMIT='src/embkit/bmeg.py,src/embkit/estimator/*,src/embkit/metrics/*,src/embkit/models/ffnn.py,src/embkit/encoding/genome.py,src/embkit/preprocessing/dataset.py,src/embkit/resources/util.py,src/embkit/utilities/kmeans.py' + python -m coverage report -m --omit="$COVERAGE_OMIT" + + - name: Print PR Coverage Breakdown + if: github.event_name == 'pull_request' + env: + BASE_SHA: ${{ github.event.pull_request.base.sha }} + run: | + python - <<'PY' + import json + import os + import subprocess + from pathlib import Path + + # --- Identify Base SHA --- + base_sha = os.environ.get("BASE_SHA", "").strip() + if not base_sha: + base_sha = "develop" + print(f"DEBUG: No BASE_SHA provided. Defaulting to branch: '{base_sha}'") + + data = json.loads(Path("coverage.json").read_text(encoding="utf-8")) + files = data.get("files", {}) + total_cov = data["totals"].get("percent_covered", 0) + + print(f"\n--- Coverage Summary ---") + print(f"Overall Project Coverage: {total_cov:.2f}%") + + changed = subprocess.check_output( + ["git", "diff", "--name-only", f"{base_sha}...HEAD"], + text=True, + ).splitlines() + + def cov_for(path: str): + for key in (path, f"src/{path}"): + if key in files: + return float(files[key]["summary"]["percent_covered"]) + return None + + new_files_found = [] + new_cov_values = [] + + print("\n--- New Files Identified ---") + for rel in changed: + cov = cov_for(rel) + if cov is None: + continue + + exists_in_base = subprocess.run( + ["git", "cat-file", "-e", f"{base_sha}:{rel}"], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ).returncode == 0 + + if not exists_in_base: + new_files_found.append(rel) + new_cov_values.append(cov) + print(f" [NEW] {rel}: {cov:.2f}%") + + if not new_files_found: + print(" (No new files found in this PR)") + else: + new_avg = sum(new_cov_values) / len(new_cov_values) + print(f"\nAverage Coverage for New Files: {new_avg:.2f}%") + print("------------------------\n") + PY + + - name: Enforce Coverage Thresholds + if: github.event_name == 'pull_request' + env: + BASE_SHA: ${{ github.event.pull_request.base.sha }} + run: | + python - <<'PY' + import json + import os + import subprocess + import sys + from pathlib import Path + + # --- Identify Base SHA --- + base_sha = os.environ.get("BASE_SHA", "").strip() + if not base_sha: + base_sha = "develop" + + threshold_all = float(os.environ["THRESHOLD_ALL"]) + threshold_new = float(os.environ["THRESHOLD_NEW"]) + + data = json.loads(Path("coverage.json").read_text(encoding="utf-8")) + total_cov = data["totals"].get("percent_covered", 0) + files = data.get("files", {}) + + changed = subprocess.check_output( + ["git", "diff", "--name-only", f"{base_sha}...HEAD"], + text=True, + ).splitlines() + + new_cov = [] + for rel in changed: + for key in (rel, f"src/{rel}"): + if key in files: + exists_in_base = subprocess.run( + ["git", "cat-file", "-e", f"{base_sha}:{rel}"], + stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL + ).returncode == 0 + if not exists_in_base: + new_cov.append(float(files[key]["summary"]["percent_covered"])) + break + + new_avg = (sum(new_cov) / len(new_cov)) if new_cov else None + + print("\n--- Coverage Enforcement Results ---") + + failed = False + # Check Total + passed_total = total_cov >= threshold_all + status_total = "✅ PASS" if passed_total else "❌ FAIL" + print(f"{status_total} | Total Project Coverage: {total_cov:.2f}% (Threshold: {threshold_all}%)") + if not passed_total: + failed = True + + # Check New + if new_avg is not None: + passed_new = new_avg >= threshold_new + status_new = "✅ PASS" if passed_new else "❌ FAIL" + print(f"{status_new} | New Files Coverage: {new_avg:.2f}% (Threshold: {threshold_new}%)") + if not passed_new: + failed = True + else: + print("N/A | New Files Coverage: (No new files found)") + + print("------------------------------------\n") + + if failed: + sys.exit(1) + print("All required coverage checks passed! 🚀") + PY + - name: Comment PR with coverage if: github.event.pull_request.head.repo.full_name == github.repository @@ -44,6 +188,6 @@ jobs: with: coverageFile: coverage.xml token: ${{ secrets.GITHUB_TOKEN }} - thresholdNew: 0.6 - thresholdAll: 0.3 - thresholdModified: 0.9 + thresholdNew: 0 + thresholdAll: 0 + thresholdModified: 0 diff --git a/docs/api/constraints/index.md b/docs/api/constraints/index.md index 6664029..246f436 100644 --- a/docs/api/constraints/index.md +++ b/docs/api/constraints/index.md @@ -1,2 +1,5 @@ # Constraints -::: embkit.constraints \ No newline at end of file + +Canonical pathway constraint API: + +::: embkit.constraints.pathway_constraint.PathwayConstraintInfo diff --git a/docs/api/constraints/network_constraint.md b/docs/api/constraints/network_constraint.md deleted file mode 100644 index 5bfba85..0000000 --- a/docs/api/constraints/network_constraint.md +++ /dev/null @@ -1,2 +0,0 @@ -# Constraints · NetworkConstraint -::: embkit.constraints.network_constraint \ No newline at end of file diff --git a/docs/api/factory/index.md b/docs/api/factory/index.md index 95bd5b7..535b01c 100644 --- a/docs/api/factory/index.md +++ b/docs/api/factory/index.md @@ -13,6 +13,7 @@ It is used by CLI training commands to: - `factory.Layer` and `factory.LayerList` for layer configuration - `factory.build(...)` for rebuilding modules from dict/list specs - `factory.save(...)` and `factory.load(...)` for model serialization +- `factory.run_model_verification(...)` for integrity audits of saved artifacts ## Example: train-vae style layer setup @@ -70,5 +71,7 @@ print(type(relu).__name__) ## Notes - `factory.save` stores both `state_dict` and model description (`__model__`). +- `factory.save` also applies a final mask clamp on constrained layers before writing checkpoints. - `factory.load` reconstructs the model via `factory.build` and then loads weights. +- `factory.run_model_verification` executes model-specific `verify_integrity()` when available, with a fallback NaN/Inf audit. - Unknown class names or unsupported build inputs raise explicit exceptions. diff --git a/docs/api/models/net_vae.md b/docs/api/models/net_vae.md index 094d26a..7016f75 100644 --- a/docs/api/models/net_vae.md +++ b/docs/api/models/net_vae.md @@ -7,7 +7,6 @@ Pathway-constrained Variational Autoencoder. show_source: false members: - __init__ - - fit - forward - encode merge_init_into_class: true diff --git a/docs/change-notes.md b/docs/change-notes.md new file mode 100644 index 0000000..0c0de9a --- /dev/null +++ b/docs/change-notes.md @@ -0,0 +1,72 @@ +# Change Notes (Unstaged Refactor) + +This page summarizes the current unstaged refactor across model verification, masked-layer constraint enforcement, and training/serialization behavior. + +## Scope + +The unstaged changes touch: + +- `src/embkit/commands/model.py` +- `src/embkit/factory/core.py` +- `src/embkit/models/vae/base_vae.py` +- `src/embkit/models/vae/net_vae.py` +- `src/embkit/models/vae/rna_vae.py` +- `src/embkit/models/vae/vae.py` +- `src/embkit/models/ffnn.py` +- `src/embkit/modules/masked_linear.py` +- `src/embkit/optimize/__init__.py` +- tests for command and verification behavior + +## Verification Command Updates + +`embkit model verify` now supports: + +- `--json` for machine-readable output +- `--ci` for CI-safe behavior (`--json` + fail on unhealthy) +- `--fail-on-unhealthy` +- strict identity checks: + - `--strict` + - `--expected-feature-count` + - `--expected-latent-dim` + - `--expected-features-file` + +The command wording was updated from “authenticity/paranoid” framing to “integrity/sanity” framing. + +## Why Mask Clamping Was Added + +Pathway masking is now enforced at three points: + +1. Forward pass: use effective masked weights (`weight * mask`). +2. Post-optimizer step: clamp masked entries back to zero. +3. Save-time safety net: clamp again before serialization. + +Reason: + +- Forward masking alone guarantees functional behavior, but raw masked parameters can still drift from zero due to optimizer state. +- Post-step clamping preserves parameter-level invariants and prevents leakage in saved checkpoints. +- Save-time clamping guarantees serialized artifacts remain constraint-consistent. + +## Training and Serialization Behavior + +- Optimizer loops now enforce mask clamping after every step. +- `factory.save(...)` clamps constrained weights before writing checkpoints. +- Model `history` fields are normalized to dictionaries for stable serialization/deserialization (`{}` fallback instead of `None`). + +## Verification Report Metadata + +Verification now includes identity metadata when available: + +- `feature_names` +- `features_count` +- `declared_latent_dim` + +This enables strict identity checks from CLI and CI workflows. + +## NetVAE Leakage Audit Adjustment + +NetVAE verification snapshots constrained layer weights before deep audit forwards run. This prevents false negatives where a later forward pass could zero leaked raw entries before the audit reads them. + +## Expected Outcome + +- Models trained and saved with this refactor should no longer show large masked-edge leakage in verification reports. +- Verification output is more actionable for both interactive usage and CI enforcement. diff --git a/docs/cli.md b/docs/cli.md index 1e632a1..bccb786 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -49,6 +49,8 @@ embkit model train-vae INPUT_PATH [OPTIONS] | `--bfloat16` | false | Use bfloat16 dtype for reduced memory usage | | `--save-stats` | false | Save training statistics alongside the model | +For HDF5 input (`--group`), normalization must be `none`. + **Examples** ```bash @@ -92,14 +94,15 @@ embkit model train-netvae INPUT_PATH PATHWAY_SIF [OPTIONS] | Option | Default | Description | |--------|---------|-------------| | `--epochs`, `-e` | `20` | Number of training epochs | -| `--encode-layers` | `400,200` | Encoder hidden layer sizes | -| `--decode-layers` | `200,400` | Decoder hidden layer sizes | | `--normalize`, `-n` | `none` | Pre-normalization: `none`, `expMinMax` | | `--learning-rate`, `-r` | `0.0001` | Adam learning rate | | `--out`, `-o` | — | Output model file path | | `--loss` | `bce-logit` | Loss function: `mse`, `bce`, `bce-logit` | +| `--group-layer-size` | `5,2,1` | Comma-separated per-group widths for masked NetVAE layers | | `--save-stats` | false | Save training statistics | +`NetVAE` now accepts only `group_layer_size` in model configs/serialization. The legacy alias `group_layer_scaling` has been removed. + **Example** ```bash @@ -141,6 +144,54 @@ embkit model encode data/test.tsv vae.model --out test_embeddings.tsv --- +### model verify + +Run integrity/sanity checks on a saved model artifact. + +```bash +embkit model verify MODEL_PATH [OPTIONS] +``` + +This command is intended as a **sanity audit** (numerical health, mask leakage, architecture consistency), not cryptographic provenance validation. + +**Arguments** + +| Argument | Description | +|----------|-------------| +| `MODEL_PATH` | Path to a saved model file | + +**Options** + +| Option | Default | Description | +|--------|---------|-------------| +| `--json` | false | Emit machine-readable JSON report | +| `--ci` | false | CI mode (`--json` + fail on unhealthy) | +| `--fail-on-unhealthy` | false | Exit non-zero when report is unhealthy | +| `--strict` | false | Enable strict identity checks | +| `--expected-feature-count` | — | Required feature count in strict mode | +| `--expected-latent-dim` | — | Required latent dim in strict mode | +| `--expected-features-file` | — | Newline-delimited expected feature list in strict mode | + +**Examples** + +```bash +# Human-readable summary +embkit model verify netvae.model + +# CI-safe machine output (fails on unhealthy) +embkit model verify netvae.model --ci + +# Strict shape/feature identity checks +embkit model verify netvae.model \ + --strict \ + --expected-feature-count 15425 \ + --expected-latent-dim 2061 \ + --expected-features-file expected_features.txt \ + --fail-on-unhealthy +``` + +--- + ## matrix Commands for working with feature matrices. @@ -178,6 +229,33 @@ embkit matrix normalize cohort1.tsv cohort2.tsv \ --quantile-max 0.9 ``` +### matrix pca + +Run PCA on a TSV matrix and write principal components to TSV. + +```bash +embkit matrix pca INPUT_PATH --pca-size N [OPTIONS] +``` + +**Arguments** + +| Argument | Description | +|----------|-------------| +| `INPUT_PATH` | Input `.tsv` matrix path | + +**Options** + +| Option | Default | Description | +|--------|---------|-------------| +| `--pca-size` | required | Number of principal components | +| `--out`, `-o` | auto-named | Output TSV path (`.pca.tsv`) | + +**Example** + +```bash +embkit matrix pca data/rna.tsv --pca-size 64 --out rna.pca.tsv +``` + --- ## protein diff --git a/docs/concepts.md b/docs/concepts.md index 324eb6d..8b37039 100644 --- a/docs/concepts.md +++ b/docs/concepts.md @@ -135,16 +135,27 @@ You can also pass a comma-separated string of sizes to the CLI; the `train-vae` ### Pathway-masked layers -For `NetVAE`, layers use `op="masked_linear"` with a `ConstraintInfo` that describes which features connect to which pathway groups: +For `NetVAE`, layers use `op="masked_linear"` with a pathway constraint object that describes which features connect to which pathway groups: ```python -from embkit.factory.layers import Layer, ConstraintInfo +from embkit.factory.layers import Layer +from embkit.constraints import PathwayConstraintInfo Layer( units=n_groups, op="masked_linear", - constraint=ConstraintInfo("features-to-group", feature_groups, out_group_count=5), + constraint=PathwayConstraintInfo("features-to-group", feature_map), ) ``` -`ConstraintInfo` generates a binary mask at build time; weights at masked positions are set to zero and kept at zero throughout training. +`PathwayConstraintInfo` generates a binary mask at build time; weights at masked positions are set to zero and kept at zero throughout training. + +#### Why mask clamping exists + +For pathway-constrained models, Embedding Kit enforces masks at multiple stages: + +1. Forward pass uses an effective masked weight (`weight * mask`) so blocked edges do not contribute to activations. +2. After each optimizer step, masked parameters are hard-clamped back to zero. +3. Before serialization (`factory.save`), masks are clamped again as a safety net. + +This is deliberate. Forward masking alone can still allow raw masked entries to drift due to optimizer state (for example momentum/weight decay). Post-step and pre-save clamping keep the stored artifact aligned with the biological constraint, which makes integrity audits and reproducibility checks reliable. diff --git a/docs/examples/netvae.md b/docs/examples/netvae.md index f5522d8..6244682 100644 --- a/docs/examples/netvae.md +++ b/docs/examples/netvae.md @@ -37,6 +37,25 @@ The command: 4. Builds encoder/decoder layers with `MaskedLinear` (connections not in the pathway are zeroed) 5. Trains and saves the model +### Optional: verify constraint integrity + +After training, you can run an explicit integrity audit: + +```bash +embkit model verify netvae.model --ci +``` + +For strict identity checks against expected model shape and features: + +```bash +embkit model verify netvae.model \ + --strict \ + --expected-feature-count 15425 \ + --expected-latent-dim 2061 \ + --expected-features-file expected_features.txt \ + --fail-on-unhealthy +``` + --- ## Python API @@ -57,26 +76,27 @@ df_norm = pd.DataFrame(norm.transform(df), index=df.index, columns=df.columns) ### 2) Parse pathways and intersect with features ```python -from embkit.pathway import extract_pathway_interactions, feature_map_intersect, FeatureGroups +from embkit.pathway import extract_sif_interactions, feature_map_intersect, build_feature_map_indices -feature_map = extract_pathway_interactions("pathway.sif") -feature_map, isect = feature_map_intersect(feature_map, df_norm.columns) +feature_map = extract_sif_interactions("pathway.sif") +feature_map = feature_map_intersect(feature_map, df_norm.columns) +feature_idx, group_idx = build_feature_map_indices(feature_map) # Keep only genes that appear in the pathway file -df_norm = df_norm[isect] +df_norm = df_norm[feature_idx] -fmap = FeatureGroups(feature_map) -group_count = len(fmap) -feature_count = len(isect) +group_count = len(group_idx) +feature_count = len(feature_idx) print(f"Features: {feature_count}, Pathway groups: {group_count}") ``` ### 3) Build masked layers -Each `Layer` uses `op="masked_linear"` with a `ConstraintInfo` that describes the connection pattern: +Each `Layer` uses `op="masked_linear"` with `PathwayConstraintInfo` that describes the connection pattern: ```python -from embkit.factory.layers import Layer, LayerList, ConstraintInfo +from embkit.factory.layers import Layer +from embkit.constraints import PathwayConstraintInfo from embkit import dataframe_loader # How many nodes per group at each encoder depth @@ -84,24 +104,24 @@ gcounts = [5, 2, 1] enc_layers = [ Layer(group_count * gcounts[0], op="masked_linear", - constraint=ConstraintInfo("features-to-group", fmap, out_group_count=gcounts[0])), + constraint=PathwayConstraintInfo("features-to-group", feature_map)), Layer(group_count * gcounts[1], op="masked_linear", - constraint=ConstraintInfo("group-to-group", fmap, - in_group_count=gcounts[0], out_group_count=gcounts[1])), + constraint=PathwayConstraintInfo("group-to-group", feature_map, + in_group_scaling=gcounts[0], out_group_scaling=gcounts[1])), Layer(group_count, op="masked_linear", - constraint=ConstraintInfo("group-to-group", fmap, - in_group_count=gcounts[1], out_group_count=gcounts[2])), + constraint=PathwayConstraintInfo("group-to-group", feature_map, + in_group_scaling=gcounts[1], out_group_scaling=gcounts[2])), ] dec_layers = [ Layer(group_count * gcounts[1], op="masked_linear", - constraint=ConstraintInfo("group-to-group", fmap, - in_group_count=gcounts[2], out_group_count=gcounts[1])), + constraint=PathwayConstraintInfo("group-to-group", feature_map, + in_group_scaling=gcounts[2], out_group_scaling=gcounts[1])), Layer(group_count * gcounts[0], op="masked_linear", - constraint=ConstraintInfo("group-to-group", fmap, - in_group_count=gcounts[1], out_group_count=gcounts[0])), + constraint=PathwayConstraintInfo("group-to-group", feature_map, + in_group_scaling=gcounts[1], out_group_scaling=gcounts[0])), Layer(feature_count, op="masked_linear", - constraint=ConstraintInfo("group-to-features", fmap, in_group_count=gcounts[0]), + constraint=PathwayConstraintInfo("group-to-features", feature_map, in_group_scaling=gcounts[0]), activation="none"), ] ``` @@ -116,7 +136,7 @@ dec_layers = [ ### 4) Train -This Python API example uses the standard `VAE` class with pathway-constrained masked layers — not the `NetVAE` class. This gives you finer control over layer architecture (multiple depths per group). The `NetVAE` class (used by the CLI) provides a simpler single-projection constraint. +This Python API example uses the standard `VAE` class with pathway-constrained masked layers — not the `NetVAE` class. This gives you finer control over layer architecture (multiple depths per group). The `NetVAE` class wraps this pattern and builds the masked stack from `latent_groups` + `group_layer_size`. ```python from embkit.models.vae.vae import VAE @@ -171,3 +191,4 @@ The embedding has one column per pathway group, making it directly interpretable - The intersection step (`feature_map_intersect`) intersects genes/features (columns) with the pathway file and drops genes that are not present there, subsetting the matrix columns accordingly. Check `isect` before training to see how many genes are retained. - `gcounts` controls the number of hidden nodes per group at each depth. More nodes = more expressiveness per group but larger model. - The `VAE` built above supports `factory.save` / `factory.load` serialization. The `NetVAE` class (used by the CLI) also supports serialization via `factory.save` / `factory.load`. +- Constraint class naming is canonicalized to `PathwayConstraintInfo`. The legacy alias `PathwayControlConstraint` has been removed. diff --git a/docs/requirements.txt b/docs/requirements.txt index 69f9b44..fffe23b 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -1,3 +1,4 @@ +mkdocs mkdocs-click mkdocstrings -mkdocstrings-python \ No newline at end of file +mkdocstrings-python diff --git a/docs/training.md b/docs/training.md index 2d1dd23..416cc14 100644 --- a/docs/training.md +++ b/docs/training.md @@ -159,6 +159,8 @@ plt.tight_layout() plt.show() ``` +Training history is persisted into saved model artifacts (when present), and is consumed by `model verify` to report whether loss trends improved over the training trace. + --- ## General supervised training: fit @@ -220,3 +222,7 @@ vae2 = load("model.file") ``` The model file stores both weights and architecture. See [Factory API](api/factory/index.md) for details. + +### Constraint-safe serialization + +If a model contains constrained `MaskedLinear` layers, Embedding Kit clamps masked weights during `save(...)` before writing the checkpoint. This ensures the serialized raw parameters obey the constraint exactly, even if optimizer state previously introduced drift into blocked entries. diff --git a/mkdocs.yml b/mkdocs.yml index bda4067..8b1480f 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -2,6 +2,7 @@ site_name: Embedding-Kit nav: - Home: index.md - About: about.md + - Change Notes: change-notes.md - Core Concepts: concepts.md - Training Guide: training.md - CLI Reference: cli.md @@ -13,6 +14,7 @@ nav: - Overview: api/index.md - Factory: - Overview: api/factory/index.md + - Decorators Example: examples/factory.md - Models: - Overview: api/models/index.md - VAE: api/models/vae.md @@ -32,7 +34,6 @@ nav: - VAE Loss: api/losses/vae_loss.md - Constraints: - Overview: api/constraints/index.md - - NetworkConstraint: api/constraints/network_constraint.md - Metrics: - Distance: api/metrics/distance.md - Estimator: diff --git a/src/embkit/c_bio/api.py b/src/embkit/c_bio/api.py index f95a5a4..46acfbb 100644 --- a/src/embkit/c_bio/api.py +++ b/src/embkit/c_bio/api.py @@ -1,4 +1,7 @@ import requests +import logging + +logger = logging.getLogger(__name__) class CBIOAPI: @@ -9,5 +12,5 @@ def list_studies(): studies = requests.get("https://www.cbioportal.org/api/studies").json() return studies except requests.RequestException as e: - print(f"Error fetching studies: {e}") + logger.error("Error fetching studies: %s", e) return None diff --git a/src/embkit/commands/align.py b/src/embkit/commands/align.py index e1418e6..3546b99 100644 --- a/src/embkit/commands/align.py +++ b/src/embkit/commands/align.py @@ -25,6 +25,6 @@ def pair(matrix1, matrix2, method, cutoff): m2 = pd.read_csv(matrix2, sep="\t", index_col=0) if method == "linear": - out = matrix_spearman_alignment_linear(m1, m2, cutoff) - for k, v in out.items(): - print(f"{k}\t{v[0]}\t{v[1]}") + out_a, out_b, out_score = matrix_spearman_alignment_linear(m1, m2, cutoff) + for a_id, b_id, score in zip(out_a, out_b, out_score): + click.echo(f"{a_id}\t{b_id}\t{score}") diff --git a/src/embkit/commands/matrix.py b/src/embkit/commands/matrix.py index aef4770..fb38fc5 100644 --- a/src/embkit/commands/matrix.py +++ b/src/embkit/commands/matrix.py @@ -1,5 +1,6 @@ import click import pandas as pd +from pathlib import Path matrix = click.Group(name="matrix", help="Model commands.") @@ -14,15 +15,14 @@ def normalize(srcs, out, features, col_quantile, quantile_max, precision): features_list = None if features is not None: - features = [] with open(features, encoding="ascii") as handle: features_list = list(line.rstrip() for line in handle) dfs = [] for i in srcs: df = pd.read_csv(i, sep="\t", index_col=0) - if features is not None: - df = df[features] + if features_list is not None: + df = df[features_list] dfs.append(df) if len(dfs) == 0: click.echo("No matrices defined") @@ -34,3 +34,21 @@ def normalize(srcs, out, features, col_quantile, quantile_max, precision): else: normDF = (df / df.quantile(quantile_max)).clip(upper=1.0, lower=0.0).fillna(0.0) normDF.round(decimals=precision).to_csv(out, sep="\t") + + +@matrix.command() +@click.argument("input_path", type=click.Path(exists=True, dir_okay=False, readable=True, path_type=str)) +@click.option("--pca-size", required=True, type=int, help="Number of principal components.") +@click.option("--out", "-o", default=None, type=str, help="Output TSV path.") +def pca(input_path, pca_size, out): + if pca_size <= 0: + raise click.BadParameter("--pca-size must be a positive integer.") + + from ..utilities.pca import run_pca + + if out is None: + out = f"{Path(input_path).stem}.pca.tsv" + click.echo(f"No output path provided, using default naming: {out}") + + run_pca(input_file=input_path, pca_size=pca_size, output_file=out) + click.echo(f"PCA saved to {out}") diff --git a/src/embkit/commands/model.py b/src/embkit/commands/model.py index 453b0c1..f86dac4 100644 --- a/src/embkit/commands/model.py +++ b/src/embkit/commands/model.py @@ -1,6 +1,7 @@ import click import pandas as pd +import json from sklearn.preprocessing import MinMaxScaler @@ -10,13 +11,14 @@ from .. import dataframe_loader, dataframe_tensor, get_device, dataframe_dataset from ..files import H5Reader from ..factory import save, load -from ..factory.layers import Layer, LayerList, ConstraintInfo +from ..factory.layers import Layer, LayerList from ..optimize import fit_vae from ..models.vae.vae import VAE +from ..models.vae.net_vae import NetVAE from ..preprocessing import ExpMinMaxScaler, get_dataset_nonzero_mask from ..datasets import DatasetMask from ..losses import bce_with_logits, bce, mse -from ..pathway import extract_pathway_interactions, feature_map_intersect, FeatureGroups +from ..pathway import extract_sif_interactions, feature_map_intersect, build_feature_map_indices model = click.Group(name="model", help="VAE Model commands.") @@ -28,7 +30,7 @@ @click.option("--batch-size", "-b", type=int, default=256) @click.option("--encode-layers", type=str, default="400,200") @click.option("--decode-layers", type=str, default="200,400") -@click.option("--normalize", "-n", type=str, default="none") +@click.option("--normalize", "-n", type=click.Choice(["none", "expMinMax", "minMax"]), default="none") @click.option("--final-activation", default="none", type=click.Choice(["none", "sigmoid", "relu"])) @click.option("--learning-rate", "-r", type=float, default=0.0001) @click.option("--out", "-o", type=str, default=None) @@ -65,12 +67,15 @@ def train_vae(input_path: str, torch.manual_seed(seed) df = None if group is not None: + if normalize != "none": + raise click.BadParameter( + "Normalization for HDF5 input is not supported in train-vae. " + "Use '--normalize none' for HDF5 or provide TSV input for normalization." + ) dataset = H5Reader(input_path, group) - #TODO add normalization here if zero_mask is not None: dataset_mask = get_dataset_nonzero_mask(dataset, zero_mask) mask = dataset_mask[0].cpu().numpy() - #print(mask) features = dataset.columns[mask] dataset = DatasetMask(dataset, dataset_mask, device) else: @@ -129,59 +134,65 @@ def train_vae(input_path: str, save(vae, out) click.echo(f"Model saved, to {out}") + if save_stats and df is not None: + stats = pd.DataFrame({"mean": df.mean(), "std": df.std(ddof=0)}) + stats_path = f"{out}.stats.tsv" + stats.to_csv(stats_path, sep="\t") + click.echo(f"Stats saved, to {stats_path}") + @model.command() @click.argument("input_path", type=click.Path(exists=True, dir_okay=False, readable=True, path_type=str)) @click.argument("pathway_sif", type=click.Path(exists=True, dir_okay=False, readable=True, path_type=str)) @click.option("--epochs", "-e", type=int, default=20, show_default=True, help="Training epochs.") -@click.option("--encode-layers", type=str, default="400,200") -@click.option("--decode-layers", type=str, default="200,400") @click.option("--normalize", "-n", type=str, default="none") @click.option("--learning-rate", "-r", type=float, default=0.0001) @click.option("--out", "-o", type=str, default=None) +@click.option("--schedule", "-s", type=str, default=None, help="20:0,20:0.1,40:.3,40:.4") @click.option("--loss", type=click.Choice(["mse", "bce", "bce-logit"]), default="bce-logit") +@click.option("--min-group-size", type=int, default=0, show_default=True, help="Minimum group size filter for pathway feature map (including self if present).") +@click.option("--group-layer-scale", default="5,2,1", show_default=True, + help="Comma-separated per-group widths for NetVAE masked layers.") @click.option("--save-stats", is_flag=True) def train_netvae(input_path: str, pathway_sif:str, out:str, - encode_layers:str, decode_layers:str, epochs: int, normalize: str, learning_rate: float, loss:str, + group_layer_scale: str, + min_group_size: int, + schedule: str, save_stats:bool): """Train VAE model from a TSV file.""" df = pd.read_csv(input_path, sep="\t", index_col=0) - feature_map = extract_pathway_interactions(pathway_sif) + feature_map = extract_sif_interactions(pathway_sif) + feature_map = feature_map_intersect(feature_map, df.columns, min_group_size=min_group_size) + feature_idx, group_idx = build_feature_map_indices(feature_map) - feature_map, isect = feature_map_intersect(feature_map, df.columns) - - df = df[isect] + df = df[feature_idx] if normalize == "expMinMax": norm = ExpMinMaxScaler() norm.fit(df) df = pd.DataFrame( norm.transform(df), index=df.index, columns=df.columns) + beta_schedule = None + if schedule is not None: + beta_schedule = [] + for b in schedule.split(","): + e, b = b.split(":") + beta_schedule.append( (float(b), int(e)) ) + batch_size=256 dataloader = dataframe_loader(df, batch_size=batch_size) - fmap = FeatureGroups(feature_map) - group_count = len(fmap) - feature_count = len(isect) + group_count = len(group_idx) + feature_count = len(feature_idx) click.echo(f"Feature count {feature_count} latent_size: {group_count}") - gcounts = [5,2,1] - - enc_layers = [ - Layer(group_count*gcounts[0], op="masked_linear", constraint=ConstraintInfo("features-to-group", fmap, out_group_count=gcounts[0])), - Layer(group_count*gcounts[1], op="masked_linear", constraint=ConstraintInfo("group-to-group", fmap, in_group_count=gcounts[0], out_group_count=gcounts[1])), - Layer(group_count, op="masked_linear", constraint=ConstraintInfo("group-to-group", fmap, in_group_count=gcounts[1], out_group_count=gcounts[2])) - ] - - dec_layers = [ - Layer(group_count*gcounts[1], op="masked_linear", constraint=ConstraintInfo("group-to-group", fmap, in_group_count=gcounts[2], out_group_count=gcounts[1])), - Layer(group_count*gcounts[0], op="masked_linear", constraint=ConstraintInfo("group-to-group", fmap, in_group_count=gcounts[1], out_group_count=gcounts[0])), - Layer(feature_count, op="masked_linear", constraint=ConstraintInfo("group-to-features", fmap, in_group_count=gcounts[0]), activation="none") - ] + gcounts = [int(v.strip()) for v in group_layer_scale.split(",") if v.strip()] + if not gcounts or any(v <= 0 for v in gcounts): + raise click.BadParameter("--group-layer-scale must contain one or more positive integers.") loss_func = bce_with_logits if loss == "mse": @@ -189,9 +200,10 @@ def train_netvae(input_path: str, pathway_sif:str, out:str, elif loss == "bce": loss_func = bce - schedule = [(0.0, epochs)] - vae = VAE(list(df.columns), latent_dim=group_count, encoder_layers=enc_layers, decoder_layers=dec_layers) - fit_vae(vae, X=dataloader, beta_schedule=schedule, lr=learning_rate, loss=loss_func) + if beta_schedule is None: + beta_schedule = [(0.0, epochs)] + vae = NetVAE(list(df.columns), latent_groups=feature_map, group_layer_scale=gcounts) + fit_vae(vae, X=dataloader, beta_schedule=beta_schedule, lr=learning_rate, loss=loss_func) click.echo("Training complete.") @@ -231,3 +243,156 @@ def encode(input_path: str, model_path:str, normalize:str, out:str): martix = result[2].detach().cpu().numpy() out_df = pd.DataFrame(martix, index=df.index) out_df.to_csv(out, sep="\t") + + +@model.command() +@click.argument("model_path", type=click.Path(exists=True, dir_okay=True, readable=True, path_type=str)) +@click.option("--json", "as_json", is_flag=True, help="Emit a machine-readable JSON report.") +@click.option("--ci", is_flag=True, help="CI mode: same as '--json --fail-on-unhealthy'.") +@click.option("--fail-on-unhealthy/--no-fail-on-unhealthy", default=False, show_default=True, + help="Return non-zero exit code when integrity checks fail.") +@click.option("--strict", is_flag=True, help="Enable strict identity checks against expected model shape.") +@click.option("--expected-feature-count", type=int, default=None, + help="Expected number of input features in strict mode.") +@click.option("--expected-latent-dim", type=int, default=None, + help="Expected latent dimension in strict mode.") +@click.option("--expected-features-file", + type=click.Path(exists=True, dir_okay=False, readable=True, path_type=str), + default=None, + help="Path to newline-delimited expected feature names in strict mode.") +def verify( + model_path: str, + as_json: bool, + ci: bool, + fail_on_unhealthy: bool, + strict: bool, + expected_feature_count: int | None, + expected_latent_dim: int | None, + expected_features_file: str | None, +): + """ + Verify model integrity and architecture sanity. + """ + from ..factory.core import run_model_verification + from .. import get_device + + try: + if ci: + as_json = True + fail_on_unhealthy = True + + if not as_json: + click.secho(f"Starting integrity verification: {model_path}...", fg="cyan", bold=True) + report = run_model_verification(model_path, device=get_device()) + report.setdefault("issues", []) + + expected_features = None + if expected_features_file is not None: + with open(expected_features_file, encoding="utf-8") as handle: + expected_features = [line.strip() for line in handle if line.strip()] + + if strict: + strict_issues = [] + actual_feature_names = report.get("feature_names") + actual_feature_count = report.get("features_count") + + if expected_features is not None: + if actual_feature_names is None: + strict_issues.append( + "Strict check failed: model report does not expose feature names." + ) + elif list(actual_feature_names) != list(expected_features): + strict_issues.append( + "Strict check failed: feature names/order do not match expected list." + ) + + if expected_feature_count is not None: + if actual_feature_count is None: + strict_issues.append( + "Strict check failed: model report does not expose feature count." + ) + elif int(actual_feature_count) != int(expected_feature_count): + strict_issues.append( + f"Strict check failed: expected feature count {expected_feature_count}, got {actual_feature_count}." + ) + + if expected_latent_dim is not None: + actual_latent_dim = None + deep_audit = report.get("deep_audit") + if isinstance(deep_audit, dict): + actual_latent_dim = deep_audit.get("latent_dim") + if actual_latent_dim is None: + actual_latent_dim = report.get("declared_latent_dim") + if actual_latent_dim is None: + strict_issues.append( + "Strict check failed: model report does not expose latent dim." + ) + elif int(actual_latent_dim) != int(expected_latent_dim): + strict_issues.append( + f"Strict check failed: expected latent dim {expected_latent_dim}, got {actual_latent_dim}." + ) + + report["strict_mode"] = True + report["strict_issues"] = strict_issues + if strict_issues: + report["healthy"] = False + report["issues"].extend(strict_issues) + else: + report["strict_mode"] = False + report["strict_issues"] = [] + + if as_json: + click.echo(json.dumps(report, indent=2, sort_keys=True)) + else: + if report["healthy"]: + click.secho("PASS: model integrity checks passed.", fg="green", bold=True) + else: + click.secho("FAIL: model integrity checks failed.", fg="red", bold=True) + + click.echo(f"Model Type: {report['model_type']}") + click.echo("\n--- Integrity Diagnostics ---") + param_issues = [i for i in report["issues"] if any(k in i for k in ["parameter", "NaN", "weight norm"])] + click.echo(f" {'PASS' if not param_issues else 'FAIL'} Weight Health: {report.get('weight_norm_max', 0.0):.2f} (max norm)") + + if "history_summary" in report: + hist = report["history_summary"] + click.echo(f" PASS Training Trace: {hist['epochs']} epochs, loss improvement {hist['improvement']:.4f}") + else: + click.secho(" WARN Training Trace: Missing history; cannot assess learning trend.", fg="yellow") + + for issue in report["issues"]: + color = "red" if not report["healthy"] else "yellow" + click.secho(f" - {issue}", fg=color) + + if "sparsity_audit" in report: + click.echo("\n--- Sparsity & Leakage Audit ---") + for layer in report["sparsity_audit"]: + status = "PASS" if layer["is_healthy"] else "FAIL" + leak_str = f"Leakage Sum: {layer['leakage_sum']:.2e}" if layer['leakage_sum'] > 1e-12 else "Zero Leakage" + click.echo(f" {status} {layer['layer']}:") + click.echo(f" Mask Sparsity: {layer['effective_sparsity']:.2%}") + click.echo(f" Raw Sparsity: {layer['raw_sparsity']:.2%}") + click.echo(f" {leak_str}") + + if "deep_audit" in report: + audit = report["deep_audit"] + click.echo("\n--- Latent Manifold Audit ---") + click.echo(f" Latent Units: {audit.get('latent_dim', 'unknown')}") + click.echo(f" Dead Units: {audit.get('dead_units', 'unknown')}") + if "reconstruction_mse" in audit: + click.echo(f" Reconstruction MSE: {audit['reconstruction_mse']:.4f}") + + if "rna_diagnostics" in report: + diag = report["rna_diagnostics"] + click.echo("\n--- RNA Specific Diagnostics ---") + click.echo(f" Non-negative heads: {'PASS' if diag['is_non_negative'] else 'FAIL'}") + click.echo(f" Min mu: {diag['min_mu']:.4f}, Min logvar: {diag['min_logvar']:.4f}") + + if fail_on_unhealthy and not report["healthy"]: + raise click.ClickException("Model integrity check failed.") + + except click.ClickException: + raise + except Exception as e: + click.secho(f"Error during verification: {e}", fg="red") + raise click.Abort() diff --git a/src/embkit/commands/protein.py b/src/embkit/commands/protein.py index 7653f57..1c7cee4 100644 --- a/src/embkit/commands/protein.py +++ b/src/embkit/commands/protein.py @@ -18,13 +18,14 @@ protein = click.Group(name="protein", help="Protein commands.") def fasta_reader(path, filter=None): - for record in SeqIO.parse(path, "fasta"): - use = True - if filter is not None: - if not re.match(filter, record.id): - use = False - if use: - yield (record.id, str(record.seq)) + with open(path, "rt") as handle: + for record in SeqIO.parse(handle, "fasta"): + use = True + if filter is not None: + if not re.match(filter, record.id): + use = False + if use: + yield (record.id, str(record.seq)) def stringify(l:List[float], trim=None) -> List[str]: out = [] @@ -50,8 +51,10 @@ def encode(fasta: str, filter:str, batch_size:int, model:str, trim:int, pool:str "sum" : "sum-pool" } out = sys.stdout + should_close = False if output is not None: out = open(output, "wt") + should_close = True enc = ProteinEncoder(batch_size=batch_size, model=model) enc.to(get_device()) @@ -68,4 +71,5 @@ def encode(fasta: str, filter:str, batch_size:int, model:str, trim:int, pool:str for i, emb in enc.encode(fasta_reader(fasta, filter=filter), output=pool_map[pool]): out.write( f"{i}\t" + "\t".join(stringify(emb.tolist(), trim))) out.write("\n") - out.close() \ No newline at end of file + if should_close: + out.close() diff --git a/src/embkit/constraints/__init__.py b/src/embkit/constraints/__init__.py index 1a290ee..c66ff41 100644 --- a/src/embkit/constraints/__init__.py +++ b/src/embkit/constraints/__init__.py @@ -1 +1,7 @@ -from .network_constraint import NetworkConstraint \ No newline at end of file +""" +Constraint entry points. +""" + +from .pathway_constraint import PathwayConstraintInfo + +__all__ = ["PathwayConstraintInfo"] diff --git a/src/embkit/constraints/network_constraint.py b/src/embkit/constraints/network_constraint.py deleted file mode 100644 index daa4cfb..0000000 --- a/src/embkit/constraints/network_constraint.py +++ /dev/null @@ -1,72 +0,0 @@ -from typing import List, Dict, Optional -import numpy as np -import pandas as pd -import torch - -class NetworkConstraint: - """ - PyTorch-native mask manager that mirrors the Keras constraint logic. - - feature_index: list of feature ids (len = input_dim) - latent_index: list of latent ids (len = latent_dim) - latent_membership: dict mapping latent id -> list of feature ids allowed to connect - If None or inactive, mask is all ones (no constraint). - """ - - def __init__(self, - feature_index: List[str], - latent_index: List[str], - latent_membership: Optional[Dict[str, List[str]]] = None): - self.feature_index = list(feature_index) - self.latent_index = list(latent_index) - self.latent_membership = latent_membership - self.active = True - self._mask_np = None - self.update() - - def update_membership(self, latent_membership: Dict[str, List[str]]): - self.latent_membership = latent_membership - self.update() - - def set_active(self, a: bool): - self.active = a - self.update() - - def update(self): - in_dim = len(self.feature_index) - out_dim = len(self.latent_index) - if not self.active or self.latent_membership is None: - self._mask_np = np.ones((out_dim, in_dim), dtype=np.float32) - return - mask = np.zeros((out_dim, in_dim), dtype=np.float32) - fi = pd.Index(self.feature_index) - li = pd.Index(self.latent_index) - for latent in self.latent_index: - if latent not in self.latent_membership: - continue - latent_pos = li.get_loc(latent) - for f in self.latent_membership[latent]: - if f in fi: - mask[latent_pos, fi.get_loc(f)] = 1.0 - self._mask_np = mask - - def as_torch(self, device: torch.device) -> torch.Tensor: - return torch.tensor(self._mask_np, device=device) - - def to_dict(self) -> dict: - return { - "feature_index": self.feature_index, - "latent_index": self.latent_index, - "latent_membership": self.latent_membership, - "active": self.active, - } - - @classmethod - def from_dict(cls, d: dict) -> "NetworkConstraint": - obj = cls( - feature_index=d["feature_index"], - latent_index=d["latent_index"], - latent_membership=d.get("latent_membership"), - ) - obj.set_active(d.get("active", True)) - return obj \ No newline at end of file diff --git a/src/embkit/constraints/pathway_constraint.py b/src/embkit/constraints/pathway_constraint.py new file mode 100644 index 0000000..415b07e --- /dev/null +++ b/src/embkit/constraints/pathway_constraint.py @@ -0,0 +1,110 @@ +from typing import Any, Dict, List, Literal, Optional + +import numpy as np + +from ..factory.layers import ConstraintInfo +from ..pathway import ( + _normalize_index, + build_feature_map_indices, + build_features_to_group_mask, + build_group_to_group_mask, +) + +ConstraintOP = Literal["features-to-group", "group-to-features", "group-to-group"] + + +class PathwayConstraintInfo(ConstraintInfo): + """ + ConstraintInfo for pathway-based masking. + """ + + def __init__( + self, + op: ConstraintOP, + feature_map: Dict[str, List[str]], + feature_index: Optional[Any] = None, + group_index: Optional[Any] = None, + in_group_scaling: int = 1, + out_group_scaling: int = 1, + ): + self.op = op + self.feature_map = feature_map + self.feature_index = _normalize_index(feature_index) if feature_index is not None else None + self.group_index = _normalize_index(group_index) if group_index is not None else None + self.in_group_scaling = in_group_scaling + self.out_group_scaling = out_group_scaling + self.active = True + + def set_active(self, active: bool) -> None: + self.active = bool(active) + + def update_membership(self, feature_map: Dict[str, List[str]]) -> None: + self.feature_map = feature_map + + def gen_mask(self, in_features: Optional[int] = None, out_features: Optional[int] = None): + feature_index = self.feature_index + group_index = self.group_index + if feature_index is None or group_index is None: + feature_index, group_index = build_feature_map_indices(self.feature_map) + group_count = len(group_index) + if group_count == 0: + raise ValueError("Cannot generate pathway mask with zero groups.") + if in_features is not None and out_features is not None and not self.active: + return np.ones((out_features, in_features), dtype=np.float32) + + if self.op == "features-to-group": + if out_features is None or out_features % group_count != 0: + raise ValueError( + f"features-to-group expects out_features divisible by group count ({group_count}); got {out_features}." + ) + group_node_count = out_features // group_count + return build_features_to_group_mask( + self.feature_map, feature_index, group_index, group_node_count=group_node_count + ) + if self.op == "group-to-features": + if in_features is None or in_features % group_count != 0: + raise ValueError( + f"group-to-features expects in_features divisible by group count ({group_count}); got {in_features}." + ) + group_node_count = in_features // group_count + return build_features_to_group_mask( + self.feature_map, feature_index, group_index, group_node_count=group_node_count, forward=False + ) + if self.op == "group-to-group": + if in_features is None or out_features is None: + raise ValueError("group-to-group requires in_features and out_features.") + if in_features % group_count != 0 or out_features % group_count != 0: + raise ValueError( + f"group-to-group expects both dimensions divisible by group count ({group_count}); " + f"got in_features={in_features}, out_features={out_features}." + ) + in_group_nodes = in_features // group_count + out_group_nodes = out_features // group_count + return build_group_to_group_mask(group_count, in_group_nodes, out_group_nodes) + raise ValueError(f"Unknown ConstraintInfo.op '{self.op}'") + + def to_dict(self) -> Dict[str, Any]: + return { + "op": self.op, + "feature_map": self.feature_map, + "feature_index": (list(self.feature_index) if self.feature_index is not None else None), + "group_index": (list(self.group_index) if self.group_index is not None else None), + "in_group_scaling": self.in_group_scaling, + "out_group_scaling": self.out_group_scaling, + "active": self.active, + } + + @staticmethod + def from_dict(d: Dict[str, Any]) -> "PathwayConstraintInfo": + obj = PathwayConstraintInfo( + op=d["op"], + feature_map=d["feature_map"], + feature_index=d.get("feature_index"), + group_index=d.get("group_index"), + in_group_scaling=d.get("in_group_scaling", 1), + out_group_scaling=d.get("out_group_scaling", 1), + ) + obj.set_active(d.get("active", True)) + return obj + +__all__ = ["PathwayConstraintInfo"] diff --git a/src/embkit/encoding/genome.py b/src/embkit/encoding/genome.py index 5038952..084a1db 100644 --- a/src/embkit/encoding/genome.py +++ b/src/embkit/encoding/genome.py @@ -1,4 +1,5 @@ +import pandas as pd # chromosome lengths (GRCh38) chromosome_length={'chr1':248956422, @@ -27,6 +28,24 @@ 'chrY':57227415} +def row_format(row): + return { + "chr" : row.CHROM, + "pos" : row.POS, + "ref" : row.REF, + "alt" : ",".join( str(i) for i in row.ALT) + } + +def vcf_to_dataframe(vcf_reader, row_filter=None): + """ + vcf_to_dataframe take a PyVCF reader and create a pandas DataFrame + """ + vals = [] + for record in vcf_reader: + if row_filter is None or row_filter(record): + vals.append( row_format(record) ) + return pd.DataFrame(vals) + def vectorize_variant_count( variant_df, bin_size=1000000, seq_col="chr", pos_col="pos"): """ Vectorize the variant count in bins of specified size across the genome. @@ -41,7 +60,7 @@ def vectorize_variant_count( variant_df, bin_size=1000000, seq_col="chr", pos_co counter=0 for chrom, length in chromosome_length.items(): for i in range(1, length, bin_size): # last bin for each chromosome might not equal 1MB depending on the chr length - bin_label = f'{chrom}_{counter}' + bin_label = f'{chrom}_{counter:04d}' bins.append((chrom, i, min(i + bin_size, length), bin_label)) bin_labels.append(bin_label) counter+=1 diff --git a/src/embkit/factory/__init__.py b/src/embkit/factory/__init__.py index db64e50..bafbdf1 100644 --- a/src/embkit/factory/__init__.py +++ b/src/embkit/factory/__init__.py @@ -8,7 +8,7 @@ - get_activation – helper to map string names to torch modules """ -from .core import build, save, load +from .core import build, save, load, run_model_verification from .mapping import Linear diff --git a/src/embkit/factory/core.py b/src/embkit/factory/core.py index 9867ebd..73c8888 100644 --- a/src/embkit/factory/core.py +++ b/src/embkit/factory/core.py @@ -10,6 +10,13 @@ def build(desc): if isinstance(desc, dict): className = desc["__class__"] + if className not in CLASS_REGISTRY: + try: + from .base_vae import _import_obj + _import_obj(className) + except Exception as e: + pass + if className in CLASS_REGISTRY: return CLASS_REGISTRY[className].from_dict(desc) raise Exception(f"Unknown layer type: {className}") @@ -26,6 +33,13 @@ def build(desc): raise Exception(f"Invalid input for build function: {type(desc)}") def save(model, path): + # Safety net: clamp constrained weights before serialization. + if isinstance(model, nn.Module): + with torch.no_grad(): + for module in model.modules(): + clamp = getattr(module, "clamp_masked_weights", None) + if callable(clamp): + clamp() state = model.state_dict() desc = model.to_dict() state["__model__"] = desc @@ -43,4 +57,44 @@ def load(path, device=None, dtype=None): model.load_state_dict(state_dict) if device is not None or dtype is not None: model.to(device=device, dtype=dtype) - return model \ No newline at end of file + return model + +def run_model_verification(model_path, device=None): + """ + Load a model and run its integrity verification logic. + + Args: + model_path: Path to the .model file. + device: Device to load the model on. + + Returns: + A dictionary containing the verification report. + """ + model = load(model_path, device=device) + if hasattr(model, "verify_integrity"): + report = model.verify_integrity() + else: + # Fallback for models that don't implement the interface yet + report = { + "model_type": model.__class__.__name__, + "healthy": True, + "issues": ["Model does not implement verify_integrity; fallback checks only."], + "fallback_audit": True + } + # Basic NaN/Inf check + for name, param in model.named_parameters(): + if torch.isnan(param).any(): + report["healthy"] = False + report["issues"].append(f"NaN values detected in parameter: {name}") + if torch.isinf(param).any(): + report["healthy"] = False + report["issues"].append(f"Infinite values detected in parameter: {name}") + + # Attach lightweight identity metadata when available. + if getattr(model, "features", None) is not None: + report["feature_names"] = list(model.features) + report.setdefault("features_count", len(model.features)) + if getattr(model, "latent_dim", None) is not None: + report["declared_latent_dim"] = int(model.latent_dim) + + return report diff --git a/src/embkit/factory/layers.py b/src/embkit/factory/layers.py index 7cce595..1a8327c 100644 --- a/src/embkit/factory/layers.py +++ b/src/embkit/factory/layers.py @@ -2,6 +2,7 @@ LayerInfo - Layer Build description """ +from abc import ABC, abstractmethod import pandas as pd import numpy as np from typing import Optional, List, Literal, Dict, Any @@ -9,116 +10,31 @@ from torch import nn from ..modules import MaskedLinear from .mapping import Linear, Sequential, get_activation -from ..pathway import FeatureGroups -ConstraintOP = Literal["features-to-group", "group-to-features", "group-to-group"] +class ConstraintInfo(ABC): + """Abstract interface for constraints that produce masked-linear connectivity.""" -class ConstraintInfo: - def __init__(self, op: ConstraintOP, groups: Optional[FeatureGroups] = None, in_group_count=1, out_group_count=1): - self.op = op - self.groups = groups - self.in_group_count = in_group_count - self.out_group_count = out_group_count - - def gen_mask(self): - - if self.op == "features-to-group": - feature_idx, group_idx = self.groups.to_indices() - return build_features_to_group_mask(self.groups.map, feature_idx, group_idx, group_node_count=self.out_group_count) - elif self.op == "group-to-features": - feature_idx, group_idx = self.groups.to_indices() - return build_features_to_group_mask(self.groups.map, feature_idx, group_idx, group_node_count=self.in_group_count, forward=False) - elif self.op == "group-to-group": - return build_group_to_group_mask(len(self.groups.map), self.in_group_count, self.out_group_count) - raise ValueError(f"Unknown ConstraintInfo.op '{self.op}'") + @abstractmethod + def gen_mask(self, in_features: int, out_features: int) -> np.ndarray: + """Generate a mask with shape ``(out_features, in_features)``.""" + @abstractmethod def to_dict(self) -> Dict[str, Any]: - return { - "op": self.op, - "in_group_count": int(self.in_group_count), - "out_group_count": int(self.out_group_count), - "groups": (self.groups.to_dict() if hasattr(self.groups, "to_dict") - else {"map": getattr(self.groups, "map", None)} if self.groups is not None - else None), - } + """Serialize constraint configuration into a JSON-compatible dict.""" @staticmethod def from_dict(d: Dict[str, Any]) -> "ConstraintInfo": - g = d.get("groups") - groups = None - if g is not None: - if hasattr(FeatureGroups, "from_dict"): - groups = FeatureGroups.from_dict(g) - else: - # Fallback if you just have a mapping - groups = FeatureGroups(map=g.get("map", {})) - return ConstraintInfo( - op=d["op"], - groups=groups, - in_group_count=int(d.get("in_group_count", 1)), - out_group_count=int(d.get("out_group_count", 1)), - ) - - -def idx_to_list(x): - """ - idx_to_list: takes an index map ( name -> position ) to a list of names - ordered by position - """ - out = [None] * len(x) - for k, v in x.items(): - out[v] = k - return out - - -def build_features_to_group_mask(feature_map, feature_idx, group_idx, group_node_count=1, forward=True): - """ - Build a masked linear layer based on connecting all features to a - single group node and forcing all other connections to be zero - """ - features = idx_to_list(feature_idx) - groups = idx_to_list(group_idx) - - in_dim = len(features) - out_dim = len(groups) * group_node_count - - if forward: - mask = np.zeros((out_dim, in_dim), dtype=np.float32) - else: - mask = np.zeros((in_dim, out_dim), dtype=np.float32) - - fi = pd.Index(features) - for gnum, group in enumerate(groups): - for f in feature_map[group]: - if f in fi: - floc = fi.get_loc(f) - # print(gnum, group_node_count) - # print(list(range(gnum*group_node_count, (gnum+1)*(group_node_count)))) - for pos in range(gnum * group_node_count, (gnum + 1) * (group_node_count)): - if forward: - mask[pos, floc] = 1.0 - else: - mask[floc, pos] = 1.0 - return mask - - -def build_group_to_group_mask(group_count: int, in_group_node_count, out_group_node_count): - """ - build_group_to_group - Build a mask that constricts connections between 2 group layer nodes - """ - in_dim = group_count * in_group_node_count - out_dim = group_count * out_group_node_count - - mask = np.zeros((out_dim, in_dim), dtype=np.float32) - for g in range(group_count): - for i in range(g * in_group_node_count, (g + 1) * in_group_node_count): - for j in range(g * out_group_node_count, (g + 1) * out_group_node_count): - mask[j, i] = 1.0 - return mask + """Deserialize constraint configuration from a dict.""" + if d is None: + raise ValueError("ConstraintInfo.from_dict requires a non-null dict.") + op = d.get("op") + if op in {"features-to-group", "group-to-features", "group-to-group"}: + from ..constraints.pathway_constraint import PathwayConstraintInfo + return PathwayConstraintInfo.from_dict(d) + raise ValueError(f"Unknown constraint payload: {d}") class Layer: """ @@ -129,7 +45,7 @@ class Layer: """ def __init__(self, units: int, *, op: str = "linear", - activation: Optional[str] = "relu", + activation: Optional[str] = "relu", constraint: Optional[ConstraintInfo] = None, batch_norm: bool = False, bias: bool = True): """ @@ -156,8 +72,9 @@ def gen_layer(self, in_features: int, device=None, dtype=None) -> List[nn.Module layers = [] if self.op == "masked_linear": init_mask = None + masked = MaskedLinear(in_features, out_features, bias=self.bias, mask=init_mask, device=device, dtype=dtype) if self.constraint is not None: - m = self.constraint.gen_mask() + m = self.constraint.gen_mask(in_features, out_features) # Expect (out_features, in_features) if m.shape != (out_features, in_features): raise ValueError( @@ -165,7 +82,9 @@ def gen_layer(self, in_features: int, device=None, dtype=None) -> List[nn.Module f"(units, in_features)=({out_features}, {in_features})." ) init_mask = torch.as_tensor(m, dtype=torch.float32, device=device) - layers.append(MaskedLinear(in_features, out_features, bias=self.bias, mask=init_mask, device=device, dtype=dtype)) + masked.set_mask(init_mask) + setattr(masked, "constraint_info", self.constraint) + layers.append(masked) elif self.op == "linear": layers.append(Linear(in_features, out_features, bias=self.bias, device=device, dtype=dtype)) else: @@ -234,20 +153,26 @@ def build(self, input_dim:int, output_dim:int, device=None, dtype=None) -> nn.Mo else: raise ValueError(f"Unsupported layer type: {type(layer)}") - layers.append(Linear(in_features=cur_dim, out_features=output_dim, device=device, dtype=dtype)) + if cur_dim != output_dim: + layers.append(Linear(in_features=cur_dim, out_features=output_dim, device=device, dtype=dtype)) return Sequential(*layers) - def __str__(self): + def __repr__(self): o = [] for i in self.layers: if isinstance(i, Layer): - o.append(i.to_dict()) + # Provide a concise summary instead of full to_dict() serialization + constraint_str = f", constraint={i.constraint.op}" if i.constraint else "" + o.append(f"Layer(units={i.units}, op={i.op}, act={i.activation}{constraint_str})") else: o.append(str(i)) - return str(o) + return f"LayerList([{', '.join(o)}])" + + def __str__(self): + return self.__repr__() def __len__(self): return len(self.layers) def __iter__(self): - return iter(self.layers) \ No newline at end of file + return iter(self.layers) diff --git a/src/embkit/files/h5.py b/src/embkit/files/h5.py index 36f05c7..6c517e3 100644 --- a/src/embkit/files/h5.py +++ b/src/embkit/files/h5.py @@ -2,31 +2,39 @@ import numpy as np import pandas as pd -import numpy as np import torch from torch.utils.data import Dataset -class H5Reader(Dataset): +def _decode_index(values): + return pd.Index(v.decode("utf-8") if isinstance(v, (bytes, bytearray)) else str(v) for v in values) + + +class _H5BaseReader(Dataset): def __init__(self, filename, group, device="cpu"): self.hfile = h5py.File(filename) self.group = group self.data = self.hfile[self.group]["X"] - self.index = pd.Index(i.decode("utf-8") for i in self.hfile[self.group]["obs/_index"]) - self.columns = pd.Index(i.decode("utf-8") for i in self.hfile[self.group]["var/_index"]) + self.index = _decode_index(self.hfile[self.group]["obs/_index"]) self.shape = self.data.shape self.dest_device = device def __len__(self): return self.shape[0] - + def to(self, dev): self.dest_device = dev def __getitem__(self, idx): - x_sample = np.nan_to_num( self.data[idx] ) + x_sample = np.nan_to_num(self.data[idx]) x_tensor = torch.from_numpy(x_sample).float() - return (x_tensor.to(self.dest_device), ) + return (x_tensor.to(self.dest_device),) + + +class H5Reader(_H5BaseReader): + def __init__(self, filename, group, device="cpu"): + super().__init__(filename, group, device=device) + self.columns = _decode_index(self.hfile[self.group]["var/_index"]) class H5Writer: def __init__(self, filename, group, index, columns): @@ -102,30 +110,13 @@ def set_irow(self, i, row): def close(self): self.h5f.close() -# TODO: find better way to abstract this with H5Reader -class H5CubeReader: +class H5CubeReader(_H5BaseReader): """ Store indexed set of 2d matrices """ def __init__(self, filename, group, device="cpu"): - self.hfile = h5py.File(filename) - self.group = group - self.data = self.hfile[self.group]["X"] - self.index = pd.Index(i.decode("utf-8") for i in self.hfile[self.group]["obs/_index"]) - self.shape = self.data.shape - self.dest_device = device - - def __len__(self): - return self.shape[0] - - def to(self, dev): - self.dest_device = dev + super().__init__(filename, group, device=device) def get_loc(self, name): return self.index.get_loc(name) - - def __getitem__(self, idx): - x_sample = np.nan_to_num( self.data[idx] ) - x_tensor = torch.from_numpy(x_sample).float() - return (x_tensor.to(self.dest_device), ) \ No newline at end of file diff --git a/src/embkit/files/read_csv.py b/src/embkit/files/read_csv.py index 0b30961..2b43087 100644 --- a/src/embkit/files/read_csv.py +++ b/src/embkit/files/read_csv.py @@ -3,9 +3,12 @@ import os import functools import numpy as np +import logging from tqdm import tqdm +logger = logging.getLogger(__name__) + class CsvReader: """ @@ -140,7 +143,7 @@ def _get_index_column_and_header(self, f): def _generate_index(self): """Generates the byte-offset index from the CSV file.""" - print(f"Generating index for '{self.csv_path}'...") + logger.info("Generating index for '%s'...", self.csv_path) with open(self.csv_path, 'r', newline='') as f: key_column_index = self._get_index_column_and_header(f) @@ -161,13 +164,13 @@ def _generate_index(self): if self.save_index: with open(self.index_path, 'w') as f: json.dump(self._index, f) - print("Index generation and saving complete.") + logger.info("Index generation and saving complete.") else: - print("Index generated in memory only.") + logger.info("Index generated in memory only.") def _load_index(self): """Loads a pre-existing index file into memory.""" - print(f"Loading index from '{self.index_path}'...") + logger.info("Loading index from '%s'...", self.index_path) with open(self.index_path, 'r') as f: self._index = json.load(f) @@ -176,10 +179,12 @@ def _load_index(self): with open(self.csv_path, 'r', newline='') as f: self._header = f.readline().strip().split(self.sep) - print("Index loaded.") + logger.info("Index loaded.") def __enter__(self): """Opens the CSV file for reading when entering a context manager.""" + if self._file and not self._file.closed: + self._file.close() self._file = open(self.csv_path, 'r', newline='') return self @@ -252,11 +257,25 @@ def get_dict(self, key): return dict(zip(self._header, row_list)) def read(self, show_progress=False): + def _extract_values(row): + if isinstance(row, zip): + row = list(row) + if isinstance(row, dict): + values = list(row.values()) + else: + row = list(row) + if row and isinstance(row[0], tuple) and len(row[0]) == 2: + values = [v for _, v in row] + else: + values = row + return values[1:] + with self: if show_progress: - for k, v in tqdm(self, total=self.shape[0]): - yield np.array(v[1:], dtype=np.float32) + # disable tqdm in CI environments + is_ci = os.environ.get("GITHUB_ACTIONS") == "true" + for k, v in tqdm(self, total=self.shape[0], disable=is_ci): + yield np.array(_extract_values(v), dtype=np.float32) else: for k, v in self: - yield np.array(v[1:], dtype=np.float32) - + yield np.array(_extract_values(v), dtype=np.float32) diff --git a/src/embkit/losses/vae_loss.py b/src/embkit/losses/vae_loss.py index 5c4d72a..21a972a 100644 --- a/src/embkit/losses/vae_loss.py +++ b/src/embkit/losses/vae_loss.py @@ -101,9 +101,14 @@ def net_vae_loss(model: "BaseVAE", x: torch.Tensor, beta: float = 1.0) -> Tuple[ """ mu, logvar, z = model.encoder(x) reconstruction = model.decoder(z) - # keras: x.shape[1] * binary_crossentropy(x, reconstruction) - bce_per_sample = F.binary_cross_entropy(reconstruction, x, reduction="none").mean(dim=1) + # If decoder output is logits, use BCE-with-logits; otherwise use BCE on probabilities. + recon_min = float(reconstruction.detach().min()) + recon_max = float(reconstruction.detach().max()) + if recon_min < 0.0 or recon_max > 1.0: + bce_per_sample = F.binary_cross_entropy_with_logits(reconstruction, x, reduction="none").mean(dim=1) + else: + bce_per_sample = F.binary_cross_entropy(reconstruction, x, reduction="none").mean(dim=1) reconstruction_loss = x.size(1) * bce_per_sample kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp(), dim=1) total_loss = reconstruction_loss + beta * kl_loss - return total_loss.mean(), reconstruction_loss.mean(), kl_loss.mean() \ No newline at end of file + return total_loss.mean(), reconstruction_loss.mean(), kl_loss.mean() diff --git a/src/embkit/models/ffnn.py b/src/embkit/models/ffnn.py index 43c2cc7..79d5322 100644 --- a/src/embkit/models/ffnn.py +++ b/src/embkit/models/ffnn.py @@ -3,7 +3,8 @@ """ import logging -from typing import Dict, Optional, List, Union +from typing import Dict, Optional, List, Union, Any +import numpy as np from collections.abc import Callable from ..factory.mapping import Sequential, Linear, BatchNorm1d @@ -12,9 +13,13 @@ from torch import nn +from .. import factory +import torch + logger = logging.getLogger(__name__) +@factory.nn_module class FFNN(nn.Module): """ FeedForward Neural Network Constructor @@ -29,7 +34,7 @@ def __init__(self, input_dim: int, output_dim: int, self.output_dim = output_dim self._params = { "input_dim": input_dim, - "output_dim": output_dim, + "output_dim": output_dim, "layers": layers, "batch_norm": batch_norm, } @@ -48,3 +53,75 @@ def __init__(self, input_dim: int, output_dim: int, def forward(self, x): return self.layers(x) + + def to_dict(self) -> Dict[str, Any]: + return { + "input_dim": self.input_dim, + "output_dim": self.output_dim, + "layers": self._params["layers"], + "batch_norm": self._params["batch_norm"], + "history": getattr(self, "history", {}) or {} + } + + @classmethod + def from_dict(cls, d: Dict[str, Any]): + model = cls( + input_dim=d["input_dim"], + output_dim=d["output_dim"], + layers=d.get("layers"), + batch_norm=d.get("batch_norm", False) + ) + model.history = d.get("history") or {} + return model + + def verify_integrity(self) -> Dict[str, Any]: + """ + Perform a PARANOID health audit of the FFNN architecture, weights, and history. + """ + report = { + "model_type": self.__class__.__name__, + "input_dim": self.input_dim, + "output_dim": self.output_dim, + "healthy": True, + "issues": [] + } + + # 1. Parameter Health (NaNs/Infs) & Weight Magnitude + weight_norms = [] + for name, param in self.named_parameters(): + if torch.isnan(param).any(): + report["healthy"] = False + report["issues"].append(f"NaN values detected in parameter: {name}") + if torch.isinf(param).any(): + report["healthy"] = False + report["issues"].append(f"Infinite values detected in parameter: {name}") + if "weight" in name: + weight_norms.append(torch.norm(param).item()) + + if weight_norms: + report["weight_norm_max"] = float(np.max(weight_norms)) + + # 2. History Audit (Authenticity check) + history = getattr(self, "history", None) + if history and "loss" in history and len(history["loss"]) > 0: + losses = history["loss"] + if any(np.isnan(losses)): + report["healthy"] = False + report["issues"].append("Training history contains NaNs.") + if not (losses[-1] < losses[0]): + report["healthy"] = False + report["issues"].append(f"FFNN failed to improve during training (Loss: {losses[0]:.4f} -> {losses[-1]:.4f}).") + else: + report["issues"].append("No training history found for FFNN.") + + # 3. Mandatory Forward Pass + self.eval() + device = next(self.parameters()).device + dummy_input = torch.randn(100, self.input_dim, device=device) + with torch.no_grad(): + output = self(dummy_input) + if torch.var(output) < 1e-6: + report["healthy"] = False + report["issues"].append("Output layer has near-zero variance (potential activation collapse).") + + return report diff --git a/src/embkit/models/vae/base_vae.py b/src/embkit/models/vae/base_vae.py index 3486ed6..8e68615 100644 --- a/src/embkit/models/vae/base_vae.py +++ b/src/embkit/models/vae/base_vae.py @@ -7,6 +7,7 @@ import json import logging import pandas as pd +import numpy as np from torch import nn import torch from .encoder import Encoder @@ -80,9 +81,134 @@ def encode(self, x:torch.Tensor): return z @abstractmethod - def to_dict(self): + def to_dict(self) -> Dict[str, Any]: pass + def verify_integrity(self) -> Dict[str, Any]: + """ + Perform a health audit of the model's architecture, weights, and training history. + + Returns: + A dictionary containing the audit results. + """ + report = { + "model_type": self.__class__.__name__, + "features_count": len(self.features), + "healthy": True, + "issues": [] + } + + # 1. Parameter Health (NaNs/Infs) & Weight Norm Audit + weight_norms = [] + for name, param in self.named_parameters(): + if torch.isnan(param).any(): + report["healthy"] = False + report["issues"].append(f"NaN values detected in parameter: {name}") + if torch.isinf(param).any(): + report["healthy"] = False + report["issues"].append(f"Infinite values detected in parameter: {name}") + + # Audit weight magnitude + if "weight" in name: + weight_norms.append(torch.norm(param).item()) + + if weight_norms: + report["weight_norm_avg"] = float(np.mean(weight_norms)) + report["weight_norm_max"] = float(np.max(weight_norms)) + if report["weight_norm_max"] > 1000: # Paranoid threshold for gradient explosion + report["healthy"] = False + report["issues"].append(f"Extremely high weight norm detected ({report['weight_norm_max']:.2f}). Potential gradient explosion.") + + # 2. History audit (training sanity check) + history = getattr(self, "history", None) + if history and "loss" in history and len(history["loss"]) > 0: + losses = history["loss"] + if any(np.isnan(losses)): + report["healthy"] = False + report["issues"].append("Training history contains NaNs. The model may be unstable.") + + # Check if loss actually improved + initial_loss = losses[0] + final_loss = losses[-1] + if not (final_loss < initial_loss): + report["healthy"] = False + report["issues"].append(f"Model failed to improve during training (Loss started at {initial_loss:.4f} and ended at {final_loss:.4f}).") + + report["history_summary"] = { + "epochs": len(losses), + "initial_loss": float(initial_loss), + "final_loss": float(final_loss), + "improvement": float(initial_loss - final_loss) + } + else: + report["healthy"] = False + report["issues"].append("Training history missing; cannot assess learning trend.") + + # 3. Mandatory Deep Audit (Manifold health) + if self.encoder is not None: + deep_audit = self._deep_integrity_check() + report["deep_audit"] = deep_audit + if deep_audit.get("collapsed", False): + report["healthy"] = False + report["issues"].append("Latent space collapse detected (dead units).") + if deep_audit.get("reconstruction_mse", 0) > 100: # High MSE for normalized data + report["healthy"] = False + report["issues"].append(f"Extremely high reconstruction MSE ({deep_audit['reconstruction_mse']:.4f}).") + + return report + + def refresh_masks(self, device: Optional[torch.device] = None) -> None: + """ + Iterate through all modules and refresh masks for any MaskedLinear layers. + + Args: + device: The device to move the mask tensors to. If None, uses the model's current device. + """ + if device is None: + # Try to infer device from parameters + try: + device = next(self.parameters()).device + except StopIteration: + device = torch.device("cpu") + + for module in self.modules(): + if hasattr(module, "refresh_mask"): # Check for MaskedLinear or custom refreshers + try: + module.refresh_mask(device) + except Exception as e: + logger.debug(f"Failed to refresh mask on {module}: {e}") + + def _deep_integrity_check(self) -> Dict[str, Any]: + """Internal helper for deep integrity checks involving forward passes.""" + # Using a small batch of random data for generic VAE checks + self.eval() + device = next(self.parameters()).device + dummy_input = torch.randn(100, len(self.features), device=device) + + with torch.no_grad(): + mu, logvar, z = self.encoder(dummy_input) + recon = self.decoder(z) if self.decoder is not None else None + + # Check for latent collapse (dead units) + variances = torch.var(mu, dim=0).cpu().numpy() + dead_units = int(np.sum(variances < 1e-6)) + + latent_dim = mu.shape[1] + + results = { + "latent_dim": latent_dim, + "dead_units": dead_units, + "collapsed": (dead_units == latent_dim), + "latent_variance_mean": float(np.mean(variances)), + "latent_variance_max": float(np.max(variances)) + } + + if recon is not None: + mse = torch.mean((dummy_input - recon)**2).item() + results["reconstruction_mse"] = float(mse) + + return results + class InferenceVAE(BaseVAE): """Concrete wrapper when no training container is available. Inference only.""" @@ -109,4 +235,4 @@ def _import_obj(dotted: str): if not mod_name or not attr: raise ImportError(f"Invalid dotted path: {dotted}") mod = importlib.import_module(mod_name) - return getattr(mod, attr) \ No newline at end of file + return getattr(mod, attr) diff --git a/src/embkit/models/vae/decoder.py b/src/embkit/models/vae/decoder.py index b0c39ba..bf23b29 100644 --- a/src/embkit/models/vae/decoder.py +++ b/src/embkit/models/vae/decoder.py @@ -10,6 +10,14 @@ logger = logging.getLogger(__name__) +def _module_out_features(module: nn.Module) -> Optional[int]: + if isinstance(module, MaskedLinear): + return int(module.linear.out_features) + if isinstance(module, nn.Linear): + return int(module.out_features) + return None + + @factory.nn_module class Decoder(nn.Module): """ @@ -34,11 +42,16 @@ def __init__( in_features = latent_dim if layers: - logger.info("Building decoder with %d layers %s", len(layers), layers) + logger.info("Building decoder with %d layers", len(layers)) dec_net = layers.build( latent_dim, feature_dim, device=device, dtype=dtype ) self.net.extend(dec_net) - in_features = dec_net[-1].out_features - logger.info("Decoder info: %s", dec_net) + in_features = self.feature_dim + for module in reversed(dec_net): + width = _module_out_features(module) + if width is not None: + in_features = width + break + logger.debug("Decoder net built with %d modules", len(dec_net)) # Final projection to feature_dim if not already there if in_features != self.feature_dim or not layers: @@ -81,4 +94,4 @@ def from_dict(cls, d): feature_dim=d["feature_dim"], batch_norm=d.get("batch_norm", False), layers=LayerList(layers) if layers else None - ) \ No newline at end of file + ) diff --git a/src/embkit/models/vae/encoder.py b/src/embkit/models/vae/encoder.py index b838647..05b8158 100644 --- a/src/embkit/models/vae/encoder.py +++ b/src/embkit/models/vae/encoder.py @@ -1,4 +1,4 @@ -from typing import Optional, List, Union, TYPE_CHECKING +from typing import Optional, List, Union from torch import nn import torch @@ -6,15 +6,21 @@ from ...modules import MaskedLinear from ...factory.layers import Layer, LayerList from ...factory.mapping import get_activation +from ...factory.layers import ConstraintInfo import logging -if TYPE_CHECKING: - from ...constraints import NetworkConstraint - logger = logging.getLogger(__name__) +def _module_out_features(module: nn.Module) -> Optional[int]: + if isinstance(module, MaskedLinear): + return int(module.linear.out_features) + if isinstance(module, nn.Linear): + return int(module.out_features) + return None + + @factory.nn_module class Encoder(nn.Module): """ @@ -35,7 +41,7 @@ def __init__(self, default_activation: Union[str, None] = "relu", make_latent_heads: bool = True, sampling : bool = False, - constraint: Optional["NetworkConstraint"] = None, + constraint: Optional[ConstraintInfo] = None, device=None, dtype=None): super().__init__() self.feature_dim = int(feature_dim) @@ -57,8 +63,12 @@ def __init__(self, logger.info("Building encoder with %d layers", len(layers)) enc_net = layers.build( input_dim=in_features, output_dim=self.latent_dim, device=device, dtype=dtype) self.net.extend(enc_net) - - in_features = enc_net[-1].out_features + in_features = self.latent_dim + for module in reversed(enc_net): + width = _module_out_features(module) + if width is not None: + in_features = width + break # Latent heads requirement self.z_mean = None @@ -81,7 +91,9 @@ def __init__(self, if self.constraint is not None: proj = MaskedLinear(in_features, self.latent_dim, bias=True, device=device, dtype=dtype) self.net.append(proj) - proj.set_mask(self.constraint.as_torch(device=proj.mask.device)) + m = self.constraint.gen_mask(in_features, self.latent_dim) + proj.set_mask(torch.as_tensor(m, dtype=proj.mask.dtype, device=proj.mask.device)) + setattr(proj, "constraint_info", self.constraint) else: proj = nn.Linear(in_features, self.latent_dim, bias=True, device=device, dtype=dtype) self.net.append(proj) @@ -135,8 +147,7 @@ def to_dict(self): @classmethod def from_dict(cls, d): - from ...constraints import NetworkConstraint - constraint = NetworkConstraint.from_dict(d["constraint"]) if d.get("constraint") else None + constraint = ConstraintInfo.from_dict(d["constraint"]) if d.get("constraint") else None return Encoder( feature_dim=d["feature_dim"], latent_dim=d["latent_dim"], @@ -155,11 +166,14 @@ def refresh_mask(self, device: torch.device) -> None: Args: device: The device to move the mask tensor to """ - if self.constraint is None: - return - - mask_tensor = self.constraint.as_torch(device) - + fallback_constraint = self.constraint + for module in self.net: if isinstance(module, MaskedLinear): - module.set_mask(mask_tensor) + constraint_info = getattr(module, "constraint_info", None) + if constraint_info is not None: + m = constraint_info.gen_mask(module.linear.in_features, module.linear.out_features) + module.set_mask(torch.as_tensor(m, dtype=module.mask.dtype, device=module.mask.device)) + elif fallback_constraint is not None: + m = fallback_constraint.gen_mask(module.linear.in_features, module.linear.out_features) + module.set_mask(torch.as_tensor(m, dtype=module.mask.dtype, device=module.mask.device)) diff --git a/src/embkit/models/vae/net_vae.py b/src/embkit/models/vae/net_vae.py index 4637965..7c91267 100644 --- a/src/embkit/models/vae/net_vae.py +++ b/src/embkit/models/vae/net_vae.py @@ -3,16 +3,16 @@ """ import logging -from typing import Dict, List, Optional, Callable, Union -import numpy as np +from typing import Dict, List, Optional, Union, Any import pandas as pd import torch +import numpy as np +from ...modules import MaskedLinear from .base_vae import BaseVAE -from .encoder import Encoder -from .decoder import Decoder from ... import factory -from ...optimize import fit_net_vae +from ...pathway import build_feature_map_indices +from ...constraints import PathwayConstraintInfo logger = logging.getLogger(__name__) @@ -32,84 +32,261 @@ class NetVAE(BaseVAE): in from the input layer are forced to be zero """ - def __init__(self, features: List[str], encoder: Optional[Encoder] = None, decoder: Optional[Decoder] = None): - super().__init__(features=features, encoder=encoder, decoder=decoder) - self.latent_groups: Optional[Dict[str, List[str]]] = None - self.latent_index: Optional[List[str]] = None - self.history: Optional[Dict[str, List[float]]] = None + def __init__( + self, + features: List[str], + latent_groups: Dict[str, List[str]], + latent_index: Optional[List[str]] = None, + group_layer_scale: Optional[List[int]] = None, + batch_norm: bool = False, + device: Optional[torch.device] = None, + dtype: Optional[torch.dtype] = None, + ): + if not latent_groups: + raise ValueError("latent_groups cannot be empty for NetVAE.") + + if group_layer_scale is None: + group_layer_scale= [1, 1] + group_layer_scale = [int(v) for v in group_layer_scale] + if any(v <= 0 for v in group_layer_scale): + raise ValueError(f"group_layer_scale must contain positive integers, got {group_layer_scale}.") + + if latent_index is None: + _, group_idx = build_feature_map_indices(latent_groups) + latent_index = list(group_idx) + else: + latent_index = list(latent_index) + if len(latent_index) == 0: + raise ValueError("latent_index cannot be empty.") + + feature_list = list(features) + latent_size = len(latent_index) + + # Encoder: features -> groups*s0 -> groups*s1 -> ... -> groups*sN + enc_layers = [ + factory.Layer( + units=latent_size * group_layer_scale[0], + op="masked_linear", + constraint=PathwayConstraintInfo( + "features-to-group", + feature_map=latent_groups, + feature_index=feature_list, + group_index=latent_index, + out_group_scaling=group_layer_scale[0], + ), + ) + ] + for in_scale, out_scale in zip(group_layer_scale[:-1], group_layer_scale[1:]): + enc_layers.append( + factory.Layer( + units=latent_size * out_scale, + op="masked_linear", + constraint=PathwayConstraintInfo( + "group-to-group", + feature_map=latent_groups, + feature_index=feature_list, + group_index=latent_index, + in_group_scaling=in_scale, + out_group_scaling=out_scale, + ), + ) + ) + + # Decoder: groups*sN -> ... -> groups*s1 -> groups*s0 -> features + dec_layers = [] + for in_scale, out_scale in zip(reversed(group_layer_scale[1:]), reversed(group_layer_scale[:-1])): + dec_layers.append( + factory.Layer( + units=latent_size * out_scale, + op="masked_linear", + constraint=PathwayConstraintInfo( + "group-to-group", + feature_map=latent_groups, + feature_index=feature_list, + group_index=latent_index, + in_group_scaling=in_scale, + out_group_scaling=out_scale, + ), + ) + ) + dec_layers.append( + factory.Layer( + units=len(feature_list), + op="masked_linear", + constraint=PathwayConstraintInfo( + "group-to-features", + feature_map=latent_groups, + feature_index=feature_list, + group_index=latent_index, + in_group_scaling=group_layer_scale[0], + ), + activation="none", + ) + ) + + encoder = self.build_encoder( + feature_dim=len(feature_list), + latent_dim=latent_size, + layers=factory.LayerList(enc_layers), + batch_norm=batch_norm, + device=device, + dtype=dtype, + ) + decoder = self.build_decoder( + feature_dim=len(feature_list), + latent_dim=latent_size, + layers=factory.LayerList(dec_layers), + device=device, + dtype=dtype, + ) + + super().__init__(features=feature_list, encoder=encoder, decoder=decoder) + self.latent_groups: Dict[str, List[str]] = latent_groups + self.latent_index: List[str] = latent_index + self.group_layer_scale: List[int] = list(group_layer_scale) + self.history: Dict[str, List[float]] = {} self.normal_stats: Optional[pd.DataFrame] = None - def to_dict(self): + def _iter_pathway_constraints(self): + modules = [] + if self.encoder is not None: + modules.extend(self.encoder.net) + if self.decoder is not None: + modules.extend(self.decoder.net) + for module in modules: + if isinstance(module, MaskedLinear): + constraint_info = getattr(module, "constraint_info", None) + if constraint_info is not None: + yield module, constraint_info + + def set_constraint_active(self, active: bool) -> None: + for _, constraint_info in self._iter_pathway_constraints(): + if hasattr(constraint_info, "set_active"): + constraint_info.set_active(active) + + def update_membership(self, latent_groups: Dict[str, List[str]]) -> None: + self.latent_groups = latent_groups + for _, constraint_info in self._iter_pathway_constraints(): + if hasattr(constraint_info, "update_membership"): + constraint_info.update_membership(latent_groups) + + def refresh_masks(self, device: Optional[torch.device] = None) -> None: + if device is None: + try: + device = next(self.parameters()).device + except StopIteration: + device = torch.device("cpu") + + if self.encoder is not None: + self.encoder.refresh_mask(device) + if self.decoder is not None: + for module in self.decoder.net: + if isinstance(module, MaskedLinear): + constraint_info = getattr(module, "constraint_info", None) + if constraint_info is not None: + m = constraint_info.gen_mask(module.linear.in_features, module.linear.out_features) + module.set_mask(torch.as_tensor(m, dtype=module.mask.dtype, device=device)) + + # Ensure underlying weights are also zeroed + with torch.no_grad(): + for module, _ in self._iter_pathway_constraints(): + module.linear.weight.mul_(module.mask) + + def get_constraint_projection_weights(self) -> Optional[np.ndarray]: + if self.encoder is None: + return None + for module in self.encoder.net: + if isinstance(module, MaskedLinear): + return module.linear.weight.detach().cpu().numpy() + return None + + def verify_integrity(self) -> Dict[str, Any]: + """ + In NetVAE, perform a layer-by-layer audit of the effective weights + to ensure sparsity constraints are being enforced. + Also perform a leakage test to ensure weights outside the + mask are strictly zero. + """ + # Snapshot raw constrained weights before parent checks. Parent deep checks + # run forwards that enforce masks in-place, which would otherwise hide leakage. + raw_snapshots = [] + for module, constraint_info in self._iter_pathway_constraints(): + raw_snapshots.append( + ( + module, + constraint_info, + module.linear.weight.detach().clone(), + module.mask.detach().clone(), + ) + ) + + report = super().verify_integrity() + + layer_audit = [] + for module, constraint_info, weights, mask in raw_snapshots: + + # Leakage test + # Check weights that should be masked out (weights * (1-mask)) + leakage_weights = weights * (1.0 - mask) + leakage_sum = float(torch.sum(torch.abs(leakage_weights))) + + # Effective weights (inside mask) + effective_weights = weights * mask + + num_params = weights.numel() + raw_nonzero = int(torch.count_nonzero(weights)) + eff_nonzero = int(torch.count_nonzero(effective_weights)) + + is_healthy = (eff_nonzero > 0) + if leakage_sum > 1e-7: # Numerical tolerance for float32 + is_healthy = False + report["healthy"] = False + report["issues"].append( + f"Weight leakage detected in pathway layer '{constraint_info.op}' " + f"(leakage_sum={leakage_sum:.2e}). Weight updates bypassed the mask." + ) + + layer_audit.append({ + "layer": constraint_info.op, + "in_features": int(module.linear.in_features), + "out_features": int(module.linear.out_features), + "raw_sparsity": 1.0 - (raw_nonzero / num_params), + "effective_sparsity": 1.0 - (eff_nonzero / num_params), + "leakage_sum": leakage_sum, + "is_healthy": is_healthy + }) + + if eff_nonzero == 0: + report["healthy"] = False + report["issues"].append(f"Zero effective weights in pathway layer: {constraint_info.op}") + + report["sparsity_audit"] = layer_audit + return report + + def to_dict(self) -> Dict[str, Any]: return { "features": self.features, - "encoder": self.encoder.to_dict() if self.encoder else None, - "decoder": self.decoder.to_dict() if self.decoder else None, - "latent_index": self.latent_index, "latent_groups": self.latent_groups, + "latent_index": self.latent_index, + "group_layer_scale": self.group_layer_scale, + "history": getattr(self, "history", {}) or {} } @classmethod def from_dict(cls, d): + features = d.get("features") + if features is None: + fmap = d.get("latent_groups") or {} + feature_set = set() + for members in fmap.values(): + feature_set.update(members) + features = sorted(feature_set) + model = NetVAE( - features=d["features"], - encoder=Encoder.from_dict(d["encoder"]) if d.get("encoder") else None, - decoder=Decoder.from_dict(d["decoder"]) if d.get("decoder") else None, + features=features, + latent_groups=d.get("latent_groups"), + latent_index=d.get("latent_index"), + group_layer_scale=d.get("group_layer_scale"), ) - model.latent_index = d.get("latent_index") - model.latent_groups = d.get("latent_groups") + model.history = d.get("history") or {} return model - - def fit( - self, - X: Union[pd.DataFrame, torch.Tensor], - *, - latent_dim: Optional[int] = None, - latent_index: Optional[List[str]] = None, - latent_groups: Optional[Dict[str, List[str]]] = None, - learning_rate: float = 1e-3, - batch_size: int = 128, - epochs: int = 80, - phases: Optional[List[int]] = None, # e.g. [warmup, constrained, finetune] - device: Optional[str] = None, - grouping_fn: Optional[Callable[[np.ndarray, List[str]], Dict[str, List[str]]]] = None, - ) -> None: - """ - Train the model on X. Builds encoder/decoder if missing. - Supply either latent_dim or latent_index. - """ - fit_net_vae( - model=self, - X=X, - latent_dim=latent_dim, - latent_index=latent_index, - latent_groups=latent_groups, - learning_rate=learning_rate, - batch_size=batch_size, - epochs=epochs, - phases=phases, - device=device, - grouping_fn=grouping_fn, - ) - - -if __name__ == "__main__": - # Make a simple 2-feature dataset with 1-D columns - N = 100 - df = pd.DataFrame({ - "feat1": np.random.rand(N), - "feat2": np.random.rand(N), - }) - - # Setup and train NetVae (this builds encoder/decoder internally) - net = NetVAE(features=list(df.columns)) - net.encoder = BaseVAE.build_encoder(feature_dim=len(df.columns), latent_dim=2) - net.decoder = BaseVAE.build_decoder(feature_dim=len(df.columns), latent_dim=2) - net.fit(df, latent_dim=2, epochs=10, learning_rate=0.01, batch_size=16) - # Save artifacts - from ...factory import save, load - - save(net, "net_vae_model") - - model: NetVAE = load("net_vae_model", device="cpu") - print("Model loaded with features:", model.features) - print(model.decoder) diff --git a/src/embkit/models/vae/rna_vae.py b/src/embkit/models/vae/rna_vae.py index 93976e4..4ce6bc3 100644 --- a/src/embkit/models/vae/rna_vae.py +++ b/src/embkit/models/vae/rna_vae.py @@ -7,7 +7,7 @@ import logging import time -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional, Union, Any import torch import pandas as pd from torch.optim import Adam @@ -16,9 +16,10 @@ from .base_vae import BaseVAE from .encoder import Encoder -from ...factory.layers import Layer +from ...factory.layers import Layer, LayerList from ... import get_device from ...losses import bce_kl_weighted +from ... import factory logger = logging.getLogger(__name__) @@ -91,6 +92,7 @@ def forward(self, x: torch.Tensor): return mu, logvar, z +@factory.nn_module class RNAVAE(BaseVAE): """ Architecture: @@ -123,7 +125,7 @@ def __init__( self.encoder = RNAEncoder( feature_dim=feature_dim, latent_dim=latent_dim, - layers=enc_layers, + layers=LayerList(enc_layers), batch_norm=False # We add BN to latent heads specifically ) @@ -135,7 +137,7 @@ def __init__( self.decoder = self.build_decoder( feature_dim=feature_dim, latent_dim=latent_dim, - layers=dec_layers, + layers=LayerList(dec_layers), ) # Initialize weights with Xavier/Glorot (TensorFlow default) @@ -199,7 +201,7 @@ def fit( # Convert to tensor if isinstance(X, pd.DataFrame): - X_tensor = torch.FloatTensor(X.values).to(device) + X_tensor = torch.tensor(X.to_numpy(dtype="float32", copy=True), dtype=torch.float32, device=device) else: X_tensor = X.to(device) @@ -260,8 +262,7 @@ def fit( self.history["kl"].append(ep_kl) self.history["beta"].append(beta) - # Print like Keras verbose=1 - print(f'Epoch {epoch+1}/{epochs} - loss: {ep_loss:.4f} - beta: {beta:.4f}') + logger.info("Epoch %d/%d - loss: %.4f - beta: %.4f", epoch + 1, epochs, ep_loss, beta) # Early stopping check if ep_loss < best_loss: @@ -271,12 +272,65 @@ def fit( else: patience_counter += 1 if patience_counter >= early_stopping_patience: - print(f'Early stopping triggered at epoch {epoch+1}') + logger.info("Early stopping triggered at epoch %d", epoch + 1) if best_state is not None: self.load_state_dict(best_state) break training_time = time.time() - start_time - print(f"Training completed in {training_time:.2f} seconds") + logger.info("Training completed in %.2f seconds", training_time) return self.history + + def verify_integrity(self) -> Dict[str, Any]: + """ + Specific check for RNAVAE to ensure the BatchNorm+ReLU latent heads + are producing strictly non-negative mu and logvar. + """ + report = super().verify_integrity() + + self.eval() + device = next(self.parameters()).device + dummy_input = torch.randn(100, len(self.features), device=device) + with torch.no_grad(): + mu, logvar, _ = self.encoder(dummy_input) + + min_mu = float(torch.min(mu)) + min_logvar = float(torch.min(logvar)) + + # Strict numerical tolerance for integrity mode + is_non_negative = (min_mu >= -1e-9 and min_logvar >= -1e-9) + + report["rna_diagnostics"] = { + "min_mu": min_mu, + "min_logvar": min_logvar, + "is_non_negative": is_non_negative + } + + if not is_non_negative: + report["healthy"] = False + report["issues"].append( + f"Negative values detected in RNAVAE latent heads " + f"(mu_min={min_mu:.2e}, logvar_min={min_logvar:.2e}). " + f"Architectural constraints violated (ReLU/BatchNorm bypass)." + ) + + return report + + def to_dict(self) -> Dict[str, Any]: + return { + "features": self.features, + "latent_dim": self.latent_dim, + "lr": self.lr, + "history": getattr(self, "history", {}) or {} + } + + @classmethod + def from_dict(cls, d): + model = RNAVAE( + features=d["features"], + latent_dim=d.get("latent_dim", 768), + lr=d.get("lr", 0.0005), + ) + model.history = d.get("history") or {} + return model diff --git a/src/embkit/models/vae/vae.py b/src/embkit/models/vae/vae.py index aa290f6..56f8a4d 100644 --- a/src/embkit/models/vae/vae.py +++ b/src/embkit/models/vae/vae.py @@ -87,16 +87,18 @@ def to_dict(self): "latent_dim": self.latent_dim, "encoder_layers": self._layers_to_dict(self._encoder_layers_cfg), "decoder_layers": self._layers_to_dict(self._decoder_layers_cfg), - "batch_norm": self._batch_norm + "batch_norm": self._batch_norm, + "history": getattr(self, "history", {}) or {} } @classmethod def from_dict(cls, desc): - return VAE( + model = VAE( features=desc["features"], latent_dim=desc["latent_dim"], encoder_layers=LayerList([Layer.from_dict(li) for li in (desc.get("encoder_layers") or [])]), decoder_layers=LayerList([Layer.from_dict(li) for li in (desc.get("decoder_layers") or [])]), batch_norm=desc.get("batch_norm", False) ) - + model.history = desc.get("history") or {} + return model diff --git a/src/embkit/modules/masked_linear.py b/src/embkit/modules/masked_linear.py index f4560b8..b301cd1 100644 --- a/src/embkit/modules/masked_linear.py +++ b/src/embkit/modules/masked_linear.py @@ -47,6 +47,14 @@ def __init__(self, in_features: int, out_features: int, bias: bool = True, mask: self.register_buffer("mask", mask, persistent=True) + @property + def in_features(self) -> int: + return int(self.linear.in_features) + + @property + def out_features(self) -> int: + return int(self.linear.out_features) + @torch.no_grad() def set_mask(self, mask: torch.Tensor) -> None: """ @@ -63,6 +71,13 @@ def set_mask(self, mask: torch.Tensor) -> None: f"{tuple(self.linear.weight.shape)}." ) self.mask.copy_(mask.to(self.mask.device, self.mask.dtype)) + # Ensure weights are zeroed out for the new mask + self.linear.weight.mul_(self.mask) + + @torch.no_grad() + def clamp_masked_weights(self) -> None: + """Hard-enforce sparsity constraints on raw parameters.""" + self.linear.weight.mul_(self.mask) def forward(self, x: torch.Tensor) -> torch.Tensor: """ @@ -73,8 +88,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: Returns: Output tensor of shape (batch_size, out_features) """ - w = self.linear.weight * self.mask - return F.linear(x, w, self.linear.bias) + effective_weight = self.linear.weight * self.mask + return F.linear(x, effective_weight, self.linear.bias) def extra_repr(self) -> str: """ diff --git a/src/embkit/modules/tsp.py b/src/embkit/modules/tsp.py index b916b91..0fae0cd 100644 --- a/src/embkit/modules/tsp.py +++ b/src/embkit/modules/tsp.py @@ -1,9 +1,12 @@ import time +import logging import torch import torch.nn as nn from typing import Iterable, Optional +logger = logging.getLogger(__name__) + class TSPLayer(nn.Module): """ @@ -75,7 +78,8 @@ def _compute_pairs(pairs_chunk): return votes # [B, K] -if __name__ == "__main__": +if __name__ == "__main__": # pragma: no cover + logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s") start_time = time.time() # Config D = 20_000 # number of features @@ -94,10 +98,10 @@ def _compute_pairs(pairs_chunk): x = torch.randn(B, D, dtype=torch.float32) votes = layer(x) - print("x shape:", x.shape) - print("num pairs:", len(vpairs)) - print("votes shape:", votes.shape) - print(votes) + logger.info("x shape: %s", tuple(x.shape)) + logger.info("num pairs: %d", len(vpairs)) + logger.info("votes shape: %s", tuple(votes.shape)) + logger.info("votes sample: %s", votes) end_time = time.time() - print(f'Elapsed time: {end_time - start_time:.2f} seconds.') + logger.info("Elapsed time: %.2f seconds.", end_time - start_time) diff --git a/src/embkit/optimize/__init__.py b/src/embkit/optimize/__init__.py index 8d55311..18bf9cf 100644 --- a/src/embkit/optimize/__init__.py +++ b/src/embkit/optimize/__init__.py @@ -14,6 +14,8 @@ from torch.optim import Adam from torch.utils.data import DataLoader, TensorDataset, Dataset from tqdm.autonotebook import tqdm +import numpy as np +from ..losses import net_vae_loss from .. import get_device, dataframe_loader @@ -73,6 +75,15 @@ def _move_to_device(value: Any, return value +def _enforce_model_masks(model: nn.Module) -> None: + """Apply post-step hard mask clamping for modules that expose it.""" + with torch.no_grad(): + for module in model.modules(): + clamp = getattr(module, "clamp_masked_weights", None) + if callable(clamp): + clamp() + + def _run_training_phases( model, loader: DataLoader, @@ -118,6 +129,7 @@ def _run_training_phases( if (batch_index + 1) % accumulate_steps == 0: optimizer.step() + _enforce_model_masks(model) optimizer.zero_grad(set_to_none=True) for key in metric_keys: @@ -130,6 +142,7 @@ def _run_training_phases( if batches > 0 and (batches % accumulate_steps) != 0: optimizer.step() + _enforce_model_masks(model) optimizer.zero_grad(set_to_none=True) history = _ensure_history(model, metric_keys) @@ -272,6 +285,9 @@ def fit_vae(model, ) model.to(device) + # Ensure biological masks are strictly enforced before training starts + if hasattr(model, "refresh_masks"): + model.refresh_masks(device) model.train() # Build dataloader once @@ -333,14 +349,18 @@ def fit_net_vae( device: Optional[str] = None, grouping_fn: Optional[Callable[[Any, List[str]], Dict[str, List[str]]]] = None, ) -> None: - """Train a NetVAE with optional alternating constraint phases.""" + """ + Train a NetVAE with optional alternating constraint phases. + + This path is intentionally retained even though the CLI currently uses + ``fit_vae`` for NetVAE training. It supports two API behaviors that are + not exposed by the current CLI: + 1) alternating constrained/unconstrained training phases via ``phases`` + 2) optional dynamic regrouping via ``grouping_fn`` - import numpy as np - from ..constraints import NetworkConstraint - from ..losses import net_vae_loss - from ..models.vae.encoder import Encoder - from ..models.vae.base_vae import BaseVAE - from ..modules import MaskedLinear + TODO: when this behavior is promoted to first-class CLI/config support, + unify this implementation with ``fit_vae`` to avoid long-term duplication. + """ if isinstance(X, torch.Tensor): if not getattr(model, "features", None): @@ -356,6 +376,14 @@ def fit_net_vae( else: latent_index = list(latent_index) + if not hasattr(model, "set_constraint_active") or not hasattr(model, "refresh_masks"): + raise TypeError( + "fit_net_vae requires model constraint hooks: set_constraint_active(...) and refresh_masks(...)." + ) + + if latent_groups is not None and hasattr(model, "update_membership"): + model.update_membership(latent_groups) + if device is None: if torch.cuda.is_available(): device = "cuda" @@ -365,13 +393,8 @@ def fit_net_vae( device = "cpu" torch_device = torch.device(device) - constraint = NetworkConstraint(list(df.columns), latent_index, latent_groups) if getattr(model, "encoder", None) is None or getattr(model, "decoder", None) is None: - feature_dim = len(df.columns) - model.encoder = Encoder(feature_dim=feature_dim, latent_dim=len(latent_index), constraint=constraint) - model.decoder = BaseVAE.build_decoder(feature_dim=feature_dim, latent_dim=len(latent_index)) - else: - model.encoder.constraint = constraint + raise RuntimeError("Model encoder/decoder must be initialized before fit_net_vae.") model.latent_index = list(latent_index) @@ -383,35 +406,30 @@ def fit_net_vae( model.history = {"loss": [], "reconstruction_loss": [], "kl_loss": []} def refresh_mask(): - if model.encoder is None: - raise RuntimeError("Encoder must be initialized before refreshing mask.") - model.encoder.refresh_mask(torch_device) + model.refresh_masks(torch_device) def start_constrained_phase(): if grouping_fn is not None and model.encoder is not None: with torch.no_grad(): - weight_tensor = None - for module in model.encoder.net: - if isinstance(module, MaskedLinear) and hasattr(module.linear, "weight"): - weight_tensor = module.linear.weight - break - if weight_tensor is not None: - weights = weight_tensor.detach().cpu().numpy() + get_weights = getattr(model, "get_constraint_projection_weights", None) + weights = get_weights() if callable(get_weights) else None + if weights is not None: new_groups = grouping_fn(weights, list(df.columns)) - constraint.update_membership(new_groups) + if hasattr(model, "update_membership"): + model.update_membership(new_groups) else: logger.warning( - "Could not locate encoder projection weights for constrained regrouping; " + "Could not obtain projection weights for constrained regrouping; " "skipping grouping_fn-based membership update." ) - constraint.set_active(True) + model.set_constraint_active(True) refresh_mask() def start_unconstrained_phase(): - constraint.set_active(False) + model.set_constraint_active(False) refresh_mask() - constraint.set_active(False) + model.set_constraint_active(False) refresh_mask() total_epochs = sum(phases) if phases else epochs @@ -448,12 +466,6 @@ def _append_history(key: str, value: float) -> None: _append_history("kl_loss", epoch_kl / batch_count) if epoch % 2 == 0: - print(f"Epoch {epoch + 1}/{total_epochs} | ") - print( - f"loss={model.history['loss'][-1]:.4f} | " - f"recon={model.history['reconstruction_loss'][-1]:.4f} | " - f"kl={model.history['kl_loss'][-1]:.4f}" - ) logger.info( "Epoch %d/%d | loss=%.4f | recon=%.4f | kl=%.4f", epoch + 1, @@ -463,7 +475,8 @@ def _append_history(key: str, value: float) -> None: model.history["kl_loss"][-1], ) - model.latent_groups = constraint.latent_membership + if latent_groups is not None: + model.latent_groups = latent_groups model.eval() with torch.no_grad(): @@ -475,30 +488,3 @@ def _append_history(key: str, value: float) -> None: normal_pred = pd.DataFrame(recon, index=df.index, columns=df.columns) resid = normal_pred - df model.normal_stats = pd.DataFrame({"mean": resid.mean(), "std": resid.std(ddof=0)}) - - -def fit_alt(model, loader, lr:float = 1e-5, epochs=32, accumulate_steps=8): - optimizer = Adam(model.parameters(), lr=lr) - criterion = nn.MSELoss() - - model.history = {"loss": []} - - def step_fn(batch, beta_value: float) -> Dict[str, torch.Tensor]: - del beta_value - x_tensor, y_tensor = batch - predictions = model(x_tensor) - loss = criterion(predictions, y_tensor) - return {"loss": loss} - - _run_training_phases( - model=model, - loader=loader, - optimizer=optimizer, - phases=[(1.0, int(epochs))], - metric_keys=["loss"], - step_fn=step_fn, - progress=True, - accumulate_steps=int(accumulate_steps), - ) - - return [{"epoch": i + 1, "loss": loss} for i, loss in enumerate(model.history["loss"])] \ No newline at end of file diff --git a/src/embkit/pathway.py b/src/embkit/pathway.py index 6ae45d5..e3c8776 100644 --- a/src/embkit/pathway.py +++ b/src/embkit/pathway.py @@ -2,15 +2,27 @@ Methods for opening and processing Pathway files """ from collections import OrderedDict, defaultdict -from typing import Dict, List, Tuple, Iterable, Optional +from typing import Any, Dict, List, Tuple, Iterable import numpy as np import pandas as pd +def _normalize_index(index_like: Any) -> pd.Index: + """Accept pd.Index, list-like, or name->position maps and return a pd.Index.""" + if isinstance(index_like, pd.Index): + return index_like + if isinstance(index_like, dict): + ordered = [None] * len(index_like) + for name, pos in index_like.items(): + ordered[pos] = name + return pd.Index(ordered) + return pd.Index(list(index_like)) + + # ---------- SIF parsing ---------- -def extract_pathway_interactions( +def extract_sif_interactions( sif_path: str, relation: str = "controls-expression-of", ) -> Dict[str, List[str]]: @@ -53,51 +65,48 @@ def extract_pathway_interactions( return fmap -def build_sif_mask( - sif_path: str, - src_index: Dict[str, int], - dst_index: Dict[str, int], - *, - relation: Optional[str] = "controls-expression-of", +def build_mask( + feature_map: Dict[str, List[str]], + src_index: Any, + dst_index: Any, + min_group_size: int = 1, ) -> np.ndarray: """ Build a binary mask from a SIF file using explicit source/destination indices. + Because this is a controller mapping (ie the src controls the dst) the inner layer + dimension corresponds to the source (controller) and the outer dimension corresponds + to the destination (target). This way the mask can be directly applied to a weight + matrix of shape (out_features, in_features) where out_features are the targets and + in_features are the controllers. + The returned mask has shape (len(dst_index), len(src_index)) and is compatible with weight matrices shaped (out_features, in_features) for MaskedLinear layers. Only edges that map from a known source to a known destination are set to 1. Parameters ---------- - sif_path : str - Path to the SIF file containing ``src \t relation \t dst`` rows. - src_index : Dict[str, int] - Mapping from source node name to its column index. - dst_index : Dict[str, int] - Mapping from destination node name to its row index. - relation : Optional[str] - If provided, filter rows to only those with the given relation. + feature_map : Dict[str, List[str]] + Mapping from source node to list of destination nodes (e.g. TF to target genes). + src_index : pd.Index + Index of source nodes. + dst_index : pd.Index + Index of destination nodes. Returns ------- np.ndarray Binary mask with shape (len(dst_index), len(src_index)). """ - pc = pd.read_csv( - sif_path, - sep="\t", - header=None, - names=["from", "relation", "to"], - dtype=str, - ).fillna("") - - if relation is not None: - pc = pc.loc[pc["relation"] == relation] - + src_index = _normalize_index(src_index) + dst_index = _normalize_index(dst_index) mask = np.zeros((len(dst_index), len(src_index)), dtype=np.float32) - for src, _, dst in pc[["from", "relation", "to"]].itertuples(index=False, name=None): - if src in src_index and dst in dst_index: - mask[dst_index[dst], src_index[src]] = 1.0 + for src, members in feature_map.items(): + if src in src_index: + if dst_index.intersection(members).size >= min_group_size: + for dst in members: + if dst in dst_index: + mask[dst_index.get_loc(dst), src_index.get_loc(src)] = 1.0 return mask @@ -107,121 +116,112 @@ def build_sif_mask( def feature_map_intersect( feature_map: Dict[str, List[str]], features: Iterable[str], - *, - keep_lonely_groups: bool = False, -) -> Tuple[Dict[str, List[str]], List[str]]: + min_group_size: int = 2, + include_self: bool = True, +) -> Dict[str, List[str]]: """ Subset `feature_map` to nodes present in `features`. Keeps the group's self-node - if it is in `features`. Members are deduplicated and ordered by the order - they appear in the original mapping (stable). + if it is in `features`. Members are deduplicated. if include_self is True, the + source node will be included in the members list if it is present in the features. Returns: - (subset_map, intersect_list_in_input_order) + subset_map: Dict[str, List[str]]: subset of feature_map with only features in `features` """ features_list = list(features) feature_set = set(features_list) - - # Build a stable list of all nodes (sources + members) - all_nodes: "OrderedDict[str, None]" = OrderedDict() + out_map = {} for src, members in feature_map.items(): - all_nodes[src] = None - for m in members: - all_nodes[m] = None - - # Intersection, respecting the order of features_list (caller’s column order) - isect = [f for f in features_list if f in all_nodes] - - out: Dict[str, List[str]] = {} + filtered_members: List[str] = [] + seen = set() + for member in members: + if member in feature_set and member not in seen: + filtered_members.append(member) + seen.add(member) + if include_self and src in feature_set and src not in seen: + filtered_members = [src] + filtered_members + if len(filtered_members) >= min_group_size: + out_map[src] = filtered_members + return out_map + +def feature_map_link_filter( + feature_map: Dict[str, List[str]], + min_group_size: int = 2) -> Dict[str, List[str]]: + """Filter a feature map to only include groups that have at least `min_group_size` members (including self if present)""" + out_map = {} for src, members in feature_map.items(): - # intersect members (including possible self if it was in the original list) - filtered = [m for m in members if m in feature_set] - # Also include the group key itself if present in features even when not in members - if src in feature_set and src not in filtered: - filtered = [src] + filtered + unique_members = set(members) + if len(unique_members) >= min_group_size: + out_map[src] = members + return out_map - if filtered or (keep_lonely_groups and src in feature_set): - # Keep group only if it retains any member in the intersection, - # unless keep_lonely_groups=True and the group itself is in features - out[src] = filtered +def build_feature_map_indices( + feature_map: Dict[str, List[str]]) -> Tuple[pd.Index, pd.Index]: + """ + Create feature index from a feature map + """ + feature_set = set() + group_set = sorted(feature_map.keys()) + for group in group_set: + feature_set.update(feature_map[group]) + feature_idx = pd.Index( sorted(feature_set) ) + group_idx = pd.Index(group_set) + return feature_idx, group_idx + +def idx_to_list(x): + """ + idx_to_list: takes an index map ( name -> position ) to a list of names + ordered by position + """ + if isinstance(x, pd.Index): + return list(x) + if isinstance(x, dict): + out = [None] * len(x) + for k, v in x.items(): + out[v] = k + return out + return list(x) - return out, isect +def build_features_to_group_mask(feature_map, feature_idx, group_idx, group_node_count=1, forward=True): + """ + Build a masked linear layer based on connecting all features to a + single group node and forcing all other connections to be zero + """ + features = idx_to_list(feature_idx) + groups = idx_to_list(group_idx) + + in_dim = len(features) + out_dim = len(groups) * group_node_count + + if forward: + mask = np.zeros((out_dim, in_dim), dtype=np.float32) + else: + mask = np.zeros((in_dim, out_dim), dtype=np.float32) + + fi = pd.Index(features) + for gnum, group in enumerate(groups): + for f in feature_map[group]: + if f in fi: + floc = fi.get_loc(f) + for pos in range(gnum * group_node_count, (gnum + 1) * (group_node_count)): + if forward: + mask[pos, floc] = 1.0 + else: + mask[floc, pos] = 1.0 + return mask -# ---------- FeatureGroups ---------- -class FeatureGroups: +def build_group_to_group_mask(group_count: int, in_group_node_count, out_group_node_count): """ - Thin wrapper that preserves insertion order of groups and members, - supports (de)serialization, and builds deterministic index maps. + build_group_to_group + Build a mask that constricts connections between 2 group layer nodes """ - - def __init__(self, map: Dict[str, List[str]]): - od = OrderedDict() - for k, v in map.items(): - # enforce list and preserve order while removing duplicates - seen = set() - ordered_members = [] - for m in v: - if m not in seen: - seen.add(m) - ordered_members.append(m) - od[k] = ordered_members - self.map: "OrderedDict[str, List[str]]" = od - - def to_indices( - self, - *, - feature_order: Optional[Iterable[str]] = None, - group_order: Optional[Iterable[str]] = None, - ) -> Tuple[Dict[str, int], Dict[str, int]]: - """ - Create index maps for features and groups. - - If `feature_order` is provided, features are indexed by that order (and - filtered to those present in the groups). Otherwise, insertion order - (group-by-group) is used. - - If `group_order` is provided, groups are indexed by that order; otherwise - insertion order is used. - """ - # group index - groups_iter = list(self.map.keys()) if group_order is None else [g for g in group_order if g in self.map] - group_idx = {g: i for i, g in enumerate(groups_iter)} - - # features seen in groups (preserve insertion order across groups) - seen = OrderedDict() - for g in groups_iter: - for m in self.map[g]: - seen[m] = None - - if feature_order is None: - feat_iter = list(seen.keys()) - else: - feat_iter = [f for f in feature_order if f in seen] - - feature_idx = {f: i for i, f in enumerate(feat_iter)} - return feature_idx, group_idx - - # convenience - def __len__(self) -> int: - return len(self.map) - - def items(self): - return self.map.items() - - def features(self) -> List[str]: - out = OrderedDict() - for _, members in self.map.items(): - for m in members: - out[m] = None - return list(out.keys()) - - # --- serialization helpers --- - - def to_dict(self) -> Dict[str, List[str]]: - # plain JSON-serializable dict - return {k: list(v) for k, v in self.map.items()} - - @staticmethod - def from_dict(d: Dict[str, List[str]]) -> "FeatureGroups": - return FeatureGroups(d) + in_dim = group_count * in_group_node_count + out_dim = group_count * out_group_node_count + + mask = np.zeros((out_dim, in_dim), dtype=np.float32) + for g in range(group_count): + for i in range(g * in_group_node_count, (g + 1) * in_group_node_count): + for j in range(g * out_group_node_count, (g + 1) * out_group_node_count): + mask[j, i] = 1.0 + return mask diff --git a/src/embkit/preprocessing/normalize.py b/src/embkit/preprocessing/normalize.py index da9729f..8313a15 100644 --- a/src/embkit/preprocessing/normalize.py +++ b/src/embkit/preprocessing/normalize.py @@ -10,7 +10,7 @@ from torch.utils.data import Dataset import torch -def quantile_max_norm(df: pd.DataFrame, quantile_max=0.9): +def quantile_max_norm(df: pd.DataFrame, quantile_max: float = 0.9) -> pd.DataFrame: """ Normalizes the DataFrame using Quantile Max normalization. @@ -24,7 +24,7 @@ def quantile_max_norm(df: pd.DataFrame, quantile_max=0.9): norm_df = (df.transpose() / df.quantile(quantile_max, axis=1)).transpose().clip(upper=1.0, lower=0.0).fillna(0.0) return norm_df -def exp_max_norm(df: pd.DataFrame): +def exp_max_norm(df: pd.DataFrame) -> pd.DataFrame: """ exp_max_norm @@ -41,11 +41,11 @@ class ExpMinMaxScaler(MinMaxScaler, BaseEstimator): MinMaxScaler to it. This allows for normalization of data with different scales, especially when dealing with non-negative values where logarithmic transformation can help. """ - def fit(self, X: pd.DataFrame): + def fit(self, X: pd.DataFrame) -> "ExpMinMaxScaler": return MinMaxScaler.fit(self, np.log2(X+1)) - def transform(self, X): + def transform(self, X: pd.DataFrame) -> np.ndarray: return MinMaxScaler.transform(self, np.log2(X+1)) - def inverse_transform(self, X): + def inverse_transform(self, X: np.ndarray) -> np.ndarray: return np.exp2(MinMaxScaler.inverse_transform(self, X))-1 def get_dataset_nonzero_mask(d: Dataset, threshold: float) -> List[torch.Tensor]: diff --git a/src/embkit/resources/resource.py b/src/embkit/resources/resource.py index 7660a44..3f96419 100644 --- a/src/embkit/resources/resource.py +++ b/src/embkit/resources/resource.py @@ -29,10 +29,20 @@ def __init__( self.name = name if save_path is None: - # Use default repository directory under home - default_path = Path(Path.home(), REPO_DIR) - default_path.mkdir(parents=True, exist_ok=True) - self.save_path = default_path + # Use default repository directory under home (overridable by EMBKIT_HOME). + env_home = os.environ.get("EMBKIT_HOME") + default_path = Path(env_home) if env_home else Path(Path.home(), REPO_DIR) + try: + default_path.mkdir(parents=True, exist_ok=True) + self.save_path = default_path + except PermissionError: + fallback = Path(tempfile.mkdtemp(prefix="embkit-")) + warnings.warn( + f"Could not create default resource path '{default_path}'. " + f"Falling back to temporary directory '{fallback}'.", + stacklevel=2, + ) + self.save_path = fallback else: logger.debug(f"Using save_path={save_path}") self.save_path = Path(save_path) @@ -42,8 +52,8 @@ def __init__( try: self.download() self._download_called_from_init = True - except Exception as e: - logger.error(e) + except (RuntimeError, requests.RequestException, OSError, ValueError) as e: + logger.error("Download failed for resource '%s': %s", self.name, e) else: # When not downloading, set target file based on the resolved save_path target_file: Path = Path(self.save_path, self.name) @@ -117,12 +127,15 @@ def download(self) -> bytes: # Use a temporary file with tempfile.NamedTemporaryFile(delete=False) as tmp_file: tmp_path = Path(tmp_file.name) + # disable tqdm in CI environments + is_ci = os.environ.get("GITHUB_ACTIONS") == "true" with tqdm( desc=f"Downloading {self.name}", total=total_size, unit="B", unit_scale=True, unit_divisor=1024, + disable=is_ci, ) as bar: for chunk in response.iter_content(chunk_size=8192): if chunk: diff --git a/src/embkit/utilities/pca.py b/src/embkit/utilities/pca.py index 15ea467..94e36aa 100644 --- a/src/embkit/utilities/pca.py +++ b/src/embkit/utilities/pca.py @@ -1,9 +1,12 @@ from pathlib import Path +import logging from ..files import LargeCsvReader import numpy as np import faiss import pandas as pd +logger = logging.getLogger(__name__) + def run_pca(input_file: Path | str, pca_size: int, output_file: Path | str | None = None) -> pd.DataFrame: """ Runs PCA on the data from the input file and optionally saves the result. @@ -16,7 +19,7 @@ def run_pca(input_file: Path | str, pca_size: int, output_file: Path | str | Non pd.DataFrame: DataFrame containing the PCA-transformed data. """ csvfile: LargeCsvReader = LargeCsvReader(input_file, sep="\t", - index_column=0, skip_header=True, + index_column=0, skip_header=False, cache_size=128) data = np.array(list(csvfile.read(show_progress=True))) @@ -25,14 +28,14 @@ def run_pca(input_file: Path | str, pca_size: int, output_file: Path | str | Non for k, _ in csvfile: names.append(k) - print("calculating PCA") + logger.info("Calculating PCA") mat = faiss.PCAMatrix(data.shape[1], pca_size) mat.train(data) - print("build pca matrix") + logger.info("Building PCA matrix") data_pca = mat.apply(data) df = pd.DataFrame(data_pca, index=names) if output_file is not None: df.to_csv(output_file, sep="\t") - return df \ No newline at end of file + return df diff --git a/tests/c_bio/test_c_bio_api.py b/tests/c_bio/test_c_bio_api.py index 054c387..e59f1fc 100644 --- a/tests/c_bio/test_c_bio_api.py +++ b/tests/c_bio/test_c_bio_api.py @@ -20,8 +20,10 @@ def test_list_studies_success(self, mock_get): def test_list_studies_request_exception(self, mock_get): mock_get.side_effect = requests.RequestException("Network error") - result = CBIOAPI.list_studies() - self.assertIsNone(result) + with self.assertLogs("embkit.c_bio.api", level="ERROR") as cm: + result = CBIOAPI.list_studies() + self.assertIsNone(result) + self.assertTrue(any("Error fetching studies: Network error" in log for log in cm.output)) @patch("embkit.c_bio.api.requests.get") def test_list_studies_json_exception(self, mock_get): diff --git a/tests/commands/test_matrix.py b/tests/commands/test_matrix.py new file mode 100644 index 0000000..9933863 --- /dev/null +++ b/tests/commands/test_matrix.py @@ -0,0 +1,123 @@ +import unittest +from pathlib import Path + +import pandas as pd +from click.testing import CliRunner + +from embkit.__main__ import cli_main + + +class TestMatrixCommands(unittest.TestCase): + def setUp(self): + self.runner = CliRunner() + self.df = pd.DataFrame( + [ + [1.0, 2.0, 3.0, 4.0], + [2.0, 3.0, 4.0, 5.0], + [3.0, 4.0, 5.0, 6.0], + [4.0, 5.0, 6.0, 7.0], + [5.0, 6.0, 7.0, 8.0], + ], + index=["s1", "s2", "s3", "s4", "s5"], + columns=["G1", "G2", "G3", "G4"], + ) + + def test_pca_smoke(self): + with self.runner.isolated_filesystem(): + self.df.to_csv("toy.tsv", sep="\t") + result = self.runner.invoke( + cli_main, + [ + "matrix", + "pca", + "toy.tsv", + "--pca-size", + "2", + "--out", + "toy.pca.tsv", + ], + ) + + self.assertEqual(result.exit_code, 0, msg=result.output) + out_df = pd.read_csv("toy.pca.tsv", sep="\t", index_col=0) + self.assertEqual(out_df.shape, (5, 2)) + + def test_pca_default_out_and_bad_size(self): + with self.runner.isolated_filesystem(): + self.df.to_csv("toy.tsv", sep="\t") + + bad = self.runner.invoke( + cli_main, + ["matrix", "pca", "toy.tsv", "--pca-size", "0"], + ) + self.assertNotEqual(bad.exit_code, 0) + self.assertIn("--pca-size must be a positive integer", bad.output) + + ok = self.runner.invoke( + cli_main, + ["matrix", "pca", "toy.tsv", "--pca-size", "2"], + ) + self.assertEqual(ok.exit_code, 0, msg=ok.output) + self.assertIn("No output path provided, using default naming", ok.output) + self.assertTrue(Path("toy.pca.tsv").exists()) + + def test_normalize_smoke_and_features_subset(self): + df2 = self.df * 2.0 + with self.runner.isolated_filesystem(): + self.df.to_csv("a.tsv", sep="\t") + df2.to_csv("b.tsv", sep="\t") + with open("features.txt", "w", encoding="ascii") as f: + f.write("G1\nG3\n") + + result = self.runner.invoke( + cli_main, + [ + "matrix", + "normalize", + "a.tsv", + "b.tsv", + "--out", + "norm.tsv", + "--features", + "features.txt", + "--precision", + "4", + ], + ) + self.assertEqual(result.exit_code, 0, msg=result.output) + out_df = pd.read_csv("norm.tsv", sep="\t", index_col=0) + self.assertEqual(list(out_df.columns), ["G1", "G3"]) + self.assertEqual(out_df.shape[0], 10) + self.assertTrue((out_df.values >= 0.0).all()) + self.assertTrue((out_df.values <= 1.0).all()) + + def test_normalize_col_quantile_and_empty_sources(self): + with self.runner.isolated_filesystem(): + empty = self.runner.invoke( + cli_main, + ["matrix", "normalize", "--out", "norm.tsv"], + ) + self.assertEqual(empty.exit_code, 0, msg=empty.output) + self.assertIn("No matrices defined", empty.output) + + self.df.to_csv("a.tsv", sep="\t") + colq = self.runner.invoke( + cli_main, + [ + "matrix", + "normalize", + "a.tsv", + "--out", + "colq.tsv", + "--col-quantile", + "--quantile-max", + "0.8", + ], + ) + self.assertEqual(colq.exit_code, 0, msg=colq.output) + out_df = pd.read_csv("colq.tsv", sep="\t", index_col=0) + self.assertEqual(out_df.shape, self.df.shape) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/commands/test_model.py b/tests/commands/test_model.py index e69de29..90f9966 100644 --- a/tests/commands/test_model.py +++ b/tests/commands/test_model.py @@ -0,0 +1,321 @@ +import unittest +from unittest.mock import MagicMock, patch +import importlib +from pathlib import Path +import json + +from click.testing import CliRunner +import torch + +from embkit.__main__ import cli_main +from embkit.files import H5Writer + +model_cmd = importlib.import_module("embkit.commands.model") + + +class TestModelCommands(unittest.TestCase): + def setUp(self): + self.runner = CliRunner() + + @patch.object(model_cmd, "save") + @patch.object(model_cmd, "fit_vae") + @patch.object(model_cmd, "dataframe_loader", return_value="loader") + @patch.object(model_cmd, "NetVAE") + def test_train_netvae_smoke(self, netvae_cls, loader_mock, fit_mock, save_mock): + dummy_model = MagicMock(name="netvae") + netvae_cls.return_value = dummy_model + + with self.runner.isolated_filesystem(): + with open("rna.tsv", "w", encoding="utf-8") as f: + f.write( + "sample\tG1\tG2\tG3\tG4\n" + "s1\t1\t2\t3\t4\n" + "s2\t4\t3\t2\t1\n" + ) + with open("pathway.sif", "w", encoding="utf-8") as f: + f.write( + "TF1\tcontrols-expression-of\tG1\n" + "TF1\tcontrols-expression-of\tG2\n" + "TF2\tcontrols-expression-of\tG3\n" + "TF2\tcontrols-expression-of\tG4\n" + ) + + result = self.runner.invoke( + cli_main, + [ + "model", + "train-netvae", + "rna.tsv", + "pathway.sif", + "--epochs", + "1", + "--group-layer-size", + "4,2,1", + "--out", + "netvae.model", + ], + ) + + self.assertEqual(result.exit_code, 0, msg=result.output) + netvae_cls.assert_called_once() + netvae_args = netvae_cls.call_args.args + netvae_kwargs = netvae_cls.call_args.kwargs + self.assertEqual(netvae_args[0], ["G1", "G2", "G3", "G4"]) + self.assertEqual(netvae_kwargs["group_layer_size"], [4, 2, 1]) + self.assertEqual(set(netvae_kwargs["latent_groups"].keys()), {"TF1", "TF2"}) + + loader_mock.assert_called_once() + fit_mock.assert_called_once() + self.assertEqual(fit_mock.call_args.kwargs["X"], "loader") + save_mock.assert_called_once_with(dummy_model, "netvae.model") + + def test_train_vae_h5_rejects_normalization(self): + with self.runner.isolated_filesystem(): + writer = H5Writer("matrix.h5", "rna", index=["s1", "s2"], columns=["G1", "G2"]) + writer.set_irow(0, [1.0, 2.0]) + writer.set_irow(1, [3.0, 4.0]) + writer.close() + + result = self.runner.invoke( + cli_main, + [ + "model", + "train-vae", + "matrix.h5", + "--group", + "rna", + "--normalize", + "expMinMax", + "--epochs", + "1", + ], + ) + + self.assertNotEqual(result.exit_code, 0) + self.assertIn("Normalization for HDF5 input is not supported in train-vae", result.output) + + @patch.object(model_cmd, "save") + @patch.object(model_cmd, "fit_vae") + @patch.object(model_cmd, "VAE") + def test_train_vae_tsv_branches(self, vae_cls, fit_mock, save_mock): + dummy_model = MagicMock(name="vae") + vae_cls.return_value = dummy_model + + with self.runner.isolated_filesystem(): + with open("rna.tsv", "w", encoding="utf-8") as f: + f.write( + "sample\tG1\tG2\tG3\n" + "s1\t1\t2\t3\n" + "s2\t2\t3\t4\n" + ) + + result = self.runner.invoke( + cli_main, + [ + "model", + "train-vae", + "rna.tsv", + "--normalize", + "minMax", + "--loss", + "mse", + "--schedule", + "1:0.2,1:0.4", + "--save-stats", + ], + ) + + self.assertEqual(result.exit_code, 0, msg=result.output) + self.assertIn("No output path provided, using default naming.", result.output) + self.assertIn("Stats saved, to vae_latent256_epochs20.model.stats.tsv", result.output) + fit_mock.assert_called_once() + self.assertEqual(fit_mock.call_args.kwargs["loss"], model_cmd.mse) + self.assertEqual(fit_mock.call_args.kwargs["beta_schedule"], [(0.2, 1), (0.4, 1)]) + save_mock.assert_called_once() + + @patch.object(model_cmd, "save") + @patch.object(model_cmd, "fit_vae") + @patch.object(model_cmd, "dataframe_loader", return_value="loader") + @patch.object(model_cmd, "NetVAE") + def test_train_netvae_default_out_and_stats(self, netvae_cls, loader_mock, fit_mock, save_mock): + netvae_cls.return_value = MagicMock(name="netvae") + + with self.runner.isolated_filesystem(): + with open("rna.tsv", "w", encoding="utf-8") as f: + f.write( + "sample\tG1\tG2\n" + "s1\t1\t2\n" + "s2\t2\t3\n" + ) + with open("pathway.sif", "w", encoding="utf-8") as f: + f.write("TF1\tcontrols-expression-of\tG1\n") + f.write("TF1\tcontrols-expression-of\tG2\n") + + result = self.runner.invoke( + cli_main, + [ + "model", + "train-netvae", + "rna.tsv", + "pathway.sif", + "--epochs", + "3", + "--normalize", + "expMinMax", + "--loss", + "bce", + "--save-stats", + ], + ) + + self.assertEqual(result.exit_code, 0, msg=result.output) + self.assertIn("No output path provided, using default naming.", result.output) + self.assertIn("Stats saved, to netvae_latent1_epochs3.model.stats.tsv", result.output) + self.assertEqual(fit_mock.call_args.kwargs["loss"], model_cmd.bce) + loader_mock.assert_called_once() + save_mock.assert_called_once() + + def test_train_netvae_rejects_bad_group_layer_size(self): + with self.runner.isolated_filesystem(): + with open("rna.tsv", "w", encoding="utf-8") as f: + f.write("sample\tG1\ns1\t1\n") + with open("pathway.sif", "w", encoding="utf-8") as f: + f.write("TF1\tcontrols-expression-of\tG1\n") + + result = self.runner.invoke( + cli_main, + [ + "model", + "train-netvae", + "rna.tsv", + "pathway.sif", + "--group-layer-size", + "0", + ], + ) + + self.assertNotEqual(result.exit_code, 0) + self.assertIn("--group-layer-size must contain one or more positive integers", result.output) + + @patch.object(model_cmd, "load") + @patch.object(model_cmd, "get_device", return_value=torch.device("cpu")) + def test_encode_with_expminmax(self, _device, load_mock): + class DummyModel: + def __init__(self): + self.features = ["G1", "G2"] + + def to(self, *_args, **_kwargs): + return self + + def encoder(self, x): + return (None, None, x[:, :1]) + + load_mock.return_value = DummyModel() + + with self.runner.isolated_filesystem(): + with open("rna.tsv", "w", encoding="utf-8") as f: + f.write("sample\tG1\tG2\ns1\t1\t2\ns2\t2\t3\n") + Path("dummy.model").write_text("mock", encoding="utf-8") + result = self.runner.invoke( + cli_main, + ["model", "encode", "rna.tsv", "dummy.model", "--normalize", "expMinMax", "--out", "embed.tsv"], + ) + + self.assertEqual(result.exit_code, 0, msg=result.output) + self.assertTrue(Path("embed.tsv").exists()) + + @patch("embkit.factory.core.run_model_verification") + def test_verify_json_ci_pass(self, verify_mock): + verify_mock.return_value = { + "model_type": "NetVAE", + "healthy": True, + "issues": [], + "features_count": 2, + "feature_names": ["G1", "G2"], + "deep_audit": {"latent_dim": 1}, + } + + with self.runner.isolated_filesystem(): + Path("dummy.model").write_text("mock", encoding="utf-8") + result = self.runner.invoke( + cli_main, + ["model", "verify", "dummy.model", "--ci"], + ) + + self.assertEqual(result.exit_code, 0, msg=result.output) + payload = json.loads(result.output) + self.assertTrue(payload["healthy"]) + self.assertTrue(payload["strict_mode"] is False) + + @patch("embkit.factory.core.run_model_verification") + def test_verify_strict_mismatch_fails(self, verify_mock): + verify_mock.return_value = { + "model_type": "NetVAE", + "healthy": True, + "issues": [], + "features_count": 2, + "feature_names": ["G1", "G2"], + "deep_audit": {"latent_dim": 1}, + } + + with self.runner.isolated_filesystem(): + Path("dummy.model").write_text("mock", encoding="utf-8") + Path("expected_features.txt").write_text("G1\nG9\n", encoding="utf-8") + result = self.runner.invoke( + cli_main, + [ + "model", + "verify", + "dummy.model", + "--strict", + "--expected-features-file", + "expected_features.txt", + "--expected-feature-count", + "2", + "--expected-latent-dim", + "1", + "--fail-on-unhealthy", + "--json", + ], + ) + + self.assertNotEqual(result.exit_code, 0) + self.assertIn("Model integrity check failed.", result.output) + + @patch("embkit.factory.core.run_model_verification") + def test_verify_strict_match_passes(self, verify_mock): + verify_mock.return_value = { + "model_type": "NetVAE", + "healthy": True, + "issues": [], + "features_count": 2, + "feature_names": ["G1", "G2"], + "deep_audit": {"latent_dim": 3}, + } + + with self.runner.isolated_filesystem(): + Path("dummy.model").write_text("mock", encoding="utf-8") + Path("expected_features.txt").write_text("G1\nG2\n", encoding="utf-8") + result = self.runner.invoke( + cli_main, + [ + "model", + "verify", + "dummy.model", + "--strict", + "--expected-features-file", + "expected_features.txt", + "--expected-feature-count", + "2", + "--expected-latent-dim", + "3", + "--fail-on-unhealthy", + ], + ) + + self.assertEqual(result.exit_code, 0, msg=result.output) + self.assertIn("PASS: model integrity checks passed.", result.output) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/commands/test_protein.py b/tests/commands/test_protein.py new file mode 100644 index 0000000..3efd59c --- /dev/null +++ b/tests/commands/test_protein.py @@ -0,0 +1,86 @@ +import unittest +from unittest.mock import MagicMock, patch +import importlib + +import torch +from click.testing import CliRunner + +from embkit.__main__ import cli_main + +protein_cmd = importlib.import_module("embkit.commands.protein") + + +class TestProteinCommands(unittest.TestCase): + def setUp(self): + self.runner = CliRunner() + + def test_fasta_reader_and_stringify(self): + with self.runner.isolated_filesystem(): + with open("toy.fasta", "w", encoding="utf-8") as f: + f.write(">A1\nMKT\n>B2\nAAA\n") + + rows = list(protein_cmd.fasta_reader("toy.fasta", filter="^A")) + self.assertEqual(rows, [("A1", "MKT")]) + + out = protein_cmd.stringify([1.23456, 2.0], trim=2) + self.assertEqual(out, ["1.23", "2"]) + + @patch.object(protein_cmd, "get_device", return_value=torch.device("cpu")) + @patch.object(protein_cmd, "ProteinEncoder") + def test_encode_mean_stdout(self, enc_cls, _dev): + dummy = MagicMock() + dummy.encode.return_value = [("seq1", torch.tensor([1.2, 3.4]))] + enc_cls.return_value = dummy + + with self.runner.isolated_filesystem(): + with open("toy.fasta", "w", encoding="utf-8") as f: + f.write(">seq1\nMKT\n") + + res = self.runner.invoke( + cli_main, + ["protein", "encode", "toy.fasta", "--pool", "mean", "--model", "t6"], + ) + + self.assertEqual(res.exit_code, 0, msg=res.output) + self.assertIn("seq1\t1.200000\t3.400000", res.output) + dummy.to.assert_called_once() + + @patch.object(protein_cmd, "get_device", return_value=torch.device("cpu")) + @patch.object(protein_cmd, "ProteinEncoder") + def test_encode_vector_output_file(self, enc_cls, _dev): + dummy = MagicMock() + dummy.encode.return_value = [("seq1", torch.tensor([[1.23456, 2.34567]]))] + enc_cls.return_value = dummy + + with self.runner.isolated_filesystem(): + with open("toy.fasta", "w", encoding="utf-8") as f: + f.write(">seq1\nMKT\n") + + res = self.runner.invoke( + cli_main, + [ + "protein", + "encode", + "toy.fasta", + "--pool", + "none", + "--trim", + "2", + "--output", + "vec.jsonl", + "--model", + "t6", + "--fix-len", + "8", + ], + ) + + self.assertEqual(res.exit_code, 0, msg=res.output) + with open("vec.jsonl", "r", encoding="utf-8") as fh: + text = fh.read().strip() + self.assertTrue(text.startswith("seq1\t")) + self.assertIn("[[1.23, 2.35]]", text) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/commands/test_resources.py b/tests/commands/test_resources.py new file mode 100644 index 0000000..389a87e --- /dev/null +++ b/tests/commands/test_resources.py @@ -0,0 +1,108 @@ +import importlib +import unittest +from types import SimpleNamespace +from pathlib import Path +from unittest.mock import patch + +import pandas as pd +from click.testing import CliRunner + +from embkit.__main__ import cli_main + +resources_cmd = importlib.import_module("embkit.commands.resources") + + +class TestResourcesCommands(unittest.TestCase): + def setUp(self): + self.runner = CliRunner() + + def test_gtex_rejects_unknown_dataset(self): + with self.runner.isolated_filesystem(): + result = self.runner.invoke( + cli_main, + ["resources", "gtex", "-t", "does_not_exist", "-f", "out"], + ) + + self.assertEqual(result.exit_code, 0, msg=result.output) + self.assertIn("not recognized", result.output) + + @patch.object(resources_cmd, "GTEx") + def test_gtex_download_invokes_resource(self, gtex_cls): + gtex_cls.NAMES = {"gene_tpm": "dummy"} + + with self.runner.isolated_filesystem(): + result = self.runner.invoke( + cli_main, + ["resources", "gtex", "-t", "gene_tpm", "-f", "out"], + ) + + self.assertEqual(result.exit_code, 0, msg=result.output) + gtex_cls.assert_called_once_with(data_type="gene_tpm", save_path="out") + + @patch.object(resources_cmd, "SIF") + def test_sif_download_invokes_resource(self, sif_cls): + with self.runner.isolated_filesystem(): + result = self.runner.invoke( + cli_main, + ["resources", "sif", "-f", "out"], + ) + + self.assertEqual(result.exit_code, 0, msg=result.output) + sif_cls.assert_called_once_with(save_path="out") + + @patch.object(resources_cmd, "Hugo") + def test_hugo_download_without_conversion(self, hugo_cls): + hugo_cls.return_value = SimpleNamespace(save_path="out") + + with self.runner.isolated_filesystem(): + result = self.runner.invoke( + cli_main, + ["resources", "hugo", "-f", "out"], + ) + + self.assertEqual(result.exit_code, 0, msg=result.output) + self.assertIn("Downloading hugo dataset into 'out'", result.output) + hugo_cls.assert_called_once_with(save_path="out") + + @patch.object(resources_cmd, "load_raw_hugo") + @patch.object(resources_cmd, "load_gct") + @patch.object(resources_cmd, "GTEx") + @patch.object(resources_cmd, "Hugo") + def test_hugo_gtex_conversion_writes_output(self, hugo_cls, gtex_cls, load_gct_mock, load_raw_hugo_mock): + hugo_cls.return_value = SimpleNamespace(save_path="out") + gtex_cls.return_value = "gtex.gct" + load_gct_mock.return_value = pd.DataFrame( + { + "ENSG000001.1": [1.0, 2.0], + "ENSG000099.9": [3.0, 4.0], + }, + index=["s1", "s2"], + ) + load_raw_hugo_mock.return_value = pd.DataFrame( + { + "locus_group": ["protein-coding gene", "RNA, transfer"], + "symbol": ["TP53", "TRNA1"], + "ensembl_gene_id": ["ENSG000001", "ENSG000099"], + } + ) + + with self.runner.isolated_filesystem(): + Path("out").mkdir(parents=True, exist_ok=True) + result = self.runner.invoke( + cli_main, + ["resources", "hugo", "-f", "out", "-gtex"], + ) + + out_path = Path("out/gtex.hugo.tsv") + self.assertTrue(out_path.exists()) + converted = pd.read_csv(out_path, sep="\t", index_col=0) + self.assertIn("TP53", converted.columns) + self.assertNotIn("ENSG000099.9", converted.columns) + + self.assertEqual(result.exit_code, 0, msg=result.output) + self.assertIn("Converting hugo dataset using GTEx dataset", result.output) + self.assertIn("Done.", result.output) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/constraints/test_network_constraints.py b/tests/constraints/test_network_constraints.py deleted file mode 100644 index d01d493..0000000 --- a/tests/constraints/test_network_constraints.py +++ /dev/null @@ -1,77 +0,0 @@ -import unittest -import numpy as np -import torch -from embkit.constraints import NetworkConstraint # adjust import path to match your structure - - -class TestNetworkConstraint(unittest.TestCase): - def setUp(self): - self.feature_index = ["f1", "f2", "f3", "f4"] - self.latent_index = ["z1", "z2"] - self.latent_membership = { - "z1": ["f1", "f2"], - "z2": ["f3"] - } - - def test_no_membership_active_constraint_is_all_ones(self): - constraint = NetworkConstraint(self.feature_index, self.latent_index, None) - expected = np.ones((2, 4), dtype=np.float32) - np.testing.assert_array_equal(constraint._mask_np, expected) - - def test_membership_creates_correct_mask(self): - constraint = NetworkConstraint(self.feature_index, self.latent_index, self.latent_membership) - expected = np.array([ - [1, 1, 0, 0], # z1 connects to f1, f2 - [0, 0, 1, 0] # z2 connects to f3 - ], dtype=np.float32) - np.testing.assert_array_equal(constraint._mask_np, expected) - - def test_inactive_constraint_is_all_ones(self): - constraint = NetworkConstraint(self.feature_index, self.latent_index, self.latent_membership) - constraint.set_active(False) - expected = np.ones((2, 4), dtype=np.float32) - np.testing.assert_array_equal(constraint._mask_np, expected) - - def test_update_membership_changes_mask(self): - constraint = NetworkConstraint(self.feature_index, self.latent_index, self.latent_membership) - new_membership = { - "z1": ["f4"], - "z2": ["f2", "f3"] - } - constraint.update_membership(new_membership) - expected = np.array([ - [0, 0, 0, 1], # z1 -> f4 - [0, 1, 1, 0] # z2 -> f2, f3 - ], dtype=np.float32) - np.testing.assert_array_equal(constraint._mask_np, expected) - - def test_as_torch_tensor(self): - constraint = NetworkConstraint(self.feature_index, self.latent_index, self.latent_membership) - tensor = constraint.as_torch(device=torch.device("cpu")) - self.assertIsInstance(tensor, torch.Tensor) - self.assertEqual(tensor.shape, (2, 4)) - expected = torch.tensor([ - [1, 1, 0, 0], - [0, 0, 1, 0] - ], dtype=torch.float32) - self.assertTrue(torch.equal(tensor, expected)) - - def test_latent_not_in_membership_triggers_continue(self): - feature_index = ["f1", "f2"] - latent_index = ["z1", "z2"] # z1 is included - latent_membership = { - "z2": ["f2"] # z1 is missing from the mapping - } - - constraint = NetworkConstraint(feature_index, latent_index, latent_membership) - - # z1 gets skipped entirely, so only z2 -> f2 matters - expected = np.array([ - [0, 0], # z1 skipped - [0, 1] # z2 -> f2 - ], dtype=np.float32) - - np.testing.assert_array_equal(constraint._mask_np, expected) - -if __name__ == '__main__': - unittest.main() \ No newline at end of file diff --git a/tests/constraints/test_pathway_constraints.py b/tests/constraints/test_pathway_constraints.py new file mode 100644 index 0000000..722e974 --- /dev/null +++ b/tests/constraints/test_pathway_constraints.py @@ -0,0 +1,120 @@ +import unittest +import numpy as np + +from embkit.constraints import PathwayConstraintInfo + + +class TestPathwayConstraintInfo(unittest.TestCase): + def setUp(self): + self.feature_index = ["f1", "f2", "f3", "f4"] + self.group_index = ["z1", "z2"] + self.feature_map = { + "z1": ["f1", "f2"], + "z2": ["f3"], + } + + def test_features_to_group_mask_shape_and_values(self): + c = PathwayConstraintInfo( + "features-to-group", + feature_map=self.feature_map, + feature_index=self.feature_index, + group_index=self.group_index, + ) + mask = c.gen_mask(in_features=4, out_features=2) + expected = np.array( + [ + [1, 1, 0, 0], + [0, 0, 1, 0], + ], + dtype=np.float32, + ) + np.testing.assert_array_equal(mask, expected) + + def test_inactive_returns_ones(self): + c = PathwayConstraintInfo( + "features-to-group", + feature_map=self.feature_map, + feature_index=self.feature_index, + group_index=self.group_index, + ) + c.set_active(False) + mask = c.gen_mask(in_features=4, out_features=2) + np.testing.assert_array_equal(mask, np.ones((2, 4), dtype=np.float32)) + + def test_update_membership_changes_mask(self): + c = PathwayConstraintInfo( + "features-to-group", + feature_map=self.feature_map, + feature_index=self.feature_index, + group_index=self.group_index, + ) + c.update_membership( + { + "z1": ["f4"], + "z2": ["f2", "f3"], + } + ) + mask = c.gen_mask(in_features=4, out_features=2) + expected = np.array( + [ + [0, 0, 0, 1], + [0, 1, 1, 0], + ], + dtype=np.float32, + ) + np.testing.assert_array_equal(mask, expected) + + def test_group_to_features_mask(self): + c = PathwayConstraintInfo( + "group-to-features", + feature_map=self.feature_map, + feature_index=self.feature_index, + group_index=self.group_index, + ) + mask = c.gen_mask(in_features=2, out_features=4) + self.assertEqual(mask.shape, (4, 2)) + self.assertEqual(mask[0, 0], 1.0) + self.assertEqual(mask[2, 1], 1.0) + + def test_group_to_group_mask(self): + c = PathwayConstraintInfo( + "group-to-group", + feature_map=self.feature_map, + feature_index=self.feature_index, + group_index=self.group_index, + ) + mask = c.gen_mask(in_features=4, out_features=2) + self.assertEqual(mask.shape, (2, 4)) + # group 0 only connects to group 0 inputs + self.assertEqual(mask[0, 0], 1.0) + self.assertEqual(mask[0, 1], 1.0) + self.assertEqual(mask[0, 2], 0.0) + + def test_to_from_dict_roundtrip(self): + c = PathwayConstraintInfo( + "features-to-group", + feature_map=self.feature_map, + feature_index=self.feature_index, + group_index=self.group_index, + in_group_scaling=1, + out_group_scaling=2, + ) + c.set_active(False) + payload = c.to_dict() + loaded = PathwayConstraintInfo.from_dict(payload) + self.assertEqual(loaded.op, "features-to-group") + self.assertFalse(loaded.active) + + def test_invalid_group_to_group_dimensions_raise(self): + c = PathwayConstraintInfo( + "group-to-group", + feature_map=self.feature_map, + feature_index=self.feature_index, + group_index=self.group_index, + ) + with self.assertRaises(ValueError): + c.gen_mask(in_features=3, out_features=4) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/datasets/__init__.py b/tests/datasets/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/datasets/test_datasets.py b/tests/datasets/test_datasets.py new file mode 100644 index 0000000..b4dfde1 --- /dev/null +++ b/tests/datasets/test_datasets.py @@ -0,0 +1,52 @@ +import itertools +import unittest + +import torch +from torch.utils.data import Dataset + +from embkit.datasets import BalancedMixer, DatasetMask + + +class TinyDataset(Dataset): + def __init__(self, values): + self.values = values + + def __len__(self): + return len(self.values) + + def __getitem__(self, idx): + return self.values[idx] + + +class TestDatasets(unittest.TestCase): + def test_balanced_mixer_iterates_and_recycles(self): + d1 = ["a1", "a2"] + d2 = ["b1"] + mix = BalancedMixer([d1, d2], seed=42) + + out = list(itertools.islice(iter(mix), 8)) + self.assertEqual(len(out), 8) + self.assertTrue(any(v.startswith("a") for v in out)) + self.assertTrue(any(v.startswith("b") for v in out)) + + def test_dataset_mask_applies_mask_and_device(self): + base = TinyDataset( + [ + (torch.tensor([1.0, 2.0, 3.0]), torch.tensor([5.0, 6.0])), + (torch.tensor([4.0, 5.0, 6.0]), torch.tensor([7.0, 8.0])), + ] + ) + mask = [torch.tensor([True, False, True]), torch.tensor([False, True])] + + masked = DatasetMask(base, mask, device=torch.device("cpu")) + self.assertEqual(len(masked), 2) + + x0, y0 = masked[0] + self.assertTrue(torch.equal(x0, torch.tensor([1.0, 3.0]))) + self.assertTrue(torch.equal(y0, torch.tensor([6.0]))) + self.assertEqual(x0.device.type, "cpu") + self.assertEqual(y0.device.type, "cpu") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/e2e/__init__.py b/tests/e2e/__init__.py new file mode 100644 index 0000000..ecdf82c --- /dev/null +++ b/tests/e2e/__init__.py @@ -0,0 +1 @@ +# Ensure unittest discovery descends into tests/e2e in CI. diff --git a/tests/e2e/data/toy_expr.tsv b/tests/e2e/data/toy_expr.tsv new file mode 100644 index 0000000..093a2de --- /dev/null +++ b/tests/e2e/data/toy_expr.tsv @@ -0,0 +1,11 @@ +sample G1 G2 G3 G4 G5 G6 G7 G8 G9 G10 +S1 0.10 0.20 0.30 0.40 0.50 0.60 0.70 0.80 0.90 0.10 +S2 0.20 0.30 0.40 0.50 0.60 0.70 0.80 0.90 0.10 0.20 +S3 0.30 0.40 0.50 0.60 0.70 0.80 0.90 0.10 0.20 0.30 +S4 0.40 0.50 0.60 0.70 0.80 0.90 0.10 0.20 0.30 0.40 +S5 0.50 0.60 0.70 0.80 0.90 0.10 0.20 0.30 0.40 0.50 +S6 0.60 0.70 0.80 0.90 0.10 0.20 0.30 0.40 0.50 0.60 +S7 0.70 0.80 0.90 0.10 0.20 0.30 0.40 0.50 0.60 0.70 +S8 0.80 0.90 0.10 0.20 0.30 0.40 0.50 0.60 0.70 0.80 +S9 0.90 0.10 0.20 0.30 0.40 0.50 0.60 0.70 0.80 0.90 +S10 0.15 0.25 0.35 0.45 0.55 0.65 0.75 0.85 0.95 0.05 diff --git a/tests/e2e/data/toy_pathway.sif b/tests/e2e/data/toy_pathway.sif new file mode 100644 index 0000000..379073f --- /dev/null +++ b/tests/e2e/data/toy_pathway.sif @@ -0,0 +1,10 @@ +TFA controls-expression-of G1 +TFA controls-expression-of G2 +TFA controls-expression-of G3 +TFA controls-expression-of G4 +TFB controls-expression-of G5 +TFB controls-expression-of G6 +TFB controls-expression-of G7 +TFC controls-expression-of G8 +TFC controls-expression-of G9 +TFC controls-expression-of G10 diff --git a/tests/e2e/test_cli_workflows_e2e.py b/tests/e2e/test_cli_workflows_e2e.py new file mode 100644 index 0000000..587286a --- /dev/null +++ b/tests/e2e/test_cli_workflows_e2e.py @@ -0,0 +1,167 @@ +import unittest +from pathlib import Path + +import pandas as pd +from click.testing import CliRunner + +from embkit.__main__ import cli_main +from embkit.files import H5Writer + + +class TestCLIWorkflowsE2E(unittest.TestCase): + def setUp(self): + self.runner = CliRunner() + + def test_train_vae_then_encode_tsv_e2e(self): + df = pd.DataFrame( + [ + [0.10, 0.20, 0.30, 0.40], + [0.20, 0.30, 0.40, 0.50], + [0.30, 0.40, 0.50, 0.60], + [0.40, 0.50, 0.60, 0.70], + [0.50, 0.60, 0.70, 0.80], + [0.60, 0.70, 0.80, 0.90], + ], + index=[f"s{i}" for i in range(1, 7)], + columns=["G1", "G2", "G3", "G4"], + ) + + with self.runner.isolated_filesystem(): + df.to_csv("toy.tsv", sep="\t") + + train_result = self.runner.invoke( + cli_main, + [ + "model", + "train-vae", + "toy.tsv", + "--latent", + "2", + "--epochs", + "1", + "--batch-size", + "2", + "--out", + "toy_vae.model", + "--save-stats", + ], + ) + self.assertEqual(train_result.exit_code, 0, msg=train_result.output) + self.assertTrue(Path("toy_vae.model").exists()) + self.assertTrue(Path("toy_vae.model.stats.tsv").exists()) + + encode_result = self.runner.invoke( + cli_main, + ["model", "encode", "toy.tsv", "toy_vae.model", "--out", "toy_embed.tsv"], + ) + self.assertEqual(encode_result.exit_code, 0, msg=encode_result.output) + out_df = pd.read_csv("toy_embed.tsv", sep="\t", index_col=0) + self.assertEqual(out_df.shape, (6, 2)) + + def test_train_vae_h5_normalize_guard_and_success_e2e(self): + with self.runner.isolated_filesystem(): + writer = H5Writer("toy.h5", "rna", index=["s1", "s2", "s3"], columns=["G1", "G2"]) + writer.set_irow(0, [1.0, 2.0]) + writer.set_irow(1, [2.0, 3.0]) + writer.set_irow(2, [3.0, 4.0]) + writer.close() + + bad_result = self.runner.invoke( + cli_main, + [ + "model", + "train-vae", + "toy.h5", + "--group", + "rna", + "--normalize", + "expMinMax", + "--epochs", + "1", + ], + ) + self.assertNotEqual(bad_result.exit_code, 0) + self.assertIn("Normalization for HDF5 input is not supported in train-vae", bad_result.output) + + ok_result = self.runner.invoke( + cli_main, + [ + "model", + "train-vae", + "toy.h5", + "--group", + "rna", + "--normalize", + "none", + "--latent", + "2", + "--epochs", + "1", + "--batch-size", + "2", + "--out", + "toy_h5.model", + ], + ) + self.assertEqual(ok_result.exit_code, 0, msg=ok_result.output) + self.assertTrue(Path("toy_h5.model").exists()) + + def test_train_netvae_no_overlap_fails_loudly_e2e(self): + df = pd.DataFrame( + [[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]], + index=["s1", "s2", "s3"], + columns=["G1", "G2"], + ) + + with self.runner.isolated_filesystem(): + df.to_csv("rna.tsv", sep="\t") + with open("pathway.sif", "w", encoding="utf-8") as f: + f.write("TF1\tcontrols-expression-of\tX1\n") + f.write("TF2\tcontrols-expression-of\tX2\n") + + result = self.runner.invoke( + cli_main, + [ + "model", + "train-netvae", + "rna.tsv", + "pathway.sif", + "--epochs", + "1", + "--out", + "bad.model", + ], + ) + + self.assertNotEqual(result.exit_code, 0) + self.assertIsInstance(result.exception, ValueError) + self.assertIn("latent_groups cannot be empty", str(result.exception)) + + def test_align_pair_cli_e2e(self): + a = pd.DataFrame( + [[1.0, 2.0, 3.0], [3.0, 2.0, 1.0]], + index=["A1", "A2"], + columns=["F1", "F2", "F3"], + ) + b = pd.DataFrame( + [[10.0, 20.0, 30.0], [30.0, 20.0, 10.0]], + index=["B1", "B2"], + columns=["F1", "F2", "F3"], + ) + + with self.runner.isolated_filesystem(): + a.to_csv("a.tsv", sep="\t") + b.to_csv("b.tsv", sep="\t") + + result = self.runner.invoke( + cli_main, + ["align", "pair", "a.tsv", "b.tsv", "--cutoff", "0.9"], + ) + + self.assertEqual(result.exit_code, 0, msg=result.output) + self.assertIn("A1\tB1", result.output) + self.assertIn("A2\tB2", result.output) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/e2e/test_netvae_cli_e2e.py b/tests/e2e/test_netvae_cli_e2e.py new file mode 100644 index 0000000..d4fac2f --- /dev/null +++ b/tests/e2e/test_netvae_cli_e2e.py @@ -0,0 +1,152 @@ +import unittest +from pathlib import Path +import logging + +import numpy as np +import pandas as pd +import torch +from click.testing import CliRunner + +from embkit.__main__ import cli_main +from embkit.factory import load +from embkit.losses import bce_with_logits +from embkit.modules import MaskedLinear +from embkit.pathway import ( + build_features_to_group_mask, + extract_sif_interactions, + feature_map_intersect, +) + +logger = logging.getLogger("tests.e2e.netvae_cli") + + +class TestNetVAECLIE2E(unittest.TestCase): + @classmethod + def setUpClass(cls): + logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s") + + def setUp(self): + self.runner = CliRunner() + self.data_dir = Path(__file__).resolve().parent / "data" + self.expr_tsv = self.data_dir / "toy_expr.tsv" + self.pathway_sif = self.data_dir / "toy_pathway.sif" + + def test_train_and_encode_netvae_with_constraints(self): + with self.runner.isolated_filesystem(): + model_path = "toy_netvae.model" + embed_path = "toy_embedding.tsv" + print("[E2E] Starting NetVAE train on toy data") + logger.info("E2E start: training toy NetVAE model") + + train_result = self.runner.invoke( + cli_main, + [ + "model", + "train-netvae", + str(self.expr_tsv), + str(self.pathway_sif), + "--epochs", + "2", + "--group-layer-size", + "1", + "--save-stats", + "--out", + model_path, + ], + ) + self.assertEqual(train_result.exit_code, 0, msg=train_result.output) + print("[E2E] Train command finished") + logger.info("Train command completed. Output:\n%s", train_result.output.strip()) + self.assertTrue(Path(model_path).exists()) + self.assertTrue(Path(f"{model_path}.stats.tsv").exists()) + logger.info("Artifacts created: %s and %s.stats.tsv", model_path, model_path) + + logger.info("Running encode command against trained model") + print("[E2E] Running encode command") + encode_result = self.runner.invoke( + cli_main, + [ + "model", + "encode", + str(self.expr_tsv), + model_path, + "--out", + embed_path, + ], + ) + self.assertEqual(encode_result.exit_code, 0, msg=encode_result.output) + print("[E2E] Encode command finished") + logger.info("Encode command completed. Output:\n%s", encode_result.output.strip()) + self.assertTrue(Path(embed_path).exists()) + + embedding = pd.read_csv(embed_path, sep="\t", index_col=0) + self.assertEqual(embedding.shape, (10, 3)) + print("[E2E] Embedding shape validated:", embedding.shape) + logger.info("Embedding shape validated: %s", embedding.shape) + + model = load(model_path, device=torch.device("cpu")) + self.assertIsNotNone(model.encoder) + self.assertIsNotNone(model.decoder) + logger.info("Model loaded and encoder/decoder present") + + first_masked = next(m for m in model.encoder.net if isinstance(m, MaskedLinear)) + self.assertEqual(tuple(first_masked.mask.shape), (3, 10)) + print("[E2E] Masked layer shape validated:", tuple(first_masked.mask.shape)) + logger.info("First masked encoder layer shape: %s", tuple(first_masked.mask.shape)) + + expr_df = pd.read_csv(self.expr_tsv, sep="\t", index_col=0) + fmap = extract_sif_interactions(str(self.pathway_sif)) + fmap = feature_map_intersect(fmap, expr_df.columns) + expected_mask = build_features_to_group_mask( + fmap, + feature_idx=model.features, + group_idx=model.latent_index, + group_node_count=1, + ) + + np.testing.assert_array_equal(first_masked.mask.detach().cpu().numpy(), expected_mask) + print("[E2E] Constraint mask equality check passed") + logger.info( + "Mask equality check passed. allowed_edges=%d blocked_edges=%d", + int(expected_mask.sum()), + int(expected_mask.size - expected_mask.sum()), + ) + + constraint_info = getattr(first_masked, "constraint_info", None) + self.assertIsNotNone(constraint_info) + self.assertTrue(getattr(constraint_info, "active", False)) + + effective_weight = first_masked.linear.weight.detach() * first_masked.mask.detach() + self.assertTrue(torch.all(effective_weight[first_masked.mask == 0] == 0)) + self.assertTrue(torch.any(torch.abs(effective_weight[first_masked.mask == 1]) > 0)) + print("[E2E] Effective masked weights validated") + logger.info("Effective masked weights validated (blocked=0, allowed has signal)") + + # Explicitly verify constrained edges are blocked during optimization. + x = torch.tensor(expr_df[model.features].values, dtype=torch.float32) + opt = torch.optim.Adam(model.parameters(), lr=1e-3) + model.train() + opt.zero_grad() + mu, logvar, z = model.encoder(x) + recon = model.decoder(z) + total, _, _ = bce_with_logits(recon, x, mu, logvar, beta=1.0) + total.backward() + + weight = first_masked.linear.weight + mask = first_masked.mask + self.assertIsNotNone(weight.grad) + self.assertTrue(torch.all(weight.grad[mask == 0] == 0)) + print("[E2E] Gradient blocking validated") + logger.info("Gradient blocking validated on constrained edges") + + blocked_before = weight.detach()[mask == 0].clone() + opt.step() + blocked_after = weight.detach()[mask == 0] + self.assertTrue(torch.all(blocked_after == blocked_before)) + print("[E2E] No-update invariant validated on constrained parameters") + logger.info("No-update invariant validated for constrained raw parameters") + logger.info("E2E NetVAE constraint test finished successfully") + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/encoding/__init__.py b/tests/encoding/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/encoding/test_protein_encoder.py b/tests/encoding/test_protein_encoder.py new file mode 100644 index 0000000..2c27f7f --- /dev/null +++ b/tests/encoding/test_protein_encoder.py @@ -0,0 +1,98 @@ +import unittest +from unittest.mock import patch + +import torch + +from embkit.encoding.protein import ProteinEncoder + + +class DummyAlphabet: + def __init__(self): + self.padding_idx = 0 + + def get_batch_converter(self): + def _convert(block): + labels = [x[0] for x in block] + seqs = [x[1] for x in block] + max_len = max(len(s) for s in seqs) if seqs else 0 + # +2 to simulate BOS/EOS style tokenized length. + tokens = torch.ones((len(block), max_len + 2), dtype=torch.long) + return labels, seqs, tokens + + return _convert + + +class DummyModel: + def __init__(self): + self.embed_dim = 16 + self.sent_to = None + + def eval(self): + return self + + def to(self, device): + self.sent_to = device + return self + + def __call__(self, batch_tokens, repr_layers, return_contacts): + layer = repr_layers[0] + bsz, tlen = batch_tokens.shape + reps = torch.arange(bsz * tlen * 4, dtype=torch.float32).reshape(bsz, tlen, 4) + return {"representations": {layer: reps}} + + +class TestProteinEncoder(unittest.TestCase): + def _make_pretrained_fn(self): + def _factory(): + return DummyModel(), DummyAlphabet() + + return _factory + + @patch("embkit.encoding.protein.esm.pretrained.esm2_t6_8M_UR50D") + def test_init_to_get_embed_dim_and_unknown_model(self, t6_mock): + t6_mock.side_effect = self._make_pretrained_fn() + + enc = ProteinEncoder(model="t6", batch_size=2, device=torch.device("cpu")) + self.assertEqual(enc.out_layer, 6) + self.assertEqual(enc.get_embed_dim(), 16) + self.assertEqual(enc.device.type, "cpu") + + enc.to(torch.device("cpu")) + self.assertEqual(enc.model.sent_to.type, "cpu") + + with self.assertRaises(Exception): + ProteinEncoder(model="bogus") + + @patch("embkit.encoding.protein.tqdm", side_effect=lambda x: x) + @patch("embkit.encoding.protein.esm.pretrained.esm2_t33_650M_UR50D") + def test_encode_modes_and_fix_len(self, t33_mock, _tqdm_mock): + t33_mock.side_effect = self._make_pretrained_fn() + enc = ProteinEncoder(model="t33", batch_size=2, device=None) + + data = [("p1", "AAAA"), ("p2", "AA")] + + out_sum = list(enc.encode(data, output="sum-pool", fix_len=None, verbose=False)) + self.assertEqual(len(out_sum), 2) + self.assertEqual(out_sum[0][0], "p1") + self.assertEqual(tuple(out_sum[0][1].shape), (4,)) + + out_mean = list(enc.encode(data, output="mean-pool", fix_len=None, verbose=True)) + self.assertEqual(len(out_mean), 2) + self.assertEqual(tuple(out_mean[1][1].shape), (4,)) + + out_vec = list(enc.encode(data, output="vector", fix_len=3, verbose=False)) + self.assertEqual(tuple(out_vec[0][1].shape), (3, 4)) + + @patch("embkit.encoding.protein.esm.pretrained.esm2_t12_35M_UR50D") + def test_pad(self, t12_mock): + t12_mock.side_effect = self._make_pretrained_fn() + enc = ProteinEncoder(model="t12", batch_size=1) + + tokens = torch.tensor([[1, 2, 3]], dtype=torch.long) + padded = enc.pad(tokens, 6) + self.assertEqual(tuple(padded.shape), (1, 6)) + self.assertTrue(torch.equal(padded[0, :3], torch.tensor([1, 2, 3]))) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/factory/test_layers_helpers.py b/tests/factory/test_layers_helpers.py index 185939b..0363a69 100644 --- a/tests/factory/test_layers_helpers.py +++ b/tests/factory/test_layers_helpers.py @@ -14,26 +14,12 @@ def __init__(self, out_features, in_features): self.out_features = out_features self.in_features = in_features - def gen_mask(self): + def gen_mask(self, in_features, out_features): # Return a mask matching expected shape - return np.ones((self.out_features, self.in_features), dtype=np.float32) + return np.ones((out_features, in_features), dtype=np.float32) class TestLayerHelpers(unittest.TestCase): - def test_idx_to_list_preserves_order(self): - mapping = {"c": 2, "a": 0, "b": 1} - ordered = layers.idx_to_list(mapping) - self.assertEqual(ordered, ["a", "b", "c"]) - - def test_build_features_to_group_mask_forward(self): - feature_map = {"G1": ["f0"], "G2": ["f1"]} - feature_idx = {"f0": 0, "f1": 1} - group_idx = {"G1": 0, "G2": 1} - mask = layers.build_features_to_group_mask( - feature_map, feature_idx, group_idx, group_node_count=1, forward=True - ) - expected = np.array([[1.0, 0.0], [0.0, 1.0]], dtype=np.float32) - np.testing.assert_array_equal(mask, expected) def test_layer_from_dict_and_gen_layer_linear(self): spec = { diff --git a/tests/factory/test_vae.py b/tests/factory/test_vae.py index ec71580..67d4a92 100644 --- a/tests/factory/test_vae.py +++ b/tests/factory/test_vae.py @@ -6,6 +6,9 @@ from pathlib import Path from embkit.models.vae.vae import VAE +from embkit.models.vae.net_vae import NetVAE +from embkit.constraints import PathwayConstraintInfo +from embkit.factory.layers import ConstraintInfo from embkit import factory @@ -36,3 +39,39 @@ def test_from_dict_without_layer_keys(self): self.assertEqual(model.features, ["a", "b", "c"]) self.assertEqual(model.latent_dim, 2) + def test_netvae_save_and_load_roundtrip(self): + features = ["G1", "G2", "G3"] + latent_groups = { + "TF1": ["G1", "G3"], + "TF2": ["G2"], + } + model = NetVAE(features=features, latent_groups=latent_groups, group_layer_size=[2, 1]) + + with tempfile.TemporaryDirectory() as temp_dir: + model_path = Path(temp_dir) / "netvae.pth" + factory.save(model, model_path) + loaded = factory.load(model_path) + + self.assertIsInstance(loaded, NetVAE) + self.assertEqual(loaded.features, features) + self.assertEqual(loaded.latent_groups, latent_groups) + self.assertEqual(loaded.group_layer_size, [2, 1]) + + def test_netvae_from_dict_rejects_deprecated_group_layer_scaling(self): + desc = { + "features": ["G1", "G2"], + "latent_groups": {"TF1": ["G1"], "TF2": ["G2"]}, + "group_layer_scaling": [3, 1], + } + with self.assertRaises(ValueError): + NetVAE.from_dict(desc) + + def test_constraintinfo_from_dict_pathway_dispatch(self): + payload = { + "op": "features-to-group", + "feature_map": {"TF1": ["G1"], "TF2": ["G2"]}, + "in_group_scaling": 1, + "out_group_scaling": 2, + } + constraint = ConstraintInfo.from_dict(payload) + self.assertIsInstance(constraint, PathwayConstraintInfo) diff --git a/tests/factory/test_verify.py b/tests/factory/test_verify.py new file mode 100644 index 0000000..5c0b2ca --- /dev/null +++ b/tests/factory/test_verify.py @@ -0,0 +1,92 @@ +import unittest +import torch +import os +import tempfile +import numpy as np +from embkit.models.vae.net_vae import NetVAE +from embkit.models.vae.rna_vae import RNAVAE +from embkit.models.ffnn import FFNN +from embkit.factory import save, run_model_verification + +class TestModelVerification(unittest.TestCase): + def setUp(self): + self.features = ["G1", "G2", "G3", "G4"] + self.latent_groups = {"P1": ["G1", "G2"], "P2": ["G3", "G4"]} + self.temp_dir = tempfile.TemporaryDirectory() + + def tearDown(self): + self.temp_dir.cleanup() + + def test_net_vae_verification_healthy(self): + model = NetVAE(features=self.features, latent_groups=self.latent_groups) + model.history = {"loss": [10.0, 5.0, 2.0]} + path = os.path.join(self.temp_dir.name, "netvae.model") + save(model, path) + + report = run_model_verification(path) + self.assertTrue(report["healthy"]) + self.assertEqual(report["model_type"], "NetVAE") + self.assertIn("history_summary", report) + + def test_net_vae_leakage_failure(self): + model = NetVAE(features=self.features, latent_groups=self.latent_groups) + model.history = {"loss": [10.0, 5.0, 2.0]} + # Poison a weight outside the mask + with torch.no_grad(): + module = model.encoder.net[0] + mask = module.mask + # Find a zero in the mask + idx = (mask == 0).nonzero(as_tuple=True) + if len(idx[0]) > 0: + # Add a significant weight where it should be zero + module.linear.weight[idx[0][0], idx[1][0]] = 1.0 + + # In-memory verifier should detect leakage. + in_memory_report = model.verify_integrity() + self.assertFalse(in_memory_report["healthy"]) + self.assertTrue(any("leakage" in i.lower() for i in in_memory_report["issues"])) + + # Serialization safety net clamps masked weights before save. + path = os.path.join(self.temp_dir.name, "leakage.model") + save(model, path) + report = run_model_verification(path) + self.assertTrue(report["healthy"]) + self.assertFalse(any("leakage" in i.lower() for i in report["issues"])) + + def test_rna_vae_verification(self): + model = RNAVAE(features=self.features, latent_dim=2) + model.history = {"loss": [1.0, 0.5]} + path = os.path.join(self.temp_dir.name, "rnavae.model") + save(model, path) + + report = run_model_verification(path) + self.assertTrue(report["healthy"]) + self.assertIn("rna_diagnostics", report) + self.assertTrue(report["rna_diagnostics"]["is_non_negative"]) + + def test_ffnn_history_failure(self): + model = FFNN(input_dim=4, output_dim=1) + # Loss went UP - failed to improve + model.history = {"loss": [0.1, 0.5]} + path = os.path.join(self.temp_dir.name, "bad_history.model") + save(model, path) + + report = run_model_verification(path) + self.assertFalse(report["healthy"]) + self.assertTrue(any("failed to improve" in i for i in report["issues"])) + + def test_corrupt_nan_verification(self): + model = FFNN(input_dim=4, output_dim=1) + with torch.no_grad(): + params = list(model.parameters()) + params[0][0] = float('nan') + + path = os.path.join(self.temp_dir.name, "nan.model") + save(model, path) + + report = run_model_verification(path) + self.assertFalse(report["healthy"]) + self.assertTrue(any("NaN" in issue for issue in report["issues"])) + +if __name__ == "__main__": + unittest.main() diff --git a/tests/files/test_h5.py b/tests/files/test_h5.py index 8214dfa..0f9680a 100644 --- a/tests/files/test_h5.py +++ b/tests/files/test_h5.py @@ -3,6 +3,7 @@ import numpy as np import torch from embkit.files import H5Writer, H5Reader +from embkit.files.h5 import H5CubeWriter, H5CubeReader import tempfile class TestH5(unittest.TestCase): @@ -59,3 +60,28 @@ def test_writer_set_row_by_name(self): if __name__ == "__main__": unittest.main() + + +class TestH5Extended(unittest.TestCase): + def test_writer_with_integer_columns(self): + with tempfile.TemporaryDirectory() as td: + path = os.path.join(td, "int_cols.h5") + writer = H5Writer(path, "g", ["r1", "r2"], 3) + writer.set_irow(0, [1.0, 2.0, 3.0]) + writer.set_irow(1, [4.0, 5.0, 6.0]) + writer.close() + + self.assertEqual(list(writer.columns), [0, 1, 2]) + + def test_h5_cube_roundtrip(self): + with tempfile.TemporaryDirectory() as td: + path = os.path.join(td, "cube.h5") + cube = H5CubeWriter(path, "g", ["a", "b"], xsize=2, ysize=2) + cube.set_row("a", np.array([[1.0, 2.0], [3.0, 4.0]], dtype="f")) + cube.set_irow(1, np.array([[5.0, 6.0], [7.0, 8.0]], dtype="f")) + cube.close() + + reader = H5CubeReader(path, "g") + self.assertEqual(reader.get_loc("b"), 1) + row_tensor, *_ = reader[0] + self.assertEqual(tuple(row_tensor.shape), (2, 2)) diff --git a/tests/files/test_large_csv_reader.py b/tests/files/test_large_csv_reader.py new file mode 100644 index 0000000..c4d282c --- /dev/null +++ b/tests/files/test_large_csv_reader.py @@ -0,0 +1,49 @@ +import os +import tempfile +import unittest + +import numpy as np + +from embkit.files.read_csv import LargeCsvReader + + +class TestLargeCsvReader(unittest.TestCase): + def setUp(self): + self.tmp = tempfile.TemporaryDirectory() + self.path = os.path.join(self.tmp.name, "x.tsv") + with open(self.path, "w", encoding="utf-8") as f: + f.write("id\tA\tB\n") + f.write("r1\t1\t2\n") + f.write("r2\t3\t4\n") + + def tearDown(self): + self.tmp.cleanup() + + def test_read_and_get_dict(self): + reader = LargeCsvReader(self.path, sep="\t", index_column="id", skip_header=False, save_index=False) + with reader: + row = reader.get("r1") + self.assertEqual(row, ["r1", "1", "2"]) + self.assertEqual(reader.get("missing"), None) + d = reader.get_dict("r2") + self.assertEqual(d["A"], "3") + self.assertEqual(d["B"], "4") + arr = list(reader.read(show_progress=False)) + self.assertEqual(len(arr), 2) + self.assertTrue(np.allclose(arr[0], np.array([1.0, 2.0], dtype=np.float32))) + + def test_iter_requires_context(self): + reader = LargeCsvReader(self.path, sep="\t", index_column=0, skip_header=False, save_index=False) + with self.assertRaises(RuntimeError): + list(reader) + + def test_load_existing_index(self): + reader = LargeCsvReader(self.path, sep="\t", index_column=0, skip_header=False, save_index=True) + self.assertTrue(os.path.exists(self.path + ".index")) + reader2 = LargeCsvReader(self.path, sep="\t", index_column=0, skip_header=False, save_index=True) + with reader2: + self.assertEqual(reader2.get("r2"), ["r2", "3", "4"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/files/test_loaders.py b/tests/files/test_loaders.py new file mode 100644 index 0000000..56cb7ea --- /dev/null +++ b/tests/files/test_loaders.py @@ -0,0 +1,69 @@ +import unittest +from pathlib import Path +from tempfile import TemporaryDirectory + +from embkit.files.loaders import load_gct, load_raw_hugo, load_gtex_hugo + + +class TestLoaders(unittest.TestCase): + def test_load_gct_transposes_and_drops_description(self): + with TemporaryDirectory() as td: + path = Path(td) / "x.gct" + path.write_text( + "#1.2\n" + "2\t2\n" + "Name\tDescription\tS1\tS2\n" + "ENSG1\tGene1\t1\t2\n" + "ENSG2\tGene2\t3\t4\n", + encoding="utf-8", + ) + + df = load_gct(path) + + self.assertEqual(list(df.index), ["S1", "S2"]) + self.assertEqual(list(df.columns), ["ENSG1", "ENSG2"]) + self.assertEqual(float(df.loc["S1", "ENSG1"]), 1.0) + self.assertEqual(float(df.loc["S2", "ENSG2"]), 4.0) + + def test_load_gct_with_nrows(self): + with TemporaryDirectory() as td: + path = Path(td) / "x.gct" + path.write_text( + "#1.2\n" + "2\t2\n" + "Name\tDescription\tS1\tS2\n" + "ENSG1\tGene1\t1\t2\n" + "ENSG2\tGene2\t3\t4\n", + encoding="utf-8", + ) + + df = load_gct(path, nrows=1) + + self.assertEqual(list(df.columns), ["ENSG1"]) + + def test_load_raw_hugo_and_load_gtex_hugo(self): + with TemporaryDirectory() as td: + hugo_path = Path(td) / "hugo.tsv" + hugo_path.write_text( + "id\tlocus_group\tsymbol\n" + "0\tprotein-coding gene\tTP53\n", + encoding="utf-8", + ) + gtex_hugo_path = Path(td) / "gtex.hugo.tsv" + gtex_hugo_path.write_text( + "sample\tTP53\tBRCA1\n" + "s1\t1\t2\n" + "s2\t3\t4\n", + encoding="utf-8", + ) + + raw = load_raw_hugo(hugo_path) + converted = load_gtex_hugo(gtex_hugo_path, nrows=1) + + self.assertIn("locus_group", raw.columns) + self.assertEqual(list(converted.index), ["s1"]) + self.assertEqual(list(converted.columns), ["TP53", "BRCA1"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/models/vae_models/test_encoder.py b/tests/models/vae_models/test_encoder.py index 592295e..aa5cb84 100644 --- a/tests/models/vae_models/test_encoder.py +++ b/tests/models/vae_models/test_encoder.py @@ -5,7 +5,7 @@ from embkit.models.vae.encoder import Encoder from embkit.factory.layers import Layer, LayerList from embkit.modules import MaskedLinear -from embkit.constraints import NetworkConstraint +from embkit.constraints import PathwayConstraintInfo class TestEncoder(unittest.TestCase): @@ -46,12 +46,17 @@ def test_encoder_with_masked_linear(self): def test_encoder_initializes_with_constraint_and_sets_mask(self): feature_index = ["f1", "f2"] - latent_index = ["z1", "z2"] # updated to 2 rows + latent_index = ["z1", "z2"] latent_membership = { "z1": ["f1"], - "z2": ["f2"] + "z2": ["f2"], } - constraint = NetworkConstraint(feature_index, latent_index, latent_membership) + constraint = PathwayConstraintInfo( + "features-to-group", + feature_map=latent_membership, + feature_index=feature_index, + group_index=latent_index, + ) layers = LayerList([Layer(units=2, op="masked_linear")]) # matches latent_index length enc = Encoder(feature_dim=2, latent_dim=2, layers=layers, constraint=constraint) @@ -113,4 +118,4 @@ def test_encoder_layer_batch_norm(self): if __name__ == "__main__": - unittest.main() \ No newline at end of file + unittest.main() diff --git a/tests/models/vae_models/test_net_vae.py b/tests/models/vae_models/test_net_vae.py index c01a969..105c4bb 100644 --- a/tests/models/vae_models/test_net_vae.py +++ b/tests/models/vae_models/test_net_vae.py @@ -1,13 +1,14 @@ import unittest +import torch + import numpy as np import pandas as pd -import torch from embkit.models.vae.net_vae import NetVAE from embkit.modules import MaskedLinear from embkit.losses import bce_with_logits - +from embkit.optimize import fit_vae, fit_net_vae class TestNetVAE(unittest.TestCase): def test_fit_applies_constraint_mask(self): @@ -22,23 +23,17 @@ def test_fit_applies_constraint_mask(self): "TF2": ["G2"], } - model = NetVAE(features=list(df.columns)) - model.fit( + model = NetVAE(features=list(df.columns), latent_groups=latent_groups, latent_index=latent_index) + fit_vae( + model, df, - latent_index=latent_index, - latent_groups=latent_groups, epochs=0, batch_size=2, - learning_rate=1e-3, + lr=1e-3, + loss=bce_with_logits, device="cpu", ) - constraint = model.encoder.constraint - self.assertIsNotNone(constraint) - - constraint.set_active(True) - model.encoder.refresh_mask(device=torch.device("cpu")) - masked_layers = [m for m in model.encoder.net if isinstance(m, MaskedLinear)] self.assertTrue(masked_layers) @@ -64,21 +59,17 @@ def test_masked_edges_remain_zero_during_training(self): "TF2": ["G2"], } - model = NetVAE(features=list(df.columns)) - model.fit( + model = NetVAE(features=list(df.columns), latent_groups=latent_groups, latent_index=latent_index, group_layer_size=[1,1]) + fit_vae( + model, df, - latent_index=latent_index, - latent_groups=latent_groups, epochs=0, batch_size=4, - learning_rate=1e-3, + lr=1e-3, + loss=bce_with_logits, device="cpu", ) - constraint = model.encoder.constraint - constraint.set_active(True) - model.encoder.refresh_mask(device=torch.device("cpu")) - masked_layer = next(m for m in model.encoder.net if isinstance(m, MaskedLinear)) opt = torch.optim.Adam(model.parameters(), lr=1e-3) @@ -102,6 +93,39 @@ def test_masked_edges_remain_zero_during_training(self): weight_after = weight.detach() self.assertTrue(torch.all(weight_after[mask == 0] == weight_before[mask == 0])) + def test_fit_net_vae_toggles_pathway_constraints(self): + df = pd.DataFrame( + np.random.rand(6, 3), + columns=["G1", "G2", "G3"], + ) + + latent_index = ["TF1", "TF2"] + latent_groups = { + "TF1": ["G1", "G3"], + "TF2": ["G2"], + } + + model = NetVAE(features=list(df.columns), latent_groups=latent_groups, latent_index=latent_index, group_layer_size=[1, 1]) + fit_net_vae( + model=model, + X=df, + latent_index=latent_index, + latent_groups=latent_groups, + learning_rate=1e-3, + batch_size=3, + epochs=1, + phases=[1, 1], # unconstrained then constrained + device="cpu", + ) + + model.set_constraint_active(True) + model.refresh_masks(torch.device("cpu")) + enc_masked = [m for m in model.encoder.net if isinstance(m, MaskedLinear)] + self.assertTrue(enc_masked) + constrained_mask = enc_masked[0].mask.cpu().numpy() + self.assertEqual(constrained_mask.shape, (2, 3)) + self.assertFalse(np.all(constrained_mask == 1.0)) + if __name__ == "__main__": unittest.main() diff --git a/tests/models/vae_models/test_rna_vae.py b/tests/models/vae_models/test_rna_vae.py new file mode 100644 index 0000000..3d5c1c3 --- /dev/null +++ b/tests/models/vae_models/test_rna_vae.py @@ -0,0 +1,34 @@ +import unittest + +import numpy as np +import pandas as pd +import torch + +from embkit.models.vae.rna_vae import RNAVAE + + +class TestRNAVAE(unittest.TestCase): + def test_fit_smoke(self): + df = pd.DataFrame( + np.random.rand(8, 6).astype(np.float32), + columns=[f"G{i}" for i in range(6)], + ) + + model = RNAVAE(features=list(df.columns), latent_dim=3, lr=1e-3) + history = model.fit( + df, + epochs=1, + batch_size=4, + kappa=1.0, + early_stopping_patience=2, + device=torch.device("cpu"), + progress=False, + ) + + self.assertIn("loss", history) + self.assertGreaterEqual(len(history["loss"]), 1) + self.assertEqual(model.latent_dim, 3) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/modules/test_tsp.py b/tests/modules/test_tsp.py new file mode 100644 index 0000000..90e4ce1 --- /dev/null +++ b/tests/modules/test_tsp.py @@ -0,0 +1,34 @@ +import unittest + +import torch + +from embkit.modules import TSPLayer + + +class TestTSPLayer(unittest.TestCase): + def test_soft_votes_shape(self): + layer = TSPLayer(pairs=[(0, 1), (2, 3)], beta=10.0, hard=False) + x = torch.tensor([[2.0, 1.0, 0.0, 3.0]]) + out = layer(x) + self.assertEqual(tuple(out.shape), (1, 2)) + self.assertTrue(torch.all((out > 0) & (out < 1))) + + def test_hard_votes_and_chunking(self): + layer = TSPLayer(pairs=[(0, 1), (2, 3), (1, 0)], hard=True, chunk_size=2) + x = torch.tensor([[2.0, 1.0, 0.0, 3.0]]) + out = layer(x) + self.assertEqual(tuple(out.shape), (1, 3)) + self.assertTrue(torch.equal(out, torch.tensor([[1.0, 0.0, 0.0]]))) + + def test_learnable_weight_aggregation(self): + layer = TSPLayer(pairs=[(0, 1), (1, 0)], hard=True, learnable_weights=True) + with torch.no_grad(): + layer.weights.copy_(torch.tensor([2.0, 3.0])) + x = torch.tensor([[2.0, 1.0]]) + out = layer(x) + self.assertEqual(tuple(out.shape), (1,)) + self.assertAlmostEqual(float(out.item()), 2.0, places=5) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/optimize/test_optimize_helpers.py b/tests/optimize/test_optimize_helpers.py new file mode 100644 index 0000000..714f2a8 --- /dev/null +++ b/tests/optimize/test_optimize_helpers.py @@ -0,0 +1,64 @@ +import unittest + +import pandas as pd +import torch +from torch import nn +from torch.utils.data import DataLoader, TensorDataset + +from embkit.optimize import ( + _move_to_device, + _resolve_phases, + fit, + fit_vae, +) +from embkit.losses import bce_with_logits +from embkit.models.vae.vae import VAE + + +class TinyModel(nn.Module): + def __init__(self): + super().__init__() + self.lin = nn.Linear(2, 1) + + def forward(self, x): + return self.lin(x) + + +class TestOptimizeHelpers(unittest.TestCase): + def test_move_to_device_nested(self): + x = torch.ones(2) + nested = {"a": [x, (x,)]} + out = _move_to_device(nested, torch.device("cpu")) + self.assertEqual(out["a"][0].device.type, "cpu") + self.assertEqual(out["a"][1][0].device.type, "cpu") + + def test_resolve_phases(self): + self.assertEqual(_resolve_phases(epochs=3, beta=0.5, beta_schedule=None), [(0.5, 3)]) + self.assertEqual(_resolve_phases(epochs=3, beta=0.5, beta_schedule=[(0.1, 2)]), [(0.1, 2)]) + + def test_fit_tensor_requires_y(self): + model = TinyModel() + with self.assertRaises(ValueError): + fit(model=model, X=torch.randn(3, 2), y=None, epochs=1, progress=False) + + def test_fit_vae_guards(self): + vae = VAE(features=["G1", "G2"], latent_dim=1) + df = pd.DataFrame([[0.1, 0.2], [0.2, 0.3]], columns=["G1", "G2"]) + + with self.assertRaises(ValueError): + fit_vae(vae, df, epochs=1, loss=None, progress=False) + + bad_df = pd.DataFrame([[0.1, 0.2]], columns=["X1", "X2"]) + with self.assertRaises(ValueError): + fit_vae(vae, bad_df, epochs=1, loss=bce_with_logits, progress=False) + + def test_fit_vae_accepts_dataloader(self): + vae = VAE(features=["G1", "G2"], latent_dim=1) + x = torch.tensor([[0.1, 0.2], [0.2, 0.3]], dtype=torch.float32) + loader = DataLoader(TensorDataset(x), batch_size=1, shuffle=False) + out = fit_vae(vae, loader, epochs=1, loss=bce_with_logits, progress=False) + self.assertIsInstance(out, float) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/factory/test_layer_helpers.py b/tests/pathway/test_layer_helpers.py similarity index 59% rename from tests/factory/test_layer_helpers.py rename to tests/pathway/test_layer_helpers.py index f162975..b926bce 100644 --- a/tests/factory/test_layer_helpers.py +++ b/tests/pathway/test_layer_helpers.py @@ -4,13 +4,13 @@ import numpy as np import pandas as pd -from embkit.factory import layers +from embkit import pathway class TestMaskHelpers(unittest.TestCase): def test_idx_to_list_orders_correctly(self): mapping = {"b": 0, "a": 1, "c": 2} - ordered = layers.idx_to_list(mapping) + ordered = pathway.idx_to_list(mapping) self.assertEqual(ordered, ["b", "a", "c"]) # positions 0,1,2 def test_build_features_to_group_mask_forward(self): @@ -18,7 +18,7 @@ def test_build_features_to_group_mask_forward(self): feature_map = {"G1": ["f0"], "G2": ["f1"]} feature_idx = {"f0": 0, "f1": 1} group_idx = {"G1": 0, "G2": 1} - mask = layers.build_features_to_group_mask( + mask = pathway.build_features_to_group_mask( feature_map, feature_idx, group_idx, group_node_count=1, forward=True ) expected = np.array([[1.0, 0.0], [0.0, 1.0]], dtype=np.float32) @@ -28,7 +28,7 @@ def test_build_features_to_group_mask_reverse(self): feature_map = {"G1": ["f0"], "G2": ["f1"]} feature_idx = {"f0": 0, "f1": 1} group_idx = {"G1": 0, "G2": 1} - mask = layers.build_features_to_group_mask( + mask = pathway.build_features_to_group_mask( feature_map, feature_idx, group_idx, group_node_count=1, forward=False ) # reverse shape (in_features, out_features) @@ -36,5 +36,21 @@ def test_build_features_to_group_mask_reverse(self): np.testing.assert_array_equal(mask, expected) +class TestLayerHelpers(unittest.TestCase): + def test_idx_to_list_preserves_order(self): + mapping = {"c": 2, "a": 0, "b": 1} + ordered = pathway.idx_to_list(mapping) + self.assertEqual(ordered, ["a", "b", "c"]) + + def test_build_features_to_group_mask_forward(self): + feature_map = {"G1": ["f0"], "G2": ["f1"]} + feature_idx = {"f0": 0, "f1": 1} + group_idx = {"G1": 0, "G2": 1} + mask = pathway.build_features_to_group_mask( + feature_map, feature_idx, group_idx, group_node_count=1, forward=True + ) + expected = np.array([[1.0, 0.0], [0.0, 1.0]], dtype=np.float32) + np.testing.assert_array_equal(mask, expected) + if __name__ == "__main__": unittest.main() diff --git a/tests/pathway/test_pathway_mask.py b/tests/pathway/test_pathway_mask.py index f78f451..6abaf48 100644 --- a/tests/pathway/test_pathway_mask.py +++ b/tests/pathway/test_pathway_mask.py @@ -3,21 +3,23 @@ import numpy as np -from embkit.pathway import build_sif_mask +from embkit.pathway import build_mask, extract_sif_interactions class TestPathwayMask(unittest.TestCase): def test_build_sif_mask_filters_relation_and_indices(self): sif_path = Path(__file__).resolve().parents[1] / "data" / "sample_pathway.sif" + sif_data = extract_sif_interactions(str(sif_path)) + + src_index = {"TF1": 0, "TF2": 1} dst_index = {"G1": 0, "G2": 1, "G3": 2} - mask = build_sif_mask( - str(sif_path), + mask = build_mask( + sif_data, src_index, dst_index, - relation="controls-expression-of", ) expected = np.array( diff --git a/tests/resources/test_c_bio_portal.py b/tests/resources/test_c_bio_portal.py index f5eaa36..c87a5aa 100644 --- a/tests/resources/test_c_bio_portal.py +++ b/tests/resources/test_c_bio_portal.py @@ -2,9 +2,9 @@ from unittest.mock import patch, MagicMock from pathlib import Path import tempfile +import os from requests.exceptions import RequestException import tarfile -import shutil from embkit.resources import CBIOPortal from embkit.resources.resource import REPO_DIR @@ -132,14 +132,12 @@ def test_download_warns_if_already_called_from_init(self, mock_get): dataset.download() # triggers the warning path def test_default_save_path_creation(self): - default_path = Path.home() / REPO_DIR - if default_path.exists(): - shutil.rmtree(default_path) - self.assertFalse(default_path.exists()) - - dataset = CBIOPortal(study_id=self.study_id, save_path=None, download=False) - self.assertTrue(Path(dataset.save_path).exists()) - self.assertEqual(Path(dataset.save_path).resolve(), default_path.resolve()) + with tempfile.TemporaryDirectory() as tmpdir: + default_path = Path(tmpdir) / REPO_DIR + with patch.dict(os.environ, {"EMBKIT_HOME": str(default_path)}): + dataset = CBIOPortal(study_id=self.study_id, save_path=None, download=False) + self.assertTrue(Path(dataset.save_path).exists()) + self.assertEqual(Path(dataset.save_path).resolve(), default_path.resolve()) # ✅ Covers lines 69–70 (target_file and unpacked_folder path logic) @patch("embkit.resources.c_bio_portal.requests.get") @@ -164,29 +162,27 @@ def test_download_resolves_expected_paths(self, mock_get): @patch("embkit.resources.c_bio_portal.requests.get") def test_default_embkit_path_created_if_missing(self, mock_get): - default_path = Path.home() / REPO_DIR + with tempfile.TemporaryDirectory() as tmpdir: + default_path = Path(tmpdir) / REPO_DIR + self.assertFalse(default_path.exists()) - # Clean up before test - if default_path.exists(): - shutil.rmtree(default_path) - self.assertFalse(default_path.exists()) + # Setup mock for download + mock_response = MagicMock() + mock_response.iter_content = lambda chunk_size: [b"x" * 10] + mock_response.headers = {"Content-Length": "10"} + mock_response.status_code = 200 + mock_response.raise_for_status = MagicMock() + mock_get.return_value = mock_response - # Setup mock for download - mock_response = MagicMock() - mock_response.iter_content = lambda chunk_size: [b"x" * 10] - mock_response.headers = {"Content-Length": "10"} - mock_response.status_code = 200 - mock_response.raise_for_status = MagicMock() - mock_get.return_value = mock_response + with patch.dict(os.environ, {"EMBKIT_HOME": str(default_path)}): + dataset = CBIOPortal(study_id=self.study_id, save_path=None, download=False) + self.assertEqual(Path(dataset.save_path).resolve(), default_path.resolve()) - # Run - dataset = CBIOPortal(study_id=self.study_id, save_path=None, download=False) - self.assertEqual(Path(dataset.save_path).resolve(), default_path.resolve()) + # Should be created in download + dataset.download() - # Should be created in download - dataset.download() - self.assertTrue(default_path.exists()) - self.assertTrue((default_path / f"{self.study_id}.tar.gz").exists()) + self.assertTrue(default_path.exists()) + self.assertTrue((default_path / f"{self.study_id}.tar.gz").exists()) @patch("embkit.resources.c_bio_portal.requests.get") @patch("embkit.resources.c_bio_portal.logger") @@ -261,4 +257,4 @@ def test_unpack_returns_path_if_skipped_due_to_existing_files(self): self.assertEqual(dataset.unpacked_file_path.resolve(), unpacked_folder.resolve()) if __name__ == "__main__": - unittest.main() \ No newline at end of file + unittest.main() diff --git a/tests/resources/test_resource.py b/tests/resources/test_resource.py index 5a8bed3..76c2c9d 100644 --- a/tests/resources/test_resource.py +++ b/tests/resources/test_resource.py @@ -1,6 +1,7 @@ import unittest import tempfile import logging +import os from pathlib import Path from embkit.resources.resource import Resource, REPO_DIR @@ -44,10 +45,12 @@ def test_download_exception_logged(self): self.assertFalse(dataset._download_called_from_init) def test_none_save_path_with_download_skips_creation(self): - dataset = DummyResource(name="dummy", save_path=None, download=False) - expected_path = Path.home() / REPO_DIR - self.assertEqual(Path(dataset.save_path).resolve(), expected_path.resolve()) + with tempfile.TemporaryDirectory() as tmpdir: + home = Path(tmpdir) / "embkit-home" + with unittest.mock.patch.dict(os.environ, {"EMBKIT_HOME": str(home)}): + dataset = DummyResource(name="dummy", save_path=None, download=False) + self.assertEqual(Path(dataset.save_path).resolve(), home.resolve()) if __name__ == '__main__': - unittest.main() \ No newline at end of file + unittest.main() diff --git a/tests/resources/test_resource_and_gtex.py b/tests/resources/test_resource_and_gtex.py index 6494e55..65e2bb1 100644 --- a/tests/resources/test_resource_and_gtex.py +++ b/tests/resources/test_resource_and_gtex.py @@ -18,11 +18,12 @@ def download(self): class TestDatasetBase(unittest.TestCase): def test_default_save_path_created(self): - ds = DummyDataset(name="test", save_path=None, download=False) - # Should create a .embkit directory in home - default_dir = Path(Path.home(), ".embkit") - self.assertTrue(ds.save_path == default_dir) - self.assertTrue(ds.save_path.is_dir()) + with tempfile.TemporaryDirectory() as tmpdir: + embkit_home = Path(tmpdir) / "embkit-home" + with mock.patch.dict(os.environ, {"EMBKIT_HOME": str(embkit_home)}): + ds = DummyDataset(name="test", save_path=None, download=False) + self.assertEqual(ds.save_path.resolve(), embkit_home.resolve()) + self.assertTrue(ds.save_path.is_dir()) def test_custom_save_path_created(self): with tempfile.TemporaryDirectory() as td: @@ -31,10 +32,13 @@ def test_custom_save_path_created(self): self.assertTrue(ds.save_path.is_dir()) def test_str_representation(self): - ds = DummyDataset(name="test", save_path=None, download=False) - s = str(ds) - self.assertIn(str(ds.save_path), s) - self.assertIn("Unpacked file", s) + with tempfile.TemporaryDirectory() as tmpdir: + embkit_home = Path(tmpdir) / "embkit-home" + with mock.patch.dict(os.environ, {"EMBKIT_HOME": str(embkit_home)}): + ds = DummyDataset(name="test", save_path=None, download=False) + s = str(ds) + self.assertIn(str(ds.save_path), s) + self.assertIn("Unpacked file", s) class TestGTExDownloader(unittest.TestCase): diff --git a/tests/resources/test_resource_base.py b/tests/resources/test_resource_base.py index abe394f..7de5206 100644 --- a/tests/resources/test_resource_base.py +++ b/tests/resources/test_resource_base.py @@ -2,6 +2,7 @@ import unittest import tempfile +import os from pathlib import Path from unittest import mock @@ -16,11 +17,12 @@ def download(self): class TestDatasetBase(unittest.TestCase): def test_default_save_path_created(self): - # When save_path is None, a default .embkit dir under HOME should be created - ds = DummyDataset(name="tmp", save_path=None, download=False) - self.assertTrue(ds.save_path.exists()) - # Ensure it is a directory - self.assertTrue(ds.save_path.is_dir()) + with tempfile.TemporaryDirectory() as td: + embkit_home = Path(td) / "embkit-home" + with mock.patch.dict(os.environ, {"EMBKIT_HOME": str(embkit_home)}): + ds = DummyDataset(name="tmp", save_path=None, download=False) + self.assertTrue(ds.save_path.exists()) + self.assertTrue(ds.save_path.is_dir()) def test_custom_save_path(self): with tempfile.TemporaryDirectory() as td: