Some more fiddling with the examples. Saving to try this on a gpu on …

…greene.

examples/2_Intermediate/stage_two_optimization_dipole_arrays_QA.py

@@ -121,6 +121,7 @@ def initialize_coils_QA(TEST_DIR, s):
     curves = [c.curve for c in coils]
     currents = [c.current.get_value() for c in coils]
     curves_to_vtk(curves, OUT_DIR + "curves_TF_0", I=currents,
+            close=True,
             # NetForces=np.array(bs.coil_coil_forces()),
             # NetTorques=bs.coil_coil_torques(),
             # NetSelfForces=bs.coil_self_forces(a, b)
@@ -184,12 +185,13 @@ def initialize_coils_QA(TEST_DIR, s):
 # print(np.sum(CoilCoilNetForces12(bs, bs_TF).coil_coil_forces12() ** 2, axis=-1))
 
 # exit()
-curves_to_vtk(curves, OUT_DIR + "curves_0", close=True, I=currents,
+curves_to_vtk(curves, OUT_DIR + "curves_0", I=currents,
+            close=True,
             NetForces=bs.coil_coil_forces(),
             NetTorques=bs.coil_coil_torques(),
-            MixedCoilForces=CoilCoilNetForces12(bs, bs_TF).coil_coil_forces12()[:len(curves), :],
-            MixedCoilTorques=CoilCoilNetTorques12(bs, bs_TF).coil_coil_torques12()[:len(curves), :],
-            NetSelfForces=bs.coil_self_forces(a, b)
+            # MixedCoilForces=CoilCoilNetForces12(bs, bs_TF).coil_coil_forces12()[:len(curves), :],
+            # MixedCoilTorques=CoilCoilNetTorques12(bs, bs_TF).coil_coil_torques12()[:len(curves), :],
+            # NetSelfForces=bs.coil_self_forces(a, b)
 )
 pointData = {"B_N": np.sum(btot.B().reshape((nphi, ntheta, 3)) * s.unitnormal(), axis=2)[:, :, None]}
 s.to_vtk(OUT_DIR + "surf_init_DA", extra_data=pointData)
@@ -208,10 +210,10 @@ def initialize_coils_QA(TEST_DIR, s):
 s_plot.to_vtk(OUT_DIR + "surf_full_init", extra_data=pointData)
 btot.set_points(s.gamma().reshape((-1, 3)))
 
-LENGTH_WEIGHT = Weight(0.005)
+LENGTH_WEIGHT = Weight(0.0025)
 # CURRENTS_WEIGHT = 10
 LINK_WEIGHT = 100
-LINK_WEIGHT2 = 1e-5
+LINK_WEIGHT2 = 1e-3
 CC_THRESHOLD = 0.7
 CC_WEIGHT = 400
 CS_THRESHOLD = 1.3
@@ -222,10 +224,10 @@ def initialize_coils_QA(TEST_DIR, s):
 # MSC_WEIGHT = 1e-12
 
 # Weight for the Coil Coil forces term
-FORCES_WEIGHT = 1e-20  # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
+FORCES_WEIGHT = 1e-18  # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
 # And this term weights the NetForce^2 ~ 10^10-10^12 
 
-TORQUES_WEIGHT = 1e-20  # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
+TORQUES_WEIGHT = 1e-18  # Forces are in Newtons, and typical values are ~10^5, 10^6 Newtons
 
 # TVE_WEIGHT = 1e-19
 
@@ -383,20 +385,22 @@ def fun(dofs):
 
     dipole_currents = [c.current.get_value() for c in bs.coils]
     curves_to_vtk([c.curve for c in bs.coils], OUT_DIR + "curves_{0:d}".format(i), 
+        close=True,
         I=dipole_currents,
         NetForces=np.array(bs.coil_coil_forces()),
         NetTorques=bs.coil_coil_torques(),
-        MixedCoilForces=CoilCoilNetForces12(bs, bs_TF).coil_coil_forces12()[:len(curves), :],
-        MixedCoilTorques=CoilCoilNetTorques12(bs, bs_TF).coil_coil_torques12()[:len(curves), :],
-        NetSelfForces=bs.coil_self_forces(a, b)
+        # MixedCoilForces=CoilCoilNetForces12(bs, bs_TF).coil_coil_forces12()[:len(curves), :],
+        # MixedCoilTorques=CoilCoilNetTorques12(bs, bs_TF).coil_coil_torques12()[:len(curves), :],
+        # NetSelfForces=bs.coil_self_forces(a, b)
         )
     curves_to_vtk([c.curve for c in bs_TF.coils], OUT_DIR + "curves_TF_{0:d}".format(i), 
+        close=True,
         I=[c.current.get_value() for c in bs_TF.coils],
         NetForces=np.array(bs_TF.coil_coil_forces()),
         NetTorques=bs_TF.coil_coil_torques(),
-        MixedCoilForces=CoilCoilNetForces12(bs, bs_TF).coil_coil_forces12()[len(curves):, :],
-        MixedCoilTorques=CoilCoilNetTorques12(bs, bs_TF).coil_coil_torques12()[len(curves):, :],
-        NetSelfForces=bs_TF.coil_self_forces(a, b)
+        # MixedCoilForces=CoilCoilNetForces12(bs, bs_TF).coil_coil_forces12()[len(curves):, :],
+        # MixedCoilTorques=CoilCoilNetTorques12(bs, bs_TF).coil_coil_torques12()[len(curves):, :],
+        # NetSelfForces=bs_TF.coil_self_forces(a, b)
     )
 
     btot.set_points(s_plot.gamma().reshape((-1, 3)))

examples/2_Intermediate/stage_two_optimization_dipole_arrays_QA_noForces_noTorques.py

@@ -121,6 +121,7 @@ def initialize_coils_QA(TEST_DIR, s):
     curves = [c.curve for c in coils]
     currents = [c.current.get_value() for c in coils]
     curves_to_vtk(curves, OUT_DIR + "curves_TF_0", I=currents,
+            close=True,
             # NetForces=np.array(bs.coil_coil_forces()),
             # NetTorques=bs.coil_coil_torques(),
             # NetSelfForces=bs.coil_self_forces(a, b)
@@ -184,7 +185,8 @@ def initialize_coils_QA(TEST_DIR, s):
 # print(np.sum(CoilCoilNetForces12(bs, bs_TF).coil_coil_forces12() ** 2, axis=-1))
 
 # exit()
-curves_to_vtk(curves, OUT_DIR + "curves_0", close=True, I=currents,
+curves_to_vtk(curves, OUT_DIR + "curves_0", I=currents,
+            close=True,
             NetForces=bs.coil_coil_forces(),
             NetTorques=bs.coil_coil_torques(),
             MixedCoilForces=CoilCoilNetForces12(bs, bs_TF).coil_coil_forces12()[:len(curves), :],
@@ -251,8 +253,8 @@ def initialize_coils_QA(TEST_DIR, s):
 # interlink. 
 linkNum = LinkingNumber(curves_TF)
 linkNum2 = LinkingNumber(curves)
-Jforces = CoilCoilNetForces(bs) + CoilCoilNetForces12(bs, bs_TF) + CoilCoilNetForces(bs_TF)
-Jtorques = CoilCoilNetTorques(bs) + CoilCoilNetTorques12(bs, bs_TF) + CoilCoilNetTorques(bs_TF)
+# Jforces = CoilCoilNetForces(bs) + CoilCoilNetForces12(bs, bs_TF) + CoilCoilNetForces(bs_TF)
+# Jtorques = CoilCoilNetTorques(bs) + CoilCoilNetTorques12(bs, bs_TF) + CoilCoilNetTorques(bs_TF)
 # Jtve = TotalVacuumEnergy(bs, a=a, b=b)
 # Jsf = CoilSelfNetForces(bs, a=a, b=b)
 
@@ -286,9 +288,9 @@ def initialize_coils_QA(TEST_DIR, s):
     + CS_WEIGHT * Jcsdist \
     + LINK_WEIGHT * linkNum \
     + LINK_WEIGHT2 * linkNum2 \
-    + LENGTH_WEIGHT * sum(Jls_TF) \
-    + FORCES_WEIGHT * Jforces \
-    + TORQUES_WEIGHT * Jtorques 
+    + LENGTH_WEIGHT * sum(Jls_TF) # \
+    # + FORCES_WEIGHT * Jforces \
+    # + TORQUES_WEIGHT * Jtorques 
     # + TVE_WEIGHT * Jtve
     # + SF_WEIGHT * Jsf
     # + CURRENTS_WEIGHT * DipoleJaxCurrentsObj
@@ -312,8 +314,8 @@ def fun(dofs):
     cs_val = CS_WEIGHT * Jcsdist.J()
     link_val1 = LINK_WEIGHT * linkNum.J()
     link_val2 = LINK_WEIGHT2 * linkNum2.J()
-    forces_val = FORCES_WEIGHT * Jforces.J()
-    torques_val = TORQUES_WEIGHT * Jtorques.J()
+    # forces_val = FORCES_WEIGHT * Jforces.J()
+    # torques_val = TORQUES_WEIGHT * Jtorques.J()
     # tve_val = TVE_WEIGHT * Jtve.J()
     # sf_val = SF_WEIGHT * Jsf.J()
     BdotN = np.mean(np.abs(np.sum(btot.B().reshape((nphi, ntheta, 3)) * s.unitnormal(), axis=2)))
@@ -328,8 +330,8 @@ def fun(dofs):
     valuestr += f", csObj={cs_val:.2e}" 
     valuestr += f", Lk1Obj={link_val1:.2e}" 
     valuestr += f", Lk2Obj={link_val2:.2e}" 
-    valuestr += f", forceObj={forces_val:.2e}" 
-    valuestr += f", torqueObj={torques_val:.2e}" 
+    # valuestr += f", forceObj={forces_val:.2e}" 
+    # valuestr += f", torqueObj={torques_val:.2e}" 
     # valuestr += f", tveObj={tve_val:.2e}" 
     # valuestr += f", sfObj={sf_val:.2e}" 
     # valuestr += f", currObj={curr_val:.2e}" 
@@ -341,8 +343,8 @@ def fun(dofs):
     outstr += f", C-C-Sep={Jccdist.shortest_distance():.2f}, C-S-Sep={Jcsdist.shortest_distance():.2f}"
     outstr += f", Link Number = {linkNum.J()}"
     outstr += f", Link Number 2 = {linkNum2.J()}"
-    outstr += f", C-C-Forces={Jforces.J():.1e}"
-    outstr += f", C-C-Torques={Jtorques.J():.1e}"
+    # outstr += f", C-C-Forces={Jforces.J():.1e}"
+    # outstr += f", C-C-Torques={Jtorques.J():.1e}"
     # outstr += f", TVE={Jtve.J():.1e}"
     # outstr += f", TotalSelfForces={Jsf.J():.1e}"
     outstr += f", ║∇J║={np.linalg.norm(grad):.1e}"
@@ -383,6 +385,7 @@ def fun(dofs):
 
     dipole_currents = [c.current.get_value() for c in bs.coils]
     curves_to_vtk([c.curve for c in bs.coils], OUT_DIR + "curves_{0:d}".format(i), 
+        close=True,
         I=dipole_currents,
         NetForces=np.array(bs.coil_coil_forces()),
         NetTorques=bs.coil_coil_torques(),
@@ -391,6 +394,7 @@ def fun(dofs):
         NetSelfForces=bs.coil_self_forces(a, b)
         )
     curves_to_vtk([c.curve for c in bs_TF.coils], OUT_DIR + "curves_TF_{0:d}".format(i), 
+        close=True,
         I=[c.current.get_value() for c in bs_TF.coils],
         NetForces=np.array(bs_TF.coil_coil_forces()),
         NetTorques=bs_TF.coil_coil_torques(),

src/simsopt/field/biotsavart.py

@@ -557,14 +557,13 @@ def coil_coil_forces_pure(self, curve_dofs, currents):
         eps = 1e-20  # small number to avoid blow up in the denominator when i = j
         gammas = self.get_gammas(curve_dofs)
         gammadashs = self.get_gammadashs(curve_dofs)
-        Ii_Ij = (currents[None, :] * currents[:, None])[:, :, None]
         # gamma and gammadash are shape (ncoils, nquadpoints, 3)
         r_ij = gammas[None, :, None, :, :] - gammas[:, None, :, None, :]  # Note, do not use the i = j indices
         gammadash_prod = jnp.sum(gammadashs[None, :, None, :, :] * gammadashs[:, None, :, None, :], axis=-1) 
         rij_norm3 = jnp.linalg.norm(r_ij + eps, axis=-1) ** 3
 
         # Double sum over each of the closed curves
-        F = Ii_Ij * jnp.sum(jnp.sum((gammadash_prod / rij_norm3)[:, :, :, :, None] * r_ij, axis=3), axis=2)
+        F = (currents[None, :] * currents[:, None])[:, :, None] * jnp.sum(jnp.sum((gammadash_prod / rij_norm3)[:, :, :, :, None] * r_ij, axis=3), axis=2)
         net_forces = -jnp.sum(F, axis=1) * 1e-7 / jnp.shape(gammas)[1] ** 2
         return net_forces
 
@@ -722,10 +721,8 @@ def coil_coil_torques_pure(self, curve_dofs, currents):
         gammas_i = gammas[None, :, None, :, :]
         gammas_j = gammas[:, None, :, None, :]
         r_ij = gammas_i - gammas_j  # Note, do not use the i = j indices
-        gammadash_i = gammadashs[None, :, None, :, :]
-        gammadash_j = gammadashs[:, None, :, None, :]
-        cross1 = jnp.cross(gammadash_i, r_ij)
-        cross2 = jnp.cross(gammadash_j, cross1)
+        cross1 = jnp.cross(gammadashs[None, :, None, :, :], r_ij)
+        cross2 = jnp.cross(gammadashs[:, None, :, None, :], cross1)
         cross3 = jnp.cross(gammas_i, cross2)
         rij_norm3 = jnp.linalg.norm(r_ij + eps, axis=-1) ** 3
 

src/simsopt/geo/curve.py

@@ -1194,7 +1194,8 @@ def create_equally_spaced_curves(ncurves, nfp, stellsym, R0=1.0, R1=0.5, order=6
     return curves
 
 
-def create_equally_spaced_planar_curves(ncurves, nfp, stellsym, R0=1.0, R1=0.5, order=6, numquadpoints=None, jax_flag=False):
+def create_equally_spaced_planar_curves(
+    ncurves, nfp, stellsym, R0=1.0, R1=0.5, order=6, numquadpoints=None, jax_flag=False):
     """
     Create ``ncurves`` curves of type
     :obj:`~simsopt.geo.curveplanarfourier.CurvePlanarFourier` of order

tests/field/test_jaxbiotsavart.py

@@ -646,6 +646,10 @@ def total_energy_test(self):
         dofs[1:5] = np.ones(1)
         dofs[5:9] = np.array([0.0, 0.0, 0.0])
         curve0.set_dofs(dofs)
+        from matplotlib import pyplot as plt
+        # plt.figure()
+        # plt.plot(curve0.gamma()[:, 1], curve0.gamma()[:, 2])
+        # plt.show()
         current0 = Current(I1)
         bs = JaxBiotSavart([Coil(curve0, current0)])
         E_dipoles = bs.total_vacuum_energy_pure(bs.get_curve_dofs(), bs.get_currents(), a=0.01, b=0.01)
@@ -656,6 +660,39 @@ def total_energy_test(self):
         print('J = ', tve.J())
         print('dJ = ', tve.dJ())
         sf = CoilSelfNetForces(bs)
+        print(bs.coil_self_forces(0.01, 0.01))
+        print(sf.J(), sf.dJ())
+
+    def coil_objectives_symmetrized(self):       
+        from simsopt.geo import JaxCurvePlanarFourier, create_equally_spaced_planar_curves
+        ncoils = 2
+        I1 = 5.0e6
+        R0 = 1
+        # a = 1e-5
+        curves = create_equally_spaced_planar_curves(
+            ncoils, 2, True, R0=R0, R1=0.1, order=0, numquadpoints=None, jax_flag=True)
+        # curve0 = JaxCurvePlanarFourier(200, 0)
+        # dofs = np.zeros(8)
+        # dofs[0] = R1
+        # dofs[1:5] = np.ones(1)
+        # dofs[5:9] = np.array([0.0, 0.0, 0.0])
+        # curve0.set_dofs(dofs)
+        # from matplotlib import pyplot as plt
+        # plt.figure()
+        # plt.plot(curve0.gamma()[:, 1], curve0.gamma()[:, 2])
+        # plt.show()
+        currents = [Current(I1) for c in curves]
+        # current0 = Current(I1)
+        bs = JaxBiotSavart([Coil(c, currents[i]) for i, c in enumerate(curves)])
+        E_dipoles = bs.total_vacuum_energy_pure(bs.get_curve_dofs(), bs.get_currents(), a=0.01, b=0.01)
+        print('U = ', E_dipoles)
+        dE_dipoles = bs.dtve_by_dX(bs.get_curve_dofs(), bs.get_currents(), a=0.01, b=0.01)
+        print('dU = ', dE_dipoles)
+        tve = TotalVacuumEnergy(bs)
+        print('J = ', tve.J())
+        print('dJ = ', tve.dJ())
+        sf = CoilSelfNetForces(bs)
+        print(bs.coil_self_forces(0.01, 0.01))
         print(sf.J(), sf.dJ())
 
 if __name__ == "__main__":