Merge pull request #358 from AlexWiedman/planar

Planar

Author: landreman

CMakeLists.txt

@@ -123,7 +123,7 @@ pybind11_add_module(${PROJECT_NAME}
     src/simsoptpp/biot_savart_py.cpp
     src/simsoptpp/biot_savart_vjp_py.cpp
     src/simsoptpp/regular_grid_interpolant_3d_py.cpp
-    src/simsoptpp/curve.cpp src/simsoptpp/curverzfourier.cpp src/simsoptpp/curvexyzfourier.cpp
+    src/simsoptpp/curve.cpp src/simsoptpp/curverzfourier.cpp src/simsoptpp/curvexyzfourier.cpp src/simsoptpp/curveplanarfourier.cpp
     src/simsoptpp/surface.cpp src/simsoptpp/surfacerzfourier.cpp src/simsoptpp/surfacexyzfourier.cpp
     src/simsoptpp/integral_BdotN.cpp
     src/simsoptpp/dipole_field.cpp src/simsoptpp/permanent_magnet_optimization.cpp

docs/source/simsopt.geo.rst

@@ -60,6 +60,14 @@ simsopt.geo.curverzfourier module
    :undoc-members:
    :show-inheritance:
 
+simsopt.geo.curveplanarfourier module
+-------------------------------------
+
+.. automodule:: simsopt.geo.curveplanarfourier
+   :members:
+   :undoc-members:
+   :show-inheritance:
+
 simsopt.geo.curvexyzfourier module
 ----------------------------------
 

examples/2_Intermediate/stage_two_optimization_planar_coils.py

@@ -0,0 +1,199 @@
+#!/usr/bin/env python
+r"""
+This coil optimization script is similar to stage_two_optimization.py. However
+in this version, the coils are constrained to be planar, by using the curve type
+CurvePlanarFourier. Also the LinkingNumber objective is used to prevent coils
+from becoming topologically linked with each other.
+
+In this example we solve a FOCUS like Stage II coil optimisation problem: the
+goal is to find coils that generate a specific target normal field on a given
+surface.  In this particular case we consider a vacuum field, so the target is
+just zero.
+
+The objective is given by
+
+    J = (1/2) \int |B dot n|^2 ds
+        + LENGTH_WEIGHT * (sum CurveLength)
+        + DISTANCE_WEIGHT * MininumDistancePenalty(DISTANCE_THRESHOLD)
+        + CURVATURE_WEIGHT * CurvaturePenalty(CURVATURE_THRESHOLD)
+        + MSC_WEIGHT * MeanSquaredCurvaturePenalty(MSC_THRESHOLD)
+        + LinkingNumber
+
+if any of the weights are increased, or the thresholds are tightened, the coils
+are more regular and better separated, but the target normal field may not be
+achieved as well. This example demonstrates the adjustment of weights and
+penalties via the use of the `Weight` class.
+
+The target equilibrium is the QA configuration of arXiv:2108.03711.
+"""
+
+import os
+from pathlib import Path
+import numpy as np
+from scipy.optimize import minimize
+from simsopt.field import BiotSavart, Current, coils_via_symmetries
+from simsopt.geo import (
+    CurveLength, CurveCurveDistance,
+    MeanSquaredCurvature, LpCurveCurvature, CurveSurfaceDistance, LinkingNumber,
+    SurfaceRZFourier, curves_to_vtk, create_equally_spaced_planar_curves,
+)
+from simsopt.objectives import Weight, SquaredFlux, QuadraticPenalty
+from simsopt.util import in_github_actions
+
+# Number of unique coil shapes, i.e. the number of coils per half field period:
+# (Since the configuration has nfp = 2, multiply by 4 to get the total number of coils.)
+ncoils = 4
+
+# Major radius for the initial circular coils:
+R0 = 1.0
+
+# Minor radius for the initial circular coils:
+R1 = 0.5
+
+# Number of Fourier modes describing each Cartesian component of each coil:
+order = 5
+
+# Weight on the curve lengths in the objective function. We use the `Weight`
+# class here to later easily adjust the scalar value and rerun the optimization
+# without having to rebuild the objective.
+LENGTH_WEIGHT = Weight(10)
+
+# Threshold and weight for the coil-to-coil distance penalty in the objective function:
+CC_THRESHOLD = 0.08
+CC_WEIGHT = 1000
+
+# Threshold and weight for the coil-to-surface distance penalty in the objective function:
+CS_THRESHOLD = 0.12
+CS_WEIGHT = 10
+
+# Threshold and weight for the curvature penalty in the objective function:
+CURVATURE_THRESHOLD = 10.
+CURVATURE_WEIGHT = 1e-6
+
+# Threshold and weight for the mean squared curvature penalty in the objective function:
+MSC_THRESHOLD = 10
+MSC_WEIGHT = 1e-6
+
+# Number of iterations to perform:
+MAXITER = 50 if in_github_actions else 400
+
+# File for the desired boundary magnetic surface:
+TEST_DIR = (Path(__file__).parent / ".." / ".." / "tests" / "test_files").resolve()
+filename = TEST_DIR / 'input.LandremanPaul2021_QA'
+
+# Directory for output
+OUT_DIR = "./output/"
+os.makedirs(OUT_DIR, exist_ok=True)
+
+#######################################################
+# End of input parameters.
+#######################################################
+
+# Initialize the boundary magnetic surface:
+nphi = 32
+ntheta = 32
+s = SurfaceRZFourier.from_vmec_input(filename, range="half period", nphi=nphi, ntheta=ntheta)
+
+# Create the initial coils:
+base_curves = create_equally_spaced_planar_curves(ncoils, s.nfp, stellsym=True, R0=R0, R1=R1, order=order)
+base_currents = [Current(1e5) for i in range(ncoils)]
+# Since the target field is zero, one possible solution is just to set all
+# currents to 0. To avoid the minimizer finding that solution, we fix one
+# of the currents:
+base_currents[0].fix_all()
+
+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 individual terms objective function:
+Jf = SquaredFlux(s, bs)
+Jls = [CurveLength(c) for c in base_curves]
+Jccdist = CurveCurveDistance(curves, CC_THRESHOLD, num_basecurves=ncoils)
+Jcsdist = CurveSurfaceDistance(curves, s, CS_THRESHOLD)
+Jcs = [LpCurveCurvature(c, 2, CURVATURE_THRESHOLD) for c in base_curves]
+Jmscs = [MeanSquaredCurvature(c) for c in base_curves]
+linkNum = LinkingNumber(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_WEIGHT * QuadraticPenalty(sum(Jls), 2.6*ncoils) \
+    + CC_WEIGHT * Jccdist \
+    + CS_WEIGHT * Jcsdist \
+    + CURVATURE_WEIGHT * sum(Jcs) \
+    + MSC_WEIGHT * sum(QuadraticPenalty(J, MSC_THRESHOLD) for J in Jmscs) \
+    + linkNum
+
+# 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()
+    BdotN = np.mean(np.abs(np.sum(bs.B().reshape((nphi, ntheta, 3)) * s.unitnormal(), axis=2)))
+    MaxBdotN = np.max(np.abs(np.sum(bs.B().reshape((nphi, ntheta, 3)) * s.unitnormal(), axis=2)))
+    mean_AbsB = np.mean(bs.AbsB())
+    outstr = f"J={J:.1e}, Jf={jf:.1e}, ⟨B·n⟩={BdotN:.1e}"
+    cl_string = ", ".join([f"{J.J():.1f}" for J in Jls])
+    kap_string = ", ".join(f"{np.max(c.kappa()):.1f}" for c in base_curves)
+    msc_string = ", ".join(f"{J.J():.1f}" for J in Jmscs)
+    outstr += f", Len=sum([{cl_string}])={sum(J.J() for J in Jls):.1f}, ϰ=[{kap_string}], ∫ϰ²/L=[{msc_string}]"
+    outstr += f", C-C-Sep={Jccdist.shortest_distance():.2f}, C-S-Sep={Jcsdist.shortest_distance():.2f}"
+    outstr += f", ║∇J║={np.linalg.norm(grad):.1e}"
+    outstr += f", ⟨B·n⟩/|B|={BdotN/mean_AbsB:.1e}"
+    outstr += f", (Max B·n)/|B|={MaxBdotN/mean_AbsB:.1e}"
+    outstr += f", Link Number = {linkNum.J()}"
+    print(outstr)
+    return J, 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-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}, tol=1e-15)
+curves_to_vtk(curves, OUT_DIR + "curves_opt_short")
+pointData = {"B_N": np.sum(bs.B().reshape((nphi, ntheta, 3)) * s.unitnormal(), axis=2)[:, :, None]}
+s.to_vtk(OUT_DIR + "surf_opt_short", extra_data=pointData)
+
+
+# We now use the result from the optimization as the initial guess for a
+# subsequent optimization with reduced penalty for the coil length. This will
+# result in slightly longer coils but smaller `B·n` on the surface.
+dofs = res.x
+LENGTH_WEIGHT *= 0.1
+res = minimize(fun, dofs, jac=True, method='L-BFGS-B', options={'maxiter': MAXITER, 'maxcor': 300}, tol=1e-15)
+curves_to_vtk(curves, OUT_DIR + "curves_opt_long")
+pointData = {"B_N": np.sum(bs.B().reshape((nphi, ntheta, 3)) * s.unitnormal(), axis=2)[:, :, None]}
+s.to_vtk(OUT_DIR + "surf_opt_long", extra_data=pointData)
+
+# Save the optimized coil shapes and currents so they can be loaded into other scripts for analysis:
+bs.save(OUT_DIR + "biot_savart_opt.json")

examples/run_serial_examples

@@ -16,6 +16,7 @@ set -ex
 ./2_Intermediate/QSC.py
 ./2_Intermediate/boozer.py
 ./2_Intermediate/stage_two_optimization.py
+./2_Intermediate/stage_two_optimization_planar_coils.py
 ./2_Intermediate/stage_two_optimization_stochastic.py
 ./2_Intermediate/stage_two_optimization_finite_beta.py
 ./2_Intermediate/strain_optimization.py

src/simsopt/geo/__init__.py

@@ -8,6 +8,7 @@
 from .curvexyzfourier import *
 from .curveperturbed import *
 from .curveobjectives import *
+from .curveplanarfourier import *
 from .framedcurve import *
 from .finitebuild import *
 from .plotting import *
@@ -28,6 +29,7 @@
 __all__ = (curve.__all__ + curvehelical.__all__ +
            curverzfourier.__all__ + curvexyzfourier.__all__ +
            curveperturbed.__all__ + curveobjectives.__all__ +
+           curveplanarfourier.__all__ +
            finitebuild.__all__ + plotting.__all__ +
            boozersurface.__all__ + qfmsurface.__all__ +
            surface.__all__ +

src/simsopt/geo/curve.py

@@ -11,7 +11,7 @@
 from .jit import jit
 from .plotting import fix_matplotlib_3d
 
-__all__ = ['Curve', 'RotatedCurve', 'curves_to_vtk', 'create_equally_spaced_curves']
+__all__ = ['Curve', 'RotatedCurve', 'curves_to_vtk', 'create_equally_spaced_curves', 'create_equally_spaced_planar_curves']
 
 
 @jit
@@ -878,3 +878,46 @@ def create_equally_spaced_curves(ncurves, nfp, stellsym, R0=1.0, R1=0.5, order=6
         curve.x = curve.x  # need to do this to transfer data to C++
         curves.append(curve)
     return curves
+
+
+def create_equally_spaced_planar_curves(ncurves, nfp, stellsym, R0=1.0, R1=0.5, order=6, numquadpoints=None):
+    """
+    Create ``ncurves`` curves of type
+    :obj:`~simsopt.geo.curveplanarfourier.CurvePlanarFourier` of order
+    ``order`` that will result in circular equally spaced coils (major
+    radius ``R0`` and minor radius ``R1``) after applying
+    :obj:`~simsopt.field.coil.coils_via_symmetries`.
+    """
+
+    if numquadpoints is None:
+        numquadpoints = 15 * order
+    curves = []
+    from simsopt.geo.curveplanarfourier import CurvePlanarFourier
+    for k in range(ncurves):
+        angle = (k+0.5)*(2*np.pi) / ((1+int(stellsym))*nfp*ncurves)
+        curve = CurvePlanarFourier(numquadpoints, order, nfp, stellsym)
+
+        rcCoeffs = np.zeros(order+1)
+        rcCoeffs[0] = R1
+        rsCoeffs = np.zeros(order)
+        center = [R0 * cos(angle), R0 * sin(angle), 0]
+        rotation = [1, -cos(angle), -sin(angle), 0]
+        dofs = np.zeros(len(curve.get_dofs()))
+
+        j = 0
+        for i in rcCoeffs:
+            dofs[j] = i
+            j += 1
+        for i in rsCoeffs:
+            dofs[j] = i
+            j += 1
+        for i in rotation:
+            dofs[j] = i
+            j += 1
+        for i in center:
+            dofs[j] = i
+            j += 1
+
+        curve.set_dofs(dofs)
+        curves.append(curve)
+    return curves

src/simsopt/geo/curveobjectives.py

@@ -520,7 +520,7 @@ def J(self):
                     integrals = sopp.linkNumber(R1, R2, dR1, dR2) * dS * dT
                     linkNum[i-1][j-1] = 1/(4*np.pi) * (integrals)
         linkNumSum = sum(sum(abs(linkNum)))
-        return linkNumSum
+        return round(linkNumSum)
 
     @derivative_dec
     def dJ(self):

src/simsopt/geo/curveplanarfourier.py

@@ -0,0 +1,75 @@
+import numpy as np
+
+import simsoptpp as sopp
+from .curve import Curve
+
+__all__ = ['CurvePlanarFourier']
+
+
+class CurvePlanarFourier(sopp.CurvePlanarFourier, Curve):
+    r"""
+    ``CurvePlanarFourier`` is a curve that is restricted to lie in a plane. The
+    shape of the curve within the plane is represented by a Fourier series in
+    polar coordinates. The resulting planar curve is then rotated in three
+    dimensions using a quaternion, and finally a translation is applied. The
+    Fourier series in polar coordinates is
+
+    .. math::
+
+       r(\phi) = \sum_{m=0}^{\text{order}} r_{c,m}\cos(m \phi) + \sum_{m=1}^{\text{order}} r_{s,m}\sin(m \phi).
+
+    The rotation quaternion is
+
+    .. math::
+
+       \bf{q} &= [q_0,q_i,q_j,q_k]
+
+       &= [\cos(\theta / 2), \hat{x}\sin(\theta / 2), \hat{y}\sin(\theta / 2), \hat{z}\sin(\theta / 2)]
+
+    where :math:`\theta` is the counterclockwise rotation angle about a unit axis
+    :math:`(\hat{x},\hat{y},\hat{z})`. Details of the quaternion rotation can be
+    found for example in pages 575-576 of
+    https://www.cis.upenn.edu/~cis5150/ws-book-Ib.pdf.
+
+
+    A quaternion is used for rotation rather than other methods for rotation to
+    prevent gimbal locking during optimization. The quaternion is normalized
+    before being applied to prevent scaling of the curve. The dofs themselves are not normalized. This
+    results in a redundancy in the optimization, where several different sets of 
+    dofs may correspond to the same normalized quaternion. Normalizing the dofs 
+    directly would create a dependence between the quaternion dofs, which may cause 
+    issues during optimization.
+
+    The dofs are stored in the order
+
+    .. math::
+       [r_{c,0}, \cdots, r_{c,\text{order}}, r_{s,1}, \cdots, r_{s,\text{order}}, q_0, q_i, q_j, q_k, x_{\text{center}}, y_{\text{center}}, z_{\text{center}}]
+
+
+    """
+
+    def __init__(self, quadpoints, order, nfp, stellsym, dofs=None):
+        if isinstance(quadpoints, int):
+            quadpoints = list(np.linspace(0, 1., quadpoints, endpoint=False))
+        elif isinstance(quadpoints, np.ndarray):
+            quadpoints = list(quadpoints)
+        sopp.CurvePlanarFourier.__init__(self, quadpoints, order, nfp, stellsym)
+        if dofs is None:
+            Curve.__init__(self, external_dof_setter=CurvePlanarFourier.set_dofs_impl,
+                           x0=self.get_dofs())
+        else:
+            Curve.__init__(self, external_dof_setter=CurvePlanarFourier.set_dofs_impl,
+                           dofs=dofs)
+
+    def get_dofs(self):
+        """
+        This function returns the dofs associated to this object.
+        """
+        return np.asarray(sopp.CurvePlanarFourier.get_dofs(self))
+
+    def set_dofs(self, dofs):
+        """
+        This function sets the dofs associated to this object.
+        """
+        self.local_x = dofs
+        sopp.CurvePlanarFourier.set_dofs(self, dofs)