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

risk/neighborhoods/community.py

@@ -4,186 +4,364 @@ risk/neighborhoods/community
 """
 
 import community as community_louvain
+import igraph as ig
+import markov_clustering as mc
 import networkx as nx
 import numpy as np
-import markov_clustering as mc
-from networkx.algorithms.community import asyn_lpa_communities, greedy_modularity_communities
+from leidenalg import find_partition, RBConfigurationVertexPartition
+from networkx.algorithms.community import greedy_modularity_communities
 
+from risk.log import logger
 
-def calculate_greedy_modularity_neighborhoods(network: nx.Graph) -> np.ndarray:
+
+def calculate_greedy_modularity_neighborhoods(
+    network: nx.Graph, fraction_shortest_edges: float = 1.0
+) -> np.ndarray:
     """Calculate neighborhoods using the Greedy Modularity method.
 
     Args:
-        network (nx.Graph): The network graph to analyze for community structure.
+        network (nx.Graph): The network graph.
+        fraction_shortest_edges (float, optional): Shortest edge rank fraction threshold for creating
+            subgraphs before clustering.
 
     Returns:
         np.ndarray: A binary neighborhood matrix where nodes in the same community have 1, and others have 0.
     """
+    # Create a subgraph with the shortest edges based on the rank fraction
+    subnetwork = _create_percentile_limited_subgraph(
+        network, fraction_shortest_edges=fraction_shortest_edges
+    )
     # Detect communities using the Greedy Modularity method
-    communities = greedy_modularity_communities(network)
+    communities = greedy_modularity_communities(subnetwork)
+    # Get the list of nodes in the original NetworkX graph
+    nodes = list(network.nodes())
+    node_index_map = {node: idx for idx, node in enumerate(nodes)}
     # Create a binary neighborhood matrix
-    n_nodes = network.number_of_nodes()
-    neighborhoods = np.zeros((n_nodes, n_nodes), dtype=int)
-    # Create a mapping from node to index in the matrix
-    node_index = {node: i for i, node in enumerate(network.nodes())}
+    num_nodes = len(nodes)
+    # Initialize neighborhoods with zeros and set self-self entries to 1
+    neighborhoods = np.eye(num_nodes, dtype=int)
     # Fill in the neighborhood matrix for nodes in the same community
     for community in communities:
         # Iterate through all pairs of nodes in the same community
         for node_i in community:
-            idx_i = node_index[node_i]
             for node_j in community:
-                idx_j = node_index[node_j]
+                idx_i = node_index_map[node_i]
+                idx_j = node_index_map[node_j]
                 # Set them as neighbors (1) in the binary matrix
                 neighborhoods[idx_i, idx_j] = 1
 
     return neighborhoods
 
 
-def calculate_label_propagation_neighborhoods(network: nx.Graph) -> np.ndarray:
+def calculate_label_propagation_neighborhoods(
+    network: nx.Graph, fraction_shortest_edges: float = 1.0
+) -> np.ndarray:
     """Apply Label Propagation to the network to detect communities.
 
     Args:
         network (nx.Graph): The network graph.
+        fraction_shortest_edges (float, optional): Shortest edge rank fraction threshold for creating
+            subgraphs before clustering.
 
     Returns:
