Merge branch 'planar_coil_arrays' of https://github.com/hiddenSymmetr…

…ies/simsopt into planar_coil_arrays

.github/workflows/docs_test.yml

@@ -9,7 +9,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: [3.9]
+        python-version: ["3.11"]
 
     steps:
     - uses: actions/checkout@v4

.readthedocs.yaml

@@ -4,7 +4,7 @@ version: 2
 build:
   os: ubuntu-22.04
   tools:
-    python: "3.8"
+    python: "3.11"
 
 # Build documentation in the docs/ directory with Sphinx
 sphinx:

examples/3_Advanced/QA_reactorscale_fixed_orientations.py

@@ -159,6 +159,7 @@ def initialize_coils_QA(TEST_DIR, s):
 import warnings
 
 keep_inds = []
+dists = np.zeros(len(base_curves))
 for ii in range(len(base_curves)):
     counter = 0
     for i in range(base_curves[0].gamma().shape[0]):
@@ -174,9 +175,19 @@ def initialize_coils_QA(TEST_DIR, s):
                         'of a TF coil. Deleting these PSCs now.')
                     counter += 1
                     break
+    dists[ii] = np.min(np.linalg.norm(base_curves[ii].gamma(), axis=-1))
     if counter == 0:
         keep_inds.append(ii)
 
+# Chop off the dipole coils in the tight inboard side
+# since these coils often have very high forces and prevent the TF
+# coils from moving around much 
+dists = dists[keep_inds]
+argdists = np.argsort(dists)
+keep_inds = np.array(keep_inds)[argdists]
+for i in range(4):
+    keep_inds = np.delete(keep_inds, [0])
+
 print(keep_inds)
 base_curves = np.array(base_curves)[keep_inds]
 
@@ -279,7 +290,11 @@ def pointData_forces_torques(coils, allcoils, aprimes, bprimes, nturns_list):
 CS_THRESHOLD = 1.5
 CS_WEIGHT = 1e2
 # Weight for the Coil Coil forces term
-FORCE_WEIGHT = Weight(1e-21) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
+# FORCE_WEIGHT = Weight(1e-19) # 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-24) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
+# TORQUE_WEIGHT2 = Weight(1e-24) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
+FORCE_WEIGHT = Weight(1e-19) # 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-24) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
 TORQUE_WEIGHT2 = Weight(1e-24) # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
@@ -365,9 +380,9 @@ 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=1e5 * 40) 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=2, threshold=4e5 * 100) for i, c in enumerate(base_coils + base_coils_TF)])
 Jforce2 = sum([SquaredMeanForce(c, coils + coils_TF) for c in (base_coils + base_coils_TF)])
-Jtorque = sum([LpCurveTorque(c, coils + coils_TF, regularization_rect(a_list[i], b_list[i]), p=2, threshold=1e5 * 40) 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)]
@@ -541,7 +556,7 @@ def fun(dofs):
 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': 400}, tol=1e-15)
+        options={'maxiter': MAXITER, 'maxcor': 100}, tol=1e-15)
     dofs = res.x
 
     dipole_currents = [c.current.get_value() for c in bs.coils]

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 = 5
+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(1e-19) # 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(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',