tfmindi 1.0.0__py3-none-any.whl

tfmindi/__init__.py

@@ -0,0 +1,46 @@
+"""Package initialization for tfmindi."""
+
+import warnings
+from importlib.metadata import version
+
+# Suppress numba hashing warning from tangermeme
+warnings.filterwarnings("ignore", message=".*FNV hashing is not implemented in Numba.*", category=UserWarning)
+
+# Pre-import rapids_singlecell to establish proper CUDA environment
+# This must happen BEFORE any other CUDA libraries are imported
+try:
+    import rapids_singlecell  # type: ignore
+except ImportError:
+    pass
+
+from tfmindi import pl, pp, tl  # noqa: E402
+from tfmindi.backends import get_backend, is_gpu_available, set_backend  # noqa: E402
+from tfmindi.datasets import (  # noqa: E402
+    fetch_motif_annotations,
+    fetch_motif_collection,
+    load_motif_annotations,
+    load_motif_collection,
+    load_motif_to_dbd,
+)
+from tfmindi.io import load_h5ad, save_h5ad  # noqa: E402
+from tfmindi.types import Pattern, Seqlet  # noqa: E402
+
+__all__ = [
+    "pl",
+    "pp",
+    "tl",
+    "Pattern",
+    "Seqlet",
+    "get_backend",
+    "set_backend",
+    "is_gpu_available",
+    "fetch_motif_collection",
+    "fetch_motif_annotations",
+    "load_motif_collection",
+    "load_motif_annotations",
+    "load_motif_to_dbd",
+    "save_h5ad",
+    "load_h5ad",
+]
+
+__version__ = version("tfmindi")

tfmindi/backends/__init__.py

@@ -0,0 +1,151 @@
+"""Backend detection and configuration for CPU/GPU acceleration."""
+
+from __future__ import annotations
+
+import os
+import warnings
+from typing import Any, Literal
+
+Backend = Literal["cpu", "gpu"]
+
+# Global backend state
+_backend: Backend | None = None
+_gpu_available: bool | None = None
+
+
+def _check_gpu_availability() -> bool:
+    """Check if GPU acceleration packages are available."""
+    global _gpu_available
+    if _gpu_available is not None:
+        return _gpu_available
+
+    try:
+        import cupy as cp  # type: ignore
+
+        device_count = cp.cuda.runtime.getDeviceCount()
+        _gpu_available = device_count > 0
+        return _gpu_available  # type: ignore
+    except ImportError:
+        _gpu_available = False
+        return False
+
+
+def get_backend() -> str:
+    """
+    Get the current computational backend.
+
+    Returns the backend based on the following priority:
+    1. Explicitly set backend via set_backend()
+    2. Environment variable TFMINDI_BACKEND
+    3. Automatic detection based on GPU availability
+
+    Returns
+    -------
+    Backend type: "cpu" or "gpu"
+    """
+    global _backend
+
+    if _backend is not None:
+        return _backend
+
+    # Check environment variable
+    env_backend = os.getenv("TFMINDI_BACKEND", "").lower()
+    if env_backend in ["cpu", "gpu"]:
+        _backend = env_backend  # type: ignore
+        if _backend == "gpu" and not _check_gpu_availability():
+            warnings.warn(
+                "GPU backend requested but GPU packages not available. "
+                "Install with 'pip install tfmindi[gpu]'. Falling back to CPU.",
+                UserWarning,
+                stacklevel=2,
+            )
+            _backend = "cpu"
+        return _backend  # type: ignore
+
+    # Auto-detect based on availability
+    _backend = "gpu" if _check_gpu_availability() else "cpu"
+    return _backend
+
+
+def set_backend(backend: Backend) -> None:
+    """
+    Explicitly set the computational backend.
+
+    Parameters
+    ----------
+    backend
+        Backend type: "cpu" or "gpu"
+
+    Raises
+    ------
+    ValueError
+        If backend is not supported
+    ImportError
+        If GPU backend is requested but packages are not available
+    """
+    global _backend
+
+    if backend not in ["cpu", "gpu"]:
+        raise ValueError(f"Invalid backend: {backend}. Must be 'cpu' or 'gpu'.")
+
+    if backend == "gpu" and not _check_gpu_availability():
+        raise ImportError(
+            "GPU backend requested but required packages not available. Install with 'pip install tfmindi[gpu]'."
+        )
+
+    _backend = backend
+
+
+def is_gpu_available() -> bool:
+    """Check if GPU acceleration is available."""
+    return _check_gpu_availability()
+
+
+def get_array_module() -> Any:
+    """Get the appropriate array module (numpy or cupy) based on backend."""
+    if get_backend() == "gpu":
+        try:
+            import cupy as cp
+
+            return cp
+        except ImportError:
+            warnings.warn("CuPy not available, falling back to NumPy", UserWarning, stacklevel=2)
+            import numpy as np
+
+            return np
+    else:
+        import numpy as np
+
+        return np
+
+
+def to_cpu(array: Any) -> Any:
+    """Transfer array to CPU memory if it's on GPU."""
+    if hasattr(array, "get"):  # CuPy array
+        return array.get()
+    return array
+
+
+def to_gpu(array: Any) -> Any:
+    """Transfer array to GPU memory if GPU backend is active."""
+    if get_backend() == "gpu":
+        try:
+            import cupy as cp
+
+            if hasattr(array, "get"):  # Already on GPU
+                return array
+            return cp.asarray(array)
+        except ImportError:
+            warnings.warn("CuPy not available, keeping array on CPU", UserWarning, stacklevel=2)
+    return array
+
+
+__all__ = [
+    "Backend",
+    "get_backend",
+    "set_backend",
+    "is_gpu_available",
+    "get_array_module",
+    "to_cpu",
+    "to_gpu",
+]

