risk-network 0.0.8b18__py3-none-any.whl → 0.0.9b26__py3-none-any.whl

risk/risk.py

@@ -3,34 +3,21 @@ risk/risk
 ~~~~~~~~~
 """
 
-from typing import Any, Dict, Tuple, Union
+from risk.network import NetworkIO
+from risk.annotations import AnnotationsIO
+from risk.neighborhoods import NeighborhoodsAPI
+from risk.network.graph import GraphAPI
+from risk.network.plotter import PlotterAPI
 
-import networkx as nx
-import numpy as np
-import pandas as pd
+from risk.log import params, set_global_verbosity
 
-from risk.annotations import AnnotationsIO, define_top_annotations
-from risk.log import params, logger, log_header, set_global_verbosity
-from risk.neighborhoods import (
-    define_domains,
-    get_network_neighborhoods,
-    process_neighborhoods,
-    trim_domains_and_top_annotations,
-)
-from risk.network import NetworkIO, NetworkGraph, NetworkPlotter
-from risk.stats import (
-    calculate_significance_matrices,
-    compute_hypergeom_test,
-    compute_permutation_test,
-    compute_poisson_test,
-)
 
-
-class RISK(NetworkIO, AnnotationsIO):
+class RISK(NetworkIO, AnnotationsIO, NeighborhoodsAPI, GraphAPI, PlotterAPI):
     """RISK: A class for network analysis and visualization.
 
     The RISK class integrates functionalities for loading networks, processing annotations,
-    and performing network-based statistical analysis, such as neighborhood significance testing.
+    performing network-based statistical analysis to quantify neighborhood relationships,
+    and visualizing networks and their properties.
     """
 
     def __init__(self, verbose: bool = True):
@@ -41,451 +28,6 @@ class RISK(NetworkIO, AnnotationsIO):
         """
         # Set global verbosity for logging
         set_global_verbosity(verbose)
-        # Initialize and log network parameters
-        params.initialize()
+        # Provide public access to network parameters
+        self.params = params
         super().__init__()
