embkit 0.1.0__tar.gz

embkit-0.1.0/PKG-INFO

@@ -0,0 +1,42 @@
+Metadata-Version: 2.4
+Name: embkit
+Version: 0.1.0
+Summary: Embedding Kit for embedding models
+Author-email: Kyle Ellrott <ellrott@ohsu.edu>, Raphael Kirchgaessner <kirchgae@ohsu.edu>
+License: 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
+
+
+# Embedding Kit
+
+Methods for data normalizatio, embedding, synthesis and transformation
+
+
+## Installation
+```
+python setup.py build
+python setup.py install
+```
+
+# Training a model
+```
+python -m embkit model train ./experiments/tcga/tumor.normalized.tsv --epochs 120
+```
+
+
+
+
+## Notebook cleaning
+
+To avoid checking in large notebooks with lots of intermediate figures, please follow the instructions at https://gist.github.com/33eyes/431e3d432f73371509d176d0dfb95b6e
+
+
+Running build
+

embkit-0.1.0/README.md

@@ -0,0 +1,27 @@
+
+# Embedding Kit
+
+Methods for data normalizatio, embedding, synthesis and transformation
+
+
+## Installation
+```
+python setup.py build
+python setup.py install
+```
+
+# Training a model
+```
+python -m embkit model train ./experiments/tcga/tumor.normalized.tsv --epochs 120
+```
+
+
+
+
+## Notebook cleaning
+
+To avoid checking in large notebooks with lots of intermediate figures, please follow the instructions at https://gist.github.com/33eyes/431e3d432f73371509d176d0dfb95b6e
+
+
+Running build
+

embkit-0.1.0/pyproject.toml

@@ -0,0 +1,29 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "embkit"
+version = "0.1.0"
+description = "Embedding Kit for embedding models"
+readme = "README.md"
+requires-python = ">=3.8"
+license = { text = "MIT" }
+authors = [
+    { name = "Kyle Ellrott", email = "ellrott@ohsu.edu" },
+    { name = "Raphael Kirchgaessner", email = "kirchgae@ohsu.edu" }
+]
+keywords = ["embeddings", "machine learning", "deep learning"]
+dependencies = [
+    "torch>=2.0",
+    "pandas",
+    "numpy",
+    "click",
+    "scikit-learn"
+]
+
+[tool.setuptools]
+package-dir = { "" = "src" }
+
+[tool.setuptools.packages.find]
+where = ["src"]

embkit-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

embkit-0.1.0/setup.py

@@ -0,0 +1,10 @@
+from setuptools import setup, find_packages
+
+setup(
+   name='embedding-kit',
+   version='0.1',
+   packages=find_packages(),
+   install_requires=[
+       # List your dependencies here!
+   ],
+)

embkit-0.1.0/src/embkit/__main__.py

@@ -0,0 +1,14 @@
+
+  # my_project/cli.py
+import click
+from .commands import model
+
+@click.group()
+def cli():
+    """A multi-module CLI application."""
+    pass
+
+cli.add_command(model)
+
+if __name__ == '__main__':
+    cli()

embkit-0.1.0/src/embkit/align.py

@@ -0,0 +1,59 @@
+import numpy as np
+import pandas as pd
+from scipy.stats import spearmanr
+from scipy.optimize import linear_sum_assignment
+
+def calc_rmsd(array1, array2):
+    if len(array1) != len(array2):
+        raise ValueError("Arrays must be of the same length")
+
+    diff = array1 - array2
+    squared_diff = diff ** 2
+    mean_squared_diff = np.mean(squared_diff)
+    rmsd = np.sqrt(mean_squared_diff)
+
+    return rmsd
+
+
+def matrix_spearman_alignment_set(a,b,cuttoff=0.0):
+    """
+    matrix_spearman_alignment_set
+    
+    """
+    # identify overlapping columns
+    isect = a.columns.intersection( b.columns )
+    # run spearman calculation on common features
+    o = spearmanr( a[isect], b[isect], axis=1 )
+    # use the corner of the correlation matrix that is the intersection of the 2 matrices
+    a_count=a.shape[0]
+    sdf = pd.DataFrame( o.correlation[a_count:,:a_count], index=b.index, columns=a.index)
+    # the matrix represents the edge weights in a bi-partite graph, use linear sum assignment
+    # to obtain the maximal set
+    mapping = linear_sum_assignment( sdf, True )
+    # remap output to dict
+    out = {}
+    for k, v in zip(mapping[0], mapping[1]):
+        c = sdf.iloc[ k, v ]
+        if c >= cuttoff:
+            out[ a.index[v] ] = (b.index[k], c)
+    return out
+
+
+def procrustes(X, Y):
+    """
+    Computes the Procrustes transformation between two matrices.
+
+    Args:
+        X: The first matrix.
+        Y: The second matrix.
+
+    Returns:
+        Rotation matrix
+    """
+    # Ensure matrices have the same dimensions
+    assert X.shape == Y.shape
+    # Compute the singular value decomposition (SVD)
+    U, S, Vt = np.linalg.svd(np.dot(X.T, Y))
+    # Determine the rotation matrix
+    R = np.dot(Vt.T, U.T)
+    return R.T

