got force minimization working to some extent.

examples/2_Intermediate/stage_two_optimization.py

@@ -174,6 +174,7 @@ def fun(dofs):
 # 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.
+print('Switching to longer coils: ')
 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)

examples/2_Intermediate/stage_two_optimization_jax.py

@@ -22,6 +22,7 @@
 """
 
 import os
+import shutil
 from pathlib import Path
 import numpy as np
 from scipy.optimize import minimize
@@ -48,7 +49,7 @@
 # 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(1e-6)
+LENGTH_WEIGHT = Weight(1e-5)
 
 # Threshold and weight for the coil-to-coil distance penalty in the objective function:
 CC_THRESHOLD = 0.1
@@ -67,7 +68,8 @@
 MSC_WEIGHT = 1e-6
 
 # Weight for the Coil Coil forces term
-FORCES_WEIGHT = 1e-9  # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
+FORCES_WEIGHT = 1e-14  # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
+# And this term weights the NetForce^2 ~ 10^10-10^12 
 
 # Number of iterations to perform:
 MAXITER = 100
@@ -78,6 +80,8 @@
 
 # Directory for output
 OUT_DIR = "./stage_two_optimization_jax/"
+if os.path.exists(OUT_DIR):
+    shutil.rmtree(OUT_DIR)
 os.makedirs(OUT_DIR, exist_ok=True)
 
 #######################################################
@@ -89,6 +93,17 @@
 ntheta = 32
 s = SurfaceRZFourier.from_vmec_input(filename, range="half period", nphi=nphi, ntheta=ntheta)
 
+qphi = nphi * 2
+qtheta = ntheta * 2
+quadpoints_phi = np.linspace(0, 1, qphi, endpoint=True)
+quadpoints_theta = np.linspace(0, 1, qtheta, endpoint=True)
+# Make high resolution, full torus version of the plasma boundary for plotting
+s_plot = SurfaceRZFourier.from_vmec_input(
+    filename, 
+    quadpoints_phi=quadpoints_phi, 
+    quadpoints_theta=quadpoints_theta
+)
+
 # Create the initial coils:
 base_curves = create_equally_spaced_curves(ncoils, s.nfp, stellsym=True, R0=R0, R1=R1, order=order, jax_flag=True)
 base_currents = [JaxCurrent(1e5) for i in range(ncoils)]
@@ -99,13 +114,23 @@
 
 coils = coils_via_symmetries(base_curves, base_currents, s.nfp, True)
 curves = [c.curve for c in coils]
-curves_to_vtk(curves, OUT_DIR + "curves_init")
+currents = [c.current.get_value() for c in coils]
 
 bs = JaxBiotSavart(coils)
 bs.set_points(s.gamma().reshape((-1, 3)))
+curves_to_vtk(curves, OUT_DIR + "curves_0", I=np.array(currents), NetForces=np.array(bs.coil_coil_forces()))
 
 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)
+s.to_vtk(OUT_DIR + "surf_0", extra_data=pointData)
+
+bs.set_points(s_plot.gamma().reshape((-1, 3)))
+pointData = {"B_N": np.sum(bs.B().reshape((qphi, qtheta, 3)) * s_plot.unitnormal(), axis=2)[:, :, None]}
+s_plot.to_vtk(OUT_DIR + "surf_full_0", extra_data=pointData)
+
+pointData = {"B_N / B": (np.sum(bs.B().reshape((qphi, qtheta, 3)) * s_plot.unitnormal(), axis=2
+    ) / np.linalg.norm(bs.B().reshape(qphi, qtheta, 3), axis=-1))[:, :, None]}
+s_plot.to_vtk(OUT_DIR + "surf_full_normalizedBn_0", extra_data=pointData)
+bs.set_points(s.gamma().reshape((-1, 3)))
 
 # Define the individual terms objective function:
 Jf = SquaredFlux(s, bs)
@@ -140,13 +165,15 @@ def fun(dofs):
     length_val = LENGTH_WEIGHT.value * sum(J.J() for J in Jls)
     cc_val = CC_WEIGHT * Jccdist.J()
     cs_val = CS_WEIGHT * Jcsdist.J()
+    curv_val = CURVATURE_WEIGHT * sum(J.J() for J in Jcs)
     forces_val = FORCES_WEIGHT * Jforces.J()
     BdotN = np.mean(np.abs(np.sum(bs.B().reshape((nphi, ntheta, 3)) * s.unitnormal(), axis=2)))
     outstr = f"J={J:.1e}, Jf={jf:.1e}, ⟨B·n⟩={BdotN:.1e}"
     valuestr = f"J={J:.2e}, Jf={jf:.2e}"
     valuestr += f", LenObj={length_val:.2e}" 
     valuestr += f", ccObj={cc_val:.2e}" 
     valuestr += f", csObj={cs_val:.2e}" 
+    valuestr += f", curvObj={curv_val:.2e}" 
     valuestr += f", forceObj={forces_val:.2e}" 
     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)
@@ -182,21 +209,38 @@ def fun(dofs):
 ### 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)
+n_saves = 40
+MAXITER = 10
+for i in range(1, n_saves + 1):
+    print('Iteration ' + str(i) + ' / ' + str(n_saves))
+    res = minimize(fun, dofs, jac=True, method='L-BFGS-B', options={'maxiter': MAXITER, 'maxcor': 300}, tol=1e-15)
+    dofs = res.x
+
+    curves_to_vtk([c.curve for c in bs.coils], OUT_DIR + "curves_{0:d}".format(i), 
+                  I=[c.current.get_value() for c in bs.coils], NetForces=bs.coil_coil_forces())
+
+    bs.set_points(s_plot.gamma().reshape((-1, 3)))
+    pointData = {"B_N": np.sum(bs.B().reshape((qphi, qtheta, 3)) * s_plot.unitnormal(), axis=2)[:, :, None]}
+    s_plot.to_vtk(OUT_DIR + "surf_full_{0:d}".format(i), extra_data=pointData)
+
+    pointData = {"B_N / B": (np.sum(bs.B().reshape((qphi, qtheta, 3)) * s_plot.unitnormal(), axis=2
+        ) / np.linalg.norm(bs.B().reshape(qphi, qtheta, 3), axis=-1))[:, :, None]}
+    s_plot.to_vtk(OUT_DIR + "surf_full_normalizedBn_{0:d}".format(i), extra_data=pointData)
+
+    bs.set_points(s.gamma().reshape((-1, 3)))
+    print('Max I = ', np.max(np.abs([c.current.get_value() for c in bs.coils])))
+    print('Min I = ', np.min(np.abs([c.current.get_value() for c in bs.coils])))
 
 
 # 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")
+# 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")

src/simsopt/geo/curve.py

@@ -822,7 +822,7 @@ def dgammadashdashdash_by_dcoeff_vjp(self, v):
     def flip(self):
         return True if self.rotmat[2][2] == -1 else False
 
-def curves_to_vtk(curves, filename, close=False, scalar_data=None):
+def curves_to_vtk(curves, filename, close=False, I=None, NetForces=None):
     """
     Export a list of Curve objects in VTK format, so they can be
     viewed using Paraview. This function requires the python package ``pyevtk``,
@@ -849,14 +849,33 @@ def wrap(data):
         z = np.concatenate([c.gamma()[:, 2] for c in curves])
         ppl = np.asarray([c.gamma().shape[0] for c in curves])
     data = np.concatenate([i*np.ones((ppl[i], )) for i in range(len(curves))])
-    if scalar_data is None:
-        polyLinesToVTK(str(filename), x, y, z, pointsPerLine=ppl, pointData={'idx': data})
-    else:
+    pointData={'idx': data}
+    # cellData={}
+    contig = np.ascontiguousarray
+    if I is not None:
         coil_data = np.zeros(data.shape)
-        for i in range(len(scalar_data)):
-            coil_data[i * ppl[i]: (i + 1) * ppl[i]] = scalar_data[i]
+        for i in range(len(I)):
+            coil_data[i * ppl[i]: (i + 1) * ppl[i]] = I[i]
+        coil_data = np.ascontiguousarray(coil_data)
+        pointData['I'] = coil_data
+        pointData['I_mag'] = contig(np.abs(coil_data))
+        # cellData['I'] = contig(I)
+        # cellData['I_mag'] = contig(np.abs(I))
+    if NetForces is not None:
+        coil_data = np.zeros((data.shape[0], 3))
+        for i in range(len(NetForces)):
+            coil_data[i * ppl[i]: (i + 1) * ppl[i], :] = NetForces[i, :]
         coil_data = np.ascontiguousarray(coil_data)
-        polyLinesToVTK(str(filename), x, y, z, pointsPerLine=ppl, pointData={'idx': data, 'I': coil_data, 'I_mag': np.abs(coil_data)})
+        pointData['NetForces'] = (contig(coil_data[:, 0]), 
+                                contig(coil_data[:, 1]),
+                                contig(coil_data[:, 2]))
+        pointData['NetForces_mag'] = contig(np.linalg.norm(coil_data, axis=-1))
+        # cellData['NetForces'] = (contig(NetForces[:, 0]), 
+        #                         contig(NetForces[:, 1]),
+        #                         contig(NetForces[:, 2]))
+        # cellData['NetForces_mag'] = contig(np.linalg.norm(NetForces, axis=-1))
+
+    polyLinesToVTK(str(filename), x, y, z, pointsPerLine=ppl, pointData=pointData) #, cellData=cellData)
 
 def setup_uniform_grid(s, s_inner, s_outer, Nx, Ny, Nz, coil_coil_flag):
     # Get (X, Y, Z) coordinates of the two boundaries