Source code for haskoning_atr_tools.signal_processing_tool.utils.plot

### ===================================================================================================================
###   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 ### ===================================================================================================================