embkit 0.2__tar.gz → 0.4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: embkit
-Version: 0.2
+Version: 0.4
 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.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "embkit"
-version = "0.2"
+version = "0.4"
 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.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.2 → embkit-0.4}/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.4}/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.4}/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.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.2 → embkit-0.4}/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.4}/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":
-        out = matrix_spearman_alignment_linear(m1, m2, cutoff)
-        for k, v in out.items():
-            print(f"{k}\t{v[0]}\t{v[1]}")
+    if method == "linear":
+        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.2 → embkit-0.4}/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.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}")