embkit 0.4__tar.gz → 0.5__tar.gz

embkit-0.5/PKG-INFO

@@ -0,0 +1,89 @@
+Metadata-Version: 2.4
+Name: embkit
+Version: 0.5
+Summary: Embedding Kit for embedding models
+Author-email: Kyle Ellrott <ellrott@ohsu.edu>, Raphael Kirchgaessner <kirchgae@ohsu.edu>
+License-Expression: MIT
+Keywords: embeddings,machine learning,deep learning
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: torch>=2.0
+Requires-Dist: pandas
+Requires-Dist: numpy
+Requires-Dist: click
+Requires-Dist: scikit-learn
+Requires-Dist: scipy
+Requires-Dist: gripql
+Requires-Dist: tqdm
+Requires-Dist: biopython
+Requires-Dist: fair-esm
+Requires-Dist: faiss-cpu
+Requires-Dist: h5py
+
+
+# Embedding Kit
+
+Methods for data normalization, embedding, synthesis and transformation in computational biology workflows.
+
+Embedding Kit (embkit) is a toolkit for building and applying embedding models. It combines:
+
+- CLI commands for repeatable model training and encoding
+- PyTorch VAE model components for custom pipelines
+- Utilities for loading, normalizing, and aligning large molecular datasets
+- ESM2-based protein sequence embeddings
+
+Use Embedding Kit when you want to move from tabular molecular data (RNA-seq, proteomics, methylation) to trainable latent representations that can be reused for downstream analysis.
+
+## Features
+
+- **VAE & NetVAE training** — Train variational autoencoders, including pathway-constrained NetVAE models, from tabular or HDF5 matrices.
+- **Beta-KL scheduling** — Schedule the KL-regularization weight across training epochs for stable convergence.
+- **Protein embeddings** — Generate sequence embeddings from FASTA files using ESM2 models.
+- **Normalization utilities** — Min-max and exponential min-max normalization for expression matrices.
+- **Device auto-detection** — Runs on CPU, CUDA, or Apple Metal GPUs.
+- **Python API** — Composable layers, losses, and model factories for custom pipelines beyond the CLI.
+
+## Installation
+
+```bash
+pip install embkit
+```
+
+
+## Training in Python
+
+```python
+from embkit import dataframe_loader
+from embkit.models.vae import VAE
+from embkit.factory.layers import Layer
+from embkit.losses import BCEWithLogitsVAELoss
+from embkit.factory import save, load
+from embkit import optimize
+
+loader = dataframe_loader(df_norm, batch_size=256)
+
+vae = VAE(
+    features=list(df_norm.columns),
+    latent_dim=128,
+    encoder_layers=[Layer(512, activation="relu"), Layer(256, activation="relu")],
+    decoder_layers=[Layer(512, activation="relu")],
+)
+
+optimize.fit_vae(vae, X=loader, epochs=60, lr=1e-3, loss=BCEWithLogitsVAELoss())
+
+save(vae, "vae.model")
+```
+
+## Documentation
+
+Full documentation, including core concepts, the training guide, the CLI reference, and API docs, is available in the [docs folder](docs/). To build it locally, see [DEV.md](DEV.md).
+
+## License
+
+MIT
+
+## Authors
+
+- Kyle Ellrott — ellrott@ohsu.edu
+- Raphael Kirchgaessner — kirchgae@ohsu.edu
+

embkit-0.5/README.md

@@ -0,0 +1,67 @@
+
+# Embedding Kit
+
+Methods for data normalization, embedding, synthesis and transformation in computational biology workflows.
+
+Embedding Kit (embkit) is a toolkit for building and applying embedding models. It combines:
+
+- CLI commands for repeatable model training and encoding
+- PyTorch VAE model components for custom pipelines
+- Utilities for loading, normalizing, and aligning large molecular datasets
+- ESM2-based protein sequence embeddings
+
+Use Embedding Kit when you want to move from tabular molecular data (RNA-seq, proteomics, methylation) to trainable latent representations that can be reused for downstream analysis.
+
+## Features
+
+- **VAE & NetVAE training** — Train variational autoencoders, including pathway-constrained NetVAE models, from tabular or HDF5 matrices.
+- **Beta-KL scheduling** — Schedule the KL-regularization weight across training epochs for stable convergence.
+- **Protein embeddings** — Generate sequence embeddings from FASTA files using ESM2 models.
+- **Normalization utilities** — Min-max and exponential min-max normalization for expression matrices.
+- **Device auto-detection** — Runs on CPU, CUDA, or Apple Metal GPUs.
+- **Python API** — Composable layers, losses, and model factories for custom pipelines beyond the CLI.
+
+## Installation
+
+```bash
+pip install embkit
+```
+
+
+## Training in Python
+
+```python
+from embkit import dataframe_loader
+from embkit.models.vae import VAE
+from embkit.factory.layers import Layer
+from embkit.losses import BCEWithLogitsVAELoss
+from embkit.factory import save, load
+from embkit import optimize
+
+loader = dataframe_loader(df_norm, batch_size=256)
+
+vae = VAE(
+    features=list(df_norm.columns),
+    latent_dim=128,
+    encoder_layers=[Layer(512, activation="relu"), Layer(256, activation="relu")],
+    decoder_layers=[Layer(512, activation="relu")],
+)
+
+optimize.fit_vae(vae, X=loader, epochs=60, lr=1e-3, loss=BCEWithLogitsVAELoss())
+
+save(vae, "vae.model")
+```
+
+## Documentation
+
+Full documentation, including core concepts, the training guide, the CLI reference, and API docs, is available in the [docs folder](docs/). To build it locally, see [DEV.md](DEV.md).
+
+## License
+
+MIT
+
+## Authors
+
+- Kyle Ellrott — ellrott@ohsu.edu
+- Raphael Kirchgaessner — kirchgae@ohsu.edu
+

{embkit-0.4 → embkit-0.5}/pyproject.toml

@@ -4,11 +4,11 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "embkit"
-version = "0.4"
+version = "0.5"
 description = "Embedding Kit for embedding models"
 readme = "README.md"
 requires-python = ">=3.8"
-license = { text = "MIT" }
+license = "MIT"
 authors = [
     { name = "Kyle Ellrott", email = "ellrott@ohsu.edu" },
     { name = "Raphael Kirchgaessner", email = "kirchgae@ohsu.edu" }

{embkit-0.4 → embkit-0.5}/setup.py

@@ -2,7 +2,7 @@ from setuptools import setup, find_packages
 
 setup(
    name='embedding-kit',
-   version='0.4',
+   version='0.5',
    packages=find_packages(),
    install_requires=[
        # List your dependencies here!

{embkit-0.4 → embkit-0.5}/src/embkit/align.py

@@ -58,8 +58,7 @@ def matrix_spearman_alignment_linear(a, b, cutoff=0.0):
 
 def procrustes(X, Y):
     """
-    Computes the Procrustes transformation (optimal PURE rotation)
-    to align X to Y, correcting for reflections.
+    Computes the Procrustes transformation (optimal PURE rotation) to align X to Y, correcting for reflections.
 
     Args:
         X: The first matrix (N_points, N_dims).
@@ -67,6 +66,7 @@ def procrustes(X, Y):
 
     Returns:
         R: The optimal rotation matrix (guaranteed det(R) = +1).
+        S: Singular values of X.T @ Y
     """
     assert X.shape == Y.shape
 
@@ -74,37 +74,65 @@ def procrustes(X, Y):
 
     U, S, Vt = np.linalg.svd(M)
 
-    R = np.dot(U, Vt)
+    R = np.dot(U, Vt) # R can be optimized indep. of scalar
 
     if np.linalg.det(R) < 0:
         Vt_corrected = Vt.copy()
         Vt_corrected[-1, :] *= -1
         R = np.dot(U, Vt_corrected)
 
-    return R
-
+    return R,S
 
 def procrustes_scale(X, Y):
     """
-    Compute the procrustes transformation (R), then compute a factor (k) to rescale the src matrix. 
+    Compute the procrustes transformation (R), then a single (scalar) scale factor k = ||Y||_F / ||X||_F (equivalently ||Y||_F / ||X @ R||_F, since rotation preserves Frobenius norm). 
+    Note: Scaling is calculated without Least-Squares, unlike more traditional approaches, to match Y's spread directly and avoid shrinking under noisy correspondances.
     
-    To apply transformtion:
+    To apply transformation:
     src.dot(R) * k
 
-    
     Args:
         X: The first matrix (N_points, N_dims).
         Y: The second matrix (N_points, N_dims).
 
     Returns:
         R: The optimal rotation matrix (guaranteed det(R) = +1).
-        k: Scaling factor
+        k: Scalar scaling factor, ||Y||_F / ||X||_F
+        S: Singular values of X.T @ Y
 
     """
-    R = procrustes(X,Y)
-    A = np.array(X.dot(R))
-    B = np.array(Y)
-    numerators = np.sum(A * B, axis=0)
-    denominators = np.sum(A * A, axis=0)
-    k = np.divide(numerators, denominators, out=np.zeros_like(numerators), where=denominators!=0)
-    return R, k
+    R, S = procrustes(X, Y)
+    Xnorm = np.linalg.norm(X)
+    k = np.linalg.norm(Y) / Xnorm if Xnorm != 0 else 0.0 # Unlike Least Squares (superceded), we dont bias towards shrinking under noisy correspondence.
+    return R, k, S
+
+def procrustes_scale_centered(X, Y):
+    """
+    Mean-centers X and Y, fits the rotation, then scales: k = ||Yc||_F / ||Xc||_F (equivalently ||Yc||_F / ||Xc @ R||_F, since rotation preserves Frobenius norm).
+    Note: Scaling is calculated without Least-Squares, unlike more traditional approaches, to match Y's spread directly and avoid shrinking under noisy correspondances.
+
+    To apply transformation:
+    (src - Xmean).dot(R) * k + Ymean
+
+    Args:
+        X: The first matrix (N_points, N_dims).
+        Y: The second matrix (N_points, N_dims).
+
+    Returns:
+        R: The optimal rotation matrix (guaranteed det(R) = +1).
+        k: Scalar scaling factor, ||Yc||_F / ||Xc||_F
+        Xmean: Per-dimension mean of X (shape: (N_dims,)), subtract before applying R/k
+        Ymean: Per-dimension mean of Y (shape: (N_dims,)), add back after applying R/k
+        S: Singular values of (X - Xmean).T @ (Y - Ymean)
+
+    """
+    Xmean = np.array(X).mean(axis=0)
+    Ymean = np.array(Y).mean(axis=0)
+    Xc = np.array(X) - Xmean
+    Yc = np.array(Y) - Ymean
+    R, k, S = procrustes_scale(Xc, Yc)
+    return R, k, Xmean, Ymean, S
+
+
+
+    

embkit-0.5/src/embkit/c_bio/api.py

@@ -0,0 +1,31 @@
+import requests
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class CBIOAPI:
+    """
+    Small client for the cBioPortal public API.
+
+    cBioPortal is a public catalog of cancer genomics studies. This helper
+    wraps the lightweight metadata endpoints (like listing available studies)
+    that don't require downloading study data — use the
+    :class:`~embkit.resources.c_bio_portal.CBIOPortal` resource for fetching an
+    actual study's molecular profile data.
+    """
+
+    @staticmethod
+    def list_studies():
+        """
+        List the cBioPortal studies available for download.
+
+        Returns:
+            The API's JSON list of studies, or ``None`` if the request failed.
+        """
+        try:
+            studies = requests.get("https://www.cbioportal.org/api/studies").json()
+            return studies
+        except requests.RequestException as e:
+            logger.error("Error fetching studies: %s", e)
+            return None

embkit-0.5/src/embkit/commands/matrix.py

@@ -0,0 +1,97 @@
+"""
+embkit matrix — preprocessing of expression/count matrices.
+
+The ``matrix`` command group prepares the tab-separated matrices that become the
+input to an embedding model. Its subcommands implement the "normalize" and "pca"
+major operations: rescaling raw feature values to a comparable ``[0, 1]`` range,
+and projecting a matrix onto its leading principal components. These are the
+data-preparation steps you run before (and independently of) training.
+"""
+import click
+import pandas as pd
+from pathlib import Path
+from ..utilities.pca import run_pca
+
+
+matrix = click.Group(name="matrix", help="Model commands.")
+
+@matrix.command()
+@click.argument("srcs", nargs=-1)
+@click.option("--out", "-o", required=True)
+@click.option("--features", default=None, type=click.Path(exists=True, dir_okay=False, readable=True, path_type=str))
+@click.option("--col-quantile", "col_quantile", default=False, is_flag=True)
+@click.option("--quantile-max", "quantile_max", default=0.9, type=float)
+@click.option("--precision", "-p", type=int, default=5)
+def normalize(srcs, out, features, col_quantile, quantile_max, precision):
+    """
+    embkit matrix normalize — rescale one or more matrices for embedding training.
+
+    Reads tab-separated expression/count matrices, optionally subsets them to a
+    list of features, concatenates them, and writes a single normalized matrix in
+    which each sample's values fall in ``[0, 1]``. This is the preprocessing step
+    that puts the input on a comparable scale before a VAE / NetVAE is trained.
+
+    By default each *sample* (row) is divided by its own ``quantile_max``
+    quantile; passing ``--col-quantile`` divides each *feature* (column) by its
+    quantile instead.
+
+    Args:
+        srcs: One or more input TSV matrices (index column required).
+        out: Path to write the normalized TSV.
+        features: Optional file listing feature names to keep.
+        col_quantile: Normalize per feature (column) rather than per sample.
+        quantile_max: The quantile used as the divisor (default 0.9).
+        precision: Number of decimals to round to in the output.
+    """
+
+    features_list = None
+    if features is not None:
+        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_list is not None:
+            df = df[features_list]
+        dfs.append(df)
+    if len(dfs) == 0:
+        click.echo("No matrices defined")
+        return
+    df = pd.concat( dfs )
+
+    if not col_quantile:
+        normDF = (df.transpose() / df.quantile(quantile_max, axis=1)).transpose().clip(upper=1.0, lower=0.0).fillna(0.0)
+    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):
+    """
+    embkit matrix pca — project a matrix onto its leading principal components.
+
+    Runs principal component analysis on a tab-separated matrix and writes the
+    first ``pca_size`` components to a TSV. Useful for reducing dimensionality
+    before downstream analysis or visualization (e.g. inspecting whether a
+    normalized matrix separates samples as expected).
+
+    If no ``--out`` is given, the output defaults to ``<input_stem>.pca.tsv``.
+
+    Args:
+        input_path: Input TSV matrix.
+        pca_size: Number of principal components to retain (must be > 0).
+        out: Output TSV path (optional; defaults to ``<input_stem>.pca.tsv``).
+    """
+    if pca_size <= 0:
+        raise click.BadParameter("--pca-size must be a positive integer.")
+
+    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}")

{embkit-0.4 → embkit-0.5}/src/embkit/commands/model.py

@@ -13,11 +13,11 @@ from ..files import H5Reader
 from ..factory import save, load
 from ..factory.layers import Layer, LayerList
 from ..optimize import fit_vae
-from ..models.vae.vae import VAE
+from ..models.vae.vae import BaseVAE
 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 ..losses import get_vae_loss
 from ..pathway import extract_sif_interactions, feature_map_intersect, build_feature_map_indices
 
 model = click.Group(name="model", help="VAE Model commands.")
@@ -37,6 +37,7 @@ model = click.Group(name="model", help="VAE Model commands.")
 @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("--save-stats", is_flag=True)
+@click.option("--save-latent", is_flag=True)
 @click.option("--zero-mask", default=None, type=float)
 @click.option("--seed", default=42, type=int)
 @click.option("--bfloat16", is_flag=True)
@@ -54,6 +55,7 @@ def train_vae(input_path: str,
               schedule:str,
               zero_mask: float,
               save_stats: bool,
+              save_latent:bool,
               seed: int,
               bfloat16: bool,
               sampling: bool
@@ -100,12 +102,11 @@ def train_vae(input_path: str,
     
     dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
     
-
     layer_sizes = list( int(i) for i in encode_layers.split(",") )
     enc_layers_list = LayerList( layer_sizes )
 
     layer_sizes = list( int(i) for i in decode_layers.split(",") )
-    dec_layers_list = LayerList( layer_sizes, end_activation=final_activation )
+    dec_layers_list = LayerList( layer_sizes, output_activation=final_activation )
 
     beta_schedule = None
     if schedule is not None:
@@ -113,18 +114,14 @@ def train_vae(input_path: str,
         for b in schedule.split(","):
             e, b = b.split(":")
             beta_schedule.append( (float(b), int(e)) )
-    vae = VAE(features=features,
+    vae = BaseVAE(features=features,
               latent_dim=latent,
               encoder_layers=enc_layers_list,
               decoder_layers=dec_layers_list,
               sampling=sampling,
               device=device, dtype=dtype)
 
-    loss_func = bce_with_logits
-    if loss == "mse":
-        loss_func = mse
-    elif loss == "bce":
-        loss_func = bce
+    loss_func = get_vae_loss(loss)
 
     fit_vae(vae, dataloader, epochs=epochs,
             beta_schedule=beta_schedule, lr=learning_rate, loss=loss_func)
@@ -143,6 +140,46 @@ def train_vae(input_path: str,
         stats.to_csv(stats_path, sep="\t")
         click.echo(f"Stats saved, to {stats_path}")
 
+    if save_latent:
+        vae.eval() # Makes output deterministic. Otherwise, model in training.
+        losses_df = pd.DataFrame({
+            "epoch": range(1, len(vae.history["loss"]) + 1),
+            "loss": vae.history["loss"],
+            "recon": vae.history["recon"],
+            "kl": vae.history["kl"]})
+        
+        losses_df.to_csv(f"{out}.losses_stats.tsv", sep="\t", index=False)
+        click.echo(f"Training losses saved, to {out}.losses_stats.tsv")
+
+        exportloader = DataLoader(dataset, batch_size=batch_size, shuffle=False)
+        all_mu = []
+        all_logvar = []
+
+        with torch.no_grad():
+            for batch in exportloader:
+                    x_tensor = batch[0] if isinstance(batch, (tuple, list)) else batch
+                    x_tensor = x_tensor.to(device=device, dtype=dtype)
+                    _, mu, logvar, _ = vae(x_tensor)
+                    all_mu.append(mu.cpu())
+                    all_logvar.append(logvar.cpu())
+
+        mu = torch.cat(all_mu, dim=0)
+        logvar = torch.cat(all_logvar, dim=0)
+        std = torch.sqrt(torch.exp(logvar)) # sigma
+
+        index = df.index if df is not None else None
+
+        mu_np = mu.to(device="cpu", dtype=torch.float32).detach().numpy()
+        std_np = std.to(device="cpu", dtype=torch.float32).detach().numpy()
+        mu_df = pd.DataFrame(mu_np, index=index, columns=[f"mu_{i}" for i in range(mu.shape[1])])
+        std_df = pd.DataFrame(std_np, index=index, columns=[f"std_{i}" for i in range(std.shape[1])])
+        mu_path = f"{out}.latent_mu.tsv"
+        std_path = f"{out}.latent_std.tsv"
+        mu_df.to_csv(mu_path, sep="\t")
+        std_df.to_csv(std_path, sep="\t")
+        click.echo(f"Latent mu saved, to {mu_path}")
+        click.echo(f"Latent std (sigma) saved, to {std_path}")
+
 
 @model.command()
 @click.argument("input_path", type=click.Path(exists=True, dir_okay=False, readable=True, path_type=str))
@@ -153,7 +190,7 @@ def train_vae(input_path: str,
 @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("--min-group-size", type=int, default=2, 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)
@@ -197,11 +234,7 @@ def train_netvae(input_path: str, pathway_sif:str, out:str,
     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":
-        loss_func = mse
-    elif loss == "bce":
-        loss_func = bce
+    loss_func = get_vae_loss(loss)
 
     if beta_schedule is None:
         beta_schedule = [(0.0, epochs)]
@@ -243,7 +276,7 @@ def encode(input_path: str, model_path:str, normalize:str, out:str):
     m.to(get_device())
     result = m.encoder(df_tensor)
     
-    matrix = result[0].detach().cpu().numpy()
+    matrix = result[2].detach().cpu().numpy()
     out_df = pd.DataFrame(matrix, index=df.index)
     out_df.to_csv(out, sep="\t")
 

{embkit-0.4 → embkit-0.5}/src/embkit/constraints/pathway_constraint.py

@@ -2,6 +2,7 @@ from typing import Any, Dict, List, Literal, Optional
 
 import numpy as np
 
+from ..factory import nn_module
 from ..factory.layers import ConstraintInfo
 from ..pathway import (
     _normalize_index,
@@ -12,7 +13,7 @@ from ..pathway import (
 
 ConstraintOP = Literal["features-to-group", "group-to-features", "group-to-group"]
 
-
+@nn_module
 class PathwayConstraintInfo(ConstraintInfo):
     """
     ConstraintInfo for pathway-based masking.