Source code for haskoning_atr_tools.signal_processing.utils.plot

### ===================================================================================================================
###  Plotting functions for Signal Processing tool
### ===================================================================================================================
# Copyright ©2026 Haskoning Nederland B.V.

### ===================================================================================================================
###  1. Import modules
### ===================================================================================================================

# General imports
import re
import textwrap
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
import matplotlib.ticker as mticker


### ===================================================================================================================
###  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. """ def make_compact_legend(ax, lines, *, title=None, place="right", max_cols=3, min_fontsize=3, max_fontsize=9, facecolor="#F9FAFB", edgecolor="#D0D7DE", framealpha=0.95, handlelength=2.5, label_max_chars=40, prefer_first=("total", "mep limit", "nn limit"), alphabetical_rest=True,): """ Improve legend readability and style. """ # 1) Deduplicate: keep LAST style for each label by_label = {} for ln in lines: lab = ln.get_label() if lab and not lab.startswith("_nolegend_"): by_label[lab] = ln if not by_label: return None labels = list(by_label.keys()) handles = [by_label[lab] for lab in labels] # 2) Reorder: preferred labels first, then alphabetical if requested def rank(label): low = label.lower() for i, key in enumerate(prefer_first): if key in low: return (0, i, low) # top group, preserve defined order return (1, 0, low) # rest next order = sorted(range(len(labels)), key=lambda i: rank(labels[i])) if alphabetical_rest: # re-sort only the "rest" portion alphabetically after the preferred block preferred = [i for i in order if rank(labels[i])[0] == 0] rest = sorted([i for i in order if rank(labels[i])[0] == 1], key=lambda i: labels[i].lower()) order = preferred + rest labels = [labels[i] for i in order] handles = [handles[i] for i in order] # 3) Shorten very long labels to avoid wrapping disaster def shorten(s): # collapse whitespace and truncate s = re.sub(r"\s+", " ", s).strip() if len(s) > label_max_chars: return textwrap.shorten(s, width=label_max_chars, placeholder="…") return s labels = [shorten(lab) for lab in labels] # 4) Choose columns dynamically n = len(labels) # Heuristic: more entries → more columns (capped by max_cols) if n <= 10: ncol = 1 elif n <= 20: ncol = min(2, max_cols) else: ncol = max_cols # 5) Legend placement if place == "right": loc = "upper left"; bbox = (1.02, 1.0) anchor_outside = True elif place == "bottom": loc = "lower center"; bbox = (0.5, -0.02) anchor_outside = True else: # inside loc = "best"; bbox = None anchor_outside = False # 6) Font size fontsize = max(min_fontsize, min(max_fontsize, 8 if n > 12 else 9)) # 7) Create legend leg = ax.legend( handles, labels, title=title, loc=loc, bbox_to_anchor=bbox, ncol=ncol, fontsize=fontsize, frameon=True, framealpha=framealpha, facecolor=facecolor, edgecolor=edgecolor, borderpad=0.45, handlelength=handlelength, handletextpad=0.6, columnspacing=1.0, labelcolor=None, ) # 8) Title style if leg.get_title(): leg.get_title().set_fontsize(fontsize + 1) leg.get_title().set_weight("bold") # 9) Round caps for nicer samples # Version-agnostic access to legend handles handles_attr = getattr(leg, "legend_handles", None) or getattr(leg, "legendHandles", []) for h in handles_attr: if isinstance(h, Line2D): h.set_solid_capstyle('round') h.set_solid_joinstyle('round') # If legend is outside, avoid clipping when saving return leg def set_log_ticks_plain_with_label_cap(ax, max_ticks=10, base=10, tol=1e-12, rotation=0, fontsize=8): """ Keep dense log ticks (major at 1×base^n and minor at 2..9×base^n), label at most `max_ticks` of them as plain numbers, and leave the rest unlabeled. Grid lines remain for all ticks. Uses tolerance to match tick positions robustly. """ # 1) Dense locators: keep both major and minor ticks across decades major_loc = mticker.LogLocator(base=base, numticks=1000) minor_loc = mticker.LogLocator(base=base, subs=np.arange(2, 10) * 0.1, numticks=1000) ax.xaxis.set_major_locator(major_loc) ax.xaxis.set_minor_locator(minor_loc) # 2) Collect candidate ticks within current x-limits xmin, xmax = ax.get_xbound() if xmin <= 0: xmin = np.nextafter(0, 1) # guard: log axis must be > 0 def _safe_tick_values(locator): try: return np.array(locator.tick_values(xmin, xmax)) except Exception: return np.array([]) dense_major = _safe_tick_values(major_loc) dense_minor = _safe_tick_values(minor_loc) dense = np.unique(np.concatenate([dense_major, dense_minor])) dense = dense[(dense > 0) & (dense >= xmin) & (dense <= xmax)] # 3) Choose up to max_ticks to label (evenly across range) if dense.size > max_ticks: idx = np.linspace(0, dense.size - 1, num=max_ticks).round().astype(int) labeled_array = dense[idx] else: labeled_array = dense # Build a tolerance-aware membership check by storing log10 and value labeled_log = np.log10(labeled_array) labeled_vals = labeled_array # 4) Plain-number formatter that shows labels only for selected ticks (with tolerance) def _plain_log_fmt(x, pos): # Compare either in log space or linear with tolerance # (log space is more stable across decades) lx = np.log10(x) if x > 0 else -np.inf # If any labeled log is close to this tick's log, label it if np.any(np.isclose(lx, labeled_log, rtol=0, atol=tol)) or np.any( np.isclose(x, labeled_vals, rtol=0, atol=tol)): return ('%g' % x).rstrip('.') return '' # keep tick mark & grid, but hide text fmt = mticker.FuncFormatter(_plain_log_fmt) ax.xaxis.set_major_formatter(fmt) ax.xaxis.set_minor_formatter(fmt) # Ensure ticks exist for both which='major' and which='minor' ax.minorticks_on() # Make minor tick labels and marks visible (some versions auto-hide) ax.tick_params(axis='x', which='both', labelsize=fontsize, rotation=rotation) ax.tick_params(axis='x', which='minor', length=3) for lbl in ax.xaxis.get_minorticklabels(): lbl.set_visible(True) # Optional: style gridlines ax.grid(True, which='minor', alpha=0.3) ax.grid(True, which='major', alpha=0.6) # 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): color = None if 'Total' in label or "vibration limit" in label: lw = line_width * 2 ls = '-' elif 'NN limit' in label: lw = line_width * 2 ls = '--' color = 'black' else: lw = line_width ls = '--' # Pass color only if defined if color: line, = plt.plot(x_val, y_val, label=label, linewidth=lw, linestyle=ls, color=color, alpha=0.8) else: line, = plt.plot(x_val, y_val, label=label, linewidth=lw, linestyle=ls, 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.1 if min_y < 0 else min_y * 0.9, max_y * 0.9 if max_y < 0 else max_y * 1.1) 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.1 if min_y < 0 else min_y * 0.9, max_y * 0.9 if max_y < 0 else max_y * 1.1) ax = plt.gca() # Only do "plain numbers" on linear x-axis if not log_scale_x: # Option A: label every x data point all_x = np.concatenate([np.asarray(x, dtype=float) for x in x_values]) xticks = np.unique(all_x) ax.set_xticks(xticks) # Ensure plain numeric labels (no scientific notation) ax.ticklabel_format(axis='x', style='plain', useOffset=False) ax.xaxis.set_major_formatter(mticker.FormatStrFormatter('%g')) # Thin ticks: show at most ~15 ticks max_ticks = 10 if xticks.size > max_ticks: step = max(1, xticks.size // max_ticks) ax.set_xticks(xticks[::step]) else: set_log_ticks_plain_with_label_cap(ax, max_ticks=12, base=10) # 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 if len(lines) > 20: ncol = 2 min_fontsize=2 max_fontsize=4 else: ncol = 1 min_fontsize=4 max_fontsize=8 make_compact_legend( ax, lines, title=None, place="right", max_cols=ncol, min_fontsize=min_fontsize, max_fontsize=max_fontsize, ) 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 ### ===================================================================================================================