Minor updates, mostly trying to directly calculate psi and L to figur…

…e out the symmetries.

examples/2_Intermediate/QH_psc_example.py

@@ -22,12 +22,12 @@
     nphi = 32  # nphi = ntheta >= 64 needed for accurate full-resolution runs
     ntheta = 32
     # Make higher resolution surface for plotting Bnormal
-    qphi = 2 * nphi
+    qphi = nphi * 8
     quadpoints_phi = np.linspace(0, 1, qphi, endpoint=True)
-    quadpoints_theta = np.linspace(0, 1, ntheta, endpoint=True)
+    quadpoints_theta = np.linspace(0, 1, ntheta * 4, endpoint=True)
 
-coff = 1.2  # PM grid starts offset ~ 10 cm from the plasma surface
-poff = 0.8  # PM grid end offset ~ 15 cm from the plasma surface
+coff = 1.5  # PM grid starts offset ~ 10 cm from the plasma surface
+poff = 1.0  # PM grid end offset ~ 15 cm from the plasma surface
 input_name = 'input.LandremanPaul2021_QH_reactorScale_lowres'
 
 # Read in the plas/ma equilibrium file
@@ -37,7 +37,8 @@
 s = SurfaceRZFourier.from_vmec_input(
     surface_filename, range=range_param, nphi=nphi, ntheta=ntheta
 )
-print('s.R = ', s.get_dofs(), s.get_rc(0, 0))
+print('s.R = ', s.get_rc(0, 0))
+print('s.r = ', s.get_rc(1, 0))
 
 # input_name = 'tests/test_files/input.circular_tokamak'
 # s_inner = SurfaceRZFourier(
@@ -100,51 +101,54 @@
 make_Bnormal_plots(bs, s_plot, out_dir, "biot_savart_TF_optimized")
 
 # check after-optimization average on-axis magnetic field strength
-calculate_on_axis_B(bs, s)
+B_axis = calculate_on_axis_B(bs, s)
+make_Bnormal_plots(bs, s_plot, out_dir, "biot_savart_TF_optimized", B_axis)
 
 # Finally, initialize the psc class
-kwargs_geo = {"Nx": 10, "out_dir": out_str, "random_initialization": True} 
+kwargs_geo = {"Nx": 4, "out_dir": out_str, "random_initialization": True} 
 # note Bnormal below is additional Bnormal (not from the TF coils or the PSCs)
 psc_array = PSCgrid.geo_setup_between_toroidal_surfaces(
     s, np.zeros(Bnormal.shape), bs, s_inner, s_outer,  **kwargs_geo
 )
 
 # print('Currents = ', psc_array.I)
 psc_array.setup_psc_biotsavart()
-make_Bnormal_plots(psc_array.B_PSC, s_plot, out_dir, "biot_savart_PSC_initial")
-make_Bnormal_plots(psc_array.B_TF, s_plot, out_dir, "biot_savart_InterpolatedTF")
-make_Bnormal_plots(bs + psc_array.B_PSC, s_plot, out_dir, "PSC_and_TF_initial")
+make_Bnormal_plots(psc_array.B_PSC, s_plot, out_dir, "biot_savart_PSC_initial", B_axis)
+make_Bnormal_plots(psc_array.B_PSC_all, s_plot, out_dir, "biot_savart_PSC_initial_all", B_axis)
+# make_Bnormal_plots(psc_array.B_PSC_direct, s_plot, out_dir, "biot_savart_PSC_initial_direct")
+# make_Bnormal_plots(psc_array.B_TF, s_plot, out_dir, "biot_savart_InterpolatedTF", B_axis)
+make_Bnormal_plots(bs + psc_array.B_PSC, s_plot, out_dir, "PSC_and_TF_initial", B_axis)
 
 # Try and optimize with scipy
