From 58d2a2dfb1fe043b39864aa1ffc39bc14e7a8243 Mon Sep 17 00:00:00 2001 From: George Holt Date: Mon, 12 Aug 2024 14:35:39 -0700 Subject: [PATCH 1/2] style: Black and isort formatting --- 01-freeboundary.py | 35 ++-- 02-read-geqdsk.py | 7 +- 03-mast.py | 48 +++-- 04-read-mast-geqdsk.py | 8 +- 05-fixed-boundary.py | 29 +-- 06-xpoints.py | 32 +-- 07-increase-resolution.py | 8 +- 08-mast-upgrade.py | 116 ++++++----- 09-metal-wall.py | 49 +++-- 10-mastu-connection.py | 12 +- 11-optimise-coils.py | 63 +++--- docs/conf.py | 67 ++++--- freegs4e/__init__.py | 14 +- freegs4e/_aeqdsk.py | 32 ++- freegs4e/_divgeo.py | 1 + freegs4e/_geqdsk.py | 16 +- freegs4e/bilinear_interpolation.py | 63 +++--- freegs4e/boundary.py | 44 ++++- freegs4e/coil.py | 29 +-- freegs4e/control.py | 7 +- freegs4e/critical.py | 278 ++++++++++++++++----------- freegs4e/critical_old.py | 265 +++++++++++++++---------- freegs4e/divgeo.py | 3 +- freegs4e/dump.py | 38 ++-- freegs4e/equilibrium.py | 100 ++++++---- freegs4e/faster_shape.py | 8 +- freegs4e/fieldtracer.py | 28 ++- freegs4e/geqdsk.py | 63 +++--- freegs4e/gradshafranov.py | 38 ++-- freegs4e/jtor.py | 299 +++++++++++++++++------------ freegs4e/machine.py | 123 +++++++++--- freegs4e/multi_coil.py | 24 +-- freegs4e/multigrid.py | 44 +++-- freegs4e/optimise.py | 27 ++- freegs4e/optimiser.py | 15 +- freegs4e/picard.py | 12 +- freegs4e/plotting.py | 58 ++++-- freegs4e/quadrature.py | 40 +++- freegs4e/shaped_coil.py | 12 +- freegs4e/test_aeqdsk.py | 3 +- freegs4e/test_critical.py | 10 +- freegs4e/test_equilibrium.py | 21 +- freegs4e/test_geqdsk.py | 4 +- freegs4e/test_jtor.py | 2 +- freegs4e/test_linearsolve.py | 8 +- freegs4e/test_machine.py | 8 +- freegs4e/test_multi_coil.py | 14 +- freegs4e/test_optimise.py | 3 +- freegs4e/test_optimiser.py | 9 +- freegs4e/test_polygons.py | 19 +- freegs4e/test_quadrature.py | 22 ++- freegs4e/test_readwrite.py | 18 +- freegs4e/test_shaped_coil.py | 18 +- test-01-compare.py | 74 ++++--- test-convergence.py | 13 +- test_fileutils.py | 46 +++-- 56 files changed, 1511 insertions(+), 936 deletions(-) diff --git a/01-freeboundary.py b/01-freeboundary.py index b0cad0d..b904a7b 100644 --- a/01-freeboundary.py +++ b/01-freeboundary.py @@ -8,19 +8,24 @@ tokamak = freegs4e.machine.TestTokamak() -eq = freegs4e.Equilibrium(tokamak=tokamak, - Rmin=0.1, Rmax=2.0, # Radial domain - Zmin=-1.0, Zmax=1.0, # Height range - nx=65, ny=65, # Number of grid points - boundary=freegs4e.boundary.freeBoundaryHagenow) # Boundary condition +eq = freegs4e.Equilibrium( + tokamak=tokamak, + Rmin=0.1, + Rmax=2.0, # Radial domain + Zmin=-1.0, + Zmax=1.0, # Height range + nx=65, + ny=65, # Number of grid points + boundary=freegs4e.boundary.freeBoundaryHagenow, +) # Boundary condition ######################################### # Plasma profiles -profiles = freegs4e.jtor.ConstrainPaxisIp(1e3, # Plasma pressure on axis [Pascals] - 2e5, # Plasma current [Amps] - 2.0) # Vacuum f=R*Bt +profiles = freegs4e.jtor.ConstrainPaxisIp( + 1e3, 2e5, 2.0 # Plasma pressure on axis [Pascals] # Plasma current [Amps] +) # Vacuum f=R*Bt ######################################### # Coil current constraints @@ -28,19 +33,20 @@ # Specify locations of the X-points # to use to constrain coil currents -xpoints = [(1.1, -0.6), # (R,Z) locations of X-points - (1.1, 0.8)] +xpoints = [(1.1, -0.6), (1.1, 0.8)] # (R,Z) locations of X-points -isoflux = [(1.1,-0.6, 1.1,0.6)] # (R1,Z1, R2,Z2) pair of locations +isoflux = [(1.1, -0.6, 1.1, 0.6)] # (R1,Z1, R2,Z2) pair of locations constrain = freegs4e.control.constrain(xpoints=xpoints, isoflux=isoflux) ######################################### # Nonlinear solve -freegs4e.solve(eq, # The equilibrium to adjust - profiles, # The toroidal current profile function - constrain) # Constraint function to set coil currents +freegs4e.solve( + eq, # The equilibrium to adjust + profiles, # The toroidal current profile function + constrain, +) # Constraint function to set coil currents # eq now contains the solution @@ -78,6 +84,7 @@ # Safety factor import matplotlib.pyplot as plt + plt.plot(*eq.q()) plt.xlabel(r"Normalised $\psi$") plt.ylabel("Safety factor") diff --git a/02-read-geqdsk.py b/02-read-geqdsk.py index 11bf5f1..64acea5 100644 --- a/02-read-geqdsk.py +++ b/02-read-geqdsk.py @@ -1,11 +1,10 @@ -from freegs4e import geqdsk -from freegs4e import machine +from freegs4e import geqdsk, machine from freegs4e.plotting import plotEquilibrium -#tokamak = machine.MAST_sym() +# tokamak = machine.MAST_sym() tokamak = machine.TestTokamak() -#with open("g014220.00200") as f: +# with open("g014220.00200") as f: with open("lsn.geqdsk") as f: eq = geqdsk.read(f, tokamak, show=True) diff --git a/03-mast.py b/03-mast.py index f3e5f89..e84910e 100644 --- a/03-mast.py +++ b/03-mast.py @@ -8,17 +8,22 @@ tokamak = freegs4e.machine.MAST() -eq = freegs4e.Equilibrium(tokamak=tokamak, - Rmin=0.1, Rmax=2.0, # Radial domain - Zmin=-2.0, Zmax=2.0, # Height range - nx=65, ny=65) # Number of grid points +eq = freegs4e.Equilibrium( + tokamak=tokamak, + Rmin=0.1, + Rmax=2.0, # Radial domain + Zmin=-2.0, + Zmax=2.0, # Height range + nx=65, + ny=65, +) # Number of grid points ######################################### # Plasma profiles -profiles = freegs4e.jtor.ConstrainPaxisIp(3e3, # Plasma pressure on axis [Pascals] - 7e5, # Plasma current [Amps] - 0.4) # vacuum f = R*Bt +profiles = freegs4e.jtor.ConstrainPaxisIp( + 3e3, 7e5, 0.4 # Plasma pressure on axis [Pascals] # Plasma current [Amps] +) # vacuum f = R*Bt ######################################### # Coil current constraints @@ -26,26 +31,30 @@ # Specify locations of the X-points # to use to constrain coil currents -xpoints = [(0.7, -1.1), # (R,Z) locations of X-points - (0.7, 1.1)] +xpoints = [(0.7, -1.1), (0.7, 1.1)] # (R,Z) locations of X-points -isoflux = [(0.7,-1.1, 1.45, 0.0) # Outboard midplane, lower X-point - ,(0.7,1.1, 1.45, 0.0) # Outboard midplane, upper X-point -# ,(0.7,-1.1, 1.5,-1.9) # Lower X-point, lower outer leg -# ,(0.7,1.1, 1.5, 1.9) # Upper X-point, upper outer leg - ] +isoflux = [ + (0.7, -1.1, 1.45, 0.0), # Outboard midplane, lower X-point + (0.7, 1.1, 1.45, 0.0), # Outboard midplane, upper X-point + # ,(0.7,-1.1, 1.5,-1.9) # Lower X-point, lower outer leg + # ,(0.7,1.1, 1.5, 1.9) # Upper X-point, upper outer leg +] -constrain = freegs4e.control.constrain(xpoints=xpoints, gamma=1e-12, isoflux=isoflux) +constrain = freegs4e.control.constrain( + xpoints=xpoints, gamma=1e-12, isoflux=isoflux +) constrain(eq) ######################################### # Nonlinear solve -freegs4e.solve(eq, # The equilibrium to adjust - profiles, # The plasma profiles - constrain, # Plasma control constraints - show=True) # Shows results at each nonlinear iteration +freegs4e.solve( + eq, # The equilibrium to adjust + profiles, # The plasma profiles + constrain, # Plasma control constraints + show=True, +) # Shows results at each nonlinear iteration # eq now contains the solution @@ -68,4 +77,5 @@ # Call matplotlib show so plot pauses import matplotlib.pyplot as plt + plt.show() diff --git a/04-read-mast-geqdsk.py b/04-read-mast-geqdsk.py index 93a6690..737e25d 100644 --- a/04-read-mast-geqdsk.py +++ b/04-read-mast-geqdsk.py @@ -1,13 +1,11 @@ -from freegs4e import geqdsk -from freegs4e import machine +from freegs4e import geqdsk, machine from freegs4e.plotting import plotEquilibrium # Reading MAST equilibrium, up-down symmetric coils tokamak = machine.MAST() -#with open("g014220.00200") as f: +# with open("g014220.00200") as f: with open("mast.geqdsk") as f: - eq = geqdsk.read(f, tokamak, show=True) + eq = geqdsk.read(f, tokamak, show=True) plotEquilibrium(eq) - diff --git a/05-fixed-boundary.py b/05-fixed-boundary.py index dcb34bd..7d4d4c7 100644 --- a/05-fixed-boundary.py +++ b/05-fixed-boundary.py @@ -10,18 +10,24 @@ # Boundary conditions import freegs4e.boundary as boundary -profiles = freegs4e.jtor.ConstrainPaxisIp(1e3, # Plasma pressure on axis [Pascals] - 1e5, # Plasma current [Amps] - 1.0) # fvac = R*Bt - -eq = freegs4e.Equilibrium(Rmin=0.1, Rmax=2.0, - Zmin=-1.0, Zmax=1.0, - nx=65, ny=65, - boundary=boundary.fixedBoundary) +profiles = freegs4e.jtor.ConstrainPaxisIp( + 1e3, 1e5, 1.0 # Plasma pressure on axis [Pascals] # Plasma current [Amps] +) # fvac = R*Bt + +eq = freegs4e.Equilibrium( + Rmin=0.1, + Rmax=2.0, + Zmin=-1.0, + Zmax=1.0, + nx=65, + ny=65, + boundary=boundary.fixedBoundary, +) # Nonlinear solver for Grad-Shafranov equation -freegs4e.solve(eq, # The equilibrium to adjust - profiles) # The toroidal current profile function +freegs4e.solve( + eq, profiles # The equilibrium to adjust +) # The toroidal current profile function print("Done!") @@ -31,6 +37,5 @@ # Plot equilibrium from freegs4e.plotting import plotEquilibrium -plotEquilibrium(eq) - +plotEquilibrium(eq) diff --git a/06-xpoints.py b/06-xpoints.py index c303c2f..e111c00 100644 --- a/06-xpoints.py +++ b/06-xpoints.py @@ -2,23 +2,22 @@ # # Example demonstrating functions for creating and finding X-points +import matplotlib.pyplot as plt + import freegs4e # Plotting routines -from freegs4e.plotting import plotEquilibrium, plotCoils, plotConstraints - -import matplotlib.pyplot as plt +from freegs4e.plotting import plotCoils, plotConstraints, plotEquilibrium tokamak = freegs4e.machine.TestTokamak() -eq = freegs4e.Equilibrium(tokamak=tokamak, nx=256,ny=256) +eq = freegs4e.Equilibrium(tokamak=tokamak, nx=256, ny=256) ########################################################## # Calculate currents in coils to create X-points # in specified locations -# +# -xpoints = [(1.1, -0.8), # (R,Z) locations of X-points - (1.1, 0.8)] +xpoints = [(1.1, -0.8), (1.1, 0.8)] # (R,Z) locations of X-points control = freegs4e.control.constrain(xpoints=xpoints) control(eq) # Apply control to Equilibrium eq @@ -34,27 +33,28 @@ ########################################################## # Find critical points (O- and X-points) -# -# +# +# import freegs4e.critical as critical + opt, xpt = critical.find_critical(eq.R, eq.Z, psi) print("=> Found O- and X-points") ax = plotEquilibrium(eq, show=False, oxpoints=False) -for r,z,_ in xpt: - ax.plot(r,z,'ro') -for r,z,_ in opt: - ax.plot(r,z,'go') +for r, z, _ in xpt: + ax.plot(r, z, "ro") +for r, z, _ in opt: + ax.plot(r, z, "go") psi_bndry = xpt[0][2] -sep_contour=ax.contour(eq.R, eq.Z,psi, levels=[psi_bndry], colors='r') +sep_contour = ax.contour(eq.R, eq.Z, psi, levels=[psi_bndry], colors="r") plt.show() ########################################################## # Create a mask array, 1 in the core and 0 outside -# -# +# +# mask = critical.core_mask(eq.R, eq.Z, psi, opt, xpt) diff --git a/07-increase-resolution.py b/07-increase-resolution.py index 695a159..a97bf7b 100644 --- a/07-increase-resolution.py +++ b/07-increase-resolution.py @@ -1,4 +1,3 @@ - import freegs4e from freegs4e import geqdsk from freegs4e.equilibrium import refine @@ -7,9 +6,9 @@ # Reading MAST equilibrium, up-down symmetric coils tokamak = freegs4e.machine.MAST() -#with open("g014220.00200") as f: +# with open("g014220.00200") as f: with open("mast.geqdsk") as f: - eq = geqdsk.read(f, tokamak, show=True) + eq = geqdsk.read(f, tokamak, show=True) # Increase resolution by a factor of 2 eq2 = refine(eq) @@ -19,7 +18,6 @@ # Save to G-EQDSK with open("mast-highres.geqdsk", "w") as f: - geqdsk.write(eq2, f) + geqdsk.write(eq2, f) plotEquilibrium(eq2) - diff --git a/08-mast-upgrade.py b/08-mast-upgrade.py index fb2ee52..f01037d 100644 --- a/08-mast-upgrade.py +++ b/08-mast-upgrade.py @@ -10,19 +10,26 @@ tokamak = freegs4e.machine.MASTU() -eq = freegs4e.Equilibrium(tokamak=tokamak, - Rmin=0.1, Rmax=2.0, # Radial domain - Zmin=-2.1, Zmax=2.1, # Height range - nx=65, ny=65) # Number of grid points +eq = freegs4e.Equilibrium( + tokamak=tokamak, + Rmin=0.1, + Rmax=2.0, # Radial domain + Zmin=-2.1, + Zmax=2.1, # Height range + nx=65, + ny=65, +) # Number of grid points ######################################### # Plasma profiles -profiles = freegs4e.jtor.ConstrainPaxisIp(6e4, # Plasma pressure on axis [Pascals] - 1e6, # Plasma current [Amps] - 0.65, # vacuum f = R*Bt - alpha_m = 1.0, - alpha_n = 2.0) +profiles = freegs4e.jtor.ConstrainPaxisIp( + 6e4, # Plasma pressure on axis [Pascals] + 1e6, # Plasma current [Amps] + 0.65, # vacuum f = R*Bt + alpha_m=1.0, + alpha_n=2.0, +) ######################################### # Coil current constraints @@ -33,56 +40,58 @@ Rx = 0.509 Zx = 1.291 -Rmid = 1.34 # Outboard midplane +Rmid = 1.34 # Outboard midplane Rin = 0.3581 # Inboard midplane -xpoints = [(Rx, -Zx), # (R,Z) locations of X-points - (Rx, Zx)] +xpoints = [(Rx, -Zx), (Rx, Zx)] # (R,Z) locations of X-points -isoflux = [(Rx,-Zx, Rmid, 0.0) # Outboard midplane, lower X-point - ,(Rx,Zx, Rmid, 0.0) # Outboard midplane, upper X-point +isoflux = [ + (Rx, -Zx, Rmid, 0.0), # Outboard midplane, lower X-point + (Rx, Zx, Rmid, 0.0) # Outboard midplane, upper X-point + # Link inner and outer midplane locations + , + (Rmid, 0.0, Rin, 0.0) + # Separatrix in the divertor chamber + , + (Rx, -Zx, 0.95, -1.77), + (Rx, Zx, 0.95, 1.77), +] - # Link inner and outer midplane locations - ,(Rmid, 0.0, Rin, 0.0) - - # Separatrix in the divertor chamber - ,(Rx,-Zx, 0.95, -1.77) - ,(Rx, Zx, 0.95, 1.77) - ] - -constrain = freegs4e.control.constrain(xpoints=xpoints, gamma=8e-6, isoflux=isoflux) +constrain = freegs4e.control.constrain( + xpoints=xpoints, gamma=8e-6, isoflux=isoflux +) constrain(eq) ######################################### # Nonlinear solve -freegs4e.solve(eq, # The equilibrium to adjust - profiles, # The plasma profiles - constrain, # Plasma control constraints - show=True) # Shows results at each nonlinear iteration +freegs4e.solve( + eq, # The equilibrium to adjust + profiles, # The plasma profiles + constrain, # Plasma control constraints + show=True, +) # Shows results at each nonlinear iteration ######################################### # Now adjust the equilibrium manually -# - -isoflux = [(Rx,-Zx, Rmid, 0.0) # Outboard midplane, lower X-point - ,(Rx,Zx, Rmid, 0.0) # Outboard midplane, upper X-point - - ,(Rmid, 0.0, Rin, 0.0) - - ,(Rx,-Zx, 0.95, -1.77) - ,(Rx, Zx, 0.95, 1.77) - - ,(Rx,-Zx, 0.76, -1.58) - ,(Rx, Zx, 0.76, 1.58) - - ,(Rx,-Zx, 1.25, -1.8) - ,(Rx, Zx, 1.25, 1.8) - - ] +# -constrain = freegs4e.control.constrain(xpoints=xpoints, gamma=1e-12, isoflux=isoflux) +isoflux = [ + (Rx, -Zx, Rmid, 0.0), # Outboard midplane, lower X-point + (Rx, Zx, Rmid, 0.0), # Outboard midplane, upper X-point + (Rmid, 0.0, Rin, 0.0), + (Rx, -Zx, 0.95, -1.77), + (Rx, Zx, 0.95, 1.77), + (Rx, -Zx, 0.76, -1.58), + (Rx, Zx, 0.76, 1.58), + (Rx, -Zx, 1.25, -1.8), + (Rx, Zx, 1.25, 1.8), +] + +constrain = freegs4e.control.constrain( + xpoints=xpoints, gamma=1e-12, isoflux=isoflux +) # Turn off feedback control for all coils for label, coil in tokamak.coils: @@ -90,7 +99,7 @@ # Centre column coil tokamak["Pc"].current = -4e4 - + # Turn on vertical feedback control tokamak["P6"].control = True @@ -106,16 +115,18 @@ tokamak["D1"].current = 1000.0 -# Coil in outer corner -tokamak["D5"].current = 1500. +# Coil in outer corner +tokamak["D5"].current = 1500.0 # Coil at bottom centre tokamak["D3"].current = 2800 -freegs4e.solve(eq, # The equilibrium to adjust - profiles, # The plasma profiles - constrain, # Plasma control constraints - show=True) # Shows results at each nonlinear iteration +freegs4e.solve( + eq, # The equilibrium to adjust + profiles, # The plasma profiles + constrain, # Plasma control constraints + show=True, +) # Shows results at each nonlinear iteration # eq now contains the solution @@ -142,4 +153,5 @@ # Call matplotlib show so plot pauses import matplotlib.pyplot as plt + plt.show() diff --git a/09-metal-wall.py b/09-metal-wall.py index ff4f1e2..5939633 100644 --- a/09-metal-wall.py +++ b/09-metal-wall.py @@ -4,19 +4,20 @@ # This is done by creating a ring of coils, with feedback control setting # the poloidal flux to zero at the location of each coil. -import freegs4e import numpy as np +import freegs4e + ######################################### # Create a circular metal wall by using a ring of coils and psi constraints -R0 = 1.0 # Middle of the circle +R0 = 1.0 # Middle of the circle rwall = 0.5 # Radius of the circular wall -npoints = 200 # Number of points on the wall +npoints = 200 # Number of points on the wall # Poloidal angles -thetas = np.linspace(0, 2*np.pi, npoints, endpoint=False) +thetas = np.linspace(0, 2 * np.pi, npoints, endpoint=False) # Points on the wall Rwalls = R0 + rwall * np.cos(thetas) @@ -26,40 +27,49 @@ # Create the machine, which specifies coil locations # and equilibrium, specifying the domain to solve over -coils = [ ("wall_"+str(theta), freegs4e.machine.Coil(R, Z)) - for theta, R, Z in zip(thetas, Rwalls, Zwalls) ] +coils = [ + ("wall_" + str(theta), freegs4e.machine.Coil(R, Z)) + for theta, R, Z in zip(thetas, Rwalls, Zwalls) +] tokamak = freegs4e.machine.Machine(coils) -eq = freegs4e.Equilibrium(tokamak=tokamak, - Rmin=0.1, Rmax=2.0, # Radial domain - Zmin=-1.0, Zmax=1.0, # Height range - nx=65, ny=65, # Number of grid points - boundary=freegs4e.boundary.freeBoundaryHagenow) # Boundary condition +eq = freegs4e.Equilibrium( + tokamak=tokamak, + Rmin=0.1, + Rmax=2.0, # Radial domain + Zmin=-1.0, + Zmax=1.0, # Height range + nx=65, + ny=65, # Number of grid points + boundary=freegs4e.boundary.freeBoundaryHagenow, +) # Boundary condition ######################################### # Plasma profiles -profiles = freegs4e.jtor.ConstrainPaxisIp(1e4, # Plasma pressure on axis [Pascals] - 1e6, # Plasma current [Amps] - 2.0) # Vacuum f=R*Bt +profiles = freegs4e.jtor.ConstrainPaxisIp( + 1e4, 1e6, 2.0 # Plasma pressure on axis [Pascals] # Plasma current [Amps] +) # Vacuum f=R*Bt ######################################### # Coil current constraints # # Same location as the coils -psivals = [ (R, Z, 0.0) for R, Z in zip(Rwalls, Zwalls) ] +psivals = [(R, Z, 0.0) for R, Z in zip(Rwalls, Zwalls)] constrain = freegs4e.control.constrain(psivals=psivals) ######################################### # Nonlinear solve -freegs4e.solve(eq, # The equilibrium to adjust - profiles, # The toroidal current profile function - constrain, # Constraint function to set coil currents - psi_bndry=0.0) # Because no X-points, specify the separatrix psi +freegs4e.solve( + eq, # The equilibrium to adjust + profiles, # The toroidal current profile function + constrain, # Constraint function to set coil currents + psi_bndry=0.0, +) # Because no X-points, specify the separatrix psi # eq now contains the solution @@ -81,4 +91,3 @@ axis = eq.plot(show=False) constrain.plot(axis=axis, show=True) - diff --git a/10-mastu-connection.py b/10-mastu-connection.py index 5d8c085..748dcaf 100644 --- a/10-mastu-connection.py +++ b/10-mastu-connection.py @@ -1,29 +1,29 @@ # Read MAST-U GEQDSK file, calculate connection length -from freegs4e import geqdsk -from freegs4e import machine +from freegs4e import geqdsk, machine tokamak = machine.MASTU() with open("mast-upgrade.geqdsk") as f: eq = geqdsk.read(f, tokamak, show=False) -from freegs4e import fieldtracer import matplotlib.pyplot as plt +from freegs4e import fieldtracer + # Plot equilibrium axis = eq.plot(show=False) # Trace field lines both directions along the magnetic field -# By passing axis +# By passing axis forward, backward = fieldtracer.traceFieldLines(eq, axis=axis, nturns=50) plt.savefig("mast-upgrade-fieldtrace.pdf") plt.show() # Plot field line length from midplane to target -plt.plot(forward.R[0,:], forward.length[-1,:], label="Forward") -plt.plot(backward.R[0,:], backward.length[-1,:], label="Backward") +plt.plot(forward.R[0, :], forward.length[-1, :], label="Forward") +plt.plot(backward.R[0, :], backward.length[-1, :], label="Backward") plt.legend() plt.xlabel("Starting major radius [m]") plt.ylabel("Parallel connection length [m]") diff --git a/11-optimise-coils.py b/11-optimise-coils.py index 87f9872..c75f330 100644 --- a/11-optimise-coils.py +++ b/11-optimise-coils.py @@ -8,19 +8,24 @@ tokamak = freegs4e.machine.TestTokamak() -eq = freegs4e.Equilibrium(tokamak=tokamak, - Rmin=0.1, Rmax=2.0, # Radial domain - Zmin=-1.0, Zmax=1.0, # Height range - nx=65, ny=65, # Number of grid points - boundary=freegs4e.boundary.freeBoundaryHagenow) # Boundary condition +eq = freegs4e.Equilibrium( + tokamak=tokamak, + Rmin=0.1, + Rmax=2.0, # Radial domain + Zmin=-1.0, + Zmax=1.0, # Height range + nx=65, + ny=65, # Number of grid points + boundary=freegs4e.boundary.freeBoundaryHagenow, +) # Boundary condition ######################################### # Plasma profiles -profiles = freegs4e.jtor.ConstrainPaxisIp(1e3, # Plasma pressure on axis [Pascals] - 2e5, # Plasma current [Amps] - 2.0) # Vacuum f=R*Bt +profiles = freegs4e.jtor.ConstrainPaxisIp( + 1e3, 2e5, 2.0 # Plasma pressure on axis [Pascals] # Plasma current [Amps] +) # Vacuum f=R*Bt ######################################### # Coil current constraints @@ -28,23 +33,27 @@ # Specify locations of the X-points # to use to constrain coil currents -xpoints = [(1.1, -0.6), # (R,Z) locations of X-points - (1.1, 0.6)] +xpoints = [(1.1, -0.6), (1.1, 0.6)] # (R,Z) locations of X-points -isoflux = [(1.1,-0.6, 1.1,0.6), # (R1,Z1, R2,Z2) pair of locations - (1.7, 0.0, 0.84, 0.0)] +isoflux = [ + (1.1, -0.6, 1.1, 0.6), # (R1,Z1, R2,Z2) pair of locations + (1.7, 0.0, 0.84, 0.0), +] -constrain = freegs4e.control.constrain(xpoints=xpoints, isoflux=isoflux, gamma = 1e-17) +constrain = freegs4e.control.constrain( + xpoints=xpoints, isoflux=isoflux, gamma=1e-17 +) ######################################### # Nonlinear solve -freegs4e.solve(eq, # The equilibrium to adjust - profiles, # The toroidal current profile function - constrain) # Constraint function to set coil currents +freegs4e.solve( + eq, # The equilibrium to adjust + profiles, # The toroidal current profile function + constrain, +) # Constraint function to set coil currents - # Currents in the coils tokamak.printCurrents() @@ -60,19 +69,19 @@ print("Starting optimisation") -best_eq = opt.optimise(eq, # Starting equilibrium - # List of controls - [opt.CoilRadius("P2U"), - opt.CoilRadius("P2L"), opt.CoilHeight("P2L")], - # The function to minimise - opt.weighted_sum(opt.max_coil_force, opt.no_wall_intersection), - N=10, # Number of solutions in each generation - maxgen=20, # How many generations - monitor=opt.PlotMonitor()) # Plot the best in each generation +best_eq = opt.optimise( + eq, # Starting equilibrium + # List of controls + [opt.CoilRadius("P2U"), opt.CoilRadius("P2L"), opt.CoilHeight("P2L")], + # The function to minimise + opt.weighted_sum(opt.max_coil_force, opt.no_wall_intersection), + N=10, # Number of solutions in each generation + maxgen=20, # How many generations + monitor=opt.PlotMonitor(), +) # Plot the best in each generation print("Finished optimisation") # Forces on the coils best_eq.printForces() best_eq.plot() - diff --git a/docs/conf.py b/docs/conf.py index 58b7b9c..419a30e 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -18,16 +18,19 @@ # import os import sys -sys.path.insert(0, os.path.abspath('..')) + +sys.path.insert(0, os.path.abspath("..")) from unittest.mock import MagicMock + class Mock(MagicMock): @classmethod def __getattr__(cls, name): return MagicMock() -MOCK_MODULES = ['_tkinter'] + +MOCK_MODULES = ["_tkinter"] sys.modules.update((mod_name, Mock()) for mod_name in MOCK_MODULES) # -- General configuration ------------------------------------------------ @@ -39,36 +42,38 @@ def __getattr__(cls, name): # Add any Sphinx extension module names here, as strings. They can be # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom # ones. -extensions = ['sphinx.ext.autodoc', - 'sphinx.ext.coverage', - 'sphinx.ext.imgmath', - 'sphinx.ext.viewcode'] +extensions = [ + "sphinx.ext.autodoc", + "sphinx.ext.coverage", + "sphinx.ext.imgmath", + "sphinx.ext.viewcode", +] # Add any paths that contain templates here, relative to this directory. -templates_path = ['_templates'] +templates_path = ["_templates"] # The suffix(es) of source filenames. # You can specify multiple suffix as a list of string: # # source_suffix = ['.rst', '.md'] -source_suffix = '.rst' +source_suffix = ".rst" # The master toctree document. -master_doc = 'index' +master_doc = "index" # General information about the project. -project = u'FreeGS' -copyright = u'2017-2019, Ben Dudson' -author = u'Ben Dudson' +project = "FreeGS" +copyright = "2017-2019, Ben Dudson" +author = "Ben Dudson" # The version info for the project you're documenting, acts as replacement for # |version| and |release|, also used in various other places throughout the # built documents. # # The short X.Y version. -version = u'0.2' +version = "0.2" # The full version, including alpha/beta/rc tags. -release = u'0.2.0' +release = "0.2.0" # The language for content autogenerated by Sphinx. Refer to documentation # for a list of supported languages. @@ -80,10 +85,10 @@ def __getattr__(cls, name): # List of patterns, relative to source directory, that match files and # directories to ignore when looking for source files. # This patterns also effect to html_static_path and html_extra_path -exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store'] +exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] # The name of the Pygments (syntax highlighting) style to use. -pygments_style = 'sphinx' +pygments_style = "sphinx" # If true, `todo` and `todoList` produce output, else they produce nothing. todo_include_todos = False @@ -94,7 +99,7 @@ def __getattr__(cls, name): # The theme to use for HTML and HTML Help pages. See the documentation for # a list of builtin themes. # -html_theme = 'alabaster' +html_theme = "alabaster" # Theme options are theme-specific and customize the look and feel of a theme # further. For a list of options available for each theme, see the @@ -105,13 +110,13 @@ def __getattr__(cls, name): # Add any paths that contain custom static files (such as style sheets) here, # relative to this directory. They are copied after the builtin static files, # so a file named "default.css" will overwrite the builtin "default.css". -html_static_path = ['_static'] +html_static_path = ["_static"] # -- Options for HTMLHelp output ------------------------------------------ # Output file base name for HTML help builder. -htmlhelp_basename = 'FreeGSdoc' +htmlhelp_basename = "FreeGSdoc" # -- Options for LaTeX output --------------------------------------------- @@ -120,15 +125,12 @@ def __getattr__(cls, name): # The paper size ('letterpaper' or 'a4paper'). # # 'papersize': 'letterpaper', - # The font size ('10pt', '11pt' or '12pt'). # # 'pointsize': '10pt', - # Additional stuff for the LaTeX preamble. # # 'preamble': '', - # Latex figure (float) alignment # # 'figure_align': 'htbp', @@ -138,8 +140,7 @@ def __getattr__(cls, name): # (source start file, target name, title, # author, documentclass [howto, manual, or own class]). latex_documents = [ - (master_doc, 'FreeGS.tex', u'FreeGS Documentation', - u'Ben Dudson', 'manual'), + (master_doc, "FreeGS.tex", "FreeGS Documentation", "Ben Dudson", "manual"), ] @@ -147,10 +148,7 @@ def __getattr__(cls, name): # One entry per manual page. List of tuples # (source start file, name, description, authors, manual section). -man_pages = [ - (master_doc, 'freegs', u'FreeGS Documentation', - [author], 1) -] +man_pages = [(master_doc, "freegs", "FreeGS Documentation", [author], 1)] # -- Options for Texinfo output ------------------------------------------- @@ -159,8 +157,13 @@ def __getattr__(cls, name): # (source start file, target name, title, author, # dir menu entry, description, category) texinfo_documents = [ - (master_doc, 'FreeGS', u'FreeGS Documentation', - author, 'FreeGS', 'One line description of project.', - 'Miscellaneous'), + ( + master_doc, + "FreeGS", + "FreeGS Documentation", + author, + "FreeGS", + "One line description of project.", + "Miscellaneous", + ), ] - diff --git a/freegs4e/__init__.py b/freegs4e/__init__.py index 1e5e468..24edf10 100644 --- a/freegs4e/__init__.py +++ b/freegs4e/__init__.py @@ -31,20 +31,12 @@ along with FreeGS. If not, see . """ + from importlib import metadata __version__ = metadata.version("freegs4e") +from . import control, jtor, machine, plotting +from .dump import OutputFile from .equilibrium import Equilibrium - -from . import jtor - -from . import machine - -from . import control - from .picard import solve - -from .dump import OutputFile - -from . import plotting diff --git a/freegs4e/_aeqdsk.py b/freegs4e/_aeqdsk.py index cf36cde..03fdefd 100644 --- a/freegs4e/_aeqdsk.py +++ b/freegs4e/_aeqdsk.py @@ -4,10 +4,10 @@ """ -from . import _fileutils as fu - import warnings +from . import _fileutils as fu + # List of file data variables, default values, and documentation # This is used in both reader and writer fields = [ @@ -46,7 +46,11 @@ ("otop", 10.0, "plasma top gap in cm"), ("obott", 10.0, "plasma bottom gap in cm"), ("qpsib", 5.0, "q at 95% of poloidal flux"), - ("vertn", 1.0, "vacuum field (index? -- seems to be float) at current centroid"), + ( + "vertn", + 1.0, + "vacuum field (index? -- seems to be float) at current centroid", + ), # fmt_1040 = r '^\s*' + 4 * r '([\s\-]\d+\.\d+[Ee][\+\-]\d\d)' # read(neqdsk, 1040)(rco2v(k, jj), k = 1, mco2v) (None, None, None), # New line @@ -185,8 +189,16 @@ ("rvsout", 0.0, "major radius of vessel outer hit spot in cm"), ("zvsout", 0.0, "Z of vessel outer hit spot in cm"), ("vsurfa", 0.0, "plasma surface loop voltage in volt, E EQDSK only"), - ("wpdot", 0.0, "time derivative of plasma stored energy in Watt, E EQDSK only"), - ("wbdot", 0.0, "time derivative of poloidal magnetic energy in Watt, E EQDSK only"), + ( + "wpdot", + 0.0, + "time derivative of plasma stored energy in Watt, E EQDSK only", + ), + ( + "wbdot", + 0.0, + "time derivative of poloidal magnetic energy in Watt, E EQDSK only", + ), ("slantu", 0.0, ""), ("slantl", 0.0, ""), ("zuperts", 0.0, ""), @@ -214,7 +226,11 @@ "radial distance in cm between x point and external field line at ZNOSE", ), ("ssi95", 0.0, "magnetic shear at 95% of normalized poloidal flux"), - ("rqqmin", 0.0, "normalized radius of qmin , square root of normalized volume"), + ( + "rqqmin", + 0.0, + "normalized radius of qmin , square root of normalized volume", + ), ("cjor99", 0.0, ""), ( "cj1ave", @@ -322,7 +338,9 @@ def read(fh): value = next(values) if isinstance(default, int) and not isinstance(value, int): # Expecting an integer, but didn't get one - warnings.warn("Expecting an integer for '" + key + "' in aeqdsk file") + warnings.warn( + "Expecting an integer for '" + key + "' in aeqdsk file" + ) break data[key] = value diff --git a/freegs4e/_divgeo.py b/freegs4e/_divgeo.py index 950e8d4..70b6327 100644 --- a/freegs4e/_divgeo.py +++ b/freegs4e/_divgeo.py @@ -19,6 +19,7 @@ along with FreeGS4E. If not, see . """ + import numpy as np from ._fileutils import ChunkOutput, write_1d, write_2d diff --git a/freegs4e/_geqdsk.py b/freegs4e/_geqdsk.py index 5bbe484..ba6dea1 100644 --- a/freegs4e/_geqdsk.py +++ b/freegs4e/_geqdsk.py @@ -21,9 +21,10 @@ """ from datetime import date -from numpy import zeros, pi -from ._fileutils import f2s, ChunkOutput, write_1d, write_2d, next_value +from numpy import pi, zeros + +from ._fileutils import ChunkOutput, f2s, next_value, write_1d, write_2d def write(data, fh, label=None, shot=None, time=None): @@ -60,7 +61,9 @@ def write(data, fh, label=None, shot=None, time=None): label = "FREEGS" if len(label) > 11: label = label[0:12] - print("WARNING: label too long, it will be shortened to {}".format(label)) + print( + "WARNING: label too long, it will be shortened to {}".format(label) + ) creation_date = date.today().strftime("%d/%m/%Y") @@ -117,7 +120,12 @@ def write(data, fh, label=None, shot=None, time=None): # 5th line fh.write( - f2s(data["zmagx"]) + f2s(0.0) + f2s(data["sibdry"]) + f2s(0.0) + f2s(0.0) + "\n" + f2s(data["zmagx"]) + + f2s(0.0) + + f2s(data["sibdry"]) + + f2s(0.0) + + f2s(0.0) + + "\n" ) # SCENE has actual ff' and p' data so can use that diff --git a/freegs4e/bilinear_interpolation.py b/freegs4e/bilinear_interpolation.py index ad22215..3eec69d 100644 --- a/freegs4e/bilinear_interpolation.py +++ b/freegs4e/bilinear_interpolation.py @@ -1,8 +1,10 @@ import numpy as np -# + + +# def biliint(R, Z, psi, points): - R1d = R[:,:1] - Z1d = Z[:1,:] + R1d = R[:, :1] + Z1d = Z[:1, :] dR = R[1, 0] - R[0, 0] dZ = Z[0, 1] - Z[0, 0] dRdZ = dR * dZ @@ -11,28 +13,43 @@ def biliint(R, Z, psi, points): points = points.reshape(2, -1) len_points = np.shape(points)[1] - points_R = R1d - points[:1,:] - points_Z = Z1d.T - points[1:2,:] + points_R = R1d - points[:1, :] + points_Z = Z1d.T - points[1:2, :] idxs_R = np.sum(points_R < 0, axis=0) idxs_Z = np.sum(points_Z < 0, axis=0) - iR = idxs_R[:,np.newaxis,np.newaxis] - iZ = idxs_Z[:,np.newaxis,np.newaxis] - qq = psi[np.concatenate((np.concatenate((iR-1, iR-1), axis=2), - np.concatenate((iR, iR), axis=2)), axis=1), - np.concatenate((np.concatenate((iZ-1, iZ), axis=2), - np.concatenate((iZ-1, iZ), axis=2)), axis=1)] - - iR = idxs_R[:,np.newaxis] - iZ = idxs_Z[:,np.newaxis] - xx = points_R[np.concatenate((iR, iR-1), axis=1),np.arange(len_points)[:,np.newaxis]]*(np.array([[1,-1]])) - yy = points_Z[np.concatenate((iZ, iZ-1), axis=1),np.arange(len_points)[:,np.newaxis]]*(np.array([[1,-1]])) - - vals = np.sum(np.sum(qq*yy[:,np.newaxis,:], axis=-1)*xx, axis=-1)/dRdZ + iR = idxs_R[:, np.newaxis, np.newaxis] + iZ = idxs_Z[:, np.newaxis, np.newaxis] + qq = psi[ + np.concatenate( + ( + np.concatenate((iR - 1, iR - 1), axis=2), + np.concatenate((iR, iR), axis=2), + ), + axis=1, + ), + np.concatenate( + ( + np.concatenate((iZ - 1, iZ), axis=2), + np.concatenate((iZ - 1, iZ), axis=2), + ), + axis=1, + ), + ] + + iR = idxs_R[:, np.newaxis] + iZ = idxs_Z[:, np.newaxis] + xx = points_R[ + np.concatenate((iR, iR - 1), axis=1), + np.arange(len_points)[:, np.newaxis], + ] * (np.array([[1, -1]])) + yy = points_Z[ + np.concatenate((iZ, iZ - 1), axis=1), + np.arange(len_points)[:, np.newaxis], + ] * (np.array([[1, -1]])) + + vals = ( + np.sum(np.sum(qq * yy[:, np.newaxis, :], axis=-1) * xx, axis=-1) / dRdZ + ) return vals.reshape(points_shape[1:]) - - - - - diff --git a/freegs4e/boundary.py b/freegs4e/boundary.py index 20341e0..92aa461 100644 --- a/freegs4e/boundary.py +++ b/freegs4e/boundary.py @@ -20,11 +20,11 @@ """ -from .gradshafranov import Greens from numpy import concatenate, sqrt - from scipy.integrate import romb # Romberg integration +from .gradshafranov import Greens + def fixedBoundary(eq, Jtor, psi): """ @@ -141,14 +141,21 @@ def freeBoundaryHagenow(eq, Jtor, psi): # Note: normal is out of the domain, so this is -dU/dx # Fourth-order one-sided differences - coeffs = [(0, 25.0 / 12), (1, -4.0), (2, 3.0), (3, -16.0 / 12), (4, 1.0 / 4)] + coeffs = [ + (0, 25.0 / 12), + (1, -4.0), + (2, 3.0), + (3, -16.0 / 12), + (4, 1.0 / 4), + ] dUdn_L = ( sum([weight * psi_fixed[index, :] for index, weight in coeffs]) / dR ) # left boundary dUdn_R = ( - sum([weight * psi_fixed[-(1 + index), :] for index, weight in coeffs]) / dR + sum([weight * psi_fixed[-(1 + index), :] for index, weight in coeffs]) + / dR ) # Right boundary dUdn_D = ( @@ -156,30 +163,47 @@ def freeBoundaryHagenow(eq, Jtor, psi): ) # Down boundary dUdn_U = ( - sum([weight * psi_fixed[:, -(1 + index)] for index, weight in coeffs]) / dZ + sum([weight * psi_fixed[:, -(1 + index)] for index, weight in coeffs]) + / dZ ) # Upper boundary - dd = sqrt(dR ** 2 + dZ ** 2) # Diagonal spacing + dd = sqrt(dR**2 + dZ**2) # Diagonal spacing # Left down corner dUdn_L[0] = dUdn_D[0] = ( - sum([weight * psi_fixed[index, index] for index, weight in coeffs]) / dd + sum([weight * psi_fixed[index, index] for index, weight in coeffs]) + / dd ) # Left upper corner dUdn_L[-1] = dUdn_U[0] = ( - sum([weight * psi_fixed[index, -(1 + index)] for index, weight in coeffs]) / dd + sum( + [ + weight * psi_fixed[index, -(1 + index)] + for index, weight in coeffs + ] + ) + / dd ) # Right down corner dUdn_R[0] = dUdn_D[-1] = ( - sum([weight * psi_fixed[-(1 + index), index] for index, weight in coeffs]) / dd + sum( + [ + weight * psi_fixed[-(1 + index), index] + for index, weight in coeffs + ] + ) + / dd ) # Right upper corner dUdn_R[-1] = dUdn_U[-1] = ( sum( - [weight * psi_fixed[-(1 + index), -(1 + index)] for index, weight in coeffs] + [ + weight * psi_fixed[-(1 + index), -(1 + index)] + for index, weight in coeffs + ] ) / dd ) diff --git a/freegs4e/coil.py b/freegs4e/coil.py index 658d063..cca809e 100644 --- a/freegs4e/coil.py +++ b/freegs4e/coil.py @@ -24,11 +24,14 @@ along with FreeGS4E. If not, see . """ -from .gradshafranov import Greens, GreensBr, GreensBz, mu0 -import numpy as np import numbers -from matplotlib.patches import Rectangle + import matplotlib.pyplot as plt +import numpy as np +from matplotlib.patches import Rectangle + +from .gradshafranov import Greens, GreensBr, GreensBz, mu0 + class AreaCurrentLimit: """ @@ -205,7 +208,7 @@ def getForces(self, equilibrium): # Force per unit length. # In cgs units f = I^2/(c^2 * R) * (ln(8*R/a) - 1 + xi/2) # In SI units f = mu0 * I^2 / (4*pi*R) * (ln(8*R/a) - 1 + xi/2) - self_fr = (mu0 * total_current ** 2 / (4.0 * np.pi * self.R)) * ( + self_fr = (mu0 * total_current**2 / (4.0 * np.pi * self.R)) * ( np.log(8.0 * self.R / minor_radius) - 1 + self_inductance / 2.0 ) @@ -240,7 +243,8 @@ def to_numpy_array(self): Helper method for writing output """ return np.array( - (self.R, self.Z, self.current, self.turns, self.control), dtype=self.dtype + (self.R, self.Z, self.current, self.turns, self.control), + dtype=self.dtype, ) @classmethod @@ -276,7 +280,7 @@ def plot(self, axis=None, show=False): The area of the coil is used to set the radius """ - + try: axis = self.plot_nke(axis, show) except: @@ -287,7 +291,6 @@ def plot(self, axis=None, show=False): minor_radius = np.sqrt(self.area / np.pi) - circle = plt.Circle((self.R, self.Z), minor_radius, color="b") axis.add_artist(circle) @@ -298,11 +301,12 @@ def plot_nke(self, axis=None, show=False): Plot the coil location, using axis if given """ - self.rectangle = Rectangle((self.R - self.dR/2, - self.Z - self.dZ/2), - width=self.dR, - height=self.dZ, - facecolor = 'b') + self.rectangle = Rectangle( + (self.R - self.dR / 2, self.Z - self.dZ / 2), + width=self.dR, + height=self.dZ, + facecolor="b", + ) if axis is None: fig = plt.figure() @@ -310,4 +314,3 @@ def plot_nke(self, axis=None, show=False): axis.add_patch(self.rectangle) return axis - diff --git a/freegs4e/control.py b/freegs4e/control.py index 0725202..b841c09 100644 --- a/freegs4e/control.py +++ b/freegs4e/control.py @@ -4,10 +4,9 @@ Use constraints to adjust coil currents """ -from numpy import dot, transpose, eye, array -from numpy.linalg import inv import numpy as np - +from numpy import array, dot, eye, transpose +from numpy.linalg import inv from scipy import optimize from . import critical @@ -117,7 +116,7 @@ def __call__(self, eq): # Calculate the change in coil current current_change = dot( - inv(dot(transpose(A), A) + self.gamma ** 2 * eye(ncontrols)), + inv(dot(transpose(A), A) + self.gamma**2 * eye(ncontrols)), dot(transpose(A), b), ) # print("Current changes: " + str(current_change)) diff --git a/freegs4e/critical.py b/freegs4e/critical.py index 70580af..e2de74d 100644 --- a/freegs4e/critical.py +++ b/freegs4e/critical.py @@ -20,38 +20,41 @@ """ - +import warnings from unittest import makeSuite -from scipy import interpolate -from numpy import zeros -from numpy.linalg import inv + +import numpy as np from numpy import ( - dot, - linspace, + abs, + amax, + arctan2, argmax, argmin, - abs, clip, - sin, cos, + dot, + linspace, pi, - amax, - arctan2, + sin, sqrt, sum, + zeros, ) -import numpy as np -import warnings +from numpy.linalg import inv +from scipy import interpolate try: from numba import njit except ImportError: warnings.warn("Numba not found, using slower version") + def njit(*args, **kwargs): return lambda f: f + from warnings import warn -from . import bilinear_interpolation + +from . import bilinear_interpolation def find_critical_old(R, Z, psi, discard_xpoints=True): @@ -83,13 +86,15 @@ def find_critical_old(R, Z, psi, discard_xpoints=True): f = interpolate.RectBivariateSpline(R[:, 0], Z[0, :], psi) # Find candidate locations, based on minimising Bp^2 - Bp2 = (f(R, Z, dx=1, grid=False) ** 2 + f(R, Z, dy=1, grid=False) ** 2) / R ** 2 + Bp2 = ( + f(R, Z, dx=1, grid=False) ** 2 + f(R, Z, dy=1, grid=False) ** 2 + ) / R**2 # Get grid resolution, which determines a reasonable tolerance # for the Newton iteration search area dR = R[1, 0] - R[0, 0] dZ = Z[0, 1] - Z[0, 0] - radius_sq = 9 * (dR ** 2 + dZ ** 2) + radius_sq = 9 * (dR**2 + dZ**2) # Find local minima @@ -128,23 +133,25 @@ def find_critical_old(R, Z, psi, discard_xpoints=True): Br = -f(R1, Z1, dy=1, grid=False) / R1 Bz = f(R1, Z1, dx=1, grid=False) / R1 - if Br ** 2 + Bz ** 2 < 1e-6: + if Br**2 + Bz**2 < 1e-6: # Found a minimum. Classify as either # O-point or X-point dR = R[1, 0] - R[0, 0] dZ = Z[0, 1] - Z[0, 0] - d2dr2 = (psi[i + 2, j] - 2.0 * psi[i, j] + psi[i - 2, j]) / ( - 2.0 * dR - ) ** 2 - d2dz2 = (psi[i, j + 2] - 2.0 * psi[i, j] + psi[i, j - 2]) / ( - 2.0 * dZ - ) ** 2 + d2dr2 = ( + psi[i + 2, j] - 2.0 * psi[i, j] + psi[i - 2, j] + ) / (2.0 * dR) ** 2 + d2dz2 = ( + psi[i, j + 2] - 2.0 * psi[i, j] + psi[i, j - 2] + ) / (2.0 * dZ) ** 2 d2drdz = ( - (psi[i + 2, j + 2] - psi[i + 2, j - 2]) / (4.0 * dZ) - - (psi[i - 2, j + 2] - psi[i - 2, j - 2]) / (4.0 * dZ) + (psi[i + 2, j + 2] - psi[i + 2, j - 2]) + / (4.0 * dZ) + - (psi[i - 2, j + 2] - psi[i - 2, j - 2]) + / (4.0 * dZ) ) / (4.0 * dR) - D = d2dr2 * d2dz2 - d2drdz ** 2 + D = d2dr2 * d2dz2 - d2drdz**2 if D < 0.0: # Found X-point @@ -171,7 +178,9 @@ def find_critical_old(R, Z, psi, discard_xpoints=True): count += 1 # If (R1,Z1) is too far from (R0,Z0) then discard # or if we've taken too many iterations - if ((R1 - R0) ** 2 + (Z1 - Z0) ** 2 > radius_sq) or (count > 100): + if ((R1 - R0) ** 2 + (Z1 - Z0) ** 2 > radius_sq) or ( + count > 100 + ): # Discard this point break @@ -244,65 +253,70 @@ def remove_dup(points): return opoint, xpoint + # Remove duplicates def remove_dup(points): - result = [] - for n, p in enumerate(points): - dup = False - for p2 in result: - if (p[0] - p2[0]) ** 2 + (p[1] - p2[1]) ** 2 < 1e-5: - dup = True # Duplicate - break - if not dup: - result.append(p) # Add to the list - return result - -def find_critical(R, Z, psi, mask_inside_limiter=None, signIp=1, discard_xpoints=True): + result = [] + for n, p in enumerate(points): + dup = False + for p2 in result: + if (p[0] - p2[0]) ** 2 + (p[1] - p2[1]) ** 2 < 1e-5: + dup = True # Duplicate + break + if not dup: + result.append(p) # Add to the list + return result + + +def find_critical( + R, Z, psi, mask_inside_limiter=None, signIp=1, discard_xpoints=True +): # if old: # opoint, xpoint = find_critical_old(R,Z,psi, discard_xpoints) # else: opoint, xpoint = fastcrit(R, Z, psi, mask_inside_limiter) - if len(xpoint)>0 and (signIp is not None): + if len(xpoint) > 0 and (signIp is not None): # select xpoint with the correct ordering wrt Ip - xpoint = xpoint[((xpoint[:,2] - opoint[:1,2])*signIp) < 0] - if len(xpoint)>0: - # check distance to opoint and in case discard xpoints on non-monotonic LOS + xpoint = xpoint[((xpoint[:, 2] - opoint[:1, 2]) * signIp) < 0] + if len(xpoint) > 0: + # check distance to opoint and in case discard xpoints on non-monotonic LOS # closer_xpoint = np.argmin(np.linalg.norm((xpoint-opoint[:1])[:,:2], axis=-1)) - # if closer_xpoint != 0: - # f = interpolate.RectBivariateSpline(R[:, 0], Z[0, :], psi) + # if closer_xpoint != 0: + # f = interpolate.RectBivariateSpline(R[:, 0], Z[0, :], psi) result = False while result is False: result = discard_xpoints_f(R, Z, psi, opoint[0], xpoint[0]) if result is False: xpoint = xpoint[1:] - result = len(xpoint)<1 + result = len(xpoint) < 1 # print(xpoint) return opoint, xpoint + # # this is 10x faster if the numba import works; otherwise, @njit is the identity and fastcrit is 3x faster anyways @njit(fastmath=True, cache=True) def scan_for_crit(R, Z, psi): dR = R[1, 0] - R[0, 0] dZ = Z[0, 1] - Z[0, 0] - Bp2=np.zeros_like(psi) - psiR=Bp2.copy() - psiZ=Bp2.copy() - psiR[1:-1,1:-1] = 0.5*(psi[2:,1:-1]-psi[:-2,1:-1])/dR - psiZ[1:-1,1:-1] = 0.5*(psi[1:-1,2:]-psi[1:-1,:-2])/dZ + Bp2 = np.zeros_like(psi) + psiR = Bp2.copy() + psiZ = Bp2.copy() + psiR[1:-1, 1:-1] = 0.5 * (psi[2:, 1:-1] - psi[:-2, 1:-1]) / dR + psiZ[1:-1, 1:-1] = 0.5 * (psi[1:-1, 2:] - psi[1:-1, :-2]) / dZ # - # psiR[0,:]=(psi[1,:]-psi[0,:])/dR -# psiR[-1,:]=(psi[-1,:]-psi[-2,:])/dR -# psiR[1:-1,0]=(psi[1:,0]-psi[:-1,0])/dR -# psiR[1:-1,-1]=(psi[1:,-1]-psi[:-1,-0])/dR + # psiR[0,:]=(psi[1,:]-psi[0,:])/dR + # psiR[-1,:]=(psi[-1,:]-psi[-2,:])/dR + # psiR[1:-1,0]=(psi[1:,0]-psi[:-1,0])/dR + # psiR[1:-1,-1]=(psi[1:,-1]-psi[:-1,-0])/dR # - Bp2[:,:]= (psiR**2 + psiZ**2)#/R[:,:]**2 + Bp2[:, :] = psiR**2 + psiZ**2 # /R[:,:]**2 # - xpoint = [(-999.0,-999.0,-999.0)] - opoint = [(-999.0,-999.0,-999.0)] + xpoint = [(-999.0, -999.0, -999.0)] + opoint = [(-999.0, -999.0, -999.0)] # start off by finding coarse values of Bp2 closest to 0.0 as in Ben Dudsons's routine - for i in range(1,len(Bp2)-1): - for j in range(1,len(Bp2[0])-1): + for i in range(1, len(Bp2) - 1): + for j in range(1, len(Bp2[0]) - 1): if ( (Bp2[i, j] < Bp2[i + 1, j + 1]) and (Bp2[i, j] < Bp2[i + 1, j]) @@ -316,55 +330,70 @@ def scan_for_crit(R, Z, psi): # Found local minimum R0 = R[i, j] Z0 = Z[i, j] - fR , fZ = psiR[i,j], psiZ[i,j] - fRR = (psi[i+1,j]-2*psi[i,j]+psi[i-1,j])/dR**2 - fZZ = (psi[i,j+1]-2*psi[i,j]+psi[i,j-1])/dZ**2 - fRZ = 0.5*(psi[i+1,j+1] + psi[i-1,j-1] -psi[i-1,j+1] - psi[i+1,j-1] )/(dR*dZ) - det = fRR*fZZ-0.25*fRZ**2 + fR, fZ = psiR[i, j], psiZ[i, j] + fRR = (psi[i + 1, j] - 2 * psi[i, j] + psi[i - 1, j]) / dR**2 + fZZ = (psi[i, j + 1] - 2 * psi[i, j] + psi[i, j - 1]) / dZ**2 + fRZ = ( + 0.5 + * ( + psi[i + 1, j + 1] + + psi[i - 1, j - 1] + - psi[i - 1, j + 1] + - psi[i + 1, j - 1] + ) + / (dR * dZ) + ) + det = fRR * fZZ - 0.25 * fRZ**2 # - if det!=0: - delta_R = -(fR*fZZ-0.5*fRZ*fZ)/det - delta_Z = -(fZ*fRR-0.5*fRZ*fR)/det - if (np.abs(delta_R)<=dR and np.abs(delta_Z)<=dZ): - est_psi = psi[i,j]+0.5*(fR*delta_R + fZ*delta_Z) #+ 0.5*(fRR*delta_R**2 + fZZ*delta_Z**2 + fRZ*delta_R*delta_Z) - crpoint = (R0+delta_R , Z0+delta_Z , est_psi) - if det>0.0 : opoint = [crpoint] + opoint - else: xpoint = [crpoint] + xpoint + if det != 0: + delta_R = -(fR * fZZ - 0.5 * fRZ * fZ) / det + delta_Z = -(fZ * fRR - 0.5 * fRZ * fR) / det + if np.abs(delta_R) <= dR and np.abs(delta_Z) <= dZ: + est_psi = psi[i, j] + 0.5 * ( + fR * delta_R + fZ * delta_Z + ) # + 0.5*(fRR*delta_R**2 + fZZ*delta_Z**2 + fRZ*delta_R*delta_Z) + crpoint = (R0 + delta_R, Z0 + delta_Z, est_psi) + if det > 0.0: + opoint = [crpoint] + opoint + else: + xpoint = [crpoint] + xpoint xpoint = np.array(xpoint) opoint = np.array(opoint) # do NOT remove the "pop" command below, the lists were initialised with (-999.,-999.) so that numba could compile # xpoint.pop() # opoint.pop() - xpoint = xpoint[xpoint[:,0]>-990] - opoint = opoint[opoint[:,0]>-990] + xpoint = xpoint[xpoint[:, 0] > -990] + opoint = opoint[opoint[:, 0] > -990] return opoint, xpoint def fastcrit(R, Z, psi, mask_inside_limiter): - opoint , xpoint = scan_for_crit(R, Z, psi) + opoint, xpoint = scan_for_crit(R, Z, psi) len_opoint = len(opoint) if len_opoint == 0: # Can't order primary O-point, X-point so return - raise ValueError('No opoints found!') + raise ValueError("No opoints found!") # return opoint, xpoint elif mask_inside_limiter is not None: # remove any opoint outside the limiter - posR = np.argmin((R[:,:1].T - opoint[:,:1])**2, axis=1) - posZ = np.argmin((Z[:1,:] - opoint[:,1:2])**2, axis=1) - opoint = opoint[mask_inside_limiter[posR,posZ]] + posR = np.argmin((R[:, :1].T - opoint[:, :1]) ** 2, axis=1) + posZ = np.argmin((Z[:1, :] - opoint[:, 1:2]) ** 2, axis=1) + opoint = opoint[mask_inside_limiter[posR, posZ]] if len_opoint == 0: # Can't order primary O-point, X-point so return - raise ValueError('No opoints found!') + raise ValueError("No opoints found!") elif len_opoint > 1: # Find primary O-point by sorting by distance from middle of domain Rmid = 0.5 * (R[-1, 0] + R[0, 0]) Zmid = 0.5 * (Z[0, -1] + Z[0, 0]) - opoint_ordering = np.argsort((opoint[:,0]-Rmid)**2 + (opoint[:,1]-Zmid)**2) + opoint_ordering = np.argsort( + (opoint[:, 0] - Rmid) ** 2 + (opoint[:, 1] - Zmid) ** 2 + ) opoint = opoint[opoint_ordering] - + # # check that primary opoint is inside the limiter # result = False # while result is False: @@ -375,7 +404,7 @@ def fastcrit(R, Z, psi, mask_inside_limiter): # if len(opoint) == 0: # raise ValueError('No valid opoints found!') # else: - # result = True + # result = True psi_axis = opoint[0][2] # opoint.sort(key=lambda x: (x[0] - Rmid) ** 2 + (x[1] - Zmid) ** 2) @@ -418,13 +447,14 @@ def fastcrit(R, Z, psi, mask_inside_limiter): # xpoint = xpt_keep # Sort X-points by distance to primary O-point in psi space - if len(xpoint)>1: - xpoint_ordering = np.argsort((xpoint[:,2]-psi_axis)**2) + if len(xpoint) > 1: + xpoint_ordering = np.argsort((xpoint[:, 2] - psi_axis) ** 2) xpoint = xpoint[xpoint_ordering] # xpoint.sort(key=lambda x: (x[2] - psi_axis) ** 2) # return opoint, xpoint + def discard_xpoints_f(R, Z, psi, opoint, xpt): # Here opoint and xpt are individual critical points dR = R[1, 0] - R[0, 0] @@ -432,18 +462,20 @@ def discard_xpoints_f(R, Z, psi, opoint, xpt): Ro, Zo, Po = opoint # The primary O-point result = False - + # for xpt in xpoint: Rx, Zx, Px = xpt - num = int(max((np.abs(Rx-Ro))/dR, (np.abs(Zx-Zo))/dZ) + 1) + num = int(max((np.abs(Rx - Ro)) / dR, (np.abs(Zx - Zo)) / dZ) + 1) # print('num ', num) # print('num') - rline = linspace(Ro, Rx, num)#(np.abs(Rx-Ro)//dR + 1)) - zline = linspace(Zo, Zx, num)#(np.abs(Zx-Zo)//dZ + 1)) + rline = linspace(Ro, Rx, num) # (np.abs(Rx-Ro)//dR + 1)) + zline = linspace(Zo, Zx, num) # (np.abs(Zx-Zo)//dZ + 1)) - pline = bilinear_interpolation.biliint(R, Z, psi, np.array([rline, zline]))#, grid=False) + pline = bilinear_interpolation.biliint( + R, Z, psi, np.array([rline, zline]) + ) # , grid=False) if Px < Po: pline *= -1.0 # Reverse, so pline is maximum at X-point @@ -467,8 +499,6 @@ def discard_xpoints_f(R, Z, psi, opoint, xpt): return result - - def core_mask(R, Z, psi, opoint, xpoint=[], psi_bndry=None): """ Mark the parts of the domain which are in the core @@ -535,7 +565,11 @@ def core_mask(R, Z, psi, opoint, xpoint=[], psi_bndry=None): while True: mask[i, j] = 1 # Mark as in the core - if (i < nx - 1) and (psin[i + 1, j] < 1.0) and (mask[i + 1, j] < 0.5): + if ( + (i < nx - 1) + and (psin[i + 1, j] < 1.0) + and (mask[i + 1, j] < 0.5) + ): stack.append((i + 1, j)) if (i > 0) and (psin[i - 1, j] < 1.0) and (mask[i - 1, j] < 0.5): stack.append((i - 1, j)) @@ -557,12 +591,16 @@ def core_mask(R, Z, psi, opoint, xpoint=[], psi_bndry=None): return mask + # def core_mask(R, Z, psi, opoint, xpoint=[], psi_bndry=None, old=False): # if old: return core_mask_old(R, Z, psi, opoint, xpoint, psi_bndry) # else: return inside_mask(R, Z, psi, opoint, xpoint, psi_bndry) + @njit(fastmath=True, cache=True) -def inside_mask(R, Z, psi, opoint, xpoint=[], mask_outside_limiter=None, psi_bndry=None): +def inside_mask( + R, Z, psi, opoint, xpoint=[], mask_outside_limiter=None, psi_bndry=None +): """ Similar to core_mask_old above, except: (1) it's all stuff that can be JIT-compiled @@ -590,10 +628,10 @@ def inside_mask(R, Z, psi, opoint, xpoint=[], mask_outside_limiter=None, psi_bnd ix = argmin(abs(R[:, 0] - rx)) jx = argmin(abs(Z[0, :] - zx)) xpt_inds.append((ix, jx)) - # Fill this point and all around with '2'# - ilo , ihi = min(max(ix-1,0),nx-1) , max(0,min(ix+1,nx-1)) - jlo , jhi = min(max(jx-1,0),ny-1) , max(0,min(jx+1,ny-1)) - mask[ilo:ihi+1,jlo:jhi+1] = 2 + # Fill this point and all around with '2'# + ilo, ihi = min(max(ix - 1, 0), nx - 1), max(0, min(ix + 1, nx - 1)) + jlo, jhi = min(max(jx - 1, 0), ny - 1), max(0, min(jx + 1, ny - 1)) + mask[ilo : ihi + 1, jlo : jhi + 1] = 2 # # Find nearest index to start rind = argmin(abs(R[:, 0] - Ro)) @@ -608,7 +646,7 @@ def inside_mask(R, Z, psi, opoint, xpoint=[], mask_outside_limiter=None, psi_bnd if (j > 0) and (psin[i, j - 1] < 1.0) and (mask[i, j - 1] < 0.5): stack.append((i, j - 1)) # Check the point above - if (j < ny -1) and (psin[i, j + 1] < 1.0) and (mask[i, j + 1] < 0.5): + if (j < ny - 1) and (psin[i, j + 1] < 1.0) and (mask[i, j + 1] < 0.5): stack.append((i, j + 1)) # # Scan along a row to the right @@ -621,16 +659,21 @@ def inside_mask(R, Z, psi, opoint, xpoint=[], mask_outside_limiter=None, psi_bnd # # Now return to X-point locations for ix, jx in xpt_inds: - ilo , ihi = min(max(ix-1,0),nx-1) , max(0,min(ix+1,nx-1)) - jlo , jhi = min(max(jx-1,0),ny-1) , max(0,min(jx+1,ny-1)) - mask[ilo:ihi+1,jlo:jhi+1] = 1*(mask[ilo:ihi+1,jlo:jhi+1]==1)*(psin[ilo:ihi+1,jlo:jhi+1]<1.0) + ilo, ihi = min(max(ix - 1, 0), nx - 1), max(0, min(ix + 1, nx - 1)) + jlo, jhi = min(max(jx - 1, 0), ny - 1), max(0, min(jx + 1, ny - 1)) + mask[ilo : ihi + 1, jlo : jhi + 1] = ( + 1 + * (mask[ilo : ihi + 1, jlo : jhi + 1] == 1) + * (psin[ilo : ihi + 1, jlo : jhi + 1] < 1.0) + ) # # remove effect of mask_outside_limiter - mask = (mask==1) + mask = mask == 1 return mask + def find_psisurface(eq, psifunc, r0, z0, r1, z1, psival=1.0, n=100, axis=None): """ eq - Equilibrium object @@ -714,7 +757,9 @@ def find_separatrix( # Avoid putting theta grid points exactly on the X-points xpoint_theta = arctan2(xpoint[0][0] - r0, xpoint[0][1] - z0) - xpoint_theta = xpoint_theta * (xpoint_theta >= 0) + (xpoint_theta + 2 * pi) * ( + xpoint_theta = xpoint_theta * (xpoint_theta >= 0) + ( + xpoint_theta + 2 * pi + ) * ( xpoint_theta < 0 ) # let's make it between 0 and 2*pi # How close in theta to allow theta grid points to the X-point @@ -742,7 +787,14 @@ def find_separatrix( def find_safety( - eq, npsi=1, psinorm=None, ntheta=128, psi=None, opoint=None, xpoint=None, axis=None + eq, + npsi=1, + psinorm=None, + ntheta=128, + psi=None, + opoint=None, + xpoint=None, + axis=None, ): """Find the safety factor for each value of psi Calculates equally spaced flux surfaces. Points on @@ -771,7 +823,9 @@ def find_safety( else: psinormal = (psi - opoint[0][2]) / (xpoint[0][2] - opoint[0][2]) - psifunc = interpolate.RectBivariateSpline(eq.R[:, 0], eq.Z[0, :], psinormal) + psifunc = interpolate.RectBivariateSpline( + eq.R[:, 0], eq.Z[0, :], psinormal + ) r0, z0 = opoint[0][0:2] @@ -780,7 +834,9 @@ def find_safety( # Avoid putting theta grid points exactly on the X-points xpoint_theta = arctan2(xpoint[0][0] - r0, xpoint[0][1] - z0) - xpoint_theta = xpoint_theta * (xpoint_theta >= 0) + (xpoint_theta + 2 * pi) * ( + xpoint_theta = xpoint_theta * (xpoint_theta >= 0) + ( + xpoint_theta + 2 * pi + ) * ( xpoint_theta < 0 ) # let's make it between 0 and 2*pi # How close in theta to allow theta grid points to the X-point @@ -826,19 +882,19 @@ def find_safety( fpol = eq.fpol(psirange[:]).reshape(npsi, 1) Br = eq.Br(r, z) Bz = eq.Bz(r, z) - Bthe = sqrt(Br ** 2 + Bz ** 2) + Bthe = sqrt(Br**2 + Bz**2) # Differentiate location w.r.t. index dr_di = (np.roll(r, 1, axis=1) - np.roll(r, -1, axis=1)) / 2.0 dz_di = (np.roll(z, 1, axis=1) - np.roll(z, -1, axis=1)) / 2.0 # Distance between points - dl = sqrt(dr_di ** 2 + dz_di ** 2) + dl = sqrt(dr_di**2 + dz_di**2) # Integrand - Btor/(R*Bthe) = Fpol/(R**2*Bthe) - qint = fpol / (r ** 2 * Bthe) + qint = fpol / (r**2 * Bthe) # Integral q = sum(qint * dl, axis=1) / (2 * pi) - return q \ No newline at end of file + return q diff --git a/freegs4e/critical_old.py b/freegs4e/critical_old.py index 6acbcb6..92dbd3d 100644 --- a/freegs4e/critical_old.py +++ b/freegs4e/critical_old.py @@ -20,34 +20,36 @@ """ - from unittest import makeSuite -from scipy import interpolate -from numpy import zeros -from numpy.linalg import inv + +import numpy as np from numpy import ( - dot, - linspace, + abs, + amax, + arctan2, argmax, argmin, - abs, clip, - sin, cos, + dot, + linspace, pi, - amax, - arctan2, + sin, sqrt, sum, + zeros, ) -import numpy as np +from numpy.linalg import inv +from scipy import interpolate try: from numba import njit except ImportError: + def njit(*args, **kwargs): return lambda f: f + from warnings import warn @@ -80,13 +82,15 @@ def find_critical_old(R, Z, psi, discard_xpoints=True): f = interpolate.RectBivariateSpline(R[:, 0], Z[0, :], psi) # Find candidate locations, based on minimising Bp^2 - Bp2 = (f(R, Z, dx=1, grid=False) ** 2 + f(R, Z, dy=1, grid=False) ** 2) / R ** 2 + Bp2 = ( + f(R, Z, dx=1, grid=False) ** 2 + f(R, Z, dy=1, grid=False) ** 2 + ) / R**2 # Get grid resolution, which determines a reasonable tolerance # for the Newton iteration search area dR = R[1, 0] - R[0, 0] dZ = Z[0, 1] - Z[0, 0] - radius_sq = 9 * (dR ** 2 + dZ ** 2) + radius_sq = 9 * (dR**2 + dZ**2) # Find local minima @@ -125,23 +129,25 @@ def find_critical_old(R, Z, psi, discard_xpoints=True): Br = -f(R1, Z1, dy=1, grid=False) / R1 Bz = f(R1, Z1, dx=1, grid=False) / R1 - if Br ** 2 + Bz ** 2 < 1e-6: + if Br**2 + Bz**2 < 1e-6: # Found a minimum. Classify as either # O-point or X-point dR = R[1, 0] - R[0, 0] dZ = Z[0, 1] - Z[0, 0] - d2dr2 = (psi[i + 2, j] - 2.0 * psi[i, j] + psi[i - 2, j]) / ( - 2.0 * dR - ) ** 2 - d2dz2 = (psi[i, j + 2] - 2.0 * psi[i, j] + psi[i, j - 2]) / ( - 2.0 * dZ - ) ** 2 + d2dr2 = ( + psi[i + 2, j] - 2.0 * psi[i, j] + psi[i - 2, j] + ) / (2.0 * dR) ** 2 + d2dz2 = ( + psi[i, j + 2] - 2.0 * psi[i, j] + psi[i, j - 2] + ) / (2.0 * dZ) ** 2 d2drdz = ( - (psi[i + 2, j + 2] - psi[i + 2, j - 2]) / (4.0 * dZ) - - (psi[i - 2, j + 2] - psi[i - 2, j - 2]) / (4.0 * dZ) + (psi[i + 2, j + 2] - psi[i + 2, j - 2]) + / (4.0 * dZ) + - (psi[i - 2, j + 2] - psi[i - 2, j - 2]) + / (4.0 * dZ) ) / (4.0 * dR) - D = d2dr2 * d2dz2 - d2drdz ** 2 + D = d2dr2 * d2dz2 - d2drdz**2 if D < 0.0: # Found X-point @@ -168,7 +174,9 @@ def find_critical_old(R, Z, psi, discard_xpoints=True): count += 1 # If (R1,Z1) is too far from (R0,Z0) then discard # or if we've taken too many iterations - if ((R1 - R0) ** 2 + (Z1 - Z0) ** 2 > radius_sq) or (count > 100): + if ((R1 - R0) ** 2 + (Z1 - Z0) ** 2 > radius_sq) or ( + count > 100 + ): # Discard this point break @@ -241,65 +249,72 @@ def remove_dup(points): return opoint, xpoint + # Remove duplicates def remove_dup(points): - result = [] - for n, p in enumerate(points): - dup = False - for p2 in result: - if (p[0] - p2[0]) ** 2 + (p[1] - p2[1]) ** 2 < 1e-5: - dup = True # Duplicate - break - if not dup: - result.append(p) # Add to the list - return result - -def find_critical(R, Z, psi, mask_inside_limiter=None, signIp=1, discard_xpoints=True): + result = [] + for n, p in enumerate(points): + dup = False + for p2 in result: + if (p[0] - p2[0]) ** 2 + (p[1] - p2[1]) ** 2 < 1e-5: + dup = True # Duplicate + break + if not dup: + result.append(p) # Add to the list + return result + + +def find_critical( + R, Z, psi, mask_inside_limiter=None, signIp=1, discard_xpoints=True +): # if old: # opoint, xpoint = find_critical_old(R,Z,psi, discard_xpoints) # else: opoint, xpoint = fastcrit(R, Z, psi, mask_inside_limiter) - if len(xpoint)>0 and (signIp is not None): + if len(xpoint) > 0 and (signIp is not None): # select xpoint with the correct ordering wrt Ip - xpoint = xpoint[((xpoint[:,2] - opoint[:1,2])*signIp) < 0] - if len(xpoint)>1: - # check distance to opoint and in case discard xpoints on non-monotonic LOS - closer_xpoint = np.argmin(np.linalg.norm((xpoint-opoint[:1])[:,:2], axis=-1)) - if closer_xpoint != 0: + xpoint = xpoint[((xpoint[:, 2] - opoint[:1, 2]) * signIp) < 0] + if len(xpoint) > 1: + # check distance to opoint and in case discard xpoints on non-monotonic LOS + closer_xpoint = np.argmin( + np.linalg.norm((xpoint - opoint[:1])[:, :2], axis=-1) + ) + if closer_xpoint != 0: f = interpolate.RectBivariateSpline(R[:, 0], Z[0, :], psi) result = False while result is False: result = discard_xpoints_f(opoint[0], xpoint[0], f) if result is False: xpoint = xpoint[1:] - result = len(xpoint)<1 + result = len(xpoint) < 1 print(xpoint) return opoint, xpoint + # # this is 10x faster if the numba import works; otherwise, @njit is the identity and fastcrit is 3x faster anyways @njit(fastmath=True, cache=True) def scan_for_crit(R, Z, psi): dR = R[1, 0] - R[0, 0] dZ = Z[0, 1] - Z[0, 0] - Bp2=np.zeros_like(psi) - psiR=Bp2.copy() - psiZ=Bp2.copy() - psiR[1:-1,1:-1] = 0.5*(psi[2:,1:-1]-psi[:-2,1:-1])/dR - psiZ[1:-1,1:-1] = 0.5*(psi[1:-1,2:]-psi[1:-1,:-2])/dZ + Bp2 = np.zeros_like(psi) + psiR = Bp2.copy() + psiZ = Bp2.copy() + psiR[1:-1, 1:-1] = 0.5 * (psi[2:, 1:-1] - psi[:-2, 1:-1]) / dR + psiZ[1:-1, 1:-1] = 0.5 * (psi[1:-1, 2:] - psi[1:-1, :-2]) / dZ # - # psiR[0,:]=(psi[1,:]-psi[0,:])/dR -# psiR[-1,:]=(psi[-1,:]-psi[-2,:])/dR -# psiR[1:-1,0]=(psi[1:,0]-psi[:-1,0])/dR -# psiR[1:-1,-1]=(psi[1:,-1]-psi[:-1,-0])/dR + # psiR[0,:]=(psi[1,:]-psi[0,:])/dR + # psiR[-1,:]=(psi[-1,:]-psi[-2,:])/dR + # psiR[1:-1,0]=(psi[1:,0]-psi[:-1,0])/dR + # psiR[1:-1,-1]=(psi[1:,-1]-psi[:-1,-0])/dR # - Bp2[:,:]= (psiR**2 + psiZ**2)#/R[:,:]**2 + Bp2[:, :] = psiR**2 + psiZ**2 # /R[:,:]**2 # - xpoint = [(-999.0,-999.0,-999.0)] - opoint = [(-999.0,-999.0,-999.0)] + xpoint = [(-999.0, -999.0, -999.0)] + opoint = [(-999.0, -999.0, -999.0)] # start off by finding coarse values of Bp2 closest to 0.0 as in Ben Dudsons's routine - for i in range(1,len(Bp2)-1): - for j in range(1,len(Bp2[0])-1): + for i in range(1, len(Bp2) - 1): + for j in range(1, len(Bp2[0]) - 1): if ( (Bp2[i, j] < Bp2[i + 1, j + 1]) and (Bp2[i, j] < Bp2[i + 1, j]) @@ -313,55 +328,70 @@ def scan_for_crit(R, Z, psi): # Found local minimum R0 = R[i, j] Z0 = Z[i, j] - fR , fZ = psiR[i,j], psiZ[i,j] - fRR = (psi[i+1,j]-2*psi[i,j]+psi[i-1,j])/dR**2 - fZZ = (psi[i,j+1]-2*psi[i,j]+psi[i,j-1])/dZ**2 - fRZ = 0.5*(psi[i+1,j+1] + psi[i-1,j-1] -psi[i-1,j+1] - psi[i+1,j-1] )/(dR*dZ) - det = fRR*fZZ-0.25*fRZ**2 + fR, fZ = psiR[i, j], psiZ[i, j] + fRR = (psi[i + 1, j] - 2 * psi[i, j] + psi[i - 1, j]) / dR**2 + fZZ = (psi[i, j + 1] - 2 * psi[i, j] + psi[i, j - 1]) / dZ**2 + fRZ = ( + 0.5 + * ( + psi[i + 1, j + 1] + + psi[i - 1, j - 1] + - psi[i - 1, j + 1] + - psi[i + 1, j - 1] + ) + / (dR * dZ) + ) + det = fRR * fZZ - 0.25 * fRZ**2 # - if det!=0: - delta_R = -(fR*fZZ-0.5*fRZ*fZ)/det - delta_Z = -(fZ*fRR-0.5*fRZ*fR)/det - if (np.abs(delta_R)<=dR and np.abs(delta_Z)<=dZ): - est_psi = psi[i,j]+0.5*(fR*delta_R + fZ*delta_Z) #+ 0.5*(fRR*delta_R**2 + fZZ*delta_Z**2 + fRZ*delta_R*delta_Z) - crpoint = (R0+delta_R , Z0+delta_Z , est_psi) - if det>0.0 : opoint = [crpoint] + opoint - else: xpoint = [crpoint] + xpoint + if det != 0: + delta_R = -(fR * fZZ - 0.5 * fRZ * fZ) / det + delta_Z = -(fZ * fRR - 0.5 * fRZ * fR) / det + if np.abs(delta_R) <= dR and np.abs(delta_Z) <= dZ: + est_psi = psi[i, j] + 0.5 * ( + fR * delta_R + fZ * delta_Z + ) # + 0.5*(fRR*delta_R**2 + fZZ*delta_Z**2 + fRZ*delta_R*delta_Z) + crpoint = (R0 + delta_R, Z0 + delta_Z, est_psi) + if det > 0.0: + opoint = [crpoint] + opoint + else: + xpoint = [crpoint] + xpoint xpoint = np.array(xpoint) opoint = np.array(opoint) # do NOT remove the "pop" command below, the lists were initialised with (-999.,-999.) so that numba could compile # xpoint.pop() # opoint.pop() - xpoint = xpoint[xpoint[:,0]>-990] - opoint = opoint[opoint[:,0]>-990] + xpoint = xpoint[xpoint[:, 0] > -990] + opoint = opoint[opoint[:, 0] > -990] return opoint, xpoint def fastcrit(R, Z, psi, mask_inside_limiter): - opoint , xpoint = scan_for_crit(R, Z, psi) + opoint, xpoint = scan_for_crit(R, Z, psi) len_opoint = len(opoint) if len_opoint == 0: # Can't order primary O-point, X-point so return - raise ValueError('No opoints found!') + raise ValueError("No opoints found!") # return opoint, xpoint elif mask_inside_limiter is not None: # remove any opoint outside the limiter - posR = np.argmin((R[:,:1].T - opoint[:,:1])**2, axis=1) - posZ = np.argmin((Z[:1,:] - opoint[:,1:2])**2, axis=1) - opoint = opoint[mask_inside_limiter[posR,posZ]] + posR = np.argmin((R[:, :1].T - opoint[:, :1]) ** 2, axis=1) + posZ = np.argmin((Z[:1, :] - opoint[:, 1:2]) ** 2, axis=1) + opoint = opoint[mask_inside_limiter[posR, posZ]] if len_opoint == 0: # Can't order primary O-point, X-point so return - raise ValueError('No opoints found!') + raise ValueError("No opoints found!") elif len_opoint > 1: # Find primary O-point by sorting by distance from middle of domain Rmid = 0.5 * (R[-1, 0] + R[0, 0]) Zmid = 0.5 * (Z[0, -1] + Z[0, 0]) - opoint_ordering = np.argsort((opoint[:,0]-Rmid)**2 + (opoint[:,1]-Zmid)**2) + opoint_ordering = np.argsort( + (opoint[:, 0] - Rmid) ** 2 + (opoint[:, 1] - Zmid) ** 2 + ) opoint = opoint[opoint_ordering] - + # # check that primary opoint is inside the limiter # result = False # while result is False: @@ -372,7 +402,7 @@ def fastcrit(R, Z, psi, mask_inside_limiter): # if len(opoint) == 0: # raise ValueError('No valid opoints found!') # else: - # result = True + # result = True psi_axis = opoint[0][2] # opoint.sort(key=lambda x: (x[0] - Rmid) ** 2 + (x[1] - Zmid) ** 2) @@ -415,18 +445,19 @@ def fastcrit(R, Z, psi, mask_inside_limiter): # xpoint = xpt_keep # Sort X-points by distance to primary O-point in psi space - if len(xpoint)>1: - xpoint_ordering = np.argsort((xpoint[:,2]-psi_axis)**2) + if len(xpoint) > 1: + xpoint_ordering = np.argsort((xpoint[:, 2] - psi_axis) ** 2) xpoint = xpoint[xpoint_ordering] # xpoint.sort(key=lambda x: (x[2] - psi_axis) ** 2) # return opoint, xpoint + def discard_xpoints_f(opoint, xpt, f): # Here opoint and xpt are individual critical points Ro, Zo, Po = opoint # The primary O-point result = False - + # for xpt in xpoint: Rx, Zx, Px = xpt @@ -457,8 +488,6 @@ def discard_xpoints_f(opoint, xpt, f): return result - - def core_mask(R, Z, psi, opoint, xpoint=[], psi_bndry=None): """ Mark the parts of the domain which are in the core @@ -525,7 +554,11 @@ def core_mask(R, Z, psi, opoint, xpoint=[], psi_bndry=None): while True: mask[i, j] = 1 # Mark as in the core - if (i < nx - 1) and (psin[i + 1, j] < 1.0) and (mask[i + 1, j] < 0.5): + if ( + (i < nx - 1) + and (psin[i + 1, j] < 1.0) + and (mask[i + 1, j] < 0.5) + ): stack.append((i + 1, j)) if (i > 0) and (psin[i - 1, j] < 1.0) and (mask[i - 1, j] < 0.5): stack.append((i - 1, j)) @@ -547,12 +580,16 @@ def core_mask(R, Z, psi, opoint, xpoint=[], psi_bndry=None): return mask + # def core_mask(R, Z, psi, opoint, xpoint=[], psi_bndry=None, old=False): # if old: return core_mask_old(R, Z, psi, opoint, xpoint, psi_bndry) # else: return inside_mask(R, Z, psi, opoint, xpoint, psi_bndry) + @njit(fastmath=True, cache=True) -def inside_mask(R, Z, psi, opoint, xpoint=[], mask_outside_limiter=None, psi_bndry=None): +def inside_mask( + R, Z, psi, opoint, xpoint=[], mask_outside_limiter=None, psi_bndry=None +): """ Similar to core_mask_old above, except: (1) it's all stuff that can be JIT-compiled @@ -580,10 +617,10 @@ def inside_mask(R, Z, psi, opoint, xpoint=[], mask_outside_limiter=None, psi_bnd ix = argmin(abs(R[:, 0] - rx)) jx = argmin(abs(Z[0, :] - zx)) xpt_inds.append((ix, jx)) - # Fill this point and all around with '2'# - ilo , ihi = min(max(ix-1,0),nx-1) , max(0,min(ix+1,nx-1)) - jlo , jhi = min(max(jx-1,0),ny-1) , max(0,min(jx+1,ny-1)) - mask[ilo:ihi+1,jlo:jhi+1] = 2 + # Fill this point and all around with '2'# + ilo, ihi = min(max(ix - 1, 0), nx - 1), max(0, min(ix + 1, nx - 1)) + jlo, jhi = min(max(jx - 1, 0), ny - 1), max(0, min(jx + 1, ny - 1)) + mask[ilo : ihi + 1, jlo : jhi + 1] = 2 # # Find nearest index to start rind = argmin(abs(R[:, 0] - Ro)) @@ -598,7 +635,7 @@ def inside_mask(R, Z, psi, opoint, xpoint=[], mask_outside_limiter=None, psi_bnd if (j > 0) and (psin[i, j - 1] < 1.0) and (mask[i, j - 1] < 0.5): stack.append((i, j - 1)) # Check the point above - if (j < ny -1) and (psin[i, j + 1] < 1.0) and (mask[i, j + 1] < 0.5): + if (j < ny - 1) and (psin[i, j + 1] < 1.0) and (mask[i, j + 1] < 0.5): stack.append((i, j + 1)) # # Scan along a row to the right @@ -611,16 +648,21 @@ def inside_mask(R, Z, psi, opoint, xpoint=[], mask_outside_limiter=None, psi_bnd # # Now return to X-point locations for ix, jx in xpt_inds: - ilo , ihi = min(max(ix-1,0),nx-1) , max(0,min(ix+1,nx-1)) - jlo , jhi = min(max(jx-1,0),ny-1) , max(0,min(jx+1,ny-1)) - mask[ilo:ihi+1,jlo:jhi+1] = 1*(mask[ilo:ihi+1,jlo:jhi+1]==1)*(psin[ilo:ihi+1,jlo:jhi+1]<1.0) + ilo, ihi = min(max(ix - 1, 0), nx - 1), max(0, min(ix + 1, nx - 1)) + jlo, jhi = min(max(jx - 1, 0), ny - 1), max(0, min(jx + 1, ny - 1)) + mask[ilo : ihi + 1, jlo : jhi + 1] = ( + 1 + * (mask[ilo : ihi + 1, jlo : jhi + 1] == 1) + * (psin[ilo : ihi + 1, jlo : jhi + 1] < 1.0) + ) # # remove effect of mask_outside_limiter - mask = (mask==1) + mask = mask == 1 return mask + def find_psisurface(eq, psifunc, r0, z0, r1, z1, psival=1.0, n=100, axis=None): """ eq - Equilibrium object @@ -704,7 +746,9 @@ def find_separatrix( # Avoid putting theta grid points exactly on the X-points xpoint_theta = arctan2(xpoint[0][0] - r0, xpoint[0][1] - z0) - xpoint_theta = xpoint_theta * (xpoint_theta >= 0) + (xpoint_theta + 2 * pi) * ( + xpoint_theta = xpoint_theta * (xpoint_theta >= 0) + ( + xpoint_theta + 2 * pi + ) * ( xpoint_theta < 0 ) # let's make it between 0 and 2*pi # How close in theta to allow theta grid points to the X-point @@ -732,7 +776,14 @@ def find_separatrix( def find_safety( - eq, npsi=1, psinorm=None, ntheta=128, psi=None, opoint=None, xpoint=None, axis=None + eq, + npsi=1, + psinorm=None, + ntheta=128, + psi=None, + opoint=None, + xpoint=None, + axis=None, ): """Find the safety factor for each value of psi Calculates equally spaced flux surfaces. Points on @@ -761,7 +812,9 @@ def find_safety( else: psinormal = (psi - opoint[0][2]) / (xpoint[0][2] - opoint[0][2]) - psifunc = interpolate.RectBivariateSpline(eq.R[:, 0], eq.Z[0, :], psinormal) + psifunc = interpolate.RectBivariateSpline( + eq.R[:, 0], eq.Z[0, :], psinormal + ) r0, z0 = opoint[0][0:2] @@ -770,7 +823,9 @@ def find_safety( # Avoid putting theta grid points exactly on the X-points xpoint_theta = arctan2(xpoint[0][0] - r0, xpoint[0][1] - z0) - xpoint_theta = xpoint_theta * (xpoint_theta >= 0) + (xpoint_theta + 2 * pi) * ( + xpoint_theta = xpoint_theta * (xpoint_theta >= 0) + ( + xpoint_theta + 2 * pi + ) * ( xpoint_theta < 0 ) # let's make it between 0 and 2*pi # How close in theta to allow theta grid points to the X-point @@ -816,19 +871,19 @@ def find_safety( fpol = eq.fpol(psirange[:]).reshape(npsi, 1) Br = eq.Br(r, z) Bz = eq.Bz(r, z) - Bthe = sqrt(Br ** 2 + Bz ** 2) + Bthe = sqrt(Br**2 + Bz**2) # Differentiate location w.r.t. index dr_di = (np.roll(r, 1, axis=1) - np.roll(r, -1, axis=1)) / 2.0 dz_di = (np.roll(z, 1, axis=1) - np.roll(z, -1, axis=1)) / 2.0 # Distance between points - dl = sqrt(dr_di ** 2 + dz_di ** 2) + dl = sqrt(dr_di**2 + dz_di**2) # Integrand - Btor/(R*Bthe) = Fpol/(R**2*Bthe) - qint = fpol / (r ** 2 * Bthe) + qint = fpol / (r**2 * Bthe) # Integral q = sum(qint * dl, axis=1) / (2 * pi) - return q \ No newline at end of file + return q diff --git a/freegs4e/divgeo.py b/freegs4e/divgeo.py index d7faa98..434ebd1 100644 --- a/freegs4e/divgeo.py +++ b/freegs4e/divgeo.py @@ -1,5 +1,4 @@ -from . import _divgeo -from . import geqdsk +from . import _divgeo, geqdsk def write(eq, fh, oxpoints=None): diff --git a/freegs4e/dump.py b/freegs4e/dump.py index 010de94..17a8889 100644 --- a/freegs4e/dump.py +++ b/freegs4e/dump.py @@ -36,12 +36,11 @@ import numpy as np +from . import boundary, machine from .equilibrium import Equilibrium -from .machine import Coil, Circuit, Solenoid, Wall, Machine -from .shaped_coil import ShapedCoil +from .machine import Circuit, Coil, Machine, Solenoid, Wall from .multi_coil import MultiCoil -from . import boundary -from . import machine +from .shaped_coil import ShapedCoil class OutputFormatNotAvailableError(Exception): @@ -127,7 +126,9 @@ def write_equilibrium(self, equilibrium): Solenoid.dtype: self.handle["solenoid_dtype"], } - equilibrium_group = self.handle.require_group(self.EQUILIBRIUM_GROUP_NAME) + equilibrium_group = self.handle.require_group( + self.EQUILIBRIUM_GROUP_NAME + ) equilibrium_group.create_dataset("Rmin", data=equilibrium.Rmin) equilibrium_group.create_dataset("Rmax", data=equilibrium.Rmax) @@ -139,10 +140,14 @@ def write_equilibrium(self, equilibrium): equilibrium_group.create_dataset("Z_1D", data=equilibrium.Z_1D) equilibrium_group.create_dataset("Z", data=equilibrium.Z) - equilibrium_group.create_dataset("current", data=equilibrium.plasmaCurrent()) + equilibrium_group.create_dataset( + "current", data=equilibrium.plasmaCurrent() + ) equilibrium_group["current"].attrs["title"] = "Plasma current [Amps]" - psi_id = equilibrium_group.create_dataset("psi", data=equilibrium.psi()) + psi_id = equilibrium_group.create_dataset( + "psi", data=equilibrium.psi() + ) psi_id.dims[0].label = "R" psi_id.dims[1].label = "Z" psi_id.dims.create_scale(equilibrium_group["R_1D"], "R") @@ -165,8 +170,12 @@ def write_equilibrium(self, equilibrium): tokamak_group = equilibrium_group.create_group(self.MACHINE_GROUP_NAME) if equilibrium.tokamak.wall is not None: - tokamak_group.create_dataset("wall_R", data=equilibrium.tokamak.wall.R) - tokamak_group.create_dataset("wall_Z", data=equilibrium.tokamak.wall.Z) + tokamak_group.create_dataset( + "wall_R", data=equilibrium.tokamak.wall.R + ) + tokamak_group.create_dataset( + "wall_Z", data=equilibrium.tokamak.wall.Z + ) coils_group = tokamak_group.create_group(self.COILS_GROUP_NAME) for label, coil in equilibrium.tokamak.coils: @@ -198,7 +207,9 @@ def read_equilibrium(self): # HDF5 reads strings as bytes by default, so convert to string def toString(s): try: - return str(s, "utf-8") # Convert bytes to string, using encoding + return str( + s, "utf-8" + ) # Convert bytes to string, using encoding except TypeError: return s # Probably already a string @@ -206,12 +217,15 @@ def toString(s): # freegs.machine and then call the `from_numpy_array` class # method def make_coil_set(thing): - return machine.__dict__[thing.attrs["freegs type"]].from_numpy_array(thing) + return machine.__dict__[ + thing.attrs["freegs type"] + ].from_numpy_array(thing) # Unfortunately this creates the coils in lexographical order # by label, losing the origin coils = [ - (toString(label), make_coil_set(coil)) for label, coil in coil_group.items() + (toString(label), make_coil_set(coil)) + for label, coil in coil_group.items() ] if "wall_R" in tokamak_group: diff --git a/freegs4e/equilibrium.py b/freegs4e/equilibrium.py index 690db18..e9a9f4c 100644 --- a/freegs4e/equilibrium.py +++ b/freegs4e/equilibrium.py @@ -20,23 +20,17 @@ along with FreeGS4E. If not, see . """ -from numpy import pi, meshgrid, linspace, exp, array import numpy as np +from numpy import array, exp, linspace, meshgrid, pi from scipy import interpolate from scipy.integrate import romb # Romberg integration +# Multigrid solver +from . import critical, machine, multigrid, polygons from .boundary import fixedBoundary, freeBoundary -from . import critical - -from . import polygons # Operators which define the G-S equation -from .gradshafranov import mu0, GSsparse, GSsparse4thOrder - -# Multigrid solver -from . import multigrid - -from . import machine +from .gradshafranov import GSsparse, GSsparse4thOrder, mu0 class Equilibrium: @@ -132,7 +126,9 @@ def __init__( generator = GSsparse4thOrder(Rmin, Rmax, Zmin, Zmax) else: raise ValueError( - "Invalid choice of order ({}). Valid values are 2 or 4.".format(order) + "Invalid choice of order ({}). Valid values are 2 or 4.".format( + order + ) ) self.order = order @@ -140,11 +136,14 @@ def __init__( nx, ny, generator, nlevels=1, ncycle=1, niter=2, direct=True ) - def create_psi_plasma_default(self, ): - """Creates a Gaussian starting guess for plasma_psi - """ - nx,ny = np.shape(self.R) - xx, yy = meshgrid(linspace(0, 1, nx), linspace(0, 1, ny), indexing="ij") + def create_psi_plasma_default( + self, + ): + """Creates a Gaussian starting guess for plasma_psi""" + nx, ny = np.shape(self.R) + xx, yy = meshgrid( + linspace(0, 1, nx), linspace(0, 1, ny), indexing="ij" + ) psi = exp(-((xx - 0.35) ** 2 + (yy - 0.5) ** 2)) psi[0, :] = 0.0 @@ -225,7 +224,9 @@ def poloidalBeta(self): dZ = self.Z[0, 1] - self.Z[0, 0] # Normalised psi - psi_norm = (self.psi() - self.psi_axis) / (self.psi_bndry - self.psi_axis) + psi_norm = (self.psi() - self.psi_axis) / ( + self.psi_bndry - self.psi_axis + ) # Plasma pressure pressure = self.pressure(psi_norm) @@ -288,7 +289,9 @@ def Btor(self, R, Z): Toroidal magnetic field """ # Normalised psi - psi_norm = (self.psiRZ(R, Z) - self.psi_axis) / (self.psi_bndry - self.psi_axis) + psi_norm = (self.psiRZ(R, Z) - self.psi_axis) / ( + self.psi_bndry - self.psi_axis + ) # Get f = R * Btor in the core. May be invalid outside the core fpol = self.fpol(psi_norm) @@ -378,9 +381,9 @@ def separatrix(self, ntheta=20): Returns an array of ntheta (R, Z) coordinates of the separatrix, equally spaced in geometric poloidal angle. """ - return array(critical.find_separatrix(self, ntheta=ntheta, psi=self.psi()))[ - :, 0:2 - ] + return array( + critical.find_separatrix(self, ntheta=ntheta, psi=self.psi()) + )[:, 0:2] def solve(self, profiles, Jtor=None, psi=None, psi_bndry=None): """ @@ -462,9 +465,11 @@ def _updatePlasmaPsi(self, plasma_psi): # if opt: self.psi_axis = opt[0][2] - if len(xpt)>0: + if len(xpt) > 0: self.psi_bndry = xpt[0][2] - self.mask = critical.inside_mask(self.R, self.Z, psi, opt, xpt, self.mask_outside_limiter) + self.mask = critical.inside_mask( + self.R, self.Z, psi, opt, xpt, self.mask_outside_limiter + ) # Use interpolation to find if a point is in the core. self.mask_func = interpolate.RectBivariateSpline( @@ -554,8 +559,10 @@ def innerOuterSeparatrix(self, Z=0.0): # Separatrix should now be between Rindex_inner and Rindex_inner+1 # Linear interpolation R_sep_in = ( - self.R[Rindex_inner, Zindex] * (1.0 - psinorm[Rindex_inner + 1]) - + self.R[Rindex_inner + 1, Zindex] * (psinorm[Rindex_inner] - 1.0) + self.R[Rindex_inner, Zindex] + * (1.0 - psinorm[Rindex_inner + 1]) + + self.R[Rindex_inner + 1, Zindex] + * (psinorm[Rindex_inner] - 1.0) ) / (psinorm[Rindex_inner] - psinorm[Rindex_inner + 1]) # Outer separatrix @@ -569,8 +576,10 @@ def innerOuterSeparatrix(self, Z=0.0): # Separatrix should now be between Rindex_outer-1 and Rindex_outer R_sep_out = ( - self.R[Rindex_outer, Zindex] * (1.0 - psinorm[Rindex_outer - 1]) - + self.R[Rindex_outer - 1, Zindex] * (psinorm[Rindex_outer] - 1.0) + self.R[Rindex_outer, Zindex] + * (1.0 - psinorm[Rindex_outer - 1]) + + self.R[Rindex_outer - 1, Zindex] + * (psinorm[Rindex_outer] - 1.0) ) / (psinorm[Rindex_outer] - psinorm[Rindex_outer - 1]) return R_sep_in, R_sep_out @@ -582,7 +591,9 @@ def intersectsWall(self): separatrix = self.separatrix() # Array [:,2] wall = self.tokamak.wall # Wall object with R and Z members (lists) - return polygons.intersect(separatrix[:, 0], separatrix[:, 1], wall.R, wall.Z) + return polygons.intersect( + separatrix[:, 0], separatrix[:, 1], wall.R, wall.Z + ) def magneticAxis(self): """Returns the location of the magnetic axis as a list [R,Z,psi]""" @@ -637,7 +648,9 @@ def geometricElongation(self, npoints=20): def aspectRatio(self, npoints=20): """Calculates the plasma aspect ratio""" - return self.Rgeometric(npoints=npoints) / self.minorRadius(npoints=npoints) + return self.Rgeometric(npoints=npoints) / self.minorRadius( + npoints=npoints + ) def effectiveElongation(self, R_wall_inner, R_wall_outer, npoints=300): """Calculates plasma effective elongation using the plasma volume""" @@ -732,7 +745,9 @@ def poloidalBeta2(self): dV = 2.0 * np.pi * R * dR * dZ # Normalised psi - psi_norm = (self.psi() - self.psi_axis) / (self.psi_bndry - self.psi_axis) + psi_norm = (self.psi() - self.psi_axis) / ( + self.psi_bndry - self.psi_axis + ) # Plasma pressure pressure = self.pressure(psi_norm) @@ -761,7 +776,9 @@ def toroidalBeta(self): dV = 2.0 * np.pi * R * dR * dZ # Normalised psi - psi_norm = (self.psi() - self.psi_axis) / (self.psi_bndry - self.psi_axis) + psi_norm = (self.psi() - self.psi_axis) / ( + self.psi_bndry - self.psi_axis + ) # Plasma pressure pressure = self.pressure(psi_norm) @@ -779,7 +796,9 @@ def toroidalBeta(self): def totalBeta(self): """Calculate plasma total beta""" - return 1.0 / ((1.0 / self.poloidalBeta2()) + (1.0 / self.toroidalBeta())) + return 1.0 / ( + (1.0 / self.poloidalBeta2()) + (1.0 / self.toroidalBeta()) + ) def refine(eq, nx=None, ny=None): @@ -851,7 +870,9 @@ def coarsen(eq): return result -def newDomain(eq, Rmin=None, Rmax=None, Zmin=None, Zmax=None, nx=None, ny=None): +def newDomain( + eq, Rmin=None, Rmax=None, Zmin=None, Zmax=None, nx=None, ny=None +): """Creates a new Equilibrium, solving in a different domain. The domain size (Rmin, Rmax, Zmin, Zmax) and resolution (nx,ny) are taken from the input equilibrium eq if not specified. @@ -871,7 +892,13 @@ def newDomain(eq, Rmin=None, Rmax=None, Zmin=None, Zmax=None, nx=None, ny=None): # Create a new equilibrium with the new domain result = Equilibrium( - tokamak=eq.tokamak, Rmin=Rmin, Rmax=Rmax, Zmin=Zmin, Zmax=Zmax, nx=nx, ny=ny + tokamak=eq.tokamak, + Rmin=Rmin, + Rmax=Rmax, + Zmin=Zmin, + Zmax=Zmax, + nx=nx, + ny=ny, ) # Calculate the current on the old grid @@ -908,10 +935,9 @@ def newDomain(eq, Rmin=None, Rmax=None, Zmin=None, Zmax=None, nx=None, ny=None): # Test the different spline interpolation routines - from numpy import ravel - import matplotlib.pyplot as plt - import machine + import matplotlib.pyplot as plt + from numpy import ravel tokamak = machine.TestTokamak() diff --git a/freegs4e/faster_shape.py b/freegs4e/faster_shape.py index a973d58..f079a12 100644 --- a/freegs4e/faster_shape.py +++ b/freegs4e/faster_shape.py @@ -45,7 +45,9 @@ def innerOuterSeparatrix(eq, profiles, Z=0.0): Zindex = np.argmin(abs(eq.Z[0, :] - Z)) # Normalise psi at this Z index - psinorm = (eq.psi()[:, Zindex] - eq.psi_axis) / (eq.psi_bndry - eq.psi_axis) + psinorm = (eq.psi()[:, Zindex] - eq.psi_axis) / ( + eq.psi_bndry - eq.psi_axis + ) # Start from the magnetic axis Rindex_axis = np.argmin(abs(eq.R[:, 0] - profiles.opt[0][0])) @@ -164,7 +166,9 @@ def Separatrix(eq, profiles, ntheta, psival=1.0): # Avoid putting theta grid points exactly on the X-points xpoint_theta = arctan2(xpoint[0][0] - r0, xpoint[0][1] - z0) - xpoint_theta = xpoint_theta * (xpoint_theta >= 0) + (xpoint_theta + 2 * pi) * ( + xpoint_theta = xpoint_theta * (xpoint_theta >= 0) + ( + xpoint_theta + 2 * pi + ) * ( xpoint_theta < 0 ) # let's make it between 0 and 2*pi # How close in theta to allow theta grid points to the X-point diff --git a/freegs4e/fieldtracer.py b/freegs4e/fieldtracer.py index 2888fb3..8a88d79 100644 --- a/freegs4e/fieldtracer.py +++ b/freegs4e/fieldtracer.py @@ -13,13 +13,13 @@ You should have received a copy of the GNU Lesser General Public License along with FreeGS4E. If not, see . -"""# Field line tracing +""" # Field line tracing from builtins import object import numpy as np -from scipy.integrate import odeint from scipy import interpolate +from scipy.integrate import odeint from . import critical @@ -49,8 +49,12 @@ def eqDomain(self, R, Z, evolving): evolving[ np.logical_or( - np.logical_or((R < self._eq.Rmin + eps), (R > self._eq.Rmax - eps)), - np.logical_or((Z < self._eq.Zmin + eps), (Z > self._eq.Zmax - eps)), + np.logical_or( + (R < self._eq.Rmin + eps), (R > self._eq.Rmax - eps) + ), + np.logical_or( + (Z < self._eq.Zmin + eps), (Z > self._eq.Zmax - eps) + ), ) ] = 0.0 return evolving @@ -117,7 +121,7 @@ def fieldDirection(self, pos, toroidal_angle, evolving, backward): Bz = self._eq.Bz(R, Z) Btor = self._eq.Btor(R, Z) - B = np.sqrt(Br ** 2 + Bz ** 2 + Btor ** 2) + B = np.sqrt(Br**2 + Bz**2 + Btor**2) # Detect when the boundary has been reached self.updateEvolving(R, Z, evolving) @@ -152,10 +156,16 @@ def follow(self, Rstart, Zstart, angles, rtol=None, backward=False): evolving = np.ones(array_shape) # (R,Z,length) with length=0 initially - position = np.column_stack((Rstart, Zstart, np.zeros(array_shape))).flatten() + position = np.column_stack( + (Rstart, Zstart, np.zeros(array_shape)) + ).flatten() result = odeint( - self.fieldDirection, position, angles, args=(evolving, backward), rtol=rtol + self.fieldDirection, + position, + angles, + args=(evolving, backward), + rtol=rtol, ) return result.reshape(angles.shape + array_shape + (3,)) @@ -176,7 +186,9 @@ def __init__(self, R, Z, length): self.length = length -def traceFieldLines(eq, solwidth=0.03, nlines=10, nturns=50, npoints=200, axis=None): +def traceFieldLines( + eq, solwidth=0.03, nlines=10, nturns=50, npoints=200, axis=None +): """Trace field lines from the outboard midplane Inputs diff --git a/freegs4e/geqdsk.py b/freegs4e/geqdsk.py index 40df809..53c903a 100755 --- a/freegs4e/geqdsk.py +++ b/freegs4e/geqdsk.py @@ -23,28 +23,18 @@ """ -from . import _geqdsk -from . import critical -from .equilibrium import Equilibrium -from .machine import Wall -from . import jtor -from . import control -from . import picard -from .gradshafranov import mu0 - -from scipy import interpolate -from numpy import ( - linspace, - reshape, - ravel, - zeros, - clip, -) import math -import numpy as np +import numpy as np +from numpy import clip, linspace, ravel, reshape, zeros +from scipy import interpolate from scipy.integrate import romb +from . import _geqdsk, control, critical, jtor, picard +from .equilibrium import Equilibrium +from .gradshafranov import mu0 +from .machine import Wall + def write(eq, fh, label=None, oxpoints=None, fileformat=_geqdsk.write): """ @@ -80,8 +70,10 @@ def write(eq, fh, label=None, oxpoints=None, fileformat=_geqdsk.write): data = { "nx": nx, "ny": ny, - "rdim": rmax - rmin, # Horizontal dimension in meter of computational box - "zdim": zmax - zmin, # Vertical dimension in meter of computational box + "rdim": rmax + - rmin, # Horizontal dimension in meter of computational box + "zdim": zmax + - zmin, # Vertical dimension in meter of computational box "rcentr": R0, # R in meter of vacuum toroidal magnetic field BCENTR "bcentr": fvac / R0, # Vacuum magnetic field at rcentr "rleft": rmin, # Minimum R in meter of rectangular computational box @@ -95,7 +87,9 @@ def write(eq, fh, label=None, oxpoints=None, fileformat=_geqdsk.write): data["cpasma"] = eq.plasmaCurrent() # Plasma current [A] - psinorm = linspace(0.0, 1.0, nx, endpoint=False) # Does not include separatrix + psinorm = linspace( + 0.0, 1.0, nx, endpoint=False + ) # Does not include separatrix data["fpol"] = eq.fpol(psinorm) data["pres"] = eq.pressure(psinorm) @@ -122,7 +116,9 @@ def write(eq, fh, label=None, oxpoints=None, fileformat=_geqdsk.write): # rbdry, zbdry contain the boundary of the plasma isoflux = np.array( - critical.find_separatrix(eq, ntheta=101, opoint=opoint, xpoint=xpoint, psi=psi) + critical.find_separatrix( + eq, ntheta=101, opoint=opoint, xpoint=xpoint, psi=psi + ) ) ind = np.argmin(isoflux[:, 1]) @@ -220,7 +216,9 @@ def read( import matplotlib.pyplot as plt if fit_sol and domain: - raise ValueError("Sorry, fit_sol cannot be used with the domain keyword") + raise ValueError( + "Sorry, fit_sol cannot be used with the domain keyword" + ) if axis is not None: show = True @@ -241,7 +239,8 @@ def read( if not (isPow2(nx - 1) and isPow2(ny - 1)): print( - "Warning: Input grid size %d x %d has sizes which are not 2^n+1" % (nx, ny) + "Warning: Input grid size %d x %d has sizes which are not 2^n+1" + % (nx, ny) ) rin = linspace(0, 1, nx) @@ -360,13 +359,21 @@ def q_func(psinorm): print("Changing resolution: {} x {}".format(newnx, newny)) # Create an interpolation function for Jtor and the input psi - Jtor_func = interpolate.RectBivariateSpline(eq.R[:, 0], eq.Z[0, :], Jtor) + Jtor_func = interpolate.RectBivariateSpline( + eq.R[:, 0], eq.Z[0, :], Jtor + ) psi_func = interpolate.RectBivariateSpline(eq.R[:, 0], eq.Z[0, :], psi) # Create a new Equilibrium object # (replacing previous 'eq') eq = Equilibrium( - tokamak=machine, Rmin=Rmin, Rmax=Rmax, Zmin=Zmin, Zmax=Zmax, nx=nx, ny=ny + tokamak=machine, + Rmin=Rmin, + Rmax=Rmax, + Zmin=Zmin, + Zmax=Zmax, + nx=nx, + ny=ny, ) # Interpolate Jtor and psi onto new grid @@ -480,6 +487,8 @@ def q_func(psinorm): ) # Save the control system to eq - eq.control = control.constrain(xpoints=xpoint, isoflux=isoflux, gamma=1e-14) + eq.control = control.constrain( + xpoints=xpoint, isoflux=isoflux, gamma=1e-14 + ) return eq diff --git a/freegs4e/gradshafranov.py b/freegs4e/gradshafranov.py index 16b7ec3..b3adb8a 100644 --- a/freegs4e/gradshafranov.py +++ b/freegs4e/gradshafranov.py @@ -20,14 +20,11 @@ """ -from numpy import zeros +from numpy import clip, pi, sqrt, zeros +from scipy.sparse import eye, lil_matrix # Elliptic integrals of first and second kind (K and E) -from scipy.special import ellipk, ellipe - -from numpy import sqrt, pi, clip - -from scipy.sparse import lil_matrix, eye +from scipy.special import ellipe, ellipk # Physical constants mu0 = 4e-7 * pi @@ -71,8 +68,8 @@ def __call__(self, psi, dR, dZ): b = zeros([nx, ny]) - invdR2 = 1.0 / dR ** 2 - invdZ2 = 1.0 / dZ ** 2 + invdR2 = 1.0 / dR**2 + invdZ2 = 1.0 / dZ**2 for x in range(1, nx - 1): R = self.Rmin + dR * x # Major radius of this point @@ -92,7 +89,7 @@ def diag(self, dR, dZ): Return the diagonal elements """ - return -2.0 / dR ** 2 - 2.0 / dZ ** 2 + return -2.0 / dR**2 - 2.0 / dZ**2 class GSsparse: @@ -122,8 +119,8 @@ def __call__(self, nx, ny): # Create a linked list sparse matrix A = eye(N, format="lil") - invdR2 = 1.0 / dR ** 2 - invdZ2 = 1.0 / dZ ** 2 + invdR2 = 1.0 / dR**2 + invdZ2 = 1.0 / dZ**2 for x in range(1, nx - 1): R = self.Rmin + dR * x # Major radius of this point @@ -159,7 +156,12 @@ class GSsparse4thOrder: # Coefficients for first derivatives # (index offset, weight) - centred_1st = [(-2, 1.0 / 12), (-1, -8.0 / 12), (1, 8.0 / 12), (2, -1.0 / 12)] + centred_1st = [ + (-2, 1.0 / 12), + (-1, -8.0 / 12), + (1, 8.0 / 12), + (2, -1.0 / 12), + ] offset_1st = [ (-1, -3.0 / 12), @@ -210,8 +212,8 @@ def __call__(self, nx, ny): # Create a linked list sparse matrix A = lil_matrix((N, N)) - invdR2 = 1.0 / dR ** 2 - invdZ2 = 1.0 / dZ ** 2 + invdR2 = 1.0 / dR**2 + invdZ2 = 1.0 / dZ**2 for x in range(1, nx - 1): R = self.Rmin + dR * x # Major radius of this point @@ -300,7 +302,9 @@ def GreensBz(Rc, Zc, R, Z, eps=1e-3): Bz = (1/R) d psi/dR """ - return (Greens(Rc, Zc, R + eps, Z) - Greens(Rc, Zc, R - eps, Z)) / (2.0 * eps * R) + return (Greens(Rc, Zc, R + eps, Z) - Greens(Rc, Zc, R - eps, Z)) / ( + 2.0 * eps * R + ) def GreensBr(Rc, Zc, R, Z, eps=1e-3): @@ -311,4 +315,6 @@ def GreensBr(Rc, Zc, R, Z, eps=1e-3): Br = -(1/R) d psi/dZ """ - return (Greens(Rc, Zc, R, Z - eps) - Greens(Rc, Zc, R, Z + eps)) / (2.0 * eps * R) + return (Greens(Rc, Zc, R, Z - eps) - Greens(Rc, Zc, R, Z + eps)) / ( + 2.0 * eps * R + ) diff --git a/freegs4e/jtor.py b/freegs4e/jtor.py index 804d8ea..cbd2e44 100644 --- a/freegs4e/jtor.py +++ b/freegs4e/jtor.py @@ -35,16 +35,14 @@ along with FreeGS4E. If not, see . """ -from scipy.integrate import romb, quad # Romberg integration -from . import critical -from .gradshafranov import mu0 - -from numpy import clip, zeros, reshape, sqrt, pi import numpy as np - +from numpy import clip, pi, reshape, sqrt, zeros +from scipy.integrate import quad, romb # Romberg integration from scipy.special import beta as spbeta from scipy.special import betainc as spbinc +from . import critical +from .gradshafranov import mu0 class Profile(object): @@ -128,7 +126,7 @@ def fpol(self, psinorm, out=None): ovals[i] = sqrt(2.0 * val + self.fvac() ** 2) return reshape(ovals, psinorm.shape) - + def Jtor_part1(self, R, Z, psi, psi_bndry=None, mask_outside_limiter=None): """ Similar code as original Jtor method, limited to identifying critical points. @@ -145,7 +143,7 @@ def Jtor_part1(self, R, Z, psi, psi_bndry=None, mask_outside_limiter=None): Value of the poloidal field flux at the boundary of the plasma (last closed flux surface), by default None mask_outside_limiter : np.ndarray Mask of points outside the limiter, if any, optional - + Returns ------- critical points @@ -157,25 +155,30 @@ def Jtor_part1(self, R, Z, psi, psi_bndry=None, mask_outside_limiter=None): """ # Analyse the equilibrium, finding O- and X-points - opt, xpt = critical.find_critical(R, Z, psi, self.mask_inside_limiter, self.Ip) + opt, xpt = critical.find_critical( + R, Z, psi, self.mask_inside_limiter, self.Ip + ) - if psi_bndry is not None: - diverted_core_mask = critical.inside_mask(R, Z, psi, opt, xpt, mask_outside_limiter, psi_bndry) - elif len(xpt)>0: + diverted_core_mask = critical.inside_mask( + R, Z, psi, opt, xpt, mask_outside_limiter, psi_bndry + ) + elif len(xpt) > 0: psi_bndry = xpt[0][2] self.psi_axis = opt[0][2] # # check correct sorting between psi_axis and psi_bndry - if (self.psi_axis-psi_bndry)*self.Ip < 0: - raise ValueError("Incorrect critical points! Likely due to not suitable psi_plasma") - diverted_core_mask = critical.inside_mask(R, Z, psi, opt, xpt, mask_outside_limiter, psi_bndry) + if (self.psi_axis - psi_bndry) * self.Ip < 0: + raise ValueError( + "Incorrect critical points! Likely due to not suitable psi_plasma" + ) + diverted_core_mask = critical.inside_mask( + R, Z, psi, opt, xpt, mask_outside_limiter, psi_bndry + ) else: # No X-points psi_bndry = psi[0, 0] diverted_core_mask = None - return opt, xpt, diverted_core_mask, psi_bndry - - + return opt, xpt, diverted_core_mask, psi_bndry class ConstrainBetapIp(Profile): @@ -260,7 +263,6 @@ def __init__(self, betap, Ip, fvac, alpha_m=1.0, alpha_n=2.0, Raxis=1.0): # psi_bndry = psi[0, 0] # mask = None - # dR = R[1, 0] - R[0, 0] # dZ = Z[0, 1] - Z[0, 0] @@ -300,7 +302,7 @@ def __init__(self, betap, Ip, fvac, alpha_m=1.0, alpha_n=2.0, Raxis=1.0): # # for j in range(1, ny - 1): # # if (psi_norm[i, j] >= 0.0) and (psi_norm[i, j] < 1.0): # # pfunc[i, j] = pshape(psi_norm[i, j]) - + # pfunc=shapeintegral0*(psi_bndry - psi_axis)*(1.0-spbinc(1./self.alpha_m , 1.0+self.alpha_n, np.clip(psi_norm,0.0001,0.9999)**(1/self.alpha_m))) # if mask is not None: @@ -312,37 +314,36 @@ def __init__(self, betap, Ip, fvac, alpha_m=1.0, alpha_n=2.0, Raxis=1.0): # intp = np.sum(pfunc)*dR*dZ # romb(romb(pfunc)) * dR * dZ - # LBeta0 = -self.betap * (mu0 / (8.0 * pi)) * self.Raxis * self.Ip ** 2 / intp + # LBeta0 = -self.betap * (mu0 / (8.0 * pi)) * self.Raxis * self.Ip ** 2 / intp - # # Integrate current components - # IR = np.sum(jtorshape * R /self.Raxis)*dR*dZ # romb(romb(jtorshape * R / self.Raxis)) * dR * dZ - # I_R = np.sum(jtorshape * self.Raxis / R)*dR*dZ # romb(romb(jtorshape * self.Raxis / R)) * dR * dZ + # # Integrate current components + # IR = np.sum(jtorshape * R /self.Raxis)*dR*dZ # romb(romb(jtorshape * R / self.Raxis)) * dR * dZ + # I_R = np.sum(jtorshape * self.Raxis / R)*dR*dZ # romb(romb(jtorshape * self.Raxis / R)) * dR * dZ - # # Toroidal plasma current Ip is - # # - # # Ip = L * (Beta0 * IR + (1-Beta0)*I_R) - # # = L*Beta0*(IR - I_R) + L*I_R - # # + # # Toroidal plasma current Ip is + # # + # # Ip = L * (Beta0 * IR + (1-Beta0)*I_R) + # # = L*Beta0*(IR - I_R) + L*I_R + # # - # L = self.Ip / I_R - LBeta0 * (IR / I_R - 1) - # Beta0 = LBeta0 / L + # L = self.Ip / I_R - LBeta0 * (IR / I_R - 1) + # Beta0 = LBeta0 / L - # # print("Constraints: L = %e, Beta0 = %e" % (L, Beta0)) + # # print("Constraints: L = %e, Beta0 = %e" % (L, Beta0)) - # # Toroidal current - # Jtor = L * (Beta0 * R / self.Raxis + (1 - Beta0) * self.Raxis / R) * jtorshape - # self.jtor = Jtor + # # Toroidal current + # Jtor = L * (Beta0 * R / self.Raxis + (1 - Beta0) * self.Raxis / R) * jtorshape + # self.jtor = Jtor - # self.L = L - # self.Beta0 = Beta0 - # self.psi_bndry = psi_bndry - # self.psi_axis = psi_axis + # self.L = L + # self.Beta0 = Beta0 + # self.psi_bndry = psi_bndry + # self.psi_axis = psi_axis - # self.xpt = xpt - # self.opt = opt + # self.xpt = xpt + # self.opt = opt - # return Jtor - + # return Jtor # def Jtor_part1(self, R, Z, psi, psi_bndry=None): # """ @@ -376,11 +377,10 @@ def __init__(self, betap, Ip, fvac, alpha_m=1.0, alpha_n=2.0, Raxis=1.0): # raise ValueError("No O-points found!") # return opt, xpt - - + def Jtor_part2(self, R, Z, psi, psi_axis, psi_bndry, mask): """ - Same code as original Jtor method, just split into two parts to enable + Same code as original Jtor method, just split into two parts to enable identification of limiter plasma configurations. In part 2 psi_axis is replaced by self.psi_axis @@ -413,9 +413,9 @@ def Jtor_part2(self, R, Z, psi, psi_axis, psi_bndry, mask): ValueError Raises ValueError if it cannot find an O-point """ - + if psi_bndry is None: - psi_bndry = psi[0,0] + psi_bndry = psi[0, 0] self.psi_bndry = psi_bndry dR = R[1, 0] - R[0, 0] @@ -427,7 +427,9 @@ def Jtor_part2(self, R, Z, psi, psi_axis, psi_bndry, mask): psi_norm = (psi - psi_axis) / (psi_bndry - psi_axis) # Current profile shape - jtorshape = (1.0 - np.clip(psi_norm, 0.0, 1.0) ** self.alpha_m) ** self.alpha_n + jtorshape = ( + 1.0 - np.clip(psi_norm, 0.0, 1.0) ** self.alpha_m + ) ** self.alpha_n if mask is not None: # If there is a masking function (X-points, limiters) @@ -449,7 +451,9 @@ def Jtor_part2(self, R, Z, psi, psi_axis, psi_bndry, mask): # p(psinorm) = - (L*Beta0/Raxis) * pshape(psinorm) # - shapeintegral0 = spbeta(1./self.alpha_m , 1.0+self.alpha_n)/self.alpha_m + shapeintegral0 = ( + spbeta(1.0 / self.alpha_m, 1.0 + self.alpha_n) / self.alpha_m + ) nx, ny = psi_norm.shape pfunc = zeros((nx, ny)) @@ -457,8 +461,19 @@ def Jtor_part2(self, R, Z, psi, psi_axis, psi_bndry, mask): # for j in range(1, ny - 1): # if (psi_norm[i, j] >= 0.0) and (psi_norm[i, j] < 1.0): # pfunc[i, j] = pshape(psi_norm[i, j]) - - pfunc=shapeintegral0*(psi_bndry - psi_axis)*(1.0-spbinc(1./self.alpha_m , 1.0+self.alpha_n, np.clip(psi_norm,0.0001,0.9999)**(1/self.alpha_m))) + + pfunc = ( + shapeintegral0 + * (psi_bndry - psi_axis) + * ( + 1.0 + - spbinc( + 1.0 / self.alpha_m, + 1.0 + self.alpha_n, + np.clip(psi_norm, 0.0001, 0.9999) ** (1 / self.alpha_m), + ) + ) + ) if mask is not None: pfunc *= mask @@ -467,13 +482,19 @@ def Jtor_part2(self, R, Z, psi, psi_axis, psi_bndry, mask): # betap = (8pi/mu0) * int(p)dRdZ / Ip^2 # = - (8pi/mu0) * (L*Beta0/Raxis) * intp / Ip^2 - intp = np.sum(pfunc)*dR*dZ # romb(romb(pfunc)) * dR * dZ + intp = np.sum(pfunc) * dR * dZ # romb(romb(pfunc)) * dR * dZ - LBeta0 = -self.betap * (mu0 / (8.0 * pi)) * self.Raxis * self.Ip ** 2 / intp + LBeta0 = ( + -self.betap * (mu0 / (8.0 * pi)) * self.Raxis * self.Ip**2 / intp + ) # Integrate current components - IR = np.sum(jtorshape * R /self.Raxis)*dR*dZ # romb(romb(jtorshape * R / self.Raxis)) * dR * dZ - I_R = np.sum(jtorshape * self.Raxis / R)*dR*dZ # romb(romb(jtorshape * self.Raxis / R)) * dR * dZ + IR = ( + np.sum(jtorshape * R / self.Raxis) * dR * dZ + ) # romb(romb(jtorshape * R / self.Raxis)) * dR * dZ + I_R = ( + np.sum(jtorshape * self.Raxis / R) * dR * dZ + ) # romb(romb(jtorshape * self.Raxis / R)) * dR * dZ # Toroidal plasma current Ip is # @@ -487,15 +508,17 @@ def Jtor_part2(self, R, Z, psi, psi_axis, psi_bndry, mask): # print("Constraints: L = %e, Beta0 = %e" % (L, Beta0)) # Toroidal current - Jtor = L * (Beta0 * R / self.Raxis + (1 - Beta0) * self.Raxis / R) * jtorshape + Jtor = ( + L + * (Beta0 * R / self.Raxis + (1 - Beta0) * self.Raxis / R) + * jtorshape + ) self.jtor = Jtor self.L = L self.Beta0 = Beta0 return Jtor - - # Profile functions def pprime(self, pn): @@ -549,7 +572,7 @@ def __init__(self, paxis, Ip, fvac, alpha_m=1.0, alpha_n=2.0, Raxis=1.0): # parameter to indicate that this is coming from FreeGS4E self.fast = True - + # def Jtor(self, R, Z, psi, psi_bndry=None): # """Calculate toroidal plasma current @@ -659,12 +682,11 @@ def __init__(self, paxis, Ip, fvac, alpha_m=1.0, alpha_n=2.0, Raxis=1.0): # return Jtor - # def Jtor_part1(self, R, Z, psi, psi_bndry=None): # """ - # Same code as original Jtor method, just split into two parts to enable + # Same code as original Jtor method, just split into two parts to enable # identification of limiter plasma configurations. - + # Calculate toroidal plasma current # Jtor = L * (Beta0*R/Raxis + (1-Beta0)*Raxis/R)*jtorshape @@ -716,7 +738,6 @@ def __init__(self, paxis, Ip, fvac, alpha_m=1.0, alpha_n=2.0, Raxis=1.0): # # if (psi_axis-psi_bndry)*self.Ip < 0: # # raise ValueError("Incorrect critical points! Likely due to not suitable psi_plasma") - # # # added with respect to original Jtor # # self.xpt = xpt # # self.opt = opt @@ -725,10 +746,9 @@ def __init__(self, paxis, Ip, fvac, alpha_m=1.0, alpha_n=2.0, Raxis=1.0): # return opt, xpt - def Jtor_part2(self, R, Z, psi, psi_axis, psi_bndry, mask): """ - Same code as original Jtor method, just split into two parts to enable + Same code as original Jtor method, just split into two parts to enable identification of limiter plasma configurations. In part 2 psi_axis is replaced by self.psi_axis @@ -762,20 +782,21 @@ def Jtor_part2(self, R, Z, psi, psi_axis, psi_bndry, mask): Raises ValueError if it cannot find an O-point """ if psi_bndry is None: - psi_bndry = psi[0,0] + psi_bndry = psi[0, 0] self.psi_bndry = psi_bndry - + dR = R[1, 0] - R[0, 0] dZ = Z[0, 1] - Z[0, 0] - # Calculate normalised psi. # 0 = magnetic axis # 1 = plasma boundary psi_norm = (psi - psi_axis) / (psi_bndry - psi_axis) # Current profile shape - jtorshape = (1.0 - np.clip(psi_norm, 0.0, 1.0) ** self.alpha_m) ** self.alpha_n + jtorshape = ( + 1.0 - np.clip(psi_norm, 0.0, 1.0) ** self.alpha_m + ) ** self.alpha_n if mask is not None: # If there is a masking function (X-points, limiters) @@ -788,7 +809,9 @@ def Jtor_part2(self, R, Z, psi, psi_axis, psi_bndry, mask): # shapeintegral, _ = quad( # lambda x: (1.0 - x ** self.alpha_m) ** self.alpha_n, 0.0, 1.0 # ) - shapeintegral = spbeta(1./self.alpha_m , 1.0+self.alpha_n)/self.alpha_m + shapeintegral = ( + spbeta(1.0 / self.alpha_m, 1.0 + self.alpha_n) / self.alpha_m + ) shapeintegral *= psi_bndry - psi_axis # Pressure on axis is @@ -797,8 +820,12 @@ def Jtor_part2(self, R, Z, psi, psi_axis, psi_bndry, mask): # # Integrate current components - IR = np.sum(jtorshape * R /self.Raxis)*dR*dZ #romb(romb(jtorshape * R / self.Raxis)) * dR * dZ - I_R = np.sum(jtorshape * self.Raxis / R)*dR*dZ #romb(romb(jtorshape * self.Raxis / R)) * dR * dZ + IR = ( + np.sum(jtorshape * R / self.Raxis) * dR * dZ + ) # romb(romb(jtorshape * R / self.Raxis)) * dR * dZ + I_R = ( + np.sum(jtorshape * self.Raxis / R) * dR * dZ + ) # romb(romb(jtorshape * self.Raxis / R)) * dR * dZ # Toroidal plasma current Ip is # @@ -814,7 +841,11 @@ def Jtor_part2(self, R, Z, psi, psi_axis, psi_bndry, mask): # print("Constraints: L = %e, Beta0 = %e" % (L, Beta0)) # Toroidal current - Jtor = L * (Beta0 * R / self.Raxis + (1 - Beta0) * self.Raxis / R) * jtorshape + Jtor = ( + L + * (Beta0 * R / self.Raxis + (1 - Beta0) * self.Raxis / R) + * jtorshape + ) self.jtor = Jtor self.L = L @@ -822,8 +853,6 @@ def Jtor_part2(self, R, Z, psi, psi_axis, psi_bndry, mask): return Jtor - - # Profile functions def pprime(self, pn): """ @@ -845,7 +874,6 @@ def fvac(self): return self._fvac - class Fiesta_Topeol(Profile): """ As in Fiesta. Implements profile function as in Jeon arxiv:1503.03135 eq. 5 @@ -879,12 +907,11 @@ def __init__(self, Beta0, Ip, fvac, alpha_m=1.0, alpha_n=2.0, Raxis=1.0): # parameter to indicate that this is coming from FreeGS4E self.fast = True - # def Jtor_part1(self, R, Z, psi, psi_bndry=None): # """ - # Same code as original Jtor method, just split into two parts to enable + # Same code as original Jtor method, just split into two parts to enable # identification of limiter plasma configurations. - + # Calculate toroidal plasma current # Jtor = L * (Beta0*R/Raxis + (1-Beta0)*Raxis/R)*jtorshape @@ -936,7 +963,6 @@ def __init__(self, Beta0, Ip, fvac, alpha_m=1.0, alpha_n=2.0, Raxis=1.0): # # if (psi_axis-psi_bndry)*self.Ip < 0: # # raise ValueError("Incorrect critical points! Likely due to not suitable psi_plasma") - # # # added with respect to original Jtor # # self.xpt = xpt # # self.opt = opt @@ -945,10 +971,9 @@ def __init__(self, Beta0, Ip, fvac, alpha_m=1.0, alpha_n=2.0, Raxis=1.0): # return opt, xpt - def Jtor_part2(self, R, Z, psi, psi_axis, psi_bndry, mask): """ - Same code as original Jtor method, just split into two parts to enable + Same code as original Jtor method, just split into two parts to enable identification of limiter plasma configurations. In part 2 psi_axis is replaced by self.psi_axis @@ -982,9 +1007,9 @@ def Jtor_part2(self, R, Z, psi, psi_axis, psi_bndry, mask): Raises ValueError if it cannot find an O-point """ if psi_bndry is None: - psi_bndry = psi[0,0] + psi_bndry = psi[0, 0] self.psi_bndry = psi_bndry - + dR = R[1, 0] - R[0, 0] dZ = Z[0, 1] - Z[0, 0] @@ -994,26 +1019,26 @@ def Jtor_part2(self, R, Z, psi, psi_axis, psi_bndry, mask): psi_norm = (psi - psi_axis) / (psi_bndry - psi_axis) # Current profile shape - jtorshape = (1.0 - np.clip(psi_norm, 0.0, 1.0) ** self.alpha_m) ** self.alpha_n + jtorshape = ( + 1.0 - np.clip(psi_norm, 0.0, 1.0) ** self.alpha_m + ) ** self.alpha_n if mask is not None: # If there is a masking function (X-points, limiters) jtorshape *= mask - - # Toroidal current - Jtor = (self.Beta0 * R / self.Raxis + (1 - self.Beta0) * self.Raxis / R) * jtorshape - L = self.Ip/(np.sum(Jtor)*dR*dZ) - self.jtor = L*Jtor + Jtor = ( + self.Beta0 * R / self.Raxis + (1 - self.Beta0) * self.Raxis / R + ) * jtorshape + L = self.Ip / (np.sum(Jtor) * dR * dZ) + self.jtor = L * Jtor self.L = L # self.Beta0 = Beta0 return self.jtor - - # Profile functions def pprime(self, pn): """ @@ -1040,13 +1065,23 @@ class Lao85(Profile): Implements Lao profile as from eqs. 2,4,5 in Lao et al 1985 Nucl.Fus.25 J = \lambda * (R/R_axis P' + Raxis/R FF'/mu0) - - with P' = sum(alpha_i x^i) - sum(alpha_i) x^(n_P+1) - FF' = sum(beta_i x^i) - sum(beta_i) x^(n_F+1) + + with P' = sum(alpha_i x^i) - sum(alpha_i) x^(n_P+1) + FF' = sum(beta_i x^i) - sum(beta_i) x^(n_F+1) """ - def __init__(self, Ip, fvac, alpha, beta, alpha_logic=True, beta_logic=True, Raxis=1, Ip_logic=True): + def __init__( + self, + Ip, + fvac, + alpha, + beta, + alpha_logic=True, + beta_logic=True, + Raxis=1, + Ip_logic=True, + ): """ Ip - Plasma current [Amps] fvac - Vacuum f = R*Bt @@ -1059,19 +1094,18 @@ def __init__(self, Ip, fvac, alpha, beta, alpha_logic=True, beta_logic=True, Rax Raxis - R used in p' and ff' components """ - # Set parameters for later use self.alpha = np.array(alpha) self.alpha_logic = alpha_logic if alpha_logic: self.alpha = np.concatenate((self.alpha, [-np.sum(self.alpha)])) - self.alpha_exp = np.arange(0,len(self.alpha)) + self.alpha_exp = np.arange(0, len(self.alpha)) self.beta = np.array(beta) self.beta_logic = beta_logic if beta_logic: self.beta = np.concatenate((self.beta, [-np.sum(self.beta)])) - self.beta_exp = np.arange(0,len(self.beta)) + self.beta_exp = np.arange(0, len(self.beta)) self.Ip = Ip self.Ip_logic = Ip_logic @@ -1081,14 +1115,12 @@ def __init__(self, Ip, fvac, alpha, beta, alpha_logic=True, beta_logic=True, Rax # parameter to indicate that this is coming from FreeGS4E self.fast = True - - # def Jtor_part1(self, R, Z, psi, psi_bndry=None): # """ - # Same code as original Jtor method, just split into two parts to enable + # Same code as original Jtor method, just split into two parts to enable # identification of limiter plasma configurations. - + # Calculate toroidal plasma current # Parameters @@ -1132,7 +1164,6 @@ def __init__(self, Ip, fvac, alpha, beta, alpha_logic=True, beta_logic=True, Rax # # # check correct sorting between psi_axis and psi_bndry # # # if (psi_axis-psi_bndry)*self.Ip < 0: # # # raise ValueError("Incorrect critical points! Likely due to not suitable psi_plasma") - # # # added with respect to original Jtor # # self.xpt = xpt @@ -1142,10 +1173,9 @@ def __init__(self, Ip, fvac, alpha, beta, alpha_logic=True, beta_logic=True, Rax # return opt, xpt - def Jtor_part2(self, R, Z, psi, psi_axis, psi_bndry, mask): """ - + Parameters ---------- R : np.ndarray @@ -1168,13 +1198,12 @@ def Jtor_part2(self, R, Z, psi, psi_axis, psi_bndry, mask): Raises ValueError if it cannot find an O-point """ if psi_bndry is None: - psi_bndry = psi[0,0] + psi_bndry = psi[0, 0] self.psi_bndry = psi_bndry - + dR = R[1, 0] - R[0, 0] dZ = Z[0, 1] - Z[0, 0] - # Calculate normalised psi. # 0 = magnetic axis # 1 = plasma boundary @@ -1183,15 +1212,21 @@ def Jtor_part2(self, R, Z, psi, psi_axis, psi_bndry, mask): # Current profile shape # Pprime - pprime_term = psi_norm[np.newaxis, :, :]**self.alpha_exp[:, np.newaxis, np.newaxis] - pprime_term *= self.alpha[:, np.newaxis, np.newaxis] + pprime_term = ( + psi_norm[np.newaxis, :, :] + ** self.alpha_exp[:, np.newaxis, np.newaxis] + ) + pprime_term *= self.alpha[:, np.newaxis, np.newaxis] pprime_term = np.sum(pprime_term, axis=0) - pprime_term *= R/self.Raxis + pprime_term *= R / self.Raxis # FFprime - ffprime_term = psi_norm[np.newaxis, :, :]**self.beta_exp[:, np.newaxis, np.newaxis] - ffprime_term *= self.beta[:, np.newaxis, np.newaxis] + ffprime_term = ( + psi_norm[np.newaxis, :, :] + ** self.beta_exp[:, np.newaxis, np.newaxis] + ) + ffprime_term *= self.beta[:, np.newaxis, np.newaxis] ffprime_term = np.sum(ffprime_term, axis=0) - ffprime_term *= self.Raxis/R + ffprime_term *= self.Raxis / R ffprime_term /= mu0 # Sum together Jtor = pprime_term + ffprime_term @@ -1204,38 +1239,44 @@ def Jtor_part2(self, R, Z, psi, psi_axis, psi_bndry, mask): jtorIp = np.sum(Jtor) if jtorIp == 0: self.problem_psi = psi - raise ValueError("Total plasma current is zero! Cannot renormalise.") - L = self.Ip/(jtorIp*dR*dZ) - Jtor = L*Jtor + raise ValueError( + "Total plasma current is zero! Cannot renormalise." + ) + L = self.Ip / (jtorIp * dR * dZ) + Jtor = L * Jtor else: - L = 1. + L = 1.0 self.jtor = Jtor.copy() self.L = L return self.jtor - - # Profile functions def pprime(self, pn): """ dp/dpsi as a function of normalised psi. 0 outside core Calculate pprimeshape inside the core only """ - shape = np.clip(np.array(pn), 0.0, 1.0)[np.newaxis, :]**self.alpha_exp[:, np.newaxis] + shape = ( + np.clip(np.array(pn), 0.0, 1.0)[np.newaxis, :] + ** self.alpha_exp[:, np.newaxis] + ) shape *= self.alpha[:, np.newaxis] - shape = np.sum(shape, axis=0) - return self.L * shape / self.Raxis + shape = np.sum(shape, axis=0) + return self.L * shape / self.Raxis def ffprime(self, pn): """ f * df/dpsi as a function of normalised psi. 0 outside core. Calculate ffprimeshape inside the core only. """ - shape = np.clip(np.array(pn), 0.0, 1.0)[np.newaxis, :]**self.beta_exp[:, np.newaxis] + shape = ( + np.clip(np.array(pn), 0.0, 1.0)[np.newaxis, :] + ** self.beta_exp[:, np.newaxis] + ) shape *= self.beta[:, np.newaxis] - shape = np.sum(shape, axis=0) + shape = np.sum(shape, axis=0) return self.L * shape * self.Raxis def fvac(self): @@ -1250,7 +1291,9 @@ class ProfilesPprimeFfprime: """ - def __init__(self, pprime_func, ffprime_func, fvac, p_func=None, f_func=None): + def __init__( + self, pprime_func, ffprime_func, fvac, p_func=None, f_func=None + ): """ pprime_func(psi_norm) - A function which returns dp/dpsi at given normalised flux ffprime_func(psi_norm) - A function which returns f*df/dpsi at given normalised flux (f = R*Bt) diff --git a/freegs4e/machine.py b/freegs4e/machine.py index 170ae76..6433cbc 100644 --- a/freegs4e/machine.py +++ b/freegs4e/machine.py @@ -20,14 +20,13 @@ along with FreeGS4E. If not, see . """ -from .gradshafranov import Greens, GreensBr, GreensBz - -from numpy import linspace import numpy as np +from numpy import linspace -from .coil import Coil, AreaCurrentLimit -from .shaped_coil import ShapedCoil +from .coil import AreaCurrentLimit, Coil +from .gradshafranov import Greens, GreensBr, GreensBz from .multi_coil import MultiCoil +from .shaped_coil import ShapedCoil # We need this for the `label` part of the Circuit dtype for writing # to HDF5 files. See the following for information: @@ -215,7 +214,9 @@ def from_numpy_array(cls, value): # HDF5 reads strings as bytes by default, so convert to string def toString(s): try: - return str(s, "utf-8") # Convert bytes to string, using encoding + return str( + s, "utf-8" + ) # Convert bytes to string, using encoding except TypeError: return s # Probably already a string @@ -245,7 +246,13 @@ def plot(self, axis=None, show=False): def MirroredCoil( - R, Z, current=0.0, turns=1, control=True, area=AreaCurrentLimit(), symmetric=True + R, + Z, + current=0.0, + turns=1, + control=True, + area=AreaCurrentLimit(), + symmetric=True, ): """ Create a pair of coils, at +/- Z @@ -256,12 +263,26 @@ def MirroredCoil( [ ( "U", - Coil(R, Z, current=current, turns=turns, control=control, area=area), + Coil( + R, + Z, + current=current, + turns=turns, + control=control, + area=area, + ), 1.0, ), ( "L", - Coil(R, Z, current=current, turns=turns, control=control, area=area), + Coil( + R, + Z, + current=current, + turns=turns, + control=control, + area=area, + ), 1.0 if symmetric else -1.0, ), ] @@ -378,7 +399,12 @@ def getForces(self, equilibrium): def __repr__(self): return "Solenoid(Rs={0}, Zsmin={1}, Zsmax={2}, current={3}, Ns={4}, control={5})".format( - self.Rs, self.Zsmin, self.Zsmax, self.current, self.Ns, self.control + self.Rs, + self.Zsmin, + self.Zsmax, + self.current, + self.Ns, + self.control, ) def __eq__(self, other): @@ -399,7 +425,14 @@ def to_numpy_array(self): Helper method for writing output """ return np.array( - (self.Rs, self.Zsmin, self.Zsmax, self.Ns, self.current, self.control), + ( + self.Rs, + self.Zsmin, + self.Zsmax, + self.Ns, + self.current, + self.control, + ), dtype=self.dtype, ) @@ -467,7 +500,10 @@ def __repr__(self): def __eq__(self, other): # Other Machine might be equivalent except for order of # coils. Assume this doesn't actually matter - return sorted(self.coils) == sorted(other.coils) and self.wall == other.wall + return ( + sorted(self.coils) == sorted(other.coils) + and self.wall == other.wall + ) def __ne__(self, other): return not self == other @@ -476,7 +512,9 @@ def __getitem__(self, name): for label, coil in self.coils: if label == name: return coil - raise KeyError("Machine does not contain coil with label '{0}'".format(name)) + raise KeyError( + "Machine does not contain coil with label '{0}'".format(name) + ) def psi(self, R, Z): """ @@ -534,21 +572,27 @@ def controlBr(self, R, Z): Returns a list of control responses for Br at the given (R,Z) location(s). """ - return [coil.controlBr(R, Z) for label, coil in self.coils if coil.control] + return [ + coil.controlBr(R, Z) for label, coil in self.coils if coil.control + ] def controlBz(self, R, Z): """ Returns a list of control responses for Bz at the given (R,Z) location(s) """ - return [coil.controlBz(R, Z) for label, coil in self.coils if coil.control] + return [ + coil.controlBz(R, Z) for label, coil in self.coils if coil.control + ] def controlPsi(self, R, Z): """ Returns a list of control responses for psi at the given (R,Z) location(s) """ - return [coil.controlPsi(R, Z) for label, coil in self.coils if coil.control] + return [ + coil.controlPsi(R, Z) for label, coil in self.coils if coil.control + ] def controlAdjust(self, current_change): """ @@ -639,15 +683,23 @@ def TestTokamak(): coils = [ ( "P1L", - ShapedCoil([(0.95, -1.15), (0.95, -1.05), (1.05, -1.05), (1.05, -1.15)]), + ShapedCoil( + [(0.95, -1.15), (0.95, -1.05), (1.05, -1.05), (1.05, -1.15)] + ), + ), + ( + "P1U", + ShapedCoil( + [(0.95, 1.15), (0.95, 1.05), (1.05, 1.05), (1.05, 1.15)] + ), ), - ("P1U", ShapedCoil([(0.95, 1.15), (0.95, 1.05), (1.05, 1.05), (1.05, 1.15)])), ("P2L", Coil(1.75, -0.6)), ("P2U", Coil(1.75, 0.6)), ] wall = Wall( - [0.75, 0.75, 1.5, 1.8, 1.8, 1.5], [-0.85, 0.85, 0.85, 0.25, -0.25, -0.85] # R + [0.75, 0.75, 1.5, 1.8, 1.8, 1.5], + [-0.85, 0.85, 0.85, 0.25, -0.25, -0.85], # R ) # Z return Machine(coils, wall) @@ -714,20 +766,43 @@ def MAST_sym(): coils = [ ( "P2", - Circuit([("P2U", Coil(0.49, 1.76), 1.0), ("P2L", Coil(0.49, -1.76), 1.0)]), + Circuit( + [ + ("P2U", Coil(0.49, 1.76), 1.0), + ("P2L", Coil(0.49, -1.76), 1.0), + ] + ), + ), + ( + "P3", + Circuit( + [("P3U", Coil(1.1, 1.1), 1.0), ("P3L", Coil(1.1, -1.1), 1.0)] + ), ), - ("P3", Circuit([("P3U", Coil(1.1, 1.1), 1.0), ("P3L", Coil(1.1, -1.1), 1.0)])), ( "P4", Circuit( - [("P4U", Coil(1.51, 1.095), 1.0), ("P4L", Coil(1.51, -1.095), 1.0)] + [ + ("P4U", Coil(1.51, 1.095), 1.0), + ("P4L", Coil(1.51, -1.095), 1.0), + ] ), ), ( "P5", - Circuit([("P5U", Coil(1.66, 0.52), 1.0), ("P5L", Coil(1.66, -0.52), 1.0)]), + Circuit( + [ + ("P5U", Coil(1.66, 0.52), 1.0), + ("P5L", Coil(1.66, -0.52), 1.0), + ] + ), + ), + ( + "P6", + Circuit( + [("P6U", Coil(1.5, 0.9), 1.0), ("P6L", Coil(1.5, -0.9), -1.0)] + ), ), - ("P6", Circuit([("P6U", Coil(1.5, 0.9), 1.0), ("P6L", Coil(1.5, -0.9), -1.0)])), ("P1", Solenoid(0.15, -1.45, 1.45, 100)), ] diff --git a/freegs4e/multi_coil.py b/freegs4e/multi_coil.py index c0202de..f1a3d80 100644 --- a/freegs4e/multi_coil.py +++ b/freegs4e/multi_coil.py @@ -22,11 +22,12 @@ along with FreeGS4E. If not, see . """ +import matplotlib.pyplot as plt import numpy as np -from .coil import Coil, AreaCurrentLimit -from .gradshafranov import Greens, GreensBr, GreensBz from matplotlib.patches import Rectangle -import matplotlib.pyplot as plt + +from .coil import AreaCurrentLimit, Coil +from .gradshafranov import Greens, GreensBr, GreensBz class MultiCoil(Coil): @@ -259,7 +260,7 @@ def plot(self, axis=None, show=False): plt.plot(self.Rfil, self.Zfil, "bo") return axis - + def plot_nke(self, axis=None, show=False): """ Plot the location of each filament, using axis if given @@ -271,13 +272,14 @@ def plot_nke(self, axis=None, show=False): axis = fig.add_subplot(111) for i in np.arange(len(self.Rfil)): - self.rectangle = Rectangle((self.Rfil[i] - self.dR/2, - self.Zfil[i] - self.dZ/2), - width=self.dR, - height=self.dZ, - facecolor = 'dodgerblue', - edgecolor='b', - linewidth=1) + self.rectangle = Rectangle( + (self.Rfil[i] - self.dR / 2, self.Zfil[i] - self.dZ / 2), + width=self.dR, + height=self.dZ, + facecolor="dodgerblue", + edgecolor="b", + linewidth=1, + ) axis.add_patch(self.rectangle) return axis diff --git a/freegs4e/multigrid.py b/freegs4e/multigrid.py index c310ff8..dca9fc1 100644 --- a/freegs4e/multigrid.py +++ b/freegs4e/multigrid.py @@ -27,9 +27,9 @@ """ -from numpy import zeros, max, abs, reshape -from scipy.sparse.linalg import factorized +from numpy import abs, max, reshape, zeros from scipy.sparse import eye +from scipy.sparse.linalg import factorized class MGDirect: @@ -109,7 +109,9 @@ def __call__(self, xi, bi, ncycle=None, niter=None): return x.reshape(xi.shape) -def createVcycle(nx, ny, generator, nlevels=4, ncycle=1, niter=10, direct=True): +def createVcycle( + nx, ny, generator, nlevels=4, ncycle=1, niter=10, direct=True +): """ Create a hierarchy of solvers in a multigrid V-cycle @@ -277,7 +279,9 @@ def smoothVcycle(A, x, b, dx, dy, niter=10, sublevels=0, direct=True): # smooth this error Cx = zeros(Cerror.shape) - Cx = smoothVcycle(A, Cx, Cerror, dx * 2.0, dy * 2.0, niter, sublevels - 1) + Cx = smoothVcycle( + A, Cx, Cerror, dx * 2.0, dy * 2.0, niter, sublevels - 1 + ) # Prolong the solution xupdate = interpolate(Cx) @@ -325,14 +329,14 @@ def __call__(self, f, dx, dy): for y in range(1, ny - 1): # Loop over points in the domain - b[x, y] = (f[x - 1, y] - 2 * f[x, y] + f[x + 1, y]) / dx ** 2 + ( + b[x, y] = (f[x - 1, y] - 2 * f[x, y] + f[x + 1, y]) / dx**2 + ( f[x, y - 1] - 2 * f[x, y] + f[x, y + 1] - ) / dy ** 2 + ) / dy**2 return b def diag(self, dx, dy): - return -2.0 / dx ** 2 - 2.0 / dy ** 2 + return -2.0 / dx**2 - 2.0 / dy**2 class LaplaceSparse: @@ -350,19 +354,19 @@ def __call__(self, nx, ny): for x in range(1, nx - 1): for y in range(1, ny - 1): row = x * ny + y - A[row, row] = -2.0 / dx ** 2 - 2.0 / dy ** 2 + A[row, row] = -2.0 / dx**2 - 2.0 / dy**2 # y-1 - A[row, row - 1] = 1.0 / dy ** 2 + A[row, row - 1] = 1.0 / dy**2 # y+1 - A[row, row + 1] = 1.0 / dy ** 2 + A[row, row + 1] = 1.0 / dy**2 # x-1 - A[row, row - ny] = 1.0 / dx ** 2 + A[row, row - ny] = 1.0 / dx**2 # x+1 - A[row, row + ny] = 1.0 / dx ** 2 + A[row, row + ny] = 1.0 / dx**2 # Convert to Compressed Sparse Row (CSR) format return A.tocsr() @@ -371,11 +375,11 @@ def __call__(self, nx, ny): # Test case - from numpy import meshgrid, exp, linspace - import matplotlib.pyplot as plt - from timeit import default_timer as timer + import matplotlib.pyplot as plt + from numpy import exp, linspace, meshgrid + nx = 65 ny = 65 @@ -384,7 +388,7 @@ def __call__(self, nx, ny): xx, yy = meshgrid(linspace(0, 1, nx), linspace(0, 1, ny)) - rhs = exp(-((xx - 0.5) ** 2 + (yy - 0.5) ** 2) / 0.4 ** 2) + rhs = exp(-((xx - 0.5) ** 2 + (yy - 0.5) ** 2) / 0.4**2) rhs[0, :] = 0.0 rhs[:, 0] = 0.0 @@ -428,7 +432,13 @@ def __call__(self, nx, ny): start = timer() solver = createVcycle( - nx, ny, LaplaceSparse(1.0, 1.0), ncycle=2, niter=5, nlevels=4, direct=True + nx, + ny, + LaplaceSparse(1.0, 1.0), + ncycle=2, + niter=5, + nlevels=4, + direct=True, ) start_solve = timer() diff --git a/freegs4e/optimise.py b/freegs4e/optimise.py index 6a79ae0..fcaba52 100644 --- a/freegs4e/optimise.py +++ b/freegs4e/optimise.py @@ -21,13 +21,12 @@ """ -from . import optimiser -from . import picard +from math import sqrt import matplotlib.pyplot as plt from freegs.plotting import plotEquilibrium -from math import sqrt +from . import optimiser, picard # Measures which operate on Equilibrium objects @@ -42,7 +41,9 @@ def max_abs_coil_current(eq): def max_coil_force(eq): forces = eq.tokamak.getForces() # dictionary - return max(sqrt(force[0] ** 2 + force[1] ** 2) for force in forces.values()) + return max( + sqrt(force[0] ** 2 + force[1] ** 2) for force in forces.values() + ) def no_wall_intersection(eq): @@ -132,6 +133,7 @@ def get(self, eq): # Monitor optimisation solutions + # Plot and save the best equilibrium each generation class PlotMonitor: """ @@ -149,12 +151,16 @@ def __call__(self, generation, best, population): # Update the canvas and pause # Note, a short pause is needed to force drawing update self.fig.canvas.draw() - self.axis.set_title("Generation: {} Score: {}".format(generation, best[0])) + self.axis.set_title( + "Generation: {} Score: {}".format(generation, best[0]) + ) self.fig.savefig("generation_{}.pdf".format(generation)) plt.pause(0.5) -def optimise(eq, controls, measure, maxgen=10, N=10, CR=0.3, F=1.0, monitor=None): +def optimise( + eq, controls, measure, maxgen=10, N=10, CR=0.3, F=1.0, monitor=None +): """Use Differential Evolution to optimise an Equilibrium https://en.wikipedia.org/wiki/Differential_evolution @@ -203,5 +209,12 @@ def solve_and_measure(eq): # Call the generic optimiser, return optimiser.optimise( - eq, controls, solve_and_measure, maxgen=maxgen, N=N, CR=CR, F=F, monitor=monitor + eq, + controls, + solve_and_measure, + maxgen=maxgen, + N=N, + CR=CR, + F=F, + monitor=monitor, ) diff --git a/freegs4e/optimiser.py b/freegs4e/optimiser.py index 0e540ab..ee748a1 100644 --- a/freegs4e/optimiser.py +++ b/freegs4e/optimiser.py @@ -24,9 +24,9 @@ along with FreeGS4E. If not, see . """ -import random -import copy import bisect +import copy +import random def mutate(obj, controls): @@ -66,7 +66,9 @@ def pickUnique(N, m, e): return others -def optimise(obj, controls, measure, maxgen=10, N=10, CR=0.3, F=1.0, monitor=None): +def optimise( + obj, controls, measure, maxgen=10, N=10, CR=0.3, F=1.0, monitor=None +): """Use Differential Evolution to optimise an object https://en.wikipedia.org/wiki/Differential_evolution @@ -119,13 +121,16 @@ def optimise(obj, controls, measure, maxgen=10, N=10, CR=0.3, F=1.0, monitor=Non others = [population[index][1] for index in pickUnique(N, 3, [ai])] new_obj = copy.deepcopy(agent[1]) - R = random.randint(0, len(controls) - 1) # Pick a random control to modify + R = random.randint( + 0, len(controls) - 1 + ) # Pick a random control to modify for i, control in enumerate(controls): if i == R or random.random() < CR: control.set( new_obj, control.get(others[0]) - + F * (control.get(others[1]) - control.get(others[2])), + + F + * (control.get(others[1]) - control.get(others[2])), ) score = measure(new_obj) if score < agent[0]: diff --git a/freegs4e/picard.py b/freegs4e/picard.py index ec7b3f4..f96064d 100644 --- a/freegs4e/picard.py +++ b/freegs4e/picard.py @@ -17,7 +17,7 @@ along with FreeGS4E. If not, see . """ -from numpy import amin, amax, array +from numpy import amax, amin, array def solve( @@ -62,6 +62,7 @@ def solve( if show: import matplotlib.pyplot as plt + from .plotting import plotEquilibrium if pause > 0.0 and axis is None: @@ -128,11 +129,8 @@ def solve( raise RuntimeError( "Picard iteration failed to converge (too many iterations)" ) - - eq._profiles = profiles - - if convergenceInfo: - return array(psi_maxchange_iterations),\ - array(psi_relchange_iterations) + eq._profiles = profiles + if convergenceInfo: + return array(psi_maxchange_iterations), array(psi_relchange_iterations) diff --git a/freegs4e/plotting.py b/freegs4e/plotting.py index f0ac7ad..af25fb6 100644 --- a/freegs4e/plotting.py +++ b/freegs4e/plotting.py @@ -18,7 +18,8 @@ """ -from numpy import linspace, amin, amax +from numpy import amax, amin, linspace + from . import critical @@ -68,7 +69,9 @@ def plotConstraints(control, axis=None, show=True): return axis -def plotEquilibrium(eq, axis=None, show=True, oxpoints=True, wall=True, limiter=True): +def plotEquilibrium( + eq, axis=None, show=True, oxpoints=True, wall=True, limiter=True +): """ Plot the equilibrium flux surfaces @@ -96,13 +99,13 @@ def plotEquilibrium(eq, axis=None, show=True, oxpoints=True, wall=True, limiter= try: eq._profiles - + if oxpoints: # Add O- and X-points # opt, xpt = critical.find_critical(eq.R, eq.Z, psi) opt = eq._profiles.opt xpt = eq._profiles.xpt - + for r, z, _ in xpt: axis.plot(r, z, "ro") for r, z, _ in opt: @@ -112,16 +115,31 @@ def plotEquilibrium(eq, axis=None, show=True, oxpoints=True, wall=True, limiter= psi_bndry = xpt[0][2] if eq._profiles.flag_limiter: # axis.contour(eq.R, eq.Z, psi, levels=[eq._profiles.psi_bndry], colors="k") - axis.contour(eq.R, eq.Z, psi, levels=[psi_bndry], colors="r", linestyles = 'dashed') - cs = plt.contour(eq.R, eq.Z, psi, levels=[eq._profiles.psi_bndry], alpha=0) + axis.contour( + eq.R, + eq.Z, + psi, + levels=[psi_bndry], + colors="r", + linestyles="dashed", + ) + cs = plt.contour( + eq.R, + eq.Z, + psi, + levels=[eq._profiles.psi_bndry], + alpha=0, + ) paths = cs.collections[0].get_paths() for path in paths: vertices = path.vertices - if np.sum(vertices[0] == vertices[-1])>1: - axis.plot(vertices[:,0], vertices[:,1], 'k') + if np.sum(vertices[0] == vertices[-1]) > 1: + axis.plot(vertices[:, 0], vertices[:, 1], "k") else: - axis.contour(eq.R, eq.Z, psi, levels=[psi_bndry], colors="r") + axis.contour( + eq.R, eq.Z, psi, levels=[psi_bndry], colors="r" + ) # Add legend axis.plot([], [], "ro", label="X-points") @@ -130,9 +148,10 @@ def plotEquilibrium(eq, axis=None, show=True, oxpoints=True, wall=True, limiter= axis.plot([], [], "go", label="O-points") except: - print('This equilibrium has not been solved: the separatrix can not be drawn.') - print('Please solve first for a plot of the critical points.') - + print( + "This equilibrium has not been solved: the separatrix can not be drawn." + ) + print("Please solve first for a plot of the critical points.") if wall and eq.tokamak.wall and len(eq.tokamak.wall.R): axis.plot( @@ -144,7 +163,8 @@ def plotEquilibrium(eq, axis=None, show=True, oxpoints=True, wall=True, limiter= axis.plot( list(eq.tokamak.limiter.R) + [eq.tokamak.limiter.R[0]], list(eq.tokamak.limiter.Z) + [eq.tokamak.limiter.Z[0]], - "k--",lw=.5 + "k--", + lw=0.5, ) if show: @@ -154,9 +174,9 @@ def plotEquilibrium(eq, axis=None, show=True, oxpoints=True, wall=True, limiter= return axis - import numpy as np + def make_broad_mask(mask, layer_size=1): """Enlarges a mask with the points just outside the input, with a width=`layer_size` @@ -171,11 +191,15 @@ def make_broad_mask(mask, layer_size=1): Mask of the points outside the limiter within a distance of `layer_size` from the limiter """ nx, ny = np.shape(mask) - layer_mask = np.zeros(np.array([nx, ny]) + 2 * np.array([layer_size, layer_size])) + layer_mask = np.zeros( + np.array([nx, ny]) + 2 * np.array([layer_size, layer_size]) + ) for i in np.arange(-layer_size, layer_size + 1) + layer_size: for j in np.arange(-layer_size, layer_size + 1) + layer_size: layer_mask[i : i + nx, j : j + ny] += mask - layer_mask = layer_mask[layer_size : layer_size + nx, layer_size : layer_size + ny] + layer_mask = layer_mask[ + layer_size : layer_size + nx, layer_size : layer_size + ny + ] layer_mask = (layer_mask > 0).astype(bool) - return layer_mask \ No newline at end of file + return layer_mask diff --git a/freegs4e/quadrature.py b/freegs4e/quadrature.py index 1030fc2..ab1c48e 100644 --- a/freegs4e/quadrature.py +++ b/freegs4e/quadrature.py @@ -65,12 +65,36 @@ def triangle_quad(triangle, n=6): d = 0.5 * (1.0 - c) return [ - ((a * r1 + b * r2 + b * r3), (a * z1 + b * z2 + b * z3), 0.109951743655322), - ((b * r1 + a * r2 + b * r3), (b * z1 + a * z2 + b * z3), 0.109951743655322), - ((b * r1 + b * r2 + a * r3), (b * z1 + b * z2 + a * z3), 0.109951743655322), - ((c * r1 + d * r2 + d * r3), (c * z1 + d * z2 + d * z3), 0.223381589678011), - ((d * r1 + c * r2 + d * r3), (d * z1 + c * z2 + d * z3), 0.223381589678011), - ((d * r1 + d * r2 + c * r3), (d * z1 + d * z2 + c * z3), 0.223381589678011), + ( + (a * r1 + b * r2 + b * r3), + (a * z1 + b * z2 + b * z3), + 0.109951743655322, + ), + ( + (b * r1 + a * r2 + b * r3), + (b * z1 + a * z2 + b * z3), + 0.109951743655322, + ), + ( + (b * r1 + b * r2 + a * r3), + (b * z1 + b * z2 + a * z3), + 0.109951743655322, + ), + ( + (c * r1 + d * r2 + d * r3), + (c * z1 + d * z2 + d * z3), + 0.223381589678011, + ), + ( + (d * r1 + c * r2 + d * r3), + (d * z1 + c * z2 + d * z3), + 0.223381589678011, + ), + ( + (d * r1 + d * r2 + c * r3), + (d * z1 + d * z2 + c * z3), + 0.223381589678011, + ), ] else: raise ValueError("Quadrature not available for n={}".format(n)) @@ -99,7 +123,9 @@ def polygon_quad(polygon, n=6): quadrature = [] # List of all points for triangle, area in zip(triangles, areas): - points = triangle_quad(triangle, n=n) # Quadrature points for this triangle + points = triangle_quad( + triangle, n=n + ) # Quadrature points for this triangle quadrature += [ (r, z, w * area / total_area) for r, z, w in points ] # Modify the weights diff --git a/freegs4e/shaped_coil.py b/freegs4e/shaped_coil.py index 33b8b40..75dc36b 100644 --- a/freegs4e/shaped_coil.py +++ b/freegs4e/shaped_coil.py @@ -23,13 +23,11 @@ along with FreeGS4E. If not, see . """ -from . import quadrature -from . import polygons -from .gradshafranov import Greens, GreensBr, GreensBz - import numpy as np +from . import polygons, quadrature from .coil import Coil +from .gradshafranov import Greens, GreensBr, GreensBz class ShapedCoil(Coil): @@ -118,8 +116,10 @@ def controlBz(self, R, Z): return result def __repr__(self): - return "ShapedCoil({0}, current={1:.1f}, turns={2}, control={3})".format( - self.shape, self.current, self.turns, self.control + return ( + "ShapedCoil({0}, current={1:.1f}, turns={2}, control={3})".format( + self.shape, self.current, self.turns, self.control + ) ) @property diff --git a/freegs4e/test_aeqdsk.py b/freegs4e/test_aeqdsk.py index e1982be..9722967 100644 --- a/freegs4e/test_aeqdsk.py +++ b/freegs4e/test_aeqdsk.py @@ -1,6 +1,7 @@ -import numpy from io import StringIO +import numpy + from . import _aeqdsk diff --git a/freegs4e/test_critical.py b/freegs4e/test_critical.py index 5e63883..d8140b6 100644 --- a/freegs4e/test_critical.py +++ b/freegs4e/test_critical.py @@ -16,7 +16,7 @@ def test_one_opoint(): # This has one O-point at (r0,z0) and no x-points def psi_func(R, Z): - return np.exp(-((R - r0) ** 2 + (Z - z0) ** 2) / 0.3 ** 2) + return np.exp(-((R - r0) ** 2 + (Z - z0) ** 2) / 0.3**2) opoints, xpoints = critical.find_critical(r2d, z2d, psi_func(r2d, z2d)) @@ -62,9 +62,9 @@ def test_doublet(): # This has two O-points, and one x-point at (r0, z0) def psi_func(R, Z): - return np.exp(-((R - r0) ** 2 + (Z - z0 - 0.3) ** 2) / 0.3 ** 2) + np.exp( - -((R - r0) ** 2 + (Z - z0 + 0.3) ** 2) / 0.3 ** 2 - ) + return np.exp( + -((R - r0) ** 2 + (Z - z0 - 0.3) ** 2) / 0.3**2 + ) + np.exp(-((R - r0) ** 2 + (Z - z0 + 0.3) ** 2) / 0.3**2) opoints, xpoints = critical.find_critical(r2d, z2d, psi_func(r2d, z2d)) @@ -88,7 +88,7 @@ def test_mask_zero_psi_bndry(): # This has one O-point at (r0,z0) and no x-points # Range from around -0.5 to +0.5 def psi_func(R, Z): - return np.exp(-((R - r0) ** 2 + (Z - z0) ** 2) / 0.3 ** 2) - 0.5 + return np.exp(-((R - r0) ** 2 + (Z - z0) ** 2) / 0.3**2) - 0.5 psi = psi_func(r2d, z2d) diff --git a/freegs4e/test_equilibrium.py b/freegs4e/test_equilibrium.py index 68d7bd4..54bc43d 100644 --- a/freegs4e/test_equilibrium.py +++ b/freegs4e/test_equilibrium.py @@ -1,17 +1,16 @@ -from . import equilibrium -from . import boundary -from . import jtor -from . import picard - import numpy as np +from . import boundary, equilibrium, jtor, picard + def test_inoutseparatrix(): - eq = equilibrium.Equilibrium(Rmin=0.1, Rmax=2.0, Zmin=-1.0, Zmax=1.0, nx=65, ny=65) + eq = equilibrium.Equilibrium( + Rmin=0.1, Rmax=2.0, Zmin=-1.0, Zmax=1.0, nx=65, ny=65 + ) # Two O-points, one X-point half way between them - psi = np.exp((-((eq.R - 1.0) ** 2) - eq.Z ** 2) * 3) + np.exp( + psi = np.exp((-((eq.R - 1.0) ** 2) - eq.Z**2) * 3) + np.exp( (-((eq.R - 1.0) ** 2) - (eq.Z + 1) ** 2) * 3 ) @@ -27,7 +26,9 @@ def test_fixed_boundary_psi(): # This is adapted from example 5 profiles = jtor.ConstrainPaxisIp( - 1e3, 1e5, 1.0 # Plasma pressure on axis [Pascals] # Plasma current [Amps] + 1e3, + 1e5, + 1.0, # Plasma pressure on axis [Pascals] # Plasma current [Amps] ) # fvac = R*Bt eq = equilibrium.Equilibrium( @@ -51,7 +52,9 @@ def test_fixed_boundary_psi(): def test_setSolverVcycle(): - eq = equilibrium.Equilibrium(Rmin=0.1, Rmax=2.0, Zmin=-1.0, Zmax=1.0, nx=65, ny=65) + eq = equilibrium.Equilibrium( + Rmin=0.1, Rmax=2.0, Zmin=-1.0, Zmax=1.0, nx=65, ny=65 + ) oldsolver = eq._solver eq.setSolverVcycle(nlevels=2, ncycle=1, niter=5) diff --git a/freegs4e/test_geqdsk.py b/freegs4e/test_geqdsk.py index 9010e8d..df22613 100644 --- a/freegs4e/test_geqdsk.py +++ b/freegs4e/test_geqdsk.py @@ -1,7 +1,7 @@ -import numpy - from io import StringIO +import numpy + from . import _geqdsk diff --git a/freegs4e/test_jtor.py b/freegs4e/test_jtor.py index a54d40d..f7e5660 100644 --- a/freegs4e/test_jtor.py +++ b/freegs4e/test_jtor.py @@ -15,7 +15,7 @@ def test_psinorm_range(): R, Z = np.meshgrid( np.linspace(0.5, 1.5, 33), np.linspace(-1, 1, 33), indexing="ij" ) - psi = np.exp((-((R - 1.0) ** 2) - Z ** 2) * 3) + np.exp( + psi = np.exp((-((R - 1.0) ** 2) - Z**2) * 3) + np.exp( (-((R - 1.0) ** 2) - (Z + 1) ** 2) * 3 ) diff --git a/freegs4e/test_linearsolve.py b/freegs4e/test_linearsolve.py index 8b1cd37..2720ec7 100644 --- a/freegs4e/test_linearsolve.py +++ b/freegs4e/test_linearsolve.py @@ -2,10 +2,10 @@ Tests of the linear solver """ -from . import multigrid - import numpy as np +from . import multigrid + # Laplacian in 2D @@ -21,7 +21,7 @@ def test_direct_laplacian(): xx, yy = np.meshgrid(np.linspace(0, Lx, nx), np.linspace(0, Ly, ny)) - solution = np.exp(-((xx - 0.5) ** 2 + (yy - 0.5) ** 2) / 0.4 ** 2) + solution = np.exp(-((xx - 0.5) ** 2 + (yy - 0.5) ** 2) / 0.4**2) A = multigrid.LaplacianOp() rhs = A(solution, dx, dy) @@ -50,7 +50,7 @@ def test_multigrid_laplacian(): xx, yy = np.meshgrid(np.linspace(0, Lx, nx), np.linspace(0, Ly, ny)) - solution = np.exp(-((xx - 0.5) ** 2 + (yy - 0.5) ** 2) / 0.4 ** 2) + solution = np.exp(-((xx - 0.5) ** 2 + (yy - 0.5) ** 2) / 0.4**2) A = multigrid.LaplacianOp() rhs = A(solution, dx, dy) diff --git a/freegs4e/test_machine.py b/freegs4e/test_machine.py index 8c5e222..fc33bb6 100644 --- a/freegs4e/test_machine.py +++ b/freegs4e/test_machine.py @@ -2,10 +2,12 @@ # Test calculation of magnetic field due to coils # and forces between coils -from . import machine -import numpy as np import math +import numpy as np + +from . import machine + mu0 = 4e-7 * np.pi @@ -19,7 +21,7 @@ def test_coil_axis(): coil = machine.Coil(Rcoil, 1.0, current=current) def analytic_Bz(dZ): - return (mu0 / 2) * Rcoil ** 2 * current / (dZ ** 2 + Rcoil ** 2) ** 1.5 + return (mu0 / 2) * Rcoil**2 * current / (dZ**2 + Rcoil**2) ** 1.5 # Note: Can't evaluate at R=0, assert math.isclose(coil.Br(0.0001, 2.0), 0.0, abs_tol=1e-8) diff --git a/freegs4e/test_multi_coil.py b/freegs4e/test_multi_coil.py index 8479783..6ebb24d 100644 --- a/freegs4e/test_multi_coil.py +++ b/freegs4e/test_multi_coil.py @@ -1,8 +1,8 @@ -from .multi_coil import MultiCoil +import numpy as np + from .coil import Coil from .machine import Circuit - -import numpy as np +from .multi_coil import MultiCoil def test_single(): @@ -66,7 +66,9 @@ def test_move_R(): dR = 0.6 coil1 = MultiCoil([1.1, 0.2], [1.2, -0.3], current=100.0, mirror=False) - coil2 = MultiCoil([1.1 + dR, 0.2 + dR], [1.2, -0.3], current=100.0, mirror=False) + coil2 = MultiCoil( + [1.1 + dR, 0.2 + dR], [1.2, -0.3], current=100.0, mirror=False + ) # Shift coil1 to same location as coil2 coil1.R += dR @@ -85,7 +87,9 @@ def test_move_Z(): dZ = 0.4 coil1 = MultiCoil([1.1, 0.2], [1.2, -0.3], current=100.0, mirror=False) - coil2 = MultiCoil([1.1, 0.2], [1.2 + dZ, -0.3 + dZ], current=100.0, mirror=False) + coil2 = MultiCoil( + [1.1, 0.2], [1.2 + dZ, -0.3 + dZ], current=100.0, mirror=False + ) # Shift coil1 to same location as coil2 coil1.Z += dZ diff --git a/freegs4e/test_optimise.py b/freegs4e/test_optimise.py index 90138ac..340d743 100644 --- a/freegs4e/test_optimise.py +++ b/freegs4e/test_optimise.py @@ -1,6 +1,7 @@ -from . import optimise import numpy as np +from . import optimise + class DummyCoil: def __init__(self): diff --git a/freegs4e/test_optimiser.py b/freegs4e/test_optimiser.py index 6dbe535..22cee42 100644 --- a/freegs4e/test_optimiser.py +++ b/freegs4e/test_optimiser.py @@ -41,7 +41,10 @@ def test_quadratic(): # The optimiser can control two coefficients, # and should use calculate_score to evaluate solutions result = optimiser.optimise( - start_values, [ControlIndex(0), ControlIndex(1)], calculate_score, maxgen=300 + start_values, + [ControlIndex(0), ControlIndex(1)], + calculate_score, + maxgen=300, ) # Answer should be close to (1,2) @@ -141,7 +144,9 @@ def monitor(generation, best, pop): best_point = axis.scatter([best[1][0]], [best[1][1]], c="red") axis.figure.canvas.draw() - axis.set_title("Generation: {}, Best score: {}".format(generation, best[0])) + axis.set_title( + "Generation: {}, Best score: {}".format(generation, best[0]) + ) plt.pause(0.5) result = optimiser.optimise( diff --git a/freegs4e/test_polygons.py b/freegs4e/test_polygons.py index 41e0f79..148f4be 100644 --- a/freegs4e/test_polygons.py +++ b/freegs4e/test_polygons.py @@ -1,6 +1,7 @@ -from . import polygons import numpy as np +from . import polygons + def test_nointersect(): assert not polygons.intersect([0, 1], [0, 0], [0, 1], [1, 1]) @@ -15,7 +16,10 @@ def test_lineintersect(): # Two squares def test_squareintersect(): assert polygons.intersect( - [0, 1, 1, 0], [0, 0, 1, 1], [-0.5, 0.5, 0.5, -0.5], [0.5, 0.5, 1.5, 1.5] + [0, 1, 1, 0], + [0, 0, 1, 1], + [-0.5, 0.5, 0.5, -0.5], + [0.5, 0.5, 1.5, 1.5], ) @@ -37,8 +41,12 @@ def test_clockwise(): def test_triangulate_three(): # Always returns a triangle in clockwise order - assert polygons.triangulate([(0, 0), (0, 1), (1, 0)]) == [[(0, 0), (0, 1), (1, 0)]] - assert polygons.clockwise(polygons.triangulate([(0, 0), (1, 0), (0, 1)])[0]) + assert polygons.triangulate([(0, 0), (0, 1), (1, 0)]) == [ + [(0, 0), (0, 1), (1, 0)] + ] + assert polygons.clockwise( + polygons.triangulate([(0, 0), (1, 0), (0, 1)])[0] + ) def test_triangulate_four(): @@ -51,5 +59,6 @@ def test_triangulate_four(): assert polygons.area(result[1]) > 0.0 # Sum of areas equal to original area assert np.isclose( - polygons.area(square), polygons.area(result[0]) + polygons.area(result[1]) + polygons.area(square), + polygons.area(result[0]) + polygons.area(result[1]), ) diff --git a/freegs4e/test_quadrature.py b/freegs4e/test_quadrature.py index 50a8bfd..22ba130 100644 --- a/freegs4e/test_quadrature.py +++ b/freegs4e/test_quadrature.py @@ -18,24 +18,30 @@ def test_integral(): y0, y1 = -0.2, 3.1 def func(x, y): - return x ** 2 - y ** 3 + x * y + return x**2 - y**3 + x * y exact_integral = ( - (x1 ** 3 - x0 ** 3) * (y1 - y0) / 3 - - (y1 ** 4 - y0 ** 4) * (x1 - x0) / 4 - + (x1 ** 2 - x0 ** 2) * (y1 ** 2 - y0 ** 2) / 4 + (x1**3 - x0**3) * (y1 - y0) / 3 + - (y1**4 - y0**4) * (x1 - x0) / 4 + + (x1**2 - x0**2) * (y1**2 - y0**2) / 4 ) # A 1st order method can't integrate this polynomial exactly - quad1 = quadrature.polygon_quad([(x0, y0), (x0, y1), (x1, y1), (x1, y0)], n=1) + quad1 = quadrature.polygon_quad( + [(x0, y0), (x0, y1), (x1, y1), (x1, y0)], n=1 + ) assert not np.isclose( - quadrature.average(func, quad1), exact_integral / ((x1 - x0) * (y1 - y0)) + quadrature.average(func, quad1), + exact_integral / ((x1 - x0) * (y1 - y0)), ) for n in [3, 6]: # Higher order methods can - quad1 = quadrature.polygon_quad([(x0, y0), (x0, y1), (x1, y1), (x1, y0)], n=n) + quad1 = quadrature.polygon_quad( + [(x0, y0), (x0, y1), (x1, y1), (x1, y0)], n=n + ) assert len(quad1) == 2 * n assert np.isclose( - quadrature.average(func, quad1), exact_integral / ((x1 - x0) * (y1 - y0)) + quadrature.average(func, quad1), + exact_integral / ((x1 - x0) * (y1 - y0)), ) diff --git a/freegs4e/test_readwrite.py b/freegs4e/test_readwrite.py index a8277a3..6399707 100644 --- a/freegs4e/test_readwrite.py +++ b/freegs4e/test_readwrite.py @@ -1,13 +1,17 @@ -import freegs4e - import io + from numpy import allclose +import freegs4e + def test_readwrite(): """Test reading/writing to a file round-trip""" - for tokamak in [freegs4e.machine.TestTokamak(), freegs4e.machine.MAST_sym()]: + for tokamak in [ + freegs4e.machine.TestTokamak(), + freegs4e.machine.MAST_sym(), + ]: eq = freegs4e.Equilibrium( tokamak=tokamak, @@ -26,9 +30,13 @@ def test_readwrite(): # X-points are on different flux surfaces. xpoints = [(1.1, -0.6), (1.1, 0.8)] isoflux = [(1.1, -0.6, 1.1, 0.6)] - constrain = freegs4e.control.constrain(xpoints=xpoints, isoflux=isoflux) + constrain = freegs4e.control.constrain( + xpoints=xpoints, isoflux=isoflux + ) - freegs4e.solve(eq, profiles, constrain, maxits=25, atol=1e-3, rtol=1e-1) + freegs4e.solve( + eq, profiles, constrain, maxits=25, atol=1e-3, rtol=1e-1 + ) memory_file = io.BytesIO() diff --git a/freegs4e/test_shaped_coil.py b/freegs4e/test_shaped_coil.py index fda8aaf..f1a6fe8 100644 --- a/freegs4e/test_shaped_coil.py +++ b/freegs4e/test_shaped_coil.py @@ -1,7 +1,7 @@ -from .shaped_coil import ShapedCoil - import numpy as np +from .shaped_coil import ShapedCoil + def test_area(): coil1 = ShapedCoil( @@ -20,7 +20,12 @@ def test_move_R(): [(0.95, 0.1), (0.95, 0.2), (1.05, 0.2), (1.05, 0.1)], current=100.0 ) coil2 = ShapedCoil( - [(0.95 + dR, 0.1), (0.95 + dR, 0.2), (1.05 + dR, 0.2), (1.05 + dR, 0.1)], + [ + (0.95 + dR, 0.1), + (0.95 + dR, 0.2), + (1.05 + dR, 0.2), + (1.05 + dR, 0.1), + ], current=100.0, ) @@ -44,7 +49,12 @@ def test_move_Z(): [(0.95, 0.1), (0.95, 0.2), (1.05, 0.2), (1.05, 0.1)], current=100.0 ) coil2 = ShapedCoil( - [(0.95, 0.1 + dZ), (0.95, 0.2 + dZ), (1.05, 0.2 + dZ), (1.05, 0.1 + dZ)], + [ + (0.95, 0.1 + dZ), + (0.95, 0.2 + dZ), + (1.05, 0.2 + dZ), + (1.05, 0.1 + dZ), + ], current=100.0, ) diff --git a/test-01-compare.py b/test-01-compare.py index 7cf4708..5ae1222 100644 --- a/test-01-compare.py +++ b/test-01-compare.py @@ -3,41 +3,57 @@ # and compares the result -import freegs4e - -from numpy import linspace, amin, amax, meshgrid, exp - import matplotlib.pyplot as plt +from numpy import amax, amin, exp, linspace, meshgrid + +import freegs4e tokamak = freegs4e.machine.TestTokamak() -eq1 = freegs4e.Equilibrium(tokamak=tokamak, - Rmin=0.1, Rmax=2.0, # Radial domain - Zmin = -2.0, Zmax = 2.0, - nx=65, ny=129) # Number of grid points +eq1 = freegs4e.Equilibrium( + tokamak=tokamak, + Rmin=0.1, + Rmax=2.0, # Radial domain + Zmin=-2.0, + Zmax=2.0, + nx=65, + ny=129, +) # Number of grid points tokamak = freegs4e.machine.TestTokamak() -eq2 = freegs4e.Equilibrium(tokamak=tokamak, - Rmin=0.1, Rmax=2.0, # Radial domain - Zmin=-1.0, Zmax=1.0, # Height range - nx=65, ny=65) # Number of grid points +eq2 = freegs4e.Equilibrium( + tokamak=tokamak, + Rmin=0.1, + Rmax=2.0, # Radial domain + Zmin=-1.0, + Zmax=1.0, # Height range + nx=65, + ny=65, +) # Number of grid points # Set the initial psi -xx, yy = meshgrid(linspace(eq1.Rmin,eq1.Rmax,65), linspace(eq1.Zmin,eq1.Zmax,129), indexing='ij') -psi = exp( - ( (xx - 1.0)**2 + (yy)**2 ) / 0.3**2 ) +xx, yy = meshgrid( + linspace(eq1.Rmin, eq1.Rmax, 65), + linspace(eq1.Zmin, eq1.Zmax, 129), + indexing="ij", +) +psi = exp(-((xx - 1.0) ** 2 + (yy) ** 2) / 0.3**2) eq1._updatePlasmaPsi(psi) -xx, yy = meshgrid(linspace(eq2.Rmin,eq2.Rmax,65), linspace(eq2.Zmin,eq2.Zmax,65), indexing='ij') -psi = exp( - ( (xx - 1.0)**2 + (yy)**2 ) / 0.3**2 ) +xx, yy = meshgrid( + linspace(eq2.Rmin, eq2.Rmax, 65), + linspace(eq2.Zmin, eq2.Zmax, 65), + indexing="ij", +) +psi = exp(-((xx - 1.0) ** 2 + (yy) ** 2) / 0.3**2) eq2._updatePlasmaPsi(psi) -xpoints = [(1.1, -0.6), # (R,Z) locations of X-points - (1.1, 0.8)] +xpoints = [(1.1, -0.6), (1.1, 0.8)] # (R,Z) locations of X-points -isoflux = [(1.1,-0.6, 1.1,0.6)] +isoflux = [(1.1, -0.6, 1.1, 0.6)] constrain = freegs4e.control.constrain(xpoints=xpoints, isoflux=isoflux) @@ -56,24 +72,20 @@ eq2.solve(profiles) # Nonlinear solve -#freegs4e.solve(eq1, profiles, constrain) -#freegs4e.solve(eq2, profiles, constrain) +# freegs4e.solve(eq1, profiles, constrain) +# freegs4e.solve(eq2, profiles, constrain) psi1 = eq1.psi() psi2 = eq2.psi() -#psi1 = eq1.tokamak.psi(eq1.R, eq1.Z) -#psi2 = eq2.tokamak.psi(eq2.R, eq2.Z) +# psi1 = eq1.tokamak.psi(eq1.R, eq1.Z) +# psi2 = eq2.tokamak.psi(eq2.R, eq2.Z) -#psi1 = eq1.plasma_psi -#psi2 = eq2.plasma_psi +# psi1 = eq1.plasma_psi +# psi2 = eq2.plasma_psi levels = linspace(amin(psi1), amax(psi1), 40) -plt.contour(eq1.R, eq1.Z, psi1, levels=levels, colors='k') -plt.contour(eq2.R, eq2.Z, psi2, levels=levels, colors='r') +plt.contour(eq1.R, eq1.Z, psi1, levels=levels, colors="k") +plt.contour(eq2.R, eq2.Z, psi2, levels=levels, colors="r") plt.show() - - - - diff --git a/test-convergence.py b/test-convergence.py index a3478a1..5cc0438 100644 --- a/test-convergence.py +++ b/test-convergence.py @@ -1,8 +1,9 @@ #!/usr/bin/env python -import freegs4e -import numpy as np import matplotlib.pyplot as plt +import numpy as np + +import freegs4e start_resolution = 17 nrefinements = 5 # Number of refinements. Minimum 2 @@ -129,9 +130,13 @@ def plot_convergence(axis, title, values=None, diffs=None): axis.set_xticks([], minor=True) -plot_convergence(axes[0, 0], r"Change in $\psi$ at {}".format(location), values=psivals) +plot_convergence( + axes[0, 0], r"Change in $\psi$ at {}".format(location), values=psivals +) plot_convergence(axes[0, 1], r"$\psi$ difference $l_2$ norm", diffs=l2vals) -plot_convergence(axes[0, 2], r"$\psi$ difference $l_\infty$ norm", diffs=linfvals) +plot_convergence( + axes[0, 2], r"$\psi$ difference $l_\infty$ norm", diffs=linfvals +) plot_convergence(axes[1, 0], "Br at {}".format(location), values=brvals) plot_convergence(axes[1, 1], "Plasma volume", values=volumevals) plot_convergence(axes[1, 2], "P1L coil current", values=coilcurrents) diff --git a/test_fileutils.py b/test_fileutils.py index b59a50d..6d6916b 100644 --- a/test_fileutils.py +++ b/test_fileutils.py @@ -1,24 +1,44 @@ -from freegs4e._fileutils import next_value from numpy import allclose +from freegs4e._fileutils import next_value + def test_next_value(): data = [ - '+9.500000000e-01+2.000000000e+00+1.000000000e+00+2.500000000e+01+0.000000000e+00\n', - ' 9.500000000E-01 2.000000000E+00 1.000000000E+00 2.500000000E+01 0.000000000E+00\n', - '+3.563359524e+02+2.337058846e-02-1.212203732e-02+1.953790839e-03-7.116987284e-03', - ' 3.563359524E+02 2.337058846e-02-1.212203732E-02 1.953790839E-03-7.116987284e-03', - '0 0 0\n', - ' 0 0\n', + "+9.500000000e-01+2.000000000e+00+1.000000000e+00+2.500000000e+01+0.000000000e+00\n", + " 9.500000000E-01 2.000000000E+00 1.000000000E+00 2.500000000E+01 0.000000000E+00\n", + "+3.563359524e+02+2.337058846e-02-1.212203732e-02+1.953790839e-03-7.116987284e-03", + " 3.563359524E+02 2.337058846e-02-1.212203732E-02 1.953790839E-03-7.116987284e-03", + "0 0 0\n", + " 0 0\n", ] expected = [ - 0.95, 2.0, 1.0, 25.0, 0.0, - 0.95, 2.0, 1.0, 25.0, 0.0, - 356.3359524, 0.02337058846, -0.01212203732, 0.001953790839, -0.007116987284, - 356.3359524, 0.02337058846, -0.01212203732, 0.001953790839, -0.007116987284, - 0, 0, 0, - 0, 0 + 0.95, + 2.0, + 1.0, + 25.0, + 0.0, + 0.95, + 2.0, + 1.0, + 25.0, + 0.0, + 356.3359524, + 0.02337058846, + -0.01212203732, + 0.001953790839, + -0.007116987284, + 356.3359524, + 0.02337058846, + -0.01212203732, + 0.001953790839, + -0.007116987284, + 0, + 0, + 0, + 0, + 0, ] values = next_value(data) From bee8e69f94e0b4f6f19d72416506821921b680f3 Mon Sep 17 00:00:00 2001 From: George Holt Date: Mon, 12 Aug 2024 14:37:49 -0700 Subject: [PATCH 2/2] chore: Remove flake8 hooks and config --- .flake8 | 5 ----- .pre-commit-config.yaml | 5 ----- 2 files changed, 10 deletions(-) delete mode 100644 .flake8 diff --git a/.flake8 b/.flake8 deleted file mode 100644 index fdfa5d0..0000000 --- a/.flake8 +++ /dev/null @@ -1,5 +0,0 @@ -exclude = - .git - __pycache__ - build - dist \ No newline at end of file diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 02867fc..2fe8c19 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,9 +1,4 @@ repos: - - repo: https://github.com/PyCQA/flake8 - rev: 7.1.1 - hooks: - - id: flake8 - - repo: https://github.com/PyCQA/isort rev: 5.13.2 hooks: