genie3_grn_pipeline.py

#!/usr/bin/env python3
"""
GENIE3 Gene Regulatory Network Pipeline
======================================

Author: Taufia Hussain
Project: DataLens.Tools / GENIE3 GRN Framework

Description
-----------
This script infers a Gene Regulatory Network (GRN) from a normalized RNA-seq
expression matrix using the GENIE3 algorithm implemented in arboreto.

Pipeline steps:
1. Load normalized RNA-seq expression matrix
2. Filter low-expression genes
3. Load curated transcription factor list
4. Run GENIE3
5. Rank TF-target regulatory links
6. Export top edges
7. Build a network graph
8. Identify candidate hub transcription factors

Input format
------------
Expression matrix CSV:
    Rows    = genes
    Columns = samples
    Values  = normalized expression values, e.g. TPM, CPM, log2(CPM + 1), or VST

TF list TXT:
    One transcription factor gene name per line

Important note
--------------
GENIE3 infers candidate regulatory relationships from expression data.
It does not prove causality. Top-ranked edges should be validated using
literature, ChIP-seq, perturbation experiments, or curated databases.

References
----------
Huynh-Thu et al. (2010), PLoS ONE
Moerman et al. (2019), Bioinformatics
"""

import argparse
from pathlib import Path

import pandas as pd
import networkx as nx


def parse_arguments():
    parser = argparse.ArgumentParser(
        description="Infer Gene Regulatory Networks from RNA-seq data using GENIE3."
    )

    parser.add_argument(
        "--expression",
        required=True,
        help="Path to normalized expression matrix CSV. Rows = genes, columns = samples."
    )

    parser.add_argument(
        "--tf_list",
        required=True,
        help="Path to transcription factor list TXT file. One TF gene symbol per line."
    )

    parser.add_argument(
        "--output_dir",
        default="genie3_results",
        help="Directory where output files will be saved."
    )

    parser.add_argument(
        "--min_mean_expression",
        type=float,
        default=1.0,
        help="Minimum mean expression threshold for filtering genes. Default: 1.0"
    )

    parser.add_argument(
        "--top_edges",
        type=int,
        default=1000,
        help="Number of top-ranked regulatory edges to export. Default: 1000"
    )

    parser.add_argument(
        "--method",
        choices=["genie3", "grnboost2"],
        default="genie3",
        help="Network inference method. Default: genie3"
    )

    return parser.parse_args()


def load_expression_matrix(expression_path):
    """Load expression matrix with genes as rows and samples as columns."""
    expression_path = Path(expression_path)

    if not expression_path.exists():
        raise FileNotFoundError(f"Expression matrix not found: {expression_path}")

    expr_df = pd.read_csv(expression_path, index_col=0)

    if expr_df.empty:
        raise ValueError("Expression matrix is empty.")

    expr_df = expr_df.apply(pd.to_numeric, errors="coerce")

    if expr_df.isna().any().any():
        raise ValueError(
            "Expression matrix contains non-numeric or missing values. "
            "Please clean the matrix before running GENIE3."
        )

    print(f"[INFO] Loaded expression matrix: {expr_df.shape[0]} genes x {expr_df.shape[1]} samples")
    return expr_df


def filter_low_expression_genes(expr_df, min_mean_expression):
    """Remove genes with mean expression below threshold."""
    mean_expression = expr_df.mean(axis=1)
    filtered_df = expr_df.loc[mean_expression > min_mean_expression].copy()

    removed = expr_df.shape[0] - filtered_df.shape[0]

    print(f"[INFO] Removed {removed} low-expression genes")
    print(f"[INFO] Remaining genes: {filtered_df.shape[0]}")

    if filtered_df.empty:
        raise ValueError("No genes remain after filtering. Try lowering --min_mean_expression.")

    return filtered_df


def load_tf_list(tf_path, available_genes):
    """Load transcription factor list and keep only TFs found in expression matrix."""
    tf_path = Path(tf_path)

    if not tf_path.exists():
        raise FileNotFoundError(f"TF list not found: {tf_path}")

    tf_df = pd.read_csv(tf_path, header=None)
    tf_list = tf_df.iloc[:, 0].astype(str).str.strip().tolist()

    tf_list = [tf for tf in tf_list if tf in available_genes]

    print(f"[INFO] TFs found in expression matrix: {len(tf_list)}")

    if len(tf_list) == 0:
        raise ValueError(
            "No transcription factors from the TF list were found in the expression matrix. "
            "Check gene naming conventions, e.g. HGNC symbols vs Ensembl IDs."
        )

    return tf_list


