diff --git a/src/simsopt/field/currentpotential.py b/src/simsopt/field/currentpotential.py index dc5f7319f..04ecdbcb7 100644 --- a/src/simsopt/field/currentpotential.py +++ b/src/simsopt/field/currentpotential.py @@ -335,32 +335,59 @@ def from_netcdf(cls, filename: str, coil_ntheta_res=1.0, coil_nzeta_res=1.0): ntor_coil = int(np.max(xn_coil)/nfp) - quadpoints_phi, quadpoints_theta = Surface.get_quadpoints( - ntheta=ntheta_coil, nphi=nzeta_coil, nfp=nfp, range='full torus') - s_coil = SurfaceRZFourier(nfp=nfp, mpol=mpol_coil, ntor=ntor_coil, stellsym=stellsym_surf, - quadpoints_phi=quadpoints_phi,quadpoints_theta=quadpoints_theta) + # quadpoints_phi, quadpoints_theta = Surface.get_quadpoints( + # ntheta=ntheta_coil, nphi=nzeta_coil, nfp=nfp, range='full torus') + + s_coil = SurfaceRZFourier(nfp=nfp, mpol=mpol_coil, ntor=ntor_coil, stellsym=stellsym_surf) + s_coil = s_coil.from_nphi_ntheta(nfp=nfp, ntheta=ntheta_coil, nphi=nzeta_coil * nfp, + mpol=mpol_coil, ntor=ntor_coil, stellsym=stellsym_surf, range='full torus') + + # s_coil = SurfaceRZFourier(nfp=nfp, mpol=mpol_coil, ntor=ntor_coil, stellsym=stellsym_surf, + # quadpoints_phi=quadpoints_phi,quadpoints_theta=quadpoints_theta) s_coil.set_dofs(0*s_coil.get_dofs()) - s_coil.rc = 0*s_coil.rc - s_coil.zs = 0*s_coil.zs + # s_coil.rc = 0*s_coil.rc + # s_coil.zs = 0*s_coil.zs for im in range(len(xm_coil)): + s_coil.set_rc(xm_coil[im], int(xn_coil[im]/nfp), rmnc_coil[im]) + s_coil.set_zs(xm_coil[im], int(xn_coil[im]/nfp), zmns_coil[im]) m = int(xm_coil[im]) n = int(xn_coil[im] / nfp) - if (m == 0 and n<0): - s_coil.set_rc(m, -n, rmnc_coil[im]) - s_coil.set_zs(m, -n, -zmns_coil[im]) - else: - s_coil.set_rc(m, n, rmnc_coil[im]) - s_coil.set_zs(m, n, zmns_coil[im]) + # if (m == 0 and n<0): + # s_coil.set_rc(m, -n, rmnc_coil[im]) + # s_coil.set_zs(m, -n, -zmns_coil[im]) + # else: + # s_coil.set_rc(m, n, rmnc_coil[im]) + # s_coil.set_zs(m, n, zmns_coil[im]) if not stellsym_surf: - if (m == 0 and n<0): - s_coil.set_rs(m, -n, -rmns_coil[im]) - s_coil.set_zc(m, -n, zmnc_coil[im]) - else: - s_coil.set_rs(m, n, rmns_coil[im]) - s_coil.set_zc(m, n, zmnc_coil[im]) + s_coil.set_rs(xm_coil[im], int(xn_coil[im]/nfp), rmns_coil[im]) + s_coil.set_zc(xm_coil[im], int(xn_coil[im]/nfp), zmnc_coil[im]) + # if (m == 0 and n<0): + # s_coil.set_rs(m, -n, -rmns_coil[im]) + # s_coil.set_zc(m, -n, zmnc_coil[im]) + # else: + # s_coil.set_rs(m, n, rmns_coil[im]) + # s_coil.set_zc(m, n, zmnc_coil[im]) s_coil.local_full_x = s_coil.get_dofs() + # for im in range(len(xm_coil)): + # m = int(xm_coil[im]) + # n = int(xn_coil[im] / nfp) + # if (m == 0 and n<0): + # s_coil.set_rc(m, -n, rmnc_coil[im]) + # s_coil.set_zs(m, -n, -zmns_coil[im]) + # else: + # s_coil.set_rc(m, n, rmnc_coil[im]) + # s_coil.set_zs(m, n, zmns_coil[im]) + # if not stellsym_surf: + # if (m == 0 and n<0): + # s_coil.set_rs(m, -n, -rmns_coil[im]) + # s_coil.set_zc(m, -n, zmnc_coil[im]) + # else: + # s_coil.set_rs(m, n, rmns_coil[im]) + # s_coil.set_zc(m, n, zmnc_coil[im]) + + # s_coil.local_full_x = s_coil.get_dofs() cp = cls(s_coil, mpol=mpol_potential, ntor=ntor_potential, net_poloidal_current_amperes=net_poloidal_current_amperes, diff --git a/src/simsopt/field/currentpotentialsolve.py b/src/simsopt/field/currentpotentialsolve.py index 302503163..2b31f55ef 100644 --- a/src/simsopt/field/currentpotentialsolve.py +++ b/src/simsopt/field/currentpotentialsolve.py @@ -443,7 +443,11 @@ def solve_tikhonov(self, lam=0, record_history=True): b_e = self.b_e Ak_times_phi = self.fj @ phi_mn_opt f_B = 0.5 * np.linalg.norm(A_times_phi - b_e) ** 2 * nfp - f_K = 0.5 * np.linalg.norm(Ak_times_phi - self.d) ** 2 + # extra normN factor needed here because fj and d don't have it + # K^2 has 1/normn^2 factor, the sum over the winding surface has factor of normn, + # for total factor of 1/normn + normN = np.linalg.norm(self.winding_surface.normal().reshape(-1, 3), axis=-1) + f_K = 0.5 * np.linalg.norm((Ak_times_phi - self.d) / np.sqrt(normN[:, None])) ** 2 if record_history: self.ilambdas_l2.append(lam) @@ -484,12 +488,16 @@ def solve_lasso(self, lam=0, max_iter=1000, acceleration=True): # Set up some matrices _, _ = self.B_matrix_and_rhs() normN = np.linalg.norm(self.plasma_surface.normal().reshape(-1, 3), axis=-1) + ws_normN = np.linalg.norm(self.winding_surface.normal().reshape(-1, 3), axis=-1) A_matrix = self.gj for i in range(self.gj.shape[0]): A_matrix[i, :] *= (1.0 / np.sqrt(normN[i])) b_e = self.b_e Ak_matrix = self.fj.reshape(self.fj.shape[0] * 3, self.fj.shape[-1]) d = np.ravel(self.d) + # for i in range(3): + # Ak_matrix[i * len(ws_normN): (i+1) * len(ws_normN), :] *= 1.0 / np.sqrt(ws_normN)[:, None] + # d[i * len(ws_normN): (i+1) * len(ws_normN)] *= 1.0 / np.sqrt(ws_normN) nfp = self.plasma_surface.nfp # Ak is non-square so pinv required. Careful with rcond parameter @@ -505,9 +513,16 @@ def solve_lasso(self, lam=0, max_iter=1000, acceleration=True): # if alpha << 1, want to use initial guess from the Tikhonov solve, # which is exact since it comes from a matrix inverse. phi0, _, _, = self.solve_tikhonov(lam=lam, record_history=False) - z0 = Ak_matrix @ phi0 - d + + # L1 norm here should already include the contributions from the winding surface discretization + # and factor of 1 / ws_normN from the K, cancelling the factor of ws_normN from the surface + z0 = (Ak_matrix @ phi0 - d) #* np.sqrt(ws_normN) z_opt, z_history = self._FISTA(A=A_new, b=b_new, alpha=l1_reg, max_iter=max_iter, acceleration=acceleration, xi0=z0) + # Need to put back in the 1 / ws_normN dependence in K + for i in range(3): + z_opt[i * len(ws_normN): (i + 1) * len(ws_normN)] *= ws_normN + # Compute the history of values from the optimizer phi_history = [] fB_history = [] @@ -520,10 +535,9 @@ def solve_lasso(self, lam=0, max_iter=1000, acceleration=True): # Remember, Lasso solved for z = A_k * phi_mn - b_k so need to convert back phi_mn_opt = Ak_inv @ (z_opt + d) - self.current_potential.set_dofs(phi_mn_opt) f_B = 0.5 * np.linalg.norm(A_matrix @ phi_mn_opt - b_e) ** 2 * nfp - f_K = np.linalg.norm(Ak_matrix @ phi_mn_opt - d, ord=1) + f_K = np.linalg.norm(Ak_matrix @ phi_mn_opt - d, ord=1) self.ilambdas_l1.append(lam) self.dofs_l1.append(phi_mn_opt) # REGCOIL only uses 1 / 2 nfp of the winding surface diff --git a/src/simsoptpp/winding_surface.cpp b/src/simsoptpp/winding_surface.cpp index f54e28d8f..5c7bcf1ad 100644 --- a/src/simsoptpp/winding_surface.cpp +++ b/src/simsoptpp/winding_surface.cpp @@ -470,21 +470,21 @@ std::tuple winding_surface_field_K2_matrices(Array& dr_dzeta_coil, double ny = normal_coil(j, 1); double nz = normal_coil(j, 2); double normN = sqrt(nx * nx + ny * ny + nz * nz); - d(j, 0) = (G * dr_dtheta_coil(j, 0) - I * dr_dzeta_coil(j, 0)) / sqrt(normN) / (2 * M_PI); - d(j, 1) = (G * dr_dtheta_coil(j, 1) - I * dr_dzeta_coil(j, 1)) / sqrt(normN) / (2 * M_PI); - d(j, 2) = (G * dr_dtheta_coil(j, 2) - I * dr_dzeta_coil(j, 2)) / sqrt(normN) / (2 * M_PI); + d(j, 0) = (G * dr_dtheta_coil(j, 0) - I * dr_dzeta_coil(j, 0)) / (2 * M_PI); + d(j, 1) = (G * dr_dtheta_coil(j, 1) - I * dr_dzeta_coil(j, 1)) / (2 * M_PI); + d(j, 2) = (G * dr_dtheta_coil(j, 2) - I * dr_dzeta_coil(j, 2)) / (2 * M_PI); for (int k = 0; k < m.size(); k++) { double angle = 2 * M_PI * m(k) * theta_coil(j) - 2 * M_PI * n(k) * zeta_coil(j) * nfp; double cphi = std::cos(angle); double sphi = std::sin(angle); - fj(j, 0, k) = cphi * (m(k) * dr_dzeta_coil(j, 0) + nfp * n(k) * dr_dtheta_coil(j, 0)) / sqrt(normN); - fj(j, 1, k) = cphi * (m(k) * dr_dzeta_coil(j, 1) + nfp * n(k) * dr_dtheta_coil(j, 1)) / sqrt(normN); - fj(j, 2, k) = cphi * (m(k) * dr_dzeta_coil(j, 2) + nfp * n(k) * dr_dtheta_coil(j, 2)) / sqrt(normN); + fj(j, 0, k) = cphi * (m(k) * dr_dzeta_coil(j, 0) + nfp * n(k) * dr_dtheta_coil(j, 0)); + fj(j, 1, k) = cphi * (m(k) * dr_dzeta_coil(j, 1) + nfp * n(k) * dr_dtheta_coil(j, 1)); + fj(j, 2, k) = cphi * (m(k) * dr_dzeta_coil(j, 2) + nfp * n(k) * dr_dtheta_coil(j, 2)); if (! stellsym) { - fj(j, 0, k+m.size()) = -sphi * (m(k) * dr_dzeta_coil(j, 0) + nfp * n(k) * dr_dtheta_coil(j, 0)) / sqrt(normN); - fj(j, 1, k+m.size()) = -sphi * (m(k) * dr_dzeta_coil(j, 1) + nfp * n(k) * dr_dtheta_coil(j, 1)) / sqrt(normN); - fj(j, 2, k+m.size()) = -sphi * (m(k) * dr_dzeta_coil(j, 2) + nfp * n(k) * dr_dtheta_coil(j, 2)) / sqrt(normN); + fj(j, 0, k+m.size()) = -sphi * (m(k) * dr_dzeta_coil(j, 0) + nfp * n(k) * dr_dtheta_coil(j, 0)); + fj(j, 1, k+m.size()) = -sphi * (m(k) * dr_dzeta_coil(j, 1) + nfp * n(k) * dr_dtheta_coil(j, 1)); + fj(j, 2, k+m.size()) = -sphi * (m(k) * dr_dzeta_coil(j, 2) + nfp * n(k) * dr_dtheta_coil(j, 2)); } } } diff --git a/tests/field/test_currentpotential.py b/tests/field/test_currentpotential.py index a956490ce..0722f8579 100644 --- a/tests/field/test_currentpotential.py +++ b/tests/field/test_currentpotential.py @@ -42,8 +42,8 @@ def test_compare_K_with_regcoil(self): K = cp.K() K2 = np.sum(K*K, axis=2) K2_average = np.mean(K2, axis=(0, 1)) - # print(K2[0:int(len(cp.quadpoints_phi)/cp.nfp), :]/K2_average, K2.shape, K2_average) - # print(K2_regcoil/K2_average, K2_regcoil.shape) + print(K2[0:int(len(cp.quadpoints_phi)/cp.nfp), :]/K2_average, K2.shape, K2_average) + print(K2_regcoil[0:int(len(cp.quadpoints_phi)/cp.nfp), :]/K2_average, K2_regcoil.shape) assert np.allclose(K2[0:int(len(cp.quadpoints_phi)/cp.nfp), :]/K2_average, K2_regcoil/K2_average) diff --git a/tests/field/test_regcoil.py b/tests/field/test_regcoil.py index e8edbbf9d..63251dc44 100644 --- a/tests/field/test_regcoil.py +++ b/tests/field/test_regcoil.py @@ -203,6 +203,7 @@ def test_regcoil_K_solve(self): against the REGCOIL variables. """ for filename in ['regcoil_out.w7x_infty.nc', 'regcoil_out.li383_infty.nc']: + print(filename) filename = TEST_DIR / filename cpst = CurrentPotentialSolve.from_netcdf(filename) # initialize a solver object for the cp CurrentPotential @@ -240,6 +241,7 @@ def test_regcoil_K_solve(self): current_potential_thetazeta = f.variables['single_valued_current_potential_thetazeta'][()][ilambda, :, :] f.close() Bnormal_single_valued = Bnormal_regcoil_total - Bnormal_from_plasma_current - Bnormal_from_net_coil_currents + print('ilambda index = ', ilambda, lambda_regcoil) assert np.allclose(b_rhs_regcoil, b_rhs_simsopt) assert np.allclose(k_rhs, k_rhs_regcoil) @@ -254,7 +256,8 @@ def test_regcoil_K_solve(self): optimized_phi_mn_lasso, f_B_lasso, f_K_lasso, _, _ = cpst.solve_lasso(lam=lambda_regcoil) optimized_phi_mn, f_B, f_K = cpst.solve_tikhonov(lam=lambda_regcoil) assert np.allclose(single_valued_current_potential_mn, optimized_phi_mn) - assert np.isclose(f_B_lasso, f_B) + print(f_B, f_B_lasso, f_B_regcoil) + assert np.isclose(f_B, f_B_regcoil) # assert np.isclose(f_K_lasso, f_K) assert np.allclose(optimized_phi_mn_lasso, optimized_phi_mn) @@ -330,6 +333,7 @@ def test_regcoil_K_solve(self): normal = s_coil.normal().reshape(-1, 3) normN = np.linalg.norm(normal, axis=-1) f_K_direct = 0.5 * np.sum(np.ravel(K2) * normN) / (normal.shape[0]) + # print(f_K_regcoil, f_K_direct, f_K) assert np.isclose(f_K_regcoil, f_K_direct) assert np.isclose(f_K_regcoil, f_K) @@ -462,6 +466,17 @@ def test_winding_surface_regcoil(self): norm_normal_coil_simsopt = np.linalg.norm(s_coil.normal(), axis=-1) assert np.allclose(norm_normal_coil*2*np.pi*2*np.pi, norm_normal_coil_simsopt[0:nzeta_coil, :]) + # Compare two different ways of computing K() + K = cp.K().reshape(-1, 3) + winding_surface = cp.winding_surface + normal_vec = winding_surface.normal().reshape(-1, 3) + dzeta_coil = (winding_surface.quadpoints_phi[1] - winding_surface.quadpoints_phi[0]) + dtheta_coil = (winding_surface.quadpoints_theta[1] - winding_surface.quadpoints_theta[0]) + normn = np.sqrt(np.sum(normal_vec**2, axis=-1)) # |N| + K_2 = -(cpst.fj @ cp.get_dofs() - cpst.d) / \ + (np.sqrt(dzeta_coil * dtheta_coil) * normn[:, None]) + assert np.allclose(K, K_2) + # Compare field from net coil currents cp_GI = CurrentPotentialFourier.from_netcdf(filename) Bfield = WindingSurfaceField(cp_GI) @@ -573,6 +588,115 @@ def test_winding_surface_regcoil(self): Bnormal_REGCOIL += B_GI_winding_surface assert np.allclose(Bnormal_REGCOIL, np.ravel(Bnormal_regcoil)) + def test_K_calculations(self): + from simsopt import load + winding_surface, plasma_surface = load(TEST_DIR / 'winding_surface_test.json') + cp = CurrentPotentialFourier( + winding_surface, mpol=4, ntor=4, + net_poloidal_current_amperes=11884578.094260072, + net_toroidal_current_amperes=0, + stellsym=True) + cp.set_dofs(np.array([ + 235217.63668779, -700001.94517193, 1967024.36417348, + -1454861.01406576, -1021274.81793687, 1657892.17597651, + -784146.17389912, 136356.84602536, -670034.60060171, + 194549.6432583 , 1006169.72177152, -1677003.74430119, + 1750470.54137804, 471941.14387043, -1183493.44552104, + 1046707.62318593, -334620.59690486, 658491.14959397, + -1169799.54944824, -724954.843765 , 1143998.37816758, + -2169655.54190455, -106677.43308896, 761983.72021537, + -986348.57384563, 532788.64040937, -600463.7957275 , + 1471477.22666607, 1009422.80860728, -2000273.40765417, + 2179458.3105468 , -55263.14222144, -315581.96056445, + 587702.35409154, -637943.82177418, 609495.69135857, + -1050960.33686344, -970819.1808181 , 1467168.09965404, + -198308.0580687 + ])) + cpst = CurrentPotentialSolve(cp, plasma_surface, np.zeros(1024)) + assert np.allclose(cpst.current_potential.get_dofs(), cp.get_dofs()) + # Pre-compute some important matrices + cpst.B_matrix_and_rhs() + + # Copied over from the packaged grid K operator. + winding_surface = cp.winding_surface + normal_vec = winding_surface.normal() + normn = np.sqrt(np.sum(normal_vec**2, axis=-1)) # |N| + + test_K_1 = ( + cp.winding_surface.gammadash2() + *(cp.Phidash1()+cp.net_poloidal_current_amperes)[:, :, None] + - cp.winding_surface.gammadash1() + *(cp.Phidash2()+cp.net_toroidal_current_amperes)[:, :, None])/normn[:, :, None] + + test_K_3 = cp.K() + + normn = normn.reshape(-1) + dzeta_coil = (winding_surface.quadpoints_phi[1] - winding_surface.quadpoints_phi[0]) + dtheta_coil = (winding_surface.quadpoints_theta[1] - winding_surface.quadpoints_theta[0]) + + # Notice Equation A.13 for the current in Matt L's regcoil paper has factor of 1/nnorm in it + # But cpst.fj and cpst.d have factor of only 1/sqrt(normn) + test_K_2 = -(cpst.fj @ cp.get_dofs() - cpst.d) / \ + (np.sqrt(dzeta_coil * dtheta_coil) * normn[:, None]) + nzeta_coil = cpst.nzeta_coil + test_K_2 = test_K_2.reshape(nzeta_coil, nzeta_coil // cp.nfp, 3) + + plt.figure(1) + plt.subplot(3, 3, 1) + plt.pcolor(test_K_1[:, :, 0]) + plt.colorbar() + plt.subplot(3, 3, 2) + plt.pcolor(test_K_1[:, :, 1]) + plt.colorbar() + plt.subplot(3, 3, 3) + plt.pcolor(test_K_1[:, :, 2]) + plt.colorbar() + plt.subplot(3, 3, 4) + plt.pcolor(test_K_2[:, :, 0]) + plt.colorbar() + plt.subplot(3, 3, 5) + plt.pcolor(test_K_2[:, :, 1]) + plt.colorbar() + plt.subplot(3, 3, 6) + plt.pcolor(test_K_2[:, :, 2]) + plt.colorbar() + plt.subplot(3, 3, 7) + plt.pcolor(test_K_3[:, :, 0]) + plt.colorbar() + plt.subplot(3, 3, 8) + plt.pcolor(test_K_3[:, :, 1]) + plt.colorbar() + plt.subplot(3, 3, 9) + plt.pcolor(test_K_3[:, :, 2]) + plt.colorbar() + + + # In[21]: + + plt.figure(2) + # Errors + plt.subplot(3, 3, 1) + plt.pcolor(test_K_1[:, :, 0] - test_K_3[:, :, 0]) + plt.colorbar() + plt.subplot(3, 3, 2) + plt.pcolor(test_K_1[:, :, 1] - test_K_3[:, :, 1]) + plt.colorbar() + plt.subplot(3, 3, 3) + plt.pcolor(test_K_1[:, :, 2] - test_K_3[:, :, 2]) + plt.colorbar() + plt.subplot(3, 3, 4) + plt.pcolor(test_K_2[:, :, 0] - test_K_3[:, :, 0]) + plt.colorbar() + plt.subplot(3, 3, 5) + plt.pcolor(test_K_2[:, :, 1] - test_K_3[:, :, 1]) + plt.colorbar() + plt.subplot(3, 3, 6) + plt.pcolor(test_K_2[:, :, 2] - test_K_3[:, :, 2]) + plt.colorbar() + plt.show() + + assert np.allclose(test_K_1, test_K_2) + assert np.allclose(test_K_1, test_K_3) if __name__ == "__main__": unittest.main()