risk-network 0.0.8b18__py3-none-any.whl → 0.0.8b20__py3-none-any.whl

risk/__init__.py

@@ -7,4 +7,4 @@ RISK: RISK Infers Spatial Kinships
 
 from risk.risk import RISK
 
-__version__ = "0.0.8-beta.18"
+__version__ = "0.0.8-beta.20"

risk/annotations/annotations.py

@@ -36,10 +36,11 @@ def load_annotations(network: nx.Graph, annotations_input: Dict[str, Any]) -> Di
     """Convert annotations input to a DataFrame and reindex based on the network's node labels.
 
     Args:
-        annotations_input (dict): A dictionary with annotations.
+        network (nx.Graph): The network graph.
+        annotations_input (Dict[str, Any]): A dictionary with annotations.
 
     Returns:
-        dict: A dictionary containing ordered nodes, ordered annotations, and the binary annotations matrix.
+        Dict[str, Any]: A dictionary containing ordered nodes, ordered annotations, and the binary annotations matrix.
     """
     # Flatten the dictionary to a list of tuples for easier DataFrame creation
     flattened_annotations = [
@@ -255,7 +256,7 @@ def _generate_coherent_description(words: List[str]) -> str:
     If there is only one unique entry, return it directly.
 
     Args:
-        words (list): A list of words or numerical string values.
+        words (List): A list of words or numerical string values.
 
     Returns:
         str: A coherent description formed by arranging the words in a logical sequence.

risk/annotations/io.py

@@ -33,7 +33,7 @@ class AnnotationsIO:
             filepath (str): Path to the JSON annotations file.
 
         Returns:
-            dict: A dictionary containing ordered nodes, ordered annotations, and the annotations matrix.
+            Dict[str, Any]: A dictionary containing ordered nodes, ordered annotations, and the annotations matrix.
         """
         filetype = "JSON"
         # Log the loading of the JSON file
@@ -158,10 +158,10 @@ class AnnotationsIO:
 
         Args:
             network (NetworkX graph): The network to which the annotations are related.
-            content (dict): The annotations dictionary to load.
+            content (Dict[str, Any]): The annotations dictionary to load.
 
         Returns:
-            dict: A dictionary containing ordered nodes, ordered annotations, and the annotations matrix.
+            Dict[str, Any]: A dictionary containing ordered nodes, ordered annotations, and the annotations matrix.
         """
         # Ensure the input content is a dictionary
         if not isinstance(content, dict):

risk/log/params.py

@@ -159,7 +159,7 @@ class Params:
         """Load and process various parameters, converting any np.ndarray values to lists.
 
         Returns:
-            dict: A dictionary containing the processed parameters.
+            Dict[str, Any]: A dictionary containing the processed parameters.
         """
         log_header("Loading parameters")
         return _convert_ndarray_to_list(
@@ -174,14 +174,14 @@ class Params:
         )
 
 
-def _convert_ndarray_to_list(d: Any) -> Any:
+def _convert_ndarray_to_list(d: Dict[str, Any]) -> Dict[str, Any]:
     """Recursively convert all np.ndarray values in the dictionary to lists.
 
     Args:
-        d (dict): The dictionary to process.
+        d (Dict[str, Any]): The dictionary to process.
 
     Returns:
-        dict: The processed dictionary with np.ndarray values converted to lists.
+        Dict[str, Any]: The processed dictionary with np.ndarray values converted to lists.
     """
     if isinstance(d, dict):
         # Recursively process each value in the dictionary
@@ -193,5 +193,5 @@ def _convert_ndarray_to_list(d: Any) -> Any:
         # Convert numpy arrays to lists
         return d.tolist()
     else:
-        # Return the value unchanged if it's not a dict, list, or ndarray
+        # Return the value unchanged if it's not a dict, List, or ndarray
         return d

risk/neighborhoods/community.py

@@ -21,15 +21,20 @@ def calculate_greedy_modularity_neighborhoods(network: nx.Graph) -> np.ndarray:
     """
     # Detect communities using the Greedy Modularity method
     communities = greedy_modularity_communities(network)
-    # Create a mapping from node to community
-    community_dict = {node: idx for idx, community in enumerate(communities) for node in community}
     # Create a binary neighborhood matrix
-    neighborhoods = np.zeros((network.number_of_nodes(), network.number_of_nodes()), dtype=int)
+    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())}
-    for node_i, community_i in community_dict.items():
-        for node_j, community_j in community_dict.items():
-            if community_i == community_j:
-                neighborhoods[node_index[node_i], node_index[node_j]] = 1
+    # 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]
+                # Set them as neighbors (1) in the binary matrix
+                neighborhoods[idx_i, idx_j] = 1
 
     return neighborhoods
 
@@ -43,22 +48,20 @@ def calculate_label_propagation_neighborhoods(network: nx.Graph) -> np.ndarray:
     Returns:
         np.ndarray: Binary neighborhood matrix on Label Propagation.
     """
-    # Apply Label Propagation
+    # Apply Label Propagation for community detection
     communities = nx.algorithms.community.label_propagation.label_propagation_communities(network)
-    # Create a mapping from node to community
-    community_dict = {}
-    for community_id, community in enumerate(communities):
-        for node in community:
-            community_dict[node] = community_id
-
     # 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
-    for node_i, community_i in community_dict.items():
-        for node_j, community_j in community_dict.items():
-            if community_i == community_j:
-                neighborhoods[node_i, node_j] = 1
+    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
 
     return neighborhoods
 
@@ -81,12 +84,22 @@ def calculate_louvain_neighborhoods(
         network, resolution=resolution, random_state=random_seed
     )
     # Create a binary neighborhood matrix
-    neighborhoods = np.zeros((network.number_of_nodes(), network.number_of_nodes()), dtype=int)
+    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())}
+    # Group nodes by community
+    community_groups = {}
+    for node, community in partition.items():
+        community_groups.setdefault(community, []).append(node)
+
     # Assign neighborhoods based on community partitions
-    for node_i, community_i in partition.items():
-        for node_j, community_j in partition.items():
-            if community_i == community_j:
-                neighborhoods[node_i, node_j] = 1
+    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]
+                neighborhoods[idx_i, idx_j] = 1
 
     return neighborhoods
 
@@ -102,24 +115,22 @@ def calculate_markov_clustering_neighborhoods(network: nx.Graph) -> np.ndarray:
     """
     # Convert the graph to an adjacency matrix
     adjacency_matrix = nx.to_numpy_array(network)
-    # Run Markov Clustering
-    result = mc.run_mcl(adjacency_matrix)  # Run MCL with default parameters
-    # Get clusters
+    # Run Markov Clustering (MCL)
+    result = mc.run_mcl(adjacency_matrix)  # MCL with default parameters
+    # Get clusters (communities) from MCL result
     clusters = mc.get_clusters(result)
-    # Create a community label for each node
-    community_dict = {}
-    for community_id, community in enumerate(clusters):
-        for node in community:
-            community_dict[node] = community_id
-
     # Create a binary neighborhood matrix
     num_nodes = network.number_of_nodes()
     neighborhoods = np.zeros((num_nodes, num_nodes), dtype=int)
-    # Assign neighborhoods based on community labels
-    for node_i, community_i in community_dict.items():
-        for node_j, community_j in community_dict.items():
-            if community_i == community_j:
-                neighborhoods[node_i, node_j] = 1
+    # 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
+    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
 
     return neighborhoods
 
@@ -133,22 +144,20 @@ def calculate_spinglass_neighborhoods(network: nx.Graph) -> np.ndarray:
     Returns:
         np.ndarray: Binary neighborhood matrix on Spin Glass communities.
     """
-    # Use the asynchronous label propagation algorithm as a proxy for Spin Glass
+    # Apply Asynchronous Label Propagation (LPA)
     communities = asyn_lpa_communities(network)
-    # Create a community label for each node
-    community_dict = {}
-    for community_id, community in enumerate(communities):
-        for node in community:
-            community_dict[node] = community_id
-
     # Create a binary neighborhood matrix
     num_nodes = network.number_of_nodes()
     neighborhoods = np.zeros((num_nodes, num_nodes), dtype=int)
-    # Assign neighborhoods based on community labels
-    for node_i, community_i in community_dict.items():
-        for node_j, community_j in community_dict.items():
-            if community_i == community_j:
-                neighborhoods[node_i, node_j] = 1
+    # 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
 
     return neighborhoods
 
