Merge pull request #209 from hiddenSymmetries/ml/virtual_casing

Virtual casing

.github/workflows/ci.yml

@@ -93,6 +93,9 @@ jobs:
     - name: Install booz_xform
       run: pip install -v git+https://github.com/hiddenSymmetries/booz_xform
 
+    - name: Install virtual_casing
+      run: pip install -v git+https://github.com/hiddenSymmetries/virtual-casing
+
     # Checking out SPEC is a tricky because it is a private repository.
     # See https://github.community/t/best-way-to-clone-a-private-repo-during-script-run-of-private-github-action/16116/7
     # https://stackoverflow.com/questions/57612428/cloning-private-github-repository-within-organisation-in-actions

.github/workflows/extensive_ci.yml

@@ -104,6 +104,9 @@ jobs:
       if: contains(matrix.packages, 'vmec') || contains(matrix.packages, 'all')
       run: pip install -v git+https://github.com/hiddenSymmetries/booz_xform
 
+    - name: Install virtual_casing
+      run: pip install -v git+https://github.com/hiddenSymmetries/virtual-casing
+
     # Checking out SPEC is a tricky because it is a private repository.
     # See https://github.community/t/best-way-to-clone-a-private-repo-during-script-run-of-private-github-action/16116/7
     # https://stackoverflow.com/questions/57612428/cloning-private-github-repository-within-organisation-in-actions

ci/Dockerfile.ubuntu

@@ -59,6 +59,7 @@ RUN /venv/bin/pip install h5py pyoculus py_spec
 RUN /venv/bin/pip install vtk==9.0.1 pyqt5 matplotlib pyevtk plotly
 RUN /venv/bin/pip install mayavi
 RUN /venv/bin/pip install  git+https://github.com/hiddenSymmetries/booz_xform
+RUN /venv/bin/pip install  git+https://github.com/hiddenSymmetries/virtual-casing
 
 ENV CI=True
 RUN git clone --recurse-submodules https://github.com/hiddenSymmetries/simsopt.git /src/simsopt && \

ci/singularity.def

@@ -74,6 +74,7 @@ Stage: devel
     /venv/bin/pip install vtk==9.0.1 pyqt5 matplotlib pyevtk plotly
     /venv/bin/pip install mayavi
     /venv/bin/pip install  git+https://github.com/hiddenSymmetries/booz_xform
+    /venv/bin/pip install  git+https://github.com/hiddenSymmetries/virtual-casing
 
     CI=True
     git clone --recurse-submodules https://github.com/hiddenSymmetries/simsopt.git simsopt && \

docs/source/index.rst

@@ -45,6 +45,8 @@ repositories. These separate modules include
   equilibrium.
 - `booz_xform <https://hiddensymmetries.github.io/booz_xform/>`_, for
   Boozer coordinates and quasisymmetry.