-
-    @property
-    def params(self) -> params:
-        """Access the logged network parameters.
-
-        Returns:
-            Params: An instance of the Params class with logged parameters and methods to access or update them.
-        """
-        return params
-
-    def load_neighborhoods_by_hypergeom(
-        self,
-        network: nx.Graph,
-        annotations: Dict[str, Any],
-        distance_metric: str = "louvain",
-        louvain_resolution: float = 0.1,
-        edge_length_threshold: float = 0.5,
-        null_distribution: str = "network",
-        random_seed: int = 888,
-    ) -> Dict[str, Any]:
-        """Load significant neighborhoods for the network using the hypergeometric test.
-
-        Args:
-            network (nx.Graph): The network graph.
-            annotations (dict): The annotations associated with the network.
-            distance_metric (str, optional): Distance metric for neighborhood analysis. Defaults to "louvain".
-            louvain_resolution (float, optional): Resolution parameter for Louvain clustering. Defaults to 0.1.
-            edge_length_threshold (float, optional): Edge length threshold for neighborhood analysis. Defaults to 0.5.
-            null_distribution (str, optional): Type of null distribution ('network' or 'annotations'). Defaults to "network".
-            random_seed (int, optional): Seed for random number generation. Defaults to 888.
-
-        Returns:
-            dict: Computed significance of neighborhoods.
-        """
-        log_header("Running hypergeometric test")
-        # Log neighborhood analysis parameters
-        params.log_neighborhoods(
-            distance_metric=distance_metric,
-            louvain_resolution=louvain_resolution,
-            edge_length_threshold=edge_length_threshold,
-            statistical_test_function="hypergeom",
-            null_distribution=null_distribution,
-            random_seed=random_seed,
-        )
-
-        # Load neighborhoods based on the network and distance metric
-        neighborhoods = self._load_neighborhoods(
-            network,
-            distance_metric,
-            louvain_resolution=louvain_resolution,
-            edge_length_threshold=edge_length_threshold,
-            random_seed=random_seed,
-        )
-        # Run hypergeometric test to compute neighborhood significance
-        neighborhood_significance = compute_hypergeom_test(
-            neighborhoods=neighborhoods,
-            annotations=annotations["matrix"],
-            null_distribution=null_distribution,
-        )
-
-        # Return the computed neighborhood significance
-        return neighborhood_significance
-
-    def load_neighborhoods_by_poisson(
-        self,
-        network: nx.Graph,
-        annotations: Dict[str, Any],
-        distance_metric: str = "louvain",
-        louvain_resolution: float = 0.1,
-        edge_length_threshold: float = 0.5,
-        null_distribution: str = "network",
-        random_seed: int = 888,
-    ) -> Dict[str, Any]:
-        """Load significant neighborhoods for the network using the Poisson test.
-
-        Args:
-            network (nx.Graph): The network graph.
-            annotations (dict): The annotations associated with the network.
-            distance_metric (str, optional): Distance metric for neighborhood analysis. Defaults to "louvain".
-            louvain_resolution (float, optional): Resolution parameter for Louvain clustering. Defaults to 0.1.
-            edge_length_threshold (float, optional): Edge length threshold for neighborhood analysis. Defaults to 0.5.
-            null_distribution (str, optional): Type of null distribution ('network' or 'annotations'). Defaults to "network".
-            random_seed (int, optional): Seed for random number generation. Defaults to 888.
-
-        Returns:
-            dict: Computed significance of neighborhoods.
-        """
-        log_header("Running Poisson test")
-        # Log neighborhood analysis parameters
-        params.log_neighborhoods(
-            distance_metric=distance_metric,
-            louvain_resolution=louvain_resolution,
-            edge_length_threshold=edge_length_threshold,
-            statistical_test_function="poisson",
-            null_distribution=null_distribution,
-            random_seed=random_seed,
-        )
-
-        # Load neighborhoods based on the network and distance metric
-        neighborhoods = self._load_neighborhoods(
-            network,
-            distance_metric,
-            louvain_resolution=louvain_resolution,
-            edge_length_threshold=edge_length_threshold,
-            random_seed=random_seed,
-        )
-        # Run Poisson test to compute neighborhood significance
-        neighborhood_significance = compute_poisson_test(
-            neighborhoods=neighborhoods,
-            annotations=annotations["matrix"],
-            null_distribution=null_distribution,
-        )
-
-        # Return the computed neighborhood significance
-        return neighborhood_significance
-
-    def load_neighborhoods_by_permutation(
-        self,
-        network: nx.Graph,
-        annotations: Dict[str, Any],
-        distance_metric: str = "louvain",
-        louvain_resolution: float = 0.1,
-        edge_length_threshold: float = 0.5,
-        score_metric: str = "sum",
-        null_distribution: str = "network",
-        num_permutations: int = 1000,
-        random_seed: int = 888,
-        max_workers: int = 1,
-    ) -> Dict[str, Any]:
-        """Load significant neighborhoods for the network using the permutation test.
-
-        Args:
-            network (nx.Graph): The network graph.
-            annotations (dict): The annotations associated with the network.
-            distance_metric (str, optional): Distance metric for neighborhood analysis. Defaults to "louvain".
-            louvain_resolution (float, optional): Resolution parameter for Louvain clustering. Defaults to 0.1.
-            edge_length_threshold (float, optional): Edge length threshold for neighborhood analysis. Defaults to 0.5.
-            score_metric (str, optional): Scoring metric for neighborhood significance. Defaults to "sum".
-            null_distribution (str, optional): Type of null distribution ('network' or 'annotations'). Defaults to "network".
-            num_permutations (int, optional): Number of permutations for significance testing. Defaults to 1000.
-            random_seed (int, optional): Seed for random number generation. Defaults to 888.
-            max_workers (int, optional): Maximum number of workers for parallel computation. Defaults to 1.
-
-        Returns:
-            dict: Computed significance of neighborhoods.
-        """
-        log_header("Running permutation test")
-        # Log neighborhood analysis parameters
-        params.log_neighborhoods(
-            distance_metric=distance_metric,
-            louvain_resolution=louvain_resolution,
-            edge_length_threshold=edge_length_threshold,
-            statistical_test_function="permutation",
-            score_metric=score_metric,
-            null_distribution=null_distribution,
-            num_permutations=num_permutations,
-            random_seed=random_seed,
-            max_workers=max_workers,
-        )
-
-        # Load neighborhoods based on the network and distance metric
-        neighborhoods = self._load_neighborhoods(
-            network,
-            distance_metric,
-            louvain_resolution=louvain_resolution,
-            edge_length_threshold=edge_length_threshold,
-            random_seed=random_seed,
-        )
-
-        # Log and display permutation test settings
-        logger.debug(f"Neighborhood scoring metric: '{score_metric}'")
-        logger.debug(f"Null distribution: '{null_distribution}'")
-        logger.debug(f"Number of permutations: {num_permutations}")
-        logger.debug(f"Maximum workers: {max_workers}")
-        # Run permutation test to compute neighborhood significance
-        neighborhood_significance = compute_permutation_test(
-            neighborhoods=neighborhoods,
-            annotations=annotations["matrix"],
-            score_metric=score_metric,
-            null_distribution=null_distribution,
-            num_permutations=num_permutations,
-            random_seed=random_seed,
-            max_workers=max_workers,
-        )
-
-        # Return the computed neighborhood significance
-        return neighborhood_significance
-
-    def load_graph(
-        self,
-        network: nx.Graph,
-        annotations: Dict[str, Any],
-        neighborhoods: Dict[str, Any],
-        tail: str = "right",  # OPTIONS: "right" (enrichment), "left" (depletion), "both"
-        pval_cutoff: float = 0.01,  # OPTIONS: Any value between 0 to 1
-        fdr_cutoff: float = 0.9999,  # OPTIONS: Any value between 0 to 1
-        impute_depth: int = 0,
-        prune_threshold: float = 0.0,
-        linkage_criterion: str = "distance",
-        linkage_method: str = "average",
-        linkage_metric: str = "yule",
-        min_cluster_size: int = 5,
-        max_cluster_size: int = 1000,
-    ) -> NetworkGraph:
-        """Load and process the network graph, defining top annotations and domains.
-
-        Args:
-            network (nx.Graph): The network graph.
-            annotations (pd.DataFrame): DataFrame containing annotation data for the network.
-            neighborhoods (dict): Neighborhood enrichment data.
-            tail (str, optional): Type of significance tail ("right", "left", "both"). Defaults to "right".
-            pval_cutoff (float, optional): p-value cutoff for significance. Defaults to 0.01.
-            fdr_cutoff (float, optional): FDR cutoff for significance. Defaults to 0.9999.
-            impute_depth (int, optional): Depth for imputing neighbors. Defaults to 0.
-            prune_threshold (float, optional): Distance threshold for pruning neighbors. Defaults to 0.0.
-            linkage_criterion (str, optional): Clustering criterion for defining domains. Defaults to "distance".
-            linkage_method (str, optional): Clustering method to use. Defaults to "average".
-            linkage_metric (str, optional): Metric to use for calculating distances. Defaults to "yule".
-            min_cluster_size (int, optional): Minimum size for clusters. Defaults to 5.
-            max_cluster_size (int, optional): Maximum size for clusters. Defaults to 1000.
-
-        Returns:
-            NetworkGraph: A fully initialized and processed NetworkGraph object.
-        """
-        # Log the parameters and display headers
-        log_header("Finding significant neighborhoods")
-        params.log_graph(
-            tail=tail,
-            pval_cutoff=pval_cutoff,
-            fdr_cutoff=fdr_cutoff,
-            impute_depth=impute_depth,
-            prune_threshold=prune_threshold,
-            linkage_criterion=linkage_criterion,
-            linkage_method=linkage_method,
-            linkage_metric=linkage_metric,
-            min_cluster_size=min_cluster_size,
-            max_cluster_size=max_cluster_size,
-        )
-
-        logger.debug(f"p-value cutoff: {pval_cutoff}")
-        logger.debug(f"FDR BH cutoff: {fdr_cutoff}")
-        logger.debug(
-            f"Significance tail: '{tail}' ({'enrichment' if tail == 'right' else 'depletion' if tail == 'left' else 'both'})"
-        )
-        # Calculate significant neighborhoods based on the provided parameters
-        significant_neighborhoods = calculate_significance_matrices(
-            neighborhoods["depletion_pvals"],
-            neighborhoods["enrichment_pvals"],
-            tail=tail,
-            pval_cutoff=pval_cutoff,
-            fdr_cutoff=fdr_cutoff,
-        )
-
-        log_header("Processing neighborhoods")
-        # Process neighborhoods by imputing and pruning based on the given settings
-        processed_neighborhoods = process_neighborhoods(
-            network=network,
-            neighborhoods=significant_neighborhoods,
-            impute_depth=impute_depth,
-            prune_threshold=prune_threshold,
-        )
-
-        log_header("Finding top annotations")
-        logger.debug(f"Min cluster size: {min_cluster_size}")
-        logger.debug(f"Max cluster size: {max_cluster_size}")
-        # Define top annotations based on processed neighborhoods
-        top_annotations = self._define_top_annotations(
-            network=network,
-            annotations=annotations,
-            neighborhoods=processed_neighborhoods,
-            min_cluster_size=min_cluster_size,
-            max_cluster_size=max_cluster_size,
-        )
-
-        log_header("Optimizing distance threshold for domains")
-        # Define domains in the network using the specified clustering settings
-        domains = self._define_domains(
-            neighborhoods=processed_neighborhoods,
-            top_annotations=top_annotations,
-            linkage_criterion=linkage_criterion,
-            linkage_method=linkage_method,
-            linkage_metric=linkage_metric,
-        )
-        # Trim domains and top annotations based on cluster size constraints
-        top_annotations, domains, trimmed_domains = trim_domains_and_top_annotations(
-            domains=domains,
-            top_annotations=top_annotations,
-            min_cluster_size=min_cluster_size,
-            max_cluster_size=max_cluster_size,
-        )
-
-        # Prepare node mapping and enrichment sums for the final NetworkGraph object
-        ordered_nodes = annotations["ordered_nodes"]
-        node_label_to_id = dict(zip(ordered_nodes, range(len(ordered_nodes))))
-        node_enrichment_sums = processed_neighborhoods["node_enrichment_sums"]
-
-        # Return the fully initialized NetworkGraph object
-        return NetworkGraph(
-            network=network,
-            top_annotations=top_annotations,
-            domains=domains,
-            trimmed_domains=trimmed_domains,
-            node_label_to_node_id_map=node_label_to_id,
-            node_enrichment_sums=node_enrichment_sums,
-        )
-
-    def load_plotter(
-        self,
-        graph: NetworkGraph,
-        figsize: Tuple = (10, 10),
-        background_color: str = "white",
-        background_alpha: Union[float, None] = 1.0,
-        pad: float = 0.3,
-    ) -> NetworkPlotter:
-        """Get a NetworkPlotter object for plotting.
-
-        Args:
-            graph (NetworkGraph): The graph to plot.
-            figsize (tuple, optional): Size of the figure. Defaults to (10, 10).
-            background_color (str, optional): Background color of the plot. Defaults to "white".
-            background_alpha (float, None, optional): Transparency level of the background color. If provided, it overrides
-                any existing alpha values found in background_color. Defaults to 1.0.
-            pad (float, optional): Padding value to adjust the axis limits. Defaults to 0.3.
-
-        Returns:
-            NetworkPlotter: A NetworkPlotter object configured with the given parameters.
-        """
-        log_header("Loading plotter")
-
-        # Initialize and return a NetworkPlotter object
-        return NetworkPlotter(
-            graph,
-            figsize=figsize,
-            background_color=background_color,
-            background_alpha=background_alpha,
-            pad=pad,
-        )
-
-    def _load_neighborhoods(
-        self,
-        network: nx.Graph,
-        distance_metric: str = "louvain",
-        louvain_resolution: float = 0.1,
-        edge_length_threshold: float = 0.5,
-        random_seed: int = 888,
-    ) -> np.ndarray:
-        """Load significant neighborhoods for the network.
-
-        Args:
-            network (nx.Graph): The network graph.
-            annotations (pd.DataFrame): The matrix of annotations associated with the network.
-            distance_metric (str, optional): Distance metric for neighborhood analysis. Defaults to "louvain".
-            louvain_resolution (float, optional): Resolution parameter for Louvain clustering. Defaults to 0.1.
-            edge_length_threshold (float, optional): Edge length threshold for neighborhood analysis. Defaults to 0.5.
-            random_seed (int, optional): Seed for random number generation. Defaults to 888.
-
-        Returns:
-            np.ndarray: Neighborhood matrix calculated based on the selected distance metric.
-        """
-        # Display the chosen distance metric
-        if distance_metric == "louvain":
-            for_print_distance_metric = f"louvain (resolution={louvain_resolution})"
-        else:
-            for_print_distance_metric = distance_metric
-        # Log and display neighborhood settings
-        logger.debug(f"Distance metric: '{for_print_distance_metric}'")
-        logger.debug(f"Edge length threshold: {edge_length_threshold}")
-        logger.debug(f"Random seed: {random_seed}")
-
-        # Compute neighborhoods based on the network and distance metric
-        neighborhoods = get_network_neighborhoods(
-            network,
-            distance_metric,
-            edge_length_threshold,
-            louvain_resolution=louvain_resolution,
-            random_seed=random_seed,
-        )
-
-        # Return the computed neighborhoods
-        return neighborhoods
-
-    def _define_top_annotations(
-        self,
-        network: nx.Graph,
-        annotations: Dict[str, Any],
-        neighborhoods: Dict[str, Any],
-        min_cluster_size: int = 5,
-        max_cluster_size: int = 1000,
-    ) -> pd.DataFrame:
-        """Define top annotations for the network.
-
-        Args:
-            network (nx.Graph): The network graph.
-            annotations (dict): Annotations data for the network.
-            neighborhoods (dict): Neighborhood enrichment data.
-            min_cluster_size (int, optional): Minimum size for clusters. Defaults to 5.
-            max_cluster_size (int, optional): Maximum size for clusters. Defaults to 1000.
-
-        Returns:
-            dict: Top annotations identified within the network.
-        """
-        # Extract necessary data from annotations and neighborhoods
-        ordered_annotations = annotations["ordered_annotations"]
-        neighborhood_enrichment_sums = neighborhoods["neighborhood_enrichment_counts"]
-        neighborhoods_binary_enrichment_matrix = neighborhoods["binary_enrichment_matrix"]
-        # Call external function to define top annotations
-        return define_top_annotations(
-            network=network,
-            ordered_annotation_labels=ordered_annotations,
-            neighborhood_enrichment_sums=neighborhood_enrichment_sums,
-            binary_enrichment_matrix=neighborhoods_binary_enrichment_matrix,
-            min_cluster_size=min_cluster_size,
-            max_cluster_size=max_cluster_size,
-        )
-
-    def _define_domains(
-        self,
-        neighborhoods: Dict[str, Any],
-        top_annotations: pd.DataFrame,
-        linkage_criterion: str,
-        linkage_method: str,
-        linkage_metric: str,
-    ) -> pd.DataFrame:
-        """Define domains in the network based on enrichment data.
-
-        Args:
-            neighborhoods (dict): Enrichment data for neighborhoods.
-            top_annotations (pd.DataFrame): Enrichment matrix for top annotations.
-            linkage_criterion (str): Clustering criterion for defining domains.
-            linkage_method (str): Clustering method to use.
-            linkage_metric (str): Metric to use for calculating distances.
-
-        Returns:
-            pd.DataFrame: Matrix of defined domains.
-        """
-        # Extract the significant enrichment matrix from the neighborhoods data
-        significant_neighborhoods_enrichment = neighborhoods["significant_enrichment_matrix"]
-        # Call external function to define domains based on the extracted data
-        return define_domains(
-            top_annotations=top_annotations,
-            significant_neighborhoods_enrichment=significant_neighborhoods_enrichment,
-            linkage_criterion=linkage_criterion,
-            linkage_method=linkage_method,
-            linkage_metric=linkage_metric,
-        )