embkit-0.1.0/src/embkit/bmeg.py

@@ -0,0 +1,24 @@
+import gripql
+
+def connect(url="https://bmeg.io/grip", graph="rc6", cred_file="./bmeg_credentials.json"):
+    """
+    Establish a connection to the BMeg server and return the specified graph.
+
+    Parameters:
+        url (str): The URL of the BMeg server. Defaults to "https://bmeg.io/grip".
+        graph (str): The name of the graph to retrieve. Defaults to "rc6".
+        cred_file (str): The path to the JSON file containing the credentials for the BMeg server. Defaults to "./bmeg_credentials.json".
+
+    Returns:
+        gripql.Graph: A Graph object from the specified BMeg server and graph.
+    """
+    conn = gripql.Connection(url=url, credential_file=cred_file)
+    G = conn.graph(graph)
+    return G
+
+def get_gene_map(G):
+    """
+    Build mapping from Ensemb Gene IDs to Hugo Symbols
+    """
+    gm = dict( (i[0],i[1]) for i in G.query().V().hasLabel("Gene").render(["gene_id", "symbol"]).execute() )
+    return gm

embkit-0.1.0/src/embkit/commands.py

@@ -0,0 +1,27 @@
+import click
+import pandas as pd
+from .vae import VAE
+
+@click.group()
+def model():
+    """
+    model - model grouping
+    """
+    pass
+
+
+@model.command()
+@click.argument("input")
+@click.option("--latent", default=256)
+@click.option("--epochs", default=20)
+def train(input, latent, epochs):
+    """
+    Train VAE model
+    """
+
+    df = pd.read_csv(input, sep="\t", index_col=0)
+
+    vae = VAE(input_dim=df.shape[1], hidden_dim=400, latent_dim=latent)
+
+    vae.fit(df, epochs=epochs)
+

embkit-0.1.0/src/embkit/constraints/__init__.py

@@ -0,0 +1 @@
+from .network_constraint import NetworkConstraint

embkit-0.1.0/src/embkit/constraints/network_constraint.py

@@ -0,0 +1,54 @@
+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)

embkit-0.1.0/src/embkit/estimator/vae_estimator.py

@@ -0,0 +1,67 @@
+from typing import Optional
+from sklearn.base import BaseEstimator
+
+
+class VAEEstimator(BaseEstimator):
+    """
+    Scikit-learn style wrapper for quick experiments.
+    """
+
+    def __init__(self, latent_dim: int = 100, learning_rate: float = 1e-4,
+                 batch_size: int = 30, epochs: int = 130, beta: float = 0.0,
+                 device: Optional[str] = None):
+        self.latent_dim = latent_dim
+        self.learning_rate = learning_rate
+        self.batch_size = batch_size
+        self.epochs = epochs
+        self.beta = beta
+        self.history = None
+        self.model: Optional[VAE] = None
+        self.device = device
+
+    def fit(self, X: pd.DataFrame):
+        feature_dim = X.shape[1]
+        features = list(X.columns)
+
+        # Build unconstrained encoder/decoder
+        encoder = build_encoder(feature_dim, self.latent_dim, constraint=None)
+        decoder = build_decoder(feature_dim, self.latent_dim)
+        vae = VAE(features, encoder, decoder)
+
+        device = self.device or ('cuda' if torch.cuda.is_available() else 'cpu')
+        device = torch.device(device)
+        vae.to(device)
+
+        x = torch.tensor(X.values, dtype=torch.float32, device=device)
+        loader = DataLoader(TensorDataset(x), batch_size=self.batch_size, shuffle=True)
+        opt = torch.optim.Adam(vae.parameters(), lr=self.learning_rate)
+
+        loss_hist = []
+        for _ in range(self.epochs):
+            vae.train()
+            tot = 0.0
+            n = 0
+            for (batch_x,) in loader:
+                opt.zero_grad()
+                total, _, _ = vae_loss_from_model(vae, batch_x)
+                total.backward()
+                opt.step()
+                tot += float(total.item())
+                n += 1
+            loss_hist.append(tot / max(1, n))
+
+        self.history = {"loss": loss_hist}
+        self.model = vae
+        return self
+
+    def score(self, X: pd.DataFrame):
+        if self.model is None:
+            raise RuntimeError("Model not fitted.")
+        self.model.eval()
+        with torch.no_grad():
+            x = torch.tensor(X.values, dtype=torch.float32, device=next(self.model.parameters()).device)
+            mu, _, _ = self.model.encoder(x)
+            recon = self.model.decoder(mu).cpu().numpy()
+        return -mean_squared_error(X.values, recon)
+
+

