simsopt.mhd package

Submodules

simsopt.mhd.boozer module

This module provides a class that handles the transformation to Boozer coordinates, and an optimization target for quasisymmetry.

class simsopt.mhd.boozer.Boozer(equil: simsopt.mhd.vmec.Vmec, mpol: int = 32, ntor: int = 32)

Bases: simsopt._core.graph_optimizable.Optimizable

This class handles the transformation to Boozer coordinates.

A Boozer instance maintains a set “s”, which is a registry of the surfaces on which other objects want Boozer-coordinate data. When the run() method is called, the Boozer transformation is carried out on all these surfaces. The registry can be cleared at any time by setting the s attribute to {}.

__init__(equil: simsopt.mhd.vmec.Vmec, mpol: int = 32, ntor: int = 32) None

Constructor

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.

register(s: Union[float, Iterable[float]]) None

This function is called by objects that depend on this Boozer object, to indicate that they will want Boozer data on the given set of surfaces.

Parameters

s – 1 or more surfaces on which Boozer data will be requested.

run()

Run booz_xform on all the surfaces that have been registered.

class simsopt.mhd.boozer.Quasisymmetry(boozer: simsopt.mhd.boozer.Boozer, s: Union[float, Iterable[float]], m: int, n: int, normalization: str = 'B00', weight: str = 'even')

Bases: simsopt._core.graph_optimizable.Optimizable

This class is used to compute the departure from quasisymmetry on a given flux surface based on the Boozer spectrum.

Parameters
  • boozer – A Boozer object on which the calculation will be based.

  • s – The normalized toroidal magnetic flux for the flux surface to analyze. Should be in the range [0, 1].

  • m – The poloidal mode number of the symmetry you want to achive. The departure from symmetry B(m * theta - nfp * n * zeta) will be reported.

  • n – The toroidal mode number of the symmetry you want to achieve. The departure from symmetry B(m * theta - nfp * n * zeta) will be reported.

  • normalization – A uniform normalization applied to all bmnc harmonics. If "B00", the symmetry-breaking modes will be divided by the m=n=0 mode amplitude on the same surface. If "symmetric", the symmetry-breaking modes will be divided by the square root of the sum of the squares of all the symmetric modes on the same surface. This is the normalization used in stellopt.

  • weight – An option for a m- or n-dependent weight to be applied to the bmnc amplitudes.

J() numpy.ndarray

Carry out the calculation of the quasisymmetry error.

Returns

1D numpy array listing all the normalized mode amplitudes of symmetry-breaking Fourier modes of |B|.

__init__(boozer: simsopt.mhd.boozer.Boozer, s: Union[float, Iterable[float]], m: int, n: int, normalization: str = 'B00', weight: str = 'even') None

Constructor

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] = {'J': <function Quasisymmetry.J>}

simsopt.mhd.spec module

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

class simsopt.mhd.spec.Residue(spec, pp, qq, vol=1, theta=0, s_guess=None, s_min=- 1.0, s_max=1.0, rtol=1e-09)

Bases: simsopt._core.graph_optimizable.Optimizable

Greene’s residue, evaluated from a Spec equilibrum

Parameters
  • spec – a Spec object

  • pp – Numerator and denominator for the resonant iota = pp / qq

  • qq – Numerator and denominator for the resonant iota = pp / qq

  • vol – Index of the Spec volume to consider

  • theta – Spec’s theta coordinate at the periodic field line

  • s_guess – Guess for the value of Spec’s s coordinate at the periodic field line

  • s_min – bounds on s for the search

  • s_max – bounds on s for the search

  • rtol – the relative tolerance of the integrator

J()

Run Spec if needed, find the periodic field line, and return the residue

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.

simsopt.mhd.vmec module

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

class simsopt.mhd.vmec.Vmec(filename: Optional[str] = None, mpi: Optional[simsopt.util.mpi.MpiPartition] = None, keep_all_files: bool = False, ntheta=50, nphi=50)

Bases: simsopt._core.graph_optimizable.Optimizable

This class represents the VMEC equilibrium code.

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, 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 set_dofs() is called. However, need_to_run_code is not automatically set to True when entries of indata are modified, or when boundary is modified.

Once VMEC has run at least once, 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.

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 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.

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
get_dofs()
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

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.

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

simsopt.mhd.vmec_diagnostics module

This module contains functions that can postprocess VMEC output.

simsopt.mhd.vmec_diagnostics.B_cartesian(vmec)

Computes Cartesian vector components of magnetic field on boundary on a grid in the vmec toroidal and poloidal angles. This is required to compute adjoint-based shape gradients.

Parameters

vmec – instance of Vmec

Returns: 3 element tuple containing (Bx, By, Bz)
Bx, By, Bz2d arrays of size (numquadpoints_phi,numquadpoints_theta)

defining Cartesian components of magnetic field on vmec.boundary

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

Bases: simsopt._core.graph_optimizable.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: simsopt._core.graph_optimizable.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: simsopt.mhd.vmec.Vmec, surfaces: Union[float, Sequence[numbers.Real], nptyping.types._ndarray.NDArray[None, nptyping.types._number.Float[float, numpy.floating]]], helicity_m: int = 1, helicity_n: int = 0, weights: Optional[Union[Sequence[numbers.Real], nptyping.types._ndarray.NDArray[None, nptyping.types._number.Float[float, numpy.floating]]]] = None, ntheta: int = 63, nphi: int = 64)

Bases: simsopt._core.graph_optimizable.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: simsopt._core.graph_optimizable.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\)

Module contents