Skip to content

Commit

Permalink
Fix lots of sphinx warnings
Browse files Browse the repository at this point in the history 
Mostly fixed by trying to consistently use numpydoc style
  • Loading branch information
@ZedThree
ZedThree committed Jun 28, 2024
1 parent 99e3e02 commit d8fedcb
Show file tree
Hide file tree
Showing 11 changed files with 354 additions and 264 deletions.
39 changes: 21 additions & 18 deletions freegs/critical.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,22 +49,22 @@ def find_critical(R, Z, psi, discard_xpoints=True):
Inputs
------
R - R(nr, nz) 2D array of major radii
Z - Z(nr, nz) 2D array of heights
psi - psi(nr, nz) 2D array of psi values
R:
R(nr, nz) 2D array of major radii
Z:
Z(nr, nz) 2D array of heights
psi:
psi(nr, nz) 2D array of psi values
Returns
-------
opoint, xpoint:
Two lists of critical points.
Two lists of critical points
opoint, xpoint
Each of these is a list of tuples with (R, Z, psi) points
Each of these is a list of tuples with (R, Z, psi) points
The first tuple is the primary O-point (magnetic axis)
and primary X-point (separatrix)
The first tuple is the primary O-point (magnetic axis)
and primary X-point (separatrix)
"""

Expand Down Expand Up @@ -241,19 +241,22 @@ def core_mask(R, Z, psi, opoint, xpoint=[], psi_bndry=None):
Inputs
------
R[nx,ny] - 2D array of major radius (R) values
Z[nx,ny] - 2D array of height (Z) values
psi[nx,ny] - 2D array of poloidal flux
opoint, xpoint - Values returned by find_critical
R[nx,ny]:
2D array of major radius (R) values
Z[nx,ny]:
2D array of height (Z) values
psi[nx,ny]:
2D array of poloidal flux
opoint, xpoint :
Values returned by find_critical
If psi_bndry is not None, then that is used to find the
separatrix, not the X-points.
Returns
-------
A 2D array [nx,ny] which is 1 inside the core, 0 outside
array:
A 2D array [nx,ny] which is 1 inside the core, 0 outside
"""

Expand Down
2 changes: 1 addition & 1 deletion freegs/dump.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -110,7 +110,7 @@ def __exit__(self, type, value, traceback):

def write_equilibrium(self, equilibrium):
"""
Write `equilbrium` to file
Write ``equilbrium`` to file
"""

self.handle["coil_dtype"] = Coil.dtype
Expand Down
167 changes: 91 additions & 76 deletions freegs/equilibrium.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,8 @@
state, including plasma and coil currents
"""

from typing import Optional

from numpy import pi, meshgrid, linspace, exp, array
import numpy as np
from scipy import interpolate
Expand All @@ -29,62 +31,68 @@ class Equilibrium:
Represents the equilibrium state, including
plasma and coil currents
Data members
------------
These can be read, but should not be modified directly
R[nx,ny]
Z[nx,ny]
Rmin, Rmax
Zmin, Zmax
tokamak - The coils and circuits
Private data members
_applyBoundary()
_solver - Grad-Shafranov elliptic solver
_profiles An object which calculates the toroidal current
_constraints Control system which adjusts coil currents to meet constraints
e.g. X-point location and flux values
Attributes
----------
R[nx,ny]:
Read-only
Z[nx,ny]:
Read-only
Rmin, Rmax:
Read-only
Zmin, Zmax:
Read-only
tokamak:
The coils and circuits
_applyBoundary:
_solver:
Grad-Shafranov elliptic solver
_profiles:
An object which calculates the toroidal current
_constraints:
Control system which adjusts coil currents to meet constraints
e.g. X-point location and flux values
"""

def __init__(
self,
tokamak=machine.EmptyTokamak(),
Rmin=0.1,
Rmax=2.0,
Zmin=-1.0,
Zmax=1.0,
nx=65,
ny=65,
tokamak: machine.Machine = machine.EmptyTokamak(),
Rmin: float = 0.1,
Rmax: float = 2.0,
Zmin: float = -1.0,
Zmax: float = 1.0,
nx: int = 65,
ny: int = 65,
boundary=freeBoundary,
psi=None,
current=0.0,
order=4,
check_limited=False,
psi: Optional[float] = None,
current: float = 0.0,
order: int = 4,
check_limited: bool = False,
):
"""Initialises a plasma equilibrium
Rmin, Rmax - Range of major radius R [m]
Zmin, Zmax - Range of height Z [m]
nx - Resolution in R. This must be 2^n + 1
ny - Resolution in Z. This must be 2^m + 1
boundary - The boundary condition, either freeBoundary or fixedBoundary
psi - Magnetic flux. If None, use concentric circular flux
surfaces as starting guess
current - Plasma current (default = 0.0)
order - The order of the differential operators to use.
Valid values are 2 or 4.
check_limited - Boolean, checks if the plasma is limited.
Parameters
----------
tokamak:
The set of coils and currents
Rmin, Rmax:
Range of major radius R [m]
Zmin, Zmax:
Range of height Z [m]
nx:
Resolution in R. This must be ``2^n + 1``
ny:
Resolution in Z. This must be ``2^m + 1``
boundary:
The boundary condition, either `freeBoundary` or `fixedBoundary`
psi:
Magnetic flux. If None, use concentric circular flux
surfaces as starting guess
current:
Plasma current
order:
The order of the differential operators to use. Valid values are 2 or 4
check_limited:
Checks if the plasma is limited
"""