embkit-0.1.0/src/embkit/layers/__init__.py

@@ -0,0 +1 @@
+from .masked_linear import MaskedLinear

embkit-0.1.0/src/embkit/layers/masked_linear.py

@@ -0,0 +1,28 @@
+from torch import nn
+from typing import Optional
+import torch
+import torch.nn.functional as F
+
+class MaskedLinear(nn.Module):
+    """
+    Linear layer whose weight is elementwise-multiplied by a mask at forward time.
+    PyTorch Linear weight shape: (out_features, in_features). Mask must match this shape.
+    """
+
+    def __init__(self, in_features: int, out_features: int, bias: bool = True,
+                 mask: Optional[torch.Tensor] = None):
+        super().__init__()
+        self.linear = nn.Linear(in_features, out_features, bias=bias)
+        if mask is None:
+            mask = torch.ones(out_features, in_features)
+        self.register_buffer("mask", mask)
+
+    def set_mask(self, mask: torch.Tensor):
+        assert mask.shape == self.linear.weight.shape
+        self.mask = mask
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        w = self.linear.weight * self.mask
+        return F.linear(x, w, self.linear.bias)
+
+

embkit-0.1.0/src/embkit/losses/__init__.py

@@ -0,0 +1 @@
+from .vae_loss import vae_loss_from_model

embkit-0.1.0/src/embkit/losses/vae_loss.py

@@ -0,0 +1,16 @@
+import torch
+from src.embkit.models import VAE
+import torch.nn.functional as F
+
+
+
+
+def vae_loss_from_model(model: VAE, x: torch.Tensor):
+    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)
+    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 + kl_loss
+    return total_loss.mean(), reconstruction_loss.mean(), kl_loss.mean()

embkit-0.1.0/src/embkit/models/__init__.py

@@ -0,0 +1 @@
+from .vae_model import VAE

embkit-0.1.0/src/embkit/models/vae_model/__init__.py

@@ -0,0 +1,3 @@
+from .base_vae import VAE
+from .decoder import Decoder
+from .encoder import Encoder

embkit-0.1.0/src/embkit/models/vae_model/base_vae.py