tfmindi/datasets.py

@@ -0,0 +1,416 @@
+"""Dataset functions for TF-MInDi: fetching and loading motif collections and annotations."""
+
+from __future__ import annotations
+
+import os
+import zipfile
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import pooch
+
+_motif_index = None
+
+
+def _get_motif_index():
+    """
+    Set up the pooch motif collection registry from pycistarget.
+
+    Returns
+    -------
+    pooch.Pooch
+        The motif collection registry.
+    """
+    global _motif_index
+
+    if _motif_index is None:
+        _motif_index = pooch.create(
+            path=pooch.os_cache("tfmindi"),
+            base_url="https://resources.aertslab.org/cistarget/",
+            env="TFMINDI_DATA_DIR",
+            registry={
+                # Motif collections
+                "motif_collections/v10nr_clust_public/v10nr_clust_public.zip": "sha256:70dab42794f42471a3c22f5efe78ec2c8af96127656607cb0a929b4adffc2b97",
+                # Motif annotations (motifs-v{version}-nr.{species}-m0.001-o0.0.tbl)
+                # v8 (only Drosophila)
+                "motif2tf/motifs-v8-nr.flybase-m0.001-o0.0.tbl": "sha256:ffc550325334507c8ab2f6f79170fd0dfbcbecb33aee96040bf1f386fb3f982d",
+                # v9 annotations
+                "motif2tf/motifs-v9-nr.flybase-m0.001-o0.0.tbl": "sha256:db3458fdb38616758ea54543a42c7fc0e0ebb292cbf928d2648aefbf2cd59314",
+                "motif2tf/motifs-v9-nr.hgnc-m0.001-o0.0.tbl": "sha256:9e085e8e6ecd6f73a47fe435ed50d583b92bf7814d1bf38ae358e73c6207da2b",
+                "motif2tf/motifs-v9-nr.mgi-m0.001-o0.0.tbl": "sha256:cfab1245dbe770b1073f8048b1316285c9ce4b1a55d8ebca78962a2f406a172d",
+                # v10nr_clust annotations
+                "motif2tf/motifs-v10nr_clust-nr.chicken-m0.001-o0.0.tbl": "sha256:1ede59e4a737d822b7d6713243b83bf8f618f978c4fae1810fab265a30dfe3ba",
+                "motif2tf/motifs-v10nr_clust-nr.flybase-m0.001-o0.0.tbl": "sha256:91284e94b0317b764dc2f8d8147d30db707605df757ff7de16fb0953c63fda2a",
+                "motif2tf/motifs-v10nr_clust-nr.hgnc-m0.001-o0.0.tbl": "sha256:81eb754118e27e854974301b1400fcf519489f8be5249239671fb288cb501c31",
+                "motif2tf/motifs-v10nr_clust-nr.mgi-m0.001-o0.0.tbl": "sha256:5b64aad9df9804d50c50484c92d5192bdd5d2056cb105bdd343c0af2f94cce83",
+            },
+        )
+
+    return _motif_index
+
+
+def fetch_motif_collection() -> str:
+    """
+    Download motif collection (motif names and PWM) from aertslab' pycistarget resources.
+
+    Returns
+    -------
+    Path to downloaded motif collection folder
+
+    Examples
+    --------
+    >>> motif_dir = fetch_motif_collection()
+    >>> print(motif_dir)
+    """
+    # Mapping of collection names to registry keys
+    name = "v10nr_clust"  # only one collection name
+    collection_mapping = {
+        "v10nr_clust": "motif_collections/v10nr_clust_public/v10nr_clust_public.zip",
+    }
+
+    registry_key = collection_mapping[name]
+
+    def _extract_singletons(fname, action, pooch_obj):
+        """Extract only the singletons folder from the zip file."""
+        extract_dir = Path(fname).parent / Path(fname).stem / "singletons"
+        extract_dir.mkdir(parents=True, exist_ok=True)
+
+        with zipfile.ZipFile(fname, "r") as zip_file:
+            singletons_files = [
+                f for f in zip_file.namelist() if f.startswith("v10nr_clust_public/singletons/") and not f.endswith("/")
+            ]
+            for file_path in singletons_files:
+                file_name = Path(file_path).name
+                with zip_file.open(file_path) as source:
+                    with open(extract_dir / file_name, "wb") as target:
+                        target.write(source.read())
+        return str(extract_dir)
+
+    motif_dir = _get_motif_index().fetch(registry_key, processor=_extract_singletons, progressbar=True)
+
+    return motif_dir
+
+
+def fetch_motif_annotations(species: str = "hgnc", version: str = "v10nr_clust") -> str:
+    """
+    Download motif annotations from aertslab resources.
+
+    Parameters
+    ----------
+    species
+        Species name. Available options:
+        - 'hgnc' (human): v9, v10nr_clust
+        - 'mgi' (mouse): v9, v10nr_clust
+        - 'flybase' (fly): v8, v9, v10nr_clust
+        - 'chicken': v10nr_clust only
+    version
+        Motif collection version. Available options: 'v8', 'v9', 'v10nr_clust'
+
+    Returns
+    -------
+    Path to downloaded annotations file
+
+    Examples
+    --------
+    >>> annotations_file = fetch_motif_annotations("hgnc", "v10nr_clust")
+    >>> print(annotations_file)
+    /path/to/cache/motifs-v10nr_clust-nr.hgnc-m0.001-o0.0.tbl
+    """
+    # Mapping of species and versions to registry keys
+    annotation_mapping = {
+        ("flybase", "v8"): "motif2tf/motifs-v8-nr.flybase-m0.001-o0.0.tbl",
+        ("flybase", "v9"): "motif2tf/motifs-v9-nr.flybase-m0.001-o0.0.tbl",
+        ("hgnc", "v9"): "motif2tf/motifs-v9-nr.hgnc-m0.001-o0.0.tbl",
+        ("mgi", "v9"): "motif2tf/motifs-v9-nr.mgi-m0.001-o0.0.tbl",
+        ("chicken", "v10nr_clust"): "motif2tf/motifs-v10nr_clust-nr.chicken-m0.001-o0.0.tbl",
+        ("flybase", "v10nr_clust"): "motif2tf/motifs-v10nr_clust-nr.flybase-m0.001-o0.0.tbl",
+        ("hgnc", "v10nr_clust"): "motif2tf/motifs-v10nr_clust-nr.hgnc-m0.001-o0.0.tbl",
+        ("mgi", "v10nr_clust"): "motif2tf/motifs-v10nr_clust-nr.mgi-m0.001-o0.0.tbl",
+    }
+
+    key = (species, version)
+    assert key in annotation_mapping, (
+        f"Species {species} with version {version} is not recognised. "
+        f"Available combinations: {list(annotation_mapping.keys())}"
+    )
+
+    registry_key = annotation_mapping[key]
+
+    annotations_file = _get_motif_index().fetch(registry_key, progressbar=True)
+
+    return annotations_file
+
+
+def load_motif(file_name: str) -> dict[tuple[str, str], np.ndarray]:
+    """
+    Load motif(s) from a single .cb file.
+
+    Scales the PWM values to sum to 1 across each position (PPM).
+
+    Parameters
+    ----------
+    file_name
+        Path to the .cb motif file
+
+    Returns
+    -------
+    dict[tuple[str, str], np.ndarray]
+        Dictionary mapping motif file names and names to PWM matrices (4 x length)
+
+    Examples
+    --------
+    >>> motifs = load_motif("./motif1.cb")
+    >>> print(list(motifs.keys()))
+    [("filename_1", "motif_1"), ("filename_2", "motif_2")]
+    >>> print(motifs[("filename_1", "motif_1")].shape)
+    (4, 12)
+    """
+    motifs: dict[tuple[str, str], np.ndarray] = {}
+    with open(file_name) as f:
+        name = f.readline().strip()
+        if not name.startswith(">"):
+            raise ValueError(f"First line of {file_name} does not start with '>'.")
+        name = name.replace(">", "")
+        key = (os.path.basename(file_name).replace(".cb", ""), name)
+        pwm: list[list[float]] = []
+        for line in f:
+            line = line.strip()
+            if line.startswith(">"):
+                # we are at the start of a new motif
+                motifs[key] = np.array(pwm)
+                # scale values of motif
+                motifs[key] = motifs[key].T / motifs[key].sum(1)
+                # reset pwm and read new name
+                name = line.replace(">", "")
+                key = (os.path.basename(file_name).replace(".cb", ""), name)
+                pwm = []
+            else:
+                # we are in the middle of reading the pwm values
+                pwm.append([float(v) for v in line.split()])
+        # add last motif
+        motifs[key] = np.array(pwm)
+        # scale values of motif
+        motifs[key] = motifs[key].T / motifs[key].sum(1)
+    return motifs
+
+
+def load_motif_collection(motif_dir: str, motif_names: list[str] | None = None) -> dict[tuple[str, str], np.ndarray]:
+    """
+    Load motif collection from directory of .cb files.
+
+    Converts motif PWM matrices to PPM (position probability matrix) format.
+
+    Parameters
+    ----------
+    motif_dir
+        Directory path containing .cb motif files
+    motif_names
+        Optional list of specific motif names to load. If None, loads all motifs.
+
+    Returns
+    -------
+    dict[tuple[str, str], np.ndarray]
+        Dictionary mapping motif names to PWM matrices (4 x length)
+
+    Examples
+    --------
+    >>> motifs = load_motif_collection("./motif_collection/")
+    >>> print(list(motifs.keys()))
+    [("filename_1", "motif_1"), ("filename_1", "motif_2"), ("filename_1", "motif_4")]
+    >>> print(motifs[("filename_1", "motif_1")].shape)
+    (4, 12)
+
+    >>> # Load only specific motifs
+    >>> selected_motifs = load_motif_collection("./motif_collection/", ["motif1", "motif2"])
+    >>> print(list(selected_motifs.keys()))
+    ['motif1', 'motif2']
+    """
+    motif_dir_path = Path(motif_dir)
+
+    if not motif_dir_path.exists():
+        raise FileNotFoundError(f"Directory {motif_dir_path} does not exist")
+
+    motifs: dict[tuple[str, str], np.ndarray] = {}
+
+    cb_files = list(motif_dir_path.glob("*.cb"))
+
+    for cb_file in cb_files:
+        try:
+            # Load all motifs from this file
+            file_motifs = load_motif(str(cb_file))
+            motifs.update(file_motifs)
+        except (ValueError, IndexError):
+            # Skip malformed files
+            continue
+
+    # Filter motifs if specific names are provided
+    if motif_names is not None:
+        motif_names_set = set(motif_names)
+        motifs = {(filename, name): pwm for (filename, name), pwm in motifs.items() if name in motif_names_set}
+
+    return motifs
+
+
+def load_motif_annotations(
+    annotations_file: str,
+    motif_similarity_fdr: float = 0.001,
+    orthologous_identity_threshold: float = 0.0,
+    column_names: tuple[str, ...] = (
+        "#motif_id",
+        "gene_name",
+        "motif_similarity_qvalue",
+        "orthologous_identity",
+        "description",
+    ),
+) -> pd.DataFrame:
+    """
+    Load motif annotations from a motif2TF TSV file with filtering and categorization.
+
+    Parameters
+    ----------
+    annotations_file
+        Path to the annotations TSV file
+    motif_similarity_fdr
+        Maximum False Discovery Rate for enriched motifs (default: 0.001)
+    orthologous_identity_threshold
+        Minimum orthologous identity for enriched motifs (default: 0.0)
+    column_names
+        Column names to load from the TSV file
+
+    Returns
+    -------
+    DataFrame with motif annotations categorized by annotation type:
+    - Direct_annot: Direct gene annotations
+    - Motif_similarity_annot: Annotations by motif similarity
+    - Orthology_annot: Annotations by orthology
+    - Motif_similarity_and_Orthology_annot: Combined annotations
+
+    Examples
+    --------
+    >>> annotations = load_motif_annotations("./annotations.tbl")
+    >>> print(annotations.columns.tolist())
+    ['Direct_annot', 'Motif_similarity_annot', 'Orthology_annot', 'Motif_similarity_and_Orthology_annot']
+    """
+    # Load as pandas DataFrame
+    df = pd.read_csv(annotations_file, sep="\t", usecols=column_names)
+    df.rename(
+        columns={
+            "#motif_id": "MotifID",
+            "gene_name": "TF",
+            "motif_similarity_qvalue": "MotifSimilarityQvalue",
+            "orthologous_identity": "OrthologousIdentity",
+            "description": "Annotation",
+        },
+        inplace=True,
+    )
+
+    # Filter based on thresholds
+    df = df[
+        (df["MotifSimilarityQvalue"] <= motif_similarity_fdr)
+        & (df["OrthologousIdentity"] >= orthologous_identity_threshold)
+    ]
+
+    # Direct annotation
+    df_direct_annot = df[df["Annotation"] == "gene is directly annotated"]
+    df_direct_annot = df_direct_annot.groupby(["MotifID"])["TF"].apply(lambda x: ", ".join(list(set(x)))).reset_index()
+    df_direct_annot = df_direct_annot.set_index("MotifID")
+    df_direct_annot = pd.DataFrame(df_direct_annot["TF"])
+    df_direct_annot.columns = ["Direct_annot"]
+
+    # Indirect annotation - by motif similarity
+    motif_similarity_annot = df[
+        df["Annotation"].str.contains("similar") & ~df["Annotation"].str.contains("orthologous")
+    ]
+    motif_similarity_annot = (
+        motif_similarity_annot.groupby(["MotifID"])["TF"].apply(lambda x: ", ".join(list(set(x)))).reset_index()
+    )
+    motif_similarity_annot = motif_similarity_annot.set_index("MotifID")
+    motif_similarity_annot = pd.DataFrame(motif_similarity_annot["TF"])
+    motif_similarity_annot.columns = ["Motif_similarity_annot"]
+
+    # Indirect annotation - by orthology
+    orthology_annot = df[~df["Annotation"].str.contains("similar") & df["Annotation"].str.contains("orthologous")]
+    orthology_annot = orthology_annot.groupby(["MotifID"])["TF"].apply(lambda x: ", ".join(list(set(x)))).reset_index()
+    orthology_annot = orthology_annot.set_index("MotifID")
+    orthology_annot = pd.DataFrame(orthology_annot["TF"])
+    orthology_annot.columns = ["Orthology_annot"]
+
+    # Indirect annotation - by motif similarity and orthology
+    motif_similarity_and_orthology_annot = df[
+        df["Annotation"].str.contains("similar") & df["Annotation"].str.contains("orthologous")
+    ]
+    motif_similarity_and_orthology_annot = (
+        motif_similarity_and_orthology_annot.groupby(["MotifID"])["TF"]
+        .apply(lambda x: ", ".join(list(set(x))))
+        .reset_index()
+    )
+    motif_similarity_and_orthology_annot = motif_similarity_and_orthology_annot.set_index("MotifID")
+    motif_similarity_and_orthology_annot = pd.DataFrame(motif_similarity_and_orthology_annot["TF"])
+    motif_similarity_and_orthology_annot.columns = ["Motif_similarity_and_Orthology_annot"]
+
+    # Combine all annotation types
+    result = pd.concat(
+        [df_direct_annot, motif_similarity_annot, orthology_annot, motif_similarity_and_orthology_annot],
+        axis=1,
+        sort=False,
+    )
+
+    return result
+
+
+def load_motif_to_dbd(motif_annotations: pd.DataFrame) -> dict[str, str]:
+    """
+    Create motif-to-DNA-binding-domain mapping for human TFs.
+
+    Takes motif annotations and maps motifs to their DNA-binding domains
+    based on TF annotations and human TF database information.
+
+    Parameters
+    ----------
+    motif_annotations
+        DataFrame with motif annotations as returned by load_motif_annotations()
+
+    Returns
+    -------
+    dict[str, str]
+        Dictionary mapping motif IDs to DNA-binding domain names
+
+    Examples
+    --------
+    >>> annotations_file = fetch_motif_annotations("hgnc", "v10nr_clust")
+    >>> motif_annotations = load_motif_annotations(annotations_file)
+    >>> motif_to_dbd = load_motif_to_dbd(motif_annotations)
+    >>> print(motif_to_dbd["hocomoco__FOXO1_HUMAN.H11MO.0.A"])
+    'Forkhead'
+    """
+    motif_to_tf = motif_annotations.copy()
+
+    # Flatten all TF annotations into individual TF names
+    motif_to_tf = (
+        motif_to_tf.apply(lambda row: ", ".join(row.dropna()), axis=1)
+        .str.split(", ")
+        .explode()
+        .reset_index()
+        .rename({0: "TF"}, axis=1)
+    )
+
+    # Download human TF annotations with DNA-binding domains
+    human_tf_annot = pd.read_csv(
+        "https://humantfs.ccbr.utoronto.ca/download/v_1.01/DatabaseExtract_v_1.01.csv",
+        index_col=0,
+    )[["HGNC symbol", "DBD"]]
+
+    motif_to_tf = motif_to_tf.merge(right=human_tf_annot, how="left", left_on="TF", right_on="HGNC symbol")
+
+    # For each motif, take the most common (mode) DBD annotation
+    motif_to_dbd = (
+        motif_to_tf.dropna()
+        .groupby("MotifID")["DBD"]
+        .agg(lambda x: x.mode().iat[0])  # take the first mode if there's a tie
+        .reset_index()
+    )
+
+    motif_to_dbd = motif_to_dbd.set_index("MotifID")["DBD"].to_dict()
+
+    return motif_to_dbd