From 3ae381c20e2580b3cbac9f41d4d73be8d30dad84 Mon Sep 17 00:00:00 2001
From: Clair Mould <86794332+clmould@users.noreply.github.com>
Date: Mon, 3 Aug 2026 18:15:10 +0100
Subject: [PATCH 1/2] process/core/io/plot missing docstrings and raised
 exceptions

---
 process/core/io/plot/__init__.py        |  2 +
 process/core/io/plot/cli.py             |  2 +-
 process/core/io/plot/costs/__init__.py  |  2 +
 process/core/io/plot/costs/cli.py       |  2 +
 process/core/io/plot/images/__init__.py |  1 +
 process/core/io/plot/sankey.py          |  9 ++--
 process/core/io/plot/scans.py           |  8 +++-
 process/core/io/plot/stress_tf.py       | 12 ++++++
 process/core/io/plot/summary.py         | 57 +++++++++++++++++++------
 9 files changed, 78 insertions(+), 17 deletions(-)

diff --git a/process/core/io/plot/__init__.py b/process/core/io/plot/__init__.py
index c5cfb5c292..c732b9e61e 100644
--- a/process/core/io/plot/__init__.py
+++ b/process/core/io/plot/__init__.py
@@ -1,3 +1,5 @@
+"""Module containing plotting routines"""
+
 from process.core.io.plot.sankey import plot_sankey_plotly
 from process.core.io.plot.scans import plot_scan
 from process.core.io.plot.solutions import (
diff --git a/process/core/io/plot/cli.py b/process/core/io/plot/cli.py
index 38d70f106d..8ac687659a 100644
--- a/process/core/io/plot/cli.py
+++ b/process/core/io/plot/cli.py
@@ -1,4 +1,4 @@
-"""PROCESS cli"""
+"""PROCESS plotting cli"""
 
 from pathlib import Path
 
diff --git a/process/core/io/plot/costs/__init__.py b/process/core/io/plot/costs/__init__.py
index 17a8b988a2..6239ce5584 100644
--- a/process/core/io/plot/costs/__init__.py
+++ b/process/core/io/plot/costs/__init__.py
@@ -1,3 +1,5 @@
+"""Module containing cost plotting routines"""
+
 from process.core.io.plot.costs.costs_bar import cost_comp_1990, cost_comp_2014
 from process.core.io.plot.costs.costs_pie import cost_model_1990, cost_model_2014
 
diff --git a/process/core/io/plot/costs/cli.py b/process/core/io/plot/costs/cli.py
index 3eb6c6058e..d66bbc2488 100644
--- a/process/core/io/plot/costs/cli.py
+++ b/process/core/io/plot/costs/cli.py
@@ -1,3 +1,5 @@
+"""PROCESS costs plotting CLI"""
+
 import sys
 
 import click
diff --git a/process/core/io/plot/images/__init__.py b/process/core/io/plot/images/__init__.py
index e69de29bb2..ae61e58cce 100644
--- a/process/core/io/plot/images/__init__.py
+++ b/process/core/io/plot/images/__init__.py
@@ -0,0 +1 @@
+"""Module containing images for plotting routines"""
diff --git a/process/core/io/plot/sankey.py b/process/core/io/plot/sankey.py
index 51bc753201..0b1165e902 100644
--- a/process/core/io/plot/sankey.py
+++ b/process/core/io/plot/sankey.py
@@ -23,6 +23,7 @@
 
 
 def plot_sankey_plotly(m_file: Path):
+    """Plot the Sankey diagram"""
     if not PLOT_SANKEY:
         print(
             "\nPlotly is not installed, unable to create sankey diagram!\n"
@@ -34,6 +35,7 @@ def plot_sankey_plotly(m_file: Path):
 
 
 def power_balance_sankey(m_file: Path):
+    """Set up the power balance Sankey diagram"""
     m_file: MFile = MFile(m_file)
     p_hcd_injected_total_mw = m_file.get("p_hcd_injected_total_mw", scan=-1)
     p_plasma_ohmic_mw = m_file.get("p_plasma_ohmic_mw", scan=-1)
@@ -362,6 +364,7 @@ def power_balance_sankey(m_file: Path):
 
 
 def plotly(sankey_dict, m_file: Path):
+    """Set up the Sankey diagram"""
     fig = go.Figure(data=[sankey_dict])
 
     fig.update_layout({
@@ -404,6 +407,7 @@ def add(
         rotation: float = 0,
         **kwargs,
     ):
+        """Add a Sankey diagram"""
         __doc__ = super().__doc__  # noqa: F841
         # Here we first check if the "add" method has received arguments that
         # the Sankey class can't handle.
@@ -564,9 +568,8 @@ def minimise_dxdy(x_opt):
         return result.x
 
 
-def plot_sankey(
-    mfilename=Path("MFILE.DAT"), format_: str = "pdf"
-):  # Plot simplified power flow Sankey Diagram
+def plot_sankey(mfilename=Path("MFILE.DAT"), format_: str = "pdf"):
+    """Plot simplified power flow Sankey Diagram"""
     # Pulling values from the MFILE
     mfilename = Path(mfilename)
     m_file = MFile(mfilename)
diff --git a/process/core/io/plot/scans.py b/process/core/io/plot/scans.py
index 054cc38ba6..ac89dfb7d2 100644
--- a/process/core/io/plot/scans.py
+++ b/process/core/io/plot/scans.py
@@ -191,7 +191,13 @@ def plot_scan(
     twod_contour: bool = False,
     stack_plots: bool = False,
 ):
-    """Main plot scans script."""
+    """Main plot scans script.
+
+    Raises
+    ------
+    ValueError
+        If stack_plots does not have at least two output variables
+    """
     outputdir = outputdir or Path.cwd()
     input_files = mfiles if isinstance(mfiles, Iterable) else [mfiles]
 
diff --git a/process/core/io/plot/stress_tf.py b/process/core/io/plot/stress_tf.py
index f14efdde11..72fbe217bd 100644
--- a/process/core/io/plot/stress_tf.py
+++ b/process/core/io/plot/stress_tf.py
@@ -16,6 +16,11 @@
 
 @dataclass
 class StressPlotConfig:
+    """
+    Dataclass containing TF stress plot
+    configuration parameters
+    """
+
     axis_font_size: float
     axis_tick_size: int = 16
     legend_size: int = 12
@@ -24,6 +29,7 @@ class StressPlotConfig:
     outdir: Path | None = None
 
     def __post_init__(self):
+        """Set the output directory"""
         if self.outdir is None:
             self.outdir = Path.cwd()
         if not Path(self.outdir).is_dir():
@@ -38,6 +44,7 @@ def plot_stress(
     input_file,
     plot_conf: StressPlotConfig | dict | None = None,
 ):
+    """Create the relevant stress plots"""
     if plot_conf is None:
         plot_conf = StressPlotConfig(axis_font_size)
     elif isinstance(plot_conf, dict):
@@ -300,6 +307,7 @@ def terminal_output(
     toroidal_strain,
     vertical_strain,
 ):
+    """Output TF stress details to terminal"""
     ii_ins = 0
     ii_mids = int(0.5 * float(n_radial_array_layer))
     ii_outs = n_radial_array_layer - 1
@@ -385,6 +393,7 @@ def stress_summary(
     save_format,
     plot_conf,
 ):
+    """Plot structural stress"""
     lw = plot_conf.line_width
     ms = plot_conf.mark_size
     ats = plot_conf.axis_tick_size
@@ -458,6 +467,7 @@ def smeared_stress_summary(
     save_format,
     plot_conf,
 ):
+    """Plot smeared stress"""
     lw = plot_conf.line_width
     ms = plot_conf.mark_size
     ats = plot_conf.axis_tick_size
@@ -547,6 +557,7 @@ def strain_summary(
     save_format,
     plot_conf,
 ):
+    """Plot TF strain"""
     lw = plot_conf.line_width
     ms = plot_conf.mark_size
     ats = plot_conf.axis_tick_size
@@ -598,6 +609,7 @@ def strain_summary(
 
 
 def displacement(n_layers, radius, radial_displacement, save_format, plot_conf):
+    """Plot the radial displacement"""
     plt.plot(radius[0], radial_displacement[0], color="dodgerblue")
     for ii in range(1, n_layers):
         plt.plot(radius[ii], radial_displacement[ii], color="dodgerblue")
diff --git a/process/core/io/plot/summary.py b/process/core/io/plot/summary.py
index e92d96ff09..135ce418b8 100644
--- a/process/core/io/plot/summary.py
+++ b/process/core/io/plot/summary.py
@@ -94,6 +94,8 @@
 
 @dataclass
 class RadialBuild:
+    """Dataclass containing radial build dictionaries"""
+
     upper: dict[str, float]
     lower: dict[str, float]
     radial: dict[str, float]
@@ -3798,6 +3800,7 @@ def calc_xy(rt, e=e):
 
 
 def TF_outboard(axis: plt.Axes, item, n_tf_coils, r3, r4, w, facecolor):
+    """Plot outboard TF coils"""
     spacing = 2 * np.pi / n_tf_coils
     ang = item * spacing
     dx = w * np.sin(ang)
@@ -4495,6 +4498,7 @@ def read_imprad_data(_skiprows, data_path):
 
 
 def profiles_with_pedestal(mfile, scan: int):
+    """Calculate profiles with pedestal"""
     alphan = mfile.get("alphan", scan=scan)
     alphat = mfile.get("alphat", scan=scan)
     nd_plasma_electron_on_axis = mfile.get("nd_plasma_electron_on_axis", scan=scan)
@@ -5270,6 +5274,7 @@ def plot_blanket(
 
 
 def plot_first_wall_top_down_cross_section(axis: plt.Axes, mfile: MFile, scan: int):
+    """Plot first wall top down cross-section"""
     # Import required variables
     radius_fw_channel = mfile.get("radius_fw_channel", scan=scan) * 100
     dr_fw_wall = mfile.get("dr_fw_wall", scan=scan) * 100
@@ -5381,6 +5386,7 @@ def plot_first_wall_top_down_cross_section(axis: plt.Axes, mfile: MFile, scan: i
 
 
 def plot_first_wall_poloidal_cross_section(axis: plt.Axes, mfile: MFile, scan: int):
+    """Plot first wall poloidal cross-section"""
     # Import required variables
     radius_fw_channel = mfile.get("radius_fw_channel", scan=scan)
     dr_fw_wall = mfile.get("dr_fw_wall", scan=scan)
@@ -6842,6 +6848,7 @@ def plot_resistive_tf_wp(axis: plt.Axes, mfile: MFile, scan: int, fig):
 
 
 def plot_resistive_tf_info(axis: plt.Axes, mfile: MFile, scan: int, fig):
+    """Plot info about the resistive TF coils"""
     # Add info about the steel casing surrounding the WP
     textstr_casing = (
         f"$\\mathbf{{Casing:}}$\n \n"
@@ -10185,6 +10192,7 @@ def plot_cs_stress_time_profile(axis: plt.Axes, mfile: MFile, scan: int) -> None
 
 
 def plot_cs_turn_structure(axis: plt.Axes, fig, mfile: MFile, scan: int):
+    """Plot the CS turn structure"""
     a_cs_turn = mfile.get("a_cs_turn", scan=scan)
     dz_cs_turn = mfile.get("dz_cs_turn", scan=scan)
     dr_cs_turn = mfile.get("dr_cs_turn", scan=scan)
@@ -10291,7 +10299,7 @@ def plot_cs_turn_structure(axis: plt.Axes, fig, mfile: MFile, scan: int):
 
 
 def plot_tf_coil_structure(axis: plt.Axes, mfile: MFile, scan: int, colour_scheme=1):
-    # Plot the TF coil poloidal cross-section
+    """Plot the TF coil poloidal cross-section"""
     plot_tf_coils(axis, mfile, scan, colour_scheme)
 
     x1 = mfile.get("r_tf_arc(1)", scan=scan)
@@ -11481,7 +11489,7 @@ def plot_fw_90_deg_pipe_bend(ax, m_file, scan: int):
 
 
 def plot_fusion_rate_profiles(axis: plt.Axes, fig, mfile: MFile, scan: int):
-    # Plot the fusion rate profiles on the given axis
+    """Plot the fusion rate profiles on the given axis"""
     fusrat_plasma_dt_profile = []
     fusrat_plasma_dd_triton_profile = []
     fusrat_plasma_dd_helion_profile = []
@@ -12051,7 +12059,7 @@ def plot_cover_page(
 
 
 def plot_plasma_pressure_profiles(axis: plt.Axes, mfile: MFile, scan: int):
-    # Plot the plasma pressure profiles on the given axis
+    """Plot the plasma pressure profiles on the given axis"""
     n_plasma_profile_elements = int(mfile.get("n_plasma_profile_elements", scan=scan))
 
     pres_plasma_profile = [
@@ -12308,6 +12316,7 @@ def plot_max_normalised_beta_comparison(axis: plt.Axes, mfile: MFile, scan: int)
 
 
 def plot_plasma_pressure_gradient_profiles(axis: plt.Axes, mfile: MFile, scan: int):
+    """Plot plasma pressure gradient profiles"""
     # Get the plasma pressure profiles
     n_plasma_profile_elements = int(mfile.get("n_plasma_profile_elements", scan=scan))
 
@@ -12432,6 +12441,7 @@ def plot_plasma_poloidal_pressure_contours(axis: plt.Axes, mfile: MFile, scan: i
 
 
 def interp1d_profile(profile, mfile: MFile, scan: int):
+    """Interpolate profile over a grid"""
     # Get plasma geometry and boundary
     pg = plasma_geometry(
         rmajor=mfile.get("rmajor", scan=scan),
@@ -12610,7 +12620,7 @@ def plot_hts_tape_geometry(
     dx_hts_tape_hastelloy: float,
     show_legend: bool = True,
 ):
-
+    """Plot HTS tape geometry"""
     legend_label = None if show_legend else "_nolegend_"
     # Plot a rectangular tape stack in the middle
     rect = Rectangle(
@@ -12693,7 +12703,7 @@ def plot_hts_tape_geometry(
 
 
 def plot_tf_corc_cable_summary_box(axis, fig, mfile: MFile, scan: int):
-
+    """Plot TF CORC cable summary box"""
     textstr_cable = (
         f"$\\mathbf{{CroCo \\ Cable:}}$\n \n"
         f"Cable diameter: {mfile.get('dia_tf_turn_croco_cable', scan=scan) * 1e3:,.4f} mm\n"
@@ -12737,6 +12747,7 @@ def reaction_plot_grid(
     fractions=(0.25, 0.5, 0.75),
     colours=("blue", "yellow", "red"),
 ):
+    """Plot fusion reaction rate"""
     # Mask points outside the plasma boundary (optional, but grid is inside by construction)
     # Plot filled contour
 
@@ -12800,6 +12811,7 @@ def plot_fusion_rate_contours(
     mfile: MFile,
     scan: int,
 ):
+    """Plot fusion rate contours"""
     fusrat_plasma_dt_profile = []
     fusrat_plasma_dd_triton_profile = []
     fusrat_plasma_dd_helion_profile = []
@@ -12860,8 +12872,7 @@ def plot_fusion_rate_contours(
 
 
 def plot_magnetic_fields_in_plasma(axis: plt.Axes, mfile: MFile, scan: int):
-    # Plot magnetic field profiles inside the plasma boundary
-
+    """Plot magnetic field profiles inside the plasma boundary"""
     n_plasma_profile_elements = int(mfile.get("n_plasma_profile_elements", scan=scan))
 
     # Get toroidal magnetic field profile (in Tesla)
@@ -12940,8 +12951,7 @@ def plot_magnetic_fields_in_plasma(axis: plt.Axes, mfile: MFile, scan: int):
 
 
 def plot_beta_profiles(axis: plt.Axes, mfile: MFile, scan: int):
-    # Plot the beta profiles on the given axis
-
+    """Plot the beta profiles on the given axis"""
     n_plasma_profile_elements = int(mfile.get("n_plasma_profile_elements", scan=scan))
 
     beta_plasma_toroidal_profile = [
@@ -12976,6 +12986,7 @@ def plot_beta_profiles(axis: plt.Axes, mfile: MFile, scan: int):
 
 
 def plot_plasma_outboard_toroidal_ripple_map(fig, mfile: MFile, scan: int):
+    """Plot plasma outboard toroidal ripple map"""
     r_tf_outboard_mid = mfile.get("r_tf_outboard_mid", scan=scan)
     n_tf_coils = mfile.get("n_tf_coils", scan=scan)
     rmajor = mfile.get("rmajor", scan=scan)
@@ -13290,6 +13301,7 @@ def plot_plasma_outboard_toroidal_ripple_map(fig, mfile: MFile, scan: int):
 
 
 def plot_plasma_effective_charge_profile(axis: plt.Axes, mfile: MFile, scan: int):
+    """Plot plasma effective charge profile"""
     n_plasma_profile_elements = int(mfile.get("n_plasma_profile_elements", scan=scan))
 
     n_charge_plasma_effective_vol_avg = mfile.get(
@@ -13325,6 +13337,7 @@ def plot_plasma_effective_charge_profile(axis: plt.Axes, mfile: MFile, scan: int
 
 
 def plot_ion_charge_profile(axis: plt.Axes, mfile: MFile, scan: int):
+    """Plot ion charge profile"""
     n_plasma_profile_elements = int(mfile.get("n_plasma_profile_elements", scan=scan))
 
     # find impurity densities
@@ -13394,7 +13407,7 @@ def plot_ion_charge_profile(axis: plt.Axes, mfile: MFile, scan: int):
 
 
 def plot_ebw_ecrh_coupling_graph(axis: plt.Axes, mfile: MFile, scan: int):
-    # Plot EBW and ECRH coupling efficiency graph
+    """Plot EBW and ECRH coupling efficiency graph"""
     ebw = ElectronBernstein(plasma_profile=0)
     ecrg = ElectronCyclotron(plasma_profile=0)
     b_on_axis = mfile.get("b_plasma_toroidal_on_axis", scan=scan)
@@ -14973,6 +14986,11 @@ def plot_quench_time_evolution(
     show:
         Whether to display the plot with Matplotlib. Defaults to False to avoid
         GUI backend warnings in non-interactive environments.
+
+    Raises
+    ------
+    ValueError
+        If only one set of axes is provided, instead of both or neither
     """
     figure = None
     if axes_1 is None and axes_2 is None:
@@ -15279,6 +15297,7 @@ def plot_cs_radial_stress_profile(
     j_cs: float,
     b_cs_inner: float,
 ):
+    """Plot CS radial stress profile"""
     r_cs_inner = mfile.get("r_cs_inner", scan=scan)
     r_cs_outer = mfile.get("r_cs_outer", scan=scan)
 
@@ -15320,6 +15339,7 @@ def plot_cs_hoop_stress_profile(
     j_cs: float,
     b_cs_inner: float,
 ):
+    """Plot CS hoop stress profile"""
     r_cs_inner = mfile.get("r_cs_inner", scan=scan)
     r_cs_outer = mfile.get("r_cs_outer", scan=scan)
 
@@ -15363,6 +15383,7 @@ def plot_cs_radial_stress_contour_profile(
     b_cs_inner: float,
     colorbar_axis: plt.Axes | None = None,
 ):
+    """Plot CS radial stress contour profile"""
     r_cs_inner = mfile.get("r_cs_inner", scan=scan)
     r_cs_outer = mfile.get("r_cs_outer", scan=scan)
     dz_cs_full = mfile.get("dz_cs_full", scan=scan)
@@ -15454,6 +15475,7 @@ def plot_cs_hoop_stress_contour_profile(
     b_cs_inner: float,
     colorbar_axis: plt.Axes | None = None,
 ):
+    """Plot CS hoop stress contour profile"""
     r_cs_inner = mfile.get("r_cs_inner", scan=scan)
     r_cs_outer = mfile.get("r_cs_outer", scan=scan)
     dz_cs_full = mfile.get("dz_cs_full", scan=scan)
@@ -15544,6 +15566,7 @@ def plot_cs_vertical_stress_profile(
     mfile: MFile,
     scan: int,
 ):
+    """Plot CS vertical stress profile"""
     dz_cs_full = mfile.get("dz_cs_full", scan=scan)
 
     stress_z_profile = np.array([
@@ -15576,6 +15599,7 @@ def plot_vertical_stress_contour_profile(
     scan: int,
     colorbar_axis: plt.Axes | None = None,
 ):
+    """Vertical stress contour plot"""
     dz_cs_full = mfile.get("dz_cs_full", scan=scan)
     r_cs_inner = mfile.get("r_cs_inner", scan=scan)
     r_cs_outer = mfile.get("r_cs_outer", scan=scan)
@@ -15657,6 +15681,7 @@ def plot_cs_tresca_2d_contour(
     scan: int,
     colorbar_axis: plt.Axes | None = None,
 ):
+    """CS Tresca stress contour plot"""
     dz_cs_full = mfile.get("dz_cs_full", scan=scan)
     r_cs_inner = mfile.get("r_cs_inner", scan=scan)
     r_cs_outer = mfile.get("r_cs_outer", scan=scan)
@@ -15767,6 +15792,7 @@ def plot_cs_von_mises_2d_contour(
     scan: int,
     colorbar_axis: plt.Axes | None = None,
 ):
+    """CS Von Mises stress contour plot"""
     dz_cs_full = mfile.get("dz_cs_full", scan=scan)
     r_cs_inner = mfile.get("r_cs_inner", scan=scan)
     r_cs_outer = mfile.get("r_cs_outer", scan=scan)
@@ -15918,7 +15944,13 @@ def main_plot(
     pages = {}
 
     def _add_page(name: str | None = None):
-        """Add a page to the dictionary of pages. If no name is provided, then assign the lowest unused number."""
+        """Add a page to the dictionary of pages. If no name is provided, then assign the lowest unused number.
+
+        Raises
+        ------
+        KeyError
+            If a page number has already been used
+        """
         if name is None:
             prev_index = max((int(k) for k in pages if k.isnumeric()), default=0)
             name = str(prev_index + 1)
@@ -16481,7 +16513,7 @@ def _add_page(name: str | None = None):
 
 
 def create_thickness_builds(m_file, scan: int):
-    # Build the dictionaries of radial and vertical build values and cumulative values
+    """Create the dictionaries of radial and vertical build values and cumulative values"""
     if int(m_file.get("i_single_null", scan=scan)) == 0:
         vertical_upper = [
             "z_plasma_xpoint_upper",
@@ -16557,6 +16589,7 @@ def plot_summary(
     output_format: str = "pdf",
     show: bool = False,
 ):
+    """Create the summary.pdf"""
 
     def add_page_footer(
         fig: plt.Figure, page_number: int, total_pages: int, run_label: str

From 947b603ae6df2b0a75653b2d5cf1885c7ee71d55 Mon Sep 17 00:00:00 2001
From: clmould <86794332+clmould@users.noreply.github.com>
Date: Tue, 4 Aug 2026 11:32:52 +0100
Subject: [PATCH 2/2] Apply suggestion from @timothy-nunn

Co-authored-by: Timothy <75321887+timothy-nunn@users.noreply.github.com>
---
 process/core/io/plot/stress_tf.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/process/core/io/plot/stress_tf.py b/process/core/io/plot/stress_tf.py
index 72fbe217bd..ccbfb321ef 100644
--- a/process/core/io/plot/stress_tf.py
+++ b/process/core/io/plot/stress_tf.py
@@ -16,8 +16,7 @@
 
 @dataclass
 class StressPlotConfig:
-    """
-    Dataclass containing TF stress plot
+    """Dataclass containing TF stress plot
     configuration parameters
     """