self.tokamak = tokamak
Expand Down Expand Up @@ -113,7 +121,7 @@ def __init__(
if psi is None:
# Starting guess for psi
xx, yy = meshgrid(linspace(0, 1, nx), linspace(0, 1, ny), indexing="ij")
psi = exp(-((xx - 0.5) ** 2 + (yy - 0.5) ** 2) / 0.4 ** 2)
psi = exp(-((xx - 0.5) ** 2 + (yy - 0.5) ** 2) / 0.4**2)

psi[0, :] = 0.0
psi[:, 0] = 0.0
Expand Down Expand Up @@ -183,8 +191,11 @@ def callSolver(self, psi, rhs):
"""
Calls the psi solver, passing the initial guess and RHS arrays
psi Initial guess for the solution (used if iterative)
rhs
Parameters
----------
psi:
Initial guess for the solution (used if iterative)
rhs:
Returns
-------
Expand Down Expand Up @@ -352,13 +363,13 @@ def q(self, psinorm=None, npsi=100):
psinorm_inner = np.linspace(0.01, 0.99, num=npsi)
q_inner = critical.find_safety(self, psinorm=psinorm_inner)

interp = interpolate.interp1d(psinorm_inner, q_inner,
kind = "quadratic",
fill_value = "extrapolate")
interp = interpolate.interp1d(
psinorm_inner, q_inner, kind="quadratic", fill_value="extrapolate"
)
result = interp(psinorm)
else:
result = critical.find_safety(self, psinorm=psinorm)

# Convert to a scalar if only one result
if np.size(result) == 1:
return float(result)
Expand All @@ -385,7 +396,8 @@ def rhotor(self, psi=None):
"""
Calculates normalised toroidal flux at specified values of
poloidal flux.
>>> rhotor = sqrt ( tor_flux/max(tor_flux)).
>>> rhotor = sqrt ( tor_flux/max(tor_flux)).
Maximum toroidal flux shoud be at LCFS.
"""
Expand Down Expand Up @@ -437,20 +449,24 @@ def solve(self, profiles, Jtor=None, psi=None, psi_bndry=None):
This performs the linear Grad-Shafranov solve
profiles - An object describing the plasma profiles.
At minimum this must have methods:
.Jtor(R, Z, psi, psi_bndry) -> [nx, ny]
.pprime(psinorm)
.ffprime(psinorm)
.pressure(psinorm)
.fpol(psinorm)
Parameters
----------
profiles:
An object describing the plasma profiles.
At minimum this must have methods:
- `.Jtor(R, Z, psi, psi_bndry) -> [nx, ny]`
- `.pprime(psinorm)`
- `.ffprime(psinorm)`
- `.pressure(psinorm)`
- `.fpol(psinorm)`
Jtor : 2D array
If supplied, specifies the toroidal current at each (R,Z) point
If not supplied, Jtor is calculated from profiles by finding O,X-points
psi_bndry - Poloidal flux to use as the separatrix (plasma boundary)
If not given then X-point locations are used.
psi_bndry:
Poloidal flux to use as the separatrix (plasma
boundary). If not given then X-point locations are used.
"""

self._profiles = profiles
Expand Down Expand Up @@ -630,18 +646,18 @@ def _updatePlasmaPsi(self, plasma_psi):
# Update the plasma axis and boundary flux as well as mask
self._updateBoundaryPsi()

def plot(self, axis=None, show=True, oxpoints=True):
def plot(self, axis=None, show=True, oxpoints=True) -> plt.Axes:
"""
Plot the equilibrium flux surfaces
axis - Specify the axis on which to plot
show - Call matplotlib.pyplot.show() before returning
oxpoints - Plot X points as red circles, O points as green circles
Returns
-------
axis object from Matplotlib
Parameters
----------
axis :
Specify the axis on which to plot
show :
Call matplotlib.pyplot.show() before returning
oxpoints :
Plot X points as red circles, O points as green circles
"""
from .plotting import plotEquilibrium
Expand Down Expand Up @@ -1374,7 +1390,6 @@ def newDomain(eq, Rmin=None, Rmax=None, Zmin=None, Zmax=None, nx=None, ny=None):


if __name__ == "__main__":

# Test the different spline interpolation routines

from numpy import ravel
Expand Down
33 changes: 18 additions & 15 deletions freegs/fieldtracer.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -166,26 +166,29 @@ def traceFieldLines(eq, solwidth=0.03, nlines=10, nturns=50, npoints=200, axis=N
Inputs
------
eq Equilibrium object
solwidth The width of the SOL in meters
nlines Number of field lines to follow
nturns Number of times around the tokamak to follow
npoints Maximum number of points per line. May hit a wall
axis Matplotlib figure Axis. If given, field lines
are plotted on the axis
eq:
Equilibrium object
solwidth:
The width of the SOL in meters
nlines:
Number of field lines to follow
nturns:
Number of times around the tokamak to follow
npoints:
Maximum number of points per line. May hit a wall
axis:
Matplotlib figure Axis. If given, field lines are plotted on the axis
Returns
-------
[LineCoordinates, LineCoordinates]:
The forward and backward field line coordinates
stored in LineCoordinates objects
The forward and backward field line coordinates
stored in LineCoordinates objects
>>> forward, backward = traceFieldLines(eq)
>>> forward, backward = traceFieldLines(eq)
forward and backward have data members
R, Z, length 2D arrays of shape (npoints, nlines)
forward and backward have data members
R, Z, length 2D arrays of shape (npoints, nlines)
"""
ft = FieldTracer(eq)
Expand Down
Loading

0 comments on commit d8fedcb

Please sign in to comment.