Tutorial 3: SIMPLE inside an optimisation loop

This tutorial wraps a SIMPLE loss computation into a Python function that takes a VMEC wout file and returns the loss fraction with its error, ready to be called by an optimiser after each equilibrium solve. It shows how to make the objective reproducible with common random numbers, when a change is significant, and how orbit classification serves as a cheaper proxy.

Prerequisites: pysimple from setup-venv.sh (Installation) and the equilibria of Tutorial 1: Loss fraction from a VMEC equilibrium in tutorial1/qa and tutorial1/qh. Runtime: about one minute.

The objective function

docs/tutorials/scripts/loss_objective.py
"""Alpha-particle loss fraction as an objective function with pysimple.

Usage: python loss_objective.py WOUT_QA WOUT_QH
"""
import sys
import time

import numpy as np
import pysimple

TRACE_TIME = 1e-2  # s
N_PARTICLES = 128
S_START = 0.6


def loss_fraction(wout, **settings):
    """Return the loss fraction and its standard error for one equilibrium.

    ``settings`` are further &config keys. The same seed gives the same start
    points for every call (common random numbers).
    """
    pysimple.init(wout, isw_field_type=2, trace_time=TRACE_TIME,
                  deterministic=True, **settings)
    start = pysimple.sample_surface(N_PARTICLES, s=S_START)
    assert np.allclose(start[0], S_START), "start points are not on the surface"
    result = pysimple.trace_parallel(start)
    t = result["loss_times"]
    lost = np.isfinite(t) & (t > 0) & (t < TRACE_TIME)
    n_resolved = np.isfinite(t).sum()
    if n_resolved == 0:
        raise RuntimeError("no numerically resolved particles")
    f = lost.sum() / n_resolved
    return f, np.sqrt(f * (1 - f) / n_resolved)


if __name__ == "__main__":
    wout_qa, wout_qh = sys.argv[1], sys.argv[2]
    for scale in (0.5, 0.75, 1.0):
        t0 = time.perf_counter()
        f, err = loss_fraction(wout_qa, vmec_B_scale=scale)
        print(f"QA, B x {scale:4.2f}: loss fraction {f:.3f} +- {err:.3f}"
              f"  ({time.perf_counter() - t0:.0f} s)")
    f, err = loss_fraction(wout_qh)
    print(f"QH, B x 1.00: loss fraction {f:.3f} +- {err:.3f}")

pysimple.init resets all &config keys to their defaults, applies the keyword arguments and loads the equilibrium given as first argument (used as netcdffile, field_input and coord_input). One Python process can therefore evaluate any number of candidates.

trace_parallel traces all markers with OpenMP and returns per marker the loss time, trace_time if confined, NaN if numerically unresolved and -1 if skipped as strongly passing. The function counts unresolved orbits neither as lost nor as confined.

S_START = 0.6 differs from the default sbeg = 0.5 on purpose (limitation in Python interface); keep the two different when changing settings.

Common random numbers

sample_surface places the markers on one field line of the surface S_START. With deterministic=True the generator is reseeded in every init, so two calls on the same equilibrium use the same start points and two similar equilibria nearly the same ones. An optimiser compares objective values; for two candidates with loss indicators \(X_i, Y_i \in \{0,1\}\) on the same markers,

\[\operatorname{Var}(\hat f_X - \hat f_Y) = \frac{\operatorname{Var}X + \operatorname{Var}Y - 2\operatorname{Cov}(X,Y)}{N},\]

which is smaller than for independent samples whenever the outcomes are positively correlated. Estimate it from the paired per-marker outcomes, not by adding the two single-run errors.

Run

source .venv/bin/activate
mkdir -p tutorial3
cd tutorial3
python ../docs/tutorials/scripts/loss_objective.py \
    ../tutorial1/qa/wout.nc ../tutorial1/qh/wout.nc

After the SIMPLE log the script prints:

QA, B x 0.50: loss fraction 0.344 +- 0.042  (11 s)
QA, B x 0.75: loss fraction 0.242 +- 0.038  (12 s)
QA, B x 1.00: loss fraction 0.141 +- 0.031  (12 s)
QH, B x 1.00: loss fraction 0.000 +- 0.000

The first three lines scan the field strength through vmec_B_scale, which multiplies the VMEC field: a weaker field means wider orbits and more losses. A few seconds of each evaluation are field setup. For the zero-loss QH line the printed error is not an upper bound; use \(3/N\) (Tutorial 1: Loss fraction from a VMEC equilibrium).

In an optimisation:

  • keep the seed (deterministic=True) and N fixed across evaluations;

  • increase N as the optimisation converges and differences shrink;

  • do not trust differences below the single-evaluation error unless repeated with another ran_seed;

  • use trace_time = 1d-2 (cost scales with it; most losses of good configurations occur in the first milliseconds) and longer runs for the final candidates.

A cheaper proxy: orbit classification

The classifiers of Albert et al. (2023) decide after a few bounce periods whether a trapped orbit is regular or chaotic (Orbit classification). classify_fast stops after nturns recurrences of the banana tip. Regular motion implies confinement only when the invariant torus lies inside the loss boundary.

docs/tutorials/scripts/classify_proxy.py
"""Fast regular/chaotic classification of trapped orbits with pysimple.

Usage: python classify_proxy.py WOUT
"""
import sys
import time

import numpy as np
import pysimple

pysimple.init(sys.argv[1], isw_field_type=2, trace_time=1e-2, deterministic=True)
start = pysimple.sample_surface(128, s=0.6)
t0 = time.perf_counter()
c = pysimple.classify_fast(start, nturns=8)
trapped = ~c["passing"]
print(f"classified in {time.perf_counter() - t0:.0f} s")
print(f"trapped: {trapped.sum()}, J_par regular: {(c['jpar'][trapped] == 1).sum()}, "
      f"chaotic: {(c['jpar'][trapped] == 2).sum()}, "
      f"unclassified: {(c['jpar'][trapped] == 0).sum()}")
print(f"lost during classification: {c['lost'].sum()}")

python ../docs/tutorials/scripts/classify_proxy.py ../tutorial1/qa/wout.nc

classified in 4 s
trapped: 43, J_par regular: 2, chaotic: 17, unclassified: 24
lost during classification: 18

The fraction of chaotic trapped orbits is a loss metric at a fraction of the tracing cost. Orbits lost before nturns recurrences stay unclassified.

Embedding in an optimiser

Call loss_fraction with the wout file written by the equilibrium solver and add the result with a weight to the objective.

  • The objective is noisy and not differentiable in the equilibrium. Use derivative-free optimisers, or smooth proxies (for example quasi-symmetry residuals) for the main optimisation with the loss fraction as a constraint or modestly weighted term.

  • SIMPLE uses all OpenMP threads of the process; set OMP_NUM_THREADS when the optimiser runs evaluations in parallel.

  • A Fortran stop inside the library ends the Python process. For long unattended runs, evaluate each candidate in a separate process (multiprocessing or subprocess).

For large marker numbers, the CUDA backend computes the segmented loss objective of Landreman et al. on a GPU (GPU tracing).