@@ -162,21 +171,19 @@ def calculate_walktrap_neighborhoods(network: nx.Graph) -> np.ndarray:
     Returns:
         np.ndarray: Binary neighborhood matrix on Walktrap communities.
     """
-    # Use the asynchronous label propagation algorithm as a proxy for Walktrap
+    # Apply Asynchronous Label Propagation (LPA)
     communities = asyn_lpa_communities(network)
-    # Create a community label for each node
-    community_dict = {}
-    for community_id, community in enumerate(communities):
-        for node in community:
-            community_dict[node] = community_id
-
     # Create a binary neighborhood matrix
     num_nodes = network.number_of_nodes()
     neighborhoods = np.zeros((num_nodes, num_nodes), dtype=int)
-    # Assign neighborhoods based on community labels
-    for node_i, community_i in community_dict.items():
-        for node_j, community_j in community_dict.items():
-            if community_i == community_j:
-                neighborhoods[node_i, node_j] = 1
+    # 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
 
     return neighborhoods

risk/neighborhoods/domains.py

@@ -76,6 +76,10 @@ def define_domains(
     t_idxmax = node_to_domain.loc[:, 1:].idxmax(axis=1)
     t_idxmax[t_max == 0] = 0
 
+    # Assign all domains where the score is greater than 0
+    node_to_domain["all domains"] = node_to_domain.loc[:, 1:].apply(
+        lambda row: list(row[row > 0].index), axis=1
+    )
     # Assign primary domain
     node_to_domain["primary domain"] = t_idxmax
 
@@ -97,7 +101,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:
+        Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]: A tuple containing:
             - Trimmed annotations (pd.DataFrame)
             - Trimmed domains (pd.DataFrame)
             - A DataFrame with domain labels (pd.DataFrame)
@@ -154,7 +158,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]: A tuple containing:
             - Best linkage method (str)
             - Best linkage metric (str)
             - Best threshold (float)
@@ -208,7 +212,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]: A tuple containing:
             - Best threshold (float): The threshold that yields the best silhouette score.
             - Best silhouette score (float): The highest silhouette score achieved.
     """

risk/neighborhoods/neighborhoods.py

@@ -5,7 +5,7 @@ risk/neighborhoods/neighborhoods
 
 import random
 import warnings
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Tuple, Union
 
 import networkx as nx
 import numpy as np
@@ -28,50 +28,82 @@ warnings.filterwarnings(action="ignore", category=DataConversionWarning)
 
 def get_network_neighborhoods(
     network: nx.Graph,
-    distance_metric: str = "louvain",
-    edge_length_threshold: float = 1.0,
+    distance_metric: Union[str, List, Tuple, np.ndarray] = "louvain",
+    edge_length_threshold: Union[float, List, Tuple, np.ndarray] = 1.0,
     louvain_resolution: float = 1.0,
     random_seed: int = 888,
 ) -> np.ndarray:
-    """Calculate the neighborhoods for each node in the network based on the specified distance metric.
+    """Calculate the combined neighborhoods for each node based on the specified community detection algorithm(s).
 
     Args:
         network (nx.Graph): The network graph.
-        distance_metric (str): The distance metric to use ('greedy_modularity', 'louvain', 'label_propagation',
-            'markov_clustering', 'walktrap', 'spinglass').
-        edge_length_threshold (float): The edge length threshold for the neighborhoods.
+        distance_metric (str, List, Tuple, or np.ndarray, optional): The distance metric(s) to use. Can be a string for one
+            metric or a list/tuple/ndarray of metrics ('greedy_modularity', 'louvain', 'label_propagation',
+            'markov_clustering', 'walktrap', 'spinglass'). Defaults to 'louvain'.
+        edge_length_threshold (float, List, Tuple, or np.ndarray, optional): Edge length threshold(s) for creating subgraphs.
+            Can be a single float for one threshold or a list/tuple of floats corresponding to multiple thresholds.
+            Defaults to 1.0.
         louvain_resolution (float, optional): Resolution parameter for the Louvain method. Defaults to 1.0.
         random_seed (int, optional): Random seed for methods requiring random initialization. Defaults to 888.
 
     Returns:
-        np.ndarray: Neighborhood matrix calculated based on the selected distance metric.
+        np.ndarray: Summed neighborhood matrix from all selected algorithms.
     """
-    # Set random seed for reproducibility in all methods besides Louvain, which requires a separate seed
+    # Set random seed for reproducibility
     random.seed(random_seed)
     np.random.seed(random_seed)
 
-    # Create a subgraph based on the edge length percentile threshold
-    network = _create_percentile_limited_subgraph(
-        network, edge_length_percentile=edge_length_threshold
-    )
+    # Ensure distance_metric is a list/tuple for multi-algorithm handling
+    if isinstance(distance_metric, (str, np.ndarray)):
+        distance_metric = [distance_metric]
+    # Ensure edge_length_threshold is a list/tuple for multi-threshold handling
+    if isinstance(edge_length_threshold, (float, int)):
+        edge_length_threshold = [edge_length_threshold] * len(distance_metric)
+    # Check that the number of distance metrics matches the number of edge length thresholds
+    if len(distance_metric) != len(edge_length_threshold):
+        raise ValueError(
+            "The number of distance metrics must match the number of edge length thresholds."
+        )
 
-    if distance_metric == "louvain":
-        return calculate_louvain_neighborhoods(network, louvain_resolution, random_seed=random_seed)
-    if distance_metric == "greedy_modularity":
-        return calculate_greedy_modularity_neighborhoods(network)
-    if distance_metric == "label_propagation":
-        return calculate_label_propagation_neighborhoods(network)
-    if distance_metric == "markov_clustering":
-        return calculate_markov_clustering_neighborhoods(network)
-    if distance_metric == "walktrap":
-        return calculate_walktrap_neighborhoods(network)
-    if distance_metric == "spinglass":
-        return calculate_spinglass_neighborhoods(network)
-
-    raise ValueError(
-        "Incorrect distance metric specified. Please choose from 'greedy_modularity', 'louvain',"
-        "'label_propagation', 'markov_clustering', 'walktrap', 'spinglass'."
-    )
+    # Initialize combined neighborhood matrix
+    num_nodes = network.number_of_nodes()
+    combined_neighborhoods = np.zeros((num_nodes, num_nodes), dtype=int)
+
+    # Loop through each distance metric and corresponding edge length threshold
+    for metric, threshold in zip(distance_metric, edge_length_threshold):
+        # Create a subgraph based on the specific edge length threshold for this algorithm
+        subgraph = _create_percentile_limited_subgraph(network, edge_length_percentile=threshold)
+        # Call the appropriate neighborhood function based on the metric
+        if metric == "louvain":
+            neighborhoods = calculate_louvain_neighborhoods(
+                subgraph, louvain_resolution, random_seed=random_seed
+            )
+        elif metric == "greedy_modularity":
+            neighborhoods = calculate_greedy_modularity_neighborhoods(subgraph)
+        elif metric == "label_propagation":
+            neighborhoods = calculate_label_propagation_neighborhoods(subgraph)
+        elif metric == "markov_clustering":
+            neighborhoods = calculate_markov_clustering_neighborhoods(subgraph)
+        elif metric == "walktrap":
+            neighborhoods = calculate_walktrap_neighborhoods(subgraph)
+        elif metric == "spinglass":
+            neighborhoods = calculate_spinglass_neighborhoods(subgraph)
+        else:
+            raise ValueError(
+                "Incorrect distance metric specified. Please choose from 'greedy_modularity', 'louvain',"
+                "'label_propagation', 'markov_clustering', 'walktrap', 'spinglass'."
+            )
+
+        # Sum the neighborhood matrices
+        combined_neighborhoods += neighborhoods
+
+    # Ensure that the maximum value in each row is set to 1
+    # This ensures that for each row, only the strongest relationship (the maximum value) is retained,
+    # while all other values are reset to 0. This transformation simplifies the neighborhood matrix by
+    # focusing on the most significant connection per row.
+    combined_neighborhoods = _set_max_to_one(combined_neighborhoods)
+
+    return combined_neighborhoods
 
 
 def _create_percentile_limited_subgraph(G: nx.Graph, edge_length_percentile: float) -> nx.Graph:
@@ -110,6 +142,25 @@ def _create_percentile_limited_subgraph(G: nx.Graph, edge_length_percentile: flo
     return subgraph
 
 
+def _set_max_to_one(matrix: np.ndarray) -> np.ndarray:
+    """For each row in the input matrix, set the maximum value(s) to 1 and all other values to 0.
+
+    Args:
+        matrix (np.ndarray): A 2D numpy array representing the neighborhood matrix.
+
+    Returns:
+        np.ndarray: The modified matrix where only the maximum value(s) in each row is set to 1, and others are set to 0.
+    """
+    # Find the maximum value in each row (column-wise max operation)
+    max_values = np.max(matrix, axis=1, keepdims=True)
+    # Create a boolean mask where elements are True if they are the max value in their row
+    max_mask = matrix == max_values
+    # Set all elements to 0, and then set the maximum value positions to 1
+    matrix[:] = 0  # Set everything to 0
+    matrix[max_mask] = 1  # Set only the max values to 1
+    return matrix
+
+
 def process_neighborhoods(
     network: nx.Graph,
     neighborhoods: Dict[str, Any],
@@ -120,12 +171,12 @@ def process_neighborhoods(
 
     Args:
         network (nx.Graph): The network data structure used for imputing and pruning neighbors.
-        neighborhoods (dict): Dictionary containing 'enrichment_matrix', 'binary_enrichment_matrix', and 'significant_enrichment_matrix'.
+        neighborhoods (Dict[str, Any]): Dictionary containing 'enrichment_matrix', 'binary_enrichment_matrix', and 'significant_enrichment_matrix'.
         impute_depth (int, optional): Depth for imputing neighbors. Defaults to 0.
         prune_threshold (float, optional): Distance threshold for pruning neighbors. Defaults to 0.0.
 
     Returns:
-        dict: Processed neighborhoods data, including the updated matrices and enrichment counts.
+        Dict[str, Any]: Processed neighborhoods data, including the updated matrices and enrichment counts.
     """
     enrichment_matrix = neighborhoods["enrichment_matrix"]
     binary_enrichment_matrix = neighborhoods["binary_enrichment_matrix"]
