From 8fb5d362efb611150661c59897ce10ee811241c8 Mon Sep 17 00:00:00 2001 From: Alan Kaptanoglu Date: Sun, 10 Mar 2024 19:00:00 -0400 Subject: [PATCH] Caught the unit tests up with the changes to the main code. Checking now that stellarator symmetry and field period symmetry is imposed correctly when calculating the magnetic fields. So far not quite getting the right answer, but all the machinery is there. Need to convince myself about what symmetries do to the inductance matrix, since I am not sure about how the current matrix transforms. Usually the currents in the coils are imposed, so you can simply take a set of N coils and apply the symmetries to get the rest. Here, the currents in the coils are themselves dependent on the coils that you get after you symmetrize. --- examples/2_Intermediate/psc_example.py | 34 ++- src/simsopt/geo/psc_grid.py | 386 +++++++++++-------------- src/simsoptpp/psc.cpp | 20 +- src/simsoptpp/psc.h | 4 +- src/simsoptpp/python.cpp | 2 +- tests/geo/test_psc_grid.py | 165 +++++++---- 6 files changed, 317 insertions(+), 294 deletions(-) diff --git a/examples/2_Intermediate/psc_example.py b/examples/2_Intermediate/psc_example.py index a1b517c81..36a3b7096 100644 --- a/examples/2_Intermediate/psc_example.py +++ b/examples/2_Intermediate/psc_example.py @@ -19,8 +19,8 @@ nphi = 4 # nphi = ntheta >= 64 needed for accurate full-resolution runs ntheta = nphi else: - nphi = 64 # nphi = ntheta >= 64 needed for accurate full-resolution runs - ntheta = 64 + nphi = 16 # nphi = ntheta >= 64 needed for accurate full-resolution runs + ntheta = 16 # Make higher resolution surface for plotting Bnormal qphi = 2 * nphi quadpoints_phi = np.linspace(0, 1, qphi, endpoint=True) @@ -101,7 +101,7 @@ calculate_on_axis_B(bs, s) # Finally, initialize the psc class -kwargs_geo = {"Nx": 10, "out_dir": out_str, "random_initialization": True} +kwargs_geo = {"Nx": 6, "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 @@ -121,24 +121,28 @@ print('Time for call to least-squares = ', t2 - t1) psc_array.least_squares(x0 + np.random.rand(len(x0))) -# 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 +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 # print(x0) -# x_opt = minimize(psc_array.least_squares, x0, options=options) -# print(x_opt) +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() +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_initial") -make_Bnormal_plots(bs + psc_array.B_PSC_all, s_plot, out_dir, "PSC_and_TF_initial") +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() diff --git a/src/simsopt/geo/psc_grid.py b/src/simsopt/geo/psc_grid.py index d14a370f1..eba9d0497 100644 --- a/src/simsopt/geo/psc_grid.py +++ b/src/simsopt/geo/psc_grid.py @@ -214,6 +214,10 @@ def geo_setup_between_toroidal_surfaces( psc_grid.nfp = plasma_boundary.nfp psc_grid.stellsym = plasma_boundary.stellsym psc_grid.symmetry = plasma_boundary.nfp * (plasma_boundary.stellsym + 1) + if psc_grid.stellsym: + psc_grid.stell_list = [1, -1] + else: + psc_grid.stell_list = [1] psc_grid.nphi = len(psc_grid.plasma_boundary.quadpoints_phi) psc_grid.ntheta = len(psc_grid.plasma_boundary.quadpoints_theta) psc_grid.plasma_unitnormals = psc_grid.plasma_boundary.unitnormal().reshape(-1, 3) @@ -328,22 +332,7 @@ def geo_setup_between_toroidal_surfaces( # psc_grid.coil_normals[:, 1], psc_grid.coil_normals[:, 0] # ) - nn = len(psc_grid.grid_xyz) - psc_grid.grid_xyz_all = np.zeros((nn * psc_grid.symmetry, 3)) - psc_grid.coil_normals_all = np.zeros((nn * psc_grid.symmetry, 3)) - q = 0 - for fp in range(psc_grid.nfp): - for stell in [1, -1]: - phi0 = (2 * np.pi / psc_grid.nfp) * fp - # get new locations by flipping the y and z components, then rotating by phi0 - psc_grid.grid_xyz_all[nn * q: nn * (q + 1), 0] = psc_grid.grid_xyz[:, 0] * np.cos(phi0) - psc_grid.grid_xyz[:, 1] * np.sin(phi0) * stell - psc_grid.grid_xyz_all[nn * q: nn * (q + 1), 1] = psc_grid.grid_xyz[:, 0] * np.sin(phi0) + psc_grid.grid_xyz[:, 1] * np.cos(phi0) * stell - psc_grid.grid_xyz_all[nn * q: nn * (q + 1), 2] = psc_grid.grid_xyz[:, 2] * stell - # get new normal vectors by flipping the x component, then rotating by phi0 - psc_grid.coil_normals_all[nn * q: nn * (q + 1), 0] = psc_grid.coil_normals[:, 0] * np.cos(phi0) * stell - psc_grid.coil_normals[:, 1] * np.sin(phi0) - psc_grid.coil_normals_all[nn * q: nn * (q + 1), 1] = psc_grid.coil_normals[:, 0] * np.sin(phi0) * stell + psc_grid.coil_normals[:, 1] * np.cos(phi0) - psc_grid.coil_normals_all[nn * q: nn * (q + 1), 2] = psc_grid.coil_normals[:, 2] - q = q + 1 + psc_grid.setup_full_grid() t2 = time.time() print('Initialize grid time = ', t2 - t1) @@ -408,15 +397,23 @@ def geo_setup_manual( TEST_DIR = (Path(__file__).parent / ".." / ".." / ".." / "tests" / "test_files").resolve() surface_filename = TEST_DIR / input_name default_surf = SurfaceRZFourier.from_vmec_input( - surface_filename, range='full torus', nphi=16, ntheta=16 + surface_filename, range='half period', nphi=16, ntheta=16 ) psc_grid.plasma_boundary = kwargs.pop("plasma_boundary", default_surf) - psc_grid.nfp = plasma_boundary.nfp - psc_grid.stellsym = plasma_boundary.stellsym + psc_grid.nfp = psc_grid.plasma_boundary.nfp + psc_grid.stellsym = psc_grid.plasma_boundary.stellsym psc_grid.nphi = len(psc_grid.plasma_boundary.quadpoints_phi) psc_grid.ntheta = len(psc_grid.plasma_boundary.quadpoints_theta) psc_grid.plasma_unitnormals = psc_grid.plasma_boundary.unitnormal().reshape(-1, 3) psc_grid.plasma_points = psc_grid.plasma_boundary.gamma().reshape(-1, 3) + psc_grid.symmetry = psc_grid.plasma_boundary.nfp * (psc_grid.plasma_boundary.stellsym + 1) + if psc_grid.stellsym: + psc_grid.stell_list = [1, -1] + else: + psc_grid.stell_list = [1] + Ngrid = psc_grid.nphi * psc_grid.ntheta + Nnorms = np.ravel(np.sqrt(np.sum(psc_grid.plasma_boundary.normal() ** 2, axis=-1))) + psc_grid.grid_normalization = np.sqrt(Nnorms / Ngrid) # generate planar TF coils ncoils = 2 R0 = 1.0 @@ -482,6 +479,7 @@ def geo_setup_manual( # Initialize curve objects corresponding to each PSC coil for # plotting in 3D + psc_grid.setup_full_grid() psc_grid.setup_orientations(psc_grid.alphas, psc_grid.deltas) psc_grid.setup_curves() psc_grid.plot_curves() @@ -496,14 +494,28 @@ def setup_currents_and_fields(self): self.L_inv = np.linalg.inv(self.L) self.I = -self.L_inv @ self.psi contig = np.ascontiguousarray - self.Bn_PSC = (sopp.Bn_PSC( - contig(self.grid_xyz), - contig(self.plasma_points), - contig(self.alphas), - contig(self.deltas), - contig(self.plasma_unitnormals), - self.R - ) @ self.I) * self.grid_normalization + # A_matrix has shape (num_plasma_points, num_coils) + A_matrix = np.zeros((self.nphi * self.ntheta, self.num_psc)) + # Need to rotate and flip it + for fp in range(self.nfp): + for stell in self.stell_list: + phi0 = (2.0 * np.pi / self.nfp) * fp + # get new locations by flipping the y and z components, then rotating by phi0 + ox = self.grid_xyz[:, 0] * np.cos(phi0) - self.grid_xyz[:, 1] * np.sin(phi0) * stell + oy = self.grid_xyz[:, 0] * np.sin(phi0) + self.grid_xyz[:, 1] * np.cos(phi0) * stell + oz = self.grid_xyz[:, 2] * stell + xyz = np.array([ox, oy, oz]).T + A_matrix += sopp.A_matrix( + contig(xyz), + contig(self.plasma_points), + contig(self.alphas), + contig(self.deltas), + contig(self.plasma_unitnormals), + self.R, + phi0, + stell + ) * stell # accounts for sign change of the currents + self.Bn_PSC = A_matrix @ self.I def setup_psc_biotsavart(self): from simsopt.field import coils_via_symmetries, Current @@ -526,197 +538,119 @@ def setup_psc_biotsavart(self): def setup_curves(self, symmetrized=False): """ """ from . import CurvePlanarFourier - from simsopt.field import Current, coils_via_symmetries + from simsopt.field import apply_symmetries_to_curves order = 1 ncoils = self.num_psc - if symmetrized: - curves = [CurvePlanarFourier( - order*self.ppp, order, nfp=1, stellsym=False - ) for i in range(ncoils * self.symmetry)] - q = 0 - for fp in range(self.nfp): - for stell in [1.0, -1.0]: - for ic in range(ncoils): - xyz = self.grid_xyz[ic, :] - a = self.alphas[ic] - d = self.deltas[ic] - phi0 = (2 * np.pi / self.nfp) * fp - # get new locations by flipping the y and z components, then rotating by phi0 - # x' = R_fp R_s x - x = xyz[0] * np.cos(phi0) - xyz[1] * np.sin(phi0) * stell - y = xyz[0] * np.sin(phi0) + xyz[1] * np.cos(phi0) * stell - z = xyz[2] * stell - dofs = np.zeros(10) - dofs[0] = self.R - dofs[1] = 0.0 - dofs[2] = 0.0 - - # get new normal vectors by flipping the x component (R_s), - # then rotating by phi0 (R_fp) - # Need to write down rotation matrix R_fp R_s R_i - # Take product of rotation quaternion with quaternion [0, 1, 0, 0] - # x' = R_s^T R_fp^T R_i x = (R_fp R_s)^T) R_i x - phi0 = -phi0 - q1 = np.array([np.cos(a / 2.0), np.sin(a / 2.0), 0, 0]) - q2 = np.array([np.cos(d / 2.0), 0, 0, np.sin(d / 2.0)]) - q3 = np.array([np.cos(phi0 / 2.0), 0, 0, np.sin(phi0 / 2.0)]) - if stell < 0: - q4 = np.array([0, 0, 1, 0]) - q5 = np.array([0, 0, 0, 1]) - dofs[3:7] = hamilton_product(q5, hamilton_product(hamilton_product( - q4, hamilton_product(hamilton_product(q3, hamilton_product(q2, q1)), q4) - ), q5)) - else: - dofs[3:7] = hamilton_product(q3, hamilton_product(q2, q1)) - # dofs[3] = np.cos( - # a / 2.0) * np.cos( - # d / 2.0) * np.cos(phi0 / 2.0) + np.sin( - # a / 2.0) * np.sin( - # d / 2.0) * np.sin(phi0 / 2.0) - # dofs[4] = np.sin( - # a / 2.0) * np.cos( - # d / 2.0) * np.cos( - # phi0 / 2.0) - np.cos( - # a / 2.0) * np.sin( - # d / 2.0) * np.sin( - # phi0 / 2.0) - # dofs[5] = np.cos( - # a / 2.0) * np.sin( - # d / 2.0) * np.cos( - # phi0 / 2.0) + np.sin( - # a / 2.0) * np.cos( - # d / 2.0) * np.sin( - # phi0 / 2.0) - # dofs[6] = np.cos( - # a / 2.0) * np.cos( - # d / 2.0) * np.sin( - # phi0 / 2.0) - np.sin( - # a / 2.0) * np.sin( - # d / 2.0) * np.cos( - # phi0 / 2.0) - # if stell < 0: - # temp3 = dofs[3] - # temp4 = dofs[4] - # temp5 = dofs[5] - # temp6 = dofs[6] - # dofs[3] = -temp4 - # dofs[4] = temp3 - # dofs[5] = temp6 - # dofs[6] = -temp5 - # Now specify the center - dofs[7] = x - dofs[8] = y - dofs[9] = z - curves[q].set_dofs(dofs) - q += 1 - self.all_curves = curves - else: - curves = [CurvePlanarFourier(order*self.ppp, order, nfp=1, stellsym=False) for i in range(ncoils)] - for ic in range(ncoils): - xyz = self.grid_xyz[ic, :] - dofs = np.zeros(10) - dofs[0] = self.R - dofs[1] = 0.0 - dofs[2] = 0.0 - # Conversion from Euler angles in 3-2-1 body sequence to - # quaternions: - # https://en.wikipedia.org/wiki/Conversion_between_quaternions_and_Euler_angles - dofs[3] = np.cos(self.alphas[ic] / 2.0) * np.cos(self.deltas[ic] / 2.0) - dofs[4] = np.sin(self.alphas[ic] / 2.0) * np.cos(self.deltas[ic] / 2.0) - dofs[5] = np.cos(self.alphas[ic] / 2.0) * np.sin(self.deltas[ic] / 2.0) - dofs[6] = -np.sin(self.alphas[ic] / 2.0) * np.sin(self.deltas[ic] / 2.0) - # Now specify the center - dofs[7] = xyz[0] - dofs[8] = xyz[1] - dofs[9] = xyz[2] - curves[ic].set_dofs(dofs) - self.curves = curves + curves = [CurvePlanarFourier(order*self.ppp, order, nfp=1, stellsym=False) for i in range(ncoils)] + for ic in range(ncoils): + xyz = self.grid_xyz[ic, :] + dofs = np.zeros(10) + dofs[0] = self.R + dofs[1] = 0.0 + dofs[2] = 0.0 + # Conversion from Euler angles in 3-2-1 body sequence to + # quaternions: + # https://en.wikipedia.org/wiki/Conversion_between_quaternions_and_Euler_angles + dofs[3] = np.cos(self.alphas[ic] / 2.0) * np.cos(self.deltas[ic] / 2.0) + dofs[4] = np.sin(self.alphas[ic] / 2.0) * np.cos(self.deltas[ic] / 2.0) + dofs[5] = np.cos(self.alphas[ic] / 2.0) * np.sin(self.deltas[ic] / 2.0) + dofs[6] = -np.sin(self.alphas[ic] / 2.0) * np.sin(self.deltas[ic] / 2.0) + # Now specify the center + dofs[7] = xyz[0] + dofs[8] = xyz[1] + dofs[9] = xyz[2] + curves[ic].set_dofs(dofs) + self.curves = curves + self.all_curves = apply_symmetries_to_curves(curves, self.nfp, self.stellsym) - def plot_curves(self): + def plot_curves(self, filename=''): from pyevtk.hl import pointsToVTK from . import curves_to_vtk from simsopt.field import InterpolatedField, Current, coils_via_symmetries - curves_to_vtk(self.curves, self.out_dir + "psc_curves", close=True, scalar_data=self.I) + # curves_to_vtk(self.curves, self.out_dir + "psc_curves", close=True, scalar_data=self.I) + # contig = np.ascontiguousarray + # if isinstance(self.B_TF, InterpolatedField): + # self.B_TF.set_points(self.grid_xyz) + # B = self.B_TF.B() + # pointsToVTK(self.out_dir + 'curve_centers', + # contig(self.grid_xyz[:, 0]), + # contig(self.grid_xyz[:, 1]), + # contig(self.grid_xyz[:, 2]), + # data={"n": (contig(self.coil_normals[:, 0]), + # contig(self.coil_normals[:, 1]), + # contig(self.coil_normals[:, 2])), + # "psi": contig(self.psi), + # "I": contig(self.I), + # "B_TF": (contig(B[:, 0]), + # contig(B[:, 1]), + # contig(B[:, 2])), + # }, + # ) + # else: + # pointsToVTK(self.out_dir + 'curve_centers', + # contig(self.grid_xyz[:, 0]), + # contig(self.grid_xyz[:, 1]), + # contig(self.grid_xyz[:, 2]), + # data={"n": (contig(self.coil_normals[:, 0]), + # contig(self.coil_normals[:, 1]), + # contig(self.coil_normals[:, 2])), + # }, + # ) + # 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) + 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 + q = q + 1 + + 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])) + 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) contig = np.ascontiguousarray if isinstance(self.B_TF, InterpolatedField): - self.B_TF.set_points(self.grid_xyz) + self.B_TF.set_points(self.grid_xyz_all) B = self.B_TF.B() - pointsToVTK(self.out_dir + 'curve_centers', - contig(self.grid_xyz[:, 0]), - contig(self.grid_xyz[:, 1]), - contig(self.grid_xyz[:, 2]), - data={"n": (contig(self.coil_normals[:, 0]), - contig(self.coil_normals[:, 1]), - contig(self.coil_normals[:, 2])), - "psi": contig(self.psi), - "I": contig(self.I), + pointsToVTK(self.out_dir + filename + 'all_curve_centers', + contig(self.grid_xyz_all[:, 0]), + contig(self.grid_xyz_all[:, 1]), + contig(self.grid_xyz_all[:, 2]), + data={"n": (contig(self.coil_normals_all[:, 0]), + contig(self.coil_normals_all[:, 1]), + contig(self.coil_normals_all[:, 2])), + "psi": contig(self.psi_all), + "I": contig(self.I_all), "B_TF": (contig(B[:, 0]), contig(B[:, 1]), contig(B[:, 2])), }, ) else: - pointsToVTK(self.out_dir + 'curve_centers', - contig(self.grid_xyz[:, 0]), - contig(self.grid_xyz[:, 1]), - contig(self.grid_xyz[:, 2]), - data={"n": (contig(self.coil_normals[:, 0]), - contig(self.coil_normals[:, 1]), - contig(self.coil_normals[:, 2])), + pointsToVTK(self.out_dir + filename + 'all_curve_centers', + contig(self.grid_xyz_all[:, 0]), + contig(self.grid_xyz_all[:, 1]), + contig(self.grid_xyz_all[:, 2]), + data={"n": (contig(self.coil_normals_all[:, 0]), + contig(self.coil_normals_all[:, 1]), + contig(self.coil_normals_all[:, 2])), }, ) - 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) - for i in range(self.symmetry): - self.psi_all[i * self.num_psc: (i + 1) * self.num_psc] = self.psi - self.I_all[i * self.num_psc: (i + 1) * self.num_psc] = self.I - - 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])) - 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 + "all_psc_curves", close=True, scalar_data=self.I_all) - contig = np.ascontiguousarray - if isinstance(self.B_TF, InterpolatedField): - self.B_TF.set_points(self.grid_xyz_all) - B = self.B_TF.B() - pointsToVTK(self.out_dir + 'all_curve_centers', - contig(self.grid_xyz_all[:, 0]), - contig(self.grid_xyz_all[:, 1]), - contig(self.grid_xyz_all[:, 2]), - data={"n": (contig(self.coil_normals_all[:, 0]), - contig(self.coil_normals_all[:, 1]), - contig(self.coil_normals_all[:, 2])), - "psi": contig(self.psi_all), - "I": contig(self.I_all), - "B_TF": (contig(B[:, 0]), - contig(B[:, 1]), - contig(B[:, 2])), - }, - ) - else: - pointsToVTK(self.out_dir + 'all_curve_centers', - contig(self.grid_xyz_all[:, 0]), - contig(self.grid_xyz_all[:, 1]), - contig(self.grid_xyz_all[:, 2]), - data={"n": (contig(self.coil_normals_all[:, 0]), - contig(self.coil_normals_all[:, 1]), - contig(self.coil_normals_all[:, 2])), - }, - ) def b_vector(self): Bn_target = self.Bn.reshape(-1) @@ -730,7 +664,7 @@ def least_squares(self, kappas): alphas = kappas[:len(kappas) // 2] deltas = kappas[len(kappas) // 2:] self.setup_orientations(alphas, deltas) - BdotN2 = 0.5 * np.sum((self.Bn_PSC.reshape(-1) + self.b_opt) ** 2) + BdotN2 = 0.5 * np.sum((self.Bn_PSC.reshape(-1) * self.grid_normalization + self.b_opt) ** 2) outstr = f"||Bn||^2 = {BdotN2:.2e} " # for i in range(len(kappas)): # outstr += f"kappas[{i:d}] = {kappas[i]:.2e} " @@ -758,7 +692,7 @@ def setup_orientations(self, alphas, deltas): # np.cos(deltas[i]) * np.sin(alphas[i]), # np.cos(deltas[i]) * np.cos(alphas[i])]]) # self.rotation_matrix = rotation_matrix - t1 = time.time() + # t1 = time.time() flux_grid = sopp.flux_xyz( contig(self.grid_xyz), contig(alphas), @@ -768,24 +702,24 @@ def setup_orientations(self, alphas, deltas): contig(self.coil_normals) ) self.B_TF.set_points(contig(np.array(flux_grid).reshape(-1, 3))) - t2 = time.time() - print('Flux set points time = ', t2 - t1) + # t2 = time.time() + # print('Flux set points time = ', t2 - t1) N = len(self.rho) Nflux = len(self.flux_phi) # t1 = time.time() # B = # t2 = time.time() # print('Flux call B time = ', t2 - t1) - t1 = time.time() + # t1 = time.time() self.psi = sopp.flux_integration( contig(self.B_TF.B().reshape(len(alphas), N, Nflux, 3)), contig(self.rho), contig(self.coil_normals) ) self.psi *= self.dphi * self.drho - t2 = time.time() - print('Flux integration time = ', t2 - t1) - t1 = time.time() + # t2 = time.time() + # print('Flux integration time = ', t2 - t1) + # t1 = time.time() ####### Inductance C++ calculation below L = sopp.L_matrix( contig(self.grid_xyz), @@ -794,21 +728,39 @@ def setup_orientations(self, alphas, deltas): contig(self.phi), self.R, ) * self.dphi ** 2 / (4 * np.pi) - t2 = time.time() - print('Inductances c++ total time = ', t2 - t1) + # t2 = time.time() + # print('Inductances c++ total time = ', t2 - t1) L = (L + L.T) np.fill_diagonal(L, np.log(8.0 * self.R / self.a) - 2.0) self.L = L * self.mu0 * self.R * self.Nt ** 2 - t1 = time.time() + # t1 = time.time() self.setup_currents_and_fields() - t2 = time.time() - print('Setup fields time = ', t2 - t1) + # t2 = time.time() + # print('Setup fields time = ', t2 - t1) + + def setup_full_grid(self): + nn = self.num_psc + self.grid_xyz_all = np.zeros((nn * self.symmetry, 3)) + self.coil_normals_all = np.zeros((nn * self.symmetry, 3)) + q = 0 + for fp in range(self.nfp): + for stell in self.stell_list: + phi0 = (2 * np.pi / self.nfp) * fp + # get new locations by flipping the y and z components, then rotating by phi0 + self.grid_xyz_all[nn * q: nn * (q + 1), 0] = self.grid_xyz[:, 0] * np.cos(phi0) - self.grid_xyz[:, 1] * np.sin(phi0) * stell + self.grid_xyz_all[nn * q: nn * (q + 1), 1] = self.grid_xyz[:, 0] * np.sin(phi0) + self.grid_xyz[:, 1] * np.cos(phi0) * stell + self.grid_xyz_all[nn * q: nn * (q + 1), 2] = self.grid_xyz[:, 2] * stell + # get new normal vectors by flipping the x component, then rotating by phi0 + 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] + q = q + 1 -def hamilton_product(q1, q2): - q3_0 = q1[0] * q2[0] - q1[1] * q2[1] - q1[2] * q2[2] - q1[3] * q2[3] - q3_1 = q1[0] * q2[1] + q1[1] * q2[0] + q1[1] * q2[2] - q1[3] * q2[2] - q3_2 = q1[0] * q2[2] - q1[1] * q2[3] + q1[1] * q2[0] + q1[3] * q2[1] - q3_3 = q1[0] * q2[3] + q1[1] * q2[2] - q1[2] * q2[1] + q1[3] * q2[0] - return np.array([q3_0, q3_1, q3_2, q3_3]) +# def hamilton_product(q1, q2): +# q3_0 = q1[0] * q2[0] - q1[1] * q2[1] - q1[2] * q2[2] - q1[3] * q2[3] +# q3_1 = q1[0] * q2[1] + q1[1] * q2[0] + q1[1] * q2[2] - q1[3] * q2[2] +# q3_2 = q1[0] * q2[2] - q1[1] * q2[3] + q1[1] * q2[0] + q1[3] * q2[1] +# q3_3 = q1[0] * q2[3] + q1[1] * q2[2] - q1[2] * q2[1] + q1[3] * q2[0] +# return np.array([q3_0, q3_1, q3_2, q3_3]) \ No newline at end of file diff --git a/src/simsoptpp/psc.cpp b/src/simsoptpp/psc.cpp index e81374488..dd544d87a 100644 --- a/src/simsoptpp/psc.cpp +++ b/src/simsoptpp/psc.cpp @@ -153,7 +153,7 @@ Array flux_integration(Array& B, Array& rho, Array& normal) return Psi; } -Array Bn_PSC(Array& points, Array& plasma_points, Array& alphas, Array& deltas, Array& plasma_normal, double R) +Array A_matrix(Array& points, Array& plasma_points, Array& alphas, Array& deltas, Array& plasma_normal, double R, double phi0, double stell) { // warning: row_major checks below do NOT throw an error correctly on a compute node on Cori if(points.layout() != xt::layout_type::row_major) @@ -231,13 +231,17 @@ Array Bn_PSC(Array& points, Array& plasma_points, Array& alphas, Array& deltas, auto Bx_rot = Bx * nxx + By * nxy + Bz * nxz; auto By_rot = Bx * nyx + By * nyy + Bz * nyz; auto Bz_rot = Bx * nzx + By * nzy + Bz * nzz; - Bn_PSC(i, j) = Bx_rot * nx + By_rot * ny + Bz_rot * nz; + // now apply R_{fp}R_s flipping the x component, then rotating by phi0 + auto Bx_rot_flip = Bx_rot * cos(phi0) * stell - By_rot * sin(phi0); + auto By_rot_flip = Bx_rot * sin(phi0) * stell + By_rot * cos(phi0); + auto Bz_rot_flip = Bz_rot; + Bn_PSC(i, j) = Bx_rot_flip * nx + By_rot_flip * ny + Bz_rot_flip * nz; } } return Bn_PSC * fac; } -Array B_PSC(Array& points, Array& plasma_points, Array& alphas, Array& deltas, Array& psc_currents, double R) +Array B_PSC(Array& points, Array& plasma_points, Array& alphas, Array& deltas, Array& psc_currents, double R, double phi0, double stell) { // warning: row_major checks below do NOT throw an error correctly on a compute node on Cori if(points.layout() != xt::layout_type::row_major) @@ -307,9 +311,13 @@ Array B_PSC(Array& points, Array& plasma_points, Array& alphas, Array& deltas, A auto By = current * y * z * Eplus / (rho2 * beta_gamma2); auto Bz = current * Eminus / beta_gamma2; // Need to rotate the vector - B_PSC(i, 0) += Bx * nxx + By * nxy + Bz * nxz; - B_PSC(i, 1) += Bx * nyx + By * nyy + Bz * nyz; - B_PSC(i, 2) += Bx * nzx + By * nzy + Bz * nzz; + auto Bx_rot = Bx * nxx + By * nxy + Bz * nxz; + auto By_rot = Bx * nyx + By * nyy + Bz * nyz; + auto Bz_rot = Bx * nzx + By * nzy + Bz * nzz; + // now apply R_{fp}R_s flipping the x component, then rotating by phi0 + B_PSC(i, 0) += Bx_rot * cos(phi0) * stell - By_rot * sin(phi0); + B_PSC(i, 1) += Bx_rot * sin(phi0) * stell + By_rot * cos(phi0); + B_PSC(i, 2) += Bz_rot; } } return B_PSC * fac; diff --git a/src/simsoptpp/psc.h b/src/simsoptpp/psc.h index 42a197b38..e47d17657 100644 --- a/src/simsoptpp/psc.h +++ b/src/simsoptpp/psc.h @@ -20,9 +20,9 @@ Array flux_xyz(Array& points, Array& alphas, Array& deltas, Array& rho, Array& p Array flux_integration(Array& B, Array& rho, Array& normal); -Array Bn_PSC(Array& points, Array& plasma_points, Array& alphas, Array& deltas, Array& plasma_normal, double R); +Array A_matrix(Array& points, Array& plasma_points, Array& alphas, Array& deltas, Array& plasma_normal, double R, double phi0, double stell); -Array B_PSC(Array& points, Array& plasma_points, Array& alphas, Array& deltas, Array& psc_currents, double R); +Array B_PSC(Array& points, Array& plasma_points, Array& alphas, Array& deltas, Array& psc_currents, double R, double phi0, double stell); // // Array A_matrix(Array& plasma_surface_normals, Array& B, Array& alphas, Array& deltas); \ No newline at end of file diff --git a/src/simsoptpp/python.cpp b/src/simsoptpp/python.cpp index 32c35ef19..1c680944d 100644 --- a/src/simsoptpp/python.cpp +++ b/src/simsoptpp/python.cpp @@ -63,7 +63,7 @@ PYBIND11_MODULE(simsoptpp, m) { m.def("L_matrix" , &L_matrix); m.def("flux_xyz" , &flux_xyz); m.def("flux_integration" , &flux_integration); - m.def("Bn_PSC" , &Bn_PSC); + m.def("A_matrix" , &A_matrix); m.def("B_PSC" , &B_PSC); // Functions below are implemented for permanent magnet optimization diff --git a/tests/geo/test_psc_grid.py b/tests/geo/test_psc_grid.py index ab9207a5a..c70dacad2 100644 --- a/tests/geo/test_psc_grid.py +++ b/tests/geo/test_psc_grid.py @@ -4,7 +4,7 @@ import numpy as np from monty.tempfile import ScratchDir -from simsopt.geo import PSCgrid +from simsopt.geo import PSCgrid, SurfaceRZFourier from simsopt.field import BiotSavart, coils_via_symmetries, Current, CircularCoil import simsoptpp as sopp @@ -89,6 +89,18 @@ def test_L(self): # Only true because R << 1 # assert(np.isclose(psc_array.psi[0], np.pi * psc_array.R ** 2 * Bz_center)) + input_name = 'input.LandremanPaul2021_QA_lowres' + TEST_DIR = (Path(__file__).parent / ".." / ".." / "tests" / "test_files").resolve() + surface_filename = TEST_DIR / input_name + surf1 = SurfaceRZFourier.from_vmec_input( + surface_filename, range='full torus', nphi=16, ntheta=16 + ) + surf1.nfp = 1 + surf1.stellsym = False + surf2 = SurfaceRZFourier.from_vmec_input( + surface_filename, range='half period', nphi=16, ntheta=16 + ) + # Test that inductance and flux calculations for wide set of # scenarios a = 1e-4 @@ -106,60 +118,107 @@ def test_L(self): (np.random.rand(3) - 0.5) * 40])]: for alphas in [np.zeros(2), np.random.rand(2) * 2 * np.pi]: for deltas in [np.zeros(2), np.random.rand(2) * 2 * np.pi]: - psc_array = PSCgrid.geo_setup_manual( - points, R=R, a=a, alphas=alphas, deltas=deltas, - ) - coils = coils_via_symmetries([psc_array.curves[1]], [Current(I)], nfp=1, stellsym=False) - bs = BiotSavart(coils) - kwargs = {"B_TF": bs, "ppp": 2000} - psc_array = PSCgrid.geo_setup_manual( - points, R=R, a=a, alphas=alphas, deltas=deltas, **kwargs - ) - L = psc_array.L - - # This is not a robust check but it only converges when N >> 1 - # points are used to do the integrations. Can easily check that - # can increase rtol as you increase N - # print(psc_array.psi[0] / I * 1e10, L[1, 0] * 1e10) - assert(np.isclose(psc_array.psi[0] / I * 1e10, L[1, 0] * 1e10, rtol=1e-1)) - - psc_array.setup_psc_biotsavart() - contig = np.ascontiguousarray - B_PSC = sopp.B_PSC( - contig(psc_array.grid_xyz), - contig(psc_array.plasma_boundary.gamma().reshape(-1, 3)), - contig(psc_array.alphas), - contig(psc_array.deltas), - contig(psc_array.I), - psc_array.R - ) - Bn_PSC = sopp.Bn_PSC( - contig(psc_array.grid_xyz), - contig(psc_array.plasma_boundary.gamma().reshape(-1, 3)), - contig(psc_array.alphas), - contig(psc_array.deltas), - contig(psc_array.plasma_boundary.unitnormal().reshape(-1, 3)), - psc_array.R - ) @ psc_array.I - B_circular_coils = np.zeros(B_PSC.shape) - Bn_circular_coils = np.zeros(psc_array.Bn_PSC.shape) - for i in range(len(psc_array.alphas)): - PSC = CircularCoil( - psc_array.R, - psc_array.grid_xyz[i, :], - psc_array.I[i], - psc_array.coil_normals[i, :] + for surf in [surf1, surf2]: + psc_array = PSCgrid.geo_setup_manual( + points, R=R, a=a, alphas=alphas, deltas=deltas, + ) + coils = coils_via_symmetries([psc_array.curves[1]], [Current(I)], nfp=1, stellsym=False) + bs = BiotSavart(coils) + kwargs = {"B_TF": bs, "ppp": 2000, "plasma_boundary": surf} + psc_array = PSCgrid.geo_setup_manual( + points, R=R, a=a, alphas=alphas, deltas=deltas, **kwargs + ) + L = psc_array.L + + # This is not a robust check but it only converges when N >> 1 + # points are used to do the integrations. Can easily check that + # can increase rtol as you increase N + # print(psc_array.psi[0] / I * 1e10, L[1, 0] * 1e10) + assert(np.isclose(psc_array.psi[0] / I * 1e10, L[1, 0] * 1e10, rtol=1e-1)) + + psc_array.setup_psc_biotsavart() + contig = np.ascontiguousarray + # Calculate B fields like psc_array function does + B_PSC = np.zeros((psc_array.nphi * psc_array.ntheta, 3)) + Bn_PSC = np.zeros(psc_array.nphi * psc_array.ntheta) + print(psc_array.nfp, psc_array.stell_list) + for fp in range(psc_array.nfp): + for stell in psc_array.stell_list: + phi0 = (2 * np.pi / psc_array.nfp) * fp + # get new locations by flipping the y and z components, then rotating by phi0 + ox = psc_array.grid_xyz[:, 0] * np.cos(phi0) - psc_array.grid_xyz[:, 1] * np.sin(phi0) * stell + oy = psc_array.grid_xyz[:, 0] * np.sin(phi0) + psc_array.grid_xyz[:, 1] * np.cos(phi0) * stell + oz = psc_array.grid_xyz[:, 2] * stell + xyz = np.array([ox, oy, oz]).T + Bn_PSC += sopp.A_matrix( + contig(xyz), + contig(psc_array.plasma_boundary.gamma().reshape(-1, 3)), + contig(psc_array.alphas), + contig(psc_array.deltas), + contig(psc_array.plasma_boundary.unitnormal().reshape(-1, 3)), + psc_array.R, + float(phi0), + float(stell), + ) @ psc_array.I * stell + B_PSC += sopp.B_PSC( + contig(xyz), + contig(psc_array.plasma_boundary.gamma().reshape(-1, 3)), + contig(psc_array.alphas), + contig(psc_array.deltas), + contig(psc_array.I * stell), + psc_array.R, + float(phi0), + float(stell) + ) + # Calculate Bfields from CircularCoil class + B_circular_coils = np.zeros(B_PSC.shape) + Bn_circular_coils = np.zeros(psc_array.Bn_PSC.shape) + for i in range(len(psc_array.alphas)): + PSC = CircularCoil( + psc_array.R, + psc_array.grid_xyz_all[i, :], + psc_array.I_all[i], + psc_array.coil_normals_all[i, :] + ) + PSC.set_points(psc_array.plasma_boundary.gamma().reshape(-1, 3)) + B_circular_coils += PSC.B().reshape(-1, 3) + Bn_circular_coils += np.sum(PSC.B().reshape( + -1, 3) * psc_array.plasma_boundary.unitnormal().reshape(-1, 3), axis=-1) + # Calculate Bfields from direct BiotSavart + currents = [] + for i in range(len(psc_array.I)): + currents.append(Current(psc_array.I[i])) + coils = coils_via_symmetries( + psc_array.curves, currents, nfp=1, stellsym=False + ) + B_direct = BiotSavart(coils) + B_direct.set_points(psc_array.plasma_points) + Bn_direct = np.sum(B_direct.B().reshape( + -1, 3) * psc_array.plasma_boundary.unitnormal().reshape(-1, 3), axis=-1) + # Calculate Bfields from direct BiotSavart, using all the coils manually defined + currents = [] + for i in range(psc_array.num_psc * psc_array.symmetry): + currents.append(Current(psc_array.I_all[i])) + all_coils = coils_via_symmetries( + psc_array.all_curves, currents, nfp=1, stellsym=False ) - PSC.set_points(psc_array.plasma_boundary.gamma().reshape(-1, 3)) - B_circular_coils += PSC.B().reshape(-1, 3) - Bn_circular_coils += np.sum(PSC.B().reshape( + B_direct_all = BiotSavart(all_coils) + B_direct_all.set_points(psc_array.plasma_points) + Bn_direct_all = np.sum(B_direct_all.B().reshape( -1, 3) * psc_array.plasma_boundary.unitnormal().reshape(-1, 3), axis=-1) - - # Robust test of all the B and Bn calculations from circular coils - assert(np.allclose(psc_array.B_PSC.B(), B_PSC)) - assert(np.allclose(psc_array.B_PSC.B(), B_circular_coils)) - assert(np.allclose(psc_array.Bn_PSC, Bn_PSC)) - assert(np.allclose(psc_array.Bn_PSC, Bn_circular_coils)) + + # Robust test of all the B and Bn calculations from circular coils + # print('here = ', psc_array.Bn_PSC * 1e10, Bn_PSC * 1e10, Bn_circular_coils * 1e10) + # print(psc_array.B_PSC.B()[-1] * 1e10, B_PSC[-1] * 1e10, B_circular_coils[-1] * 1e10, B_direct.B()[-1] * 1e10) + # assert(np.allclose(psc_array.B_PSC.B() * 1e10, B_PSC * 1e10)) + # assert(np.allclose(psc_array.B_PSC.B() * 1e10, B_circular_coils * 1e10)) + # assert(np.allclose(psc_array.B_PSC.B() * 1e10, B_direct.B() * 1e10)) + # assert(np.allclose(psc_array.B_PSC.B() * 1e10, B_direct_all.B() * 1e10)) + print(psc_array.Bn_PSC[-1] * 1e10, Bn_PSC[-1] * 1e10, Bn_circular_coils[-1] * 1e10, Bn_direct[-1] * 1e10) + assert(np.allclose(psc_array.Bn_PSC * 1e10, Bn_PSC * 1e10)) + assert(np.allclose(psc_array.Bn_PSC * 1e10, Bn_circular_coils * 1e10)) + assert(np.allclose(psc_array.Bn_PSC * 1e10, Bn_direct * 1e10)) + assert(np.allclose(psc_array.Bn_PSC * 1e10, Bn_direct_all * 1e10)) if __name__ == "__main__": unittest.main()