def run_network_inference(expr_df, tf_list, method="genie3"):
    """
    Run GENIE3 or GRNBoost2 using arboreto.

    Arboreto expects:
        rows    = samples
        columns = genes
    """
    try:
        from arboreto.algo import genie3, grnboost2
        from dask.distributed import Client, LocalCluster
    except ImportError as exc:
        raise ImportError(
            "Required package not installed. Please install dependencies with:\n"
            "pip install arboreto dask distributed pandas networkx"
        ) from exc

    expr_matrix = expr_df.T

    print(f"[INFO] Starting {method.upper()} inference")
    print(f"[INFO] Input to arboreto: {expr_matrix.shape[0]} samples x {expr_matrix.shape[1]} genes")

    cluster = LocalCluster(
        n_workers=2,
        threads_per_worker=2,
        processes=True,
        dashboard_address=None
    )
    client = Client(cluster)

    try:
        if method == "genie3":
            network = genie3(
                expression_data=expr_matrix,
                tf_names=tf_list,
                client_or_address=client
            )
        else:
            network = grnboost2(
                expression_data=expr_matrix,
                tf_names=tf_list,
                client_or_address=client
            )
    finally:
        client.close()
        cluster.close()

    if {"TF", "target", "importance"}.issubset(network.columns):
        network = network.rename(columns={"TF": "regulator"})
    elif {"regulator", "target", "importance"}.issubset(network.columns):
        pass
    else:
        raise ValueError(f"Unexpected arboreto output columns: {network.columns.tolist()}")

    network = network[["regulator", "target", "importance"]]
    network = network.sort_values("importance", ascending=False).reset_index(drop=True)

    print(f"[INFO] Inferred regulatory links: {network.shape[0]}")
    return network


def export_results(network, output_dir, top_n):
    """Export full network, top edges, hub scores, and graph summary."""
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    full_path = output_dir / "genie3_full_ranked_edges.csv"
    top_path = output_dir / f"genie3_top_{top_n}_edges.csv"
    hub_path = output_dir / "candidate_master_regulators.csv"
    graph_summary_path = output_dir / "network_summary.txt"

    network.to_csv(full_path, index=False)

    top_edges = network.head(top_n).copy()
    top_edges.to_csv(top_path, index=False)

    G = nx.from_pandas_edgelist(
        top_edges,
        source="regulator",
        target="target",
        edge_attr="importance",
        create_using=nx.DiGraph()
    )

    hub_scores = (
        top_edges["regulator"]
        .value_counts()
        .rename_axis("transcription_factor")
        .reset_index(name="predicted_target_count")
    )

    hub_scores.to_csv(hub_path, index=False)

    with open(graph_summary_path, "w", encoding="utf-8") as f:
        f.write("GENIE3 GRN Network Summary\n")
        f.write("==========================\n\n")
        f.write(f"Total ranked edges exported: {network.shape[0]}\n")
        f.write(f"Top edges used for network: {top_edges.shape[0]}\n")
        f.write(f"Number of network nodes: {G.number_of_nodes()}\n")
        f.write(f"Number of network edges: {G.number_of_edges()}\n\n")
        f.write("Top candidate master regulators:\n")
        f.write(hub_scores.head(20).to_string(index=False))

    print("[INFO] Results exported:")
    print(f"       Full ranked edge list: {full_path}")
    print(f"       Top edge list:         {top_path}")
    print(f"       Hub TF table:          {hub_path}")
    print(f"       Network summary:       {graph_summary_path}")

    return {
        "full_edges": full_path,
        "top_edges": top_path,
        "hub_scores": hub_path,
        "summary": graph_summary_path
    }


def main():
    args = parse_arguments()

    expr_df = load_expression_matrix(args.expression)

    filtered_expr_df = filter_low_expression_genes(
        expr_df,
        min_mean_expression=args.min_mean_expression
    )

    tf_list = load_tf_list(
        args.tf_list,
        available_genes=filtered_expr_df.index
    )

    network = run_network_inference(
        filtered_expr_df,
        tf_list=tf_list,
        method=args.method
    )

    export_results(
        network,
        output_dir=args.output_dir,
        top_n=args.top_edges
    )

    print("[DONE] GRN inference pipeline completed successfully.")


if __name__ == "__main__":
    main()