Add docstrings to legendre functions and add markdown comments to tor…

…oidal notebooks

bluemira/equilibria/optimisation/harmonics/toroidal_harmonics_approx_functions.py

@@ -0,0 +1,86 @@
+# SPDX-FileCopyrightText: 2021-present M. Coleman, J. Cook, F. Franza
+# SPDX-FileCopyrightText: 2021-present I.A. Maione, S. McIntosh
+# SPDX-FileCopyrightText: 2021-present J. Morris, D. Short
+#
+# SPDX-License-Identifier: LGPL-2.1-or-later
+
+"""
+A collection of functions used to approximate toroidal harmonics.
+"""
+
+from math import factorial
+
+import numpy as np
+from scipy.special import gamma, poch
+
+
+def f_hypergeometric(a, b, c, z, n_max):
+    """Evaluates the hypergeometric power series up to n_max.
+
+    .. math::
+        F(a, b; c; z) = \\sum_0^{n_max} \\frac{(a)_{s} (b)_{s}}{Gamma(c + s) s!} z^{s}
+
+    See https://dlmf.nist.gov/15.2#E2 and https://dlmf.nist.gov/5.2#iii for more
+    information.
+    """
+    F = 0
+    for s in range(n_max + 1):
+        F += (poch(a, s) * poch(b, s)) / (gamma(c + s) * factorial(s)) * z**s
+    return F
+
+
+def my_legendre_p(lam, mu, x, n_max=20):
+    """Evaluates the associated Legendre function of the first kind of degree lambda and order
+    minus mu as a function of x. See https://dlmf.nist.gov/14.3#E18 for more information.
+
+    TODO check domain of validity? Assumed validity is 1<x<inf
+
+    .. math::
+        P_{\\lambda}^{-\\mu} = 2^{-\\mu} x^{\\lambda - \\mu} (x^2 - 1)^{\\mu/2}
+                        F(\\frac{1}{2}(\\mu - \\lambda), \\frac{1}{2}(\\mu - \\lambda + 1);
+                            \\mu + 1; 1 - \\frac{1}{x^2})
+
+        where F is the hypergeometric function defined above as f_hypergeometric.
+
+    """  # noqa: W505, E501
+    a = 1 / 2 * (mu - lam)
+    b = 1 / 2 * (mu - lam + 1)
+    c = mu + 1
+    z = 1 - 1 / (x**2)
+    F_sum = f_hypergeometric(a=a, b=b, c=c, z=z, n_max=n_max)  # noqa: N806
+    legP = 2 ** (-mu) * x ** (lam - mu) * (x**2 - 1) ** (mu / 2) * F_sum  # noqa: N806
+    return legP  # noqa: RET504
+
+
+def my_legendre_q(lam, mu, x, n_max=20):
+    """Evaluates Olver's associated Legendre function of the second kind of degree lambda
+    and order minus mu as a function of x. See https://dlmf.nist.gov/14, https://dlmf.nist.gov/14.3#E10,
+    and https://dlmf.nist.gov/14.3#E7 for more information.
+
+    TODO check domain of validity? Assumed validity is 1<x<inf
+
+    .. math::
+        Q_{\\lambda}^{-\\mu} = \\frac{\\pi^{\\frac{1}{2}} (x^2 - 1)^{\frac{\\mu}{2}}}
+                                {2^{\\lambda + 1} x^{\\lambda + \\mu + 1}}
+                                F(\\frac{1}{2}(\\lambda + \\mu)+1, \\frac{1}{2}(\\lambda
+                                + \\mu); \\lambda + \\frac{3}{2}; \\frac{1}{x^2})
+
+        where F is the hypergeometric function defined above as f_hypergeometric.
+
+    """
+    a = 1 / 2 * (lam + mu) + 1
+    b = 1 / 2 * (lam + mu + 1)
+    c = lam + 3 / 2
+    z = 1 / (x**2)
+    F_sum = f_hypergeometric(a=a, b=b, c=c, z=z, n_max=n_max)  # noqa: N806
+    legQ = (  # noqa: N806
+        (np.pi ** (1 / 2) * (x**2 - 1) ** (mu / 2))
+        / (2 ** (lam + 1) * x ** (lam + mu + 1))
+        * F_sum
+    )
+    if type(legQ) == np.float64:
+        if x == 1:
+            legQ = np.inf  # noqa: N806
+    else:
+        legQ[x == 1] = np.inf
+    return legQ