-        np.ndarray: Binary neighborhood matrix on Label Propagation.
+        np.ndarray: A binary neighborhood matrix on Label Propagation.
     """
+    # Create a subgraph with the shortest edges based on the rank fraction
+    subnetwork = _create_percentile_limited_subgraph(
+        network, fraction_shortest_edges=fraction_shortest_edges
+    )
     # Apply Label Propagation for community detection
-    communities = nx.algorithms.community.label_propagation.label_propagation_communities(network)
+    communities = nx.algorithms.community.label_propagation.label_propagation_communities(
+        subnetwork
+    )
+    # Get the list of nodes in the network
+    nodes = list(network.nodes())
+    node_index_map = {node: idx for idx, node in enumerate(nodes)}
     # Create a binary neighborhood matrix
-    num_nodes = network.number_of_nodes()
-    neighborhoods = np.zeros((num_nodes, num_nodes), dtype=int)
-    # Create a mapping from node to index in the matrix
-    node_index = {node: i for i, node in enumerate(network.nodes())}
-    # Assign neighborhoods based on community labels
+    num_nodes = len(nodes)
+    # Initialize neighborhoods with zeros and set self-self entries to 1
+    neighborhoods = np.eye(num_nodes, dtype=int)
+    # Assign neighborhoods based on community labels using the mapped indices
     for community in communities:
         for node_i in community:
-            idx_i = node_index[node_i]
             for node_j in community:
-                idx_j = node_index[node_j]
+                idx_i = node_index_map[node_i]
+                idx_j = node_index_map[node_j]
+                neighborhoods[idx_i, idx_j] = 1
+
+    return neighborhoods
+
+
+def calculate_leiden_neighborhoods(
+    network: nx.Graph,
+    resolution: float = 1.0,
+    fraction_shortest_edges: float = 1.0,
+    random_seed: int = 888,
+) -> np.ndarray:
+    """Calculate neighborhoods using the Leiden method.
+
+    Args:
+        network (nx.Graph): The network graph.
+        resolution (float, optional): Resolution parameter for the Leiden method. Defaults to 1.0.
+        fraction_shortest_edges (float, optional): Shortest edge rank fraction threshold for creating
+            subgraphs before clustering.
+        random_seed (int, optional): Random seed for reproducibility. Defaults to 888.
+
+    Returns:
+        np.ndarray: A binary neighborhood matrix where nodes in the same community have 1, and others have 0.
+    """
+    # Create a subgraph with the shortest edges based on the rank fraction
+    subnetwork = _create_percentile_limited_subgraph(
+        network, fraction_shortest_edges=fraction_shortest_edges
+    )
+    # Convert NetworkX graph to iGraph
+    igraph_network = ig.Graph.from_networkx(subnetwork)
+    # Apply Leiden algorithm using RBConfigurationVertexPartition, which supports resolution
+    partition = find_partition(
+        igraph_network,
+        partition_type=RBConfigurationVertexPartition,
+        resolution_parameter=resolution,
+        seed=random_seed,
+    )
+    # Get the list of nodes in the original NetworkX graph
+    nodes = list(network.nodes())
+    node_index_map = {node: idx for idx, node in enumerate(nodes)}
+    # Create a binary neighborhood matrix
+    num_nodes = len(nodes)
+    # Initialize neighborhoods with zeros and set self-self entries to 1
+    neighborhoods = np.eye(num_nodes, dtype=int)
+    # Assign neighborhoods based on community partitions using the mapped indices
+    for community in partition:
+        for node_i in community:
+            for node_j in community:
+                idx_i = node_index_map[igraph_network.vs[node_i]["_nx_name"]]
+                idx_j = node_index_map[igraph_network.vs[node_j]["_nx_name"]]
                 neighborhoods[idx_i, idx_j] = 1
 
     return neighborhoods
 
 
 def calculate_louvain_neighborhoods(
-    network: nx.Graph, resolution: float, random_seed: int = 888
+    network: nx.Graph,
+    resolution: float = 0.1,
+    fraction_shortest_edges: float = 1.0,
+    random_seed: int = 888,
 ) -> np.ndarray:
     """Calculate neighborhoods using the Louvain method.
 
     Args:
         network (nx.Graph): The network graph.
-        resolution (float): Resolution parameter for the Louvain method.
+        resolution (float, optional): Resolution parameter for the Louvain method. Defaults to 0.1.
+        fraction_shortest_edges (float, optional): Shortest edge rank fraction threshold for creating
+            subgraphs before clustering.
         random_seed (int, optional): Random seed for reproducibility. Defaults to 888.
 
     Returns:
-        np.ndarray: Binary neighborhood matrix on the Louvain method.
+        np.ndarray: A binary neighborhood matrix on the Louvain method.
     """
+    # Create a subgraph with the shortest edges based on the rank fraction
+    subnetwork = _create_percentile_limited_subgraph(
+        network, fraction_shortest_edges=fraction_shortest_edges
+    )
     # Apply Louvain method to partition the network
     partition = community_louvain.best_partition(
-        network, resolution=resolution, random_state=random_seed
+        subnetwork, resolution=resolution, random_state=random_seed
     )
+    # Get the list of nodes in the network and create a mapping to indices
+    nodes = list(network.nodes())
+    node_index_map = {node: idx for idx, node in enumerate(nodes)}
     # Create a binary neighborhood matrix
-    num_nodes = network.number_of_nodes()
-    neighborhoods = np.zeros((num_nodes, num_nodes), dtype=int)
-    # Create a mapping from node to index in the matrix
-    node_index = {node: i for i, node in enumerate(network.nodes())}
+    num_nodes = len(nodes)
+    # Initialize neighborhoods with zeros and set self-self entries to 1
+    neighborhoods = np.eye(num_nodes, dtype=int)
     # Group nodes by community
     community_groups = {}
     for node, community in partition.items():
         community_groups.setdefault(community, []).append(node)
 
-    # Assign neighborhoods based on community partitions
+    # Assign neighborhoods based on community partitions using the mapped indices
     for community, nodes in community_groups.items():
         for node_i in nodes:
-            idx_i = node_index[node_i]
             for node_j in nodes:
-                idx_j = node_index[node_j]
+                idx_i = node_index_map[node_i]
+                idx_j = node_index_map[node_j]
                 neighborhoods[idx_i, idx_j] = 1
 
     return neighborhoods
 
 
-def calculate_markov_clustering_neighborhoods(network: nx.Graph) -> np.ndarray:
-    """Apply Markov Clustering (MCL) to the network.
+def calculate_markov_clustering_neighborhoods(
+    network: nx.Graph, fraction_shortest_edges: float = 1.0
+) -> np.ndarray:
+    """Apply Markov Clustering (MCL) to the network and return a binary neighborhood matrix.
 
     Args:
         network (nx.Graph): The network graph.
+        fraction_shortest_edges (float, optional): Shortest edge rank fraction threshold for creating
+            subgraphs before clustering.
 
     Returns:
-        np.ndarray: Binary neighborhood matrix on Markov Clustering.
+        np.ndarray: A binary neighborhood matrix on Markov Clustering.
     """
-    # Convert the graph to an adjacency matrix
-    adjacency_matrix = nx.to_numpy_array(network)
-    # Run Markov Clustering (MCL)
-    result = mc.run_mcl(adjacency_matrix)  # MCL with default parameters
-    # Get clusters (communities) from MCL result
+    # Create a subgraph with the shortest edges based on the rank fraction
+    subnetwork = _create_percentile_limited_subgraph(
+        network, fraction_shortest_edges=fraction_shortest_edges
+    )
+    # Step 1: Convert the subnetwork to an adjacency matrix
+    subnetwork_nodes = list(subnetwork.nodes())
+    adjacency_matrix = nx.to_numpy_array(subnetwork, nodelist=subnetwork_nodes)
+    # Step 2: Run Markov Clustering (MCL) on the subnetwork's adjacency matrix
+    result = mc.run_mcl(adjacency_matrix)
     clusters = mc.get_clusters(result)
-    # Create a binary neighborhood matrix
-    num_nodes = network.number_of_nodes()
-    neighborhoods = np.zeros((num_nodes, num_nodes), dtype=int)
-    # Create a mapping from node to index in the matrix
-    node_index = {node: i for i, node in enumerate(network.nodes())}
-    # Assign neighborhoods based on MCL clusters
+    # Step 3: Prepare the original network nodes and indices
+    nodes = list(network.nodes())
+    node_index_map = {node: idx for idx, node in enumerate(nodes)}
+    num_nodes = len(nodes)
+    # Step 4: Initialize the neighborhood matrix for the original network
+    neighborhoods = np.eye(num_nodes, dtype=int)
+    # Step 5: Fill the neighborhoods matrix using the clusters from the subnetwork
     for cluster in clusters:
         for node_i in cluster:
-            idx_i = node_index[node_i]
             for node_j in cluster:
-                idx_j = node_index[node_j]
-                neighborhoods[idx_i, idx_j] = 1
+                # Map the indices back to the original network's node indices
+                original_node_i = subnetwork_nodes[node_i]
+                original_node_j = subnetwork_nodes[node_j]
+
+                if original_node_i in node_index_map and original_node_j in node_index_map:
+                    idx_i = node_index_map[original_node_i]
+                    idx_j = node_index_map[original_node_j]
+                    neighborhoods[idx_i, idx_j] = 1
 
     return neighborhoods
 
 
-def calculate_spinglass_neighborhoods(network: nx.Graph) -> np.ndarray:
-    """Apply Spin Glass Community Detection to the network.
+def calculate_spinglass_neighborhoods(
+    network: nx.Graph, fraction_shortest_edges: float = 1.0
+) -> np.ndarray:
+    """Apply Spinglass Community Detection to the network, handling disconnected components.
 
     Args:
         network (nx.Graph): The network graph.
+        fraction_shortest_edges (float, optional): Shortest edge rank fraction threshold for creating
+            subgraphs before clustering.
 
     Returns:
-        np.ndarray: Binary neighborhood matrix on Spin Glass communities.
+        np.ndarray: A binary neighborhood matrix based on Spinglass communities.
     """
-    # Apply Asynchronous Label Propagation (LPA)
-    communities = asyn_lpa_communities(network)
-    # Create a binary neighborhood matrix
-    num_nodes = network.number_of_nodes()
-    neighborhoods = np.zeros((num_nodes, num_nodes), dtype=int)
-    # Create a mapping from node to index in the matrix
-    node_index = {node: i for i, node in enumerate(network.nodes())}
-    # Assign neighborhoods based on community labels from LPA
-    for community in communities:
-        for node_i in community:
-            idx_i = node_index[node_i]
-            for node_j in community:
-                idx_j = node_index[node_j]
-                neighborhoods[idx_i, idx_j] = 1
+    # Create a subgraph with the shortest edges based on the rank fraction
+    subnetwork = _create_percentile_limited_subgraph(
+        network, fraction_shortest_edges=fraction_shortest_edges
+    )
+    # Step 1: Find connected components in the graph
+    components = list(nx.connected_components(subnetwork))
+    # Prepare to store community results
+    nodes = list(network.nodes())
+    node_index_map = {node: idx for idx, node in enumerate(nodes)}
+    num_nodes = len(nodes)
+    # Initialize neighborhoods with zeros and set self-self entries to 1
+    neighborhoods = np.eye(num_nodes, dtype=int)
+    # Step 2: Run Spinglass on each connected component
+    for component in components:
+        # Extract the subgraph corresponding to the current component
+        subgraph = network.subgraph(component)
+        # Convert the subgraph to an iGraph object
+        igraph_subgraph = ig.Graph.from_networkx(subgraph)
+        # Ensure the subgraph is connected before running Spinglass
+        if not igraph_subgraph.is_connected():
+            logger.error("Warning: Subgraph is not connected. Skipping...")
+            continue
+
+        # Apply Spinglass community detection
+        try:
+            communities = igraph_subgraph.community_spinglass()
+        except Exception as e:
+            logger.error(f"Error running Spinglass on component: {e}")
+            continue
+
+        # Step 3: Assign neighborhoods based on community labels
+        for community in communities:
+            for node_i in community:
+                for node_j in community:
+                    idx_i = node_index_map[igraph_subgraph.vs[node_i]["_nx_name"]]
+                    idx_j = node_index_map[igraph_subgraph.vs[node_j]["_nx_name"]]
+                    neighborhoods[idx_i, idx_j] = 1
 
     return neighborhoods
 
 
-def calculate_walktrap_neighborhoods(network: nx.Graph) -> np.ndarray:
+def calculate_walktrap_neighborhoods(
+    network: nx.Graph, fraction_shortest_edges: float = 1.0
+) -> np.ndarray:
     """Apply Walktrap Community Detection to the network.
 
     Args:
         network (nx.Graph): The network graph.
+        fraction_shortest_edges (float, optional): Shortest edge rank fraction threshold for creating
+            subgraphs before clustering.
 
     Returns:
-        np.ndarray: Binary neighborhood matrix on Walktrap communities.
+        np.ndarray: A binary neighborhood matrix on Walktrap communities.
     """
-    # Apply Asynchronous Label Propagation (LPA)
-    communities = asyn_lpa_communities(network)
+    # Create a subgraph with the shortest edges based on the rank fraction
+    subnetwork = _create_percentile_limited_subgraph(
+        network, fraction_shortest_edges=fraction_shortest_edges
+    )
+    # Convert NetworkX graph to iGraph
+    igraph_network = ig.Graph.from_networkx(subnetwork)
+    # Apply Walktrap community detection
+    communities = igraph_network.community_walktrap().as_clustering()
+    # Get the list of nodes in the original NetworkX graph
+    nodes = list(network.nodes())
+    node_index_map = {node: idx for idx, node in enumerate(nodes)}
     # Create a binary neighborhood matrix
-    num_nodes = network.number_of_nodes()
-    neighborhoods = np.zeros((num_nodes, num_nodes), dtype=int)
-    # Create a mapping from node to index in the matrix
-    node_index = {node: i for i, node in enumerate(network.nodes())}
-    # Assign neighborhoods based on community labels from LPA
+    num_nodes = len(nodes)
+    # Initialize neighborhoods with zeros and set self-self entries to 1
+    neighborhoods = np.eye(num_nodes, dtype=int)
+    # Assign neighborhoods based on community labels
     for community in communities:
         for node_i in community:
-            idx_i = node_index[node_i]
             for node_j in community:
-                idx_j = node_index[node_j]
+                idx_i = node_index_map[igraph_network.vs[node_i]["_nx_name"]]
+                idx_j = node_index_map[igraph_network.vs[node_j]["_nx_name"]]
                 neighborhoods[idx_i, idx_j] = 1
 
     return neighborhoods
+
+
+def _create_percentile_limited_subgraph(G: nx.Graph, fraction_shortest_edges: float) -> nx.Graph:
+    """Create a subgraph containing the shortest edges based on the specified rank fraction
+    of all edge lengths in the input graph.
+
+    Args:
+        G (nx.Graph): The input graph with 'length' attributes on edges.
+        fraction_shortest_edges (float): The rank fraction (between 0 and 1) to filter edges.
+
+    Returns:
+        nx.Graph: A subgraph with nodes and edges where the edges are within the shortest
+        specified rank fraction.
+    """
+    # Step 1: Extract edges with their lengths
+    edges_with_length = [(u, v, d) for u, v, d in G.edges(data=True) if "length" in d]
+    if not edges_with_length:
+        raise ValueError(
+            "No edge lengths found in the graph. Ensure edges have 'length' attributes."
+        )
+
+    # Step 2: Sort edges by length in ascending order
+    edges_with_length.sort(key=lambda x: x[2]["length"])
+    # Step 3: Calculate the cutoff index for the given rank fraction
+    cutoff_index = int(fraction_shortest_edges * len(edges_with_length))
+    if cutoff_index == 0:
+        raise ValueError("The rank fraction is too low, resulting in no edges being included.")
+
+    # Step 4: Create the subgraph by selecting only the shortest edges within the rank fraction
+    subgraph = nx.Graph()
+    subgraph.add_nodes_from(G.nodes(data=True))  # Retain all nodes from the original graph
+    subgraph.add_edges_from(edges_with_length[:cutoff_index])
+    # Step 5: Remove nodes with no edges
+    subgraph.remove_nodes_from(list(nx.isolates(subgraph)))
+    # Step 6: Check if the resulting subgraph has no edges and issue a warning
+    if subgraph.number_of_edges() == 0:
+        raise Warning("The resulting subgraph has no edges. Consider adjusting the rank fraction.")
+
+    return subgraph

risk/neighborhoods/domains.py

@@ -5,13 +5,13 @@ risk/neighborhoods/domains
 
 from contextlib import suppress
 from itertools import product
-from tqdm import tqdm
 from typing import Tuple
 
 import numpy as np
 import pandas as pd
 from scipy.cluster.hierarchy import linkage, fcluster
 from sklearn.metrics import silhouette_score
+from tqdm import tqdm
 
 from risk.annotations import get_weighted_description
 from risk.constants import GROUP_LINKAGE_METHODS, GROUP_DISTANCE_METRICS
@@ -20,17 +20,17 @@ from risk.log import logger
 
 def define_domains(
     top_annotations: pd.DataFrame,
-    significant_neighborhoods_enrichment: np.ndarray,
+    significant_neighborhoods_significance: np.ndarray,
     linkage_criterion: str,
     linkage_method: str,
     linkage_metric: str,
 ) -> pd.DataFrame:
-    """Define domains and assign nodes to these domains based on their enrichment scores and clustering,
+    """Define domains and assign nodes to these domains based on their significance scores and clustering,
     handling errors by assigning unique domains when clustering fails.
 
     Args:
         top_annotations (pd.DataFrame): DataFrame of top annotations data for the network nodes.
-        significant_neighborhoods_enrichment (np.ndarray): The binary enrichment matrix below alpha.
+        significant_neighborhoods_significance (np.ndarray): The binary significance matrix below alpha.
         linkage_criterion (str): The clustering criterion for defining groups.
         linkage_method (str): The linkage method for clustering.
         linkage_metric (str): The linkage metric for clustering.
@@ -39,8 +39,14 @@ def define_domains(
         pd.DataFrame: DataFrame with the primary domain for each node.
     """
     try:
+        if linkage_criterion == "off":
+            raise ValueError("Clustering is turned off.")
+
         # Transpose the matrix to cluster annotations
-        m = significant_neighborhoods_enrichment[:, top_annotations["significant_annotations"]].T
+        m = significant_neighborhoods_significance[:, top_annotations["significant_annotations"]].T
+        # Safeguard the matrix by replacing NaN, Inf, and -Inf values
+        m = _safeguard_matrix(m)
+        # Optimize silhouette score across different linkage methods and distance metrics
         best_linkage, best_metric, best_threshold = _optimize_silhouette_across_linkage_and_metrics(
             m, linkage_criterion, linkage_method, linkage_metric
         )
@@ -59,19 +65,24 @@ def define_domains(
     except ValueError:
         # If a ValueError is encountered, handle it by assigning unique domains
         n_rows = len(top_annotations)
-        logger.error(
-            f"Error encountered. Skipping clustering and assigning {n_rows} unique domains."
-        )
+        if linkage_criterion == "off":
+            logger.warning(
+                f"Clustering is turned off. Skipping clustering and assigning {n_rows} unique domains."
+            )
+        else:
+            logger.error(
+                f"Error encountered. Skipping clustering and assigning {n_rows} unique domains."
+            )
         top_annotations["domain"] = range(1, n_rows + 1)  # Assign unique domains
 
     # Create DataFrames to store domain information
-    node_to_enrichment = pd.DataFrame(
-        data=significant_neighborhoods_enrichment,
+    node_to_significance = pd.DataFrame(
+        data=significant_neighborhoods_significance,
         columns=[top_annotations.index.values, top_annotations["domain"]],
     )
-    node_to_domain = node_to_enrichment.groupby(level="domain", axis=1).sum()
+    node_to_domain = node_to_significance.T.groupby(level="domain").sum().T
 
-    # Find the maximum enrichment score for each node
+    # Find the maximum significance score for each node
     t_max = node_to_domain.loc[:, 1:].max(axis=1)
     t_idxmax = node_to_domain.loc[:, 1:].idxmax(axis=1)
     t_idxmax[t_max == 0] = 0
@@ -86,13 +97,13 @@ def define_domains(
     return node_to_domain
 
 
-def trim_domains_and_top_annotations(
+def trim_domains(
     domains: pd.DataFrame,
     top_annotations: pd.DataFrame,
     min_cluster_size: int = 5,
     max_cluster_size: int = 1000,
 ) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
-    """Trim domains and top annotations that do not meet size criteria and find outliers.
+    """Trim domains that do not meet size criteria and find outliers.
 
     Args:
         domains (pd.DataFrame): DataFrame of domain data for the network nodes.
@@ -101,8 +112,7 @@ def trim_domains_and_top_annotations(
         max_cluster_size (int, optional): Maximum size of a cluster to be retained. Defaults to 1000.
 
     Returns:
-        Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: A tuple containing:
-            - Trimmed annotations (pd.DataFrame)
+        Tuple[pd.DataFrame, pd.DataFrame]:
             - Trimmed domains (pd.DataFrame)
             - A DataFrame with domain labels (pd.DataFrame)
     """
@@ -116,30 +126,30 @@ def trim_domains_and_top_annotations(
     invalid_domain_id = 888888
     invalid_domain_ids = {0, invalid_domain_id}
     # Mark domains to be removed
-    top_annotations["domain"].replace(to_remove, invalid_domain_id, inplace=True)
+    top_annotations["domain"] = top_annotations["domain"].replace(to_remove, invalid_domain_id)
     domains.loc[domains["primary_domain"].isin(to_remove), ["primary_domain"]] = invalid_domain_id
 
-    # Normalize "num enriched neighborhoods" by percentile for each domain and scale to 0-10
+    # Normalize "num significant neighborhoods" by percentile for each domain and scale to 0-10
     top_annotations["normalized_value"] = top_annotations.groupby("domain")[
-        "significant_neighborhood_enrichment_sums"
+        "significant_neighborhood_significance_sums"
     ].transform(lambda x: (x.rank(pct=True) * 10).apply(np.ceil).astype(int))
-    # Modify the lambda function to pass both full_terms and significant_enrichment_score
+    # Modify the lambda function to pass both full_terms and significant_significance_score
     top_annotations["combined_terms"] = top_annotations.apply(
         lambda row: " ".join([str(row["full_terms"])] * row["normalized_value"]), axis=1
     )
 
-    # Perform the groupby operation while retaining the other columns and adding the weighting with enrichment scores
+    # Perform the groupby operation while retaining the other columns and adding the weighting with significance scores
     domain_labels = (
         top_annotations.groupby("domain")
         .agg(
             full_terms=("full_terms", lambda x: list(x)),
-            enrichment_scores=("significant_enrichment_score", lambda x: list(x)),
+            significance_scores=("significant_significance_score", lambda x: list(x)),
         )
         .reset_index()
     )
     domain_labels["combined_terms"] = domain_labels.apply(
         lambda row: get_weighted_description(
-            pd.Series(row["full_terms"]), pd.Series(row["enrichment_scores"])
+            pd.Series(row["full_terms"]), pd.Series(row["significance_scores"])
         ),
         axis=1,
     )
@@ -150,19 +160,39 @@ def trim_domains_and_top_annotations(
             "domain": "id",
             "combined_terms": "normalized_description",
             "full_terms": "full_descriptions",
-            "enrichment_scores": "enrichment_scores",
+            "significance_scores": "significance_scores",
         }
     ).set_index("id")
 
     # Remove invalid domains
-    valid_annotations = top_annotations[~top_annotations["domain"].isin(invalid_domain_ids)].drop(
-        columns=["normalized_value"]
-    )
     valid_domains = domains[~domains["primary_domain"].isin(invalid_domain_ids)]
     valid_trimmed_domains_matrix = trimmed_domains_matrix[
         ~trimmed_domains_matrix.index.isin(invalid_domain_ids)
     ]
-    return valid_annotations, valid_domains, valid_trimmed_domains_matrix
+    return valid_domains, valid_trimmed_domains_matrix
+
+
+def _safeguard_matrix(matrix: np.ndarray) -> np.ndarray:
+    """Safeguard the matrix by replacing NaN, Inf, and -Inf values.
+
+    Args:
+        matrix (np.ndarray): Data matrix.
+
+    Returns:
+        np.ndarray: Safeguarded data matrix.
+    """
+    # Replace NaN with column mean
+    nan_replacement = np.nanmean(matrix, axis=0)
+    matrix = np.where(np.isnan(matrix), nan_replacement, matrix)
+    # Replace Inf/-Inf with maximum/minimum finite values
+    finite_max = np.nanmax(matrix[np.isfinite(matrix)])
+    finite_min = np.nanmin(matrix[np.isfinite(matrix)])
+    matrix = np.where(np.isposinf(matrix), finite_max, matrix)
+    matrix = np.where(np.isneginf(matrix), finite_min, matrix)
+    # Ensure rows have non-zero variance (optional step)
+    row_variance = np.var(matrix, axis=1)
+    matrix = matrix[row_variance > 0]
+    return matrix
 
 
 def _optimize_silhouette_across_linkage_and_metrics(
@@ -177,7 +207,7 @@ def _optimize_silhouette_across_linkage_and_metrics(
         linkage_metric (str): Linkage metric for clustering.
 
     Returns:
-        Tuple[str, str, float]: A tuple containing:
+        Tuple[str, str, float]:
             - Best linkage method (str)
             - Best linkage metric (str)
             - Best threshold (float)
@@ -198,7 +228,8 @@ def _optimize_silhouette_across_linkage_and_metrics(
         total=total_combinations,
         bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]",
     ):
-        with suppress(Exception):
+        # Some linkage methods and metrics may not work with certain data
+        with suppress(ValueError):
             Z = linkage(m, method=method, metric=metric)
             threshold, score = _find_best_silhouette_score(Z, m, metric, linkage_criterion)
             if score > best_overall_score:
@@ -231,7 +262,7 @@ def _find_best_silhouette_score(
         resolution (float, optional): Desired resolution for the best threshold. Defaults to 0.001.
 
     Returns:
-        Tuple[float, float]: A tuple containing:
+        Tuple[float, float]:
             - Best threshold (float): The threshold that yields the best silhouette score.
             - Best silhouette score (float): The highest silhouette score achieved.
     """