@@ -408,7 +459,7 @@ def _calculate_threshold(median_distances: List, distance_threshold: float) -> f
     """Calculate the distance threshold based on the given median distances and a percentile threshold.
 
     Args:
-        median_distances (list): An array of median distances.
+        median_distances (List): An array of median distances.
         distance_threshold (float): A percentile threshold (0 to 1) used to determine the distance cutoff.
 
     Returns:

risk/network/geometry.py

@@ -68,6 +68,7 @@ def assign_edge_lengths(
             v_coords = np.append(v_coords, G_depth.nodes[v].get("z", 0))
 
         distance = compute_distance(u_coords, v_coords, is_sphere=compute_sphere)
+        # Assign edge lengths to the original graph
         if include_edge_weight:
             # Square root of the normalized weight is used to minimize the effect of large weights
             G.edges[u, v]["length"] = distance / np.sqrt(G.edges[u, v]["normalized_weight"] + 1e-6)

risk/network/graph.py

@@ -36,7 +36,7 @@ class NetworkGraph:
             top_annotations (pd.DataFrame): DataFrame containing annotations data for the network nodes.
             domains (pd.DataFrame): DataFrame containing domain data for the network nodes.
             trimmed_domains (pd.DataFrame): DataFrame containing trimmed domain data for the network nodes.
-            node_label_to_node_id_map (dict): A dictionary mapping node labels to their corresponding IDs.
+            node_label_to_node_id_map (Dict[str, Any]): A dictionary mapping node labels to their corresponding IDs.
             node_enrichment_sums (np.ndarray): Array containing the enrichment sums for the nodes.
         """
         self.top_annotations = top_annotations
@@ -46,6 +46,9 @@ class NetworkGraph:
             trimmed_domains
         )
         self.node_enrichment_sums = node_enrichment_sums