@@ -0,0 +1,226 @@
+from torch import nn
+from typing import Type, Any, List, Optional, Dict, overload, TypeVar
+from abc import ABC
+import torch
+from .encoder import Encoder
+from .decoder import Decoder
+from pathlib import Path
+import pandas as pd
+import json
+import logging
+
+logger = logging.getLogger(__name__)
+T = TypeVar("T")
+
+
+class VAE(nn.Module, ABC):
+    """
+    Minimal VAE wrapper to hold encoder/decoder and provide forward().
+    Allows late-binding of encoder/decoder by subclasses.
+    """
+
+    @staticmethod
+    def build_encoder(feature_dim: int, latent_dim: int,
+                      layers: List[Dict] = None,
+                      constraint=None,
+                      batch_norm: bool = False,
+                      activation: str = "relu") -> Encoder:
+        return Encoder(
+            feature_dim=feature_dim,
+            latent_dim=latent_dim,
+            layers=layers,
+            constraint=constraint,
+            batch_norm=batch_norm,
+            activation=activation
+        )
+
+    @staticmethod
+    def build_decoder(feature_dim: int, latent_dim: int) -> Decoder:
+        return Decoder(latent_dim, feature_dim)
+
+        # Overload 1: no model_cls provided -> returns NetVae
+
+    @overload
+    @staticmethod
+    def open_model(
+            path: str,
+            device: Optional[str] = ...,
+            model_cls: None = ...,
+            model_kwargs: Optional[Dict[str, Any]] = ...
+    ) -> "NetVae":
+        ...
+
+    # Overload 2: model_cls provided -> returns that class
+    @overload
+    @staticmethod
+    def open_model(
+            path: str,
+            device: Optional[str] = ...,
+            model_cls: Type[T] = ...,
+            model_kwargs: Optional[Dict[str, Any]] = ...
+    ) -> T:
+        ...
+
+    @staticmethod
+    def open_model(
+            path: str,
+            device: Optional[str] = None,
+            model_cls: Optional[Type[T]] = None,
+            model_kwargs: Optional[Dict[str, Any]] = None
+    ) -> T:
+        """
+        Load a model saved by NetVae.save (PyTorch-compatible).
+        You can select the container class via `model_cls` (default: NetVae).
+        The container is expected to accept at least: (features, encoder=..., decoder=...).
+
+        Works with or without net.groups.tsv:
+          - If groups exist: load names + groups
+          - Else: try latent.index
+          - Else: infer latent_dim from encoder weights and synthesize names
+        """
+        if device is None:
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        device = torch.device(device)
+
+        model_kwargs = dict(model_kwargs or {})
+
+        # --- Features (required) ---
+        index_path = Path(path, "index")
+        if not Path(index_path).exists():
+            raise FileNotFoundError(f"Missing features index file: {index_path}")
+        with open(index_path, "rt", encoding="ascii") as handle:
+            features = json.load(handle)
+        feature_dim = len(features)
+
+        # --- Determine latent names/dim ---
+        groups_path = Path(path, "net.groups.tsv")
+        latent_idx_path = Path(path, "latent.index")
+        groups: Optional[Dict[str, List[str]]] = None
+        latent_index: Optional[List[str]] = None
+        latent_dim: Optional[int] = None
+
+        if Path(groups_path).exists():
+            latent_index = []
+            groups = {}
+            with open(groups_path, "rt", encoding="ascii") as fh:
+                for line in fh:
+                    row = line.rstrip().split("\t")
+                    if not row or row[0] == "":
+                        continue
+                    latent_index.append(row[0])
+                    groups[row[0]] = row[1:]
+            latent_dim = len(latent_index)
+        elif Path(latent_idx_path).exists():
+            with open(latent_idx_path, "rt", encoding="ascii") as fh:
+                latent_index = [ln.strip() for ln in fh if ln.strip() != ""]
+            if latent_index:
+                latent_dim = len(latent_index)
+            else:
+                latent_index = None  # fall through to weight inference
+
+        if latent_dim is None:
+            # --- Infer from encoder weights ---
+            enc_state_path = Path(path, "model.enc.pt")
+            if not Path(enc_state_path).exists():
+                raise FileNotFoundError(f"Missing encoder weights: {enc_state_path}")
+            enc_state = torch.load(enc_state_path, map_location=device)
+
+            if "pathway.linear.weight" in enc_state:
+                latent_dim = int(enc_state["pathway.linear.weight"].shape[0])
+            elif "z_mean.weight" in enc_state:
+                latent_dim = int(enc_state["z_mean.weight"].shape[0])
+            else:
+                cand = None
+                for k, v in enc_state.items():
+                    if isinstance(v, torch.Tensor) and v.ndim == 2:
+                        cand = int(v.shape[0])
+                        break
+                if cand is None:
+                    raise ValueError(
+                        "Could not infer latent_dim from encoder state. "
+                        "Expected 'pathway.linear.weight' or 'z_mean.weight' in state dict."
+                    )
+                latent_dim = cand
+
+            if latent_index is None:
+                latent_index = [f"z{i}" for i in range(latent_dim)]
+
+        # --- Build fresh modules and load weights ---
+        enc = VAE.build_encoder(feature_dim, latent_dim, constraint=None)
+        dec = VAE.build_decoder(feature_dim, latent_dim)
+
+        enc.load_state_dict(torch.load(Path(path, "model.enc.pt"), map_location=device))
+        dec.load_state_dict(torch.load(Path(path, "model.dec.pt"), map_location=device))
+        enc.to(device)
+        dec.to(device)
+
+        # --- Instantiate user-provided container ---
+        try:
+            out = model_cls(features, encoder=enc, decoder=dec, **model_kwargs)
+        except TypeError as e:
+            raise TypeError(
+                f"{model_cls.__name__} could not be constructed with "
+                f"(features, encoder=..., decoder=..., **model_kwargs). "
+                f"Error: {e}"
+            ) from e
+
+        # Optional standard fields if the container supports them
+        # (we won't crash if the attributes don't exist)
+        if hasattr(out, "latent_groups"):
+            out.latent_groups = groups
+        if hasattr(out, "latent_index"):
+            out.latent_index = latent_index
+
+        stats_path = Path(path, "training.stats.tsv")
+        if Path(stats_path).exists():
+            try:
+                if hasattr(out, "normal_stats"):
+                    out.normal_stats = pd.read_csv(stats_path, sep="\t", index_col=0)
+            except Exception:
+                logger.warning(f"Could not read {stats_path}")
+                pass
+
+        return out
+
+    def __init__(self, features: List[str], encoder: Optional[Encoder] = None, decoder: Optional[Decoder] = None,
+                 **kwargs):
+        super().__init__()
+        self.features = list(features)
+        self.encoder: Optional[Encoder] = encoder
+        self.decoder: Optional[Decoder] = decoder
+        self.extra_args = kwargs
+
+    def forward(self, x: torch.Tensor):
+        if self.encoder is None or self.decoder is None:
+            raise RuntimeError(
+                "VAE encoder/decoder not initialized. Set self.encoder/self.decoder before calling forward().")
+        mu, logvar, z = self.encoder(x)
+        recon = self.decoder(z)
+        return recon, mu, logvar, z
+
+    def save(self, path: str, normal_df: Optional[pd.DataFrame] = None):
+        """Save VAE model with associated elements (PyTorch-native)."""
+        Path(path, exist_ok=True)
+
+        # Save encoder/decoder state dicts
+        torch.save(self.encoder.state_dict(), Path(path, "model.enc.pt"))
+        torch.save(self.decoder.state_dict(), Path(path, "model.dec.pt"))
+
+        # Save feature index
+        with open(Path(path, "index"), "wt", encoding="ascii") as handle:
+            handle.write(json.dumps(list(self.features)))
+
+        # Save latent groups (if present)
+        if self.latent_index is not None and self.latent_groups is not None:
+            with open(Path(path, "net.groups.tsv"), "wt", encoding="ascii") as handle:
+                for g in self.latent_index:
+                    n = self.latent_groups.get(g, [])
+                    handle.write("\t".join([str(g), *n]) + "\n")
+
+        # Save history
+        with open(Path(path, "stats.json"), "wt", encoding="ascii") as handle:
+            handle.write(json.dumps({"loss_history": self.history["loss"] if self.history else []}))
+
+        # Optionally save training stats like original
+        if normal_df is not None and self.normal_stats is not None:
+            self.normal_stats.to_csv(Path(path, "training.stats.tsv"), sep="\t")

embkit-0.1.0/src/embkit/models/vae_model/decoder.py

@@ -0,0 +1,17 @@
+from torch import nn
+import torch
+
+class Decoder(nn.Module):
+    """
+    Mirrors build_decoder (Dense with relu), but returns probabilities in [0,1] to
+    match the binary cross-entropy use in the original loss pipeline.
+    """
+
+    def __init__(self, latent_dim: int, feature_dim: int):
+        super().__init__()
+        self.fc = nn.Linear(latent_dim, feature_dim)
+
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        # Keep it simple: return sigmoid outputs for BCE
+        return torch.sigmoid(self.fc(z))
+

embkit-0.1.0/src/embkit/models/vae_model/encoder.py

@@ -0,0 +1,99 @@
+from torch import nn
+from typing import Optional, Tuple, List, Dict
+import torch
+from src.embkit.layers import MaskedLinear
+from src.embkit.constraints import NetworkConstraint
+import logging
+
+logger = logging.getLogger(__name__)
+
+class Encoder(nn.Module):
+    """
+    Dynamically built encoder:
+      input -> [optional BN] -> [list of Dense layers] -> latent dense -> BN -> (z_mean, z_log_var) -> sample z
+    forward(x) returns (z_mean, z_log_var, z)
+    """
+
+    def __init__(self, feature_dim: int, latent_dim: int,
+                 layers: List[Dict] = None,
+                 constraint: Optional[NetworkConstraint] = None,
+                 batch_norm: bool = False,
+                 activation: str = "relu"):
+        super().__init__()
+        self._constraint = constraint
+
+        self.net = nn.ModuleList()
+        in_features = feature_dim
+
+        # Optional initial batch norm
+        if batch_norm:
+            self.net.append(nn.BatchNorm1d(in_features))
+
+        # Build intermediate layers
+        if layers:
+            logger.info(f"Building encoder with {len(layers)} layers")
+            for i, layer_cfg in enumerate(layers):
+                units = layer_cfg.get("units")
+                use_mask = layer_cfg.get("masked", False)
+
+                if use_mask:
+                    init_mask = None
+                    if constraint is not None and constraint._mask_np is not None:
+                        init_mask = torch.tensor(constraint._mask_np, dtype=torch.float32)
+                    self.net.append(MaskedLinear(in_features, units, mask=init_mask))
+                else:
+                    self.net.append(nn.Linear(in_features, units))
+
+                act = self._get_activation(activation)
+                if act is not None:
+                    self.net.append(act)
+
+                if layer_cfg.get("batch_norm", False):
+                    self.net.append(nn.BatchNorm1d(units))
+
+                in_features = units
+        else:
+            logger.info(f"Building encoder with no layers")
+
+        # Embedding layer
+        self.embedding = nn.Linear(in_features, latent_dim)
+        self.embedding_act = self._get_activation(activation)
+        self.embedding_bn = nn.BatchNorm1d(latent_dim)
+
+        # Latent outputs
+        self.z_mean = nn.Linear(latent_dim, latent_dim)
+        self.z_log_var = nn.Linear(latent_dim, latent_dim)
+
+    def _get_activation(self, activation: str):
+        act_map = {
+            "relu": nn.ReLU(),
+            "tanh": nn.Tanh(),
+            "sigmoid": nn.Sigmoid(),
+            "leaky_relu": nn.LeakyReLU(),
+            "elu": nn.ELU(),
+            None: None
+        }
+        return act_map.get(activation.lower()) if activation else None
+
+    def refresh_mask(self, device: torch.device):
+        if self._constraint is not None:
+            for module in self.net:
+                if isinstance(module, MaskedLinear):
+                    module.set_mask(self._constraint.as_torch(device))
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        h = x
+        for layer in self.net:
+            h = layer(h)
+
+        h = self.embedding(h)
+        if self.embedding_act:
+            h = self.embedding_act(h)
+        h = self.embedding_bn(h)
+
+        mu = self.z_mean(h)
+        logvar = self.z_log_var(h)
+        std = torch.exp(0.5 * logvar)
+        eps = torch.randn_like(std)
+        z = mu + eps * std
+        return mu, logvar, z

embkit-0.1.0/src/embkit/models/vae_model/net_vae.py

@@ -0,0 +1,186 @@
+import os
+import json
+from typing import Dict, List, Optional, Callable
+import numpy as np
+import pandas as pd
+import torch
+
+from torch.utils.data import TensorDataset, DataLoader
+import logging
+from .base_vae import VAE
+from .encoder import Encoder
+from .decoder import Decoder
+from src.embkit.losses import vae_loss_from_model
+from src.embkit.constraints import NetworkConstraint
+
+logger = logging.getLogger(__name__)
+
+
+# ---------------------------------------------------------
+# NetVae (training with optional alternating constraint)
+# ---------------------------------------------------------
+
+class NetVae(VAE):
+    """
+    """
+
+    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
+        self.normal_stats: Optional[pd.DataFrame] = None
+
+    def run_train(self, df: pd.DataFrame, latent_index: List[str],
+                  latent_groups: Optional[Dict[str, List[str]]] = None,
+                  learning_rate: float = 1e-3, batch_size: int = 128,
+                  phases: Optional[List[int]] = None, epochs: int = 80,
+                  device: Optional[str] = None,
+                  grouping_fn: Optional[Callable[[np.ndarray, List[str]], Dict[str, List[str]]]] = None):
+
+        feature_dim = len(df.columns)
+        latent_dim = len(latent_index)
+        self.latent_index = list(latent_index)
+
+        if device is None:
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        device = torch.device(device)
+
+        # Constraint + models
+        constraint = NetworkConstraint(list(df.columns), latent_index, latent_groups)
+        constraint.set_active(False)
+
+        # Build and attach to self (no second VAE instance)
+        encoder = VAE.build_encoder(feature_dim=feature_dim, latent_dim=latent_dim, constraint=constraint)
+        decoder = VAE.build_decoder(feature_dim, latent_dim)
+        self.encoder = encoder
+        self.decoder = decoder
+
+        # Move whole module to device
+        self.to(device)
+
+        # Data
+        x = torch.tensor(df.values, dtype=torch.float32, device=device)
+        loader = DataLoader(TensorDataset(x), batch_size=batch_size, shuffle=True)
+
+        opt = torch.optim.Adam(self.parameters(), lr=learning_rate)
+
+        # History
+        total_hist: List[float] = []
+        recon_hist: List[float] = []
+        kl_hist: List[float] = []
+
+        # Determine total epochs if phases provided
+        if phases is not None:
+            total_epochs = sum(phases)
+            boundaries = np.cumsum(phases).tolist()
+        else:
+            total_epochs = epochs
+            boundaries = []
+
+        # Helper: apply constraint mask to encoder
+        def refresh_mask():
+            # encoder is attached to self
+            self.encoder.refresh_mask(device)
+
+        # Phase control
+        def start_constrained_phase():
+            if grouping_fn is not None:
+                with torch.no_grad():
+                    w = self.encoder.pathway.linear.weight.detach().cpu().numpy()  # [latent, features]
+                    new_groups = grouping_fn(w, list(df.columns))
+                    constraint.update_membership(new_groups)
+            constraint.set_active(True)
+            refresh_mask()
+
+        def start_unconstrained_phase():
+            constraint.set_active(False)
+            refresh_mask()
+
+        # Initial mask
+        refresh_mask()
+
+        # Training loop
+        for epoch in range(total_epochs):
+            if boundaries and epoch in boundaries:
+                idx = boundaries.index(epoch)
+                if idx % 2 == 0:
+                    start_constrained_phase()
+                else:
+                    start_unconstrained_phase()
+
+            self.train()
+            epoch_tot = epoch_rec = epoch_kl = 0.0
+            n_batches = 0
+
+            for (batch_x,) in loader:
+                opt.zero_grad()
+                # assumes vae_loss_from_model(model, x) calls model.forward(x)
+                total, recon, kl = vae_loss_from_model(self, batch_x)
+                total.backward()
+                opt.step()
+                epoch_tot += float(total.item())
+                epoch_rec += float(recon.item())
+                epoch_kl += float(kl.item())
+                n_batches += 1
+
+            total_hist.append(epoch_tot / max(1, n_batches))
+            recon_hist.append(epoch_rec / max(1, n_batches))
+            kl_hist.append(epoch_kl / max(1, n_batches))
+
+        # Store artifacts
+        self.latent_groups = constraint.latent_membership
+        self.history = {
+            "loss": total_hist,
+            "reconstruction_loss": recon_hist,
+            "kl_loss": kl_hist,
+        }
+
+        # normal_stats using deterministic recon from μ
+        self.eval()
+        with torch.no_grad():
+            mu, _, _ = self.encoder(x)
+            recon = self.decoder(mu).cpu().numpy()
+        normal_pred = pd.DataFrame(recon, index=df.index, columns=df.columns)
+        resid = normal_pred - df
+        self.normal_stats = pd.DataFrame({"mean": resid.mean(), "std": resid.std(ddof=0)})
+
+
+
+# ----------------
+# Load convenience
+# ----------------
+
+
+
+
+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))
+    for epoch in range(100):
+        net.run_train(
+            df=df,
+            latent_index=["latent1", "latent2"],  # latent_dim = 2
+            learning_rate=1e-3,
+            batch_size=32,
+            epochs=10,
+            device="cpu",
+            phases=None,
+        )
+        # print model training stats
+        print(
+            f"Epoch {epoch + 1}: Loss={net.history['loss'][-1]:.4f}, Recon Loss={net.history['reconstruction_loss'][-1]:.4f}, KL Loss={net.history['kl_loss'][-1]:.4f}")
+
+    # Save artifacts
+    net.save("vae_model")
+
+    model: NetVae = VAE.open_model(path="vae_model", model_cls=NetVae, device="cpu")
+    print("Model loaded with features:", model.features)
+    print(model.decoder)

