### ===================================================================================================================
### Plotting functions for Signal Processing Tool
### ===================================================================================================================
# Copyright ©2025 Haskoning Nederland B.V.
### ===================================================================================================================
### 1. Constants
### ===================================================================================================================
# General imports
import numpy as np
from pathlib import Path
from typing import Union, List, Optional, Tuple
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
from matplotlib.legend import Legend
from matplotlib.offsetbox import OffsetBox
from mpl_toolkits.mplot3d import Axes3D
### ===================================================================================================================
### 2. Helper functions and classes for plotting
### ===================================================================================================================
[docs]
class LegendHandler:
""" Custom legend handler for matplotlib plots."""
def __init__(
self, legend: Legend = None, orig_handle: Line2D = None, fontsize: float = 10, handle_box: Legend = None):
"""
Initialise the custom legend handler.
Input:
- legend (Legend): The legend instance.
- orig_handle (Line2D): The original line handle.
- fontsize (float): The font size for the legend.
- handle_box (Legend): The handle box for the legend.
"""
self.legend = legend
self.orig_handle = orig_handle
self.fontsize = fontsize
self.handle_box = handle_box
[docs]
@staticmethod
def legend_artist(legend: Legend, orig_handle: Line2D, fontsize: float, handlebox: OffsetBox) -> Line2D:
"""
Method to create a custom legend artist.
Input:
- legend (Legend): The legend instance.
- orig_handle (Line2D): The original line handle.
- fontsize (float): The font size for the legend.
- handlebox (Legend): The handle box for the legend.
Output:
- Returns the custom line for the legend.
"""
x0, y0 = handlebox.xdescent, handlebox.ydescent
width, height = handlebox.width, handlebox.height
line = Line2D(
xdata=[x0, x0 + width], ydata=[y0 + height / 2, y0 + height / 2], lw=orig_handle.get_linewidth() * 3,
color=orig_handle.get_color(), linestyle=orig_handle.get_linestyle())
handlebox.add_artist(line)
return line
[docs]
def handle_plot_saving(file: Optional[Union[Path, str]], title: str) -> Path:
""" Generic function to handle saving the plot to a file."""
if isinstance(file, str):
file = Path(file)
if file.is_dir():
file = file / f'{title}.png'
if not file.suffix:
file = file.parent / f'{file.name}.png'
if file.suffix != '.png':
raise ValueError("ERROR: Only png-files can be created.")
plt.savefig(file, format='png', dpi=300)
if not file.exists():
raise FileNotFoundError(f"ERROR: The plot file could not be created for {file.as_posix()}.")
return file
[docs]
def get_frequency_ban_level_units(level_type: str, amplitude_units: str, x_units: str) -> str:
"""
Function to get the units for frequency band levels based on the level type.
Input:
- level_type (str): Specifies the method used to compute the level of each frequency band. Options include
'peak', 'energy', 'power', and 'RMS'. Default value is 'RMS'.
- amplitude_units (str): The units of the amplitude (e.g., 'm/s^2', 'N', etc.).
- x_units (str): The units of the x-axis (e.g., 's', 'Hz', etc.).
Output:
- Returns the string for plotting the units of the frequency band levels.
"""
if not isinstance(level_type, str):
raise TypeError("ERROR: Input for level-type should be a string. Please check your input.")
level_type = level_type.lower()
if level_type not in ['peak', 'energy', 'power', 'rms']:
raise ValueError(
"ERROR: Invalid input for level-type. Please select from 'peak', 'energy', 'power', or 'RMS'.")
# Set units based on the type
if level_type == 'peak':
return f'{amplitude_units}'
elif level_type == 'rms':
return f'{amplitude_units}'
elif level_type == 'power':
return f'{amplitude_units}^2'
return f'{amplitude_units}^2*{x_units}'
### ===================================================================================================================
### 3. Functions to create plots in Signal Processing Tool
### ===================================================================================================================
[docs]
def plot_function(
x_values: List[List[float]], y_values: List[List[float]], labels: List[str],
title: str, x_label: str, y_label: str, log_scale_x: bool = False, log_scale_y: bool = False,
x_lim: Optional[List[float]] = None, y_lim: Optional[List[float]] = None,
file: Optional[Union[Path, str]] = None, show: bool = True, line_width: float = 1,
plot_max_values: bool = False, horizontal_line: Optional[float] = None) -> Optional[Path]:
"""
Generic function for plotting in the Signal Processing Tool for function with the given x and y values.
Input:
- x_values (list of lists of floats): The x-coordinates of the points on the plot.
- y_values (list of lists of floats): The y-coordinates of the points on the plot.
- labels (list of str): The labels for the data series.
- title (str): The title of the plot.
- x_label (str): The label for the x-axis.
- y_label (str): The label for the y-axis.
- log_scale_x (bool): Select to set the x-axis to a logarithmic scale. Default value is False.
- log_scale_y (bool): Select to set the y-axis to a logarithmic scale. Default value is False.
- x_lim (list of float): Optional list to set limits [bottom limit, top limit] for the x-axis. Default value is
None.
- y_lim (list of float): Optional list to set limits [bottom limit, top limit] for the y-axis. Default value is
None.
- file (Path): Provide the full path and filename for the image of the graph to be created. The image is a
png-file, if the suffix is omitted it will be added. An error will occur if the picture format is set
to a non-supported format. It is also possible to provide the folder only, the title is used as filename.
- show (bool): Select to plot the graph in the environment when running the script. Default value True.
- line_width (float): Change line width, default = 1.
- plot_max_values (bool): Select to calculate the maximum y value for each line and add it to the plot. Default
value is False.
- horizontal_line (float): Optional input to draw a horizontal line at this y-value. Default value is None.
Output:
- Returns the path of the file if the plot was selected to be saved to a file, otherwise returns None.
"""
# Prepare the environment
plt.close()
# Set the figure size
plt.figure(figsize=(16, 4))
# Plot the data
lines = []
for x_val, y_val, label in zip(x_values, y_values, labels):
line, = plt.plot(x_val, y_val, label=label, linewidth=line_width, alpha=0.8)
lines.append(line)
if plot_max_values:
max_y = max(y_val)
max_x = x_val[np.array(y_val).tolist().index(max_y)]
plt.text(max_x, max_y, f' [{max_y:.2f}, {max_x:.2f}]', fontsize=8, color=line.get_color(), weight='bold')
# Set the axes labels and title
plt.xlabel(x_label)
plt.ylabel(y_label)
plt.title(title)
# Set the axes to logarithmic scale
if log_scale_x:
plt.xscale('log')
if log_scale_y:
plt.yscale('log')
# Enable grid
plt.grid(visible=True, which='both', axis='both')
# Set the axes limits
if x_lim:
plt.xlim(*x_lim)
else:
plt.xlim(min([min(x) for x in x_values]), max([max(x) for x in x_values]))
if y_lim:
plt.ylim(*y_lim)
elif x_lim and not y_lim:
# If x-lim is provided but not y-lim, set y-lim based on x-lim range
filtered_y_values = [
y_val for x_vals, y_vals in zip(x_values, y_values)
for x_val, y_val in zip(x_vals, y_vals) if x_lim[0] <= x_val <= x_lim[1]]
min_y = min(filtered_y_values)
max_y = max(filtered_y_values)
plt.ylim(min_y * 1.2 if min_y < 0 else min_y * 0.8, max_y * 0.8 if max_y < 0 else max_y * 1.2)
else:
min_y = min(min(y) for y in y_values)
max_y = max(max(y) for y in y_values)
plt.ylim(min_y * 1.2 if min_y < 0 else min_y * 0.8, max_y * 0.8 if max_y < 0 else max_y * 1.2)
# Draw horizontal line if specified
if horizontal_line is not None:
plt.axhline(y=horizontal_line, color='red', linestyle='--', linewidth=1, alpha=0.7)
# Create a legend with the custom handler
plt.legend(
handles=lines, handler_map={Line2D: LegendHandler(legend=None, orig_handle=None, fontsize=10, handle_box=None)},
loc='upper left', bbox_to_anchor=(1, 1), fontsize=8)
plt.tight_layout()
# Save figure
if file:
file = handle_plot_saving(file=file, title=title)
# Show plot
if show:
plt.show(block=False)
plt.pause(0.01)
# Return created file
if file:
return file
return None
[docs]
def plot_signal_composition(
time_domain: Tuple[List[float], List[float]], frequency_domain: Tuple[List[float], List[float]],
significant_amplitudes: List[Tuple[float, float, float]], log_scale_frequency: bool = False,
file: Optional[Union[Path, str]] = None, show: bool = True) -> Optional[Path]:
"""
This function creates a 3D plot of the signal composition.
Input:
- time_domain (tuple of 2 lists of floats): Tuple containing time stamps and amplitudes of the time signal.
- frequency_domain (tuple of 2 lists of floats): Tuple containing frequencies and amplitudes of the frequency
domain signal.
- significant_amplitudes (list of tuples with 3 floats): List of tuples containing significant amplitudes,
frequencies, and phase angles.
- log_scale_frequency (bool): Select to set the frequency axis to a logarithmic scale. Default value is False.
- file (Path): Provide the full path and filename for the image of the graph to be created. The image is a
png-file, if the suffix is omitted it will be added. An error will occur if the picture format is set
to a non-supported format. It is also possible to provide the folder only, the title is used as filename.
- show (bool): Select to plot the graph in the environment when running the script. Default value True.
Output:
- Returns the path of the file if the plot was selected to be saved to a file, otherwise returns None.
"""
# Input data
time_stamps, time_amplitudes = time_domain
frequencies, freq_amplitudes = frequency_domain
# Prepare the environment
plt.close()
# Set the figure size
fig = plt.figure(figsize=(12, 8))
# Plot the data
ax = fig.add_axes((-0.05, 0.0, 1, .95), projection='3d')
# Plot each significant mono-frequency sinusoidal signal
for amplitude, frequency, phase_angle in significant_amplitudes:
mono_freq_signal = amplitude * np.sin(2 * np.pi * frequency * np.array(time_stamps) + phase_angle)
ax.plot(
time_stamps, [frequency] * len(time_stamps), mono_freq_signal, label=f'Freq: {frequency:.2f} Hz',
linewidth=0.8)
# Plot the original TimeDomainData on the 0 frequency plane
ax.plot(
time_stamps, [0] * len(time_stamps), time_amplitudes, label='Original Signal', color='blue', linewidth=0.5)
# Filter the frequency domain data to include only frequencies up to max_significant_frequency
max_significant_frequency = max(frequency for _, frequency, _ in significant_amplitudes)
max_index = next(i for i, freq in enumerate(frequencies) if freq > max_significant_frequency)
extended_index = min(max_index + 5, len(frequencies))
filtered_frequencies = frequencies[:extended_index]
filtered_amplitudes = freq_amplitudes[:extended_index]
# Plot the filtered frequency domain data on the last time coordinate
ax.plot(
[time_stamps[-1]] * len(filtered_frequencies), filtered_frequencies, filtered_amplitudes,
label='Frequency Domain', color='red', linewidth=1)
# Set the frequency axes limits
ax.set_xlim(0, max(time_stamps))
ax.set_ylim(0.01 if log_scale_frequency else 0, max(filtered_frequencies))
ax.set_zlim(min(time_amplitudes), max(time_amplitudes))
ax.grid(True)
ax.tick_params(axis='both', which='major', labelsize=10)
ax.set_title('Signal Composition', fontsize=15)
ax.set_xlabel('Time', fontsize=12, labelpad=12)
ax.set_ylabel('Frequency', fontsize=12)
ax.set_zlabel('Amplitude', fontsize=12)
ax.legend(loc='upper right', fontsize=10)
# Adjust the aspect ratio
ax.get_proj = lambda: np.dot(Axes3D.get_proj(ax), np.diag([1, .5, .5, .8]))
if log_scale_frequency:
ax.set_yscale('log')
# Adjust layout to reduce blank borders and extra space
fig.subplots_adjust(left=0.0, right=1, top=0.95, bottom=0.0)
# Save figure
if file:
file = handle_plot_saving(file=file, title='Signal Composition')
# Show plot
if show:
plt.show(block=False)
plt.pause(0.01)
# Return created file
if file:
return file
return None
[docs]
def bar_plot_function(
x_values: List[float], y_values: List[List[float]], labels: List[str], title: str, x_label: str,
y_label: str, log_scale_y: bool = False, x_lim: Optional[List[float]] = None,
y_lim: Optional[List[float]] = None, file: Optional[Union[Path, str]] = None, show: bool = True,
bar_width: float = 1, show_bar_values: bool = True, decimals_displayed: int = 2) -> Optional[Path]:
"""
Generic function for plotting in the Signal Processing Tool for grouped bar graph with given x and y values.
Input:
- x_values (list of lists of floats): The x-coordinates of the bars.
- y_values (list of lists of floats): The heights of the bars for each group.
- labels (list of str): The labels for each group of bars.
- title (str): The title of the plot.
- x_label (str): The label for the x-axis.
- y_label (str): The label for the y-axis.
- log_scale_y (bool): Select to set the y-axis to a logarithmic scale. Default value is False.
- x_lim (list of float): Optional list to set limits [bottom limit, top limit] for the x-axis. Default value is
None.
- y_lim (list of float): Optional list to set limits [bottom limit, top limit] for the y-axis. Default value is
None.
- file (Path): Provide the full path and filename for the image of the graph to be created. The image is a
png-file, if the suffix is omitted it will be added. An error will occur if the picture format is set
to a non-supported format. It is also possible to provide the folder only, the title is used as filename.
- show (bool): Select to plot the graph in the environment when running the script. Default value True.
- bar_width (float): The width of the bars. Default value is 1.
- show_bar_values (bool): Select to display the labels on the bars. Default value is True.
- decimals_displayed (int): The number of decimal places for formatting values. Default value is 2.
Output:
- Returns the path of the file if the plot was selected to be saved to a file, otherwise returns None.
"""
# Prepare the environment
plt.close()
# Set the figure size
plt.figure(figsize=(10, 6))
# Plot the data
num_bars = len(x_values)
indices = np.arange(num_bars)
for y_vals, group_label in zip(y_values, labels):
bars = plt.bar(indices, y_vals, width=bar_width, label=group_label, alpha=0.8, edgecolor='k')
if show_bar_values:
for bar in bars:
x_loc = bar.get_x() + bar.get_width() / 2
if x_lim is None or (x_lim[0] <= x_values[int(x_loc)] <= x_lim[1]):
y_loc = bar.get_height()
plt.text(x_loc, y_loc, f'{y_loc:.{decimals_displayed}f}', ha='center', va='bottom')
# Set the axes labels and title
plt.xlabel(x_label)
plt.ylabel(y_label)
plt.title(title)
# Set the axes to logarithmic scale
if log_scale_y:
plt.yscale('log')
# Format the x-ticks
formatted_x_values = [f'{x:.1f}' for x in x_values]
plt.xticks(indices, formatted_x_values, rotation=45, ha='right', fontsize=8)
# Set the axes limits
if x_lim:
x_lim_indices = \
[np.argmin(np.abs(np.array(x_values) - x_lim[0])), np.argmin(np.abs(np.array(x_values) - x_lim[1]))]
plt.xlim(indices[x_lim_indices[0]] - bar_width / 2, indices[x_lim_indices[1]] + bar_width / 2)
if y_lim:
plt.ylim(*y_lim)
# Create a legend
plt.legend(loc='upper left', bbox_to_anchor=(1, 1), fontsize=8)
plt.grid(visible=True, which='both', axis='y')
plt.tight_layout()
# Save figure
if file:
file = handle_plot_saving(file=file, title=title)
# Show plot
if show:
plt.show(block=False)
plt.pause(0.01)
# Return created file
if file:
return file
return None
### ===================================================================================================================
### 4. End of Script
### ===================================================================================================================