Adding documentation, delete unnecessary fn and argument (#4128)

* Adding documentation, delete unnecessary fn and argument

* remove unnecessary argument

* updating values in test

* pre commit fixes

* commit image added to docs

* edit docstring

* add currents fn back in

* reword figure caption

* change input of ToroidalHarmonicsConstraint so that it takes ToroidalHarmonicsSelectionResult and update notebook and tests

* Update documentation/source/equilibria/equilibria.rst

Co-authored-by: geograham <[email protected]>

* Update documentation/source/equilibria/equilibria.rst

Co-authored-by: geograham <[email protected]>

* Update documentation/source/equilibria/equilibria.rst

Co-authored-by: geograham <[email protected]>

* Update documentation/source/equilibria/equilibria.rst

Co-authored-by: geograham <[email protected]>

* Update documentation/source/equilibria/equilibria.rst

Co-authored-by: geograham <[email protected]>

* Update documentation/source/equilibria/equilibria.rst

Co-authored-by: geograham <[email protected]>

* Update documentation/source/equilibria/equilibria.rst

Co-authored-by: geograham <[email protected]>

* Update documentation/source/equilibria/equilibria.rst

Co-authored-by: geograham <[email protected]>

* Update documentation/source/equilibria/equilibria.rst

Co-authored-by: geograham <[email protected]>

* Additions to docs and fix broken test

* replace == with assert almost equal

---------

Co-authored-by: geograham <[email protected]>

Author: clmould

bluemira/equilibria/optimisation/harmonics/harmonics_constraints.py

@@ -27,7 +27,7 @@
     ToroidalHarmonicConstraintFunction,
 )
 from bluemira.equilibria.optimisation.harmonics.toroidal_harmonics_approx_functions import (  # noqa: E501
-    ToroidalHarmonicsParams,
+    ToroidalHarmonicsSelectionResult,
     coil_toroidal_harmonic_amplitude_matrix,
 )
 from bluemira.utilities.tools import is_num
@@ -194,17 +194,20 @@ class ToroidalHarmonicConstraint(UpdateableConstraint):
 
     def __init__(
         self,
-        ref_harmonics_cos: npt.NDArray[np.float64],
-        ref_harmonics_sin: npt.NDArray[np.float64],
-        ref_harmonics_cos_amplitudes: npt.NDArray[np.float64],
-        ref_harmonics_sin_amplitudes: npt.NDArray[np.float64],
-        th_params: ToroidalHarmonicsParams,
+        th_result: ToroidalHarmonicsSelectionResult,
         relative_tolerance_cos: float | npt.NDArray[np.float64] = 1e-3,
         relative_tolerance_sin: float | npt.NDArray[np.float64] = 1e-3,
         constraint_type: str = "equality",
         weights: float | np.ndarray = 1.0,
     ):
         self.constraint_type = constraint_type
+
+        ref_harmonics_cos = th_result.cos_m
+        ref_harmonics_sin = th_result.sin_m
+        ref_harmonics_cos_amplitudes = th_result.cos_amplitudes
+        ref_harmonics_sin_amplitudes = th_result.sin_amplitudes
+        th_params = th_result.th_params
+
         tolerance_cos = np.abs(relative_tolerance_cos * ref_harmonics_cos_amplitudes)
         tolerance_sin = np.abs(relative_tolerance_sin * ref_harmonics_sin_amplitudes)
         tolerance = np.append(tolerance_cos, tolerance_sin, axis=0)

bluemira/equilibria/optimisation/harmonics/toroidal_harmonics_approx_functions.py