-t1 = time.time()
-x0 = np.ravel(np.array([psc_array.alphas, psc_array.deltas]))
-psc_array.least_squares(x0)
-t2 = time.time()
-print('Time for call to least-squares = ', t2 - t1)
-psc_array.least_squares(x0 + np.random.rand(len(x0)))
-
-psc_array.setup_psc_biotsavart()
-make_Bnormal_plots(psc_array.B_PSC_all, s_plot, out_dir, "PSC_initial")
-make_Bnormal_plots(bs + psc_array.B_PSC_all, s_plot, out_dir, "PSC_and_TF_initial")
-
-from scipy.optimize import minimize
-print('beginning optimization: ')
-options = {"disp": True}
-# x0 = np.random.rand(len(np.ravel(np.array([psc_array.alphas, psc_array.deltas]
-#                                           )))) * 2 * np.pi
+# t1 = time.time()
+# x0 = np.ravel(np.array([psc_array.alphas, psc_array.deltas]))
+# psc_array.least_squares(x0)
+# t2 = time.time()
+# print('Time for call to least-squares = ', t2 - t1)
+# psc_array.least_squares(x0 + np.random.rand(len(x0)))
+
+# psc_array.setup_psc_biotsavart()
+# make_Bnormal_plots(psc_array.B_PSC_all, s_plot, out_dir, "PSC_initial", B_axis)
+# make_Bnormal_plots(bs + psc_array.B_PSC_all, s_plot, out_dir, "PSC_and_TF_initial", B_axis)
+
+# from scipy.optimize import minimize
+# print('beginning optimization: ')
+# options = {"disp": True}
+# # x0 = np.random.rand(len(np.ravel(np.array([psc_array.alphas, psc_array.deltas]
+# #                                           )))) * 2 * np.pi
+# # print(x0)
+# x0 = psc_array.kappas
 # print(x0)
-x0 = psc_array.kappas
-print(x0)
-x_opt = minimize(psc_array.least_squares, x0, options=options)
-print(x_opt)
-# print('currents = ', psc_array.I)
-psc_array.setup_curves(symmetrized=True)
-psc_array.plot_curves('final_')
-# Check that direct Bn calculation agrees with optimization calculation
-psc_array.setup_psc_biotsavart()
-fB = SquaredFlux(s, psc_array.B_PSC_all + bs, np.zeros((nphi, ntheta))).J()
-print(fB)
-make_Bnormal_plots(psc_array.B_PSC_all, s_plot, out_dir, "PSC_final")
-make_Bnormal_plots(bs + psc_array.B_PSC_all, s_plot, out_dir, "PSC_and_TF_final")
-print('end')
+# x_opt = minimize(psc_array.least_squares, x0, options=options)
+# print(x_opt)
+# # print('currents = ', psc_array.I)
+# psc_array.setup_curves(symmetrized=True)
+# psc_array.plot_curves('final_')
+# # Check that direct Bn calculation agrees with optimization calculation
+# psc_array.setup_psc_biotsavart()
+# fB = SquaredFlux(s, psc_array.B_PSC_all + bs, np.zeros((nphi, ntheta))).J()
+# print(fB)
+# make_Bnormal_plots(psc_array.B_PSC_all, s_plot, out_dir, "PSC_final")
+# make_Bnormal_plots(bs + psc_array.B_PSC_all, s_plot, out_dir, "PSC_and_TF_final")
+# print('end')
 # plt.show()

src/simsopt/geo/curve.py

@@ -845,8 +845,7 @@ def wrap(data):
     else:
         coil_data = np.zeros(data.shape)
         for i in range(len(scalar_data)):
-            # print(i, scalar_data[i], abs(scalar_data[i]))
-            coil_data[i * ppl[i]: (i + 1) * ppl[i]] = abs(scalar_data[i])
+            coil_data[i * ppl[i]: (i + 1) * ppl[i]] = scalar_data[i]
         coil_data = np.ascontiguousarray(coil_data)
         polyLinesToVTK(str(filename), x, y, z, pointsPerLine=ppl, pointData={'idx': data, 'I': coil_data})
 

src/simsopt/geo/psc_grid.py

@@ -80,7 +80,7 @@ def _setup_uniform_grid(self):
         # the old cells.
         #### Note that Cartesian cells can only do nfp = 2, 4, 6, ... 
         #### and correctly be rotated to have the right symmetries
-        print(self.dx, self.dy, x_max, y_max, x_min, y_min)
+        print(self.dx, self.dy, x_max, y_max, x_min, y_min, self.dz, self.R)
         # exit()
         if self.plasma_boundary.nfp > 1:
             # Throw away any points not in the section phi = [0, pi / n_p] and
@@ -299,10 +299,17 @@ def geo_setup_between_toroidal_surfaces(
         rs = np.linalg.norm(gamma_outer[:, :2], axis=-1)
         zs = gamma_outer[:, 2]
         rrange = (0, np.max(rs) + R, n)  # need to also cover the plasma
-        phirange = (0, 2 * np.pi, n * 2)
-        zrange = (np.min(zs) - R, np.max(zs) + R, n // 2)
+        if psc_grid.nfp > 1:
+            phirange = (0, 2*np.pi/psc_grid.nfp, n*2)
+        else:
+            phirange = (0, 2 * np.pi, n * 2)
+        if psc_grid.stellsym:
+            zrange = (0, np.max(zs) + R, n // 2)
+        else:
+            zrange = (np.min(zs) - R, np.max(zs) + R, n // 2)
         psc_grid.B_TF = InterpolatedField(
-            B_TF, degree, rrange, phirange, zrange, True, nfp=1, stellsym=False
+            B_TF, degree, rrange, phirange, zrange, 
+            True, nfp=psc_grid.nfp, stellsym=psc_grid.stellsym
         )
         B_TF.set_points(psc_grid.grid_xyz)
         B = B_TF.B()
@@ -536,7 +543,7 @@ def setup_psc_biotsavart(self):
             currents.append(Current(self.I[i]))
 
         coils = coils_via_symmetries(
-            self.curves, currents, nfp=1, stellsym=False
+            self.curves, currents, nfp=self.nfp, stellsym=self.stellsym
         )
         self.B_PSC = BiotSavart(coils)
         self.B_PSC.set_points(self.plasma_points)
@@ -608,12 +615,18 @@ def plot_curves(self, filename=''):
         # if hasattr(self, 'all_curves'):
         self.psi_all = np.zeros(self.num_psc * self.symmetry)
         self.I_all = np.zeros(self.num_psc * self.symmetry)
+        L_total = np.zeros((self.num_psc * self.symmetry, 
+                            self.num_psc * self.symmetry))
+        B_PSC = np.zeros(self.plasma_boundary.gamma().reshape(-1, 3).shape)
         contig = np.ascontiguousarray
+        alphas_total = []
+        deltas_total = []
+        psi_total = np.zeros(self.num_psc * self.symmetry)
         q = 0
         for fp in range(self.nfp):
             for stell in self.stell_list:
                 self.psi_all[q * self.num_psc: (q + 1) * self.num_psc] = self.psi
-                self.I_all[q * self.num_psc: (q + 1) * self.num_psc] = self.I * stell
+                # self.I_all[q * self.num_psc: (q + 1) * self.num_psc] = self.I * stell
 
                 phi0 = (2.0 * np.pi / self.nfp) * fp
                 # x' = R_p^T R_s^T x
@@ -623,9 +636,9 @@ def plot_curves(self, filename=''):
                 oz = self.grid_xyz[:, 2] * stell
                 # n' = R_s R_p n
                 # get new normal vectors by flipping the x component, then rotating by phi0
-                self.coil_normals_all[self.num_psc * q: self.num_psc * (q + 1), 0] = self.coil_normals[:, 0] * np.cos(phi0) * stell - self.coil_normals[:, 1] * np.sin(phi0) 
-                self.coil_normals_all[self.num_psc * q: self.num_psc * (q + 1), 1] = self.coil_normals[:, 0] * np.sin(phi0) * stell + self.coil_normals[:, 1] * np.cos(phi0) 
-                self.coil_normals_all[self.num_psc * q: self.num_psc * (q + 1), 2] = self.coil_normals[:, 2]
+                # self.coil_normals_all[self.num_psc * q: self.num_psc * (q + 1), 0] = self.coil_normals[:, 0] * np.cos(phi0) * stell - self.coil_normals[:, 1] * np.sin(phi0) 
+                # self.coil_normals_all[self.num_psc * q: self.num_psc * (q + 1), 1] = self.coil_normals[:, 0] * np.sin(phi0) * stell + self.coil_normals[:, 1] * np.cos(phi0) 
+                # self.coil_normals_all[self.num_psc * q: self.num_psc * (q + 1), 2] = self.coil_normals[:, 2]
                 xyz = np.array([ox, oy, oz]).T
                 normals = self.coil_normals_all[q * self.num_psc: (q + 1) * self.num_psc, :]
                 alphas = np.arcsin(-normals[:, 1])
@@ -643,6 +656,8 @@ def plot_curves(self, filename=''):
                     contig(self.flux_phi), 
                     contig(normals)
                 )
+                alphas_total.append(alphas)
+                deltas_total.append(deltas)
                 self.B_TF.set_points(contig(np.array(flux_grid).reshape(-1, 3)))
                 N = len(self.rho)
                 Nflux = len(self.flux_phi)
@@ -652,20 +667,51 @@ def plot_curves(self, filename=''):
                     contig(self.rho),
                     contig(normals)
                 )
+                psi_total[q * self.num_psc: (q + 1) * self.num_psc] = psi_check
                 print(fp, stell, psi_check)
+                B_PSC += sopp.B_PSC(
+                    contig(xyz),
+                    contig(self.plasma_boundary.gamma().reshape(-1, 3)),
+                    contig(self.alphas),
+                    contig(self.deltas),
+                    contig(self.I * stell),
+                    self.R,
+                    float(phi0),
+                    float(stell)
+                )
                 q = q + 1
+        alphas_total = np.ravel(np.array(alphas_total))
+        deltas_total = np.ravel(np.array(deltas_total))
+        self.B_PSC_direct = B_PSC
+        L_total = sopp.L_matrix(
+            contig(self.grid_xyz_all), 
+            contig(alphas_total), 
+            contig(deltas_total),
+            contig(self.phi),
+            self.R,
+        ) * self.dphi ** 2 / (4 * np.pi)
+        L_total = (L_total + L_total.T)
+        np.fill_diagonal(L_total, np.log(8.0 * self.R / self.a) - 2.0)
+        self.L_total = L_total * self.mu0 * self.R * self.Nt ** 2
+        print(self.L_total, self.L_total.shape)
+        I_total = - np.linalg.solve(L_total, psi_total)
+        print(I_total, I_total.shape)
+        self.I_all = I_total
+        print(self.I)
+              #self.num_psc, self.symmetry, alphas_total.shape,
+              #self.grid_xyz_all.shape)
         # exit()
         currents = []
-        for i in range(len(self.I)):
-            currents.append(Current(self.I[i]))
-        all_coils = coils_via_symmetries(
-            self.curves, currents, nfp=self.nfp, stellsym=self.stellsym
-        )
-        # for i in range(self.num_psc * self.symmetry):
-        #     currents.append(Current(self.I_all[i]))
+        # for i in range(len(self.I)):
+        #     currents.append(Current(self.I[i]))
         # all_coils = coils_via_symmetries(
-        #     self.all_curves, currents, nfp=1, stellsym=False
+        #     self.curves, currents, nfp=self.nfp, stellsym=self.stellsym
         # )
+        for i in range(self.symmetry * self.num_psc):
+            currents.append(Current(self.I_all[i]))
+        all_coils = coils_via_symmetries(
+            self.all_curves, currents, nfp=1, stellsym=False
+        )
         self.B_PSC_all = BiotSavart(all_coils)
 
         curves_to_vtk(self.all_curves, self.out_dir + filename + "all_psc_curves", close=True, scalar_data=self.I_all)

src/simsopt/util/permanent_magnet_helper_functions.py

@@ -313,7 +313,7 @@ def initialize_coils(config_flag, TEST_DIR, s, out_dir=''):
         # qh needs to be scaled to 0.1 T on-axis magnetic field strength
         from simsopt.mhd.vmec import Vmec
         vmec_file = 'wout_LandremanPaul2021_QH_reactorScale_lowres_reference.nc'
-        total_current = Vmec(TEST_DIR / vmec_file).external_current() / (2 * s.nfp) / 8.75 / 5.69674966667
+        total_current = Vmec(TEST_DIR / vmec_file).external_current() / (2 * s.nfp) / 1.315
         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-1)]
         total_current = Current(total_current)
@@ -541,7 +541,7 @@ def run_Poincare_plots(s_plot, bs, b_dipole, comm, filename_poincare, out_dir=''
     trace_fieldlines(bsh, 'bsh_PMs_' + filename_poincare, s_plot, comm, out_dir)
 
 
-def make_Bnormal_plots(bs, s_plot, out_dir='', bs_filename="Bnormal"):
+def make_Bnormal_plots(bs, s_plot, out_dir='', bs_filename="Bnormal", B_axis=None):
     """
     Plot Bnormal on plasma surface from a MagneticField object.
     Do this quite a bit in the permanent magnet optimization
@@ -554,11 +554,15 @@ def make_Bnormal_plots(bs, s_plot, out_dir='', bs_filename="Bnormal"):
         out_dir: Path or string for the output directory for saved files.
         bs_filename: String denoting the name of the output file. 
     """
+    from simsopt.field import BiotSavart, InterpolatedField, MagneticFieldSum
     out_dir = Path(out_dir)
     nphi = len(s_plot.quadpoints_phi)
     ntheta = len(s_plot.quadpoints_theta)
     bs.set_points(s_plot.gamma().reshape((-1, 3)))
-    pointData = {"B_N": np.sum(bs.B().reshape((nphi, ntheta, 3)) * s_plot.unitnormal(), axis=2)[:, :, None]}
+    Bn = np.sum(bs.B().reshape((nphi, ntheta, 3)) * s_plot.unitnormal(), axis=2)[:, :, None]
+    if B_axis is not None:
+        Bn = Bn / (B_axis * s_plot.area())
+    pointData = {"B_N": Bn}
     s_plot.to_vtk(out_dir / bs_filename, extra_data=pointData)