Made some more useful changes. Got the MixedLpCurve class running fas…

…ter by downsampling the calculation, but for some reason the initial compilation is very slow. Still trying to understand how to avoid spawning all the weak references to child processes during optimization with the normal LpCurveForce object, and added downsampling to this too. This is probably worth trying to figure out definitively.

examples/3_Advanced/QA_single_TFcoil.py

@@ -7,6 +7,7 @@
 from pathlib import Path
 import time
 import numpy as np
+import warnings
 from scipy.optimize import minimize
 from simsopt.field import BiotSavart, Current, coils_via_symmetries
 # from simsopt.field import CoilCoilNetForces, CoilCoilNetTorques, \
@@ -41,8 +42,8 @@
 range_param = "half period"
 nphi = 32
 ntheta = 32
-poff = 1.5
-coff = 2.0
+poff = 1.75
+coff = 2.5
 s = SurfaceRZFourier.from_vmec_input(filename, range=range_param, nphi=nphi, ntheta=ntheta)
 s_inner = SurfaceRZFourier.from_vmec_input(filename, range=range_param, nphi=nphi * 4, ntheta=ntheta * 4)
 s_outer = SurfaceRZFourier.from_vmec_input(filename, range=range_param, nphi=nphi * 4, ntheta=ntheta * 4)
@@ -88,7 +89,7 @@ def initialize_coils_QA(TEST_DIR, s):
     ncoils = 1
     R0 = s.get_rc(0, 0) * 2
     R1 = s.get_rc(1, 0) * 10
-    order = 4
+    order = 10
 
     from simsopt.mhd.vmec import Vmec
     vmec_file = 'wout_LandremanPaul2021_QA_reactorScale_lowres_reference.nc'
@@ -141,33 +142,106 @@ def initialize_coils_QA(TEST_DIR, s):
 aa = 0.05
 bb = 0.05
 
-Nx = 3
+Nx = 6
 Ny = Nx
 Nz = Nx
 # Create the initial coils:
 base_curves, all_curves = create_planar_curves_between_two_toroidal_surfaces(
     s, s_inner, s_outer, Nx, Ny, Nz, order=order, coil_coil_flag=True, jax_flag=True,
     # numquadpoints=10  # Defaults is (order + 1) * 40 so this halves it
 )
+# ncoils = len(base_curves)
+# print('Ncoils = ', ncoils)
+# for i in range(len(base_curves)):
+#     # base_curves[i].set('x' + str(2 * order + 1), np.random.rand(1) - 0.5)
+#     # base_curves[i].set('x' + str(2 * order + 2), np.random.rand(1) - 0.5)
+#     # base_curves[i].set('x' + str(2 * order + 3), np.random.rand(1) - 0.5)
+#     # base_curves[i].set('x' + str(2 * order + 4), np.random.rand(1) - 0.5)
+
+#     # Fix shape of each coil
+#     for j in range(2 * order + 1):
+#         base_curves[i].fix('x' + str(j))
+#     # Fix center points of each coil
+#     # base_curves[i].fix('x' + str(2 * order + 5))
+#     # base_curves[i].fix('x' + str(2 * order + 6))
+#     # base_curves[i].fix('x' + str(2 * order + 7))
+# base_currents = [Current(1e-1) * 2e7 for i in range(ncoils)]
+# # Fix currents in each coil
+# # for i in range(ncoils):
+# #     base_currents[i].fix_all()
+
+# coils = coils_via_symmetries(base_curves, base_currents, s.nfp, True)
+# base_coils = coils[:ncoils]
+
+
+keep_inds = []
+for ii in range(len(base_curves)):
+    counter = 0
+    for i in range(base_curves[0].gamma().shape[0]):
+        eps = 0.05
+        for j in range(len(base_curves_TF)):
+            for k in range(base_curves_TF[j].gamma().shape[0]):
+                dij = np.sqrt(np.sum((base_curves[ii].gamma()[i, :] - base_curves_TF[j].gamma()[k, :]) ** 2))
+                conflict_bool = (dij < (1.0 + eps) * base_curves[0].x[0])
+                if conflict_bool:
+                    print('bad indices = ', i, j, dij, base_curves[0].x[0])
+                    warnings.warn(
+                        'There is a PSC coil initialized such that it is within a radius'
+                        'of a TF coil. Deleting these PSCs now.')
+                    counter += 1
+                    break
+    if counter == 0:
+        keep_inds.append(ii)
+
+print(keep_inds)
+base_curves = np.array(base_curves)[keep_inds]
+
 ncoils = len(base_curves)
 print('Ncoils = ', ncoils)
+coil_normals = np.zeros((ncoils, 3))
+plasma_points = s.gamma().reshape(-1, 3)
+plasma_unitnormals = s.unitnormal().reshape(-1, 3)
+for i in range(ncoils):
+    point = (base_curves[i].get_dofs()[-3:])
+    dists = np.sum((point - plasma_points) ** 2, axis=-1)
+    min_ind = np.argmin(dists)
+    coil_normals[i, :] = plasma_unitnormals[min_ind, :]
+    # coil_normals[i, :] = (plasma_points[min_ind, :] - point)
+coil_normals = coil_normals / np.linalg.norm(coil_normals, axis=-1)[:, None]
+# alphas = np.arctan2(
+#                 -coil_normals[:, 1], 
+#                 np.sqrt(coil_normals[:, 0] ** 2 + coil_normals[:, 2] ** 2))
+# deltas = np.arcsin(coil_normals[:, 0] / \
+#                             np.sqrt(coil_normals[:, 0] ** 2 + coil_normals[:, 2] ** 2))
+alphas = np.arcsin(
+                -coil_normals[:, 1], 
+                )
+deltas = np.arctan2(coil_normals[:, 0], coil_normals[:, 2])
 for i in range(len(base_curves)):
-    # base_curves[i].set('x' + str(2 * order + 1), np.random.rand(1) - 0.5)
-    # base_curves[i].set('x' + str(2 * order + 2), np.random.rand(1) - 0.5)
-    # base_curves[i].set('x' + str(2 * order + 3), np.random.rand(1) - 0.5)
-    # base_curves[i].set('x' + str(2 * order + 4), np.random.rand(1) - 0.5)
+    alpha2 = alphas[i] / 2.0
+    delta2 = deltas[i] / 2.0
+    calpha2 = np.cos(alpha2)
+    salpha2 = np.sin(alpha2)
+    cdelta2 = np.cos(delta2)
+    sdelta2 = np.sin(delta2)
+    base_curves[i].set('x' + str(2 * order + 1), calpha2 * cdelta2)
+    base_curves[i].set('x' + str(2 * order + 2), salpha2 * cdelta2)
+    base_curves[i].set('x' + str(2 * order + 3), calpha2 * sdelta2)
+    base_curves[i].set('x' + str(2 * order + 4), -salpha2 * sdelta2)
+    # Fix orientations of each coil
+    base_curves[i].fix('x' + str(2 * order + 1))
+    base_curves[i].fix('x' + str(2 * order + 2))
+    base_curves[i].fix('x' + str(2 * order + 3))
+    base_curves[i].fix('x' + str(2 * order + 4))
 
     # Fix shape of each coil
     for j in range(2 * order + 1):
         base_curves[i].fix('x' + str(j))
     # Fix center points of each coil
-    # base_curves[i].fix('x' + str(2 * order + 5))
-    # base_curves[i].fix('x' + str(2 * order + 6))
-    # base_curves[i].fix('x' + str(2 * order + 7))
+    base_curves[i].fix('x' + str(2 * order + 5))
+    base_curves[i].fix('x' + str(2 * order + 6))
+    base_curves[i].fix('x' + str(2 * order + 7))
 base_currents = [Current(1e-1) * 2e7 for i in range(ncoils)]
-# Fix currents in each coil
-# for i in range(ncoils):
-#     base_currents[i].fix_all()
 
 coils = coils_via_symmetries(base_curves, base_currents, s.nfp, True)
 base_coils = coils[:ncoils]
@@ -210,24 +284,31 @@ def pointData_forces_torques(coils, allcoils, aprimes, bprimes, nturns_list):
 base_a_list = np.hstack((np.ones(len(base_coils)) * aa, np.ones(len(base_coils_TF)) * a))
 base_b_list = np.hstack((np.ones(len(base_coils)) * bb, np.ones(len(base_coils_TF)) * b))
 
-LENGTH_WEIGHT = Weight(0.01)
-LENGTH_TARGET = 90  # 90 works very well... can we get it down?
-LINK_WEIGHT = 1e3
+LENGTH_WEIGHT = Weight(0.001)
+LENGTH_TARGET = 80  # 90 works very well... can we get it down?
+LINK_WEIGHT = 1e2
 CC_THRESHOLD = 0.8
-CC_WEIGHT = 1e2
+CC_WEIGHT = 1e1
 CS_THRESHOLD = 1.5
-CS_WEIGHT = 1e2
+CS_WEIGHT = 1e1
 # Weight for the Coil Coil forces term
-FORCE_WEIGHT = Weight(0.0) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
+FORCE_WEIGHT = Weight(1e-37) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
 FORCE_WEIGHT2 = Weight(0.0) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
