### ===================================================================================================================
### Contour Plot Class
### ===================================================================================================================
# Copyright ©2026 Haskoning Nederland B.V.
### ===================================================================================================================
### 1. Import modules
### ===================================================================================================================
# General imports
import os
import numpy as np
from pathlib import Path
from collections import defaultdict
from typing import Tuple, Optional, Union, Literal, List, Sequence
# noinspection SpellCheckingInspection
import matplotlib.lines as mlines
from matplotlib.text import Text
from matplotlib.axes import Axes
from matplotlib.colors import ListedColormap, BoundaryNorm, LinearSegmentedColormap, Normalize, TwoSlopeNorm, Colormap
from mpl_toolkits.axes_grid1 import make_axes_locatable
from matplotlib.transforms import Bbox
# References for functions and classes in the haskoning_atr_tools package
from haskoning_atr_tools.utils import ATRBasePlottingMixin, _get_pyplot
from haskoning_atr_tools.vibration_contour_plot.mesh import Mesh, MeshNode
from haskoning_atr_tools.vibration_contour_plot.config import ATRConfigVibrationContourPlot
from haskoning_atr_tools.vibration_contour_plot.simulation.vibration_simulation import VibrationSource
from haskoning_atr_tools.vibration_contour_plot.target import Target
# Import module plotly for option to plot in plotly
try:
import plotly.graph_objects as go
import plotly.io as pio
plotly_use = True
except ImportError:
go = None
pio = None
plotly_use = False
### ===================================================================================================================
### 2. ContourPlot class
### ===================================================================================================================
[docs]
class ContourPlot(ATRBasePlottingMixin):
""" Utility for rendering filled contour maps of vibration levels on a triangular mesh. This class offers both
Matplotlib-based static plotting (`plot`) and an optional Plotly-based interactive variant (`plot_plotly`)."""
def __init__(self, mesh: Mesh, figure_size: Tuple[float, float], file_name: str, values: List[float]):
"""
Input:
- mesh (Mesh): The mesh containing nodes and triangular elements.
- figure_size (tuple of two floats): Matplotlib figure size in inches (width, height).
- file_name (str): Base file name used for titles and default output paths.
- values (list of floats): Scalar values per mesh node (same order as `nodes`). Negative infinity values in
`z` are replaced by the minimum finite value (or 0 if none). # TODO nodes?
"""
self.mesh = mesh
self.figure_size = figure_size
self.file_name = file_name
self.values = values
self.project = None
@property
def project(self):
return self.__project
@project.setter
def project(self, new_project):
self.__project = new_project
@property
def mesh(self):
""" Mesh on which contours are computed and drawn."""
return self.__mesh
@mesh.setter
def mesh(self, new_mesh):
self.__mesh = new_mesh
@property
def figure_size(self) -> Tuple[float, float]:
""" Matplotlib figure size in inches (width, height)."""
return self.__figure_size
@figure_size.setter
def figure_size(self, new_figure_size: Tuple[float, float]):
self.__figure_size = new_figure_size
@property
def values(self) -> np.ndarray:
""" List with data values for the plot."""
return self.__values
@values.setter
def values(self, new_values: List[float]):
self.__values = self.z_values(z=new_values)
[docs]
@staticmethod
def z_values(z: List[float]) -> np.ndarray:
"""
Static method to prepare values for usage in the contour plot.
Input:
- z (list of float): Scalar values per mesh node (same order as `nodes`). Negative infinity values in `z`
are replaced by the minimum finite value (or 0 if none).
Output:
- Returns the cleaned values in Numpy array.
"""
# Prepare z values: replace -inf with min finite (or 0 if all -inf
z = np.array(z, dtype=float)
mask_neg_inf = np.isneginf(z)
non_inf_z = z[~mask_neg_inf]
if non_inf_z.size > 0:
min_non_inf = np.min(non_inf_z)
z[mask_neg_inf] = min_non_inf
else:
# Handle case where all values are -inf
z[mask_neg_inf] = 0.0
return z
[docs]
@staticmethod
def continuous_colour_map(
levels: Optional[Sequence[float]] = None, v_min: Optional[float] = None, v_max: Optional[float] = None,
center: Optional[float] = None, bad: str = 'lightgray', under: str = 'white', over: str = 'black',
N: int = 256, ascending_levels: bool = True) -> Tuple[Colormap, Normalize]:
"""
Static method to build a continuous colour map and corresponding normalisation.
Input:
- levels (list of float) Optional input for sequence of contour levels to derive v-min and v-max and shape
the normalisation.
- v_min, v_max (float): Optional input for explicit range boundaries to be used when `levels`
is not provided.
- center (float): Optional input to use a two-slope normalisation. Default value is None.
- bad, under, over (str): Special colours for NaN/under/over-range values. Defaults are 'lightgrey' for bad
values, 'white' for under-range values and 'black' for over-range values.
- N (int): Number of color entries in the continuous map.
- ascending_levels (bool): Select to reverse the palette, so lower levels map to lower-intensity colours.
Default value is True.
Output:
- Returns tuple with two items:
* Continuous colour map derived from palette in config-file.
* Normalisation to apply to map data values to colours. If None, Matplotlib infers it.
"""
palette = ATRConfigVibrationContourPlot.PALETTE
# For ascending levels, reverse palette so low values map to 'low' colours consistently.
seq = palette[::-1] if ascending_levels else palette[:]
# Construct the continuous colour map
cmap = LinearSegmentedColormap.from_list('continuous_cmap', seq, N=N)
# Assign special colours for masked/out-of-range
cmap.set_bad(bad, 0.8)
cmap.set_under(under)
cmap.set_over(over)
# Build normalisation
if center is not None:
if levels is not None:
_v_min = float(min(levels))
_v_max = float(max(levels))
else:
if v_min is None or v_max is None:
raise ValueError("ERROR: Provide v-min and v-max when levels is None and center is used.")
_v_min, _v_max = float(v_min), float(v_max)
norm = TwoSlopeNorm(vmin=_v_min, vcenter=float(center), vmax=_v_max)
else:
if levels is not None:
norm = Normalize(vmin=float(min(levels)), vmax=float(max(levels)))
elif v_min is not None and v_max is not None:
norm = Normalize(vmin=float(v_min), vmax=float(v_max))
else:
norm = None # Let Matplotlib infer from data
return cmap, norm
[docs]
@staticmethod
def discrete_colour_map(levels: Sequence[float]) -> Tuple[Colormap, Normalize]:
"""
Static method to build a discrete colour map with bins defined by 'levels'.
Input:
- levels (list of float): List with the bin boundaries. For N intervals there are N-1 bins.
Output:
- Returns tuple with two items:
* Discrete colour map derived from palette in config-file.
* Boundary normalisation that maps values into the discrete bins.
"""
# For N intervals there are N-1 gaps, one color per interval, ensure len(palette) == len(levels) - 1
palette = ATRConfigVibrationContourPlot.PALETTE
if len(palette) >= len(levels) - 1:
new_palette = palette[: len(levels) - 1]
new_palette.reverse()
elif len(palette) < len(levels) - 1:
raise ValueError(f"ERROR: Not enough colours in palette ({len(palette)}) for {len(levels) - 1} levels.")
else:
new_palette = palette
discrete_cmap = ListedColormap(new_palette, name='discrete_cmap')
norm = BoundaryNorm(levels, ncolors=discrete_cmap.N, clip=False)
return discrete_cmap, norm
[docs]
def plot(
self, nodes: List[MeshNode], triangles: List[List[int]], n_levels: int = 10,
line_segments: Optional[List[List[List[float]]]] = None,
vibration_sources: Optional[List[VibrationSource]] = None, sources_colour: str = 'blue',
targets: Optional[List[Target]] = None, targets_colour: str = 'blue',
plot_origin: bool = True, show_mesh: bool = True, label_font_size: int = 8,
colour_map_type: Literal['continuous', 'discrete'] = 'discrete', min_val: Optional[float] = None,
max_val: Optional[float] = None, governing_total_level_dict: Optional[dict] = None,
colour_map: Optional = None, show: bool = False, file_path: Optional[Union[str, Path, bool]] = True) \
-> Tuple[Axes, Optional[Path]]:
"""
Render a filled contour plot using Matplotlib with optional annotations and overlays. The user can select to
show and/or save to file.
Input:
- nodes (list of MeshNode): Mesh nodes.
- triangles (list of list of int): Triangulation indices (triplets of node indices).
- n_levels (int): Number of contour levels. Default value 10.
- line_segments (list of list of list of float): Optional input for additional lines to be drawn in red.
List of 2-point segments, each segment: [[x0, y0], [x1, y1]]. Default value is None.
- vibration_sources (list of VibrationSource). Optional input to add source markers and labels at source
node coordinates to contour plot. Default value is None.
- sources_colour (str): Optional input for the colour of the source markers (if provided). Default value is
'blue'.
- targets (list of Target). Optional input to add target markers and labels at target node coordinates to
contour plot. Default value is None.
- targets_colour (str): Optional input for the colour of the target markers (if provided). Default value is
'blue'.
- plot_origin (bool): Select to draw a marker at the origin. Default value is True.
- show_mesh (bool): Draw mesh triangles as light wireframe overlays. Default value is True.
- label_font_size (int): Font size for labels (sources, targets and origin). Default value is 14.
- colour_map_type (str): Select from 'continuous' or 'discrete' for the type of colour map to be used in the
contour plot. Default value is 'discrete'.
- min_val, max_val (float): Optional override of automatic min and max from the data `z`. Default value is
None.
- governing_total_level_dict (dict): If provided, used to annotate governing frequency per node/target.
Default value is None.
- colour_map (matplotlib colormap): If provided, overrides the internally constructed colour map. Default
value is None.
- show (bool): Select to display the figure window in GUI. Default value is False.
- file_path (str, Path or bool): Optional input to save the figure as file. Default value is True, in which
case the figure is saved to the default location. Input of False or None will not save the figure. String
or Path input wille save to the specified path. `file_path=True` uses `<file_name>_contour.png` in
`project.results_folder` when available.
Output:
- Returns tuple with two items:
* Axes instance of the matplotlib plot.
* Returns the path of the saved figure, if selected.
"""
def avoid_text_overlap(
_ax: Axes, texts: List[Text], pad_points: float = 2, max_iters: int = 100, step: float = 1.0) -> None:
"""
Iteratively nudge text objects to reduce overlaps (display-coordinate heuristic).
Input:
- _ax (Axes): Instance of matplotlib Axes containing the text objects.
- texts (Text): List of instances of matplotlib Text objects to adjust in this method.
- pad_points (float): Padding in display points around each text bbox.
- max_iters (int): Maximum number of nudging iterations.
- step (float): Pixel step per iteration for the separating vector.
Output:
- Uses bounding boxes in display coordinates; updates figure canvas during iterations.
- Adjusts text positions only; arrows/annotations linked to text are not moved.
- None returned.
"""
_fig = _ax.figure
# Ensure positions are up to date
_fig.canvas.draw()
# Helper: get bbox in display coords, slightly inflated for padding
def bbox_with_pad(txt):
bb = txt.get_window_extent(_fig.canvas.get_renderer())
pad = pad_points
return Bbox.from_extents(bb.x0 - pad, bb.y0 - pad, bb.x1 + pad, bb.y1 + pad)
# Convert (x, y) data positions to display coords for directional nudging
trans = _ax.transData
for _ in range(max_iters):
moved = False
bboxes = [bbox_with_pad(text) for text in texts]
for i in range(len(texts)):
for j in range(i + 1, len(texts)):
bb_i = bboxes[i]
bb_j = bboxes[j]
if bb_i.overlaps(bb_j):
# Compute a simple separating vector
dx = (bb_i.x1 + bb_i.x0) / 2 - (bb_j.x1 + bb_j.x0) / 2
dy = (bb_i.y1 + bb_i.y0) / 2 - (bb_j.y1 + bb_j.y0) / 2
# If centers coincide, pick a random direction
if dx == 0 and dy == 0:
dx, dy = np.random.randn(2)
# Normalize and move both texts in opposite directions
norm_value = np.hypot(dx, dy)
ux, uy = dx / norm_value, dy / norm_value
# Move in display coords, then invert to data coords
# Current positions in data:
x_i, y_i = texts[i].get_position()
x_j, y_j = texts[j].get_position()
# Map to display:
Xi, Yi = trans.transform([x_i, y_i])
Xj, Yj = trans.transform([x_j, y_j])
# Nudge
Xi += ux * step
Yi += uy * step
Xj -= ux * step
Yj -= uy * step
# Back to data coords:
x_i2, y_i2 = trans.inverted().transform([Xi, Yi])
x_j2, y_j2 = trans.inverted().transform([Xj, Yj])
texts[i].set_position((x_i2, y_i2))
texts[j].set_position((x_j2, y_j2))
moved = True
if not moved:
break
_fig.canvas.draw() # update bbox for next iteration
# Setup plotting backend
plt = _get_pyplot(prefer_gui=show)
# Derive min/max values if not provided
if min_val is None:
min_val = np.min(self.values)
if max_val is None:
max_val = np.max(self.values)
# Generate levels and colour map
levels = np.linspace(min_val, max_val, n_levels)
norm = None
if not colour_map:
if colour_map_type == 'continuous':
colour_map, norm = self.continuous_colour_map(levels=levels)
elif colour_map_type == 'discrete':
colour_map, norm = self.discrete_colour_map(levels=levels)
else:
raise ValueError(
f"ERROR: Invalid colour-map-type: {colour_map_type}. Select from 'discrete' or 'continuous'.")
# Create figure and main tri-contour plot
fig, ax = plt.subplots(figsize=self.figure_size, constrained_layout=True)
contour = plt.tricontourf(
np.array([node.x for node in nodes]), np.array([node.y for node in nodes]), np.array(triangles),
self.values, levels=levels, cmap=colour_map, norm=norm)
ax.set_aspect('equal', adjustable='box')
# Colour bar with dynamic sizing relative to plot width
fig_width = fig.get_figwidth()
plot_width_inches = fig_width * ax.get_position().width
size_inches = plot_width_inches * 0.05 # 5% of plot width
final_size = max(size_inches, 1) # At least 1 inch
size_fraction = final_size / fig_width # Convert to fraction of figure width
divider = make_axes_locatable(ax)
cax = divider.append_axes('right', size=f'{size_fraction * 100}%', pad=0.4)
# Create colorbar in aligned axes
cbar = fig.colorbar(contour, cax=cax, label='Vibration Level [VdB]', ticks=levels)
cbar.ax.set_yticklabels([f'{level:.0f}' for level in levels])
# Optional mesh triangles overlay
if show_mesh:
for element in self.mesh.mesh_elements:
pts = []
for node in element.mesh_nodes:
pts.append([node.x, node.y])
polygon = plt.Polygon(pts, edgecolor='gray', facecolor='none', linewidth=0.3, alpha=0.5)
ax.add_patch(polygon)
# Optional line segments overlay
if line_segments is not None:
for seg in line_segments:
x_vals = [seg[0][0], seg[1][0]]
y_vals = [seg[0][1], seg[1][1]]
ax.plot(x_vals, y_vals, color='red', linewidth=1)
# Get mesh-dimensions for scaling offsets of text-labels
dims = self.mesh.dimensions()
y_dimension = dims[1][-1] - dims[1][0]
label_shift = y_dimension * 0.015
# Plot sources and labels
labels = []
if vibration_sources:
sources_by_node = defaultdict(list)
for source in vibration_sources:
mesh_node_id = source.mesh_node.id
sources_by_node[mesh_node_id].append(source)
for mesh_node_id, sources in sources_by_node.items():
x, y = sources[0].mesh_node.x, sources[0].mesh_node.y
ax.scatter(x, y, color=sources_colour, marker='o', s=40)
if governing_total_level_dict:
label_text = f"\nf_gov={list(governing_total_level_dict[mesh_node_id].keys())[0]} Hz"
else:
label_text = f"[{', '.join(str(s.id) for s in sources)}]"
label = ax.text(
x, y + label_shift, label_text, fontsize=label_font_size, ha='center', va='bottom')
labels.append(label)
# Plot targets and labels
if targets:
for target in targets:
x = target.coordinates[0]
y = target.coordinates[1]
mesh_node_id = self.mesh.get_mesh_node_by_coordinates_2d([x, y]).id
label_text = f"{target.name}" if target.name else f"T-{target.id}"
if governing_total_level_dict:
label_text = f"{target.name} \n f_gov={list(governing_total_level_dict[mesh_node_id].keys())[0]} Hz"
else:
label_text = f"{target.name}"
ax.scatter(x, y, color=targets_colour, marker='x', s=40)
label = ax.text(
x, y - label_shift, label_text, fontsize=label_font_size, ha='center', va='top')
labels.append(label)
# Post-process label overlap
avoid_text_overlap(ax, labels, pad_points=2, max_iters=80, step=1.5)
# Origin marker
if plot_origin:
ax.scatter(0, 0, color='blue', marker='+', s=40, label='Origin')
# Title stamp in axes coordinates
ax.text(
0.05, 0.95, f'{self.file_name}', transform=ax.transAxes, fontsize=18, fontweight='bold', ha='left',
va='top', bbox=dict(boxstyle='round,pad=0.3', fc='white', alpha=0.7))
# Consistent formatting via mixin helper
self._finalise_axes(
ax=ax, title=f'{self.file_name}', x_label='X-coordinate [m]', y_label='Y-coordinate [m]', grid=False,
tight_layout=False)
# Legend (custom handles)
legend_handles = [
mlines.Line2D([], [], color=sources_colour, marker='o', linestyle='None', markersize=8, label='Sources'),
mlines.Line2D([], [], color=targets_colour, marker='x', linestyle='None', markersize=8, label='Targets')]
if plot_origin:
legend_handles.append(
mlines.Line2D([], [], color='blue', marker='+', linestyle='None', markersize=8, label='Origin'))
ax.legend(
handles=legend_handles, loc='upper right', bbox_to_anchor=(0.985, 0.985), bbox_transform=ax.transAxes,
borderaxespad=0.0, frameon=True, framealpha=0.85, fancybox=True)
# Determine effective save path from file_path
effective_file_path = file_path
if isinstance(file_path, bool) and file_path:
# Use project.results_folder if available, else use mixin default behaviour
default_name = f'{self.file_name}_contour'
default_suffix = '.png'
if self.project and hasattr(self.project, 'results_folder') and self.project.results_folder:
os.makedirs(self.project.results_folder, exist_ok=True)
effective_file_path = Path(self.project.results_folder) / f"{default_name}{default_suffix}"
else:
# Delegate to mixin's default naming
effective_file_path = True
# Save (if requested) using mixin helper
saved_file_path = self._save_if_requested(
ax=ax, file_path=effective_file_path, dpi=500, default_name=f'{self.file_name}_contour',
default_suffix='.png')
# Show (if requested)
if show:
plt.show()
plt.close(fig)
return ax, saved_file_path
[docs]
def plot_plotly(
self, nodes: List[MeshNode], n_levels: int = 10, colour_map_name: str = 'Plasma',
line_segments: Optional[List[List[List[float]]]] = None,
vibration_sources: Optional[List[VibrationSource]] = None, min_val: Optional[float] = None,
max_val: Optional[float] = None) -> None:
"""
Draft interactive contour plot using Plotly. Saves an HTML file next to results.
.. warning:: Functionality in development.
Input:
- nodes (list of MeshNode): Mesh nodes.
- n_levels (int): Number of contour levels. Default value 10.
- colour_map_name (str): Name of colour map to use. Default value is 'Plasma'.
- line_segments (list of list of list of float): Optional input for additional lines to be drawn in red.
List of 2-point segments, each segment: [[x0, y0], [x1, y1]]. Default value is None.
- vibration_sources (list of VibrationSource). Optional input to add source markers and labels at source
node coordinates to contour plot. Default value is None.
- min_val, max_val (float): Optional override of automatic min and max from the data `z`. Default value is
None.
- governing_total_level_dict (dict): If provided, used to annotate governing frequency per node/target.
Default value is None.
- colour_map (matplotlib colormap): If provided, overrides the internally constructed colour map. Default
value is None.
- show (bool): Select to display the figure window in GUI. Default value is False.
- file_path (str, Path or bool): Optional input to save the figure as file. Default value is True, in which
case the figure is saved to the default location. Input of False or None will not save the figure. String
or Path input wille save to the specified path. `file_path=True` uses `<file_name>_contour.png` in
`project.results_folder` when available.
Output:
- Generates plotly figure in browser.
"""
if not plotly_use:
raise ModuleNotFoundError(
"ERROR: The third party module 'plotly' is not installed, please install before proceeding.")
# Prepare data arrays
x = np.array([node.x for node in nodes])
y = np.array([node.y for node in nodes])
# Derive min/max values if not provided
if min_val is None:
min_val = np.min(self.values)
if max_val is None:
max_val = np.max(self.values)
# Hover text per node
hover_text = [f"Value: {z_val:.2f}<br>Node ID: {node.id}" for z_val, node in zip(self.values, nodes)]
# Create contour figure
fig = go.Figure(data=[
go.Contour(
z=self.values, x=x, y=y, colorscale=colour_map_name, ncontours=n_levels, zmin=min_val, zmax=max_val,
hoverinfo='x+y+z+text', text=hover_text, colorbar=dict(title='Vibration Level [dB]'))])
# Add vibration sources
if vibration_sources:
fig.add_trace(go.Scatter(
x=[src.mesh_node.x for src in vibration_sources], y=[src.mesh_node.y for src in vibration_sources],
mode='markers+text', marker=dict(color='black', size=8, symbol='circle'),
text=[src.id for src in vibration_sources], textposition='top center', name='Vibration Sources'))
# Add line segments
if line_segments:
for seg in line_segments:
fig.add_trace(go.Scatter(
x=[seg[0][0], seg[1][0]], y=[seg[0][1], seg[1][1]], mode='lines', line=dict(color='red', width=2),
name='Constrained Segment'))
# Origin marker
fig.add_trace(go.Scatter(
x=[0], y=[0], mode='markers+text', marker=dict(color='blue', size=10, symbol='star'), text=['Origin'],
textposition='bottom right', name='Origin'))
fig.update_layout(
title=self.file_name, xaxis_title='X-coordinate [m]', yaxis_title='Y-coordinate [m]', autosize=True)
# Save to HTML
if self.project and hasattr(self.project, 'results_folder'):
os.makedirs(self.project.results_folder, exist_ok=True)
html_path = os.path.join(self.project.results_folder, self.file_name + '.html')
else:
html_path = self.file_name + '.html'
fig.write_html(html_path)
print(f"Interactive contour plot saved to {html_path}")
### ===================================================================================================================
### 3. End of script
### ===================================================================================================================