From f72380fb2292cfe6e8032eeea4a08caee7223e0d Mon Sep 17 00:00:00 2001
From: mn3981 <christopher.ashe@ukaea.uk>
Date: Fri, 24 Apr 2026 16:24:13 +0100
Subject: [PATCH 01/13] fix: update impurity radiation class to use
 PlasmaProfile type hint and improve comments

Co-authored-by: Copilot <copilot@github.com>
---
 process/models/physics/impurity_radiation.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/process/models/physics/impurity_radiation.py b/process/models/physics/impurity_radiation.py
index f4b536313..b70bd9462 100644
--- a/process/models/physics/impurity_radiation.py
+++ b/process/models/physics/impurity_radiation.py
@@ -14,6 +14,7 @@
 from process.core.exceptions import ProcessError, ProcessValueError
 from process.core.model import DataStructure
 from process.models.physics.plasma_profiles import PlasmaProfile
+from process.models.physics.plasma_profiles import PlasmaProfile
 
 logger = logging.getLogger(__name__)
 
@@ -561,7 +562,7 @@ class ImpurityRadiation:
     """Calculates the impurity radiation losses for given temperature and
     density profiles. The considers the  total impurity radiation from the core
     (pden_impurity_core_rad_total_mw) and total impurity radiation
-    (pden_impurity_rad_total_mw) [MW/(m^3)]. The class is used to sum the impurity
+    (pden_impurity_rad_total_mw) [MW/(m³)]. The class is used to sum the impurity
     radiation loss from each impurity element to find the total impurity radiation loss.
     """
 
@@ -655,7 +656,7 @@ def integrate_radiation_loss_profiles(self):
         """Integrate the radiation loss profiles using the Simpson rule.
         Store the total values for each aspect of impurity radiation loss.
         """
