Add function to label turbines on CutPlane plots

examples/03_making_adjustments.py

@@ -17,7 +17,7 @@
 import numpy as np
 
 from floris.tools import FlorisInterface
-from floris.tools import visualize_cut_plane #, plot_turbines_with_fi
+import floris.tools.visualization as wakeviz
 
 """
 This example makes changes to the given input file through the script.
@@ -38,53 +38,53 @@
 
 # Plot a horizatonal slice of the initial configuration
 horizontal_plane = fi.calculate_horizontal_plane(height=90.0)
-visualize_cut_plane(horizontal_plane, ax=axarr[0], title="Initial setup", min_speed=MIN_WS, max_speed=MAX_WS)
-
+wakeviz.visualize_cut_plane(horizontal_plane, ax=axarr[0], title="Initial setup", min_speed=MIN_WS, max_speed=MAX_WS)
 
 # Change the wind speed
 horizontal_plane = fi.calculate_horizontal_plane(ws=[7.0], height=90.0)
-visualize_cut_plane(horizontal_plane, ax=axarr[1], title="Wind speed at 7 m/s", min_speed=MIN_WS, max_speed=MAX_WS)
+wakeviz.visualize_cut_plane(horizontal_plane, ax=axarr[1], title="Wind speed at 7 m/s", min_speed=MIN_WS, max_speed=MAX_WS)
 
 
 # Change the wind shear, reset the wind speed, and plot a vertical slice
 fi.reinitialize( wind_shear=0.2, wind_speeds=[8.0] )
 y_plane = fi.calculate_y_plane(crossstream_dist=0.0)
-visualize_cut_plane(y_plane, ax=axarr[2], title="Wind shear at 0.2", min_speed=MIN_WS, max_speed=MAX_WS)
+wakeviz.visualize_cut_plane(y_plane, ax=axarr[2], title="Wind shear at 0.2", min_speed=MIN_WS, max_speed=MAX_WS)
 
 
-# Change the farm layout
+# # Change the farm layout
 N = 3  # Number of turbines per row and per column
 X, Y = np.meshgrid(
     5.0 * fi.floris.farm.rotor_diameters[0][0][0] * np.arange(0, N, 1),
     5.0 * fi.floris.farm.rotor_diameters[0][0][0] * np.arange(0, N, 1),
 )
 fi.reinitialize(layout_x=X.flatten(), layout_y=Y.flatten())
 horizontal_plane = fi.calculate_horizontal_plane(height=90.0)
-visualize_cut_plane(horizontal_plane, ax=axarr[3], title="3x3 Farm", min_speed=MIN_WS, max_speed=MAX_WS)
+wakeviz.visualize_cut_plane(horizontal_plane, ax=axarr[3], title="3x3 Farm", min_speed=MIN_WS, max_speed=MAX_WS)
+wakeviz.add_turbine_id_labels(fi, axarr[3], color="w", backgroundcolor="k")
 
 
 # Change the yaw angles and configure the plot differently
 yaw_angles = np.zeros((1, 1, N * N))
 
 ## First row
 yaw_angles[:,:,0] = 30.0
-yaw_angles[:,:,1] = -30.0
-yaw_angles[:,:,2] = 30.0
+yaw_angles[:,:,3] = -30.0
+yaw_angles[:,:,6] = 30.0
 
 ## Second row
-yaw_angles[:,:,3] = -30.0
+yaw_angles[:,:,1] = -30.0
 yaw_angles[:,:,4] = 30.0
-yaw_angles[:,:,5] = -30.0
+yaw_angles[:,:,7] = -30.0
 
 horizontal_plane = fi.calculate_horizontal_plane(yaw_angles=yaw_angles, height=90.0)
-visualize_cut_plane(horizontal_plane, ax=axarr[4], title="Yawesome art", cmap="PuOr", min_speed=MIN_WS, max_speed=MAX_WS)
-# plot_turbines_with_fi(axarr[8], fi)
+wakeviz.visualize_cut_plane(horizontal_plane, ax=axarr[4], title="Yawesome art", cmap="PuOr", min_speed=MIN_WS, max_speed=MAX_WS)
+wakeviz.plot_turbines_with_fi(fi, axarr[4], yaw_angles=yaw_angles, color="c")
 
 
 # Plot the cross-plane of the 3x3 configuration
 cross_plane = fi.calculate_cross_plane(yaw_angles=yaw_angles, downstream_dist=610.0)
-visualize_cut_plane(cross_plane, ax=axarr[5], title="Cross section at 610 m", min_speed=MIN_WS, max_speed=MAX_WS)
+wakeviz.visualize_cut_plane(cross_plane, ax=axarr[5], title="Cross section at 610 m", min_speed=MIN_WS, max_speed=MAX_WS)
 axarr[5].invert_xaxis()
 
 
-plt.show()
+wakeviz.show_plots()

floris/tools/visualization.py

@@ -14,14 +14,14 @@
 from __future__ import annotations
 
 from typing import Union
-from itertools import product
 
 import numpy as np
 import matplotlib as mpl
 import matplotlib.colors as mplcolors
 import matplotlib.pyplot as plt
 from matplotlib import rcParams
 
+from floris.tools.floris_interface import FlorisInterface
 
 def show_plots():
     plt.show()
@@ -54,7 +54,7 @@ def plot_turbines(ax, layout_x, layout_y, yaw_angles, rotor_diameters, color=Non
         ax.plot([x_0, x_1], [y_0, y_1], color=color)
 
 
-def plot_turbines_with_fi(fi, ax=None, color=None, yaw_angles=None):
+def plot_turbines_with_fi(fi: FlorisInterface, ax=None, color=None, yaw_angles=None):
     """
     Wrapper function to plot turbines which extracts the data
     from a FLORIS interface object
@@ -80,6 +80,21 @@ def plot_turbines_with_fi(fi, ax=None, color=None, yaw_angles=None):
     )
 
 
+def add_turbine_id_labels(fi: FlorisInterface, ax: plt.Axes, **kwargs):
+    """
+    Adds index labels to a plot based on the given FlorisInterface.
+    See the pyplot.annotate docs for more info: https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.annotate.html.
+    kwargs are passed to Text (https://matplotlib.org/stable/api/text_api.html#matplotlib.text.Text).
+
+    Args:
+        fi (FlorisInterface): Simulation object to get the layout and index information.
+        ax (plt.Axes): Axes object to add the labels.
+    """
+
+    for i in range(fi.floris.farm.n_turbines):
+        ax.annotate(i, (fi.layout_x[i], fi.layout_y[i]), xytext=(0,10), textcoords="offset points", **kwargs)
+
+
 def line_contour_cut_plane(cut_plane, ax=None, levels=None, colors=None, **kwargs):
     """
     Visualize a cut_plane as a line contour plot.
@@ -125,7 +140,7 @@ def visualize_cut_plane(
     levels=None,
     color_bar=False,
     title="",
-    kwargs={}
+    **kwargs
 ):
     """
     Generate pseudocolor mesh plot of the cut_plane.