@@ -256,6 +256,11 @@ def toroidal_harmonic_grid_and_coil_setup(
         R coordinate of the toroidal focus point in cylindrical coordinates
     Z_0:
         Z coordinate of the toroidal focus point in cylindrical coordinates
+    tau_limit:
+        How the maximum tau value is chosen. The three options are:
+            - LCFS: use the maximum extent of the LCFS
+            - COIL: use the maximum area within all the coils
+            - MANUAL: use a specified limit
     min_tau_value:
         The minimum tau for the toroidal coordinate approximation region,
         lower min tau means a larger region of space (maximum tau is at focus)
@@ -799,10 +804,10 @@ def toroidal_harmonics_to_positions(
 
     Returns
     -------
-    :
-        collocation matrix for cos components
-    :
+    harmonics2collocation_cos:
         collocation matrix for cos components
+    harmonics2collocation_sin:
+        collocation matrix for sin components
 
     Raises
     ------
@@ -918,7 +923,6 @@ def toroidal_harmonic_approximation(
     max_harmonic_mode: int = 5,
     *,
     plasma_mask: bool = False,
-    from_psi_fit: bool = True,
 ) -> ToroidalHarmonicsSelectionResult:
     """
     Calculate the toroidal harmonic (TH) amplitudes/coefficients for a given
@@ -954,10 +958,6 @@ def toroidal_harmonic_approximation(
     plasma_mask:
         Whether or not to apply a mask to the error metric (within the psi_norm flux
         surface)
-    from_psi_fit:
-        If True then the toroidal harmonics approximation of the coilset contribution
-        to psi is calculated from fit to the psi values at certain collocation points.
-        Otherwise, it is calculated from the coilset currents
 
     Returns
     -------
@@ -983,13 +983,9 @@ def toroidal_harmonic_approximation(
         max_harmonic_mode,
         len(th_params.th_coil_names),
     )
-    collocation = (
-        None
-        if not from_psi_fit
-        else collocation_points(
-            eq.get_LCFS(),
-            PointType.GRID_POINTS,
-        )
+    collocation = collocation_points(
+        eq.get_LCFS(),
+        PointType.GRID_POINTS,
     )
 
     true_coilset_psi, fixed_psi, collocation_psi = _separate_psi_contributions(
@@ -1011,26 +1007,19 @@ def toroidal_harmonic_approximation(
         sin_m_chosen = mode_values[mode_id[mode_id >= max_harmonic_mode]]
 
         # Calculate psi using the combination of poloidal mode numbers (m) selected in
-        # this iteration
-        if not from_psi_fit:
-            error_new, approximate_coilset_psi, cos_amps, sin_amps = (
-                _approximation_direct_from_currents(
-                    eq, th_params, cos_m_chosen, sin_m_chosen, true_coilset_psi, mask
-                )
-            )
-        else:
-            error_new, approximate_coilset_psi, cos_amps, sin_amps = (
-                _approximation_from_psi_fitting(
-                    th_params,
-                    n_degrees_of_freedom,
-                    collocation,
-                    mode_id,
-                    max_harmonic_mode,
-                    collocation_psi,
-                    mask,
-                    true_coilset_psi,
-                )
+        # this iteration at the collocation points
+        error_new, approximate_coilset_psi, cos_amps, sin_amps = (
+            _approximation_from_psi_fitting(
+                th_params,
+                n_degrees_of_freedom,
+                collocation,
+                mode_id,
+                max_harmonic_mode,
+                collocation_psi,
+                mask,
+                true_coilset_psi,
             )
+        )
         # If the new error is less than the previously lowest error, then select the
         # current combination of poloidal mode numbers (m), amplitudes and associated psi
         if error_new < error:

documentation/source/equilibria/equilibria.rst

@@ -659,6 +659,122 @@ calculation of the Jacobian of the field constraints. Instead, we choose
 to constrain the poloidal field at the centre of the inside edge of each
 coil, where the field is generally the highest.
 
+
+Toroidal Harmonic constraints
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+An equilibrium poloidal field has a plasma and coilset contribution. We can use Toroidal Harmonic (TH) functions to approximate the coilset contribution to the poloidal magnetic flux using equations :eq:`PoloidalFlux` and :eq:`ToroidalHarmonics`.
+
+.. math::
+   :label: PoloidalFlux
+
+   \psi = R A = \frac{R_0 \sinh \tau}{\Delta} A
+
+.. math::
+   :label: ToroidalHarmonics
+
+   \begin{aligned}
+   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}}
+   \textbf{Q}_{m-\frac{1}{2}}^{1}(\cosh \tau) \sin(m \sigma)\end{aligned}
+
+where
+
+- :math:`\psi` is poloidal flux
+
+- :math:`m` is the poloidal mode number
+
+- :math:`A_m^{\cos, \sin}` are harmonic amplitudes
+
+- :math:`\varepsilon_m = 1` for :math:`m = 0` and :math:`\varepsilon_m = 2` for :math:`m \ge 1`
+
+- :math:`\Delta = \cosh(\tau) - \cos(\sigma)`
+
+- :math:`\textbf{Q}_{\nu}^{\mu}` is Olver's definition of the associated Legendre function of the second kind.
+
+Toroidal coordinates are defined as :math:`\tau = \ln\frac{d_1}{d_2}` and :math:`\sigma  = sign(z - z_0) \arccos\frac{d_1^2 + d_2^2 - 4 R_0^2}{2 d_1 d_2}` where, :math:`d_1^2 = (R + R_0)^2 + (z - z_0)^2` and :math:`d_2^2 = (R - R_0)^2 + (z - z_0)^2`, and :math:`R_0` and :math:`Z_0` are the coordinates of the focus point.
+
+Toroidal harmonic amplitudes can be implemented as a set of constraints or targets in a coilset optimisation. The TH amplitudes are used to preserve the coilset contribution in the region occupied by the core plasma (i.e. the region characterised by closed flux surfaces), enabling modification of the poloidal magnetic configuration outside the core region without (significantly) altering the plasma equilibrium.
+
+Potential benefits of using TH as a set of constraints include:
+
+- We can choose not to re-solve for the plasma equilibrium at each optimisation 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 considering.
+
+To set up a toroidal harmonic constraint, we first find the significant poloidal modes and their amplitude values using equations :eq:`PoloidalFlux` and :eq:`ToroidalHarmonics`.
+
+The user chooses the:
+
+- Number of degrees of freedom (DoF) appropriate for the given optimisation problem – this limits the number of poloidal modes that can be used in the TH approximation. The default value will allow DoF up to the number of coilset currents being optimised.
+
+- Normalised psi value of the flux surface within which the coilset contribution to psi will be constrained.
+
+Then these values are input in `toroidal_harmonic_approximation`, which iterates through all the combinations of cos and sin modes up to the DoF limit. For each iteration:
+
+- Psi values for a set of sampled positions within the selected flux surface are found and used to determine the amplitudes of the contributing poloidal modes.
+
+- L2 norm of the error between the approximated coilset psi and the true coilset psi is calculated.
+
+The `toroidal_harmonic_approximation` returns the result as a `ToroidalHarmonicsSelectionResult` for the approximate psi with the smallest difference to the true coilset psi in the selected closed flux region. This contains all the information needed to set up a `ToroidalHarmonicConstraint`, as well as supporting information that can be used in analysis.
+
+A `ToroidalHarmonicsSelectionResult` contains the combination of modes (and their amplitude values) that gives the best approximation when compared using an L2 norm of the error across the psi map. These can be then implemented as an :math:`A\bf{x} = b` constraint or magnetic target using equation :eq:`THAmplitudeCurrentRelation`.
+
+.. math::
+   :label: THAmplitudeCurrentRelation
+
+   A_m^{\cos, \sin} = \frac{\mu_0 I_c}{2^{\frac{5}{2}}} factorial\_term \frac{\sinh(
+   \tau_c)}
+   {\Delta_c^{\frac{1}{2}}} P_{m - \frac{1}{2}}^{-1}(\cosh(\tau_c)) ^{\cos}_{\sin}(m
+   \sigma_c)
+
+where
+
+- :math:`A_m^{\cos, \sin}` are hamonic amplitudes for a single coil
+
+- subscript :math:`c` refers to a single coil
+
+- :math:`I_c, \tau_c, \sigma_c` are the coil current, and coil position in toroidal coordinates.
+
+- :math:`P_{m - \frac{1}{2}}^{-1}` is the associated Legendre function of the first kind.
+
+- :math:`\Delta_c = \cosh(\tau_c) - \cos(\sigma_c)`
+
+- :math:`factorial\_term = \prod_{i=0}^{m-1} \left( 1 + \frac{1}{2(m-i)}\right)`
+
+.. figure:: th-flux-comparison.png
+   :name: fig:th-flux-comparison
+
+   Diagram showing a comparison of coilset psi calculated in bluemira and the approximation of coilset psi found using toroidal harmonic functions.
+
+See `here`_ or F. W. J. Olver (1997b) Asymptotics and Special Functions. A. K. Peters, Wellesley, MA. for more information on the Legendre functions used in the TH approximation.
+
+.. _here: https://dlmf.nist.gov/14
+
+Note 1:
+
+The default TH approximation region is set to encompass the the maximum extent of the LCFS, with its focus point at the effective centre of the plasma. However, the function `toroidal_harmonic_grid_and_coil_setup` can be used to approximate a region with a different area and focus point. The optional arguments, which are used to specify where the approximation region is placed, are:
+
+- the maximum extent of the LCFS,
+
+- the maxmimum area that is contained within all coils,
+
+- a user-specified tau limit.
+
+Lower min tau means a larger region of space (the maximum tau is at the focus point).
+
+Note 2:
+
+If a coil is within the TH approximation region then its contribution must be kept fixed during an optimisation. The coilset 'control coils' must be set to be the same as the `th_coil_names` returned in the `ToroidalHarmonicsSelectionResult` while setting up the optimisation problem. Please see 'examples/equilibria/Toroidal_Harmonics_Optimisation_Set_Up.ex.py' for an example.
+
+
+.. figure:: th-region-and-coils.png
+   :name: fig:th-region
+
+   Diagram showing the default TH approximation region.
+
 Coil position optimisation and constraints
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 

documentation/source/equilibria/th-flux-comparison.png

@@ -106,7 +106,6 @@
     n_degrees_of_freedom=6,
     max_harmonic_mode=5,
     plasma_mask=True,
-    from_psi_fit=True,
 )
 
 # %% [markdown]
@@ -133,12 +132,8 @@
 # %%
 # Create a constraint
 th_constraint = ToroidalHarmonicConstraint(
-    ref_harmonics_cos=result.cos_m,
-    ref_harmonics_sin=result.sin_m,
-    ref_harmonics_cos_amplitudes=result.cos_amplitudes,
-    ref_harmonics_sin_amplitudes=result.sin_amplitudes,
+    th_result=result,
     constraint_type="inequality",
-    th_params=th_params,
 )
 # Ensure control coils are set to those that can be used in the toroidal
 # harmonic approximation
@@ -147,6 +142,7 @@
 # Show the constraint region
 f, ax = plt.subplots()
 th_constraint.plot(ax=ax)
+eq.coilset.plot(ax=ax)
 eq.plot(ax=ax)
 
 # %%