Skip to content

Commit

Permalink
Fixed the issue with calculating K in different ways and added it to …
Browse files Browse the repository at this point in the history 
…the unit tests. However, now need to fix the lasso problem, which has a factor of nnorm from the winding surface flying around. Testing this with the appropriate unit test and should be done soon.
  • Loading branch information
@akaptano
akaptano committed Jul 26, 2024
1 parent 54251e8 commit ef26800
Show file tree
Hide file tree
Showing 5 changed files with 199 additions and 34 deletions.
63 changes: 45 additions & 18 deletions src/simsopt/field/currentpotential.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
22 changes: 18 additions & 4 deletions src/simsopt/field/currentpotentialsolve.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand All @@ -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 = []
Expand All @@ -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
Expand Down
18 changes: 9 additions & 9 deletions src/simsoptpp/winding_surface.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -470,21 +470,21 @@ std::tuple<Array, Array> 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));
}
}
}
Expand Down
4 changes: 2 additions & 2 deletions tests/field/test_currentpotential.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)


Expand Down
126 changes: 125 additions & 1 deletion tests/field/test_regcoil.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand All @@ -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)

Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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()

0 comments on commit ef26800

Please sign in to comment.