-        # 2.0e-6 converts from W/m^3 to MW/m^3 and also accounts for both sides of the
+        # 2.0e-6 converts from W/m³ to MW/m³ and also accounts for both sides of the
         # plasma
         self.pden_impurity_rad_total_mw = 2.0e-6 * integrate.simpson(
             self.pden_impurity_rad_profile,

From 411a263f9f888c4c9cb8c202d23c23287da88e07 Mon Sep 17 00:00:00 2001
From: mn3981 <christopher.ashe@ukaea.uk>
Date: Fri, 24 Apr 2026 16:26:28 +0100
Subject: [PATCH 02/13] fix: add edge profile arrays to ImpurityRadiation class

---
 process/models/physics/impurity_radiation.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/process/models/physics/impurity_radiation.py b/process/models/physics/impurity_radiation.py
index b70bd9462..b150accba 100644
--- a/process/models/physics/impurity_radiation.py
+++ b/process/models/physics/impurity_radiation.py
@@ -587,9 +587,17 @@ def __init__(self, plasma_profile: PlasmaProfile, data_structure: DataStructure)
         self.pden_impurity_core_rad_profile = np.zeros(
             self.data.physics.n_plasma_profile_elements
         )
+        self.pden_impurity_rad_edge_profile = np.zeros(self.plasma_profile.profile_size)
 
         self.pden_impurity_rad_total_mw = 0.0
         self.pden_impurity_core_rad_total_mw = 0.0
+        self.pden_impurity_rad_edge_total_mw = 0.0
+
+    def run(self):
+        """This model isn't run"""
+
+    def output(self):
+        """This model has no output"""
 
     def run(self):
         """This model isn't run"""

From 9730ff7c41c4d09400b5e893c98795b2154fa256 Mon Sep 17 00:00:00 2001
From: mn3981 <christopher.ashe@ukaea.uk>
Date: Fri, 24 Apr 2026 16:45:51 +0100
Subject: [PATCH 03/13] fix: refactor pimpden function to improve type hints
 and documentation

Co-authored-by: Copilot <copilot@github.com>
---
 process/models/physics/impurity_radiation.py | 75 ++++++++++++--------
 1 file changed, 45 insertions(+), 30 deletions(-)

diff --git a/process/models/physics/impurity_radiation.py b/process/models/physics/impurity_radiation.py
index b150accba..904ee152d 100644
--- a/process/models/physics/impurity_radiation.py
+++ b/process/models/physics/impurity_radiation.py
@@ -459,34 +459,42 @@ def _zav_of_te_compiled(
     return zav_of_te
 
 
-def pimpden(imp_element_index, neprofile, teprofile, data: DataStructure):
-    """Calculates the impurity radiation density (W/m3)
+def pimpden(
+    imp_element_index: int,
+    nd_electron_profile: np.array | float,
+    temp_electron_profile_kev: np.array | float,
+) -> np.array | float:
+    """
+    Calculates the impurity radiation density (W/m³) based on the electron density and temperature profiles.
 
     Parameters
     ----------
-    imp_element_index : int
+    imp_element_index
         Impurity element index
-    neprofile : numpy.array
-        electron density profile
-    teprofile : numpy.array
-        electron temperature profile
+    nd_electron_profile
+        electron density profile [m⁻³]
+    temp_electron_profile_kev
+        electron temperature profile [keV]
 
     Returns
     -------
-    numpy.array
-        pimpden - total impurity radiation density (W/m3)
+    pden_impurity_profile - total impurity radiation density (W/m³)
+
+    Notes
+    -----
+     -Temperatures outside the range of the L(Z,Tₑ) table are handled by using the L(Z,Tₑ)
+      value at the closest temperature in the table,
     """
-    # less_than_imp_temp_mask = teprofile values less than impurity temperature.
-    # greater_than_imp_temp_mask = teprofile values higher than impurity temperature.
-    bins = data.impurity_radiation.temp_impurity_keV_array[imp_element_index]
-    indices = np.digitize(teprofile, bins)
+    bins = impurity_radiation_module.temp_impurity_keV_array[imp_element_index]
+    indices = np.digitize(temp_electron_profile_kev, bins)
     indices[indices >= bins.shape[0]] = bins.shape[0] - 1
     indices[indices < 0] = 0
 
-    # Use numpy.interp for linear interpolation in log-log space
-    pimpden = np.exp(
+    # Use numpy.interp for linear interpolation in log-log space to find the
+    # loss function values for the given temperature profile L(Z, Tₑ).
+    power_loss_function = np.exp(
         np.interp(
-            np.log(teprofile),
+            np.log(temp_electron_profile_kev),
             np.log(
                 data.impurity_radiation.temp_impurity_keV_array[imp_element_index, :]
             ),
@@ -498,37 +506,44 @@ def pimpden(imp_element_index, neprofile, teprofile, data: DataStructure):
         )
     )
 
-    pimpden = (
-        data.impurity_radiation.f_nd_impurity_electron_array[imp_element_index]
-        * neprofile
-        * neprofile
-        * pimpden
+    # W/m³ = nᵢ * nₑ * L(Z, Tₑ)
+    pden_impurity_profile = (
+        impurity_radiation_module.f_nd_impurity_electron_array[imp_element_index]
+        * nd_electron_profile
+        * nd_electron_profile
+        * power_loss_function
     )
 
+    # less_than_imp_temp_mask = temp_electron_profile_kev values less than impurity temperature.
+
     less_than_imp_temp_mask = (
-        teprofile
-        <= data.impurity_radiation.temp_impurity_keV_array[imp_element_index, 0]
+        temp_electron_profile_kev
+        <= impurity_radiation_module.temp_impurity_keV_array[imp_element_index, 0]
     )
-    pimpden[less_than_imp_temp_mask] = (
-        data.impurity_radiation.pden_impurity_lz_nd_temp_array[imp_element_index, 0]
+    # This is okay because line radiation will dominate at lower temp, and the L(Z,Tₑ)
+    # value at the lowest temperature in the table is likely to be an overestimate of the
+    # radiation loss at lower temperatures, so this is a conservative approach.
+    pden_impurity_profile[less_than_imp_temp_mask] = (
+        impurity_radiation_module.pden_impurity_lz_nd_temp_array[imp_element_index, 0]
     )
 
+    # greater_than_imp_temp_mask = temp_electron_profile_kev values higher than impurity temperature.
     greater_than_imp_temp_mask = (
-        teprofile
-        >= data.impurity_radiation.temp_impurity_keV_array[
+        temp_electron_profile_kev
+        >= impurity_radiation_module.temp_impurity_keV_array[
             imp_element_index,
             data.impurity_radiation.impurity_arr_len_tab[imp_element_index] - 1,
         ]
     )
     #  This is okay because Bremsstrahlung will dominate at higher temp.
-    pimpden[greater_than_imp_temp_mask] = (
-        data.impurity_radiation.pden_impurity_lz_nd_temp_array[
+    pden_impurity_profile[greater_than_imp_temp_mask] = (
+        impurity_radiation_module.pden_impurity_lz_nd_temp_array[
             imp_element_index,
             data.impurity_radiation.impurity_arr_len_tab[imp_element_index] - 1,
         ]
     )
 
-    return pimpden
+    return pden_impurity_profile
 
 
 def element2index(element: str, data: DataStructure):

From 52e75aa60a37394b2bbd2c8f38265fd316410002 Mon Sep 17 00:00:00 2001
From: mn3981 <christopher.ashe@ukaea.uk>
Date: Fri, 24 Apr 2026 16:48:49 +0100
Subject: [PATCH 04/13] =?UTF-8?q?=F0=9F=9A=9Arename=20`pimpden`=20function?=
 =?UTF-8?q?=20to=20`calculate=5Fimpurity=5Fradiation=5Fpower=5Fdensity`=20?=
 =?UTF-8?q?for=20clarity?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 process/models/physics/impurity_radiation.py         | 4 ++--
 tests/unit/models/physics/test_impurity_radiation.py | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/process/models/physics/impurity_radiation.py b/process/models/physics/impurity_radiation.py
index 904ee152d..0c9710965 100644
--- a/process/models/physics/impurity_radiation.py
+++ b/process/models/physics/impurity_radiation.py
@@ -459,7 +459,7 @@ def _zav_of_te_compiled(
     return zav_of_te
 
 
-def pimpden(
+def calculate_impurity_radiation_power_density(
     imp_element_index: int,
     nd_electron_profile: np.array | float,
     temp_electron_profile_kev: np.array | float,
@@ -641,7 +641,7 @@ def imprad_profile(self, imp_element_index):
             Sertoli, private communication
             Kallenbach et al., Plasma Phys. Control. Fus. 55(2013) 124041
         """
-        pimp = pimpden(
+        pimp = calculate_impurity_radiation_power_density(
             imp_element_index,
             self.plasma_profile.neprofile.profile_y,
             self.plasma_profile.teprofile.profile_y,
diff --git a/tests/unit/models/physics/test_impurity_radiation.py b/tests/unit/models/physics/test_impurity_radiation.py
index 335736faa..e337ee7b3 100644
--- a/tests/unit/models/physics/test_impurity_radiation.py
+++ b/tests/unit/models/physics/test_impurity_radiation.py
@@ -75,7 +75,7 @@ def test_pimpden(process_models):
         ]),
     )
 
-    pimpden = impurity_radiation.pimpden(
+    pimpden = impurity_radiation.calculate_impurity_radiation_power_density(
         pimden_parameters.imp_element_index,
         pimden_parameters.ne,
         pimden_parameters.te,

From 00cc98679914c4ecc804421352916043d57d687f Mon Sep 17 00:00:00 2001
From: mn3981 <christopher.ashe@ukaea.uk>
Date: Fri, 24 Apr 2026 17:07:50 +0100
Subject: [PATCH 05/13] fix: enhance impurity radiation calculations and
 improve type hints in imprad_profile method

Co-authored-by: Copilot <copilot@github.com>
---
 process/models/physics/impurity_radiation.py | 22 +++++++-------------
 1 file changed, 8 insertions(+), 14 deletions(-)

diff --git a/process/models/physics/impurity_radiation.py b/process/models/physics/impurity_radiation.py
index 0c9710965..beada211d 100644
--- a/process/models/physics/impurity_radiation.py
+++ b/process/models/physics/impurity_radiation.py
@@ -507,6 +507,7 @@ def calculate_impurity_radiation_power_density(
     )
 
     # W/m³ = nᵢ * nₑ * L(Z, Tₑ)
+    # nᵢ = f_nd_species_electron * nₑ
     pden_impurity_profile = (
         impurity_radiation_module.f_nd_impurity_electron_array[imp_element_index]
         * nd_electron_profile
@@ -624,31 +625,24 @@ def map_imprad_profile(self):
         """Map imprad_profile() over each impurity element index."""
         list(map(self.imprad_profile, self.imp))
 
-    def imprad_profile(self, imp_element_index):
+    def imprad_profile(self, imp_element_index: int) -> None:
         """Calculates the impurity radiation losses for given temperature
         and density profiles.
 
         Parameters
         ----------
-        imp_element_index : Int
+        imp_element_index
             Index used to access different impurity radiation elements
 
-        References
-        ----------
-            Bremsstrahlung equation from Johner, L(z) data (coronal equilibrium)
-            from Marco Sertoli, Asdex-U, ref. Kallenbach et al.
-            Johner, Fusion Science and Technology 59 (2011), pp 308-349
-            Sertoli, private communication
-            Kallenbach et al., Plasma Phys. Control. Fus. 55(2013) 124041
         """
-        pimp = calculate_impurity_radiation_power_density(
-            imp_element_index,
-            self.plasma_profile.neprofile.profile_y,
-            self.plasma_profile.teprofile.profile_y,
+        pden_impurity_radiation_profile = calculate_impurity_radiation_power_density(
+            imp_element_index=imp_element_index,
+            nd_electron_profile=self.plasma_profile.neprofile.profile_y,
+            temp_electron_profile_kev=self.plasma_profile.teprofile.profile_y,
             self.data,
         )
 
-        self.pimp_profile = np.add(self.pimp_profile, pimp)
+        self.pimp_profile = np.add(self.pimp_profile, pden_impurity_radiation_profile)
 
     def calculate_radiation_loss_profiles(self):
         """Calculate the Bremsstrahlung (radb), line radiation (radl), total impurity

From 1e0f3520d979c02028fab0cd44b3f9774078a7cc Mon Sep 17 00:00:00 2001
From: mn3981 <christopher.ashe@ukaea.uk>
Date: Fri, 24 Apr 2026 17:15:37 +0100
Subject: [PATCH 06/13] fix: rename `zav_of_te` function to
 `calculate_average_charge_at_temp` for clarity and update references

---
 process/models/physics/impurity_radiation.py     |  3 ++-
 process/models/physics/physics.py                | 16 +++++++++-------
 .../models/physics/test_impurity_radiation.py    |  2 +-
 3 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/process/models/physics/impurity_radiation.py b/process/models/physics/impurity_radiation.py
index beada211d..abdfd6aa3 100644
--- a/process/models/physics/impurity_radiation.py
+++ b/process/models/physics/impurity_radiation.py
@@ -396,7 +396,7 @@ def fradcore(rho, radius_plasma_core_norm, f_p_plasma_core_rad_reduction):
     return fradcore
 
 
-def zav_of_te(imp_element_index, teprofile, data: DataStructure):
+def calculate_average_charge_at_temp(imp_element_index, teprofile, data: DataStructure):
     """Calculates electron temperature dependent average atomic number
 
     Parameters
@@ -430,6 +430,7 @@ def _zav_of_te_compiled(
     impurity_arr_zav: np.array,
     impurity_arr_len_tab: np.array,
 ):
+
     bins = temp_impurity_keV_array[imp_element_index]
     indices = np.digitize(teprofile, bins)
     indices[indices >= bins.shape[0]] = bins.shape[0] - 1
diff --git a/process/models/physics/physics.py b/process/models/physics/physics.py
index 573ab73ac..d2c373f14 100644
--- a/process/models/physics/physics.py
+++ b/process/models/physics/physics.py
@@ -1173,7 +1173,7 @@ def plasma_composition(self):
         znimp = 0.0
         for imp in range(N_IMPURITIES):
             if self.data.impurity_radiation.impurity_arr_z[imp] > 2:
-                znimp += impurity_radiation.zav_of_te(
+                znimp += impurity_radiation.calculate_average_charge_at_temp(
                     imp,
                     np.array([self.data.physics.temp_plasma_electron_vol_avg_kev]),
                     self.data,
@@ -1284,7 +1284,7 @@ def plasma_composition(self):
         for imp in range(N_IMPURITIES):
             self.data.physics.n_charge_plasma_effective_vol_avg += (
                 self.data.impurity_radiation.f_nd_impurity_electron_array[imp]
-                * impurity_radiation.zav_of_te(
+                * impurity_radiation.calculate_average_charge_at_temp(
                     imp,
                     np.array([self.data.physics.temp_plasma_electron_vol_avg_kev]),
                     self.data,
@@ -1396,7 +1396,7 @@ def plasma_composition(self):
             if self.data.impurity_radiation.impurity_arr_z[imp] > 2:
                 self.data.physics.n_charge_plasma_effective_mass_weighted_vol_avg += (
                     self.data.impurity_radiation.f_nd_impurity_electron_array[imp]
-                    * impurity_radiation.zav_of_te(
+                    * impurity_radiation.calculate_average_charge_at_temp(
                         imp,
                         np.array([self.data.physics.temp_plasma_electron_vol_avg_kev]),
                         self.data,
@@ -1673,7 +1673,7 @@ def calculate_effective_charge_ionisation_profiles(self):
             for imp in range(N_IMPURITIES):
                 zeff_profile[i] += (
                     self.data.impurity_radiation.f_nd_impurity_electron_array[imp]
-                    * impurity_radiation.zav_of_te(
+                    * impurity_radiation.calculate_average_charge_at_temp(
                         imp,
                         np.array([self.plasma_profile.teprofile.profile_y[i]]),
                         self.data,
@@ -1690,9 +1690,11 @@ def calculate_effective_charge_ionisation_profiles(self):
         charge_profiles = np.zeros((n_impurities, n_points))
         for imp in range(n_impurities):
             for i in range(n_points):
-                charge_profiles[imp, i] = impurity_radiation.zav_of_te(
-                    imp, np.array([te_profile[i]]), self.data
-                ).squeeze()
+                charge_profiles[imp, i] = (
+                    impurity_radiation.calculate_average_charge_at_temp(
+                        imp, np.array([te_profile[i]]), self.data
+                    ).squeeze()
+                )
         self.data.impurity_radiation.n_charge_impurity_profile = charge_profiles
 
     def outplas(self):
diff --git a/tests/unit/models/physics/test_impurity_radiation.py b/tests/unit/models/physics/test_impurity_radiation.py
index e337ee7b3..abebe910c 100644
--- a/tests/unit/models/physics/test_impurity_radiation.py
+++ b/tests/unit/models/physics/test_impurity_radiation.py
@@ -174,7 +174,7 @@ def test_zav_of_te(process_models):
             1.00000000000001,
         ]),
     )
-    zav_of_te = impurity_radiation.zav_of_te(
+    zav_of_te = impurity_radiation.calculate_average_charge_at_temp(
         zavofteparam.imp_element_index, zavofteparam.te, process_models.physics.data
     )
 

From 6c6443b9928896b15933c05228560b8996d8c8af Mon Sep 17 00:00:00 2001
From: mn3981 <christopher.ashe@ukaea.uk>
Date: Fri, 24 Apr 2026 17:16:26 +0100
Subject: [PATCH 07/13] fix: rename `_zav_of_te_compiled` to
 `calculate_average_charge_at_temp_compiled` for clarity

---
 process/models/physics/impurity_radiation.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/process/models/physics/impurity_radiation.py b/process/models/physics/impurity_radiation.py
index abdfd6aa3..784551997 100644
--- a/process/models/physics/impurity_radiation.py
+++ b/process/models/physics/impurity_radiation.py
@@ -413,7 +413,7 @@ def calculate_average_charge_at_temp(imp_element_index, teprofile, data: DataStr
     """
     # less_than_imp_temp_mask = teprofile values less than impurity temperature.
     # greater_than_imp_temp_mask = teprofile values higher than impurity temperature.
-    return _zav_of_te_compiled(
+    return calculate_average_charge_at_temp_compiled(
         imp_element_index,
         teprofile,
         data.impurity_radiation.temp_impurity_keV_array,
@@ -423,7 +423,7 @@ def calculate_average_charge_at_temp(imp_element_index, teprofile, data: DataStr
 
 
 @njit(cache=True)
-def _zav_of_te_compiled(
+def calculate_average_charge_at_temp_compiled(
     imp_element_index: int,
     teprofile: np.array,
     temp_impurity_keV_array: np.array,

From a1cc25eecdb900f40b079f76ea0ae9c397e8eaa5 Mon Sep 17 00:00:00 2001
From: mn3981 <christopher.ashe@ukaea.uk>
Date: Fri, 24 Apr 2026 17:27:32 +0100
Subject: [PATCH 08/13] fix: update calculate_average_charge_at_temp function
 for clarity and accuracy in parameter types and documentation

Co-authored-by: Copilot <copilot@github.com>
---
 process/models/physics/impurity_radiation.py | 76 ++++++++++++++------
 1 file changed, 53 insertions(+), 23 deletions(-)

diff --git a/process/models/physics/impurity_radiation.py b/process/models/physics/impurity_radiation.py
index 784551997..4c6ad10de 100644
--- a/process/models/physics/impurity_radiation.py
+++ b/process/models/physics/impurity_radiation.py
@@ -396,68 +396,98 @@ def fradcore(rho, radius_plasma_core_norm, f_p_plasma_core_rad_reduction):
     return fradcore
 
 
-def calculate_average_charge_at_temp(imp_element_index, teprofile, data: DataStructure):
-    """Calculates electron temperature dependent average atomic number
+def calculate_average_charge_at_temp(
+    imp_element_index: int, temp_electron_kev: np.array | float
+) -> np.array | float:
+    """Calculates electron temperature dependent average atomic charge (Z) for a given impurity element.
 
     Parameters
     ----------
-    imp_element_index : int
+    imp_element_index
         Impurity element index
-    teprofile : numpy.array
-        temperature profile
+    temp_electron_kev
+        electron temperature in keV
 
     Returns
     -------
     numpy.array
-        zav_of_te - electron temperature dependent average atomic number
+        zav_of_te - electron temperature dependent average atomic charge
     """
-    # less_than_imp_temp_mask = teprofile values less than impurity temperature.
-    # greater_than_imp_temp_mask = teprofile values higher than impurity temperature.
     return calculate_average_charge_at_temp_compiled(
-        imp_element_index,
-        teprofile,
-        data.impurity_radiation.temp_impurity_keV_array,
-        data.impurity_radiation.impurity_arr_zav,
-        data.impurity_radiation.impurity_arr_len_tab,
+        imp_element_index=imp_element_index,
+        temp_electron_kev=temp_electron_kev,
+        temp_impurity_keV_array=data.impurity_radiation.temp_impurity_keV_array,
+        impurity_arr_zav=data.impurity_radiation.impurity_arr_zav,
+        impurity_arr_len_tab=data.impurity_radiation.impurity_arr_len_tab,
     )
 
 
 @njit(cache=True)
 def calculate_average_charge_at_temp_compiled(
     imp_element_index: int,
-    teprofile: np.array,
+    temp_electron_kev: np.array | float,
     temp_impurity_keV_array: np.array,
     impurity_arr_zav: np.array,
     impurity_arr_len_tab: np.array,
-):
+) -> np.array | float:
+    """Calculates electron temperature dependent average atomic charge (Z) for a given impurity element.
+
+    Parameters
+    ----------
+    imp_element_index
+        Impurity element index
+    temp_electron_kev
+        electron temperature in keV
+    temp_impurity_keV_array
+        2D array of impurity temperatures in keV for each impurity element
+    impurity_arr_zav
+        2D array of average charge values for each impurity element at the corresponding temperatures in temp_impurity_keV_array
+    impurity_arr_len_tab
+        1D array of the length of the temperature and average charge tables for each impurity element
 
+    Returns
+    -------
+    n_charge_impurity_average
+        electron temperature dependent average atomic charge of impurity element at the given temperature(s)
+
+    """
     bins = temp_impurity_keV_array[imp_element_index]
-    indices = np.digitize(teprofile, bins)
+    indices = np.digitize(temp_electron_kev, bins)
     indices[indices >= bins.shape[0]] = bins.shape[0] - 1
     indices[indices < 0] = 0
     # Use numpy.interp for linear interpolation in log space
-    zav_of_te = np.interp(
-        np.log(teprofile),
+    n_charge_impurity_average = np.interp(
+        np.log(temp_electron_kev),
         np.log(temp_impurity_keV_array[imp_element_index, :]),
         impurity_arr_zav[imp_element_index, :],
     )
 
-    less_than_imp_temp_mask = teprofile <= temp_impurity_keV_array[imp_element_index, 0]
+    # less_than_imp_temp_mask = temp_electron_profile_kev values less than impurity temperature
+    less_than_imp_temp_mask = (
+        temp_electron_kev <= temp_impurity_keV_array[imp_element_index, 0]
+    )
+
+    # Sets n_charge_impurity_average to the value at the lowest temperature in the table for temperatures below the lowest temperature in the table.
+    n_charge_impurity_average[less_than_imp_temp_mask] = impurity_arr_zav[
+        imp_element_index, 0
+    ]
 
-    zav_of_te[less_than_imp_temp_mask] = impurity_arr_zav[imp_element_index, 0]
+    # greater_than_imp_temp_mask = temp_electron_profile_kev values higher than impurity temperature.
     greater_than_imp_temp_mask = (
-        teprofile
+        temp_electron_kev
         >= temp_impurity_keV_array[
             imp_element_index,
             (impurity_arr_len_tab[imp_element_index]) - 1,
         ]
     )
-    zav_of_te[greater_than_imp_temp_mask] = impurity_arr_zav[
+
+    # Sets n_charge_impurity_average to the value at the highest temperature in the table for temperatures above the highest temperature in the table.
+    n_charge_impurity_average[greater_than_imp_temp_mask] = impurity_arr_zav[
         imp_element_index,
         impurity_arr_len_tab[imp_element_index] - 1,
     ]
 
-    return zav_of_te
+    return n_charge_impurity_average
 
 
 def calculate_impurity_radiation_power_density(

From 313972db502e5b48a60c136a0c621e164011ca54 Mon Sep 17 00:00:00 2001
From: mn3981 <christopher.ashe@ukaea.uk>
Date: Fri, 24 Apr 2026 17:31:41 +0100
Subject: [PATCH 09/13] fix: reorganize imports for clarity and improve type
 checking

Co-authored-by: Copilot <copilot@github.com>
---
 process/models/physics/impurity_radiation.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/process/models/physics/impurity_radiation.py b/process/models/physics/impurity_radiation.py
index 4c6ad10de..a9662587f 100644
--- a/process/models/physics/impurity_radiation.py
+++ b/process/models/physics/impurity_radiation.py
@@ -1,10 +1,12 @@
 """Module for impurity radiation calculations and data handling."""
 
+from __future__ import annotations
+
 import dataclasses
 import logging
 import re
 from importlib import resources
-from pathlib import Path
+from typing import TYPE_CHECKING
 
 import numpy as np
 from numba import njit
@@ -14,7 +16,11 @@
 from process.core.exceptions import ProcessError, ProcessValueError
 from process.core.model import DataStructure
 from process.models.physics.plasma_profiles import PlasmaProfile
-from process.models.physics.plasma_profiles import PlasmaProfile
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from process.models.physics.plasma_profiles import PlasmaProfile
 
 logger = logging.getLogger(__name__)
 

From da503e5080a7998099ca69aed45a15f204e50eff Mon Sep 17 00:00:00 2001
From: mn3981 <christopher.ashe@ukaea.uk>
Date: Fri, 24 Apr 2026 17:46:11 +0100
Subject: [PATCH 10/13] fix: update comments for effective charge calculations
 to use proper mathematical notation

Co-authored-by: Copilot <copilot@github.com>
---
 process/models/physics/physics.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/process/models/physics/physics.py b/process/models/physics/physics.py
index d2c373f14..543163143 100644
--- a/process/models/physics/physics.py
+++ b/process/models/physics/physics.py
@@ -1278,8 +1278,8 @@ def plasma_composition(self):
         # ======================================================================
 
         # Effective charge
-        # Calculation should be sum(ni.Zi^2) / sum(ni.Zi),
-        # but ne = sum(ni.Zi) through quasineutrality
+        # Calculation should be Σ(nᵢZᵢ²) / Σ(nᵢZᵢ),
+        # but ne = Σ(nᵢZᵢ) through quasineutrality
         self.data.physics.n_charge_plasma_effective_vol_avg = 0.0
         for imp in range(N_IMPURITIES):
             self.data.physics.n_charge_plasma_effective_vol_avg += (
@@ -1361,7 +1361,7 @@ def plasma_composition(self):
         # ======================================================================
 
         # Mass weighted plasma effective charge
-        # Sum of (Zi^2*n_i) / m_i
+        # Σ(Z²ᵢnᵢ) / mᵢ
         self.data.physics.n_charge_plasma_effective_mass_weighted_vol_avg = (
             (
                 self.data.physics.f_plasma_fuel_deuterium

From 31e2a67fdc90518ff0cb0262e98a2ffef9516060 Mon Sep 17 00:00:00 2001
From: mn3981 <christopher.ashe@ukaea.uk>
Date: Fri, 24 Apr 2026 17:51:47 +0100
Subject: [PATCH 11/13] fix: refactor fradcore function for improved clarity
 and type annotations

Co-authored-by: Copilot <copilot@github.com>
---
 process/models/physics/impurity_radiation.py | 31 +++++++++++---------
 1 file changed, 17 insertions(+), 14 deletions(-)

diff --git a/process/models/physics/impurity_radiation.py b/process/models/physics/impurity_radiation.py
index a9662587f..65e77f558 100644
--- a/process/models/physics/impurity_radiation.py
+++ b/process/models/physics/impurity_radiation.py
@@ -377,29 +377,32 @@ def init_imp_element(
     data.impurity_radiation.impurity_arr_zav[n_species_index - 1, :] = zav
 
 
-def fradcore(rho, radius_plasma_core_norm, f_p_plasma_core_rad_reduction):
-    """Finds the fraction of radiation from the core that is subtracted in impurity
-    radiation model.
+def fradcore(
+    rho: np.array, radius_plasma_core_norm: float, f_p_plasma_core_rad_reduction: float
+) -> np.array:
+    """
+    Creates an array of the same length as `rho` filled with the value of `f_p_plasma_core_rad_reduction`
+    for values of `rho` less than `radius_plasma_core_norm` and 0 for values of `rho`
+    greater than or equal to `radius_plasma_core_norm`.
 
     Parameters
     ----------
-    rho : numpy.array
+    rho
         normalised minor radius
-    radius_plasma_core_norm : float
+    radius_plasma_core_norm
         normalised radius defining the 'core' region
-    f_p_plasma_core_rad_reduction : float
+    f_p_plasma_core_rad_reduction
         fraction of radiation from the core region
 
     Returns
     -------
-    numpy.array
-        fradcore - array filled with the f_p_plasma_core_rad_reduction
+        f_rad_core_profile - array filled with the f_p_plasma_core_rad_reduction
     """
-    fradcore = np.zeros(len(rho))
+    f_rad_core_profile = np.zeros(len(rho))
     rho_mask = rho < radius_plasma_core_norm
-    fradcore[rho_mask] = f_p_plasma_core_rad_reduction
+    f_rad_core_profile[rho_mask] = f_p_plasma_core_rad_reduction
 
-    return fradcore
+    return f_rad_core_profile
 
 
 def calculate_average_charge_at_temp(
@@ -693,9 +696,9 @@ def calculate_radiation_loss_profiles(self):
         pden_impurity_core_rad_total = self.pimp_profile * (
             self.plasma_profile.neprofile.profile_x
             * fradcore(
-                self.plasma_profile.neprofile.profile_x,
-                self.data.impurity_radiation.radius_plasma_core_norm,
-                self.data.impurity_radiation.f_p_plasma_core_rad_reduction,
+                rho=self.plasma_profile.neprofile.profile_x,
+                radius_plasma_core_norm=self.data.impurity_radiation.radius_plasma_core_norm,
+                f_p_plasma_core_rad_reduction=self.data.impurity_radiation.f_p_plasma_core_rad_reduction,
             )
         )
 

From 4c7c3ce6e73da025b1d310e878f0764ede9aff80 Mon Sep 17 00:00:00 2001
From: mn3981 <christopher.ashe@ukaea.uk>
Date: Fri, 24 Apr 2026 17:52:48 +0100
Subject: [PATCH 12/13] fix: rename fradcore function to
 create_f_rad_core_profile for clarity and update related tests

---
 process/models/physics/impurity_radiation.py     | 16 ++++++++++++++--
 .../models/physics/test_impurity_radiation.py    |  4 ++--
 2 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/process/models/physics/impurity_radiation.py b/process/models/physics/impurity_radiation.py
index 65e77f558..7d9b19684 100644
--- a/process/models/physics/impurity_radiation.py
+++ b/process/models/physics/impurity_radiation.py
@@ -377,7 +377,19 @@ def init_imp_element(
     data.impurity_radiation.impurity_arr_zav[n_species_index - 1, :] = zav
 
 
-def fradcore(
+
+def z2index(zimp):
+    for i in range(len(impurity_radiation_module.impurity_arr_label)):
+        if zimp == impurity_radiation_module.impurity_arr_z[i]:
+            return i
+
+    # Should only get here if there is a problem
+    raise ProcessValueError(
+        "Element with the given charge is not in the impurity array", zimp=zimp
+    )
+
+
+def create_f_rad_core_profile(
     rho: np.array, radius_plasma_core_norm: float, f_p_plasma_core_rad_reduction: float
 ) -> np.array:
     """
@@ -695,7 +707,7 @@ def calculate_radiation_loss_profiles(self):
         )
         pden_impurity_core_rad_total = self.pimp_profile * (
             self.plasma_profile.neprofile.profile_x
-            * fradcore(
+            * create_f_rad_core_profile(
                 rho=self.plasma_profile.neprofile.profile_x,
                 radius_plasma_core_norm=self.data.impurity_radiation.radius_plasma_core_norm,
                 f_p_plasma_core_rad_reduction=self.data.impurity_radiation.f_p_plasma_core_rad_reduction,
diff --git a/tests/unit/models/physics/test_impurity_radiation.py b/tests/unit/models/physics/test_impurity_radiation.py
index abebe910c..0f8485d6e 100644
--- a/tests/unit/models/physics/test_impurity_radiation.py
+++ b/tests/unit/models/physics/test_impurity_radiation.py
@@ -91,7 +91,7 @@ class FradcoreParam(NamedTuple):
     expected_fradcore: np.array = np.array
 
 
-def test_fradcore():
+def test_create_f_rad_core_profile():
     """Tests `fradcore` function.
 
     :param rho: normalised minor radius
@@ -121,7 +121,7 @@ def test_fradcore():
             0.0,
         ]),
     )
-    fradcore = impurity_radiation.fradcore(
+    fradcore = impurity_radiation.create_f_rad_core_profile(
         fradcoreparam.rho,
         fradcoreparam.radius_plasma_core_norm,
         fradcoreparam.f_p_plasma_core_rad_reduction,

From 03ddd718c461bdbe32a8e6d35e66206e5f4bda0e Mon Sep 17 00:00:00 2001
From: mn3981 <christopher.ashe@ukaea.uk>
Date: Thu, 28 May 2026 15:19:03 +0100
Subject: [PATCH 13/13] Post rebase fixes

---
 process/models/physics/impurity_radiation.py | 38 ++++++++++----------
 1 file changed, 18 insertions(+), 20 deletions(-)

diff --git a/process/models/physics/impurity_radiation.py b/process/models/physics/impurity_radiation.py
index 7d9b19684..ae4ef3206 100644
--- a/process/models/physics/impurity_radiation.py
+++ b/process/models/physics/impurity_radiation.py
@@ -14,12 +14,12 @@
 
 from process.core import constants
 from process.core.exceptions import ProcessError, ProcessValueError
-from process.core.model import DataStructure
 from process.models.physics.plasma_profiles import PlasmaProfile
 
 if TYPE_CHECKING:
     from pathlib import Path
 
+    from process.core.model import DataStructure
     from process.models.physics.plasma_profiles import PlasmaProfile
 
 logger = logging.getLogger(__name__)
@@ -377,10 +377,9 @@ def init_imp_element(
     data.impurity_radiation.impurity_arr_zav[n_species_index - 1, :] = zav
 
 
-
-def z2index(zimp):
-    for i in range(len(impurity_radiation_module.impurity_arr_label)):
-        if zimp == impurity_radiation_module.impurity_arr_z[i]:
+def z2index(zimp, data: DataStructure) -> int:
+    for i in range(len(data.impurity_radiation.impurity_arr_label)):
+        if zimp == data.impurity_radiation.impurity_arr_z[i]:
             return i
 
     # Should only get here if there is a problem
@@ -418,7 +417,7 @@ def create_f_rad_core_profile(
 
 
 def calculate_average_charge_at_temp(
-    imp_element_index: int, temp_electron_kev: np.array | float
+    imp_element_index: int, temp_electron_kev: np.array | float, data: DataStructure
 ) -> np.array | float:
     """Calculates electron temperature dependent average atomic charge (Z) for a given impurity element.
 
@@ -428,6 +427,8 @@ def calculate_average_charge_at_temp(
         Impurity element index
     temp_electron_kev
         electron temperature in keV
+    data
+        DataStructure containing impurity radiation data
 
     Returns
     -------
@@ -515,6 +516,7 @@ def calculate_impurity_radiation_power_density(
     imp_element_index: int,
     nd_electron_profile: np.array | float,
     temp_electron_profile_kev: np.array | float,
+    data: DataStructure,
 ) -> np.array | float:
     """
     Calculates the impurity radiation density (W/m³) based on the electron density and temperature profiles.
@@ -537,7 +539,7 @@ def calculate_impurity_radiation_power_density(
      -Temperatures outside the range of the L(Z,Tₑ) table are handled by using the L(Z,Tₑ)
       value at the closest temperature in the table,
     """
-    bins = impurity_radiation_module.temp_impurity_keV_array[imp_element_index]
+    bins = data.impurity_radiation.temp_impurity_keV_array[imp_element_index]
     indices = np.digitize(temp_electron_profile_kev, bins)
     indices[indices >= bins.shape[0]] = bins.shape[0] - 1
     indices[indices < 0] = 0
@@ -561,7 +563,7 @@ def calculate_impurity_radiation_power_density(
     # W/m³ = nᵢ * nₑ * L(Z, Tₑ)
     # nᵢ = f_nd_species_electron * nₑ
     pden_impurity_profile = (
-        impurity_radiation_module.f_nd_impurity_electron_array[imp_element_index]
+        data.impurity_radiation.f_nd_impurity_electron_array[imp_element_index]
         * nd_electron_profile
         * nd_electron_profile
         * power_loss_function
@@ -571,26 +573,26 @@ def calculate_impurity_radiation_power_density(
 
     less_than_imp_temp_mask = (
         temp_electron_profile_kev
-        <= impurity_radiation_module.temp_impurity_keV_array[imp_element_index, 0]
+        <= data.impurity_radiation.temp_impurity_keV_array[imp_element_index, 0]
     )
     # This is okay because line radiation will dominate at lower temp, and the L(Z,Tₑ)
     # value at the lowest temperature in the table is likely to be an overestimate of the
     # radiation loss at lower temperatures, so this is a conservative approach.
     pden_impurity_profile[less_than_imp_temp_mask] = (
-        impurity_radiation_module.pden_impurity_lz_nd_temp_array[imp_element_index, 0]
+        data.impurity_radiation.pden_impurity_lz_nd_temp_array[imp_element_index, 0]
     )
 
     # greater_than_imp_temp_mask = temp_electron_profile_kev values higher than impurity temperature.
     greater_than_imp_temp_mask = (
         temp_electron_profile_kev
-        >= impurity_radiation_module.temp_impurity_keV_array[
+        >= data.impurity_radiation.temp_impurity_keV_array[
             imp_element_index,
             data.impurity_radiation.impurity_arr_len_tab[imp_element_index] - 1,
         ]
     )
     #  This is okay because Bremsstrahlung will dominate at higher temp.
     pden_impurity_profile[greater_than_imp_temp_mask] = (
-        impurity_radiation_module.pden_impurity_lz_nd_temp_array[
+        data.impurity_radiation.pden_impurity_lz_nd_temp_array[
             imp_element_index,
             data.impurity_radiation.impurity_arr_len_tab[imp_element_index] - 1,
         ]
@@ -655,7 +657,9 @@ def __init__(self, plasma_profile: PlasmaProfile, data_structure: DataStructure)
         self.pden_impurity_core_rad_profile = np.zeros(
             self.data.physics.n_plasma_profile_elements
         )
-        self.pden_impurity_rad_edge_profile = np.zeros(self.plasma_profile.profile_size)
+        self.pden_impurity_rad_edge_profile = np.zeros(
+            self.data.physics.n_plasma_profile_elements
+        )
 
         self.pden_impurity_rad_total_mw = 0.0
         self.pden_impurity_core_rad_total_mw = 0.0
@@ -667,12 +671,6 @@ def run(self):
     def output(self):
         """This model has no output"""
 
-    def run(self):
-        """This model isn't run"""
-
-    def output(self):
-        """This model has no output"""
-
     def map_imprad_profile(self):
         """Map imprad_profile() over each impurity element index."""
         list(map(self.imprad_profile, self.imp))
@@ -691,7 +689,7 @@ def imprad_profile(self, imp_element_index: int) -> None:
             imp_element_index=imp_element_index,
             nd_electron_profile=self.plasma_profile.neprofile.profile_y,
             temp_electron_profile_kev=self.plasma_profile.teprofile.profile_y,
-            self.data,
+            data=self.data,
         )
 
         self.pimp_profile = np.add(self.pimp_profile, pden_impurity_radiation_profile)