-TORQUE_WEIGHT = Weight(1e-18) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
+TORQUE_WEIGHT = Weight(0.0) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
+TORQUE_WEIGHT2 = Weight(0.0) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
+
+# FORCE_WEIGHT = Weight(0.0) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
+# FORCE_WEIGHT2 = Weight(0.0) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
+# TORQUE_WEIGHT = Weight(0.0) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
+# TORQUE_WEIGHT2 = Weight(0.0) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
 # Directory for output
 OUT_DIR = ("./QA_singleTF_n{:d}_p{:.2e}_c{:.2e}_lw{:.2e}_lt{:.2e}_lkw{:.2e}" + \
     "_cct{:.2e}_ccw{:.2e}_cst{:.2e}_csw{:.2e}_fw{:.2e}_fww{:2e}_tw{:.2e}/").format(
         ncoils, poff, coff, LENGTH_WEIGHT.value, LENGTH_TARGET, LINK_WEIGHT, 
         CC_THRESHOLD, CC_WEIGHT, CS_THRESHOLD, CS_WEIGHT, FORCE_WEIGHT.value, 
         FORCE_WEIGHT2.value,
-        TORQUE_WEIGHT.value)
+        TORQUE_WEIGHT.value,
+        TORQUE_WEIGHT2.value)
 if os.path.exists(OUT_DIR):
     shutil.rmtree(OUT_DIR)
 os.makedirs(OUT_DIR, exist_ok=True)
@@ -281,7 +362,8 @@ def pointData_forces_torques(coils, allcoils, aprimes, bprimes, nturns_list):
 
 # coil-coil and coil-plasma distances should be between all coils
 Jccdist = CurveCurveDistance(curves + curves_TF, CC_THRESHOLD, num_basecurves=len(coils + coils_TF))
-Jcsdist = CurveSurfaceDistance(curves + curves_TF, s, CS_THRESHOLD)
+Jcsdist = CurveSurfaceDistance(curves_TF, s, CS_THRESHOLD)
+Jcsdist2 = CurveSurfaceDistance(curves + curves_TF, s, CS_THRESHOLD)
 
 # While the coil array is not moving around, they cannot
 # interlink. 
@@ -298,9 +380,13 @@ def pointData_forces_torques(coils, allcoils, aprimes, bprimes, nturns_list):
 regularization_list2 = np.zeros(len(coils_TF)) * regularization_rect(a, b)
 # Jforce = MixedLpCurveForce(coils, coils_TF, regularization_list, regularization_list2) # [SquaredMeanForce2(c, coils) for c in (base_coils)]
 # Jforce = MixedSquaredMeanForce(coils, coils_TF)
-Jforce = sum([LpCurveForce(c, coils + coils_TF, regularization_rect(a_list[i], b_list[i]), p=2, threshold=2e5 * 40) for i, c in enumerate(base_coils + base_coils_TF)])
+
+###### NOTE JFORCE BELOW ONLY DOING THE DIPOLE COILS!!!!
+Jforce = sum([LpCurveForce(c, coils + coils_TF, regularization_rect(a_list[i], b_list[i]), p=4, threshold=1e5 * 100) for i, c in enumerate(base_coils + base_coils_TF)])
+# Jforce = sum([LpCurveForce(c, coils + coils_TF, regularization_rect(a_list[i], b_list[i]), p=4, threshold=4e5 * 100) for i, c in enumerate(base_coils)])
 Jforce2 = sum([SquaredMeanForce(c, coils + coils_TF) for i, c in enumerate(base_coils + base_coils_TF)])
-Jtorque = sum([LpCurveTorque(c, coils + coils_TF, regularization_rect(a_list[i], b_list[i])) for i, c in enumerate(base_coils + base_coils_TF)])
+Jtorque = sum([LpCurveTorque(c, coils + coils_TF, regularization_rect(a_list[i], b_list[i]), p=2, threshold=4e5 * 100) for i, c in enumerate(base_coils + base_coils_TF)])
+Jtorque2 = sum([SquaredMeanTorque(c, coils + coils_TF) for c in (base_coils + base_coils_TF)])
 # Jtorque = SquaredMeanTorque2(coils, coils_TF) # [SquaredMeanForce2(c, coils) for c in (base_coils)]
 # Jtorque = [SquaredMeanTorque(c, coils + coils_TF) for c in (base_coils + base_coils_TF)]
 
@@ -317,17 +403,11 @@ def pointData_forces_torques(coils, allcoils, aprimes, bprimes, nturns_list):
     JF += FORCE_WEIGHT2.value * Jforce2  #\
 
 if TORQUE_WEIGHT.value > 0.0:
-    JF += TORQUE_WEIGHT.value * Jtorque  #\
-    # + FORCE_WEIGHT2.value * Jforce2
-    # + TORQUE_WEIGHT * Jtorque
-    # + TVE_WEIGHT * Jtve
-    # + SF_WEIGHT * Jsf
-    # + CURRENTS_WEIGHT * DipoleCurrentsObj
-    # + CURVATURE_WEIGHT * sum(Jcs_TF) \
-    # + MSC_WEIGHT * sum(QuadraticPenalty(J, MSC_THRESHOLD) for J in Jmscs_TF) \
-#    + MSC_WEIGHT * sum(QuadraticPenalty(J, MSC_THRESHOLD) for J in Jmscs) \
-    # + CURVATURE_WEIGHT * sum(Jcs) \
+    JF += TORQUE_WEIGHT * Jtorque
 
+if TORQUE_WEIGHT2.value > 0.0:
+    JF += TORQUE_WEIGHT2 * Jtorque2
+
 # 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
@@ -359,6 +439,7 @@ def fun(dofs):
     forces_val = FORCE_WEIGHT.value * Jforce.J()
     forces_val2 = FORCE_WEIGHT2.value * Jforce2.J()
     torques_val = TORQUE_WEIGHT.value * Jtorque.J()
+    torques_val2 = TORQUE_WEIGHT2.value * Jtorque2.J()
     BdotN = np.mean(np.abs(np.sum(btot.B().reshape((nphi, ntheta, 3)) * s.unitnormal(), axis=2)))
     BdotN_over_B = np.mean(np.abs(np.sum(btot.B().reshape((nphi, ntheta, 3)) * s.unitnormal(), axis=2))
         ) / np.mean(btot.AbsB())
@@ -373,10 +454,12 @@ def fun(dofs):
     valuestr += f", forceObj={forces_val:.2e}" 
     valuestr += f", forceObj2={forces_val2:.2e}" 
     valuestr += f", torqueObj={torques_val:.2e}" 
+    valuestr += f", torqueObj2={torques_val2:.2e}" 
     outstr += f", F={Jforce.J():.2e}"
-    outstr += f", Fpointwise={Jforce2.J():.2e}"
+    outstr += f", Fnet={Jforce2.J():.2e}"
     outstr += f", T={Jtorque.J():.2e}"
-    outstr += f", C-C-Sep={Jccdist.shortest_distance():.2f}, C-S-Sep={Jcsdist.shortest_distance():.2f}"
+    outstr += f", Tnet={Jtorque2.J():.2e}"
+    outstr += f", C-C-Sep={Jccdist.shortest_distance():.2f}, C-S-Sep={Jcsdist2.shortest_distance():.2f}"
     outstr += f", Link Number = {linkNum.J()}"
     outstr += f", ║∇J║={np.linalg.norm(grad):.1e}"
     print(outstr)
@@ -438,33 +521,33 @@ def fun(dofs):
 Jlength.dJ()
 t2 = time.time()
 print('sum(Jls_TF) time = ', t2 - t1, ' s')
-t1 = time.time()
-Jforce.J()
-t2 = time.time()
-print('Jforces time = ', t2 - t1, ' s')
-t1 = time.time()
-Jforce.dJ()
-t2 = time.time()
-print('dJforces time = ', t2 - t1, ' s')
-t1 = time.time()
-Jforce2.J()
-t2 = time.time()
-print('Jforces2 time = ', t2 - t1, ' s')
-t1 = time.time()
-Jforce2.dJ()
-t2 = time.time()
-print('dJforces2 time = ', t2 - t1, ' s')
-t1 = time.time()
-Jtorque.J()
-t2 = time.time()
-print('Jtorques time = ', t2 - t1, ' s')
-t1 = time.time()
-Jtorque.dJ()
-t2 = time.time()
-print('dJtorques time = ', t2 - t1, ' s')
+# t1 = time.time()
+# Jforce.J()
+# t2 = time.time()
+# print('Jforces time = ', t2 - t1, ' s')
+# t1 = time.time()
+# Jforce.dJ()
+# t2 = time.time()
+# print('dJforces time = ', t2 - t1, ' s')
+# t1 = time.time()
+# Jforce2.J()
+# t2 = time.time()
+# print('Jforces2 time = ', t2 - t1, ' s')
+# t1 = time.time()
+# Jforce2.dJ()
+# t2 = time.time()
+# print('dJforces2 time = ', t2 - t1, ' s')
+# t1 = time.time()
+# Jtorque.J()
+# t2 = time.time()
+# print('Jtorques time = ', t2 - t1, ' s')
+# t1 = time.time()
+# Jtorque.dJ()
+# t2 = time.time()
+# print('dJtorques time = ', t2 - t1, ' s')
 
 n_saves = 1
-MAXITER = 400
+MAXITER = 200
 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',