embkit 0.2__tar.gz → 0.3__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: embkit
-Version: 0.2
+Version: 0.3
 Summary: Embedding Kit for embedding models
 Author-email: Kyle Ellrott <ellrott@ohsu.edu>, Raphael Kirchgaessner <kirchgae@ohsu.edu>
 License: MIT
@@ -17,6 +17,8 @@ Requires-Dist: gripql
 Requires-Dist: tqdm
 Requires-Dist: biopython
 Requires-Dist: fair-esm
+Requires-Dist: faiss-cpu
+Requires-Dist: h5py
 
 
 # Embedding Kit

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "embkit"
-version = "0.2"
+version = "0.3"
 description = "Embedding Kit for embedding models"
 readme = "README.md"
 requires-python = ">=3.8"
@@ -24,7 +24,9 @@ dependencies = [
     "gripql",
     "tqdm",
     "biopython",
-    "fair-esm"
+    "fair-esm",
+    "faiss-cpu",
+    "h5py"
 ]
 
 [tool.setuptools]
@@ -34,4 +36,4 @@ package-dir = { "" = "src" }
 where = ["src"]
 
 [project.scripts]
-embkit = "embkit.__main__:cli"
+embkit = "embkit.__main__:cli"

{embkit-0.2 → embkit-0.3}/src/embkit/__init__.py

@@ -24,7 +24,7 @@ def get_device() -> torch.device:
     logger.info("Using CPU")
     return torch.device("cpu")
 
-def dataframe_tensor(df: pd.DataFrame,device=None) -> torch.Tensor:
+def dataframe_tensor(df: pd.DataFrame, device=None, dtype=torch.float32) -> torch.Tensor:
     """
     dataframe_tensor
 
@@ -32,7 +32,7 @@ def dataframe_tensor(df: pd.DataFrame,device=None) -> torch.Tensor:
     """
     if device is None:
         device = get_device()
-    return torch.from_numpy(df.values.astype(np.float32)).to(device)
+    return torch.from_numpy(df.to_numpy(copy=True)).to(device=device, dtype=dtype)
 
 def tensor_dataframe(tn: torch.Tensor, index=None, columns=None) -> pd.DataFrame:
     """
@@ -40,11 +40,29 @@ def tensor_dataframe(tn: torch.Tensor, index=None, columns=None) -> pd.DataFrame
 
     take a torch tensor, create a pandas dataframe
     """
-    return pd.DataFrame(tn.cpu().detach().numpy(), index=index, columns=columns)
+    return pd.DataFrame(tn.to(device="cpu", dtype=torch.float32).detach().numpy(), index=index, columns=columns)
+
+def dataframe_dataset(df: pd.DataFrame, device=None, dtype=torch.float32) -> torch.utils.data.Dataset:
+    """
+    dataframe_dataset
+    
+    take a pandas dataframe and return a pytorch dataset
+
+    :param df: Input pandas Dataframe
+    :type df: pd.DataFrame
+    :param device: Device to put torch tensors on
+    :return: Description
+    :rtype: Dataset
+    """
+    if device is None:
+        device = get_device()
+    x = torch.from_numpy(df.to_numpy(copy=True)).to(device=device, dtype=dtype)
+    dataset = TensorDataset(x)
+    return dataset
 
 def dataframe_loader(df: pd.DataFrame,
                      batch_size = 256, shuffle=True,
-                     device=None) -> torch.utils.data.DataLoader:
+                     device=None, dtype=torch.float32) -> torch.utils.data.DataLoader:
     """
     dataframe_loader
 
@@ -54,9 +72,6 @@ def dataframe_loader(df: pd.DataFrame,
     :param device: Device the where tensor should be moved. If None, use get_device 
     """
 
-    if device is None:
-        device = get_device()
-    x = torch.from_numpy(df.values.astype(np.float32)).to(device)
-    dataset = TensorDataset(x)
+    dataset = dataframe_dataset(df, device=device, dtype=dtype)
     data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)
     return data_loader

{embkit-0.2 → embkit-0.3}/src/embkit/__main__.py

@@ -1,12 +1,15 @@
 import click
 import logging
 