examples/equilibria/single_wire_field_toroidal_harmonics.ex.py

@@ -0,0 +1,165 @@
+# ---
+# jupyter:
+#   jupytext:
+#     cell_metadata_filter: tags,-all
+#     notebook_metadata_filter: -jupytext.text_representation.jupytext_version
+#     text_representation:
+#       extension: .py
+#       format_name: percent
+#       format_version: '1.3'
+#   kernelspec:
+#     display_name: Python 3 (ipykernel)
+#     language: python
+#     name: python3
+# ---
+
+# %% tags=["remove-cell"]
+# SPDX-FileCopyrightText: 2021-present M. Coleman, J. Cook, F. Franza
+# SPDX-FileCopyrightText: 2021-present I.A. Maione, S. McIntosh
+# SPDX-FileCopyrightText: 2021-present J. Morris, D. Short
+#
+# SPDX-License-Identifier: LGPL-2.1-or-later
+
+"""
+Calculate solution due to a single wire as a sum of toroidal harmonics.
+"""
+
+# %% [markdown]
+# # Example of calculating the flux solution due to a single wire as a sum of toroidal
+# harmonics
+
+# %% [markdown]
+# ### Imports
+
+# %%
+from math import factorial
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from bluemira.base.constants import MU_0
+from bluemira.equilibria.optimisation.harmonics.toroidal_harmonics_approx_functions import (  # noqa: E501
+    my_legendre_p,
+    my_legendre_q,
+)
+from bluemira.equilibria.plotting import PLOT_DEFAULTS
+from bluemira.utilities.tools import cylindrical_to_toroidal, toroidal_to_cylindrical
+
+# %% [markdown]
+# First we define the location in cylindrical coordinates of the focus $(R_0, z_0)$ and
+# of the wire $(R_c, z_c)$, and the current in the wire, $I_c$.
+#
+# We then convert the location of the wire from cylindrical $(R_c, z_c)$ to toroidal
+# coordinates $(\tau_c, \sigma_c)$ using the following relations:
+# $$ \tau_c = \ln\frac{d_{1}}{d_{2}} $$
+# $$ \sigma_c = {sign}(z - z_{0}) \arccos\frac{d_{1}^2 + d_{2}^2 - 4 R_{0}^2}{2 d_{1}
+# d_{2}} $$
+# where
+# $$ d_{1}^2 = (R_{c} + R_{0})^2 + (z_c - z_{0})^2 $$
+# $$ d_{2}^2 = (R_{c} - R_{0})^2 + (z_c - z_{0})^2 $$
+
+# We need a range of $\tau$ in order to create the grid over which we want to solve. We
+# specify this using the value of $\tau$ at the wire, $\tau_c$ and the approximate
+# minimum distance from the focus. This estimation is necessary as using coordinate
+# transform functions would result in divide by zero errors.
+#
+# Once we have the grid in toroidal coordinates, we can convert this to cylindrical
+# coordinates for use later using the following relations:
+# $$ R = R_0 \frac{\sinh\tau}{\cosh \tau - \cos \sigma}$$
+# $$ z - z_0 = R_0 \frac{\sin \sigma}{\cosh \tau - \cos \sigma}$$
+# %%
+# Wire location
+R_c = 5.0
+Z_c = 1.0
+I_c = 5e6
+
+# Focus of toroidal coordinates
+R_0 = 3.0
+Z_0 = 0.2
+
+# Location of wire in toroidal coordinates
+tau_c, sigma_c = cylindrical_to_toroidal(R_0=R_0, z_0=Z_0, R=R_c, Z=Z_c)
+
+# Using approximate value for d2_min to avoid infinities
+# Approximating tau at the focus instead of using coordinate transform functions
+# (in order to avoid divide by 0 errors)
+d2_min = 0.05
+tau_max = np.log(2 * R_0 / d2_min)
+n_tau = 200
+tau = np.linspace(tau_c, tau_max, n_tau)
+n_sigma = 150
+sigma = np.linspace(-np.pi, np.pi, n_sigma)
+
+# Create grid in toroidal coordinates
+tau, sigma = np.meshgrid(tau, sigma)
+
+# Convert to cylindrical coordinates
+R, Z = toroidal_to_cylindrical(R_0=R_0, z_0=Z_0, tau=tau, sigma=sigma)
+
+
+# %% [markdown]
+# Now we want to calculate the following coefficients
+# $$ A_m = \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)) $$
+# $$ A_m^{sin} = A_m \sin(m \sigma_c) $$
+# $$ A_m^{cos} = A_m \cos(m \sigma_c) $$
+# where $$ \Delta_c = \cosh(\tau_c) - \cos(\sigma_c) $$
+# and $$ factorial\_term = \prod_{0}^{m-1} \left( 1 + \frac{1}{2(m-i)}\right) $$
+#
+
+# %%
+# Useful combinations
+Delta = np.cosh(tau) - np.cos(sigma)
+Deltac = np.cosh(tau_c) - np.cos(sigma_c)
+
+# Calculate coefficients
+m_max = 5
+Am_cos = np.zeros(m_max + 1)
+Am_sin = np.zeros(m_max + 1)
+
+for m in range(m_max + 1):
+    factorial_term = 1 if m == 0 else np.prod(1 + 0.5 / np.arange(1, m + 1))
+    A_m = (
+        (MU_0 * I_c / 2 ** (5 / 2))
+        * factorial_term
+        * (np.sinh(tau_c) / np.sqrt(Deltac))
+        * my_legendre_p(m - 1 / 2, 1, np.cosh(tau_c))
+    )
+    Am_cos[m] = A_m * np.cos(m * sigma_c)
+    Am_sin[m] = A_m * np.sin(m * sigma_c)
+
+# %% [markdown]
+# Now we can use the following
+#
+# $$ A(\tau, \sigma) = \sum_{m=0}^{\infty} A_m^{\cos} \epsilon_m m! \sqrt{\frac{2}{\pi}}
+# \Delta^{\frac{1}{2}} 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}}
+# Q_{m-\frac{1}{2}}^{1}(\cosh \tau) \sin(m \sigma) $$
+
+# along with $$\psi = R A$$
+# to calculate the solution and plot the psi graph. Here we have that
+# $ \epsilon_0 = 1$ and $\epsilon_{m\ge 1} = 2$.
+
+# %%
+# TODO equation (19) from OB document
+epsilon = 2 * np.ones(m_max + 1)
+epsilon[0] = 1
+A = np.zeros_like(R)
+
+for m in range(m_max + 1):
+    A += Am_cos[m] * epsilon[m] * factorial(m) * np.sqrt(2 / np.pi) * np.sqrt(
+        Delta
+    ) * my_legendre_q(m - 1 / 2, 1, np.cosh(tau)) * np.cos(m * sigma) + Am_sin[
+        m
+    ] * epsilon[m] * factorial(m) * np.sqrt(2 / np.pi) * np.sqrt(Delta) * my_legendre_q(
+        m - 1 / 2, 1, np.cosh(tau)
+    ) * np.sin(m * sigma)
+
+psi = R * A
+nlevels = PLOT_DEFAULTS["psi"]["nlevels"]
+cmap = PLOT_DEFAULTS["psi"]["cmap"]
+plt.contour(R, Z, psi, nlevels, cmap=cmap)
+plt.xlabel("R")
+plt.ylabel("Z")
+plt.title("Psi")
+# %%