+        self.node_id_to_domain_ids_and_enrichments_map = (
+            self._create_node_id_to_domain_ids_and_enrichments(domains)
+        )
         self.node_id_to_node_label_map = {v: k for k, v in node_label_to_node_id_map.items()}
         self.node_label_to_enrichment_map = dict(
             zip(node_label_to_node_id_map.keys(), node_enrichment_sums)
@@ -57,14 +60,14 @@ class NetworkGraph:
         self.network = _unfold_sphere_to_plane(network)
         self.node_coordinates = _extract_node_coordinates(self.network)
 
-    def _create_domain_id_to_node_ids_map(self, domains: pd.DataFrame) -> Dict[str, Any]:
+    def _create_domain_id_to_node_ids_map(self, domains: pd.DataFrame) -> Dict[int, Any]:
         """Create a mapping from domains to the list of node IDs belonging to each domain.
 
         Args:
             domains (pd.DataFrame): DataFrame containing domain information, including the 'primary domain' for each node.
 
         Returns:
-            dict: A dictionary where keys are domain IDs and values are lists of node IDs belonging to each domain.
+            Dict[int, Any]: A dictionary where keys are domain IDs and values are lists of node IDs belonging to each domain.
         """
         cleaned_domains_matrix = domains.reset_index()[["index", "primary domain"]]
         node_to_domains_map = cleaned_domains_matrix.set_index("index")["primary domain"].to_dict()
@@ -76,14 +79,14 @@ class NetworkGraph:
 
     def _create_domain_id_to_domain_terms_map(
         self, trimmed_domains: pd.DataFrame
-    ) -> Dict[str, Any]:
+    ) -> Dict[int, Any]:
         """Create a mapping from domain IDs to their corresponding terms.
 
         Args:
             trimmed_domains (pd.DataFrame): DataFrame containing domain IDs and their corresponding labels.
 
         Returns:
-            dict: A dictionary mapping domain IDs to their corresponding terms.
+            Dict[int, Any]: A dictionary mapping domain IDs to their corresponding terms.
         """
         return dict(
             zip(
@@ -92,11 +95,45 @@ class NetworkGraph:
             )
         )
 
+    def _create_node_id_to_domain_ids_and_enrichments(
+        self, domains: pd.DataFrame
+    ) -> Dict[int, Dict]:
+        """Creates a dictionary mapping each node ID to its corresponding domain IDs and enrichment values.
+
+        Args:
+            domains (pd.DataFrame): A DataFrame containing domain information for each node. Assumes the last
+                two columns are 'all domains' and 'primary domain', which are excluded from processing.
+
+        Returns:
+            Dict[int, Dict]: A dictionary where the key is the node ID (index of the DataFrame), and the value is another dictionary
+                with 'domain' (a list of domain IDs with non-zero enrichment) and 'enrichment'
+                (a dict of domain IDs and their corresponding enrichment values).
+        """
+        # Initialize an empty dictionary to store the result
+        node_id_to_domain_ids_and_enrichments = {}
+        # Get the list of domain columns (excluding 'all domains' and 'primary domain')
+        domain_columns = domains.columns[
+            :-2
+        ]  # The last two columns are 'all domains' and 'primary domain'
+        # Iterate over each row in the dataframe
+        for idx, row in domains.iterrows():
+            # Get the domains (column names) where the enrichment score is greater than 0
+            all_domains = domain_columns[row[domain_columns] > 0].tolist()
+            # Get the enrichment values for those domains
+            enrichment_values = row[all_domains].to_dict()
+            # Store the result in the dictionary with index as the key
+            node_id_to_domain_ids_and_enrichments[idx] = {
+                "domains": all_domains,  # The column names where enrichment > 0
+                "enrichments": enrichment_values,  # The actual enrichment values for those columns
+            }
+
+        return node_id_to_domain_ids_and_enrichments
+
     def _create_domain_id_to_node_labels_map(self) -> Dict[int, List[str]]:
         """Create a map from domain IDs to node labels.
 
         Returns:
-            dict: A dictionary mapping domain IDs to the corresponding node labels.
+            Dict[int, List[str]]: A dictionary mapping domain IDs to the corresponding node labels.
         """
         domain_id_to_label_map = {}
         for domain_id, node_ids in self.domain_id_to_node_ids_map.items():

risk/network/io.py

@@ -491,7 +491,7 @@ class NetworkIO:
             if "x" not in attrs or "y" not in attrs:
                 if (
                     "pos" in attrs
-                    and isinstance(attrs["pos"], (list, tuple, np.ndarray))
+                    and isinstance(attrs["pos"], (List, Tuple, np.ndarray))
                     and len(attrs["pos"]) >= 2
                 ):
                     attrs["x"], attrs["y"] = attrs["pos"][