-from .commands import model, matrix, cbio_cmd, protein, datasets, align
+from embkit.version import get_version
+
+from .commands import model, matrix, cbio_cmd, protein, resources, align
 
 @click.group(invoke_without_command=True)
 @click.option('--verbose', is_flag=True, help="Enable verbose logging")
+@click.option('--version', is_flag=True, help="Show version information")
 @click.pass_context
-def cli(ctx, verbose):
+def cli(ctx, verbose, version):
     """Embedding Kit CLI.
 
     Use 'embkit help [COMMAND]' for details on a specific command.
@@ -20,6 +23,10 @@ def cli(ctx, verbose):
     ctx.ensure_object(dict)
     ctx.obj["VERBOSE"] = verbose
 
+    if version:
+        click.echo(get_version())
+        ctx.exit()
+    
     if ctx.invoked_subcommand is None:
         click.echo(ctx.get_help())
 
@@ -29,7 +36,7 @@ cli.add_command(model)
 cli.add_command(matrix)
 cli.add_command(cbio_cmd)
 cli.add_command(protein)
-cli.add_command(datasets)
+cli.add_command(resources)
 cli.add_command(align)
 
 
@@ -62,11 +69,14 @@ def help_cmd(ctx, path):
     info_name = []
     parent = ctx.parent  # start from top-level context
     for part in path:
-        if not hasattr(cmd, "get_command"):
+        if not isinstance(cmd, click.Group):
             click.echo(cli.get_help(ctx)) # pragma: no cover
+            return
+
         nxt_cmd = cmd.get_command(parent, part)
         if nxt_cmd is None:
             raise click.UsageError(f"Unknown command: {' '.join(path)}")
+
         info_name.append(part)
         cmd = nxt_cmd
         parent = click.Context(cmd, info_name=part, parent=parent)

{embkit-0.2 → embkit-0.3}/src/embkit/align.py

@@ -2,8 +2,6 @@ import numpy as np
 import pandas as pd
 from scipy.stats import spearmanr
 from scipy.optimize import linear_sum_assignment
-from hopcroftkarp import HopcroftKarp
-
 
 def calc_rmsd(array1, array2):
     """
@@ -58,45 +56,6 @@ def matrix_spearman_alignment_linear(a, b, cutoff=0.0):
     return out_a, out_b, out_score
 
 
