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:
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)