Tidying up toroidal harmonics notebooks on develop for release (#4039)

* Tidying up toroidal harmonics notebooks on develop for release

* fix test failures

* Changes based on review comments

Author: clmould

CITATION.cff

@@ -84,6 +84,10 @@ authors:
     given-names: S.
     affiliation: United Kingdom Atomic Energy Authority, Culham Science Centre, Abingdon, Oxfordshire OX14 3DB, United Kingdom
 
+  - family-names: Mould
+    given-names: C. L.
+    affiliation: United Kingdom Atomic Energy Authority, Culham Science Centre, Abingdon, Oxfordshire OX14 3DB, United Kingdom
+
   - family-names: Saunders
     given-names: H.
     affiliation: United Kingdom Atomic Energy Authority, Culham Science Centre, Abingdon, Oxfordshire OX14 3DB, United Kingdom

bluemira/equilibria/optimisation/harmonics/toroidal_harmonics_approx_functions.py

@@ -324,6 +324,9 @@ def coil_toroidal_harmonic_amplitude_matrix(
     Legendre functions of the first kind of degree lambda and order minus mu, and :math:
     \\Delta_c = \\cosh{\\tau_c} - \\cos{\\sigma_c}.
 
+    Note: the factorial term \\frac{(2m+1)!!}{2^m m!} is equivalent to 1 if m = 0,
+    otherwise \\prod_{i=0}^{m-1} \\left( 1 + \\frac{1}{2(m-i)}\\right)
+
     Parameters
     ----------
     input_coils:
@@ -368,7 +371,7 @@ def coil_toroidal_harmonic_amplitude_matrix(
     currents2harmonics = np.zeros([max_degree, np.size(tau_c)])
 
     # TH coefficients from function of the current distribution
-    # outside of the region containing the core plamsa
+    # outside of the region containing the core plasma
     # TH coefficients = currents2harmonics @ coil currents
     degrees = np.arange(0, max_degree)[:, None]
     factorial_term = np.array([

documentation/source/examples.rst

@@ -27,9 +27,8 @@ Equilibria Examples
     examples/equilibria/coilset_opt_problem_tutorial
     examples/equilibria/anaylsis_toolbox_examples
     examples/equilibria/quick_look_at_example_coilsets
-    examples/equilibria/toroidal_harmonics_approximation_basic_use
-    examples/equilibria/toroidal_harmonics_component_function_walkthrough_and_verification
-    examples/equilibria/toroidal_harmonics_full_coilset_approximation_bluemira_comparison
+    examples/equilibria/Toroidal_Approximation_Explained
+    examples/equilibria/Toroidal_Harmonics_Optimisation_Set_Up
 
 
 Geometry Examples

…_approximation_bluemira_comparison.ex.py …a/Toroidal_Approximation_Explained.ex.pyexamples/equilibria/toroidal_harmonics_full_coilset_approximation_bluemira_comparison.ex.py renamed to examples/equilibria/Toroidal_Approximation_Explained.ex.py examples/equilibria/toroidal_harmonics_full_coilset_approximation_bluemira_comparison.ex.py renamed to examples/equilibria/Toroidal_Approximation_Explained.ex.py

@@ -30,8 +30,10 @@
 # This example illustrates the inner workings of the bluemira toroidal harmonics
 # approximation function (toroidal_harmonic_approximation) which can be used in
 # coilset current and position optimisation for conventional aspect ratio tokamaks.
-# For full details and all equations, look in the
-# notebook toroidal_harmonics_component_function_walkthrough_and_verification.ex.py.
+# For an example of how toroidal_harmonic_approximation is used
+# to create toroidal harmonics constraints for use in a coilset optimisation problem,
+# please see
+# Toroidal_Harmonics_Optimisation_Set_Up.ex.py
 #
 # ## Premise:
 #
@@ -49,21 +51,21 @@
 # There are benefits to using TH as a minimal set of constraints:
 # - We can choose not to re-solve for the plasma equilibrium at each step, since the
 # coilset contribution to the core plasma (within the LCFS) is constrained.
-# - We have a minimal set of constraints (a set of harmonic amplitudes) for the core
-# plasma contribution, which can reduce the dimensionality of the problem we are
+# - We have a small, viable set of constraints (a set of harmonic amplitudes) for the
+# core plasma contribution, which can reduce the dimensionality of the problem we are
 # considering.
 #
+# %% [markdown]
 # We get the TH amplitudes/coefficients, $A(\tau, \sigma)$, from the following equations:
 #
 # $$ A(\tau, \sigma) = \sum_{m=0}^{\infty} A_m^{\cos} \epsilon_m m! \sqrt{\frac{2}{\pi}}
-# \Delta^{\frac{1}{2}} \textbf{Q}_{m-\frac{1}{2}}^{1}(\cosh \tau) \cos(m \sigma) + A_m^
-# {\sin}
-# \epsilon_m m! \sqrt{\frac{2}{\pi}} \Delta^{\frac{1}{2}}
+# \Delta^{\frac{1}{2}} \textbf{Q}_{m-\frac{1}{2}}^{1}(\cosh \tau) \cos(m \sigma)
+# \\+ A_m^{\sin} \epsilon_m m! \sqrt{\frac{2}{\pi}} \Delta^{\frac{1}{2}}
 # \textbf{Q}_{m-\frac{1}{2}}^{1}(\cosh \tau) \sin(m \sigma) $$
 #
 # where
 #
-# $$ A_m^{\cos, \sin} = \frac{\mu_0 I_c}{2^{\frac{5}{2}}} factorial\_term \frac{\sinh(
+# $$ A_m^{\cos, \sin} = \frac{\mu_0 I_c}{2^{\frac{5}{2}}} \upsilon_{fact} \frac{\sinh(
 # \tau_c)}
 # {\Delta_c^{\frac{1}{2}}} P_{m - \frac{1}{2}}^{-1}(\cosh(\tau_c)) ^{\cos}_{\sin}(m
 # \sigma_c) $$
@@ -82,15 +84,16 @@
 # - $\varepsilon_m = 1 $ for $m = 0$ and $\varepsilon_m = 2$ for $m \ge 1$
 # - $ \Delta = \cosh(\tau) - \cos(\sigma) $
 # - $ \Delta_c = \cosh(\tau_c) - \cos(\sigma_c) $
-# - $ factorial\_term = \prod_{i=0}^{m-1} \left( 1 + \frac{1}{2(m-i)}\right) $
+# - The factorial term, $\upsilon_{fact}$ = 1 if $m = 0$, else $= \prod_{i=0}^{m-1}
+# \left( 1 + \frac{1}{2(m-i)}\right) $
 #
-# %% [markdown]
-# ## Imports
+# We can then obtain the flux, $\psi$ by using $\psi = R A$.
 
 # %%
 from copy import deepcopy
 from pathlib import Path
 
+import matplotlib.patches as patch
 import matplotlib.pyplot as plt
 import numpy as np
 
@@ -109,79 +112,98 @@
 from bluemira.geometry.coordinates import Coordinates
 
 # %% [markdown]
-# Get equilibrium from EQDSK file and plot
+# Get equilibrium from EQDSK file
 
 # %%
 # Get equilibrium
 EQDATA = get_bluemira_path("equilibria/test_data", subfolder="tests")
 
-# eq_name = "eqref_OOB.json"
-# eq_name = Path(EQDATA, eq_name)
-# eq = Equilibrium.from_eqdsk(eq_name, from_cocos=7)
-
 eq_name = "DN-DEMO_eqref.json"
 eq_name = Path(EQDATA, eq_name)
 eq = Equilibrium.from_eqdsk(eq_name, from_cocos=3, qpsi_positive=False)
 
-# Plot the equilibrium
-f, ax = plt.subplots()
-eq.plot(ax=ax)
-eq.coilset.plot(ax=ax)
-plt.show()
 
 # %% [markdown]
 # Set the focus point
 #
 # Note: there is the possibility to experiment with the position of the focus, but
-# the default is to use the plasma o point.
+# the default is to use the effective centre of the plasma.
 # %%
 # Set the focus point
 R_0, Z_0 = eq.effective_centre()
 
 # %% [markdown]
 # Use the `toroidal_harmonic_grid_and_coil_setup` to obtain the necessary
 # parameters required for the TH approximation, such as the R and Z coordinates of the
-# TH grid.
-#
-# We then use the `toroidal_harmonic_approximate_psi` function to approximate the coilset
-# contribution to the flux using toroidal harmonics. We need to provide this function
-# with the equilibrium, eq, and the ToroidalHarmonicsParams dataclass, which contains
-# necessary parameters for the TH approximation, such as the relevant coordinates
-# and coil names for use in the approximation. The
-# function returns the approx_coilset_psi array and the TH coefficient matrix A_m.
-# The default focus point is the plasma o point.
-# The white dot in the plot shows the focus point.
+# TH grid. We then plot the equilibrium and in orange we show the region over which
+# we will use our toroidal harmonic approximation.
+
 # %%
 # Approximate psi and plot
 th_params = toroidal_harmonic_grid_and_coil_setup(eq=eq, R_0=R_0, Z_0=Z_0)
 
 R_approx = th_params.R
 Z_approx = th_params.Z
 
+# Plot the equilibrium and the region over which we approximate the flux using toroidal
+# harmonics.
+f, ax = plt.subplots()
+eq.plot(ax=ax)
+eq.coilset.plot(ax=ax)
+max_R = np.max(th_params.R)  # noqa: N816
+min_R = np.min(th_params.R)  # noqa: N816
+max_Z = np.max(th_params.Z)  # noqa: N816
+min_Z = np.min(th_params.Z)  # noqa: N816
+centre_R = (max_R - min_R) / 2 + min_R  # noqa: N816
+centre_Z = (max_Z - np.abs(min_Z)) / 2  # noqa: N816
+radius = (max_R - min_R) / 2
+
+ax.add_patch(
+    patch.Circle((centre_R, centre_Z), radius, ec="orange", fill=True, fc="orange")
+)
+plt.show()
+# %% [markdown]
+#
+# We then use the `toroidal_harmonic_approximate_psi` function to approximate the coilset
+# contribution to the flux using toroidal harmonics. This makes use of the equation for
+# $A_m^{\cos, \sin}$ as displayed at the start of this notebook.
+# We need to provide this function
+# with the equilibrium, eq, and the ToroidalHarmonicsParams dataclass, which contains
+# necessary parameters for the TH approximation, such as the relevant coordinates
+# and coil names for use in the approximation. The
+# function returns the approx_coilset_psi array and the TH coefficient matrices Am_cos
+# and Am_sin.
+
+# %%
 approx_coilset_psi, _, _ = toroidal_harmonic_approximate_psi(
     eq=eq, th_params=th_params, max_degree=5
 )
 
-nlevels = PLOT_DEFAULTS["psi"]["nlevels"]
-cmap = PLOT_DEFAULTS["psi"]["cmap"]
-plt.contourf(R_approx, Z_approx, approx_coilset_psi, nlevels, cmap=cmap)
-plt.xlabel("R")
-plt.ylabel("Z")
-plt.title("TH Approximation for Coilset Psi")
-plt.show()
-
 
 # %% [markdown]
-# Now we want to compare this approximation to the solution from bluemira.
+# Now we want to compare this approximation for the coilset psi to the solution from
+# bluemira.
+# Since we assume the flux in the plasma region is kept fixed, we can calculate
+# approximate total psi = bluemira plasma psi + toroidal harmonic coilset psi
+# approximation.
+#
+# We also create a difference plot to compare our TH coilset psi approximation to the
+# bluemira coilset psi.
+# We can see zero difference in the core region, which we expect as we are constraining
+# the flux in this region, and we see differences outside of the approximation
+# region, but this is expected as we are only requiring the core region to be kept
+# fixed.
+# Since the plasma psi is kept fixed, the total psi difference plot would look
+# the same as the coilset psi difference plot shown below.
 
 # %%
-# Plot total psi using approximate coilset psi from TH, and plasma psi from bluemira
+# Approximate total psi = approximate coilset psi from TH + plasma psi from bluemira
 
 plasma_psi = eq.plasma.psi(R_approx, Z_approx)
 
 total_psi = approx_coilset_psi + plasma_psi
 
-# Find LCFS from TH approx
+# Find LCFS from TH approx to add to plot
 approx_eq = deepcopy(eq)
 o_points, x_points = approx_eq.get_OX_points(total_psi)
 
@@ -197,84 +219,63 @@
     eq.get_LCFS() if np.isclose(psi_norm, 1.0) else eq.get_flux_surface(psi_norm)
 )
 
+# Difference in coilset psi between TH approximation and bluemira
+bluemira_coilset_psi = eq.coilset.psi(R_approx, Z_approx)
+
+coilset_psi_diff = np.abs(approx_coilset_psi - bluemira_coilset_psi) / np.max(
+    np.abs(bluemira_coilset_psi)
+)
+coilset_psi_diff_plot = coilset_psi_diff
+f, axs = plt.subplots(1, 2)
+
 # Plot
-plt.contourf(R_approx, Z_approx, total_psi, nlevels, cmap=cmap)
-plt.xlabel("R")
-plt.ylabel("Z")
-plt.plot(
+nlevels = PLOT_DEFAULTS["psi"]["nlevels"]
+cmap = PLOT_DEFAULTS["psi"]["cmap"]
+axs0_plot = axs[0].contourf(R_approx, Z_approx, total_psi, nlevels, cmap=cmap)
+axs[0].set_xlabel("R")
+axs[0].set_ylabel("Z")
+axs[0].plot(
     approx_fs.x,
     approx_fs.z,
     color="red",
-    label=f"Closed Flux Surface (psi_n={psi_norm}) from TH approximation",
+    label="Approx. FS from TH",
 )
-plt.title("Total Psi using TH approximation for coilset psi")
-plt.legend(loc="upper right")
-plt.show()
-
-# %%
-# Plot bluemira total psi
-# Obtain psi from Bluemira coilset
-bm_coil_psi = np.zeros(np.shape(eq.grid.x))
-for n in eq.coilset.name:
-    bm_coil_psi = np.sum([bm_coil_psi, eq.coilset[n].psi(eq.grid.x, eq.grid.z)], axis=0)
-
-
-# Obtain total psi from Bluemira
-bluemira_total_psi = eq.psi(R_approx, Z_approx)
-
-# Plotting
-plt.contourf(R_approx, Z_approx, bluemira_total_psi, levels=nlevels, cmap=cmap)
-plt.xlabel("R")
-plt.ylabel("Z")
-plt.title("Total Bluemira Psi")
-plt.show()
-# %%
-# Plot bluemira coilset psi
-coilset_psi = eq.coilset.psi(R_approx, Z_approx)
-
-# Plotting
-plt.contourf(R_approx, Z_approx, coilset_psi, nlevels, cmap=cmap)
-plt.xlabel("R")
-plt.ylabel("Z")
-plt.title("Bluemira Coilset Psi")
-plt.show()
-# %%
-# Difference plot to compare TH approximation to Bluemira coilset psi
-# We see zero difference in the core region, which we expect as we are constraining
-# the flux in this region, and we see larger differences outside of the approximation
-# region.
-coilset_psi_diff = np.abs(approx_coilset_psi - coilset_psi) / np.max(np.abs(coilset_psi))
-coilset_psi_diff_plot = coilset_psi_diff
-f, ax = plt.subplots()
-ax.plot(approx_fs.x, approx_fs.z, color="red", label="Approximate LCFS from TH")
-ax.plot(original_fs.x, original_fs.z, color="blue", label="LCFS from Bluemira")
-im = ax.contourf(R_approx, Z_approx, coilset_psi_diff_plot, levels=nlevels, cmap=cmap)
-f.colorbar(mappable=im)
-ax.set_title("Absolute relative difference between coilset psi and TH approximation psi")
-ax.legend(loc="upper right", bbox_to_anchor=(1.1, 1.0))
-eq.coilset.plot(ax=ax)
-plt.show()
-
-
-# %%
-# Difference plot to compare TH approximation to Bluemira total psi
-total_psi_diff = np.abs(total_psi - bluemira_total_psi) / np.max(
-    np.abs(bluemira_total_psi)
+axs[0].set_title("Total Psi using TH approximation for coilset psi")
+axs[0].legend(loc="upper right")
+
+axs[1].plot(approx_fs.x, approx_fs.z, color="red", label="Approx. FS from TH")
+axs[1].plot(
+    original_fs.x,
+    original_fs.z,
+    color="c",
+    linestyle="dashed",
+    label="LCFS from Bluemira",
 )
-total_psi_diff_plot = total_psi_diff
-f, ax = plt.subplots()
-ax.plot(approx_fs.x, approx_fs.z, color="red", label="Approx FS from TH")
-ax.plot(original_fs.x, original_fs.z, color="blue", label="FS from Bluemira")
-im = ax.contourf(R_approx, Z_approx, total_psi_diff_plot, levels=nlevels, cmap=cmap)
-f.colorbar(mappable=im)
-ax.set_title("Absolute relative difference between total psi and TH approximation psi")
-ax.legend(loc="upper right")
-eq.coilset.plot(ax=ax)
+axs1_plot = axs[1].contourf(
+    R_approx, Z_approx, coilset_psi_diff_plot, levels=nlevels, cmap=cmap
+)
+f.colorbar(axs1_plot, ax=axs[1], fraction=0.05)
+axs[1].set_title(
+    "Absolute relative difference between coilset psi \nand TH approximation psi"
+)
+axs[1].legend(loc="upper right", bbox_to_anchor=(1.1, 1.0))
+eq.coilset.plot(ax=axs[1])
+axs[0].set_aspect("equal")
+axs[1].set_aspect("equal")
 plt.show()
 
+
 # %% [markdown]
 # We use a fit metric to evaluate the TH approximation.
-
+# This involves
+# finding a flux surface (here corresponding to a psi_norm=0.95) for our
+# approximate total psi, and comparing to the equivalent flux surface from
+# the bluemira total psi. The fit metric we use for this flux surface
+# comparison is as follows: fit metric value = total area within one but not
+# both flux surfaces / (bluemira flux surface area + approximation flux surface area).
+# A fit metric of 0 would correspond to a perfect match between our approximation
+# and bluemira, and a fit metric of 1 would correspond to no overlap between the
+# approximation and bluemira flux surfaces.
 # %%
 # Fit metric to evaluate TH approximation
 fit_metric_value = fs_fit_metric(original_fs, approx_fs)
@@ -287,7 +288,8 @@
 # number of degrees to use for the approximation.
 #
 # Here is an example of using the function, setting plot to True outputs a graph of the
-# difference in total psi between the TH approximation and bluemira.
+# difference in total psi between the TH approximation and bluemira. This graph
+# is the same as the one we produced in this notebook above.
 # %%
 (
     toroidal_harmonics_params,