-def matrix_spearman_alignment_hopkraft(a, b, cutoff=0.0):
-    """
-    matrix_spearman_alignment_hopkraft
-
-    Uses the HopcroftKarp maximum *cardinality* matching algorithm.
-    It finds the LARGEST NUMBER of matches, not the BEST (highest score)
-    match.
-
-    Note: This is fundamentally different from linear_sum_assignment.
-    """
-    isect = a.columns.intersection(b.columns)
-
-    combined_df = pd.concat([a[isect], b[isect]], axis=0)
-    o = spearmanr(combined_df, axis=1)
-
-    a_count = a.shape[0]
-    sdf = pd.DataFrame(o.correlation[a_count:, :a_count], index=b.index, columns=a.index)
-
-    m = {}
-    s_result = sdf.apply(lambda x: x.nlargest(10).index, axis=0)
-    for k, row in s_result.items():
-        m[k] = row.tolist()
-
-    id_map = HopcroftKarp(m).maximum_matching(keys_only=True)
-
-    out_a = []
-    out_b = []
-    out_score = []
-
-    for k, v in id_map.items():
-        c = sdf.loc[v, k]
-        if c >= cutoff:
-            out_a.append(k)
-            out_b.append(v)
-            out_score.append(c)
-
-    return out_a, out_b, out_score
-
-
 def procrustes(X, Y):
     """
     Computes the Procrustes transformation (optimal PURE rotation)

{embkit-0.2 → embkit-0.3}/src/embkit/commands/__init__.py

@@ -13,8 +13,8 @@ embkit cbio <> : cBIO methods for data querying and download
 from .model import model
 from .matrix import matrix
 from .cbio import cbio_cmd
-from .datasets import datasets
+from .resources import resources
 from .protein import protein
 from .align import align
 
-__all__ = ["model", "matrix", "datasets", "protein"]
+__all__ = ["model", "matrix", "resources", "protein"]

{embkit-0.2 → embkit-0.3}/src/embkit/commands/align.py

@@ -6,7 +6,7 @@ Command line tools for aligning embedding spaces
 
 import click
 import pandas as pd
-from ..align import matrix_spearman_alignment_linear, matrix_spearman_alignment_hopkraft
+from ..align import matrix_spearman_alignment_linear
 
 align = click.Group(name="align", help="Embedding Alignment methods")
 
@@ -24,12 +24,7 @@ def pair(matrix1, matrix2, method, cutoff):
     m1 = pd.read_csv(matrix1, sep="\t", index_col=0)
     m2 = pd.read_csv(matrix2, sep="\t", index_col=0)
     
-    if method == "hopkraft":
-        out = matrix_spearman_alignment_hopkraft(m1, m2)
-        for k, v in out.items():
-            print(f"{k}\t{v[0]}\t{v[1]}")
-
-    elif method == "linear":
+    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]}")

{embkit-0.2 → embkit-0.3}/src/embkit/commands/cbio.py

@@ -1,5 +1,5 @@
 import click
-from ..datasets import CBIOPortal
+from ..resources import CBIOPortal
 from ..c_bio import CBIOAPI
 import logging
 

{embkit-0.2 → embkit-0.3}/src/embkit/commands/model.py

@@ -1,15 +1,20 @@
 
 import click
 import pandas as pd
-import numpy as np
-import torch
 
 from sklearn.preprocessing import MinMaxScaler
 
-from .. import dataframe_loader, dataframe_tensor, get_device
-from ..layers import LayerInfo, build_layers, ConstraintInfo
-from ..models.vae.vae import VAE, BaseVAE
-from ..preprocessing import ExpMinMaxScaler
+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, ConstraintInfo
+from ..optimize import fit_vae
+from ..models.vae.vae import VAE
+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
 
@@ -17,57 +22,95 @@ 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=str, 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)
-def train_vae(input_path: str, latent: int, 
-              epochs: int, out: str, normalize:str,
+@click.option("--zero-mask", default=None, type=float)
+@click.option("--seed", default=42, type=int)
+@click.option("--bfloat16", is_flag=True)
+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, 
-              save_stats: bool
+              schedule:str,
+              zero_mask: float,
+              save_stats: bool,
+              seed: int,
+              bfloat16: bool
               ):
-    """Train VAE model from a TSV file."""
-    df = pd.read_csv(input_path, sep="\t", index_col=0)
+    """
+    Train VAE model from a TSV file.
+    """
+    device = get_device()
+    dtype = torch.float32
+    if bfloat16:
+        dtype = torch.bfloat16
 
-    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)
-
-
-    feature_count = len(df.columns)
+    torch.manual_seed(seed)
+    df = None
+    if group is not None:
+        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:
+            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(",") )
-    layer_sizes.append(latent)
-    enc_layers = build_layers( layer_sizes )
+    enc_layers_list = LayerList( layer_sizes )
 
     layer_sizes = list( int(i) for i in decode_layers.split(",") )
-    layer_sizes.append(feature_count)
-    dec_layers = build_layers( layer_sizes, end_activation="none" )
+    dec_layers_list = LayerList( layer_sizes, end_activation=final_activation )
 
     beta_schedule = None
     if schedule is not None:
-        beta_schedule = [] #[(0.0, 20), (0.1, 20), (0.3, 40), (0.4, 40)]
+        beta_schedule = []
         for b in schedule.split(","):
             e, b = b.split(":")
             beta_schedule.append( (float(b), int(e)) )
-    vae = VAE(features=df.columns,
+    vae = VAE(features=features,
               latent_dim=latent,
-              encoder_layers=enc_layers,
-              decoder_layers=dec_layers)
+              encoder_layers=enc_layers_list,
+              decoder_layers=dec_layers_list,
+              device=device, dtype=dtype)
 
     loss_func = bce_with_logits
     if loss == "mse":
@@ -75,7 +118,7 @@ def train_vae(input_path: str, latent: int,
     elif loss == "bce":
         loss_func = bce
 
-    vae.fit(df, epochs=epochs,
+    fit_vae(vae, dataloader, epochs=epochs,
             beta_schedule=beta_schedule, lr=learning_rate, loss=loss_func)
     click.echo("Training complete.")
 
@@ -83,11 +126,7 @@ def train_vae(input_path: str, latent: int,
         click.echo(f"No output path provided, using default naming.")
         out = f"vae_latent{latent}_epochs{epochs}.model"
 
-    if save_stats:
-        vae.save(out, df)
-    else:
-        vae.save(out)
-
+    save(vae, out)
     click.echo(f"Model saved, to {out}")
 
 
@@ -133,15 +172,15 @@ def train_netvae(input_path: str, pathway_sif:str, out:str,
     gcounts = [5,2,1]
 
     enc_layers = [
-        LayerInfo(group_count*gcounts[0], op="masked_linear", constraint=ConstraintInfo("features-to-group", fmap, out_group_count=gcounts[0])),
-        LayerInfo(group_count*gcounts[1], op="masked_linear", constraint=ConstraintInfo("group-to-group", fmap, in_group_count=gcounts[0], out_group_count=gcounts[1])),
-        LayerInfo(group_count, op="masked_linear", constraint=ConstraintInfo("group-to-group", fmap, in_group_count=gcounts[1], out_group_count=gcounts[2]))
+        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 = [
-        LayerInfo(group_count*gcounts[1], op="masked_linear", constraint=ConstraintInfo("group-to-group", fmap, in_group_count=gcounts[2], out_group_count=gcounts[1])),
-        LayerInfo(group_count*gcounts[0], op="masked_linear", constraint=ConstraintInfo("group-to-group", fmap, in_group_count=gcounts[1], out_group_count=gcounts[0])),
-        LayerInfo(feature_count, op="masked_linear", constraint=ConstraintInfo("group-to-features", fmap, in_group_count=gcounts[0]), activation="none")
+        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")
     ]
 
     loss_func = bce_with_logits
@@ -150,18 +189,24 @@ def train_netvae(input_path: str, pathway_sif:str, out:str,
     elif loss == "bce":
         loss_func = bce
 
-
-    #schedule = [(0.0, 20), (0.1, 20), (0.3, 40), (0.4, 40)]
-    schedule = [(0.0, 20)]
-    vae = VAE(df.columns, latent_dim=group_count, encoder_layers=enc_layers, decoder_layers=dec_layers)
-    vae.fit(X=dataloader, beta_schedule=schedule, lr=learning_rate, loss=loss_func)
+    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)
 
     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:
-        vae.save(out, df)
-    else:
-        vae.save(out)
+        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))
@@ -169,11 +214,9 @@ def train_netvae(input_path: str, pathway_sif:str, out: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 = VAE.open_model(path=model_path)
-    df = pd.read_csv(input_path, sep="\t", index_col=0)
 
-    # print(m)
-    # df = pd.read_csv(input_path, sep="\t", index_col=0)
+    m = load(model_path)
+    df = pd.read_csv(input_path, sep="\t", index_col=0)
 
     df = df[m.features]
     if normalize == "expMinMax":
@@ -188,33 +231,3 @@ 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("h5", type=str)
-@click.option("--epochs", "-e", type=int, default=5, show_default=True, help="Training epochs")
-@click.option("--batch-size", type=int, default=1024, show_default=True, help="Batch Size")
-@click.option("--learning-rate", type=float, default=1e-3, show_default=True, help="Learning Rate")
-@click.option("--seed", type=int, default=42, show_default=True, help="Batch Size")
-@click.option("--embedding-dim", type=int, default=64)
-def train_hla2vec(h5:str, epochs:int, batch_size:int, learning_rate:float, seed:int, embedding_dim:int):
-
-    #TODO: this should be put into a library function
-    np.random.seed(seed)
-    if torch.cuda.is_available():
-        torch.cuda.manual_seed_all(seed)
-    torch.manual_seed(seed)
-
-    data = load_bigmhc(h5)
-
-    model, history = fit_hla2vec(
-        **data,
-        emb_dim=embedding_dim,
-        batch_size=batch_size,
-        epochs=epochs,
-        lr=learning_rate,
-        seed=seed,
-    )
-
-    print("Final epoch:", history[-1])
-    torch.save(model.state_dict(), "hla2vec_demo.pt")
-    print("Saved weights -> hla2vec_demo.pt")

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

@@ -5,10 +5,14 @@ import re
 import sys
 from typing import List
 import click
+import json
+from ..files.json import StringTrimEncoder
 from Bio import SeqIO
 
+import torch
+
 from .. import get_device
-from ..models.protein import ProteinEncoder
+from ..encoding.protein import ProteinEncoder
 
 
 protein = click.Group(name="protein", help="Protein commands.")
@@ -37,9 +41,10 @@ def stringify(l:List[float], trim=None) -> List[str]:
 @click.option("--batch-size", type=int, default=128)
 @click.option("--trim", type=int, default=None)
 @click.option("--model", "-m", type=click.Choice(['t6', 't12', 't30', 't33', 't36', 't48']), default="t33")
-@click.option("--pool", "-p", type=click.Choice(["mean", "sum"]), default="mean")
+@click.option("--pool", "-p", type=click.Choice(["mean", "sum", "none"]), default="mean")
 @click.option("--output", "-o", type=str, default=None)
-def encode(fasta: str, filter:str, batch_size:int, model:str, trim:int, pool:str, output:str):
+@click.option("--fix-len", type=int, default=None)
+def encode(fasta: str, filter:str, batch_size:int, model:str, trim:int, pool:str, output:str, fix_len):
     pool_map = {
         "mean" : "mean-pool",
         "sum" : "sum-pool"
@@ -50,7 +55,17 @@ def encode(fasta: str, filter:str, batch_size:int, model:str, trim:int, pool:str
 
     enc = ProteinEncoder(batch_size=batch_size, model=model)
     enc.to(get_device())
-    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")
+    if pool == "none":
+        for i, emb in enc.encode(fasta_reader(fasta, filter=filter), output="vector", fix_len=fix_len):
+            out.write(f"{i}\t")
+            if trim:
+                out.write(json.dumps(emb, cls=StringTrimEncoder, trim=trim ))
+            else:
+                out.write(json.dumps(emb.tolist()))
+            #out.write( f"{i}\t" + "\t".join(stringify(emb.tolist(), trim)))
+            out.write("\n")
+    else:
+        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()

embkit-0.2/src/embkit/commands/datasets.py → embkit-0.3/src/embkit/commands/resources.py

@@ -1,12 +1,12 @@
 import click
 import pandas as pd
-from ..preprocessing import load_raw_hugo, load_gct
-from ..datasets import GTEx, Hugo, SIF
+from ..files import load_raw_hugo, load_gct
+from ..resources import GTEx, Hugo, SIF
 
-datasets = click.Group(name="datasets", help="Datasets commands.")
+resources = click.Group(name="resources", help="Resource commands.")
 
 
-@datasets.command()
+@resources.command()
 @click.option("--data_type", "-t", required=False, type=str, help="Name of the dataset to download.")
 @click.option("--output_folder", "-f", required=True, type=str, help="The folder to download the dataset into.")
 def gtex(data_type: str, output_folder: str):
@@ -20,7 +20,7 @@ def gtex(data_type: str, output_folder: str):
 
 
 
-@datasets.command()
+@resources.command()
 @click.option("--output_folder", "-f", required=True, type=str, help="The folder to download the dataset into.")
 def sif(output_folder: str):
     """Download SIF dataset."""
@@ -28,7 +28,7 @@ def sif(output_folder: str):
 
 
 
-@datasets.command()
+@resources.command()
 @click.option("--output_folder", "-f", required=True, type=str, help="The folder to download the dataset into.")
 @click.option("--gtex_conversion", "-gtex", required=False, is_flag=True, type=bool, default=False,
               help="Whether to convert the hugo data to a gtex version.")

{embkit-0.2 → embkit-0.3}/src/embkit/constraints/network_constraint.py

@@ -51,4 +51,22 @@ class NetworkConstraint:
         self._mask_np = mask
 
     def as_torch(self, device: torch.device) -> torch.Tensor:
-        return torch.tensor(self._mask_np, device=device)
+        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