+- `virtual_casing <https://github.com/hiddenSymmetries/virtual_casing>`_,
+  needed for coil optimization in the case of finite-beta plasmas.
 
 We gratefully acknowledge funding from the `Simons Foundation's Hidden
 symmetries and fusion energy project

docs/source/simsopt.mhd.rst

@@ -36,6 +36,14 @@ simsopt.mhd.spec module
    :undoc-members:
    :show-inheritance:
 
+simsopt.mhd.virtual\_casing module
+----------------------------------
+
+.. automodule:: simsopt.mhd.virtual_casing
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 simsopt.mhd.vmec module
 -----------------------
 

docs/source/simsopt.util.rst

@@ -20,6 +20,14 @@ simsopt.util.dev module
    :undoc-members:
    :show-inheritance:
 
+simsopt.util.fourier_interpolation module
+-----------------------------------------
+
+.. automodule:: simsopt.util.fourier_interpolation
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 simsopt.util.log module
 -----------------------
 

examples/2_Intermediate/B_external_normal.py

@@ -0,0 +1,52 @@
+#!/usr/bin/env python
+
+import os
+from simsopt.mhd.virtual_casing import VirtualCasing
+
+"""
+This example illustrates the virtual casing calculation, in which
+we compute the contribution to the magnetic field on the plasma
+boundary due to current inside the boundary, excluding currents
+outside the plasma. Typically this calculation is run once at the end
+of a stage-1 optimization (optimizing the plasma shape) before
+starting the stage-2 optimization (optimizing the coil shapes). It is
+only required for finite-beta plasmas, not for vacuum fields.
+
+This example requires that the python virtual_casing module is installed.
+"""
+
+filename = os.path.join(os.path.dirname(__file__), '..', '..',
+                        'tests', 'test_files', 'wout_li383_low_res_reference.nc')
+
+# The "trgt_" resolution is the resolution on which results are
+# computed, which you will then use for the stage-2 coil
+# optimization. The "src_" resolution is used internally for the
+# virtual casing calculation, and is often higher than the "trgt_"
+# resolution for better accuracy.  Typically "src_ntheta" is not
+# specified; in this case it is computed automatically from "src_nphi"
+# to minimize anisotropy of the grid.
+vc = VirtualCasing.from_vmec(filename, src_nphi=48, trgt_ntheta=32, trgt_nphi=32)
+print('automatically determined src_ntheta:', vc.src_ntheta)
+
+# The VirtualCasing.from_vmec command writes a file
+# simsopt/tests/test_files/vcasing_li383_low_res_reference.nc
+# containing the results of the virtual casing calculation.
+
+# You can generate a matplotlib plot of B_external_normal on the
+# boundary surface by uncommenting the following line:
+
+# vc.plot()
+
+# B_external_normal is now available as an attribute:
+print('B_external_normal:')
+print(vc.B_external_normal[:4, :4])  # Just print the first few elements
+
+# The saved virtual casing results can be loaded in later for stage-2
+# coil optimization, as follows:
+directory, wout_file = os.path.split(filename)
+vcasing_file = os.path.join(directory, wout_file.replace('wout', 'vcasing'))
+vc2 = VirtualCasing.load(vcasing_file)
+
+print('B_external_normal, loaded from file:')
+print(vc2.B_external_normal[:4, :4])
+

examples/2_Intermediate/stage_two_optimization_finite_beta.py

@@ -0,0 +1,175 @@
+#!/usr/bin/env python
+r"""
+In this example we solve a stage-II coil optimisation problem: the
+goal is to find coils that generate a specific target normal field on
+a given surface.  The target equilibrium is a W7X configuration with
+average beta of 4%. Since it is not a vacuum field, the target
+B_{External}·n is nonzero. A virtual casing calculation is used to
+compute this target B_{External}·n.
+
+The objective is given by
+
+    J = (1/2) ∫ |B_{BiotSavart}·n - B_{External}·n|^2 ds
+        + LENGTH_PENALTY * Σ ½(CurveLength - L0)^2
+"""
+
+import os
+from pathlib import Path
+import numpy as np
+from scipy.optimize import minimize
+from simsopt.mhd.vmec import Vmec
+from simsopt.geo.surfacerzfourier import SurfaceRZFourier
+from simsopt.objectives.fluxobjective import SquaredFlux
+from simsopt.objectives.utilities import QuadraticPenalty
+from simsopt.geo.curve import curves_to_vtk, create_equally_spaced_curves
+from simsopt.field.biotsavart import BiotSavart
+from simsopt.field.coil import Current, coils_via_symmetries, ScaledCurrent
+from simsopt.geo.curveobjectives import CurveLength
+from simsopt.mhd.virtual_casing import VirtualCasing
+
+
+# Number of unique coil shapes, i.e. the number of coils per half field period:
+# (Since the configuration has nfp = 5 and stellarator symmetry, multiply ncoils by 5 * 2 to get the total number of coils.)
+ncoils = 5
+
+# Major radius for the initial circular coils:
+R0 = 5.5
+
+# Minor radius for the initial circular coils:
+R1 = 1.25
+
+# Number of Fourier modes describing each Cartesian component of each coil:
+order = 6
+
+# Weight on the curve length penalties in the objective function:
+LENGTH_PENALTY = 1e0
+
+# Number of iterations to perform:
+ci = "CI" in os.environ and os.environ['CI'].lower() in ['1', 'true']
+MAXITER = 50 if ci else 500
+
+# File for the desired boundary magnetic surface:
+TEST_DIR = (Path(__file__).parent / ".." / ".." / "tests" / "test_files").resolve()
+filename = 'wout_W7-X_without_coil_ripple_beta0p05_d23p4_tm_reference.nc'
+vmec_file = TEST_DIR / filename
+
+# Resolution on the plasma boundary surface:
+# nphi is the number of grid points in 1/2 a field period.
+nphi = 32
+ntheta = 32
+
+# Resolution for the virtual casing calculation:
+vc_src_nphi = 80
+# (For the virtual casing src_ resolution, only nphi needs to be
+# specified; the theta resolution is computed automatically to
+# minimize anisotropy of the grid.)
+
+#######################################################
+# End of input parameters.
+#######################################################
+
+# Directory for output
+OUT_DIR = "./output/"
+os.makedirs(OUT_DIR, exist_ok=True)
+
+# Once the virtual casing calculation has been run once, the results
+# can be used for many coil optimizations. Therefore here we check to
+# see if the virtual casing output file alreadys exists. If so, load
+# the results, otherwise run the virtual casing calculation and save
+# the results.
+head, tail = os.path.split(vmec_file)
+vc_filename = os.path.join(head, tail.replace('wout', 'vcasing'))
+print('virtual casing data file:', vc_filename)
+if os.path.isfile(vc_filename):
+    print('Loading saved virtual casing result')
+    vc = VirtualCasing.load(vc_filename)
+else:
+    # Virtual casing must not have been run yet.
+    print('Running the virtual casing calculation')
+    vc = VirtualCasing.from_vmec(vmec_file, src_nphi=vc_src_nphi, trgt_nphi=nphi, trgt_ntheta=ntheta)
+
+# Initialize the boundary magnetic surface:
+s = SurfaceRZFourier.from_wout(vmec_file, range="half period", nphi=nphi, ntheta=ntheta)
+total_current = Vmec(vmec_file).external_current() / (2 * s.nfp)
+
+# Create the initial coils:
+base_curves = create_equally_spaced_curves(ncoils, s.nfp, stellsym=True, R0=R0, R1=R1, order=order, numquadpoints=128)
+# Since we know the total sum of currents, we only optimize for ncoils-1
+# currents, and then pick the last one so that they all add up to the correct
+# value.
+base_currents = [ScaledCurrent(Current(total_current / ncoils * 1e-5), 1e5) for _ in range(ncoils-1)]
+# Above, the factors of 1e-5 and 1e5 are included so the current
+# degrees of freedom are O(1) rather than ~ MA.  The optimization
+# algorithm may not perform well if the dofs are scaled badly.
+total_current = Current(total_current)
+total_current.fix_all()
+base_currents += [total_current - sum(base_currents)]
+
+coils = coils_via_symmetries(base_curves, base_currents, s.nfp, True)
+bs = BiotSavart(coils)
+bs.set_points(s.gamma().reshape((-1, 3)))
+curves = [c.curve for c in coils]
+curves_to_vtk(curves, OUT_DIR + "curves_init")
+pointData = {"B_N": np.sum(bs.B().reshape((nphi, ntheta, 3)) * s.unitnormal(), axis=2)[:, :, None]}
+s.to_vtk(OUT_DIR + "surf_init", extra_data=pointData)
+
+# Define the objective function:
+Jf = SquaredFlux(s, bs, target=vc.B_external_normal)
+Jls = [CurveLength(c) for c in base_curves]
+
+# Form the total objective function. To do this, we can exploit the
+# fact that Optimizable objects with J() and dJ() functions can be
+# multiplied by scalars and added:
+JF = Jf \
+    + LENGTH_PENALTY * sum(QuadraticPenalty(Jls[i], Jls[i].J()) for i in range(len(base_curves)))
+
+# We don't have a general interface in SIMSOPT for optimisation problems that
+# are not in least-squares form, so we write a little wrapper function that we
+# pass directly to scipy.optimize.minimize
+
+
+def fun(dofs):
+    JF.x = dofs
+    J = JF.J()
+    grad = JF.dJ()
+    jf = Jf.J()
+    Bbs = bs.B().reshape((nphi, ntheta, 3))
+    BdotN = np.abs(np.sum(Bbs * s.unitnormal(), axis=2) - vc.B_external_normal) / np.linalg.norm(Bbs, axis=2)
+    BdotN_mean = np.mean(BdotN)
+    BdotN_max = np.max(BdotN)
+    outstr = f"J={J:.1e}, Jf={jf:.1e}, ⟨|B·n|⟩={BdotN_mean:.1e}, max(|B·n|)={BdotN_max:.1e}"
+    cl_string = ", ".join([f"{J.J():.1f}" for J in Jls])
+    outstr += f", Len=sum([{cl_string}])={sum(J.J() for J in Jls):.1f}"
+    outstr += f", ║∇J║={np.linalg.norm(grad):.1e}"
+    print(outstr)
+    return 1e-4*J, 1e-4*grad
+
+
+print("""
+################################################################################
+### Perform a Taylor test ######################################################
+################################################################################
+""")
+f = fun
+dofs = JF.x
+np.random.seed(1)
+h = np.random.uniform(size=dofs.shape)
+J0, dJ0 = f(dofs)
+dJh = sum(dJ0 * h)
+for eps in [1e-2, 1e-3, 1e-4, 1e-5, 1e-6, 1e-7]:
+    J1, _ = f(dofs + eps*h)
+    J2, _ = f(dofs - eps*h)
+    print("err", (J1-J2)/(2*eps) - dJh)
+
+print("""
+################################################################################
+### Run the optimisation #######################################################
+################################################################################
+""")
+res = minimize(fun, dofs, jac=True, method='L-BFGS-B', options={'maxiter': MAXITER, 'maxcor': 300, 'ftol': 1e-20, 'gtol': 1e-20}, tol=1e-20)
+dofs = res.x
+curves_to_vtk(curves, OUT_DIR + "curves_opt")
+Bbs = bs.B().reshape((nphi, ntheta, 3))
+BdotN = np.abs(np.sum(Bbs * s.unitnormal(), axis=2) - vc.B_external_normal) / np.linalg.norm(Bbs, axis=2)
+pointData = {"B_N": BdotN[:, :, None]}
+s.to_vtk(OUT_DIR + "surf_opt", extra_data=pointData)

examples/run_serial_examples

@@ -14,4 +14,5 @@ set -ex
 ./2_Intermediate/boozer.py
 ./2_Intermediate/stage_two_optimization.py
 ./2_Intermediate/stage_two_optimization_stochastic.py
+./2_Intermediate/stage_two_optimization_finite_beta.py
 ./3_Advanced/stage_two_optimization_finitebuild.py

examples/run_vmec_examples

@@ -16,6 +16,8 @@ set -ex
 MPI_OPTIONS=${GITHUB_ACTIONS:+--oversubscribe}
 echo MPI_OPTIONS=$MPI_OPTIONS
 
+./2_Intermediate/B_external_normal.py
+
 ./stellarator_benchmarks/1DOF_circularCrossSection_varyAxis_targetIota.py
 mpiexec $MPI_OPTIONS -n 2 ./stellarator_benchmarks/1DOF_circularCrossSection_varyAxis_targetIota.py
 mpiexec $MPI_OPTIONS -n 3 ./stellarator_benchmarks/1DOF_circularCrossSection_varyAxis_targetIota.py