simsopt.mhd package

Submodules

simsopt.mhd.bootstrap module

simsopt.mhd.boozer module

simsopt.mhd.profiles module

This module provides classes to handle radial profiles of density, temperature, pressure, and other quantities that are flux functions.

class simsopt.mhd.profiles.Profile(*args, **kwargs)

Bases: Optimizable

Base class for radial profiles. This class should not be used directly - use subclasses instead.

__call__(*args, **kwargs): Shortcut for calling f(s)

plot(ax=None, show=True, n=100)

Plot the profile using matplotlib.

Parameters:

ax – The axis object on which to plot. If None, a new figure will be created.
show – Whether to call matplotlib’s show() function.
n – The number of grid points in s to show.

class simsopt.mhd.profiles.ProfilePolynomial(data)

Bases: Profile

A profile described by a polynomial in the normalized toroidal flux s. The polynomial coefficients are dofs that are fixed by default.

Parameters:: data – 1D numpy array of the polynomial coefficients. The first coefficient is the constant term, the next coefficient is the linear term, etc.

dfds(s): Return the d/ds derivative of the profile at specified points in s.

f(s): Return the value of the profile at specified points in s.

class simsopt.mhd.profiles.ProfilePressure(*args)

Bases: Profile

A Profile \(f(s)\) which is determined by other profiles \(f_j(s)\) as follows:

\[f(s) = \sum_j f_{2j}(s) f_{2j+1}(s).\]

This is useful for creating a pressure profile in terms of density and temperature profiles, with any number of species. Typical usage is as follows:

ne = ProfilePolynomial(1.0e20 * np.array([1.0, 0.0, 0.0, 0.0, -1.0]))
Te = ProfilePolynomial(8.0e3 * np.array([1.0, -1.0]))
nH = ne
TH = ProfilePolynomial(7.0e3 * np.array([1.0, -1.0]))
pressure = ProfilePressure(ne, Te, nH, TH)

This class does not have any optimizable dofs.

Parameters:: args – An even number of Profile objects.

dfds(s): Return the d/ds derivative of the profile at specified points in s.

f(s): Return the value of the profile at specified points in s.

class simsopt.mhd.profiles.ProfileScaled(base, scalefac)

Bases: Profile

A Profile which is equivalent to another Profile object but scaled by a constant. This constant is an optimizable dof, which is fixed by default.

Parameters:

base – A Profile object to scale
scalefac – A number by which the base profile will be scaled.

dfds(s): Return the d/ds derivative of the profile at specified points in s.

f(s): Return the value of the profile at specified points in s.

class simsopt.mhd.profiles.ProfileSpec(data, cumulative: bool = False, psi_edge: float | None = None)

Bases: Profile

A profile described by an array of size Nvol

Parameters:

data – 1D numpy array containing the profile value in each volume
cumulative – Set to True if the profile is cumulative, i.e. if the value in volume lvol is the integrated quantity from the axis to volume lvol. Only the toroidal flux, poloidal flux and the volume currents are cumulative quantities in SPEC input file. False by default.

dfds(lvol)

Returns the derivative of the profile w.r.t s accross interface. The derivative is returned at the interface lvol, with the innermost interface being lvol=1. (Volume lvol is bounded by interface lvol and lvol+1, with innermost volume being lvol=0)

Here \(s\) is defined as \(s = \psi_t/\psi_{edge}\). Thus,

\[dp/ds = \sum_l [[p]]_l \psi_{edge} \delta(\psi_t-\psi_{t,l})\]

with p the profile, and the sum is on the interfaces.

Parameters:: lvol – int, list or np.array of int, between 1 and Mvol-1.

f(lvol: int)

Return the value of the profile in volume lvol

Parameters:: lvol – int, list or np.array of int, between 0 and Mvol

class simsopt.mhd.profiles.ProfileSpline(s, f, degree=3)

Bases: Profile

A Profile that uses spline interpolation via scipy.interpolate.InterpolatedUnivariateSpline

The f data are optimizable dofs, which are fixed by default.

Parameters:

s – A 1d array with the x coordinates for the spline.
f – A 1d array with the y coordinates for the spline.
degree – The polynomial degree of the spline. Must be in [1, 2, 3, 4, 5].

dfds(s): Return the d/ds derivative of the profile at specified points in s.

f(s): Return the value of the profile at specified points in s.

resample(new_s, degree=None)

Return a new ProfileSpline object that has different grid points (spline nodes). The data from the old s grid will be interpolated onto the new s grid.

Parameters:

new_s – A 1d array representing the x coordinates of the new ProfileSpline.
degree – The polynomial degree used for the new ProfileSpline object. If None, the degree of the original ProfileSpline will be used.

Returns:

A new ProfileSpline object, in which the data have been resampled onto new_s.

simsopt.mhd.spec module

simsopt.mhd.virtual_casing module

This module provides routines for interacting with the virtual_casing package by D Malhotra et al, for e.g. computing the magnetic field on a surface due to currents outside the surface.

For details of the algorithm, see D Malhotra, A J Cerfon, M O’Neil, and E Toler, “Efficient high-order singular quadrature schemes in magnetic fusion”, Plasma Physics and Controlled Fusion 62, 024004 (2020).

class simsopt.mhd.virtual_casing.VirtualCasing

Bases: object

Use the virtual casing principle to compute the contribution to the total magnetic field due to current outside a bounded surface.

Usually, an instance of this class is created using the from_vmec() class method, which also drives the computationally demanding part of the calculation (solving the integral equation). In the future, interfaces to other equilibrium codes may be added.

In the standard 2-stage approach to stellarator optimization, the virtual casing calculation is run once, at the end of stage 1 (optimizing the plasma shape), with the result provided as input to stage 2 (optimizing the coil shapes). In this case, you can use the save() function or the filename argument of from_vmec() to save the results of the virtual casing calculation. These saved results can then be loaded in later using the load() class method, when needed for solving the stage-2 problem.

In order to compute the external field accurately, one requires fairly high grid resolution. For the stage-2 problem however, a lower resolution is often sufficient. To deal with this, we consider two grids, one denoted by the prefix src_ and the other one denoted by trgt_. src_nphi and src_ntheta refer to the resolution of the grid for the input data, i.e. the total magnetic field and shape of the surface. trgt_nphi and trgt_ntheta refer to the resolution of the grid where the external field is computed, i.e. the output of the virtual casing calculation that is provided as input to the stage 2 problem).

To set the grid resolutions src_nphi and src_ntheta, it can be convenient to use the function simsopt.geo.surface.best_nphi_over_ntheta().

An instance of this class has the following attributes. For all vector quantites, Cartesian coordinates are used, corresponding to array dimensions of size 3:

src_nphi: The number of grid points in the toroidal angle \(\phi\), for a half field period (or a full field period if use_stellsym=False)
src_ntheta: The number of grid points in the poloidal angle \(\theta\).
src_phi: An array of size (src_nphi,) with the grid points of \(\phi\).
src_theta: An array of size (src_ntheta,) with the grid points of \(\theta\).
trgt_nphi: The number of grid points in the toroidal angle \(\phi\), for a half field period (or a full field period if use_stellsym=False)
trgt_ntheta: The number of grid points in the poloidal angle \(\theta\).
trgt_phi: An array of size (trgt_nphi,) with the grid points of \(\phi\).
trgt_theta: An array of size (trgt_ntheta,) with the grid points of \(\theta\).
gamma: An array of size (src_nphi, src_ntheta, 3) with the position vector on the surface.
B_total: An array of size (src_nphi, src_ntheta, 3) with the total magnetic field vector on the surface.
unit_normal: An array of size (trgt_nphi, trgt_ntheta, 3) with the unit normal vector on the surface.
B_external: An array of size (trgt_nphi, trgt_ntheta, 3) with the contribution to the magnetic field due to current outside the surface.
B_external_normal: An array of size (trgt_nphi, trgt_ntheta) with the contribution to the magnetic field due to current outside the surface, taking just the component normal to the surface.

The \(\phi\) and \(\theta\) grids for these data are both uniformly spaced, and are the same as for Surface classes with range="half period" or range="full period", for the case of stellarator-symmetry or non-stellarator-symmetry respectively. (See the description of the range parameter in the documentation on Surfaces.) For the usual case of stellarator symmetry, all the virtual casing data are given on half a field period. There is no grid point at \(\phi=0\), rather the grid is shifted in \(\phi\) by half the grid spacing. Thus, the src_phi grid is np.linspace(1 / (2 * nfp * src_nphi), (src_nphi - 0.5) / (src_nphi * nfp), src_nphi) (recalling the simsopt convention that \(\phi\) and \(\theta\) have period 1, not \(2\pi\)). For a non-stellarator-symmetric calculation, the src_phi grid is np.linspace(0, 1 / nfp, src_nphi, endpoint=False). The trgt_phi grid follows the same logic as the src_phi grid. Note that for stellarator symmetry, if src_nphi != trgt_nphi, then the shift (i.e. first grid point) in src_phi and trgt_phi will be different. For both stellarator symmetry and non-stellarator-symmetry, the src_theta grid is np.linspace(0, 1, src_ntheta, endpoint=False), and the trgt_theta grid is the same but with trgt_ntheta.

In particular, B_external_normal is given on the grid that would be naturally used for stage-2 coil optimization, so no resampling is required.

classmethod from_vmec(vmec, src_nphi, src_ntheta=None, trgt_nphi=None, trgt_ntheta=None, use_stellsym=True, digits=6, filename='auto')

Given a Vmec object, compute the contribution to the total magnetic field due to currents outside the plasma.

This function requires the python virtual_casing package to be installed.

The argument src_nphi refers to the number of points around a half field period if stellarator symmetry is exploited, or a full field period if not.

To set the grid resolutions src_nphi and src_ntheta, it can be convenient to use the function simsopt.geo.surface.best_nphi_over_ntheta(). This is done automatically if you omit the src_ntheta argument.

For now, this routine only works for stellarator symmetry.

Parameters:

vmec – Either an instance of simsopt.mhd.vmec.Vmec, or the name of a Vmec input.* or wout* file.
src_nphi – Number of grid points toroidally for the input of the calculation.
src_ntheta – Number of grid points poloidally for the input of the calculation. If None, the number of grid points will be calculated automatically using simsopt.geo.surface.best_nphi_over_ntheta() to minimize the grid anisotropy, given the specified nphi.
trgt_nphi – Number of grid points toroidally for the output of the calculation. If unspecified, src_nphi will be used.
trgt_ntheta – Number of grid points poloidally for the output of the calculation. If unspecified, src_ntheta will be used.
use_stellsym – whether to exploit stellarator symmetry in the calculation.
digits – Approximate number of digits of precision for the calculation.
filename – If not None, the results of the virtual casing calculation will be saved in this file. For the default value of "auto", the filename will automatically be set to "vcasing_<extension>.nc" where <extension> is the string associated with Vmec input and output files, analogous to the Vmec output file "wout_<extension>.nc".

classmethod load(filename)

Load in the results of a previous virtual casing calculation, previously saved in NetCDF format.

Parameters:: filename – Name of the file to load.

plot(ax=None, show=True)

Plot B_external_normal, the component normal to the surface of the magnetic field generated by currents outside the surface. This routine requires matplotlib.

Parameters:

ax – The axis object on which to plot. This argument is useful when plotting multiple objects on the same axes. If equal to the default None, a new axis will be created.
show – Whether to call matplotlib’s show() function.

Returns:

An axis which could be passed to a further call to matplotlib if desired.

save(filename='vcasing.nc')

Save the results of a virtual casing calculation in a NetCDF file.

Parameters:: filename – Name of the file to create.

simsopt.mhd.vmec module

This module provides a class that handles the VMEC equilibrium code.

class simsopt.mhd.vmec.Vmec(filename: str | None = None, mpi: MpiPartition | None = None, keep_all_files: bool = False, verbose: bool = True, ntheta=50, nphi=50)

Bases: Optimizable

This class represents the VMEC equilibrium code.

You can initialize this class either from a VMEC input.<extension> file or from a wout_<extension>.nc output file. If neither is provided, a default input file is used. When this class is initialized from an input file, it is possible to modify the input parameters and run the VMEC code. When this class is initialized from a wout file, all the data from the wout file is available in memory but the VMEC code cannot be re-run, since some of the input data (e.g. radial multigrid parameters) is not available in the wout file.

The input parameters to VMEC are all accessible as attributes of the indata attribute. For example, if vmec is an instance of Vmec, then you can read or write the input resolution parameters using vmec.indata.mpol, vmec.indata.ntor, vmec.indata.ns_array, etc. However, the boundary surface is different: rbc, rbs, zbc, and zbs from the indata attribute are always ignored, and these arrays are instead taken from the simsopt surface object associated to the boundary attribute. If boundary is a surface based on some other representation than VMEC’s Fourier representation, the surface will automatically be converted to VMEC’s representation (SurfaceRZFourier) before each run of VMEC. You can replace boundary with a new surface object, of any type that implements the conversion function to_RZFourier().

VMEC is run either when the run() function is called, or when any of the output functions like aspect() or iota_axis() are called.

A caching mechanism is implemented, using the attribute need_to_run_code. Whenever VMEC is run, or if the class is initialized from a wout file, this attribute is set to False. Subsequent calls to run() or output functions like aspect() will not actually run VMEC again, until need_to_run_code is changed to True. The attribute need_to_run_code is automatically set to True whenever the state vector .x is changed, and when dofs of the boundary are changed. However, need_to_run_code is not automatically set to True when entries of indata are modified.

Once VMEC has run at least once, or if the class is initialized from a wout file, all of the quantities in the wout output file are available as attributes of the wout attribute. For example, if vmec is an instance of Vmec, then the flux surface shapes can be obtained from vmec.wout.rmnc and vmec.wout.zmns.

Since the underlying fortran implementation of VMEC uses global module variables, it is not possible to have more than one python Vmec object with different parameters; changing the parameters of one would change the parameters of the other.

An instance of this class owns just a few optimizable degrees of freedom, particularly phiedge and curtor. The optimizable degrees of freedom associated with the boundary surface are owned by that surface object.

To run VMEC, two input profiles must be specified: pressure and either iota or toroidal current. Each of these profiles can be specified in several ways. One way is to specify the profile in the input file used to initialize the Vmec object. For instance, the pressure profile is determined by the variables pmass_type, am, am_aux_s, and am_aux_f. You can also modify these variables from python via the indata attribute, e.g. vmec.indata.am = [1.0e5, -1.0e5]. Another option is to assign a simsopt.mhd.profiles.Profile object to the attributes pressure_profile, current_profile, or iota_profile. This approach allows for the profiles to be optimized, and it allows you to use profile shapes defined in python that are not available in the fortran VMEC code. To explain this approach we focus here on the pressure profile; the iota and current profiles are analogous. If the pressure_profile attribute of a Vmec object is None (the default), then a simsopt Profile object is not used, and instead the settings from Vmec.indata (initialized from the input file) are used. If a Profile object is assigned to the pressure_profile attribute, then an edge in the dependency graph is introduced, so the Vmec object then depends on the dofs of the Profile object. Whenever VMEC is run, the simsopt Profile is converted to either a polynomial (power series) or cubic spline in the normalized toroidal flux \(s\), depending on whether indata.pmass_type is "power_series" or "cubic_spline". (The current profile is different in that either "cubic_spline_ip" or "cubic_spline_i" is specified instead of "cubic_spline".) The number of terms in the power series or number of spline nodes is determined by the attributes n_pressure, n_current, and n_iota. If a cubic spline is used, the spline nodes are uniformly spaced from \(s=0\) to 1. Note that the choice of whether a polynomial or spline is used for the VMEC calculation is independent of the subclass of Profile used. Also, whether the iota or current profile is used is always determined by the indata.ncurr attribute: 0 for iota, 1 for current. Example:

from sismopt.mhd.profiles import ProfilePolynomial, ProfileSpline, ProfilePressure, ProfileScaled
from simsopt.util.constants import ELEMENTARY_CHARGE

ne = ProfilePolynomial(1.0e20 * np.array([1, 0, 0, 0, -0.9]))
Te = ProfilePolynomial(8.0e3 * np.array([1, -0.9]))
Ti = ProfileSpline([0, 0.5, 0.8, 1], 7.0e3 * np.array([1, 0.9, 0.8, 0.1]))
ni = ne
pressure = ProfilePressure(ne, Te, ni, Ti)  # p = ne * Te + ni * Ti
pressure_Pa = ProfileScaled(pressure, ELEMENTARY_CHARGE)  # Te and Ti profiles were in eV, so convert to SI here.
vmec = Vmec(filename)
vmec.pressure_profile = pressure_Pa
vmec.indata.pmass_type = "cubic_spline"
vmec.n_pressure = 8  # Use 8 spline nodes

When VMEC is run multiple times, the default behavior is that all wout output files will be deleted except for the first and most recent iteration on worker group 0. If you wish to keep all the wout files, you can set keep_all_files = True. If you want to save the wout file for a certain intermediate iteration, you can set the files_to_delete attribute to [] after that run of VMEC.

Parameters:

filename – Name of a VMEC input.<extension> file or wout_<extension>.nc output file to use for loading the initial parameters. If None, default parameters will be used.
mpi – A simsopt.util.mpi.MpiPartition instance, from which the worker groups will be used for VMEC calculations. If None, each MPI process will run VMEC independently.
keep_all_files – If False, all wout output files will be deleted except for the first and most recent ones from worker group 0. If True, all wout files will be kept.
verbose – Whether to print to stdout when running vmec.

iter: Number of times VMEC has run.

s_full_grid: The “full” grid in the radial coordinate s (normalized toroidal flux), including points at s=0 and s=1. Used for the output arrays and zmns.

s_half_grid: The “half” grid in the radial coordinate s, used for bmnc, lmns, and other output arrays. In contrast to wout files, this array has only ns-1 entries, so there is no leading 0.

ds: The spacing between grid points for the radial coordinate s.

__repr__(): Print the object in an informative way.

aspect(): Return the plasma aspect ratio.

property boundary

property current_profile

external_current()

Return the total electric current associated with external currents, i.e. the current through the “doughnut hole”. This number is useful for coil optimization, to know what the sum of the coil currents must be.

Returns:: float with the total external electric current in Amperes.

get_dofs()

get_input(): Generate a VMEC input file. The result will be returned as a string. To save a file, see the write_input() function.

get_max_mn(): Look through the rbc and zbs data in fortran to determine the largest m and n for which rbc or zbs is nonzero.

iota_axis(): Return the rotational transform on axis

iota_edge(): Return the rotational transform at the boundary

property iota_profile

load_wout(): Read in the most recent wout file created, and store all the data in a wout attribute of this Vmec object.

mean_iota(): Return the mean rotational transform. The average is taken over the normalized toroidal flux s.

mean_shear(): Return an average magnetic shear, d(iota)/ds, where s is the normalized toroidal flux. This is computed by fitting the rotational transform to a linear (plus constant) function in s. The slope of this fit function is returned.

property pressure_profile

recompute_bell(parent=None)

Function to be called whenever new DOFs input is given or if the parent Optimizable’s data changed, so the output from the current Optimizable object is invalid.

This method gets called by various DOF setters. If only the local DOFs of an object are being set, the recompute_bell method is called in that object and also in the descendent objects that have a dependency on the object, whose local DOFs are being changed. If gloabl DOFs of an object are being set, the recompute_bell method is called in the object, ancestors of the object, as well as the descendents of the object.

Need to be implemented by classes that provide a dof_setter for external handling of DOFs.

return_fn_map: Dict[str, Callable] = {'aspect': <function Vmec.aspect>, 'iota_axis': <function Vmec.iota_axis>, 'iota_edge': <function Vmec.iota_edge>, 'mean_iota': <function Vmec.mean_iota>, 'mean_shear': <function Vmec.mean_shear>, 'vacuum_well': <function Vmec.vacuum_well>, 'volume': <function Vmec.volume>}

run(): Run VMEC, if need_to_run_code is True.

set_dofs(x)

set_indata(): Transfer data from simsopt objects to Vmec’s fortran module data. Presently, this function sets the boundary shape and magnetic axis shape. In the future, the input profiles will be set here as well. This data transfer is performed before writing a Vmec input file or running Vmec. The boundary surface object converted to SurfaceRZFourier is returned.

set_profile(longname, shortname, letter): This function is used to set the pressure, current, and/or iota profiles.

update_mpi(new_mpi)

Replace the MpiPartition with a new one.

Parameters:: new_mpi – A new simsopt.util.mpi.MpiPartition object.

vacuum_well()

Compute a single number W that summarizes the vacuum magnetic well, given by the formula

W = (dV/ds(s=0) - dV/ds(s=1)) / (dV/ds(s=0)

where dVds is the derivative of the flux surface volume with respect to the radial coordinate s. Positive values of W are favorable for stability to interchange modes. This formula for W is motivated by the fact that

d^2 V / d s^2 < 0

is favorable for stability. Integrating over s from 0 to 1 and normalizing gives the above formula for W. Notice that W is dimensionless, and it scales as the square of the minor radius. To compute dV/ds, we use

dV/ds = 4 * pi**2 * abs(sqrt(g)_{0,0})

where sqrt(g) is the Jacobian of (s, theta, phi) coordinates, computed by VMEC in the gmnc array, and _{0,0} indicates the m=n=0 Fourier component. Since gmnc is reported by VMEC on the half mesh, we extrapolate by half of a radial grid point to s = 0 and 1.

volume(): Return the volume inside the VMEC last closed flux surface.

write_input(filename)

Write a VMEC input file. To just get the result as a string without saving a file, see the get_input() function.

Parameters:: filename – Name of the file to write. Selected MPI processes can pass None if you wish for these processes to not write a file.

simsopt.mhd.vmec_diagnostics module

This module contains functions that can postprocess VMEC output.

class simsopt.mhd.vmec_diagnostics.IotaTargetMetric(vmec, iota_target, adjoint_epsilon=0.1)

Bases: Optimizable

IotaTargetMetric is a class that computes a metric quantifying the deviation of the rotational transform \(\iota\) in from a prescribed target profile in a Vmec equilibrium:

\[J = \frac{1}{2} \int ds \, (\iota - \iota_{target})^2\]

where the integral is over the normalized toroidal flux \(s\), and the function \(\iota_{target}(s)\) corresponds to the argument iota_target. This class also can compute the derivatives of \(J\) using an adjoint method.

Parameters:

vmec – instance of Vmec
iota_target – function handle which takes a single argument, s, the normalized toroidal flux, and returns the target rotational transform.
adjoint_epsilon – sets the amplitude of the toroidal current perturbation required for the adjoint solve.

J(): Computes the quantity \(J\) described in the class definition.

dJ(): Computes derivatives of \(J\) with respect to surface parameters using an adjoint method.

shape_gradient()

Computes the shape gradient of the quantity \(J\) described in the class definition. For a perturbation to the surface \(\delta \vec{x}\), the resulting perturbation to the objective function is

\[\delta J(\delta \vec{x}) = \int d^2 x \, G \delta \vec{x} \cdot \vec{n}\]

where the integral is over the VMEC boundary surface, \(G\) is the shape gradient, and \(\vec{n}\) is the unit normal.

Returns:: 2d array of size (numquadpoints_phi,numquadpoints_theta)
Return type:: \(G\)

class simsopt.mhd.vmec_diagnostics.IotaWeighted(vmec, weight_function, adjoint_epsilon=0.1)

Bases: Optimizable

Computes a weighted average of the rotational transform for a VMEC configuration. The quantity computed is defined by

\[J = \frac{ \int ds \, \iota(s) w(s)} { \int ds \, w(s)}\]

where \(w(s)\) is a prescribed weight function, corresponding to the argument weight_function. This class also can compute the derivatives of \(J\) using an adjoint method.

Parameters:

vmec – instance of Vmec
weight_function – function handle which takes a single argument, s, the normalized toroidal flux
adjoint_epsilon – sets the amplitude of the toroidal current perturbation required for the adjoint solve.

J(): Computes the quantity \(J\) described in the class definition.

dJ(): Computes derivatives of \(J\) with respect to surface parameters using an adjoint method.

shape_gradient()

Computes the shape gradient of the quantity \(J\) described in the class definition. For a perturbation to the surface \(\delta \vec{x}\), the resulting perturbation to the objective function is

\[\delta J(\delta \vec{x}) = \int d^2 x \, G \delta \vec{x} \cdot \vec{n}\]

where the integral is over the VMEC boundary surface, \(G\) is the shape gradient, and \(\vec{n}\) is the unit normal.

Returns:: 2d array of size (numquadpoints_phi,numquadpoints_theta)
Return type:: \(G\)

class simsopt.mhd.vmec_diagnostics.QuasisymmetryRatioResidual(vmec: Vmec, surfaces: float | Sequence[Real] | NDArray[Any, float64], helicity_m: int = 1, helicity_n: int = 0, weights: Sequence[Real] | NDArray[Any, float64] | None = None, ntheta: int = 63, nphi: int = 64)

Bases: Optimizable

This class provides a measure of the deviation from quasisymmetry, one that can be computed without Boozer coordinates. This metric is based on the fact that for quasisymmetry, the ratio

\[(\vec{B}\times\nabla B \cdot\nabla\psi) / (\vec{B} \cdot\nabla B)\]

is constant on flux surfaces.

Specifically, this class represents the objective function

\[f = \sum_{s_j} w_j \left< \left[ \frac{1}{B^3} \left( (N - \iota M)\vec{B}\times\nabla B\cdot\nabla\psi - (MG+NI)\vec{B}\cdot\nabla B \right) \right]^2 \right>\]

where the sum is over a set of flux surfaces with normalized toroidal flux \(s_j\), the coefficients \(w_j\) are user-supplied weights, \(\left< \ldots \right>\) denotes a flux surface average, \(G(s)\) is \(\mu_0/(2\pi)\) times the poloidal current outside the surface, \(I(s)\) is \(\mu_0/(2\pi)\) times the toroidal current inside the surface, \(\mu_0\) is the permeability of free space, \(2\pi\psi\) is the toroidal flux, and \((M,N)\) are user-supplied integers that specify the desired helicity of symmetry. If the magnetic field is quasisymmetric, so \(B=B(\psi,\chi)\) where \(\chi=M\vartheta - N\varphi\) where \((\vartheta,\varphi)\) are the poloidal and toroidal Boozer angles, then \(\vec{B}\times\nabla B\cdot\nabla\psi \to -(MG+NI)(\vec{B}\cdot\nabla\varphi)\partial B/\partial \chi\) and \(\vec{B}\cdot\nabla B \to (-N+\iota M)(\vec{B}\cdot\nabla\varphi)\partial B/\partial \chi\), implying the metric \(f\) vanishes. The flux surface average is discretized using a uniform grid in the VMEC poloidal and toroidal angles \((\theta,\phi)\). In this case \(f\) can be written as a finite sum of squares:

\[f = \sum_{s_j, \theta_j, \phi_j} R(s_j, \theta_k, \phi_{\ell})^2\]

where the \(\phi_{\ell}\) grid covers a single field period. Here, each residual term is

\[R(\theta, \phi) = \sqrt{w_j \frac{n_{fp} \Delta_{\theta} \Delta_{\phi}}{V'}|\sqrt{g}|} \frac{1}{B^3} \left( (N-\iota M)\vec{B}\times\nabla B\cdot\nabla\psi - (MG+NI)\vec{B}\cdot\nabla B \right).\]

Here, \(n_{fp}\) is the number of field periods, \(\Delta_{\theta}\) and \(\Delta_{\phi}\) are the spacing of grid points in the poloidal and toroidal angles, \(\sqrt{g} = 1/(\nabla s\cdot\nabla\theta \times \nabla\phi)\) is the Jacobian of the \((s,\theta,\phi)\) coordinates, and \(V' = \int_0^{2\pi} d\theta \int_0^{2\pi}d\phi |\sqrt{g}| = dV/d\psi\) where \(V\) is the volume enclosed by a flux surface.

Parameters:

vmec – A simsopt.mhd.vmec.Vmec object from which the quasisymmetry error will be calculated.
surfaces – Value of normalized toroidal flux at which you want the quasisymmetry error evaluated, or a list of values. Each value must be in the interval [0, 1], with 0 corresponding to the magnetic axis and 1 to the VMEC plasma boundary. This parameter corresponds to \(s_j\) above.
helicity_m – Desired poloidal mode number \(M\) in the magnetic field strength \(B\), so \(B = B(s, M \vartheta - n_{fp} \hat{N} \varphi)\) where \(\vartheta\) and \(\varphi\) are Boozer angles.
helicity_n – Desired toroidal mode number \(\hat{N} = N / n_{fp}\) in the magnetic field strength \(B\), so \(B = B(s, M \vartheta - n_{fp} \hat{N} \varphi)\) where \(\vartheta\) and \(\varphi\) are Boozer angles. Note that the supplied value of helicity_n will be multiplied by the number of field periods \(n_{fp}\), so typically helicity_n should be +1 or -1 for quasi-helical symmetry.
weights – The list of weights \(w_j\) for each flux surface. If None, a weight of \(w_j=1\) will be used for all surfaces.
ntheta – Number of grid points in \(\theta\) used to discretize the flux surface average.
nphi – Number of grid points per field period in \(\phi\) used to discretize the flux surface average.

compute(): Compute the quasisymmetry metric. This function returns an object that contains (as attributes) all the intermediate quantities for the calculation. Users do not need to call this function for optimization; instead the residuals() function can be used. However, this function can be useful if users wish to inspect the quantities going into the calculation.

profile(): Return the quasisymmetry metric in terms of a 1D radial profile. The residuals \(R\) are squared and summed over theta and phi, but not over s. The total quasisymmetry error \(f\) returned by the total() function is the sum of the values in the profile returned by this function.

residuals(): Evaluate the quasisymmetry metric in terms of a 1D numpy vector of residuals, corresponding to \(R\) in the documentation for this class. This is the function to use when forming a least-squares objective function.

total(): Evaluate the quasisymmetry metric in terms of the scalar total \(f\).

class simsopt.mhd.vmec_diagnostics.WellWeighted(vmec, weight_function1, weight_function2, adjoint_epsilon=0.1)

Bases: Optimizable

WellWeighted is a class that computes a measure of magnetic well for a vmec equilibrium. The magnetic well measure is

\[J = \frac{ \int ds \, V'(s) [w_1(s) - w_2(s)]} { \int ds \, V'(s) [w_1(s) + w_2(s)]},\]

where \(w_1(s)\) and \(w_2(s)\) correspond to the arguments weight_function1 and weight_function2, and \(V(s)\) is the volume enclosed by the flux surface with normalized toroidal flux \(s\). Typically, \(w_1\) would be peaked on the edge while \(w_2\) would be peaked on the axis, such that \(J < 0\) corresonds to \(V''(s) < 0\), which is favorable for stability.

This class also provides calculations of the derivatives of \(J\) using an adjoint method.

Parameters:

vmec – instance of Vmec
weight_function1 – function handle which takes a single argument, s, the normalized toroidal flux
weight_function2 – function handle which takes a single argument, s, the normalized toroidal flux
adjoint_epsilon – sets the amplitude of the toroidal current perturbation required for the adjoint solve.

J(): Computes the quantity \(J\) described in the class definition.

dJ(): Computes derivatives of \(J\) with respect to surface parameters using an adjoint method.

shape_gradient()

Computes the shape gradient of the quantity \(J\) described in the class definition. For a perturbation to the surface \(\delta \vec{x}\), the resulting perturbation to the objective function is

\[\delta J(\delta \vec{x}) = \int d^2 x \, G \delta \vec{x} \cdot \vec{n}\]

where the integral is over the VMEC boundary surface, \(G\) is the shape gradient, and \(\vec{n}\) is the unit normal.

Returns:: 2d array of size (numquadpoints_phi,numquadpoints_theta)
Return type:: \(G\)

simsopt.mhd.vmec_diagnostics.vmec_compute_geometry(vs, s, theta, phi, phi_center=0)

Compute many geometric quantities of interest from a vmec configuration.

Some of the quantities computed by this function refer to alpha, a field line label coordinate defined by

\[\alpha = \theta_{pest} - \iota (\phi - \phi_{center}).\]

Here, \(\phi_{center}\) is a constant, usually 0, which can be set to a nonzero value if desired so the magnetic shear contribution to \(\nabla\alpha\) vanishes at a toroidal angle different than 0. Also, wherever the term psi appears in variable names in this function and the returned arrays, it means \(\psi =\) the toroidal flux divided by \(2\pi\), so

\[\vec{B} = \nabla\psi\times\nabla\theta_{pest} + \iota\nabla\phi\times\nabla\psi = \nabla\psi\times\nabla\alpha.\]

Most of the arrays that are returned by this function have shape (ns, ntheta, nphi), where ns is the number of flux surfaces, ntheta is the number of grid points in VMEC’s poloidal angle, and nphi is the number of grid points in the standard toroidal angle. For the arguments theta and phi, you can either provide 1D arrays, in which case a tensor product grid is used, or you can provide 3D arrays of shape (ns, ntheta, nphi). In this latter case, the grids are not necessarily tensor-product grids. Note that all angles in this function have period \(2\pi\), not period 1.

The output arrays are returned as attributes of the returned object. Many intermediate quantities are included, such as the Cartesian components of the covariant and contravariant basis vectors. Some of the most useful of these output arrays are (all with SI units):

phi: The standard toroidal angle \(\phi\).
theta_vmec: VMEC’s poloidal angle \(\theta_{vmec}\).
theta_pest: The straight-field-line angle \(\theta_{pest}\) associated with \(\phi\).
modB: The magnetic field magnitude \(|B|\).
B_sup_theta_vmec: \(\vec{B}\cdot\nabla\theta_{vmec}\).
B_sup_phi: \(\vec{B}\cdot\nabla\phi\).
B_cross_grad_B_dot_grad_alpha: \(\vec{B}\times\nabla|B|\cdot\nabla\alpha\).
B_cross_grad_B_dot_grad_psi: \(\vec{B}\times\nabla|B|\cdot\nabla\psi\).
B_cross_kappa_dot_grad_alpha: \(\vec{B}\times\vec{\kappa}\cdot\nabla\alpha\), where \(\vec{\kappa}=\vec{b}\cdot\nabla\vec{b}\) is the curvature and \(\vec{b}=|B|^{-1}\vec{B}\).
B_cross_kappa_dot_grad_psi: \(\vec{B}\times\vec{\kappa}\cdot\nabla\psi\).
grad_alpha_dot_grad_alpha: \(|\nabla\alpha|^2 = \nabla\alpha\cdot\nabla\alpha\).
grad_alpha_dot_grad_psi: \(\nabla\alpha\cdot\nabla\psi\).
grad_psi_dot_grad_psi: \(|\nabla\psi|^2 = \nabla\psi\cdot\nabla\psi\).
iota: The rotational transform \(\iota\). This array has shape (ns,).
shat: The magnetic shear \(\hat s= (x/q) (d q / d x)\) where \(x = \mathrm{Aminor_p} \, \sqrt{s}\) and \(q=1/\iota\). This array has shape (ns,).

The following normalized versions of these quantities used in the gyrokinetic codes stella, gs2, and GX are also returned: bmag, gbdrift, gbdrift0, cvdrift, cvdrift0, gds2, gds21, and gds22, along with L_reference and B_reference. Instead of gradpar, two variants are returned, gradpar_theta_pest and gradpar_phi, corresponding to choosing either \(\theta_{pest}\) or \(\phi\) as the parallel coordinate.

The value(s) of s provided as input need not coincide with the full grid or half grid in VMEC, as spline interpolation will be used radially.

The implementation in this routine is similar to the one in the gyrokinetic code stella.

Example usage:

import numpy as np
from simsopt.mhd import Vmec, vmec_compute_geometry

v = Vmec("wout_li383_1.4m.nc")
s = 1
theta = np.linspace(0, 2 * np.pi, 50)
phi = np.linspace(0, 2 * np.pi / 3, 60)
data = vmec_compute_geometry(v, s, theta, phi)
print(data.grad_s_dot_grad_s)

Parameters:

vs – Either an instance of simsopt.mhd.vmec.Vmec or the structure returned by vmec_splines().
s – Values of normalized toroidal flux on which to construct the field lines. You can give a single number, or a list or numpy array.
theta – Values of vmec’s poloidal angle. You can provide a float, a 1d array of size (ntheta,), or a 3d array of size (ns, ntheta, nphi).
phi – Values of the standard toroidal angle. You can provide a float, a 1d array of size (nphi,) or a 3d array of size (ns, ntheta, nphi).
phi_center – \(\phi_{center}\), an optional shift to the toroidal angle in the definition of \(\alpha\).

simsopt.mhd.vmec_diagnostics.vmec_fieldlines(vs, s, alpha, theta1d=None, phi1d=None, phi_center=0, plot=False, show=True)

Compute field lines in a vmec configuration, and compute many geometric quantities of interest along the field lines. In particular, this routine computes the geometric quantities that enter the gyrokinetic equation.

One task performed by this function is to convert between the poloidal angles \(\theta_{vmec}\) and \(\theta_{pest}\). The latter is the angle in which the field lines are straight when used in combination with the standard toroidal angle \(\phi\). Note that all angles in this function have period \(2\pi\), not period 1.

To specify the parallel extent of the field lines, you can provide either a grid of \(\theta_{pest}\) values or a grid of \(\phi\) values. If you specify both or neither, ValueError will be raised.

The geometric quanties computed by this function are the same as for vmec_compute_geometry(). See the documentation of that function for details.

Most of the arrays that are returned by this function have shape (ns, nalpha, nl), where ns is the number of flux surfaces, nalpha is the number of field lines on each flux surface, and nl is the number of grid points along each field line. In other words, ns is the size of the input s array, nalpha is the size of the input alpha array, and nl is the size of the input theta1d or phi1d array. The output arrays are returned as attributes of the returned object.

The value(s) of s provided as input need not coincide with the full grid or half grid in VMEC, as spline interpolation will be used radially.

Example usage:

import numpy as np
from simsopt.mhd import Vmec, vmec_fieldlines

v = Vmec('wout_li383_1.4m.nc')
theta = np.linspace(-np.pi, np.pi, 50)
fl = vmec_fieldlines(v, 0.5, 0, theta1d=theta)
print(fl.B_cross_grad_B_dot_grad_alpha)

Parameters:

vs – Either an instance of simsopt.mhd.vmec.Vmec or the structure returned by vmec_splines().
s – Values of normalized toroidal flux on which to construct the field lines. You can give a single number, or a list or numpy array.
alpha – Values of the field line label \(\alpha\) on which to construct the field lines. You can give a single number, or a list or numpy array.
theta1d – 1D array of \(\theta_{pest}\) values, setting the grid points along the field line and the parallel extent of the field line.
phi1d – 1D array of \(\phi\) values, setting the grid points along the field line and the parallel extent of the field line.
phi_center – \(\phi_{center}\), an optional shift to the toroidal angle in the definition of \(\alpha\).
plot – Whether to create a plot of the main geometric quantities. Only one field line will be plotted, corresponding to the leading elements of s and alpha.
show – Only matters if plot==True. Whether to call matplotlib’s show() function after creating the plot.

simsopt.mhd.vmec_diagnostics.vmec_splines(vmec)

Initialize radial splines for a VMEC equilibrium.

Parameters:: vmec – An instance of simsopt.mhd.vmec.Vmec.
Returns:: A structure with the splines as attributes.