Show seperate impurity density plot

process/core/io/plot_proc.py

@@ -3897,12 +3897,23 @@ def plot_n_profiles(prof, demo_ranges: bool, mfile: mf.MFile, scan: int):
     fgwped_out = mfile.get("fgwped_out", scan=scan)
     fgwsep_out = mfile.get("fgwsep_out", scan=scan)
     nd_plasma_electrons_vol_avg = mfile.get("nd_plasma_electrons_vol_avg", scan=scan)
+    # find impurity densities
+    imp_frac = np.array([
+        mfile.get(f"f_nd_impurity_electrons({i:02d})", scan=scan) for i in range(1, 15)
+    ])
 
     nd_plasma_separatrix_electron = mfile.get("nd_plasma_separatrix_electron", scan=scan)
 
-    prof.set_xlabel(r"$\rho \quad [r/a]$")
-    prof.set_ylabel(r"$n \ [10^{19}\ \mathrm{m}^{-3}]$")
-    prof.set_title("Density profile")
+    ax_main = prof.add_subplot(631)
+    ax_main.set_position([0.075, 0.625, 0.25, 0.325])
+    ax_impurity = prof.add_subplot(634, sharex=ax_main)
+    ax_impurity.set_position([0.075, 0.275, 0.25, 0.325])
+    ax_main.tick_params(labelbottom=False)
+
+    ax_impurity.set_xlabel(r"$\rho \quad [r/a]$")
+    ax_main.set_ylabel(r"$n \ [10^{19}\ \mathrm{m}^{-3}]$")
+    ax_impurity.set_ylabel(r"$n \ [10^{16}\ \mathrm{m}^{-3}]$")
+    ax_main.set_title("Density profile")
 
     i_plasma_pedestal = mfile.get("i_plasma_pedestal", scan=scan)
     nd_plasma_pedestal_electron = mfile.get("nd_plasma_pedestal_electron", scan=scan)
@@ -3912,28 +3923,27 @@ def plot_n_profiles(prof, demo_ranges: bool, mfile: mf.MFile, scan: int):
         "radius_plasma_pedestal_density_norm", scan=scan
     )
     ne0 = mfile.get("nd_plasma_electron_on_axis", scan=scan)
+    n_plasma_profile_elements = mfile.get("n_plasma_profile_elements", scan=scan)
 
     # build electron profile and species profiles (scale with electron profile shape)
     if i_plasma_pedestal == 1:
-        rhocore = np.linspace(0, radius_plasma_pedestal_density_norm)
-        necore = (
-            nd_plasma_pedestal_electron
-            + (ne0 - nd_plasma_pedestal_electron)
-            * (1 - rhocore**2 / radius_plasma_pedestal_density_norm**2) ** alphan
-        )
-
-        rhosep = np.linspace(radius_plasma_pedestal_density_norm, 1)
-        neesep = nd_plasma_separatrix_electron + (
-            nd_plasma_pedestal_electron - nd_plasma_separatrix_electron
-        ) * (1 - rhosep) / (1 - min(0.9999, radius_plasma_pedestal_density_norm))
-
-        rho = np.append(rhocore, rhosep)
-        ne = np.append(necore, neesep)
+        rho = np.linspace(0, 1.0, int(n_plasma_profile_elements))
+        ne = np.zeros_like(rho)
 
+        for i in range(len(rho)):
+            if rho[i] <= radius_plasma_pedestal_density_norm:
+                ne[i] = (
+                    nd_plasma_pedestal_electron
+                    + (ne0 - nd_plasma_pedestal_electron)
+                    * (1 - rho[i] ** 2 / radius_plasma_pedestal_density_norm**2)
+                    ** alphan
+                )
+            else:
+                ne[i] = nd_plasma_separatrix_electron + (
+                    nd_plasma_pedestal_electron - nd_plasma_separatrix_electron
+                ) * (1 - rho[i]) / (1 - min(0.9999, radius_plasma_pedestal_density_norm))
     else:
-        rho1 = np.linspace(0, 0.95)
-        rho2 = np.linspace(0.95, 1)
-        rho = np.append(rho1, rho2)
+        rho = np.linspace(0, 1.0, n_plasma_profile_elements)
         ne = ne0 * (1 - rho**2) ** alphan
 
     # species profiles scaled by their average fraction relative to electrons
@@ -3961,72 +3971,107 @@ def plot_n_profiles(prof, demo_ranges: bool, mfile: mf.MFile, scan: int):
     # convert to 1e19 m^-3 units for plotting (vectorised)
     density_profiles_plotting = density_profiles / 1e19
 
-    prof.plot(
+    ax_main.plot(
         rho,
         density_profiles_plotting[0],
         label=r"$n_{\text{fuel}}$",
         color="#2ca02c",
         linewidth=1.5,
     )
-    prof.plot(
+    ax_main.plot(
         rho,
         density_profiles_plotting[1],
         label=r"$n_{\alpha}$",
         color="#d62728",
         linewidth=1.5,
     )
-    prof.plot(
+    ax_impurity.plot(
         rho,
-        density_profiles_plotting[2],
+        density_profiles_plotting[2] * 1e3,
         label=r"$n_{p}$",
         color="#17becf",
         linewidth=1.5,
     )
-    prof.plot(
+    ax_impurity.plot(
         rho,
-        density_profiles_plotting[3],
-        label=r"$n_{Z}$",
+        density_profiles_plotting[3] * 1e3,
+        label=r"$n_{imp,total}$",
         color="#9467bd",
-        linewidth=1.5,
+        linewidth=2.5,
+        linestyle="dotted",
     )
-    prof.plot(
+    ax_main.plot(
         rho,
         density_profiles_plotting[4],
         label=r"$n_{i,total}$",
         color="#ff7f0e",
         linewidth=1.5,
     )
-    prof.plot(
+    ax_main.plot(
         rho, density_profiles_plotting[5], label=r"$n_{e}$", color="blue", linewidth=1.5
     )
 
-    # make legend use multiple columns (up to 4) and place it to the right to avoid overlapping the plots
-    _handles, labels = prof.get_legend_handles_labels()
-    ncol = min(4, max(1, len(labels)))
-    prof.legend(loc="upper center", bbox_to_anchor=(0.5, -0.1), ncol=ncol)
+    if imp_frac[2] > 1.0e-30:
+        ax_impurity.plot(rho, imp_frac[2] * ne / 1e16, label=r"$n_{\text{Be}}$")
+    if imp_frac[3] > 1.0e-30:
+        ax_impurity.plot(rho, imp_frac[3] * ne / 1e16, label=r"$n_{\text{C}}$")
+    if imp_frac[4] > 1.0e-30:
+        ax_impurity.plot(rho, imp_frac[4] * ne / 1e16, label=r"$n_{\text{N}}$")
+    if imp_frac[5] > 1.0e-30:
+        ax_impurity.plot(rho, imp_frac[5] * ne / 1e16, label=r"$n_{\text{O}}$")
+    if imp_frac[6] > 1.0e-30:
+        ax_impurity.plot(rho, imp_frac[6] * ne / 1e16, label=r"$n_{\text{Ne}}$")
+    if imp_frac[7] > 1.0e-30:
+        ax_impurity.plot(rho, imp_frac[7] * ne / 1e16, label=r"$n_{\text{Si}}$")
+    if imp_frac[8] > 1.0e-30:
+        ax_impurity.plot(rho, imp_frac[8] * ne / 1e16, label=r"$n_{\text{Ar}}$")
+    if imp_frac[9] > 1.0e-30:
+        ax_impurity.plot(rho, imp_frac[9] * ne / 1e16, label=r"$n_{\text{Fe}}$")
+    if imp_frac[10] > 1.0e-30:
+        ax_impurity.plot(rho, imp_frac[10] * ne / 1e16, label=r"$n_{\text{Ni}}$")
+    if imp_frac[11] > 1.0e-30:
+        ax_impurity.plot(rho, imp_frac[11] * ne / 1e16, label=r"$n_{\text{Kr}}$")
+    if imp_frac[12] > 1.0e-30:
+        ax_impurity.plot(rho, imp_frac[12] * ne / 1e16, label=r"$n_{\text{Xe}}$")
+    if imp_frac[13] > 1.0e-30:
+        ax_impurity.plot(rho, imp_frac[13] * ne / 1e16, label=r"$n_{\text{W}}$")
+
+    ax_main.legend(loc="best")
+    ax_impurity.legend(loc="best")
 
     # Ranges
     # ---
     # DEMO : Fixed ranges for comparison
-    prof.set_xlim([0, 1])
+    ax_main.set_xlim([0, 1])
+    ax_impurity.set_xlim([0, 1])
     if demo_ranges:
-        prof.set_ylim([0, 20])
+        ax_main.set_ylim([0, 20])
 
     # Adapatative ranges
     else:
-        prof.set_ylim([0, prof.get_ylim()[1]])
+        ax_main.set_ylim([0, ax_main.get_ylim()[1]])
+        # Use logarithmic scale for impurity axis if any impurity values are very small
+        impurity_data = [
+            imp_frac[i] * ne / 1e16
+            for i in range(len(imp_frac))
+            if imp_frac[i] > 1.0e-30
+        ]
+        if impurity_data and np.min(impurity_data) / np.max(impurity_data) < 0.01:
+            # If range spans more than 100x, use log scale
+            ax_impurity.set_yscale("log")
+        ax_impurity.set_ylim([1e-3, ax_impurity.get_ylim()[1]])
 
     if i_plasma_pedestal != 0:
         # Print pedestal lines
-        prof.axhline(
+        ax_main.axhline(
             y=nd_plasma_pedestal_electron / 1e19,
             xmax=radius_plasma_pedestal_density_norm,
             color="r",
             linestyle="-",
             linewidth=0.4,
             alpha=0.4,
         )
-        prof.vlines(
+        ax_main.vlines(
             x=radius_plasma_pedestal_density_norm,
             ymin=0.0,
             ymax=nd_plasma_pedestal_electron / 1e19,
@@ -4035,7 +4080,8 @@ def plot_n_profiles(prof, demo_ranges: bool, mfile: mf.MFile, scan: int):
             linewidth=0.4,
             alpha=0.4,
         )
-    prof.minorticks_on()
+    ax_main.minorticks_on()
+    ax_impurity.minorticks_on()
 
     # Add text box with density profile parameters
     textstr_density = "\n".join((
@@ -4056,11 +4102,11 @@ def plot_n_profiles(prof, demo_ranges: bool, mfile: mf.MFile, scan: int):
     ))
 
     props_density = {"boxstyle": "round", "facecolor": "wheat", "alpha": 0.5}
-    prof.text(
+    ax_main.text(
         0.0,
-        -0.25,
+        -0.175,
         textstr_density,
-        transform=prof.transAxes,
+        transform=ax_impurity.transAxes,
         fontsize=9,
         verticalalignment="top",
         bbox=props_density,
@@ -4079,23 +4125,23 @@ def plot_n_profiles(prof, demo_ranges: bool, mfile: mf.MFile, scan: int):
         f"{mfile.get('nd_plasma_protons_vol_avg', scan=scan):.3e} m$^{{-3}}$",
     ))
 
-    prof.text(
+    ax_impurity.text(
         1.2,
-        -0.725,
+        0.05,
         textstr_ions,
         fontsize=9,
         verticalalignment="bottom",
         horizontalalignment="left",
-        transform=prof.transAxes,
+        transform=ax_impurity.transAxes,
         bbox={
             "boxstyle": "round",
             "facecolor": "wheat",
             "alpha": 0.5,
         },
     )
 
-    prof.grid(True, which="both", linestyle="--", linewidth=0.5, alpha=0.2)
-
+    ax_main.grid(True, which="both", linestyle="--", linewidth=0.5, alpha=0.2)
+    ax_impurity.grid(True, which="both", linestyle="--", linewidth=0.5, alpha=0.2)
     # ---
 
 
@@ -4146,8 +4192,8 @@ def plot_jprofile(prof, mfile: mf.MFile, scan: int):
 
     props_j = {"boxstyle": "round", "facecolor": "wheat", "alpha": 0.5}
     prof.text(
-        1.1,
-        0.75,
+        0.65,
+        1.6,
         textstr_j,
         transform=prof.transAxes,
         fontsize=9,
@@ -4156,15 +4202,15 @@ def plot_jprofile(prof, mfile: mf.MFile, scan: int):
     )
 
     prof.text(
-        0.05,
+        0.35,
         0.04,
         "*Current profile is assumed to be parabolic",
         fontsize=10,
         ha="left",
         transform=plt.gcf().transFigure,
     )
     prof.text(
-        0.05,
+        0.35,
         0.02,
         "*Bootstrap profile is for representation only",
         fontsize=10,
@@ -4284,7 +4330,7 @@ def plot_t_profiles(prof, demo_ranges: bool, mfile: mf.MFile, scan: int):
     props_temperature = {"boxstyle": "round", "facecolor": "wheat", "alpha": 0.5}
     prof.text(
         0.0,
-        -0.16,
+        -0.125,
         textstr_temperature,
         transform=prof.transAxes,
         fontsize=9,
@@ -4347,16 +4393,16 @@ def plot_qprofile(prof, demo_ranges: bool, mfile: mf.MFile, scan: int):
     prof.grid(True, which="both", linestyle="--", linewidth=0.5, alpha=0.2)
     # ---
 
-    textstr_q = "\n".join((
-        r"$q_0$: " + f"{q0:.3f}\n",
-        r"$q_{95}$: " + f"{q95:.3f}\n",
+    textstr_q = " | ".join((
+        r"$q_0$: " + f"{q0:.3f}",
+        r"$q_{95}$: " + f"{q95:.3f}",
         r"$q_{\text{cyl}}$: " + f"{mfile.get('qstar', scan=scan):.3f}",
     ))
 
     props_q = {"boxstyle": "round", "facecolor": "wheat", "alpha": 0.5}
     prof.text(
-        -0.4,
-        0.75,
+        0.0,
+        1.4,
         textstr_q,
         transform=prof.transAxes,
         fontsize=9,
@@ -13651,13 +13697,11 @@ def main_plot(
     )
 
     # Plot density profiles
-    ax9 = figs[6].add_subplot(231)
-    ax9.set_position([0.075, 0.55, 0.25, 0.4])
-    plot_n_profiles(ax9, demo_ranges, m_file, scan)
+    plot_n_profiles(figs[6], demo_ranges, m_file, scan)
 
     # Plot temperature profiles
     ax10 = figs[6].add_subplot(232)
-    ax10.set_position([0.375, 0.55, 0.25, 0.4])
+    ax10.set_position([0.375, 0.575, 0.25, 0.375])
     plot_t_profiles(ax10, demo_ranges, m_file, scan)
 
     # Plot impurity profiles
@@ -13667,12 +13711,12 @@ def main_plot(
 
     # Plot current density profile
     ax12 = figs[6].add_subplot(4, 3, 10)
-    ax12.set_position([0.075, 0.105, 0.25, 0.15])
+    ax12.set_position([0.375, 0.105, 0.25, 0.15])
     plot_jprofile(ax12, m_file, scan)
 
     # Plot q profile
     ax13 = figs[6].add_subplot(4, 3, 12)
-    ax13.set_position([0.7, 0.125, 0.25, 0.15])
+    ax13.set_position([0.7, 0.105, 0.25, 0.15])
     plot_qprofile(ax13, demo_ranges, m_file, scan)
 
     plot_plasma_effective_charge_profile(figs[7].add_subplot(221), m_file, scan)

process/models/physics/physics.py

@@ -2362,6 +2362,16 @@ def outplas(self):
                     str2,
                     impurity_radiation_module.f_nd_impurity_electrons[imp],
                 )
+                if impurity_radiation_module.f_nd_impurity_electrons[imp] != 0.0:
+                    for i in range(physics_variables.n_plasma_profile_elements):
+                        po.ovarre(
+                            self.mfile,
+                            str1 + f" at point {i}",
+                            f"(f_nd_impurity_electrons{imp}_{i})",
+                            impurity_radiation_module.f_nd_impurity_electrons[imp]
+                            * self.plasma_profile.neprofile.profile_y[i],
+                            "OP ",
+                        )
 
         po.ovarre(
             self.outfile,