tnx 0.1.0

TENET eXtremely optimized

GPU-accelerated TENET algorithm for gene regulatory network inference from single-cell RNA-seq data

Documentation · Installation · Releases

TENEX computes pairwise transfer entropy (TE) on GPU to infer gene regulatory networks (GRNs) from scRNA-seq data. It is a high-performance reimplementation of FastTENET, achieving up to 2,203x speedup.

Features

• Adaptive kernel selection: automatic selection of the best CUDA kernel for the data and hardware characteristics
• Multi-GPU support: thread-based parallelism with no spawn overhead
• GPU-native preprocessing: discretization and bin remapping entirely on GPU
• Multiple inference methods: FDR, CLR, Network Deconvolution, TRACE key-driver analysis, and a surrogate-based statistical test

Quick Start

import tenex as tnx

# Load scRNA-seq data (aligned by pseudotime within selected branch)
scrna = tnx.load_scrna(
    expression="expression_data.csv",
    pseudotime="pseudotime.txt",
    branch="branch.txt",
)

# Compute pairwise TE matrix
engine = tnx.TransferEntropyEngine(
    data=scrna.data,
    variable_names=scrna.gene_names,
)
result = engine.compute(accelerator="gpu")
nw = tnx.NetWeaver(result, fdr=0.01)

Infer GRN links (FDR-based with DPI trimming):

grn, trimmed = nw.infer(method="fdr")

Surrogate test (effective TE + per-pair Gaussian z-test against the empirical null distribution from time-axis block shuffles):

sur_result = nw.infer(method="surrogate_test", n_surrogates=100)
sur_result.effective_te        # (n, n) bias-corrected TE
sur_result.p_values            # (n, n) p-values
sur_result.grn                 # BH-FDR-thresholded edges

Or use the Pipeline for one-line end-to-end compute + multiple inferences (TE matrix is computed once and reused across methods):

pipe = tnx.Pipeline(engine, fdr=0.05)

pr = pipe.run(
    methods=["fdr", "surrogate_test"],
    method_kwargs={"surrogate_test": {"n_surrogates": 100}},
)
pr.fdr                         # (grn, trimmed): FDR + DPI-trimmed edges
pr.surrogate_test              # effective TE, p-values, significant edges
pr.matrix                      # cached pairwise TE matrix

Installation

Install PyTorch first (it is a prerequisite, not a TENEX dependency), then install TENEX. Keeping PyPI as the primary index lets numpy, scipy, and the other runtime dependencies resolve normally, while the TENEX wheel index is added as an extra source.

# Step 1: install PyTorch matching your CUDA version
pip install torch --index-url https://download.pytorch.org/whl/cu132

# Step 2: install TENEX from the matching CUDA sub-index
pip install tnx --extra-index-url https://cxinsys.github.io/tenex/whl/cu132/

Use the same CUDA line in both URLs. Pre-built GPU wheels are published for Linux (cu132, cu128) and Windows (cu128, cu126); on macOS or for CPU-only use, pip install tnx from PyPI installs the universal CPU wheel.

For build-from-source, JIT fallback, CPU-only mode, and troubleshooting, see INSTALL.md.

Currently supported kernels

The counting of the 3-D joint histogram is the bottleneck of TE computation, so TENEX provides several kernels and auto_select() picks one from b_max (the largest bin count of any gene after dense remapping) and the available GPU shared memory.

Performance

Matched single-GPU comparison on NVIDIA PRO 6000 Blackwell (median of 3 runs):

Surrogate test

Wall time for nw.infer(method="surrogate_test", n_surrogates=100) on NVIDIA PRO 6000 Blackwell. The dispatch auto-selects between the loop path (GPU-accumulator) and the fused CUDA kernel based on L = T - dt.

1-GPU dispatch

The dispatch picks the fused CUDA kernel for short series and the loop path (GPU-side accumulator) for long ones, using the heuristic L < 1500 (the fused path is available only for the Full-SMEM and Adaptive-SMEM backends; the fused kwarg overrides the choice). Numerical agreement between the two paths stays within float precision: Δ mean_surrogate_te ≤ 2e-9 on every dataset listed below.

2-GPU scaling

The 2-GPU runs use the same auto-dispatched path on each GPU, with the surrogate workload partitioned between devices. The resulting accumulators are combined at the end of the run and produce results that are bit-for-bit identical to the 1-GPU run (we verified |Δ mean_surrogate_te| = 0.0 and identical significant-edge counts on every dataset).

mESC and Skin numbers are direct measurements at n_surrogates=100. Zebrafish and CeNGEN are extrapolated from n_surrogates=10 runs. The per-iteration cost is constant once the per-process setup is amortized, so the extrapolation is tight.

Citation

If you use TENEX in your work, please cite the accompanying paper, forthcoming in Bioinformatics. Until it appears, cite the software using the metadata in CITATION.cff.

License

TENEX is released under the TENEX Non-Commercial License (see LICENSE). The source code is openly available and free to use, modify, and redistribute for non-commercial purposes, provided the license is reproduced in copies and derivative works. For commercial use, contact Daewon Lee (dwlee@cau.ac.kr) for a commercial license.