risk-network 0.0.3b4__py3-none-any.whl → 0.0.4__py3-none-any.whl

risk/network/plot.py

@@ -27,7 +27,7 @@ class NetworkPlotter:
 
     def __init__(
         self,
-        network_graph: NetworkGraph,
+        graph: NetworkGraph,
         figsize: tuple = (10, 10),
         background_color: str = "white",
         plot_outline: bool = True,
@@ -37,22 +37,22 @@ class NetworkPlotter:
         """Initialize the NetworkPlotter with a NetworkGraph object and plotting parameters.
 
         Args:
-            network_graph (NetworkGraph): The network data and attributes to be visualized.
+            graph (NetworkGraph): The network data and attributes to be visualized.
             figsize (tuple, optional): Size of the figure in inches (width, height). Defaults to (10, 10).
             background_color (str, optional): Background color of the plot. Defaults to "white".
             plot_outline (bool, optional): Whether to plot the network perimeter circle. Defaults to True.
             outline_color (str, optional): Color of the network perimeter circle. Defaults to "black".
             outline_scale (float, optional): Outline scaling factor for the perimeter diameter. Defaults to 1.0.
         """
-        self.network_graph = network_graph
+        self.graph = graph
         # Initialize the plot with the specified parameters
         self.ax = self._initialize_plot(
-            network_graph, figsize, background_color, plot_outline, outline_color, outline_scale
+            graph, figsize, background_color, plot_outline, outline_color, outline_scale
         )
 
     def _initialize_plot(
         self,
-        network_graph: NetworkGraph,
+        graph: NetworkGraph,
         figsize: tuple,
         background_color: str,
         plot_outline: bool,
@@ -62,7 +62,7 @@ class NetworkPlotter:
         """Set up the plot with figure size, optional circle perimeter, and background color.
 
         Args:
-            network_graph (NetworkGraph): The network data and attributes to be visualized.
+            graph (NetworkGraph): The network data and attributes to be visualized.
             figsize (tuple): Size of the figure in inches (width, height).
             background_color (str): Background color of the plot.
             plot_outline (bool): Whether to plot the network perimeter circle.
@@ -73,7 +73,7 @@ class NetworkPlotter:
             plt.Axes: The axis object for the plot.
         """
         # Extract node coordinates from the network graph
-        node_coordinates = network_graph.node_coordinates
+        node_coordinates = graph.node_coordinates
         # Calculate the center and radius of the bounding box around the network
         center, radius = _calculate_bounding_box(node_coordinates)
         # Scale the radius by the outline_scale factor
@@ -141,10 +141,10 @@ class NetworkPlotter:
             network_node_shape=node_shape,
         )
         # Extract node coordinates from the network graph
-        node_coordinates = self.network_graph.node_coordinates
+        node_coordinates = self.graph.node_coordinates
         # Draw the nodes of the graph
         nx.draw_networkx_nodes(
-            self.network_graph.G,
+            self.graph.network,
             pos=node_coordinates,
             node_size=node_size,
             node_color=node_color,
@@ -155,7 +155,7 @@ class NetworkPlotter:
         )
         # Draw the edges of the graph
         nx.draw_networkx_edges(
-            self.network_graph.G,
+            self.graph.network,
             pos=node_coordinates,
             width=edge_width,
             edge_color=edge_color,
@@ -197,20 +197,18 @@ class NetworkPlotter:
         )
         # Filter to get node IDs and their coordinates
         node_ids = [
-            self.network_graph.node_label_to_id_map.get(node)
+            self.graph.node_label_to_id_map.get(node)
             for node in nodes
-            if node in self.network_graph.node_label_to_id_map
+            if node in self.graph.node_label_to_id_map
         ]
         if not node_ids:
             raise ValueError("No nodes found in the network graph.")
 
         # Get the coordinates of the filtered nodes
-        node_coordinates = {
-            node_id: self.network_graph.node_coordinates[node_id] for node_id in node_ids
-        }
+        node_coordinates = {node_id: self.graph.node_coordinates[node_id] for node_id in node_ids}
         # Draw the nodes in the subnetwork
         nx.draw_networkx_nodes(
-            self.network_graph.G,
+            self.graph.network,
             pos=node_coordinates,
             nodelist=node_ids,
             node_size=node_size,
@@ -221,7 +219,7 @@ class NetworkPlotter:
             ax=self.ax,
         )
         # Draw the edges between the specified nodes in the subnetwork
-        subgraph = self.network_graph.G.subgraph(node_ids)
+        subgraph = self.graph.network.subgraph(node_ids)
         nx.draw_networkx_edges(
             subgraph,
             pos=node_coordinates,
@@ -260,9 +258,9 @@ class NetworkPlotter:
             color = self.get_annotated_contour_colors(color=color)
 
         # Extract node coordinates from the network graph
-        node_coordinates = self.network_graph.node_coordinates
+        node_coordinates = self.graph.node_coordinates
         # Draw contours for each domain in the network
-        for idx, (_, nodes) in enumerate(self.network_graph.domain_to_nodes.items()):
+        for idx, (_, nodes) in enumerate(self.graph.domain_to_nodes.items()):
             if len(nodes) > 1:
                 self._draw_kde_contour(
                     self.ax,
@@ -299,23 +297,23 @@ class NetworkPlotter:
         """
         # Log the plotting parameters
         params.log_plotter(
-            contour_levels=levels,
-            contour_bandwidth=bandwidth,
-            contour_grid_size=grid_size,
-            contour_alpha=alpha,
-            contour_color="custom" if isinstance(color, np.ndarray) else color,
+            subcontour_levels=levels,
+            subcontour_bandwidth=bandwidth,
+            subcontour_grid_size=grid_size,
+            subcontour_alpha=alpha,
+            subcontour_color="custom" if isinstance(color, np.ndarray) else color,
         )
         # Filter to get node IDs and their coordinates
         node_ids = [
-            self.network_graph.node_label_to_id_map.get(node)
+            self.graph.node_label_to_id_map.get(node)
             for node in nodes
-            if node in self.network_graph.node_label_to_id_map
+            if node in self.graph.node_label_to_id_map
         ]
         if not node_ids or len(node_ids) == 1:
             raise ValueError("No nodes found in the network graph or insufficient nodes to plot.")
 
         # Draw the KDE contour for the specified nodes
-        node_coordinates = self.network_graph.node_coordinates
+        node_coordinates = self.graph.node_coordinates
         self._draw_kde_contour(
             self.ax,
             node_coordinates,
@@ -402,8 +400,10 @@ class NetworkPlotter:
         fontcolor: Union[str, np.ndarray] = "black",
         arrow_linewidth: float = 1,
         arrow_color: Union[str, np.ndarray] = "black",
-        num_words: int = 10,
+        max_labels: Union[int, None] = None,
+        max_words: int = 10,
         min_words: int = 1,
+        words_to_omit: Union[List[str], None] = None,
     ) -> None:
         """Annotate the network graph with labels for different domains, positioned around the network for clarity.
 
@@ -415,8 +415,10 @@ class NetworkPlotter:
             fontcolor (str or np.ndarray, optional): Color of the label text. Can be a string or RGBA array. Defaults to "black".
             arrow_linewidth (float, optional): Line width of the arrows pointing to centroids. Defaults to 1.
             arrow_color (str or np.ndarray, optional): Color of the arrows. Can be a string or RGBA array. Defaults to "black".
-            num_words (int, optional): Maximum number of words in a label. Defaults to 10.
+            max_labels (int, optional): Maximum number of labels to plot. Defaults to None (no limit).
+            max_words (int, optional): Maximum number of words in a label. Defaults to 10.
             min_words (int, optional): Minimum number of words required to display a label. Defaults to 1.
+            words_to_omit (List[str], optional): List of words to omit from the labels. Defaults to None.
         """
         # Log the plotting parameters
         params.log_plotter(
@@ -427,39 +429,82 @@ class NetworkPlotter:
             label_fontcolor="custom" if isinstance(fontcolor, np.ndarray) else fontcolor,
             label_arrow_linewidth=arrow_linewidth,
             label_arrow_color="custom" if isinstance(arrow_color, np.ndarray) else arrow_color,
-            label_num_words=num_words,
+            label_max_labels=max_labels,
+            label_max_words=max_words,
             label_min_words=min_words,
+            label_words_to_omit=words_to_omit,  # Log words_to_omit parameter
         )
+
         # Convert color strings to RGBA arrays if necessary
         if isinstance(fontcolor, str):
-            fontcolor = self.get_annotated_contour_colors(color=fontcolor)
+            fontcolor = self.get_annotated_label_colors(color=fontcolor)
         if isinstance(arrow_color, str):
-            arrow_color = self.get_annotated_contour_colors(color=arrow_color)
+            arrow_color = self.get_annotated_label_colors(color=arrow_color)
+        # Normalize words_to_omit to lowercase
+        if words_to_omit:
+            words_to_omit = set(word.lower() for word in words_to_omit)
 
         # Calculate the center and radius of the network
         domain_centroids = {}
-        for domain, nodes in self.network_graph.domain_to_nodes.items():
+        for domain, nodes in self.graph.domain_to_nodes.items():
             if nodes:  # Skip if the domain has no nodes
                 domain_centroids[domain] = self._calculate_domain_centroid(nodes)
 
+        # Initialize empty lists to collect valid indices
+        valid_indices = []
+        filtered_domain_centroids = {}
+        filtered_domain_terms = {}
+        # Loop through domain_centroids with index
+        for idx, (domain, centroid) in enumerate(domain_centroids.items()):
+            # Process the domain term
+            terms = self.graph.trimmed_domain_to_term[domain].split(" ")
+            # Remove words_to_omit
+            if words_to_omit:
+                terms = [term for term in terms if term.lower() not in words_to_omit]
+            # Trim to max_words
+            terms = terms[:max_words]
+            # Check if the domain passes the word count condition
+            if len(terms) >= min_words:
+                # Add to filtered_domain_centroids
+                filtered_domain_centroids[domain] = centroid
+                # Store the trimmed terms
+                filtered_domain_terms[domain] = " ".join(terms)
+                # Keep track of the valid index
+                valid_indices.append(idx)
+
+        # If max_labels is specified and less than the available labels
+        if max_labels is not None and max_labels < len(filtered_domain_centroids):
+            step = len(filtered_domain_centroids) / max_labels
+            selected_indices = [int(i * step) for i in range(max_labels)]
+            filtered_domain_centroids = {
+                k: v
+                for i, (k, v) in enumerate(filtered_domain_centroids.items())
+                if i in selected_indices
+            }
+            filtered_domain_terms = {
+                k: v
+                for i, (k, v) in enumerate(filtered_domain_terms.items())
+                if i in selected_indices
+            }
+            fontcolor = fontcolor[selected_indices]
+            arrow_color = arrow_color[selected_indices]
+
+        # Update the terms in the graph after omitting words and filtering
+        for domain, terms in filtered_domain_terms.items():
+            self.graph.trimmed_domain_to_term[domain] = terms
+
         # Calculate the bounding box around the network
         center, radius = _calculate_bounding_box(
-            self.network_graph.node_coordinates, radius_margin=perimeter_scale
+            self.graph.node_coordinates, radius_margin=perimeter_scale
         )
-
-        # Filter out domains with insufficient words for labeling
-        filtered_domains = {
-            domain: centroid
-            for domain, centroid in domain_centroids.items()
-            if len(self.network_graph.trimmed_domain_to_term[domain].split(" ")[:num_words])
-            >= min_words
-        }
         # Calculate the best positions for labels around the perimeter
-        best_label_positions = _best_label_positions(filtered_domains, center, radius, offset)
+        best_label_positions = _best_label_positions(
+            filtered_domain_centroids, center, radius, offset
+        )
         # Annotate the network with labels
         for idx, (domain, pos) in enumerate(best_label_positions.items()):
-            centroid = filtered_domains[domain]
-            annotations = self.network_graph.trimmed_domain_to_term[domain].split(" ")[:num_words]
+            centroid = filtered_domain_centroids[domain]
+            annotations = self.graph.trimmed_domain_to_term[domain].split(" ")[:max_words]
             self.ax.annotate(
                 "\n".join(annotations),
                 xy=centroid,
@@ -473,6 +518,81 @@ class NetworkPlotter:
                 arrowprops=dict(arrowstyle="->", color=arrow_color[idx], linewidth=arrow_linewidth),
             )
 
+    def plot_sublabel(
+        self,
+        nodes: list,
+        label: str,
+        radial_position: float = 0.0,
+        perimeter_scale: float = 1.05,
+        offset: float = 0.10,
+        font: str = "Arial",
+        fontsize: int = 10,
+        fontcolor: str = "black",
+        arrow_linewidth: float = 1,
+        arrow_color: str = "black",
+    ) -> None:
+        """Annotate the network graph with a single label for the given nodes, positioned at a specified radial angle.
+
+        Args:
+            nodes (List[str]): List of node labels to be used for calculating the centroid.
+            label (str): The label to be annotated on the network.
+            radial_position (float, optional): Radial angle for positioning the label, in degrees (0-360). Defaults to 0.0.
+            perimeter_scale (float, optional): Scale factor for positioning the label around the perimeter. Defaults to 1.05.
+            offset (float, optional): Offset distance for the label from the perimeter. Defaults to 0.10.
+            font (str, optional): Font name for the label. Defaults to "Arial".
+            fontsize (int, optional): Font size for the label. Defaults to 10.
+            fontcolor (str, optional): Color of the label text. Defaults to "black".
+            arrow_linewidth (float, optional): Line width of the arrow pointing to the centroid. Defaults to 1.
+            arrow_color (str, optional): Color of the arrow. Defaults to "black".
+        """
+        # Log the plotting parameters
+        params.log_plotter(
+            sublabel_perimeter_scale=perimeter_scale,
+            sublabel_offset=offset,
+            sublabel_font=font,
+            sublabel_fontsize=fontsize,
+            sublabel_fontcolor="custom" if isinstance(fontcolor, np.ndarray) else fontcolor,
+            sublabel_arrow_linewidth=arrow_linewidth,
+            sublabel_arrow_color="custom" if isinstance(arrow_color, np.ndarray) else arrow_color,
+            sublabel_radial_position=radial_position,
+        )
+
+        # Map node labels to IDs
+        node_ids = [
+            self.graph.node_label_to_id_map.get(node)
+            for node in nodes
+            if node in self.graph.node_label_to_id_map
+        ]
+        if not node_ids or len(node_ids) == 1:
+            raise ValueError("No nodes found in the network graph or insufficient nodes to plot.")
+
+        # Calculate the centroid of the provided nodes
+        centroid = self._calculate_domain_centroid(node_ids)
+        # Calculate the bounding box around the network
+        center, radius = _calculate_bounding_box(
+            self.graph.node_coordinates, radius_margin=perimeter_scale
+        )
+        # Convert radial position to radians, adjusting for a 90-degree rotation
+        radial_radians = np.deg2rad(radial_position - 90)
+        label_position = (
+            center[0] + (radius + offset) * np.cos(radial_radians),
+            center[1] + (radius + offset) * np.sin(radial_radians),
+        )
+
+        # Annotate the network with the label
+        self.ax.annotate(
+            label,
+            xy=centroid,
+            xytext=label_position,
+            textcoords="data",
+            ha="center",
+            va="center",
+            fontsize=fontsize,
+            fontname=font,
+            color=fontcolor,
+            arrowprops=dict(arrowstyle="->", color=arrow_color, linewidth=arrow_linewidth),
+        )
+
     def _calculate_domain_centroid(self, nodes: list) -> tuple:
         """Calculate the most centrally located node in .
 
@@ -483,7 +603,7 @@ class NetworkPlotter:
             tuple: A tuple containing the domain's central node coordinates.
         """
         # Extract positions of all nodes in the domain
-        node_positions = self.network_graph.node_coordinates[nodes, :]
+        node_positions = self.graph.node_coordinates[nodes, :]
         # Calculate the pairwise distance matrix between all nodes in the domain
         distances_matrix = np.linalg.norm(node_positions[:, np.newaxis] - node_positions, axis=2)
         # Sum the distances for each node to all other nodes in the domain
@@ -508,7 +628,7 @@ class NetworkPlotter:
             np.ndarray: Array of RGBA colors adjusted for enrichment status.
         """
         # Get the initial domain colors for each node
-        network_colors = self.network_graph.get_domain_colors(**kwargs, random_seed=random_seed)
+        network_colors = self.graph.get_domain_colors(**kwargs, random_seed=random_seed)
         if isinstance(nonenriched_color, str):
             # Convert the non-enriched color from string to RGBA
             nonenriched_color = mcolors.to_rgba(nonenriched_color)
@@ -535,14 +655,14 @@ class NetworkPlotter:
         """
         # Merge all enriched nodes from the domain_to_nodes dictionary
         enriched_nodes = set()
-        for _, nodes in self.network_graph.domain_to_nodes.items():
+        for _, nodes in self.graph.domain_to_nodes.items():
             enriched_nodes.update(nodes)
 
         # Initialize all node sizes to the non-enriched size
-        node_sizes = np.full(len(self.network_graph.G.nodes), nonenriched_nodesize)
+        node_sizes = np.full(len(self.graph.network.nodes), nonenriched_nodesize)
         # Set the size for enriched nodes
         for node in enriched_nodes:
-            if node in self.network_graph.G.nodes:
+            if node in self.graph.network.nodes:
                 node_sizes[node] = enriched_nodesize
 
         return node_sizes
@@ -587,12 +707,12 @@ class NetworkPlotter:
         if isinstance(color, str):
             # If a single color string is provided, convert it to RGBA and apply to all domains
             rgba_color = np.array(matplotlib.colors.to_rgba(color))
-            return np.array([rgba_color for _ in self.network_graph.domain_to_nodes])
+            return np.array([rgba_color for _ in self.graph.domain_to_nodes])
 
         # Generate colors for each domain using the provided arguments and random seed
-        node_colors = self.network_graph.get_domain_colors(**kwargs, random_seed=random_seed)
+        node_colors = self.graph.get_domain_colors(**kwargs, random_seed=random_seed)
         annotated_colors = []
-        for _, nodes in self.network_graph.domain_to_nodes.items():
+        for _, nodes in self.graph.domain_to_nodes.items():
             if len(nodes) > 1:
                 # For domains with multiple nodes, choose the brightest color (sum of RGB values)
                 domain_colors = np.array([node_colors[node] for node in nodes])
@@ -662,12 +782,12 @@ def _calculate_bounding_box(
 
 
 def _best_label_positions(
-    filtered_domains: Dict[str, Any], center: np.ndarray, radius: float, offset: float
+    filtered_domain_centroids: Dict[str, Any], center: np.ndarray, radius: float, offset: float
 ) -> Dict[str, Any]:
     """Calculate and optimize label positions for clarity.
 
     Args:
-        filtered_domains (dict): Centroids of the filtered domains.
+        filtered_domain_centroids (dict): Centroids of the filtered domains.
         center (np.ndarray): The center coordinates for label positioning.
         radius (float): The radius for positioning labels around the center.
         offset (float): The offset distance from the radius for positioning labels.
@@ -675,15 +795,16 @@ def _best_label_positions(
     Returns:
         dict: Optimized positions for labels.
     """
-    num_domains = len(filtered_domains)
+    num_domains = len(filtered_domain_centroids)
     # Calculate equidistant positions around the center for initial label placement
     equidistant_positions = _equidistant_angles_around_center(center, radius, offset, num_domains)
     # Create a mapping of domains to their initial label positions
     label_positions = {
-        domain: position for domain, position in zip(filtered_domains.keys(), equidistant_positions)
+        domain: position
+        for domain, position in zip(filtered_domain_centroids.keys(), equidistant_positions)
     }
     # Optimize the label positions to minimize distance to domain centroids
-    return _optimize_label_positions(label_positions, filtered_domains)
+    return _optimize_label_positions(label_positions, filtered_domain_centroids)
 
 
 def _equidistant_angles_around_center(