updated wind rose viz

NREL · Aug 21, 2024 · eaba6b5 · eaba6b5
1 parent 2c3be8f
commit eaba6b5
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/floris/wind_data.py b/floris/wind_data.py
@@ -5,6 +5,7 @@
 from pathlib import Path
 
 import matplotlib.pyplot as plt
+import matplotlib as mpl
 import numpy as np
 import pandas as pd
 from pandas.api.types import CategoricalDtype
@@ -692,7 +693,7 @@ def plot(
         color_map="viridis_r",
         wd_step=None,
         ws_step=None,
-        legend_kwargs={"title": "Wind speed [m/s]"},
+        legend_kwargs={"label": "Wind speed [m/s]"},
     ):
         """
         This method creates a wind rose plot showing the frequency of occurrence
@@ -738,6 +739,8 @@ def plot(
 
         # Get a color array
         color_array = plt.get_cmap(color_map, len(ws_bins))
+        norm_ws = mpl.colors.Normalize(vmin=np.min(ws_bins), vmax=np.max(ws_bins))
+        sm_ws = mpl.cm.ScalarMappable(norm=norm_ws, cmap=color_array)
 
         for wd_idx, wd in enumerate(wd_bins):
             rects = []
@@ -755,7 +758,10 @@ def plot(
                 )
 
         # Configure the plot
-        ax.legend(reversed(rects), ws_bins, **legend_kwargs)
+        # ax.legend(reversed(rects), ws_bins, **legend_kwargs)
+        cb_ws = ax.figure.colorbar(sm_ws, ax=ax, **legend_kwargs)
+        ax.figure.tight_layout()
+        cb_ws.outline.set_visible(False)
         ax.set_theta_direction(-1)
         ax.set_theta_offset(np.pi / 2.0)
         ax.set_theta_zero_location("N")