plasmacontrol/desc

Stellarator Equilibrium and Optimization Suite

Stellarator Optimization Package

DESC solves for and optimizes 3D MHD equilibria using pseudo-spectral numerical methods and automatic differentiation.

The theoretical approach and implementation details used by DESC are presented in the following papers and documented at Theory. Please cite our work if you use DESC!

• Dudt, D. & Kolemen, E. (2020). DESC: A Stellarator Equilibrium Solver. [Physics of Plasmas] [pdf]
• Panici, D. et al (2023). The DESC Stellarator Code Suite Part I: Quick and accurate equilibria computations. [Journal of Plasma Physics] [pdf]
• Conlin, R. et al. (2023). The DESC Stellarator Code Suite Part II: Perturbation and continuation methods. [Journal of Plasma Physics] [pdf]
• Dudt, D. et al. (2023). The DESC Stellarator Code Suite Part III: Quasi-symmetry optimization. [Journal of Plasma Physics] [pdf]

Quick Start

The easiest way to install DESC is from PyPI: pip install desc-opt

For more detailed instructions on installing DESC and its dependencies, see Installation.

The best place to start learning about DESC is our tutorials:

• Basic fixed boundary equilibrium: running from a VMEC input, creating an equilibrium from scratch
• Advanced equilibrium: continuation and perturbation methods.
• Free boundary equilibrium: vacuum and or finite beta with external field.
• Using DESC outputs: analysis, plotting, saving to VMEC format.
• Basic optimization: specifying objectives, fixing degrees of freedom.
• Advanced optimization: advanced constraints, precise quasi-symmetry, constrained optimization.
• Near axis constraints: loading solutions from QSC/QIC and fixing near axis expansion.

For details on the various objectives, constraints, optimizable objects and more, see the full api documentation.

If all you need is an equilibrium solution, the simplest method is through the command line by giving an input file specifying the equilibrium and solver options, this way can also can also accept VMEC input files.

The code is run using the syntax desc <path/to/inputfile> and the full list of command line options are given in Command Line Interface. (Note that you may have to prepend the command with python -m)

Refer to Inputs for documentation on how to format the input file.

The equilibrium solution is output in a HDF5 binary file, whose format is detailed in Outputs.

Repository Contents

• desc contains the source code including the main script and supplemental files. Refer to the API documentation for details on all of the available functions and classes.
• docs contains the documentation files.
• tests contains routines for automatic testing.
• publications contains PDFs of publications by the DESC group, as well as scripts and data to reproduce the results of these papers.

Contribute

• Contributing guidelines
• Issue Tracker
• Source Code
• Documentation

desc's Issues

2nd order perturbations

Autodiff hessian and solve quadratic tensor equation for dx
ways to format it as matrix equation?

Add node option to avoid axis / rational surfaces

Flag for avoiding axis could help with debugging stuff there

Also having option to place at "most rational" surfaces.

Avoid unneeded recalculation of basis functions when changing resolution with fft method

Direct method currently only computes basis fns when they're not already part of the basis, while the fft method recomputes everything every time the resolution changes

allow for ansi indexing

I added an option in the get_zern_basis_idx_dense function for either fringe or ansi indexing, but the symmetry constraint matrix still assumes fringe indexing. Can we change that to just use the zern_idx array instead of M,N? Also we'll need to add some sort of "indexing" option to the input file and parser

Set up user group / community platform

We should have some way to keep track of who is using the code, and for users to be able to ask questions.
some options:
email list
https://www.discourse.org/
https://slack.com/
https://discord.com/
https://gitter.im/

Don't compute more field components than necessary

Different objective functions need different field components, so we should add a flag so that only the needed components are computed. Might save a bit of time.

compute_bc_err_RZ all kinds of broken

It doesn't work with the bdry ratio argument to ignore errors on 3d modes.
Also hasnt been tested in a while and could probably use some speed improvements and make the signature match the others

Profiling & jit

Need to do a lot of testing to see where bottlenecks are, and see if jit can help, or if we should use np over jnp in places

Finish VMEC I/O

Currently we can read VMEC inputs (i think it automatically detects?), but it would be good if we could also read/write vmec outputs and interpolate between the two bases (for comparisons, and for using DESC outputs in places that expect VMEC format)

