feat: add plot_power_flow

Author: phumthep

src/pownet/core/visualizer.py

@@ -6,6 +6,7 @@
 import geopandas as gpd
 import matplotlib as mpl
 from matplotlib.colors import ListedColormap
+from matplotlib.lines import Line2D
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
@@ -62,7 +63,7 @@ def plot_fuelmix_bar(
             bbox_to_anchor=(0.5, -0.12),
         )
         ax.set_ylabel("Power (MW)")
-        ax.set_ylim(top=(demand[:total_timesteps].max() * 1.30).values[0])
+        ax.set_ylim(top=(demand[:total_timesteps].max() * 1.30))
 
         if output_folder is not None:
             figure_name = f"{self.model_id}_fuelmix.png"
@@ -396,3 +397,172 @@ def plot_generation_by_contracts(self, contract_generation: pd.DataFrame) -> Non
         df_to_plot.plot(ax=ax, kind="bar", linewidth=2, legend=False)
         ax.set_ylabel("Generation (MWh)")
         plt.show()
+
+    def plot_power_flow(self, 
+                        flow_variables: pd.DataFrame, 
+                        figsize_per_line: tuple = (10, 2), # Note: height is for plot area of one subplot
+                        fixed_legend_height_inches: float = 0.5) -> None:
+        """
+        Plots the power flow on transmission lines over time.
+
+        Each unique transmission line (node_a to node_b) gets its own subplot.
+        The y-axis label for each subplot is the line segment name.
+        Legend is placed at the top center of the figure, occupying a fixed absolute height.
+        Power flow is colored:
+        - Green: Positive flow
+        - Red: Negative flow
+        - Black: Zero flow
+
+        Args:
+            flow_variables (pd.DataFrame): DataFrame with simulation results. Expected columns:
+                                     'node_a', 'node_b', 'value', 'type' ('fwd' or 'bwd'), 'hour'.
+            figsize_per_line (tuple): Tuple specifying (width, height_for_each_subplot_plot_area).
+            fixed_legend_height_inches (float): Absolute height in inches for the legend area at the top.
+        """
+        lines = flow_variables[['node_a', 'node_b']].drop_duplicates().values.tolist()
+        num_lines = len(lines)
+
+        if num_lines == 0:
+            print("No transmission lines to plot.")
+            return
+
+        fig_width = figsize_per_line[0]
+        
+        # Calculate the total height needed for the plot areas of all subplots
+        plots_area_height_inches = figsize_per_line[1] * num_lines
+        if plots_area_height_inches < 0: # Ensure non-negative plot height
+            plots_area_height_inches = 0
+
+        # Calculate the total figure height, including the fixed space for the legend
+        total_figure_height_inches = plots_area_height_inches + fixed_legend_height_inches
+        
+        # Ensure total figure height is positive; if not, default to a minimum
+        if total_figure_height_inches <= 0:
+            total_figure_height_inches = max(fixed_legend_height_inches, 1.0) # Use legend height or 1 inch min
+
+        fig, axes = plt.subplots(num_lines, 1,
+                                 figsize=(fig_width, total_figure_height_inches), # Use the calculated total height
+                                 sharex=True, squeeze=False)
+        
+        # Calculate the fraction of the total figure height that the legend area will occupy
+        if total_figure_height_inches > 0: # Avoid division by zero
+            legend_top_margin_fraction = fixed_legend_height_inches / total_figure_height_inches
+        else: # Should be unreachable due to the check above
+            legend_top_margin_fraction = 0.5 # Fallback: 50% for legend if total height is still 0
+
+        # Ensure legend fraction is reasonable (e.g., not more than 80% if there are plots)
+        # This prevents plots from being overly squished if their requested height is tiny.
+        if plots_area_height_inches > 0 and legend_top_margin_fraction > 0.8:
+            legend_top_margin_fraction = 0.8
+
+
+        for i, line_nodes in enumerate(lines):
+            ax = axes[i, 0] 
+            node_a, node_b = line_nodes
+            line_segment_name = f"{node_a} to {node_b}\nPower flow (MW)"
+
+            line_df = flow_variables[(flow_variables['node_a'] == node_a) & (flow_variables['node_b'] == node_b)]
+
+            if line_df.empty:
+                ax.set_title(f"Power Flow: {line_segment_name} (No data)")
+                ax.set_ylabel(line_segment_name)
+                ax.text(0.5, 0.5, 'No data for this line',
+                        horizontalalignment='center', verticalalignment='center',
+                        transform=ax.transAxes)
+                continue
+            
+            try:
+                pivot_df = line_df.pivot_table(index='hour', columns='type', values='value', fill_value=0)
+            except Exception as e:
+                ax.set_title(f"Power Flow: {line_segment_name} (Error pivoting data)")
+                ax.set_ylabel(line_segment_name)
+                ax.text(0.5, 0.5, f'Error processing data: {e}',
+                        horizontalalignment='center', verticalalignment='center',
+                        transform=ax.transAxes)
+                continue
+
+            if 'fwd' not in pivot_df.columns:
+                pivot_df['fwd'] = 0
+            if 'bwd' not in pivot_df.columns:
+                pivot_df['bwd'] = 0
+
+            pivot_df = pivot_df.sort_index()
+            net_flow = pivot_df['fwd'] - pivot_df['bwd']
+            hours = net_flow.index
+
+            if len(hours) < 2:
+                if len(hours) == 1:
+                    y_val = net_flow.iloc[0]
+                    color = 'green' if y_val > 0 else ('red' if y_val < 0 else 'black')
+                    ax.plot(hours[0], y_val, marker='o', color=color)
+                ax.set_title(f"Power Flow: {line_segment_name} (Not enough data for line plot)")
+                ax.set_ylabel(line_segment_name)
+                if len(hours) > 0: 
+                    ax.set_xlim(hours.min() -1 if pd.api.types.is_numeric_dtype(hours) else hours.min() - pd.Timedelta(days=1), 
+                                hours.max() + 1 if pd.api.types.is_numeric_dtype(hours) else hours.max() + pd.Timedelta(days=1))
+                continue
+
+            for j in range(len(hours) - 1):
+                x1, x2 = hours[j], hours[j+1]
+                y1, y2 = net_flow.iloc[j], net_flow.iloc[j+1]
+
+                x1_num = pd.to_numeric(x1)
+                x2_num = pd.to_numeric(x2)
+                
+                if y1 == 0 and y2 == 0:
+                    ax.plot([x1, x2], [y1, y2], color='black', linestyle='-')
+                elif y1 * y2 >= 0: 
+                    if y1 == 0: 
+                        color = 'green' if y2 > 0 else ('red' if y2 < 0 else 'black')
+                    elif y2 == 0: 
+                        color = 'green' if y1 > 0 else ('red' if y1 < 0 else 'black')
+                    else: 
+                        color = 'green' if y1 > 0 else 'red'
+                    ax.plot([x1, x2], [y1, y2], color=color, linestyle='-')
+                else: 
+                    if (y2 - y1) == 0: 
+                        x_intersect_num = x1_num 
+                    else:
+                        x_intersect_num = x1_num - y1 * (x2_num - x1_num) / (y2 - y1)
+                    
+                    if isinstance(x1, (pd.Timestamp, np.datetime64)):
+                        x_intersect = pd.Timestamp(x_intersect_num)
+                    elif isinstance(x1, pd.Timedelta) or isinstance(x1, np.timedelta64):
+                         x_intersect = pd.Timedelta(x_intersect_num, unit='ns')
+                    else:
+                        x_intersect = x_intersect_num
+
+                    color1 = 'green' if y1 > 0 else 'red'
+                    ax.plot([x1, x_intersect], [y1, 0], color=color1, linestyle='-')
+                    color2 = 'green' if y2 > 0 else 'red'
+                    ax.plot([x_intersect, x2], [0, y2], color=color2, linestyle='-')
+
+            ax.axhline(0, color='gray', linestyle='--', linewidth=0.8)
+            # ax.set_title(f"Power Flow: {line_segment_name}")
+            ax.set_ylabel(line_segment_name)
+
+        if num_lines > 0:
+            axes[-1, 0].set_xlabel("Hour")
+
+        legend_elements = [Line2D([0], [0], color='green', lw=2, label='Positive Flow'),
+                           Line2D([0], [0], color='red', lw=2, label='Negative Flow'),
+                           Line2D([0], [0], color='black', lw=2, label='Zero Flow')]
+
+        fig.tight_layout()
+        
+        # Adjust the top of the subplots area to make space for the legend.
+        # The subplots will occupy the space from y=0 to y=(1 - legend_top_margin_fraction)
+        fig.subplots_adjust(top=(1 - legend_top_margin_fraction))
+
+        # Define the bounding box for the legend at the top of the figure
+        # y coordinate for bbox starts where subplots end and goes up by legend_top_margin_fraction
+        legend_bbox_y_start = 1 - legend_top_margin_fraction
+        legend_bbox = (0, legend_bbox_y_start, 1, legend_top_margin_fraction)
+
+        fig.legend(handles=legend_elements,
+                   loc='center',
+                   bbox_to_anchor=legend_bbox,
+                   ncol=3,
+                   frameon=False)
+
+        plt.show()