embkit-0.1.0/src/embkit/normalize.py

@@ -0,0 +1,41 @@
+from sklearn.preprocessing import RobustScaler, MinMaxScaler
+from sklearn.base import BaseEstimator
+import numpy as np
+
+def quantile_max_norm(df, quantile_max=0.9):
+    """
+    Normalizes the DataFrame using Quantile Max normalization.
+
+    Parameters:
+        df (DataFrame): The input DataFrame to be normalized.
+        quantile_max (float, optional): The maximum quantile value to use for normalization. Default is 0.9.
+
+    Returns:
+        DataFrame: The normalized DataFrame.
+    """
+    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):
+    """
+    exp_max_norm
+
+    """
+    t = np.log2(df+1)
+    return (t.transpose() / t.max(axis=1)).transpose()
+
+class ExpMinMaxScaler(MinMaxScaler, BaseEstimator):
+    """
+    A custom Min-Max scaler that applies the log2 transformation before scaling.
+
+    Inherits from sklearn's MinMaxScaler and implements a fit method that
+    transforms the input data using the formula np.log2(X+1) before fitting the
+    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):
+        return MinMaxScaler.fit(self, np.log2(X+1))
+    def transform(self, X):
+        return MinMaxScaler.transform(self, np.log2(X+1))
+    def inverse_transform(self, X):
+        return np.exp2(MinMaxScaler.inverse_transform(self, X))-1

embkit-0.1.0/src/embkit/pathway.py

@@ -0,0 +1,8 @@
+
+import pandas as pd
+
+def extract_pathway_interactions(sif_path, relation='controls-expression-of'):
+    pc = pd.read_csv(sif_path, sep="\t", 
+                     header=None, names=["from", "relation", "to"])
+    pc = pc[ ~(pc['from'].map(lambda x:"CHEBI" in x) | pc['to'].map(lambda x:"CHEBI" in x)) ]
+    return pc[ pc["relation"] == relation ]

embkit-0.1.0/src/embkit/vae.py

@@ -0,0 +1,97 @@
+import os
+import json
+
+import numpy as np
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+import logging
+
+
+logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
+def get_device():
+
+    if torch.cuda.is_available():
+        print("Using CUDA GPU")
+        return  torch.device("cuda")
+    elif torch.backends.mps.is_available():
+        print("Using Metal GPU")
+        return torch.device("mps")
+    else:
+        print("Using CPU")
+        return torch.device("cpu")
+    
+
+def vae_loss(recon_x, x, mu, logvar):
+    bce_per_sample = F.binary_cross_entropy(recon_x, 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)
+    return (reconstruction_loss + kl_loss).mean(), reconstruction_loss.mean(), kl_loss.mean()
+
+
+class VAE(nn.Module):
+    def __init__(self, input_dim, hidden_dim, latent_dim, features=None):
+        super().__init__()
+        self.fc1 = nn.Linear(input_dim, hidden_dim)
+        self.fc21 = nn.Linear(hidden_dim, latent_dim)  # Mean
+        self.fc22 = nn.Linear(hidden_dim, latent_dim)  # Log-variance
+        self.fc3 = nn.Linear(latent_dim, hidden_dim)
+        self.fc4 = nn.Linear(hidden_dim, input_dim)
+
+        self.features = list(features) if features is not None else None
+
+    def encode(self, x):
+        h1 = F.relu(self.fc1(x))
+        return self.fc21(h1), self.fc22(h1)
+
+    def reparameterize(self, mu, logvar):
+        std = torch.exp(0.5 * logvar)
+        eps = torch.randn_like(std)
+        return mu + eps * std
+
+    def decode(self, z):
+        h3 = F.relu(self.fc3(z))
+        return torch.sigmoid(self.fc4(h3))
+
+    def forward(self, x):
+        mu, logvar = self.encode(x)
+        z = self.reparameterize(mu, logvar)
+        recon_x = self.decode(z)
+        return recon_x, mu, logvar
+
+    def save(self, path):
+        if not os.path.exists(path):
+            os.makedirs(path)
+        torch.save(self.state_dict(), os.path.join(path, "vae_model.pth"))
+        if self.features is not None:
+            with open(os.path.join(path, "index"), "wt", encoding="ascii") as handle:
+                handle.write(json.dumps(list(self.features)))
+
+
+    def fit(self, X, epochs=20):
+        optimizer = torch.optim.Adam(self.parameters(), lr=0.001)
+
+        logger.debug(self)
+
+        device = get_device()
+
+        x_tensor = torch.from_numpy(X.values.astype(np.float32))
+        x_tensor.to(device)
+
+        for epoch in range(epochs):
+            self.train()
+            optimizer.zero_grad()
+
+            recon_x, mu, logvar = self(x_tensor)
+            total_loss, recon_loss, kl_loss = vae_loss(recon_x, x_tensor, mu, logvar)
+
+            total_loss.backward()
+            optimizer.step()
+
+            logger.debug(f"Epoch {epoch+1} | Loss: {total_loss.item():.4f} | "
+                f"Recon: {recon_loss.item():.4f} | KL: {kl_loss.item():.4f}")    

embkit-0.1.0/src/embkit.egg-info/PKG-INFO

@@ -0,0 +1,42 @@
+Metadata-Version: 2.4
+Name: embkit
+Version: 0.1.0
+Summary: Embedding Kit for embedding models
+Author-email: Kyle Ellrott <ellrott@ohsu.edu>, Raphael Kirchgaessner <kirchgae@ohsu.edu>
+License: 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
+
+
+# Embedding Kit
+
+Methods for data normalizatio, embedding, synthesis and transformation
+
+
+## Installation
+```
+python setup.py build
+python setup.py install
+```
+
+# Training a model
+```
+python -m embkit model train ./experiments/tcga/tumor.normalized.tsv --epochs 120
+```
+
+
+
+
+## Notebook cleaning
+
+To avoid checking in large notebooks with lots of intermediate figures, please follow the instructions at https://gist.github.com/33eyes/431e3d432f73371509d176d0dfb95b6e
+
+
+Running build
+

embkit-0.1.0/src/embkit.egg-info/SOURCES.txt

@@ -0,0 +1,30 @@
+README.md
+pyproject.toml
+setup.py
+src/embkit/__init__.py
+src/embkit/__main__.py
+src/embkit/align.py
+src/embkit/bmeg.py
+src/embkit/commands.py
+src/embkit/normalize.py
+src/embkit/pathway.py
+src/embkit/vae.py
+src/embkit.egg-info/PKG-INFO
+src/embkit.egg-info/SOURCES.txt
+src/embkit.egg-info/dependency_links.txt
+src/embkit.egg-info/requires.txt
+src/embkit.egg-info/top_level.txt
+src/embkit/constraints/__init__.py
+src/embkit/constraints/network_constraint.py
+src/embkit/estimator/__init__.py
+src/embkit/estimator/vae_estimator.py
+src/embkit/layers/__init__.py
+src/embkit/layers/masked_linear.py
+src/embkit/losses/__init__.py
+src/embkit/losses/vae_loss.py
+src/embkit/models/__init__.py
+src/embkit/models/vae_model/__init__.py
+src/embkit/models/vae_model/base_vae.py
+src/embkit/models/vae_model/decoder.py
+src/embkit/models/vae_model/encoder.py
+src/embkit/models/vae_model/net_vae.py

embkit-0.1.0/src/embkit.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

embkit-0.1.0/src/embkit.egg-info/requires.txt

@@ -0,0 +1,5 @@
+torch>=2.0
+pandas
+numpy
+click
+scikit-learn

embkit-0.1.0/src/embkit.egg-info/top_level.txt

@@ -0,0 +1 @@
+embkit