To Do:

• fix bug in input conversion
• load vmec solution as initial guess (ie convert vmec to desc)
• save desc solution as vmec .nc format

Add Conservative default settings for continuation method?

What are good paramters for resolution & ratios during a multistep solution?

vmec_to_desc_input

cannot handle multi-line VMEC inputs
need to remove pres_scale from DESC input file

Look into node weighting

We're currently weighting the nodes by their volume elements, so the error should be the total force balance error in the domain. But it seems like this fluctuates a lot as you vary the number of nodes and basis functions (I would assume that for a given guess the total error shouldn't depend strongly on the discretization).

Zernike transform using FFT

rather than full matrx stuff, we can use FFT in the toroidal direction (like the matlab version)
Possible speed & memory gain, but will need to ensure uniform spacing and node pattern in the toroidal direction.

Checkpointing

Need to save intermediate results from continuation method, also return output from all steps instead of just the last one

Change behavior of keyboard interupt

Right now, if the user gives the keyboard interupt Ctrl-C, it exists DESC entirely. Another option that might be more convenient for interactive use would be if it opened a small terminal menu, where the user could select options like, press c to continue, press l to exit the current loop and proceed to the next outer iteration, press e to exit all loops and go to the end of the program (but still saving and plotting data), press d to enter debugger at current position. If they really want to hard exit, they'd do Ctrl-C twice (like when exiting jupyter)

Option for "square" system with root finding method

Instead of overdetermined system, find way to make it square and use root finding methods

Docstrings

Lots of functions don't have em, and some that do end up formatted weird with Sphinx/RTD

Automated check for nested surfaces

I think this might be really useful for cases where you're running a high res case but it goes wacky somewhere early and you waste a lot of time.

matplotlib has a "path" object that we can probably use to easily check if one curve encloses another (ie, are the surfaces properly nested)

https://matplotlib.org/api/path_api.html

Separate "objective" from "constraint" and make composable objective functions

As we move from just computing equilibria to optimizing them, we will need to be clear about which objectives are really physics objectives (eg QS, particle confinement, stabilty etc) and which are really constraints (eg, force balance).

We might want to make those two things separate classes or treat them differently in some way

It would also be good to have an easy way of combining different objectives with a given weight, so that we could do something like
objective = 0.5*QS() + 10*Confinement()

`jax RuntimeError: Invalid argument: Cannot copy value to device 'gpu:0' with 'cpu' backend` on trying to run W7X example, with jax >= 0.2.2

Hi, I wanted to try running your code on the GPU, but got the following error:

Code execution command, error

dominik@dell ~/Code/DESC % XLA_FLAGS=--xla_gpu_cuda_data_dir=/opt/cuda python desc examples/DESC/W7X -g 0 -p                                                                        [1]

 ____  ____  _____   ___                                                                                                                                                                
|  _ \| ___|/  ___|/ ___|                                                                                                                                                               
| | \ | |_  | (__ | |                                                                                                                                                                   
| | | |  _| \___ \| |                                                                                                                                                                   
| |_/ | |__  ___) | |___                                                                                                                                                                
|____/|____||____/ \____|                                                                                                                                                               
                                                                                                                                                                                        

DESC version 0.2.0, using JAX backend, jax version=0.2.6, jaxlib version=0.1.57, dtype=float64
Using device: gpu:0
Reading input from /home/dominik/Code/DESC/examples/DESC/W7X
Output will be written to /home/dominik/Code/DESC/examples/DESC/W7X.output
================
Step 1/10
================
Spectral resolution (M,N,delta_lm)=(4,2,4)
Node resolution (M,N)=(6,3)
Boundary ratio = 1.0
Pressure ratio = 1.0
Zeta ratio = 1.0
Error ratio = 0.0001
Perturbation Order = 1
Function tolerance = 0.0001
Gradient tolerance = 0.0001
State vector tolerance = 1e-06
Max function evaluations = 5000
================
Precomputing Fourier-Zernike basis
Computing initial guess
Compiling objective function
Starting optimization
Warning: Desired error not necessarily achieved due to precision loss.
         Current function value: 15.013052
         Iterations: 3539
         Function evaluations: 3625
         Gradient evaluations: 3612
