Think I have the derivative calculated correctly, at least for the QA…

… discrete symmetries. Unit tests passing and WP optimization running well. Jacobian calculation still more expensive than I would like since xsimd seems not to give speedup. Saving and switching to laptop to see if intel gives the speedup I would expect.

examples/2_Intermediate/QA_psc_example.py

@@ -227,7 +227,7 @@ def fun(dofs):
 make_Bnormal_plots(bs, s_plot, out_dir, "biot_savart_TF_optimized", B_axis)
 
 # Finally, initialize the psc class
-kwargs_geo = {"Nx": 5, "out_dir": out_str, 
+kwargs_geo = {"Nx": 4, "out_dir": out_str, 
               "random_initialization": True, "poff": poff,}
               # "interpolated_field": True} 
 psc_array = PSCgrid.geo_setup_between_toroidal_surfaces(

examples/2_Intermediate/QA_wp_example.py

@@ -57,15 +57,15 @@ def coil_optimization_QA(s, bs, base_curves, curves, out_dir=''):
     ncoils = len(base_curves)
 
     # Weight on the curve lengths in the objective function:
-    LENGTH_WEIGHT = 1e-4
+    LENGTH_WEIGHT = 1e1
 
     # Threshold and weight for the coil-to-coil distance penalty in the objective function:
     CC_THRESHOLD = 0.2
     CC_WEIGHT = 1e1
 
     # Threshold and weight for the coil-to-surface distance penalty in the objective function:
-    CS_THRESHOLD = 0.5
-    CS_WEIGHT = 1
+    CS_THRESHOLD = 2.0
+    CS_WEIGHT = 1e-2
 
     # Threshold and weight for the curvature penalty in the objective function:
     CURVATURE_THRESHOLD = 10
@@ -144,14 +144,14 @@ def initialize_coils_qa(TEST_DIR, s, out_dir=''):
     from simsopt.geo import curves_to_vtk
     # generate planar TF coils
     ncoils = 1
-    R0 = 1.0
-    R1 = 0.6
+    R0 = s.get_rc(0, 0)
+    R1 = s.get_rc(1, 0) * 4
     order = 5
 
     # qa needs to be scaled to 0.1 T on-axis magnetic field strength
     from simsopt.mhd.vmec import Vmec
-    vmec_file = 'wout_LandremanPaul2021_QA_lowres.nc'
-    total_current = Vmec(TEST_DIR / vmec_file).external_current() / (2 * s.nfp) / 7.131
+    vmec_file = 'wout_LandremanPaul2021_QA_reactorScale_lowres_reference.nc'
+    total_current = Vmec(TEST_DIR / vmec_file).external_current() / (2 * s.nfp) / 7.131 * 6
     base_curves = create_equally_spaced_curves(ncoils, s.nfp, stellsym=True, R0=R0, R1=R1, order=order, numquadpoints=128)
     base_currents = [(Current(total_current / ncoils * 1e-5) * 1e5) for _ in range(ncoils)]
     base_currents[0].fix_all()
@@ -184,11 +184,11 @@ def initialize_coils_qa(TEST_DIR, s, out_dir=''):
     quadpoints_phi = np.linspace(0, 1, qphi, endpoint=True)
     quadpoints_theta = np.linspace(0, 1, ntheta * 2, endpoint=True)
 
-poff = 0.2  # WP grid will be offset 'poff' meters from the plasma surface
-coff = 0.1  # WP grid will be initialized between 1 m and 2 m from plasma
+poff = 1.0  # WP grid will be offset 'poff' meters from the plasma surface
+coff = 1.2  # WP grid will be initialized between 1 m and 2 m from plasma
 
 # Read in the plasma equilibrium file
-input_name = 'input.LandremanPaul2021_QA_lowres'
+input_name = 'input.LandremanPaul2021_QA_reactorScale_lowres'
 TEST_DIR = (Path(__file__).parent / ".." / ".." / "tests" / "test_files").resolve()
 surface_filename = TEST_DIR / input_name
 range_param = 'half period'
@@ -246,7 +246,7 @@ def initialize_coils_qa(TEST_DIR, s, out_dir=''):
 # fix all the coil shapes so only the currents are optimized
 # for i in range(ncoils):
 #     base_curves[i].fix_all()
-# bs = coil_optimization_QA(s, bs, base_curves, curves, out_dir)
+bs = coil_optimization_QA(s, bs, base_curves, curves, out_dir)
 curves_to_vtk(curves, out_dir / "TF_coils", close=True)
 bs.set_points(s.gamma().reshape((-1, 3)))
 Bnormal = np.sum(bs.B().reshape((nphi, ntheta, 3)) * s.unitnormal(), axis=2)
@@ -300,7 +300,7 @@ def initialize_coils_qa(TEST_DIR, s, out_dir=''):
 for i in range(wp_array.num_wp):
     opt_bounds.append((None, None))
 opt_bounds = tuple(opt_bounds)
-options = {"disp": True, "maxiter": 300, "iprint": 101}  #, "maxls": 100}
+options = {"disp": True, "maxiter": 500, "iprint": 101}  #, "maxls": 100}
 # print(opt_bounds)
 verbose = True
 
@@ -311,7 +311,7 @@ def initialize_coils_qa(TEST_DIR, s, out_dir=''):
                  "coils_TF" : coils
                  }
 t1 = time.time()
-I_threshold = 1e4
+I_threshold = 5e5
 STLSQ_max_iters = 10
 BdotN2_list = []
 num_wps = []

src/simsopt/geo/psc_grid.py

@@ -32,7 +32,7 @@ def __init__(self):
         self.BdotN2_list = []
         # Define a set of quadrature points and weights for the N point
         # Gaussian quadrature rule
-        num_quad = 80
+        num_quad = 20
         (quad_points_phi, 
           self.quad_weights_phi) = np.polynomial.legendre.leggauss(num_quad)
         self.quad_points_phi = quad_points_phi * np.pi + np.pi
@@ -312,6 +312,7 @@ def geo_setup_between_toroidal_surfaces(
         # Normalization of the ||A*Linv*psi - b||^2 objective 
         # representing Bnormal errors on the plasma surface
         psc_grid.normalization = B_axis ** 2 * psc_grid.plasma_boundary.area()
+        psc_grid.fac2_norm = psc_grid.fac ** 2 / psc_grid.normalization
         psc_grid.B_TF = B_TF
 
         # Random or B_TF aligned initialization of the coil orientations
@@ -460,7 +461,7 @@ def geo_setup_manual(
                                      (np.random.rand(
                                          psc_grid.num_psc) - 0.5) * 2 * np.pi
         )
-        N = 20  # Must be larger for some unit tests to converge
+        N = 50  # Must be larger for some unit tests to converge
         psc_grid.phi = np.linspace(0, 2 * np.pi, N, endpoint=False)
         psc_grid.dphi = psc_grid.phi[1] - psc_grid.phi[0]
         psc_grid.rho = np.linspace(0, R, N, endpoint=False)
@@ -512,18 +513,17 @@ def geo_setup_manual(
         # qa scaled to 0.1 T on-axis magnetic field strength?
         total_current = 1e6
         base_curves = create_equally_spaced_curves(
-            ncoils, psc_grid.plasma_boundary.nfp, 
-            stellsym=psc_grid.plasma_boundary.stellsym, 
+            ncoils, 1,  #psc_grid.plasma_boundary.nfp, 
+            stellsym=False,  #psc_grid.plasma_boundary.stellsym, 
             R0=R0, R1=R1, order=order, numquadpoints=128
         )
         base_currents = [(Current(total_current / ncoils * 1e-5) * 1e5) for _ in range(ncoils)]
-        I_TF_total = total_current
         total_current = Current(total_current)
         total_current.fix_all()
         default_coils = coils_via_symmetries(
             base_curves, base_currents, 
-            psc_grid.plasma_boundary.nfp, 
-            psc_grid.plasma_boundary.stellsym
+            1,  #psc_grid.plasma_boundary.nfp, 
+            False,  #psc_grid.plasma_boundary.stellsym
         )
         # fix all the coil shapes so only the currents are optimized
         for i in range(ncoils):
@@ -564,6 +564,7 @@ def geo_setup_manual(
         # Normalization of the ||A*Linv*psi - b||^2 objective 
         # representing Bnormal errors on the plasma surface
         psc_grid.normalization = B_axis ** 2 * psc_grid.plasma_boundary.area()
+        psc_grid.fac2_norm = psc_grid.fac ** 2 / psc_grid.normalization
         psc_grid.B_TF = B_TF
 
         # Order of the coils. For unit tests, needs > 400
@@ -634,10 +635,10 @@ def setup_currents_and_fields(self):
                     contig(self.deltas_total[q * nn: (q + 1) * nn]),
                     contig(self.plasma_unitnormals),
                     self.R,
-                )  # accounts for sign change of the currents
+                )  # accounts for sign change of the currents (does it?)
                 q = q + 1
         self.A_matrix = 2 * A_matrix
-        self.Bn_PSC = (self.A_matrix @ self.I).reshape(-1) 
+        self.Bn_PSC = (self.A_matrix @ self.I).reshape(-1) # need minus signs?
 
     def setup_curves(self):
         """ 
@@ -796,7 +797,7 @@ def least_squares(self, kappas, verbose=False):
         deltas = kappas[len(kappas) // 2:]
         self.setup_orientations(alphas, deltas)
         Ax_b = (self.Bn_PSC + self.b_opt) * self.grid_normalization
-        BdotN2 = 0.5 * Ax_b.T @ Ax_b / self.normalization * self.fac ** 2
+        BdotN2 = 0.5 * Ax_b.T @ Ax_b * self.fac2_norm
         # outstr = f"Normalized f_B = {BdotN2:.3e} "
         # for i in range(len(kappas)):
         #     outstr += f"kappas[{i:d}] = {kappas[i]:.2e} "
@@ -829,7 +830,7 @@ def least_squares_jacobian(self, kappas, verbose=False):
         deltas = kappas[len(kappas) // 2:]
         self.setup_orientations(alphas, deltas)
         # Two factors of grid normalization since it is not added to the gradients
-        Ax_b = (self.Bn_PSC + self.b_opt) / self.normalization * self.grid_normalization
+        Ax_b = (self.Bn_PSC + self.b_opt) * self.grid_normalization
         # t1 = time.time()
         A_deriv = self.A_deriv()
         Linv = self.L_inv[:self.num_psc, :self.num_psc]
@@ -855,13 +856,14 @@ def least_squares_jacobian(self, kappas, verbose=False):
         grad_alpha3 = self.A_matrix @ (Linv * psi_deriv[:self.num_psc])
         grad_delta3 = self.A_matrix @ (Linv * psi_deriv[self.num_psc:])
         grad_kappa3 = np.hstack((grad_alpha3, grad_delta3))
+        grad = self.grid_normalization[:, None] * (grad_kappa1 + grad_kappa2 + grad_kappa3)
         # t2 = time.time()
         # print('grad_A3 time = ', t2 - t1, grad_alpha3.shape, grad_kappa3.shape)
         # if verbose:
             # print('Jacobian = ', Ax_b @ (grad_kappa1 + grad_kappa2 + grad_kappa3))
         # minus sign in front because I = -Linv @ psi
-        jac = -Ax_b.T @ (grad_kappa1 + grad_kappa2 + grad_kappa3)
-        return jac
+        jac = -Ax_b.T @ grad
+        return jac * self.fac2_norm
 
     def A_deriv(self):
         """
@@ -880,7 +882,7 @@ def A_deriv(self):
         q = 0
         for fp in range(self.nfp):
             for stell in self.stell_list:
-                dA_dkappa += sopp.dA_dkappa(
+                dA = sopp.dA_dkappa(
                     contig(self.grid_xyz_all[q * nn: (q + 1) * nn, :]),
                     contig(self.plasma_points),
                     contig(self.alphas_total[q * nn: (q + 1) * nn]),
@@ -890,6 +892,14 @@ def A_deriv(self):
                     contig(self.quad_weights_phi),
                     self.R,
                 )
+                dA_dkappa[:, :self.num_psc] += dA[:, :self.num_psc] * (-1) ** fp
+                dA_dkappa[:, self.num_psc:] += dA[:, self.num_psc:] * (-1) ** fp * stell
+                # print(fp, stell, dA[0, 0], self.grid_xyz_all[q * nn: q * nn + 1, :])
+                # print(self.plasma_points[0, :], 
+                #       self.alphas_total[q * nn: q * nn + 1],
+                #       self.deltas_total[q * nn: q * nn + 1],
+                #       self.plasma_unitnormals[0, :],
+                #       )
                 q = q + 1
         return dA_dkappa * np.pi # rescale by pi for gauss-leg quadrature
 
@@ -940,7 +950,7 @@ def psi_deriv(self):
         q = 0
         for fp in range(self.nfp):
             for stell in self.stell_list:
-                psi_deriv = sopp.dpsi_dkappa(
+                psi_deriv += sopp.dpsi_dkappa(
                     contig(self.I_TF),
                     contig(self.dl_TF),
                     contig(self.gamma_TF),
@@ -1020,6 +1030,8 @@ def setup_orientations(self, alphas, deltas):
             contig(self.rho),
             contig(self.coil_normals),
         )
+        # print('B_TF at flux grid = ', self.B_TF.B().reshape(len(alphas), N, N, 3))
+
         self.psi *= self.dphi * self.drho
         t2 = time.time()
         # print('Flux2 time = ', t2 - t1)
@@ -1090,20 +1102,24 @@ def update_alphas_deltas(self):
                 self.coil_normals_all[nn * q: nn * (q + 1), 0] = self.coil_normals[:, 0] * np.cos(phi0) * stell - self.coil_normals[:, 1] * np.sin(phi0) 
                 self.coil_normals_all[nn * q: nn * (q + 1), 1] = self.coil_normals[:, 0] * np.sin(phi0) * stell + self.coil_normals[:, 1] * np.cos(phi0) 
                 self.coil_normals_all[nn * q: nn * (q + 1), 2] = self.coil_normals[:, 2]
-                normals = self.coil_normals_all[q * nn: (q + 1) * nn, :]
-                # normals = self.coil_normals
-                # get new deltas by flipping sign of deltas, then rotation alpha by phi0
-                deltas = np.arctan2(normals[:, 0], normals[:, 2]) # + np.pi
-                deltas[abs(deltas) == np.pi] = 0.0
-                if fp == 0 and stell == 1:
-                    shift_deltas = self.deltas - deltas  # +- pi
-                deltas += shift_deltas
-                alphas = -np.arctan2(normals[:, 1] * np.cos(deltas), normals[:, 2])
-                # alphas = -np.arcsin(normals[:, 1])
-                alphas[abs(alphas) == np.pi] = 0.0
-                if fp == 0 and stell == 1:
-                    shift_alphas = self.alphas - alphas  # +- pi
-                alphas += shift_alphas
+                # normals = self.coil_normals_all[q * nn: (q + 1) * nn, :]
+                # # normals = self.coil_normals
+                # # get new deltas by flipping sign of deltas, then rotation alpha by phi0
+                # deltas = np.arctan2(normals[:, 0], normals[:, 2]) # + np.pi
+                # deltas[abs(deltas) == np.pi] = 0.0
+                # if fp == 0 and stell == 1:
+                #     shift_deltas = self.deltas - deltas  # +- pi
+                # deltas += shift_deltas
+                # alphas = -np.arctan2(normals[:, 1] * np.cos(deltas), normals[:, 2])
+                # # alphas = -np.arcsin(normals[:, 1])
+                # alphas[abs(alphas) == np.pi] = 0.0
+                # if fp == 0 and stell == 1:
+                #     shift_alphas = self.alphas - alphas  # +- pi
+                # alphas += shift_alphas
+                # self.alphas_total[nn * q: nn * (q + 1)] = alphas
+                # self.deltas_total[nn * q: nn * (q + 1)] = deltas
+                alphas = self.alphas * (-1) ** fp
+                deltas = self.deltas * stell * (-1) ** fp
                 self.alphas_total[nn * q: nn * (q + 1)] = alphas
                 self.deltas_total[nn * q: nn * (q + 1)] = deltas
                 q = q + 1

src/simsopt/geo/wp_grid.py

@@ -68,7 +68,7 @@ def _setup_uniform_grid(self):
 
         # This is not a guarantee that coils will not touch but inductance
         # matrix blows up if they do so it is easy to tell when they do
-        self.R = min(Nmin / 4.0, self.poff / 4.0)
+        self.R = min(Nmin / 2.0, self.poff / 4.0)
         self.a = self.R / 10.0  # Hard-coded aspect ratio of 100 right now
         print('Major radius of the coils is R = ', self.R)
         print('Coils are spaced so that every coil of radius R '
@@ -631,7 +631,7 @@ def setup_currents_and_fields(self):
                 #     self.alphas_total[q * nn: (q + 1) * nn],
                 #     self.deltas_total[q * nn: (q + 1) * nn])
                 # t1 = time.time()
-                A_matrix += sopp.A_matrix(
+                A_matrix += sopp.A_matrix_simd(
                     contig(self.grid_xyz_all[q * nn: (q + 1) * nn, :]),
                     contig(self.plasma_points),
                     contig(self.alphas_total[q * nn: (q + 1) * nn]),
@@ -818,16 +818,16 @@ def least_squares(self, kappa_I, verbose=False):
                 array of kappas.
 
         """
-        # t1 = time.time()
+        t1 = time.time()
         ind3 = len(kappa_I) // 3
         ind3_2 = 2 * ind3
         self.I = kappa_I[ind3_2:]
         self.setup_orientations(kappa_I[:ind3], kappa_I[ind3:ind3_2])
         Ax_b = (self.Bn_WP + self.b_opt) * self.grid_normalization 
         BdotN2 = 0.5 * Ax_b.T @ Ax_b * self.fac2_norm
         self.BdotN2_list.append(BdotN2)
-        # t2 = time.time()
-        # print('obj time = ', t2 - t1)
+        t2 = time.time()
+        print('obj time = ', t2 - t1)
         return BdotN2
 
     def python_A_matrix(self, points, alphas, deltas):
@@ -923,14 +923,18 @@ def least_squares_jacobian(self, kappa_I, verbose=False):
         self.I = kappa_I[ind3_2:]
         self.setup_orientations(kappa_I[:ind3], kappa_I[ind3:ind3_2])
         Ax_b = (self.Bn_WP + self.b_opt) * self.grid_normalization
+        t1_grad = time.time()
         grad_alpha_delta = self.A_deriv() * np.hstack((self.I, self.I))
+        t2_grad = time.time()
+        print('Aderiv time = ', t2_grad - t1_grad)
         grad_I = self.A_matrix
         # Should be shape (num_plasma_points, 3 * num_wp)
         # grad_kappa = np.hstack((grad_alpha1, grad_delta1))
         grad = np.hstack((grad_alpha_delta, grad_I)) * 1.0e6
         grad = self.grid_normalization[:, None] * grad
         jac = Ax_b.T @ grad
         t2 = time.time()
+        print('jac time = ', t2 - t1)
         return jac * self.fac2_norm
 
     def A_deriv(self):
@@ -950,7 +954,7 @@ def A_deriv(self):
         q = 0
         for fp in range(self.nfp):
             for stell in self.stell_list:
-                dA_dkappa += sopp.dA_dkappa(
+                dA = sopp.dA_dkappa_simd(
                     contig(self.grid_xyz_all[q * nn: (q + 1) * nn, :]),
                     contig(self.plasma_points),
                     contig(self.alphas_total[q * nn: (q + 1) * nn]),
@@ -960,6 +964,8 @@ def A_deriv(self):
                     contig(self.quad_weights_phi),
                     self.R,
                 )
+                dA_dkappa[:, :self.num_wp] += dA[:, :self.num_wp] * (-1) ** fp
+                dA_dkappa[:, self.num_wp:] += dA[:, self.num_wp:] * (-1) ** fp * stell
                 q = q + 1
         return dA_dkappa * np.pi # rescale by pi for gauss-leg quadrature