risk/stats/__init__.py

@@ -3,7 +3,11 @@ risk/stats
 ~~~~~~~~~~
 """
 
-from .hypergeom import compute_hypergeom_test
-from .permutation import compute_permutation_test
-from .poisson import compute_poisson_test
-from .stats import calculate_significance_matrices
+from risk.stats.binom import compute_binom_test
+from risk.stats.chi2 import compute_chi2_test
+from risk.stats.hypergeom import compute_hypergeom_test
+from risk.stats.permutation import compute_permutation_test
+from risk.stats.poisson import compute_poisson_test
+from risk.stats.zscore import compute_zscore_test
+
+from risk.stats.stats import calculate_significance_matrices

risk/stats/binom.py

@@ -0,0 +1,51 @@
+"""
+risk/stats/binomial
+~~~~~~~~~~~~~~~~~~~
+"""
+
+from typing import Any, Dict
+
+from scipy.sparse import csr_matrix
+from scipy.stats import binom
+
+
+def compute_binom_test(
+    neighborhoods: csr_matrix,
+    annotations: csr_matrix,
+    null_distribution: str = "network",
+) -> Dict[str, Any]:
+    """Compute Binomial test for enrichment and depletion in neighborhoods with selectable null distribution.
+
+    Args:
+        neighborhoods (csr_matrix): Sparse binary matrix representing neighborhoods.
+        annotations (csr_matrix): Sparse binary matrix representing annotations.
+        null_distribution (str, optional): Type of null distribution ('network' or 'annotations'). Defaults to "network".
+
+    Returns:
+        Dict[str, Any]: Dictionary containing depletion and enrichment p-values.
+    """
+    # Get the total number of nodes in the network
+    total_nodes = neighborhoods.shape[1]
+
+    # Compute sums (remain sparse here)
+    neighborhood_sizes = neighborhoods.sum(axis=1)  # Row sums
+    annotation_totals = annotations.sum(axis=0)  # Column sums
+    # Compute probabilities (convert to dense)
+    if null_distribution == "network":
+        p_values = (annotation_totals / total_nodes).A.flatten()  # Dense 1D array
+    elif null_distribution == "annotations":
+        p_values = (annotation_totals / annotations.sum()).A.flatten()  # Dense 1D array
+    else:
+        raise ValueError(
+            "Invalid null_distribution value. Choose either 'network' or 'annotations'."
+        )
+
+    # Observed counts (sparse matrix multiplication)
+    annotated_counts = neighborhoods @ annotations  # Sparse result
+    annotated_counts_dense = annotated_counts.toarray()  # Convert for dense operations
+
+    # Compute enrichment and depletion p-values
+    enrichment_pvals = 1 - binom.cdf(annotated_counts_dense - 1, neighborhood_sizes.A, p_values)
+    depletion_pvals = binom.cdf(annotated_counts_dense, neighborhood_sizes.A, p_values)
+
+    return {"enrichment_pvals": enrichment_pvals, "depletion_pvals": depletion_pvals}

risk/stats/chi2.py

@@ -0,0 +1,69 @@
+"""
+risk/stats/chi2
+~~~~~~~~~~~~~~~
+"""
+
+from typing import Any, Dict
+
+import numpy as np
+from scipy.sparse import csr_matrix
+from scipy.stats import chi2
+
+
+def compute_chi2_test(
+    neighborhoods: csr_matrix,
+    annotations: csr_matrix,
+    null_distribution: str = "network",
+) -> Dict[str, Any]:
+    """Compute chi-squared test for enrichment and depletion in neighborhoods with selectable null distribution.
+
+    Args:
+        neighborhoods (csr_matrix): Sparse binary matrix representing neighborhoods.
+        annotations (csr_matrix): Sparse binary matrix representing annotations.
+        null_distribution (str, optional): Type of null distribution ('network' or 'annotations'). Defaults to "network".
+
+    Returns:
+        Dict[str, Any]: Dictionary containing depletion and enrichment p-values.
+    """
+    # Total number of nodes in the network
+    total_node_count = neighborhoods.shape[0]
+
+    if null_distribution == "network":
+        # Case 1: Use all nodes as the background
+        background_population = total_node_count
+        neighborhood_sums = neighborhoods.sum(axis=0)  # Column sums of neighborhoods
+        annotation_sums = annotations.sum(axis=0)  # Column sums of annotations
+    elif null_distribution == "annotations":
+        # Case 2: Only consider nodes with at least one annotation
+        annotated_nodes = (
+            np.ravel(annotations.sum(axis=1)) > 0
+        )  # Row-wise sum to filter nodes with annotations
+        background_population = annotated_nodes.sum()  # Total number of annotated nodes
+        neighborhood_sums = neighborhoods[annotated_nodes].sum(
+            axis=0
+        )  # Neighborhood sums for annotated nodes
+        annotation_sums = annotations[annotated_nodes].sum(
+            axis=0
+        )  # Annotation sums for annotated nodes
+    else:
+        raise ValueError(
+            "Invalid null_distribution value. Choose either 'network' or 'annotations'."
+        )
+
+    # Convert to dense arrays for downstream computations
+    neighborhood_sums = np.asarray(neighborhood_sums).reshape(-1, 1)  # Ensure column vector shape
+    annotation_sums = np.asarray(annotation_sums).reshape(1, -1)  # Ensure row vector shape
+
+    # Observed values: number of annotated nodes in each neighborhood
+    observed = neighborhoods.T @ annotations  # Shape: (neighborhoods, annotations)
+    # Expected values under the null
+    expected = (neighborhood_sums @ annotation_sums) / background_population
+    # Chi-squared statistic: sum((observed - expected)^2 / expected)
+    with np.errstate(divide="ignore", invalid="ignore"):  # Handle divide-by-zero
+        chi2_stat = np.where(expected > 0, np.power(observed - expected, 2) / expected, 0)
+
+    # Compute p-values for enrichment (upper tail) and depletion (lower tail)
+    enrichment_pvals = chi2.sf(chi2_stat, df=1)  # Survival function for upper tail
+    depletion_pvals = chi2.cdf(chi2_stat, df=1)  # Cumulative distribution for lower tail
+
+    return {"depletion_pvals": depletion_pvals, "enrichment_pvals": enrichment_pvals}