diff --git a/.github/workflows/black-fix.yml b/.github/workflows/black-fix.yml index 2785a3a..72c3fd1 100644 --- a/.github/workflows/black-fix.yml +++ b/.github/workflows/black-fix.yml @@ -11,12 +11,12 @@ jobs: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v4 with: ref: ${{ github.head_ref }} - name: Set up Python - uses: actions/setup-python@v2 + uses: actions/setup-python@v5 with: python-version: '3.x' diff --git a/.github/workflows/python_publish.yml b/.github/workflows/python_publish.yml index f6134d4..7569693 100644 --- a/.github/workflows/python_publish.yml +++ b/.github/workflows/python_publish.yml @@ -8,9 +8,9 @@ jobs: deploy: runs-on: ubuntu-latest steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v4 - name: Set up Python - uses: actions/setup-python@v2 + uses: actions/setup-python@v5 with: python-version: '3.x' - name: Install dependencies diff --git a/.github/workflows/test-python-package.yml b/.github/workflows/test-python-package.yml index fefac86..4eb673d 100644 --- a/.github/workflows/test-python-package.yml +++ b/.github/workflows/test-python-package.yml @@ -11,19 +11,19 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: ['3.8', '3.9', '3.10', '3.11'] + python-version: ["3.10", "3.11", "3.12"] steps: - - uses: actions/checkout@v2 + - uses: actions/checkout@v4 - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 + uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: Install dependencies run: | python -m pip install --upgrade pip - python -m pip install flake8 pytest - if [ -f requirements.txt ]; then pip install -r requirements.txt; fi + python -m pip install flake8 + python -m pip install .[tests] - name: Lint with flake8 run: | # stop the build if there are Python syntax errors or undefined names diff --git a/pyproject.toml b/pyproject.toml index 9b38547..03e356b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,12 +1,54 @@ [build-system] requires = [ - "setuptools >= 42", + "setuptools >= 61.0.0", "setuptools_scm[toml] >= 6.2", - "setuptools_scm_git_archive", - "wheel >= 0.29.0", ] build-backend = "setuptools.build_meta" +[project] +name = "zoidberg" +readme = "README.md" +authors = [{name = "Peter Hill", email = "peter.hill@york.ac.uk"}, {name = "BOUT++ team"}] +description = "Generate flux-coordinate independent (FCI) grids for BOUT++" +license = {text = "GNU Lesser General Public License v3 or later (LGPLv3+)"} +classifiers = [ + "Programming Language :: Python", + "Development Status :: 4 - Beta", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: GNU Lesser General Public License v3 or later (LGPLv3+)", + "Natural Language :: English", + "Topic :: Scientific/Engineering :: Visualization", +] +keywords = [ + "bout++", + "bout", + "plasma", + "physics", +] +requires-python = ">=3.10" +dependencies = [ + "boututils ~= 0.1.10", + "numpy ~= 1.24", + "sympy ~= 1.7", + "scipy ~= 1.10", + "matplotlib ~= 3.7", + "netCDF4 ~= 1.7", + "freeqdsk >= 0.4.0", +] +dynamic = ["version"] + +[project.urls] +Documentation = "https://bout-dev.readthedocs.io/en/latest" +Repository = "https://github.com/boutproject/zoidberg" +Tracker = "https://github.com/boutproject/zoidberg/issues" + +[project.optional-dependencies] +tests = ["pytest >= 3.3.0"] + +[project.scripts] +zoidberg-rotating-ellipse = "zoidberg.examples.rotating_ellipse:main" +zoidberg-screwpinch = "zoidberg.examples.screwpinch:main" + [tool.setuptools_scm] write_to = "zoidberg/_version.py" git_describe_command = "git describe --dirty --tags --first-parent" diff --git a/setup.cfg b/setup.cfg deleted file mode 100644 index 5611520..0000000 --- a/setup.cfg +++ /dev/null @@ -1,48 +0,0 @@ -[metadata] -name = zoidberg -long_description = file: README.md -long_description_content_type = text/markdown -author = Peter Hill -author_email = peter.hill@york.ac.uk -description = Generate flux-coordinate independent (FCI) grids for BOUT++ -license = LGPLv3 -license_file = LICENSE -classifiers = - Development Status :: 4 - Beta - Intended Audience :: Science/Research - Intended Audience :: Education - Intended Audience :: Developers - License :: OSI Approved :: GNU Lesser General Public License v3 or later (LGPLv3+) - Natural Language :: English - Operating System :: POSIX :: Linux - Programming Language :: Python :: 3.6 - Programming Language :: Python :: 3.7 - Programming Language :: Python :: 3.8 - Programming Language :: Python :: 3.9 - Programming Language :: Python :: 3.10 - Topic :: Scientific/Engineering :: Visualization -project_urls = - Documentation = https://bout-dev.readthedocs.io/en/latest - Source = https://github.com/boutproject/zoidberg - Tracker = https://github.com/boutproject/zoidberg/issues - -[options] -packages = find: -install_requires = - boututils ~= 0.1.10 - numpy ~= 1.14 - sympy ~= 1.7 - scipy ~= 1.0 - matplotlib ~= 3.1 - netcdf4 ~= 1.4 - importlib-metadata; python_version >= "3.6" - freeqdsk >= 0.4.0 - -[options.extras_require] -tests = - pytest >= 3.3.0 - -[options.entry_points] -console_scripts = - zoidberg-rotating-ellipse = zoidberg.examples.rotating_ellipse:main - zoidberg-screwpinch = zoidberg.examples.screwpinch:main diff --git a/setup.py b/setup.py deleted file mode 100644 index 6068493..0000000 --- a/setup.py +++ /dev/null @@ -1,3 +0,0 @@ -from setuptools import setup - -setup() diff --git a/zoidberg/field.py b/zoidberg/field.py index 2cad2fc..3d709a3 100644 --- a/zoidberg/field.py +++ b/zoidberg/field.py @@ -1,6 +1,8 @@ from math import gamma import numpy as np +from sympy import (And, Piecewise, Sum, Symbol, atan2, cos, diff, factorial, + lambdify, log, pi, sin, sqrt) from . import boundary @@ -315,624 +317,574 @@ def Rfunc(self, x, z, phi): return np.full(x.shape, self.Rmaj) -try: - from sympy import (And, Piecewise, Sum, Symbol, atan2, cos, diff, - factorial, lambdify, log, pi, sin, sqrt) +class StraightStellarator(MagneticField): + """A "rotating ellipse" stellarator without curvature - class StraightStellarator(MagneticField): - """A "rotating ellipse" stellarator without curvature + Parameters + ---------- + xcentre : float, optional + Middle of the domain in x [m] + zcentre : float, optional + Middle of the domain in z [m] + radius : float, optional + Radius of coils [meters] + yperiod : float, optional + The period over which the coils return to their original position + I_coil : float, optional + Current in each coil + + """ + + def coil(self, xcentre, zcentre, radius, angle, iota, I): + """Defines a single coil Parameters ---------- - xcentre : float, optional - Middle of the domain in x [m] - zcentre : float, optional - Middle of the domain in z [m] - radius : float, optional - Radius of coils [meters] - yperiod : float, optional - The period over which the coils return to their original position - I_coil : float, optional - Current in each coil + radius : float + Radius to coil + angle : float + Initial angle of coil + iota : float + Rotational transform of coil + I : float + Current through coil + Returns + ------- + (x, z) - x, z coordinates of coils along phi """ - def coil(self, xcentre, zcentre, radius, angle, iota, I): - """Defines a single coil - - Parameters - ---------- - radius : float - Radius to coil - angle : float - Initial angle of coil - iota : float - Rotational transform of coil - I : float - Current through coil - - Returns - ------- - (x, z) - x, z coordinates of coils along phi - """ - - return ( - xcentre + radius * cos(angle + iota * self.phi), - zcentre + radius * sin(angle + iota * self.phi), - I, + return ( + xcentre + radius * cos(angle + iota * self.phi), + zcentre + radius * sin(angle + iota * self.phi), + I, + ) + + def __init__( + self, + xcentre=0.0, + zcentre=0.0, + radius=0.8, + yperiod=np.pi, + I_coil=0.05, + smooth=False, + smooth_args={}, + ): + xcentre = float(xcentre) + zcentre = float(zcentre) + radius = float(radius) + yperiod = float(yperiod) + + iota = 2.0 * np.pi / yperiod + + self.x = Symbol("x") + self.z = Symbol("z") + self.y = Symbol("y") + self.r = Symbol("r") + self.r = (self.x**2 + self.z**2) ** (0.5) + self.phi = Symbol("phi") + + self.xcentre = xcentre + self.zcentre = zcentre + self.radius = radius + + # Four coils equally spaced, alternating direction for current + self.coil_list = [ + self.coil( + xcentre, + zcentre, + radius, + n * pi, + iota, + ((-1) ** np.mod(i, 2)) * I_coil, ) + for i, n in enumerate(np.arange(4) / 2.0) + ] - def __init__( - self, - xcentre=0.0, - zcentre=0.0, - radius=0.8, - yperiod=np.pi, - I_coil=0.05, - smooth=False, - smooth_args={}, - ): - xcentre = float(xcentre) - zcentre = float(zcentre) - radius = float(radius) - yperiod = float(yperiod) - - iota = 2.0 * np.pi / yperiod - - self.x = Symbol("x") - self.z = Symbol("z") - self.y = Symbol("y") - self.r = Symbol("r") - self.r = (self.x**2 + self.z**2) ** (0.5) - self.phi = Symbol("phi") - - self.xcentre = xcentre - self.zcentre = zcentre - self.radius = radius - - # Four coils equally spaced, alternating direction for current - self.coil_list = [ - self.coil( - xcentre, - zcentre, - radius, - n * pi, - iota, - ((-1) ** np.mod(i, 2)) * I_coil, - ) - for i, n in enumerate(np.arange(4) / 2.0) - ] + A = 0.0 + Bx = 0.0 + Bz = 0.0 - A = 0.0 - Bx = 0.0 - Bz = 0.0 + for c in self.coil_list: + xc, zc, Ic = c + r2 = (self.x - xc) ** 2 + (self.z - zc) ** 2 + theta = atan2(self.z - zc, self.x - xc) # Angle relative to coil - for c in self.coil_list: - xc, zc, Ic = c - r2 = (self.x - xc) ** 2 + (self.z - zc) ** 2 - theta = atan2(self.z - zc, self.x - xc) # Angle relative to coil + A -= Ic * 0.1 * log(r2) - A -= Ic * 0.1 * log(r2) + B = Ic * 0.2 / sqrt(r2) - B = Ic * 0.2 / sqrt(r2) + Bx += B * sin(theta) + Bz -= B * cos(theta) - Bx += B * sin(theta) - Bz -= B * cos(theta) + self.Afunc = lambdify((self.x, self.z, self.phi), A, "numpy") - self.Afunc = lambdify((self.x, self.z, self.phi), A, "numpy") + self.Bxfunc = lambdify((self.x, self.z, self.phi), Bx, "numpy") + self.Bzfunc = lambdify((self.x, self.z, self.phi), Bz, "numpy") - self.Bxfunc = lambdify((self.x, self.z, self.phi), Bx, "numpy") - self.Bzfunc = lambdify((self.x, self.z, self.phi), Bz, "numpy") - class RotatingEllipse(MagneticField): - """A "rotating ellipse" stellarator +class RotatingEllipse(MagneticField): + """A "rotating ellipse" stellarator + Parameters + ---------- + xcentre : float, optional + Middle of the domain in x [m] + zcentre : float, optional + Middle of the domain in z [m] + radius : float, optional + Radius of coils [meters] + yperiod : float, optional + The period over which the coils return to their original position + I_coil : float, optional + Current in each coil + Btor : float, optional + Toroidal magnetic field strength + """ + + def coil(self, xcentre, zcentre, radius, angle, iota, I): + """Defines a single coil Parameters ---------- - xcentre : float, optional - Middle of the domain in x [m] - zcentre : float, optional - Middle of the domain in z [m] - radius : float, optional - Radius of coils [meters] - yperiod : float, optional - The period over which the coils return to their original position - I_coil : float, optional - Current in each coil - Btor : float, optional - Toroidal magnetic field strength + radius : float + Radius to coil + angle : float + Initial angle of coil + iota : float + Rotational transform of coil + I : float + Current through coil + Returns + ------- + (x, z) - x, z coordinates of coils along phi """ - def coil(self, xcentre, zcentre, radius, angle, iota, I): - """Defines a single coil - Parameters - ---------- - radius : float - Radius to coil - angle : float - Initial angle of coil - iota : float - Rotational transform of coil - I : float - Current through coil - Returns - ------- - (x, z) - x, z coordinates of coils along phi - """ - - return ( - xcentre + radius * cos(angle + iota * self.phi), - zcentre + radius * sin(angle + iota * self.phi), - I, + return ( + xcentre + radius * cos(angle + iota * self.phi), + zcentre + radius * sin(angle + iota * self.phi), + I, + ) + + def __init__( + self, + xcentre=0.0, + zcentre=0.0, + radius=0.8, + yperiod=np.pi, + I_coil=0.05, + Btor=1.0, + smooth=False, + smooth_args={}, + ): + xcentre = float(xcentre) + zcentre = float(zcentre) + radius = float(radius) + yperiod = float(yperiod) + Btor = float(Btor) + + iota = 2.0 * np.pi / yperiod + + self.x = Symbol("x") + self.z = Symbol("z") + self.y = Symbol("y") + self.r = Symbol("r") + self.r = (self.x**2 + self.z**2) ** (0.5) + self.phi = Symbol("phi") + + self.xcentre = xcentre + self.zcentre = zcentre + self.radius = radius + + # Four coils equally spaced, alternating direction for current + self.coil_list = [ + self.coil( + xcentre, + zcentre, + radius, + n * pi, + iota, + ((-1) ** np.mod(i, 2)) * I_coil, ) + for i, n in enumerate(np.arange(4) / 2.0) + ] - def __init__( - self, - xcentre=0.0, - zcentre=0.0, - radius=0.8, - yperiod=np.pi, - I_coil=0.05, - Btor=1.0, - smooth=False, - smooth_args={}, - ): - xcentre = float(xcentre) - zcentre = float(zcentre) - radius = float(radius) - yperiod = float(yperiod) - Btor = float(Btor) - - iota = 2.0 * np.pi / yperiod - - self.x = Symbol("x") - self.z = Symbol("z") - self.y = Symbol("y") - self.r = Symbol("r") - self.r = (self.x**2 + self.z**2) ** (0.5) - self.phi = Symbol("phi") - - self.xcentre = xcentre - self.zcentre = zcentre - self.radius = radius - - # Four coils equally spaced, alternating direction for current - self.coil_list = [ - self.coil( - xcentre, - zcentre, - radius, - n * pi, - iota, - ((-1) ** np.mod(i, 2)) * I_coil, - ) - for i, n in enumerate(np.arange(4) / 2.0) - ] + A = 0.0 + Bx = 0.0 + Bz = 0.0 + + for c in self.coil_list: + xc, zc, Ic = c + rc = (xc**2 + zc**2) ** (0.5) + r2 = (self.x - xc) ** 2 + (self.z - zc) ** 2 + theta = atan2(self.z - zc, self.x - xc) # Angle relative to coil + + A -= Ic * 0.1 * log(r2) + + B = Ic * 0.2 / sqrt(r2) - A = 0.0 - Bx = 0.0 - Bz = 0.0 + Bx += B * sin(theta) + Bz -= B * cos(theta) - for c in self.coil_list: - xc, zc, Ic = c - rc = (xc**2 + zc**2) ** (0.5) - r2 = (self.x - xc) ** 2 + (self.z - zc) ** 2 - theta = atan2(self.z - zc, self.x - xc) # Angle relative to coil + By = Btor / self.x + self.Afunc = lambdify((self.x, self.z, self.phi), A, "numpy") - A -= Ic * 0.1 * log(r2) + self.Bxfunc = lambdify((self.x, self.z, self.phi), Bx, "numpy") + self.Bzfunc = lambdify((self.x, self.z, self.phi), Bz, "numpy") + self.Byfunc = lambdify((self.x, self.z, self.phi), By, "numpy") - B = Ic * 0.2 / sqrt(r2) + def Rfunc(self, x, z, phi): + return np.full(x.shape, x) - Bx += B * sin(theta) - Bz -= B * cos(theta) - By = Btor / self.x - self.Afunc = lambdify((self.x, self.z, self.phi), A, "numpy") +class DommaschkPotentials(MagneticField): + """A magnetic field generator using the Dommaschk potentials. + Parameters + ---------- + A: Coefficient matrix for the torodial and polidial harmonics. Form: (m,l,(a,b,c,d)) + R_0: major radius [m] + B_0: magnetic field on axis [T] - self.Bxfunc = lambdify((self.x, self.z, self.phi), Bx, "numpy") - self.Bzfunc = lambdify((self.x, self.z, self.phi), Bz, "numpy") - self.Byfunc = lambdify((self.x, self.z, self.phi), By, "numpy") + Important Methods + ----------------- + Bxfunc/Byfunc/Bzfunc(x,z,y): + Returns magnetic field in radial/torodial/z-direction + Sfunc(x,z,y): + Returns approximate magnetic surface invariant for Dommaschk potentials. + Use this to visualize flux surfaces - def Rfunc(self, x, z, phi): - return np.full(x.shape, x) - class DommaschkPotentials(MagneticField): - """A magnetic field generator using the Dommaschk potentials. + + """ + + def __init__(self, A, R_0=1.0, B_0=1.0): + self.R_0 = R_0 + self.B_0 = B_0 + + self.R = Symbol("R") + self.phi = Symbol("phi") + self.Z = Symbol("Z") + + self.m = Symbol("m") + self.l = Symbol("l") + self.n = Symbol("n") + self.k = Symbol("k") + + self.A = A + + self.P = ( + self.U(self.A) + .doit() + .subs([(self.R, self.R / self.R_0), (self.Z, self.Z / self.R_0)]) + ) + self.P_hat = ( + self.U_hat(self.A) + .doit() + .subs([(self.R, self.R / self.R_0), (self.Z, self.Z / self.R_0)]) + ) + + S = 0.5 * ( + log(self.R / self.R_0) ** 2 + (self.Z / self.R_0) ** 2 + ) - self.R / self.R_0 * ( + log(self.R / self.R_0) * self.R_0 * diff(self.P_hat, self.R) + + self.Z * self.R / self.R_0 * diff(self.P_hat, self.Z) + ) + + Bx = R_0 * diff(self.P, self.R) + By = R_0 / self.R * diff(self.P, self.phi) + Bz = R_0 * diff(self.P, self.Z) + + self.Sf = lambdify((self.R, self.phi, self.Z), S, "numpy") + + self.Bxf = lambdify((self.R, self.phi, self.Z), Bx, "numpy") + self.Byf = lambdify((self.R, self.phi, self.Z), By, "numpy") + self.Bzf = lambdify((self.R, self.phi, self.Z), Bz, "numpy") + + def Bxfunc(self, x, z, phi): + return self.Bxf(x, phi, z) / self.Byf(self.R_0, 0, 0) * self.B_0 + + def Byfunc(self, x, z, phi): + return self.Byf(x, phi, z) / self.Byf(self.R_0, 0, 0) * self.B_0 + + def Bzfunc(self, x, z, phi): + return self.Bzf(x, phi, z) / self.Byf(self.R_0, 0, 0) * self.B_0 + + def Sfunc(self, x, z, y): + """ Parameters ---------- - A: Coefficient matrix for the torodial and polidial harmonics. Form: (m,l,(a,b,c,d)) - R_0: major radius [m] - B_0: magnetic field on axis [T] - - Important Methods - ----------------- - Bxfunc/Byfunc/Bzfunc(x,z,y): - Returns magnetic field in radial/torodial/z-direction - Sfunc(x,z,y): - Returns approximate magnetic surface invariant for Dommaschk potentials. - Use this to visualize flux surfaces + x: radial coordinates normalized to R_0 + z: binormal coordinate + y: torodial angle normalized to 2*pi + Returns + ------- + Approximate magnetic surface invariant S at location (x,z,y). + This is from the original Dommaschk paper. + Use to visualize flux surfaces + """ + return self.Sf(x, y, z) + def Rfunc(self, x, z, phi): + """ + Parameters + ---------- + x: radial coordinates normalized to R_0 + z: binormal coordinate + y: torodial angle normalized to 2*pi + Returns + ------- + Radial coordinate x """ - def __init__(self, A, R_0=1.0, B_0=1.0): - self.R_0 = R_0 - self.B_0 = B_0 + return x - self.R = Symbol("R") - self.phi = Symbol("phi") - self.Z = Symbol("Z") + def CD(self, m, k): + """ + Parameters + ---------- + m: torodial harmonic + k: summation index in D - self.m = Symbol("m") - self.l = Symbol("l") - self.n = Symbol("n") - self.k = Symbol("k") + Returns: + -------- + Sympy function CD_mk (R) (Dirichlet boudary conditions) + """ - self.A = A + alpha = lambda n, b: ( + (-1.0) ** n / (gamma(b + n + 1) * gamma(n + 1) * 2.0 ** (2 * n + b)) + if (n >= 0) + else 0.0 + ) + alpha_st = lambda n, b: alpha(n, b) * (2 * n + b) - self.P = ( - self.U(self.A) - .doit() - .subs([(self.R, self.R / self.R_0), (self.Z, self.Z / self.R_0)]) - ) - self.P_hat = ( - self.U_hat(self.A) - .doit() - .subs([(self.R, self.R / self.R_0), (self.Z, self.Z / self.R_0)]) - ) + beta = lambda n, b: ( + gamma(b - n) / (gamma(n + 1) * 2.0 ** (2 * n - b + 1)) + if (n >= 0 and n < b) + else 0.0 + ) + beta_st = lambda n, b: beta(n, b) * (2 * n - b) - S = 0.5 * ( - log(self.R / self.R_0) ** 2 + (self.Z / self.R_0) ** 2 - ) - self.R / self.R_0 * ( - log(self.R / self.R_0) * self.R_0 * diff(self.P_hat, self.R) - + self.Z * self.R / self.R_0 * diff(self.P_hat, self.Z) + delta = lambda n, b: ( + alpha(n, b) * np.sum([1.0 / i + 1.0 / (b + i) for i in range(1, n + 1)]) / 2 + if (n > 0) + else 0.0 + ) + delta_st = lambda n, b: delta(n, b) * (2 * n + b) + + CD = log(1) + for j in range(k + 1): + CD += -( + alpha(j, m) + * ( + alpha_st(k - m - j, m) * log(self.R) + + delta_st(k - m - j, m) + - alpha(k - m - j, m) + ) + - delta(j, m) * alpha_st(k - m - j, m) + + alpha(j, m) * beta_st(k - j, m) + ) * self.R ** (2 * j + m) + beta(j, m) * alpha_st(k - j, m) * self.R ** ( + 2 * j - m ) - Bx = R_0 * diff(self.P, self.R) - By = R_0 / self.R * diff(self.P, self.phi) - Bz = R_0 * diff(self.P, self.Z) - - self.Sf = lambdify((self.R, self.phi, self.Z), S, "numpy") - - self.Bxf = lambdify((self.R, self.phi, self.Z), Bx, "numpy") - self.Byf = lambdify((self.R, self.phi, self.Z), By, "numpy") - self.Bzf = lambdify((self.R, self.phi, self.Z), Bz, "numpy") - - def Bxfunc(self, x, z, phi): - return self.Bxf(x, phi, z) / self.Byf(self.R_0, 0, 0) * self.B_0 - - def Byfunc(self, x, z, phi): - return self.Byf(x, phi, z) / self.Byf(self.R_0, 0, 0) * self.B_0 - - def Bzfunc(self, x, z, phi): - return self.Bzf(x, phi, z) / self.Byf(self.R_0, 0, 0) * self.B_0 - - def Sfunc(self, x, z, y): - """ - Parameters - ---------- - x: radial coordinates normalized to R_0 - z: binormal coordinate - y: torodial angle normalized to 2*pi - - Returns - ------- - Approximate magnetic surface invariant S at location (x,z,y). - This is from the original Dommaschk paper. - Use to visualize flux surfaces - """ - return self.Sf(x, y, z) - - def Rfunc(self, x, z, phi): - """ - Parameters - ---------- - x: radial coordinates normalized to R_0 - z: binormal coordinate - y: torodial angle normalized to 2*pi - - Returns - ------- - Radial coordinate x - """ - - return x - - def CD(self, m, k): - """ - Parameters - ---------- - m: torodial harmonic - k: summation index in D - - Returns: - -------- - Sympy function CD_mk (R) (Dirichlet boudary conditions) - """ - - alpha = lambda n, b: ( - (-1.0) ** n / (gamma(b + n + 1) * gamma(n + 1) * 2.0 ** (2 * n + b)) - if (n >= 0) - else 0.0 - ) - alpha_st = lambda n, b: alpha(n, b) * (2 * n + b) + return CD - beta = lambda n, b: ( - gamma(b - n) / (gamma(n + 1) * 2.0 ** (2 * n - b + 1)) - if (n >= 0 and n < b) - else 0.0 - ) - beta_st = lambda n, b: beta(n, b) * (2 * n - b) - - delta = lambda n, b: ( - alpha(n, b) - * np.sum([1.0 / i + 1.0 / (b + i) for i in range(1, n + 1)]) - / 2 - if (n > 0) - else 0.0 - ) - delta_st = lambda n, b: delta(n, b) * (2 * n + b) - - CD = log(1) - for j in range(k + 1): - CD += -( - alpha(j, m) - * ( - alpha_st(k - m - j, m) * log(self.R) - + delta_st(k - m - j, m) - - alpha(k - m - j, m) - ) - - delta(j, m) * alpha_st(k - m - j, m) - + alpha(j, m) * beta_st(k - j, m) - ) * self.R ** (2 * j + m) + beta(j, m) * alpha_st( - k - j, m - ) * self.R ** ( - 2 * j - m - ) + def CN(self, m, k): + """ + Parameters + ---------- + m: torodial harmonic + k: summation index in N - return CD + Returns: + -------- + Sympy function CN_mk (R) (Neumann boundary conditions) + """ - def CN(self, m, k): - """ - Parameters - ---------- - m: torodial harmonic - k: summation index in N + alpha = lambda n, b: ( + (-1.0) ** n / (gamma(b + n + 1) * gamma(n + 1) * 2.0 ** (2 * n + b)) + if (n >= 0) + else 0.0 + ) + alpha_st = lambda n, b: alpha(n, b) * (2 * n + b) - Returns: - -------- - Sympy function CN_mk (R) (Neumann boundary conditions) - """ + beta = lambda n, b: ( + gamma(b - n) / (gamma(n + 1) * 2.0 ** (2 * n - b + 1)) + if (n >= 0 and n < b) + else 0.0 + ) + beta_st = lambda n, b: beta(n, b) * (2 * n - b) - alpha = lambda n, b: ( - (-1.0) ** n / (gamma(b + n + 1) * gamma(n + 1) * 2.0 ** (2 * n + b)) - if (n >= 0) - else 0.0 + delta = lambda n, b: ( + alpha(n, b) * np.sum([1.0 / i + 1.0 / (b + i) for i in range(1, n + 1)]) / 2 + if (n > 0) + else 0.0 + ) + delta_st = lambda n, b: delta(n, b) * (2 * n + b) + + CN = log(1) + for j in range(k + 1): + CN += ( + alpha(j, m) * (alpha(k - m - j, m) * log(self.R) + delta(k - m - j, m)) + - delta(j, m) * alpha(k - m - j, m) + + alpha(j, m) * beta(k - j, m) + ) * self.R ** (2 * j + m) - beta(j, m) * alpha(k - j, m) * self.R ** ( + 2 * j - m ) - alpha_st = lambda n, b: alpha(n, b) * (2 * n + b) - beta = lambda n, b: ( - gamma(b - n) / (gamma(n + 1) * 2.0 ** (2 * n - b + 1)) - if (n >= 0 and n < b) - else 0.0 - ) - beta_st = lambda n, b: beta(n, b) * (2 * n - b) - - delta = lambda n, b: ( - alpha(n, b) - * np.sum([1.0 / i + 1.0 / (b + i) for i in range(1, n + 1)]) - / 2 - if (n > 0) - else 0.0 - ) - delta_st = lambda n, b: delta(n, b) * (2 * n + b) - - CN = log(1) - for j in range(k + 1): - CN += ( - alpha(j, m) - * (alpha(k - m - j, m) * log(self.R) + delta(k - m - j, m)) - - delta(j, m) * alpha(k - m - j, m) - + alpha(j, m) * beta(k - j, m) - ) * self.R ** (2 * j + m) - beta(j, m) * alpha(k - j, m) * self.R ** ( - 2 * j - m - ) + return CN - return CN - - def D(self, m, n): - """ - Parameters - ---------- - m: torodial mode number - n: summation index in V - - Returns: - -------- - Sympy function D_mn (R, Z) (Dirichlet boundary conditions) - """ - - D = log(1) - k_arr = np.arange(0, int(n / 2) + 1, 1) - - for k in k_arr: - D += (self.Z ** (n - 2 * k)) / factorial(n - 2 * k) * self.CD(m, k) - - return D - - def N(self, m, n): - """ - Parameters - ---------- - m: torodial mode number - n: summation index in V - - Returns: - -------- - Sympy function N_mn (R, Z) (Neumann boundary conditions) - """ - - N = log(1) - k_arr = np.arange(0, int(n / 2) + 1, 1) - - for k in k_arr: - N += (self.Z ** (n - 2 * k)) / factorial(n - 2 * k) * self.CN(m, k) - - return N - - def V(self, m, l, a, b, c, d): - """ - Parameters - ---------- - m: torodial mode number - l: polodial mode number - a,b,c,d: Coefficients for m,l-th Dommaschk potential (elements of matrix A) - - Returns: - -------- - Sympy function V_ml - """ - - V = (a * cos(m * self.phi) + b * sin(m * self.phi)) * self.D(m, l) + ( - c * cos(m * self.phi) + d * sin(m * self.phi) - ) * self.N(m, l - 1) - - return V - - def U(self, A): - """ - Parameters - ---------- - A: Coefficient matrix for the torodial and polidial harmonics. Form: (m,l,(a,b,c,d)) - - Returns - ----------------- - U: Superposition of all modes given in A - - """ - U = self.phi - for i in range(A.shape[0]): - for j in range(A.shape[1]): - if A[i, j, 0] or A[i, j, 1] or A[i, j, 2] or A[i, j, 3] != 0: - U += self.V( - i, j, A[i, j, 0], A[i, j, 1], A[i, j, 2], A[i, j, 3] - ) - - return U - - def V_hat(self, m, l, a, b, c, d): - """ - Parameters - ---------- - m: torodial mode number - l: polodial mode number - a,b,c,d: Coefficients for m,l-th Dommaschk potential (elements of matrix A) - - Returns: - -------- - Sympy function V_hat_ml; Similar to V; needed for calculation of magnetic surface invariant S - """ - - V = ( - a * cos(m * self.phi - np.pi / 2) + b * sin(m * self.phi - np.pi / 2) - ) * self.D(m, l) + ( - c * cos(m * self.phi - np.pi / 2) + d * sin(m * self.phi - np.pi / 2) - ) * self.N( - m, l - 1 - ) + def D(self, m, n): + """ + Parameters + ---------- + m: torodial mode number + n: summation index in V - return V - - def U_hat(self, A): - """ - Parameters - ---------- - A: Coefficient matrix for the torodial and polidial harmonics. Form: (m,l,(a,b,c,d)) - - Returns - ----------------- - U: Superposition of all modes given in A - - """ - - U = log(1) - for i in range(A.shape[0]): - if i != 0: - i_inv = 1 / i - else: - i_inv = np.nan - for j in range(A.shape[1]): - if A[i, j, 0] or A[i, j, 1] or A[i, j, 2] or A[i, j, 3] != 0: - U += self.V_hat( - i, j, A[i, j, 0], A[i, j, 1], A[i, j, 2], A[i, j, 3] - ) * Piecewise((self.phi, i == 0), (i_inv, i > 0)) - - return U - - class Screwpinch(MagneticField): - def __init__( - self, xcentre=1.5, zcentre=0.0, shear=0, yperiod=2 * np.pi, Btor=1.0 - ): - self.x = Symbol("x") - self.z = Symbol("z") - self.y = Symbol("y") - self.r = Symbol("r") - self.r = ((self.x - xcentre) ** 2 + (self.z - zcentre) ** 2) ** (0.5) + Returns: + -------- + Sympy function D_mn (R, Z) (Dirichlet boundary conditions) + """ - self.phi = Symbol("phi") + D = log(1) + k_arr = np.arange(0, int(n / 2) + 1, 1) - alpha = shear - self.theta = atan2(self.z - zcentre, self.x - xcentre) - A = alpha * self.r**2 - Bx = -alpha * self.r * self.r * sin(self.theta) - Bz = alpha * self.r * self.r * cos(self.theta) - By = Btor / self.x + for k in k_arr: + D += (self.Z ** (n - 2 * k)) / factorial(n - 2 * k) * self.CD(m, k) - self.Afunc = lambdify((self.x, self.z, self.phi), A, "numpy") - self.Bxfunc = lambdify((self.x, self.z, self.phi), Bx, "numpy") - self.Bzfunc = lambdify((self.x, self.z, self.phi), Bz, "numpy") - self.Byfunc = lambdify((self.x, self.z, self.phi), By, "numpy") + return D - def Rfunc(self, x, z, phi): - return np.full(x.shape, x) + def N(self, m, n): + """ + Parameters + ---------- + m: torodial mode number + n: summation index in V + + Returns: + -------- + Sympy function N_mn (R, Z) (Neumann boundary conditions) + """ + + N = log(1) + k_arr = np.arange(0, int(n / 2) + 1, 1) -except ImportError: + for k in k_arr: + N += (self.Z ** (n - 2 * k)) / factorial(n - 2 * k) * self.CN(m, k) - class StraightStellarator(MagneticField): + return N + + def V(self, m, l, a, b, c, d): """ - Invalid StraightStellarator, since no Sympy module. + Parameters + ---------- + m: torodial mode number + l: polodial mode number + a,b,c,d: Coefficients for m,l-th Dommaschk potential (elements of matrix A) - Rather than printing an error on startup, which may - be missed or ignored, this raises - an exception if StraightStellarator is ever used. + Returns: + -------- + Sympy function V_ml """ - def __init__(self, *args, **kwargs): - raise ImportError( - "No Sympy module: Can't generate StraightStellarator fields" - ) + V = (a * cos(m * self.phi) + b * sin(m * self.phi)) * self.D(m, l) + ( + c * cos(m * self.phi) + d * sin(m * self.phi) + ) * self.N(m, l - 1) + + return V + + def U(self, A): + """ + Parameters + ---------- + A: Coefficient matrix for the torodial and polidial harmonics. Form: (m,l,(a,b,c,d)) + + Returns + ----------------- + U: Superposition of all modes given in A - class RotatingEllipse(MagneticField): """ - Invalid RotatingEllipse, since no Sympy module. - Rather than printing an error on startup, which may - be missed or ignored, this raises - an exception if StraightStellarator is ever used. + U = self.phi + for i in range(A.shape[0]): + for j in range(A.shape[1]): + if A[i, j, 0] or A[i, j, 1] or A[i, j, 2] or A[i, j, 3] != 0: + U += self.V(i, j, A[i, j, 0], A[i, j, 1], A[i, j, 2], A[i, j, 3]) + + return U + + def V_hat(self, m, l, a, b, c, d): """ + Parameters + ---------- + m: torodial mode number + l: polodial mode number + a,b,c,d: Coefficients for m,l-th Dommaschk potential (elements of matrix A) - def __init__(self, *args, **kwargs): - raise ImportError("No Sympy module: Can't generate RotatingEllipse fields") + Returns: + -------- + Sympy function V_hat_ml; Similar to V; needed for calculation of magnetic surface invariant S + """ - class Screwpinch(MagneticField): + V = ( + a * cos(m * self.phi - np.pi / 2) + b * sin(m * self.phi - np.pi / 2) + ) * self.D(m, l) + ( + c * cos(m * self.phi - np.pi / 2) + d * sin(m * self.phi - np.pi / 2) + ) * self.N( + m, l - 1 + ) + + return V + + def U_hat(self, A): """ - Invalid screwpinch, since no Sympy module. - Rather than printing an error on startup, which may - be missed or ignored, this raises - an exception if StraightStellarator is ever used. + Parameters + ---------- + A: Coefficient matrix for the torodial and polidial harmonics. Form: (m,l,(a,b,c,d)) + + Returns + ----------------- + U: Superposition of all modes given in A + """ - def __init__(self, *args, **kwargs): - raise ImportError("No Sympy module: Can't generate screwpinch fields") + U = log(1) + for i in range(A.shape[0]): + if i != 0: + i_inv = 1 / i + else: + i_inv = np.nan + for j in range(A.shape[1]): + if A[i, j, 0] or A[i, j, 1] or A[i, j, 2] or A[i, j, 3] != 0: + U += self.V_hat( + i, j, A[i, j, 0], A[i, j, 1], A[i, j, 2], A[i, j, 3] + ) * Piecewise((self.phi, i == 0), (i_inv, i > 0)) + + return U + + +class Screwpinch(MagneticField): + def __init__(self, xcentre=1.5, zcentre=0.0, shear=0, yperiod=2 * np.pi, Btor=1.0): + self.x = Symbol("x") + self.z = Symbol("z") + self.y = Symbol("y") + self.r = Symbol("r") + self.r = ((self.x - xcentre) ** 2 + (self.z - zcentre) ** 2) ** (0.5) + + self.phi = Symbol("phi") + + alpha = shear + self.theta = atan2(self.z - zcentre, self.x - xcentre) + A = alpha * self.r**2 + Bx = -alpha * self.r * self.r * sin(self.theta) + Bz = alpha * self.r * self.r * cos(self.theta) + By = Btor / self.x + + self.Afunc = lambdify((self.x, self.z, self.phi), A, "numpy") + self.Bxfunc = lambdify((self.x, self.z, self.phi), Bx, "numpy") + self.Bzfunc = lambdify((self.x, self.z, self.phi), Bz, "numpy") + self.Byfunc = lambdify((self.x, self.z, self.phi), By, "numpy") + + def Rfunc(self, x, z, phi): + return np.full(x.shape, x) class VMEC(MagneticField): diff --git a/zoidberg/plot.py b/zoidberg/plot.py index fe3e96d..13bb70e 100644 --- a/zoidberg/plot.py +++ b/zoidberg/plot.py @@ -242,7 +242,7 @@ def plot_streamlines(grid, magnetic_field, y_slice=0, width=None, **kwargs): magnetic_field.bx[full_slice].T, magnetic_field.bz[full_slice].T, linewidth=linewidth, - **kwargs + **kwargs, ) ax.set_xlabel("Radius [m]", fontsize=20) diff --git a/zoidberg/rzline.py b/zoidberg/rzline.py index cc8b1b3..5146e27 100644 --- a/zoidberg/rzline.py +++ b/zoidberg/rzline.py @@ -8,13 +8,7 @@ import numpy as np from numpy import append, argmin, cos, linspace, pi, sin, sqrt - -try: - # New scipy changed the name - from scipy.integrate import cumulative_trapezoid as cumtrapz -except ImportError: - from scipy.integrate import cumtrapz - +from scipy.integrate import cumulative_trapezoid as cumtrapz from scipy.interpolate import interp1d, splev, splrep try: