embkit 0.3__tar.gz → 0.4__tar.gz

{embkit-0.3 → embkit-0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: embkit
-Version: 0.3
+Version: 0.4
 Summary: Embedding Kit for embedding models
 Author-email: Kyle Ellrott <ellrott@ohsu.edu>, Raphael Kirchgaessner <kirchgae@ohsu.edu>
 License: MIT

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "embkit"
-version = "0.3"
+version = "0.4"
 description = "Embedding Kit for embedding models"
 readme = "README.md"
 requires-python = ">=3.8"

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

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

{embkit-0.3 → embkit-0.4}/src/embkit/c_bio/api.py

@@ -1,4 +1,7 @@
 import requests
+import logging
+
+logger = logging.getLogger(__name__)
 
 
 class CBIOAPI:
@@ -9,5 +12,5 @@ class CBIOAPI:
             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

{embkit-0.3 → embkit-0.4}/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}")

{embkit-0.3 → embkit-0.4}/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}")

embkit-0.4/src/embkit/commands/model.py

@@ -0,0 +1,401 @@
+
+import click
+import pandas as pd
+import json
+
+from sklearn.preprocessing import MinMaxScaler
+
+import torch
+from torch.utils.data import DataLoader
+
+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
+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_sif_interactions, feature_map_intersect, build_feature_map_indices
+
+model = click.Group(name="model", help="VAE Model commands.")
+
+@model.command()
+@click.argument("input_path", type=click.Path(exists=True, dir_okay=False, readable=True, path_type=str))
+@click.option("--group", "-g", default=None, help="HD5F group name")
+@click.option("--latent", "-l", type=int, default=256, show_default=True, help="Latent dimension size.")
+@click.option("--epochs", "-e", type=int, default=20, show_default=True, help="Training epochs.")
+@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=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)
+@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("--zero-mask", default=None, type=float)
+@click.option("--seed", default=42, type=int)
+@click.option("--bfloat16", is_flag=True)
+@click.option("--sampling/--no-sampling", default=True, show_default=True, help="Enable reparameterization sampling during training (VAE). Use --no-sampling for a standard autoencoder.")
+def train_vae(input_path: str,
+              group: str,
+              latent: int,
+              epochs: int,
+              batch_size: int,
+              out: str, normalize:str,
+              encode_layers:str, decode_layers:str,
+              learning_rate: float,
+              final_activation: str,
+              loss: str,
+              schedule:str,
+              zero_mask: float,
+              save_stats: bool,
+              seed: int,
+              bfloat16: bool,
+              sampling: bool
+              ):
+    """
+    Train VAE model from a TSV file.
+    """
+    device = get_device()
+    dtype = torch.float32
+    if bfloat16:
+        dtype = torch.bfloat16
+
+    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)
+        if zero_mask is not None:
+            dataset_mask = get_dataset_nonzero_mask(dataset, zero_mask)
+            mask = dataset_mask[0].cpu().numpy()
+            features = dataset.columns[mask]
+            dataset = DatasetMask(dataset, dataset_mask, device)
+        else:
+            dataset.to(device)
+            features = dataset.columns
+    
+    else:
+        df = pd.read_csv(input_path, sep="\t", index_col=0)
+
+        if normalize == "expMinMax":
+            norm = ExpMinMaxScaler()
+            norm.fit(df)
+            df = pd.DataFrame( norm.transform(df), index=df.index, columns=df.columns)
+        elif normalize == "minMax":
+            norm = MinMaxScaler()
+            norm.fit(df)
+            df = pd.DataFrame( norm.transform(df), index=df.index, columns=df.columns)
+        features = df.columns
+        dataset = dataframe_dataset(df, device=device, dtype=dtype)
+    
+    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 )
+
+    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)) )
+    vae = VAE(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
+
+    fit_vae(vae, dataloader, epochs=epochs,
+            beta_schedule=beta_schedule, lr=learning_rate, loss=loss_func)
+    click.echo("Training complete.")
+
+    if out is None:
+        click.echo(f"No output path provided, using default naming.")
+        out = f"vae_latent{latent}_epochs{epochs}.model"
+
+    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("--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,
+                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_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)
+
+    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)
+
+    group_count = len(group_idx)
+    feature_count = len(feature_idx)
+
+    click.echo(f"Feature count {feature_count} latent_size: {group_count}")
+
+    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":
+        loss_func = mse
+    elif loss == "bce":
+        loss_func = bce
+
+    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.")
+
+    if out is None:
+        click.echo("No output path provided, using default naming.")
+        out = f"netvae_latent{group_count}_epochs{epochs}.model"
+
+    save(vae, out)
+    click.echo(f"Model saved, to {out}")
+
+    if save_stats:
+        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("model_path", type=click.Path(exists=True, dir_okay=True, readable=True, path_type=str))
+@click.option("--normalize", "-n", type=str, default="none")
+@click.option("--out", "-o", type=str, default="embedding.tsv")
+def encode(input_path: str, model_path:str, normalize:str, out:str):
+
+    m = load(model_path)
+    df = pd.read_csv(input_path, sep="\t", index_col=0)
+
+    df = df[m.features]
+    if normalize == "expMinMax":
+        norm = ExpMinMaxScaler()
+        norm.fit(df)
+        df = pd.DataFrame( norm.transform(df), index=df.index, columns=df.columns)
+
+    df_tensor = dataframe_tensor(df).to(get_device())
+    m.to(get_device())
+    result = m.encoder(df_tensor)
+    
+    matrix = result[0].detach().cpu().numpy()
+    out_df = pd.DataFrame(matrix, 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()

{embkit-0.3 → embkit-0.4}/src/embkit/commands/protein.py

@@ -18,13 +18,14 @@ from ..encoding.protein import ProteinEncoder
 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()
+    if should_close:
+        out.close()

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

@@ -0,0 +1,7 @@
+"""
+Constraint entry points.
+"""
+
+from .pathway_constraint import PathwayConstraintInfo
+
+__all__ = ["PathwayConstraintInfo"]

embkit-0.4/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"]