examples/equilibria/toroidal_harmonics_component_function_verification.ex.py

@@ -0,0 +1,121 @@
+# ---
+# jupyter:
+#   jupytext:
+#     cell_metadata_filter: tags,-all
+#     notebook_metadata_filter: -jupytext.text_representation.jupytext_version
+#     text_representation:
+#       extension: .py
+#       format_name: percent
+#       format_version: '1.3'
+#   kernelspec:
+#     display_name: Python 3 (ipykernel)
+#     language: python
+#     name: python3
+# ---
+
+# %% tags=["remove-cell"]
+# SPDX-FileCopyrightText: 2021-present M. Coleman, J. Cook, F. Franza
+# SPDX-FileCopyrightText: 2021-present I.A. Maione, S. McIntosh
+# SPDX-FileCopyrightText: 2021-present J. Morris, D. Short
+#
+# SPDX-License-Identifier: LGPL-2.1-or-later
+
+"""
+An example of plotting the internal harmonics in toroidal coordinates about a focus
+point.
+"""
+
+# %% [markdown]
+# # Example of plotting the internal harmonics in toroidal coordinates
+
+# This example shows how to plot the individual cos and sin toroidal harmonic
+# contributions about a focus point.
+
+# %% [markdown]
+# ### Imports
+
+# %%
+import matplotlib.pyplot as plt
+import numpy as np
+
+from bluemira.equilibria.optimisation.harmonics.toroidal_harmonics_approx_functions import (  # noqa: E501
+    my_legendre_q,
+)
+from bluemira.equilibria.plotting import PLOT_DEFAULTS
+from bluemira.utilities.tools import cylindrical_to_toroidal
+
+# %% [markdown]
+# First we need to set up the grid in cylindrical coordinates over which we will solve,
+# and define the location of the focus point.
+# We will convert this to toroidal coordinates.
+#
+#
+# To convert from cylindrical coordinates $(R, z)$ to toroidal coordinates
+# $(\tau, \sigma)$ about a focus point $(R_0, z_0)$ we have the following relations:
+# $$ \tau = \ln\frac{d_{1}}{d_{2}} $$
+# $$ \sigma = {sign}(z - z_{0}) \arccos\frac{d_{1}^2 + d_{2}^2 - 4 R_{0}^2}{2 d_{1}
+# d_{2}}  $$
+# where
+# $$ d_{1}^2 = (R + R_{0})^2 + (z - z_{0})^2 $$
+# $$ d_{2}^2 = (R - R_{0})^2 + (z - z_{0})^2 $$
+#
+# The domains for the toroidal coordinates are $0 \le \tau < \infty$ and $-\pi < \sigma
+# \le \pi$.
+
+# %%
+# Set up grid over which to solve
+r = np.linspace(0, 6, 100)
+z = np.linspace(-6, 6, 100)
+R, Z = np.meshgrid(r, z)
+
+# Define focus point
+R_0 = 3.0
+Z_0 = 0.0
+
+# Convert to toroidal coordinates
+tau, sigma = cylindrical_to_toroidal(R_0=R_0, z_0=Z_0, R=R, Z=Z)
+
+# %% [markdown]
+# Now we want to calculate and plot the internal harmonics in toroidal coordinates about
+# the focus.
+#
+# We use the following equations
+# $$ \psi_{sin} = R \sqrt{\cosh (\tau) - \cos (\sigma)} Q_{m - \frac{1}{2}}^1
+# (\cosh \tau) \sin (m \sigma) $$
+#
+# $$ \psi_{cos} = R \sqrt{\cosh (\tau) - \cos (\sigma)} Q_{m - \frac{1}{2}}^1
+# (\cosh \tau) \cos (m \sigma) $$
+#
+# We evaluate these contributions to the harmonics about the focus and plot.
+# %%
+# Calculate and plot individual contributions from toroidal harmonics
+
+nu = np.arange(0, 5)
+
+# Setting up plots
+fig_sin, axs_sin = plt.subplots(1, len(nu))
+fig_sin.suptitle("sin")
+fig_sin.supxlabel("R")
+fig_sin.supylabel("Z")
+fig_cos, axs_cos = plt.subplots(1, len(nu))
+fig_cos.suptitle("cos")
+fig_cos.supxlabel("R")
+fig_cos.supylabel("Z")
+nlevels = PLOT_DEFAULTS["psi"]["nlevels"]
+cmap = PLOT_DEFAULTS["psi"]["cmap"]
+
+for i_nu in range(len(nu)):
+    levels_m = np.linspace(-(i_nu + 1), i_nu + 1, 100 * i_nu)
+    foo = (
+        R
+        * np.sqrt(np.cosh(tau) - np.cos(sigma))
+        * my_legendre_q(nu[i_nu] - 1 / 2, 1, np.cosh(tau))
+    )
+    psi_sin = foo * np.sin(nu[i_nu] * sigma)
+    psi_cos = foo * np.cos(nu[i_nu] * sigma)
+    axs_sin[i_nu].contour(R, Z, psi_sin, levels=nlevels, cmap=cmap)
+    axs_sin[i_nu].title.set_text(f"m = {i_nu}")
+    axs_cos[i_nu].contour(R, Z, psi_cos, levels=nlevels, cmap=cmap)
+    axs_cos[i_nu].title.set_text(f"m = {i_nu}")
+
+# %%