risk-network 0.0.12b0__py3-none-any.whl → 0.0.12b1__py3-none-any.whl

risk/neighborhoods/stats/permutation/permutation.py

@@ -0,0 +1,240 @@
+"""
+risk/neighborhoods/stats/permutation/permutation
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+"""
+
+from multiprocessing import Manager, get_context
+from multiprocessing.managers import ValueProxy
+from typing import Any, Callable, Dict, List, Tuple, Union
+
+import numpy as np
+from scipy.sparse import csr_matrix
+from threadpoolctl import threadpool_limits
+from tqdm import tqdm
+
+from risk.neighborhoods.stats.permutation.test_functions import DISPATCH_TEST_FUNCTIONS
+
+
+def compute_permutation_test(
+    neighborhoods: csr_matrix,
+    annotations: csr_matrix,
+    score_metric: str = "sum",
+    null_distribution: str = "network",
+    num_permutations: int = 1000,
+    random_seed: int = 888,
+    max_workers: int = 1,
+) -> Dict[str, Any]:
+    """Compute permutation test for enrichment and depletion in neighborhoods.
+
+    Args:
+        neighborhoods (csr_matrix): Sparse binary matrix representing neighborhoods.
+        annotations (csr_matrix): Sparse binary matrix representing annotations.
+        score_metric (str, optional): Metric to use for scoring ('sum' or 'stdev'). Defaults to "sum".
+        null_distribution (str, optional): Type of null distribution ('network' or 'annotations'). Defaults to "network".
+        num_permutations (int, optional): Number of permutations to run. Defaults to 1000.
+        random_seed (int, optional): Seed for random number generation. Defaults to 888.
+        max_workers (int, optional): Number of workers for multiprocessing. Defaults to 1.
+
+    Returns:
+        Dict[str, Any]: Dictionary containing depletion and enrichment p-values.
+    """
+    # Ensure that the matrices are in the correct format and free of NaN values
+    # NOTE: Keep the data type as float32 to avoid locking issues with dot product operations
+    neighborhoods = neighborhoods.astype(np.float32)
+    annotations = annotations.astype(np.float32)
+    # Retrieve the appropriate neighborhood score function based on the metric
+    neighborhood_score_func = DISPATCH_TEST_FUNCTIONS[score_metric]
+
+    # Run the permutation test to calculate depletion and enrichment counts
+    counts_depletion, counts_enrichment = _run_permutation_test(
+        neighborhoods=neighborhoods,
+        annotations=annotations,
+        neighborhood_score_func=neighborhood_score_func,
+        null_distribution=null_distribution,
+        num_permutations=num_permutations,
+        random_seed=random_seed,
+        max_workers=max_workers,
+    )
+    # Compute p-values for depletion and enrichment
+    # If counts are 0, set p-value to 1/num_permutations to avoid zero p-values
+    depletion_pvals = np.maximum(counts_depletion, 1) / num_permutations
+    enrichment_pvals = np.maximum(counts_enrichment, 1) / num_permutations
+
+    return {
+        "depletion_pvals": depletion_pvals,
+        "enrichment_pvals": enrichment_pvals,
+    }
+
+
+def _run_permutation_test(
+    neighborhoods: csr_matrix,
+    annotations: csr_matrix,
+    neighborhood_score_func: Callable,
+    null_distribution: str = "network",
+    num_permutations: int = 1000,
+    random_seed: int = 888,
+    max_workers: int = 4,
+) -> tuple:
+    """Run the permutation test to calculate depletion and enrichment counts.
+
+    Args:
+        neighborhoods (csr_matrix): Sparse binary matrix representing neighborhoods.
+        annotations (csr_matrix): Sparse binary matrix representing annotations.
+        neighborhood_score_func (Callable): Function to calculate neighborhood scores.
+        null_distribution (str, optional): Type of null distribution ('network' or 'annotations'). Defaults to "network".
+        num_permutations (int, optional): Number of permutations. Defaults to 1000.
+        random_seed (int, optional): Seed for random number generation. Defaults to 888.
+        max_workers (int, optional): Number of workers for multiprocessing. Defaults to 4.
+
+    Returns:
+        tuple: Depletion and enrichment counts.
+
+    Raises:
+        ValueError: If an invalid null_distribution value is provided.
+    """
+    # Initialize the RNG for reproducibility
+    rng = np.random.default_rng(seed=random_seed)
+    # Determine the indices to use based on the null distribution type
+    if null_distribution == "network":
+        idxs = range(annotations.shape[0])
+    elif null_distribution == "annotations":
+        idxs = np.nonzero(annotations.getnnz(axis=1) > 0)[0]
+    else:
+        raise ValueError(
+            "Invalid null_distribution value. Choose either 'network' or 'annotations'."
+        )
+
+    # Replace NaNs with zeros in the sparse annotations matrix
+    annotations.data[np.isnan(annotations.data)] = 0
+    annotation_matrix_obsv = annotations[idxs]
+    neighborhoods_matrix_obsv = neighborhoods.T[idxs].T
+    # Calculate observed neighborhood scores
+    with np.errstate(invalid="ignore", divide="ignore"):
+        observed_neighborhood_scores = neighborhood_score_func(
+            neighborhoods_matrix_obsv, annotation_matrix_obsv
+        )
+
+    # Initialize count matrices for depletion and enrichment
+    counts_depletion = np.zeros(observed_neighborhood_scores.shape)
+    counts_enrichment = np.zeros(observed_neighborhood_scores.shape)
+    # Determine the number of permutations to run in each worker process
+    subset_size = num_permutations // max_workers
+    remainder = num_permutations % max_workers
+
+    # Use the spawn context for creating a new multiprocessing pool
+    ctx = get_context("spawn")
+    manager = Manager()
+    progress_counter = manager.Value("i", 0)
+    total_progress = num_permutations
+
+    # Generate precomputed permutations
+    permutations = [rng.permutation(idxs) for _ in range(num_permutations)]
+    # Divide permutations into batches for workers
+    batch_size = subset_size + (1 if remainder > 0 else 0)
+    permutation_batches = [
+        permutations[i * batch_size : (i + 1) * batch_size] for i in range(max_workers)
+    ]
+
+    # Execute the permutation test using multiprocessing
+    with ctx.Pool(max_workers) as pool:
+        with tqdm(total=total_progress, desc="Total progress", position=0) as progress:
+            # Prepare parameters for multiprocessing
+            params_list = [
+                (
+                    permutation_batches[i],  # Pass the batch of precomputed permutations
+                    annotations,
+                    neighborhoods_matrix_obsv,
+                    observed_neighborhood_scores,
+                    neighborhood_score_func,
+                    num_permutations,
+                    progress_counter,
+                    max_workers,
+                )
+                for i in range(max_workers)
+            ]
+
+            # Start the permutation process in parallel
+            results = pool.starmap_async(_permutation_process_batch, params_list, chunksize=1)
+
+            # Update progress bar based on progress_counter
+            while not results.ready():
+                progress.update(progress_counter.value - progress.n)
+                results.wait(0.1)  # Wait for 100ms
+            # Ensure progress bar reaches 100%
+            progress.update(total_progress - progress.n)
+
+    # Accumulate results from each worker
+    for local_counts_depletion, local_counts_enrichment in results.get():
+        counts_depletion = np.add(counts_depletion, local_counts_depletion)
+        counts_enrichment = np.add(counts_enrichment, local_counts_enrichment)
+
+    return counts_depletion, counts_enrichment
+
+
+def _permutation_process_batch(
+    permutations: Union[List, Tuple, np.ndarray],
+    annotation_matrix: csr_matrix,
+    neighborhoods_matrix_obsv: csr_matrix,
+    observed_neighborhood_scores: np.ndarray,
+    neighborhood_score_func: Callable,
+    num_permutations: int,
+    progress_counter: ValueProxy,
+    max_workers: int,
+) -> tuple:
+    """Process a batch of permutations in a worker process.
+
+    Args:
+        permutations (Union[List, Tuple, np.ndarray]): Permutation batch to process.
+        annotation_matrix (csr_matrix): Sparse binary matrix representing annotations.
+        neighborhoods_matrix_obsv (csr_matrix): Sparse binary matrix representing observed neighborhoods.
+        observed_neighborhood_scores (np.ndarray): Observed neighborhood scores.
+        neighborhood_score_func (Callable): Function to calculate neighborhood scores.
+        num_permutations (int): Number of total permutations across all subsets.
+        progress_counter (multiprocessing.managers.ValueProxy): Shared counter for tracking progress.
+        max_workers (int): Number of workers for multiprocessing.
+
+    Returns:
+        tuple: Local counts of depletion and enrichment.
+    """
+    # Initialize local count matrices for this worker
+    local_counts_depletion = np.zeros(observed_neighborhood_scores.shape)
+    local_counts_enrichment = np.zeros(observed_neighborhood_scores.shape)
+
+    # Limit the number of threads used by NumPy's BLAS implementation to 1 when more than one worker is used
+    # NOTE: This does not work for Mac M chips due to a bug in the threadpoolctl package
+    # This is currently a known issue and is being addressed by the maintainers [https://github.com/joblib/threadpoolctl/issues/135]
+    limits = None if max_workers == 1 else 1
+    with threadpool_limits(limits=limits, user_api="blas"):
+        # Initialize a local counter for batched progress updates
+        local_progress = 0
+        # Calculate the modulo value based on total permutations for 1/100th frequency updates
+        modulo_value = max(1, num_permutations // 100)
+
+        for permuted_idxs in permutations:
+            # Apply precomputed permutation
+            annotation_matrix_permut = annotation_matrix[permuted_idxs]
+            # Calculate permuted neighborhood scores
+            with np.errstate(invalid="ignore", divide="ignore"):
+                permuted_neighborhood_scores = neighborhood_score_func(
+                    neighborhoods_matrix_obsv, annotation_matrix_permut
+                )
+
+            # Update local depletion and enrichment counts
+            local_counts_depletion = np.add(
+                local_counts_depletion, permuted_neighborhood_scores <= observed_neighborhood_scores
+            )
+            local_counts_enrichment = np.add(
+                local_counts_enrichment,
+                permuted_neighborhood_scores >= observed_neighborhood_scores,
+            )
+
+            # Update progress
+            local_progress += 1
+            if local_progress % modulo_value == 0:
+                progress_counter.value += modulo_value
+
+        # Final progress update for any remaining iterations
+        if local_progress % modulo_value != 0:
+            progress_counter.value += modulo_value
+
+    return local_counts_depletion, local_counts_enrichment

risk/neighborhoods/stats/permutation/test_functions.py

@@ -0,0 +1,70 @@
+"""
+risk/neighborhoods/stats/permutation/test_functions
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+"""
+
+import numpy as np
+from scipy.sparse import csr_matrix
+
+# NOTE: Cython optimizations provided minimal performance benefits.
+# The final version with Cython is archived in the `cython_permutation` branch.
+
+# DISPATCH_TEST_FUNCTIONS can be found at the end of the file.
+
+
+def compute_neighborhood_score_by_sum(
+    neighborhoods_matrix: csr_matrix, annotation_matrix: csr_matrix
+) -> np.ndarray:
+    """Compute the sum of attribute values for each neighborhood using sparse matrices.
+
+    Args:
+        neighborhoods_matrix (csr_matrix): Sparse binary matrix representing neighborhoods.
+        annotation_matrix (csr_matrix): Sparse matrix representing annotation values.
+
+    Returns:
+        np.ndarray: Dense array of summed attribute values for each neighborhood.
+    """
+    # Calculate the neighborhood score as the dot product of neighborhoods and annotations
+    neighborhood_score = neighborhoods_matrix @ annotation_matrix  # Sparse matrix multiplication
+    # Convert the result to a dense array for downstream calculations
+    neighborhood_score_dense = neighborhood_score.toarray()
+    return neighborhood_score_dense
+
+
+def compute_neighborhood_score_by_stdev(
+    neighborhoods_matrix: csr_matrix, annotation_matrix: csr_matrix
+) -> np.ndarray:
+    """Compute the standard deviation of neighborhood scores for sparse matrices.
+
+    Args:
+        neighborhoods_matrix (csr_matrix): Sparse binary matrix representing neighborhoods.
+        annotation_matrix (csr_matrix): Sparse matrix representing annotation values.
+
+    Returns:
+        np.ndarray: Standard deviation of the neighborhood scores.
+    """
+    # Calculate the neighborhood score as the dot product of neighborhoods and annotations
+    neighborhood_score = neighborhoods_matrix @ annotation_matrix  # Sparse matrix multiplication
+    # Calculate the number of elements in each neighborhood (sum of rows)
+    N = neighborhoods_matrix.sum(axis=1).A.flatten()  # Convert to 1D array
+    # Avoid division by zero by replacing zeros in N with np.nan temporarily
+    N[N == 0] = np.nan
+    # Compute the mean of the neighborhood scores
+    M = neighborhood_score.multiply(1 / N[:, None]).toarray()  # Sparse element-wise division
+    # Compute the mean of squares (EXX) directly using squared annotation matrix
+    annotation_squared = annotation_matrix.multiply(annotation_matrix)  # Element-wise squaring
+    EXX = (neighborhoods_matrix @ annotation_squared).multiply(1 / N[:, None]).toarray()
+    # Calculate variance as EXX - M^2
+    variance = EXX - np.power(M, 2)
+    # Compute the standard deviation as the square root of the variance
+    neighborhood_stdev = np.sqrt(variance)
+    # Replace np.nan back with zeros in case N was 0 (no elements in the neighborhood)
+    neighborhood_stdev[np.isnan(neighborhood_stdev)] = 0
+    return neighborhood_stdev
+
+
+# Dictionary to dispatch statistical test functions based on the score metric
+DISPATCH_TEST_FUNCTIONS = {
+    "sum": compute_neighborhood_score_by_sum,
+    "stdev": compute_neighborhood_score_by_stdev,
+}

risk/neighborhoods/stats/tests.py

@@ -0,0 +1,275 @@
+"""
+risk/neighborhoods/stats/tests
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+"""
+
+from typing import Any, Dict
+
+import numpy as np
+from scipy.sparse import csr_matrix
+from scipy.stats import binom, chi2, hypergeom, norm, poisson
+
+
+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.
+
+    Raises:
+        ValueError: If an invalid null_distribution value is provided.
+    """
+    # 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}
+
+
+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.
+
+    Raises:
+        ValueError: If an invalid null_distribution value is provided.
+    """
+    # 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}
+
+
+def compute_hypergeom_test(
+    neighborhoods: csr_matrix,
+    annotations: csr_matrix,
+    null_distribution: str = "network",
+) -> Dict[str, Any]:
+    """Compute hypergeometric 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.
+
+    Raises:
+        ValueError: If an invalid null_distribution value is provided.
+    """
+    # Get the total number of nodes in the network
+    total_nodes = neighborhoods.shape[1]
+
+    # Compute sums
+    neighborhood_sums = neighborhoods.sum(axis=0).A.flatten()  # Convert to dense array
+    annotation_sums = annotations.sum(axis=0).A.flatten()  # Convert to dense array
+
+    if null_distribution == "network":
+        background_population = total_nodes
+    elif null_distribution == "annotations":
+        annotated_nodes = annotations.sum(axis=1).A.flatten() > 0  # Boolean mask
+        background_population = annotated_nodes.sum()
+        neighborhood_sums = neighborhoods[annotated_nodes].sum(axis=0).A.flatten()
+        annotation_sums = annotations[annotated_nodes].sum(axis=0).A.flatten()
+    else:
+        raise ValueError(
+            "Invalid null_distribution value. Choose either 'network' or 'annotations'."
+        )
+
+    # Observed counts
+    annotated_in_neighborhood = neighborhoods.T @ annotations  # Sparse result
+    annotated_in_neighborhood = annotated_in_neighborhood.toarray()  # Convert to dense
+    # Align shapes for broadcasting
+    neighborhood_sums = neighborhood_sums.reshape(-1, 1)
+    annotation_sums = annotation_sums.reshape(1, -1)
+    background_population = np.array(background_population).reshape(1, 1)
+
+    # Compute hypergeometric p-values
+    depletion_pvals = hypergeom.cdf(
+        annotated_in_neighborhood, background_population, annotation_sums, neighborhood_sums
+    )
+    enrichment_pvals = hypergeom.sf(
+        annotated_in_neighborhood - 1, background_population, annotation_sums, neighborhood_sums
+    )
+
+    return {"depletion_pvals": depletion_pvals, "enrichment_pvals": enrichment_pvals}
+
+
+def compute_poisson_test(
+    neighborhoods: csr_matrix,
+    annotations: csr_matrix,
+    null_distribution: str = "network",
+) -> Dict[str, Any]:
+    """Compute Poisson 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.
+
+    Raises:
+        ValueError: If an invalid null_distribution value is provided.
+    """
+    # Matrix multiplication to get the number of annotated nodes in each neighborhood
+    annotated_in_neighborhood = neighborhoods @ annotations  # Sparse result
+    # Convert annotated counts to dense for downstream calculations
+    annotated_in_neighborhood_dense = annotated_in_neighborhood.toarray()
+
+    # Compute lambda_expected based on the chosen null distribution
+    if null_distribution == "network":
+        # Use the mean across neighborhoods (axis=1)
+        lambda_expected = np.mean(annotated_in_neighborhood_dense, axis=1, keepdims=True)
+    elif null_distribution == "annotations":
+        # Use the mean across annotations (axis=0)
+        lambda_expected = np.mean(annotated_in_neighborhood_dense, axis=0, keepdims=True)
+    else:
+        raise ValueError(
+            "Invalid null_distribution value. Choose either 'network' or 'annotations'."
+        )
+
+    # Compute p-values for enrichment and depletion using Poisson distribution
+    enrichment_pvals = 1 - poisson.cdf(annotated_in_neighborhood_dense - 1, lambda_expected)
+    depletion_pvals = poisson.cdf(annotated_in_neighborhood_dense, lambda_expected)
+
+    return {"enrichment_pvals": enrichment_pvals, "depletion_pvals": depletion_pvals}
+
+
+def compute_zscore_test(
+    neighborhoods: csr_matrix,
+    annotations: csr_matrix,
+    null_distribution: str = "network",
+) -> Dict[str, Any]:
+    """Compute z-score 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.
+
+    Raises:
+        ValueError: If an invalid null_distribution value is provided.
+    """
+    # Total number of nodes in the network
+    total_node_count = neighborhoods.shape[1]
+
+    # Compute sums
+    if null_distribution == "network":
+        background_population = total_node_count
+        neighborhood_sums = neighborhoods.sum(axis=0).A.flatten()  # Dense column sums
+        annotation_sums = annotations.sum(axis=0).A.flatten()  # Dense row sums
+    elif null_distribution == "annotations":
+        annotated_nodes = annotations.sum(axis=1).A.flatten() > 0  # Dense boolean mask
+        background_population = annotated_nodes.sum()
+        neighborhood_sums = neighborhoods[annotated_nodes].sum(axis=0).A.flatten()
+        annotation_sums = annotations[annotated_nodes].sum(axis=0).A.flatten()
+    else:
+        raise ValueError(
+            "Invalid null_distribution value. Choose either 'network' or 'annotations'."
+        )
+
+    # Observed values
+    observed = (neighborhoods.T @ annotations).toarray()  # Convert sparse result to dense
+    # Expected values under the null
+    neighborhood_sums = neighborhood_sums.reshape(-1, 1)  # Ensure correct shape
+    annotation_sums = annotation_sums.reshape(1, -1)  # Ensure correct shape
+    expected = (neighborhood_sums @ annotation_sums) / background_population
+
+    # Standard deviation under the null
+    std_dev = np.sqrt(
+        expected
+        * (1 - annotation_sums / background_population)
+        * (1 - neighborhood_sums / background_population)
+    )
+    std_dev[std_dev == 0] = np.nan  # Avoid division by zero
+    # Compute z-scores
+    z_scores = (observed - expected) / std_dev
+
+    # Convert z-scores to depletion and enrichment p-values
+    enrichment_pvals = norm.sf(z_scores)  # Upper tail
+    depletion_pvals = norm.cdf(z_scores)  # Lower tail
+
+    return {"depletion_pvals": depletion_pvals, "enrichment_pvals": enrichment_pvals}

risk/network/__init__.py

@@ -0,0 +1,4 @@
+"""
+risk/network
+~~~~~~~~~~~~
+"""

risk/network/graph/__init__.py

@@ -0,0 +1,4 @@
+"""
+risk/network/graph
+~~~~~~~~~~~~~~~~~~
+"""