Start of Step 1:
Weighted Loss:  2.010e+14  errFrho:  1.984e+07  errFbeta:  2.922e+06  errRb:  1.192e-14  errZb:  8.984e-16  errL0:  0.000e+00
End of Step 1:
Traceback (most recent call last):
  File "desc/__main__.py", line 181, in <module>
    main(sys.argv[1:])
  File "desc/__main__.py", line 155, in main
    iterations, timer = solve_eq_continuation(
  File "/home/dominik/Code/DESC/desc/continuation.py", line 362, in solve_eq_continuation
    callback(x, *args)
  File "/home/dominik/Code/DESC/desc/objective_funs.py", line 96, in callback
    errRb, errZb = bdry_fun(
  File "/home/dominik/Code/DESC/desc/boundary_conditions.py", line 226, in compute_bdry_err_four
    lamda = eval_double_fourier(cL, lambda_idx, NFP, vartheta, zeta)
jax._src.traceback_util.FilteredStackTrace: RuntimeError: Invalid argument: Cannot copy value to device 'gpu:0' with 'cpu' backend

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/usr/lib/python3.8/runpy.py", line 194, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/usr/lib/python3.8/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "desc/__main__.py", line 181, in <module>
    main(sys.argv[1:])
  File "desc/__main__.py", line 155, in main
    iterations, timer = solve_eq_continuation(
  File "/home/dominik/Code/DESC/desc/continuation.py", line 362, in solve_eq_continuation
    callback(x, *args)
  File "/home/dominik/Code/DESC/desc/objective_funs.py", line 96, in callback
    errRb, errZb = bdry_fun(
  File "/home/dominik/Code/DESC/desc/boundary_conditions.py", line 226, in compute_bdry_err_four
    lamda = eval_double_fourier(cL, lambda_idx, NFP, vartheta, zeta)
  File "/home/dominik/.local/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 133, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/dominik/.local/lib/python3.8/site-packages/jax/api.py", line 371, in f_jitted
    return cpp_jitted_f(*args, **kwargs)
RuntimeError: Invalid argument: Cannot copy value to device 'gpu:0' with 'cpu' backend

I'm not sure whether this is jax-specific; I haven't really used it, so I don't know how to reproduce it in a MVE to send over to them.

v0.4.0 to do list

Some stuff we should get done or look at before merging back into master. Most of it is super quick, just check them off as you do them.

General

• check how super().init works with multiple parents
• test out zeta ratio
• clean up documentation (no more DESCRIPTION, make sure everything is up to date with all the recent changes)
• examples of code in use - maybe a notebook or two showing how to use it interactively

Perturbations

• optionally pass in $df/dx$
• use jvp for everything else
• make it a function + method of equilibrium
• return new equilibrium (eq.copy())
• force equilibrium to have right resolution before perturbation
• def perturb(equil, objective, dx:dict, Jx=None, etc)

Configuration

• resolution for $p$, $\iota$ same as configuration
• use np.where(x==[l,m,n]) instead of logical_and.reduce (see old transform.change_resolution)
• make configuration ABC, not public
• in config, change_resolution also changes profiles/bdry resolution
• remove setter methods for basis in config

Equilibrium / EquilibriumFamily

• force equilibrium to have right resolution before perturbation
• equil family only accepts equil as list members

BoundaryConditions

• add build method to BC to delay expensive SVD of A until we need it

• Backend/Utils

• remove cross function, just import directly from jax

• find better jax replacement for dot

• replace tristate with 'cos'/'sin'/None

Objective

• objective knows which transforms it needs, equilibrium gets info when objective assigned
• objective has attribute for string form of its name
• objectivefactory should be a function, not a class
• objective has other system derivatives (df/dc stuff for each c)
• objective function for minimizing energy
• objective function for QS (time permitting)

Grid

• quadrature node options for radial coordinate
• quadrature grid option for full quadrature setup
• replace "volumes" with 1d array of weights

Basis

• numba jit where possible in basis
• use np instead of jnp in building transforms
• possibly extend fringe to higher L

Transform

• better way to do fft without breaking up arrays (DCT/DST?)
• figure out why fft doesnt work with symmetric basis

IO

• flag to save as vmec format
• parse file name for output to figure out what format

Plotting

• logo plot function is all kinds broken
• tests for plotting: http://www.davidketcheson.info/2015/01/13/using_matplotlib_image_comparison.html

Clean up / standardize calling signatures

Make sure our notation for different variables is consistent, and the formats/containers we use are consistent.

i have issues

Add generic option for optimization method

We should at least have the option for using different solvers, whether newton or quasi, vector or scalar etc. We can add support for different options as we go, but for a start it seems like the least squares trust region and BFGS are decent options

Get symmetric nodes and fft method to work together

Right now using symmetric nodes requires using the direct method for the zernike transform, which is a lot slower.

With fft, no symmetry:

With symmetry & direct method:

All these tests were using the w7x case with m,n = 12, mnodes,nnodes=18

Using symmetry and the direct method gives slightly better compile time for the gradient, but is waay slower for pretty much everything else, both compiling and evaluating.
As it is now, using symmetry really isn't worth it, despite it reducing the dimension of the system by a half.

We could probably switch to using a different set of "symmetric" nodes that arent technically the right symmetry nodes but would work with the fft. Since we're already oversampling it's probably ok to not use the exact right nodes.
We could alternatively figure out how to redo the fft to work with an arbitrary node pattern (though from some earlier testing, doing the fourier and zernike transforms in the opposite order slows the transform part down,though it might be worth it)

Another thing to note is that the direct method uses a TON more memory when JIT compiling the first time (I kept running out even with 32 gigs) so thats another reason to use the fft, especially for high res

[BUG] Memory magagement for GPU

Running on GPU runs out of memory.

Bug in reading input for floating point values

if floats in the input file are given as .01 it gets read as 1, while 0.01 gets read correctly

(R0, Z0, r, lambda) independent variables instead of (R, Z)

We are currently using R & Z as the independent variables. This might be an inefficient discretization in spectral space because we are restricted to using the SFL coordinates.

Alternatively, we could use R0, Z0, r, lambda with:

This has two potential advantages:

1. More efficient spectral representation (fewer Fourier modes required).
2. The BC can be applied as a linear constraint (satisfied exactly at each iteration).

But it also has the potential disadvantage that the flux surfaces will be highly sensitive to the magnetic axis position (R0, Z0).

We should add this alternative discretization as an option to test if it is ever superior to the existing approach.

Calculating MHD Stability

We have the jacobian, the eigenvalues of that should be related to stability

    project back to get R_tt and Z_tt
    incompressibility constraint

Cache computed values of Configuration for plotting/analysis

When making plots we end up needing to calculate the same underlying quantities several times, so it would be efficient to compute them once and cache them.

See https://docs.python.org/3/library/functools.html#functools.lru_cache
or https://pypi.org/project/cached-property/

We'd need to break up some of the functions into smaller chunks to make things more efficient.

Note that we likely don't want to cache the actual compute functions that get called by the optimizer, so we might want to have separate wrapper functions that do the caching.

IE:
def fun(x):
` do something

@lru_cache
def fun_cached(x):
return fun(x)

Vectorize zernike polynomials over mode idx

currently they're vectorized over spatial pts but not over modes, so the modes are calculated in a python loop, which is slow. The azimuthal and toroidal terms are properly vectorized and are ~10x faster

Use iterative method for trust region solve

Currently we're using a direct method on the dense jacobian, which is super expensive at high res

Also: look into making jacobian a LinearOperator using jvp/vjp to avoid having to compute dense jacobian.
Currently its computing a dense matrix with finite differences at each step -> super innefficient
Could also look into some object that does rank 1 updates and occasional full recompute?

Set up automated testing

We really need at least some basic tests for core functionality that we can automate with pytest and travis. This will also help with detecting issues as jax changes, by running tests against multiple versions of all the dependencies: https://docs.travis-ci.com/user/languages/python/#dependency-management

Here's a basic template: https://docs.python.org/3/library/unittest.html
and an example: https://github.com/f0uriest/GASTOp/blob/master/tests/test_fitness.py

Comments in input file not parsed if at the end of a line

If the line begins with # or !, the line is correctly parsed as a comment. But if the comment is on the same line as other stuff, it seems to get parsed as if it was all valid input
eg, M = 6 #6, 8, 10 should ideally be parsed as just M=6, but it currently gets parsed as M=6,6,8,10
It would be good to have it parsed the other way, both so one can quickly comment out part of a line without erasing it, and so we can have comments at the end of a line explaining what the different options do

Enforcing stellarator symmetry

Current method is a kludge that doesn't actually reduce the problem size.
Some things to consider:
which nodes are symmetric?
do we need different ZernikeTransforms for R,Z? or one that does both?
how to deal with evaluating different errors at different locations. Do we need to call every function twice with different args?

What is the best solver?

Right now we use a least squares approach using the trust region method (levenberg-marquadt also works well). Scipy also has a bunch of other routines we should test out:
Unconstrained:
BFGS
Nelder Mead
CG / Newton-CG

Constrained (ie explicitly enforce BC rather than with penalty method):
SLSQP
trust-region constrained

Add option for symmetry nodes

helps reduce size of the system

Also, we should think about how we're currently oversampling. If we oversample in both M and N, we've sort of double counting, so we end up with oversampling by 3x instead of 2x. Could be an improvement

Support plotting of arbitrary expressions

eg num/den or term1*term2+term3

• change parse_name to handle multiple terms
• change _compute function to reproduce the arithmetic operations
• figure out way to format name in a pretty way for display
• possibly automatically cancelling or combining units

API Documentation Does Not List anything

Per title, after the last pull Request, the API Documentation only lists headers and not any of the function, but it did before the last pull request was merged I think.

Quasisymmetry function specifying M/N

@ddudt i dont know what this means

Issue with initial guess for W7X

This is the initial guess generated for the w7x test case. Couldn't find any obvious reason, but didn't look super close. Not sure if its an issue with the input file or the initialization function, @ddudt can you take a look?

plotting to specified axes

All the plotting routines should have an option to plot to a preallocated figure/axis, and should also return the fig/axis that was used, whether an existing one was used or a new one created.

Use seperate "load' method instead of keyword in standard init

Right now all the IOable classes have a keyword arg in init "load_from", which is used to instantiate from a saved file.
However, this means that all of the other args need to have default values, even ones which should be mandatory, so its a bit confusing.
Ideally we would have a regular init method, where mandatory args must always be specified, and others have some default value.
We'd then have a separate method for loading from a saved file

eg:

class foo():

def __init__(self, arg1, arg2, arg3="default"):
    self.arg1 = arg1
    self.arg2 = arg2
    self.arg3 = arg3
    
@classmethod
def load(cls, load_from):
    # create a blank object bypassing init
    self = cls.__new__(cls)
    
    # open the file and set attribute like:
    # for attr in file:
    #     setattr(self, attr)
    
    # to set other secondary stuff that wasnt saved possibly:
    self._after_init()

    return self

Support JAX 0.2.5

they changed some stuff in the latest version, so our code needs to change as well. Did some testing and looks like the major stuff is how we're using jnp.where and put to do logical indexing and stuff. Updating that should also fix NaNs in reverse mode autodiff when including a point at the axis.
Looks like there's also some stuff with the datatypes for static arguments, still trying to figure that out.

Bug in intial guess for non-symmetric cases

the initial guess for ATF looks like this:

However, if we simply change the input file from sym=1 to sym=0, we then get this:

Calculate total plasma energy

For comparison with VMEC, and as a possible optimization target (might be cheaper to evaluate the force error)

Interactive plotting

Right now we only plot stuff at the end of the program. I might be nice to have an option to plot the latest results after each outer iteration, or to plot the initial guess before starting, etc. We could make the plot flag have a value, like plot=2 with different levels, like with verbose

By default matplotlib blocks execution when plotting (ie, waiting for the user to close the plot before continuing the program), which we might want or might not. There are a few different options for handling that, but I think we probably could make it so it doesn't block and it will just display plots when ready and keep on calculating

Auto-size jacobian blocks based on available memory

For large jacobian matrices that can't be computed in memory, we have a way to break them up into blocks. However, the block size must be defined by the user. It would be better to do some profiling to find the biggest matrices we can fit in a given chunk of memory and then have it automatically determine the block size based on available memory

support JAX version 0.2.0

JAX version 0.1.77 works, but not version 0.2.0

Specify initial guess / restarting

We should have an easy way to either specify an initial guess, or to restart from a previous run at a certain point.