diff --git a/megpy/vmec_funcs/vmec2mil_fourier_funcs.py b/megpy/vmec_funcs/vmec2mil_fourier_funcs.py new file mode 100644 index 0000000..e519982 --- /dev/null +++ b/megpy/vmec_funcs/vmec2mil_fourier_funcs.py @@ -0,0 +1,245 @@ +import scipy +import math +import numpy as np +import matplotlib.pyplot as plt +import scipy.io as io +from scipy.optimize import minimize +from mpi4py import MPI +from scipy.interpolate import UnivariateSpline as sp + + +# make a function to read the file +def read_file(filename): + file2read = io.netcdf_file(filename,'r',mmap=False) + # retrieve rmnc, zmns, xn, ns, and nfp + rmnc = file2read.variables['rmnc'] + zmns = file2read.variables['zmns'] + xn = (file2read.variables['xn']).data + xm = (file2read.variables['xm']).data + ns = (file2read.variables['ns']).data + nfp = (file2read.variables['nfp']).data + # retrieve pres and iotas + pres = file2read.variables['presf'].data + iota = file2read.variables['iotaf'].data + pres = pres[1::] + iota = iota[1::] + phi_last = file2read.variables['phi'].data[-1] + edge_toroidal_flux_over_2pi = phi_last / (2*np.pi) + L_reference = file2read.variables['Aminor_p'].data + # close the file + file2read.close() + return rmnc, zmns, xn, xm, ns, nfp, pres, iota, edge_toroidal_flux_over_2pi, L_reference + + +# return (R,Z) coordinates of the vmec cross section +def vmec_RZ(filename,fit_modes=2**8): + rmnc, zmns, xn, xm, ns, nfp, pres, iota, edge_toroidal_flux_over_2pi, L_reference = read_file(filename) + # make grid of theta and zeta + nzeta= 1 + theta = np.linspace(0,2*np.pi,num=fit_modes) + zeta = np.linspace(0,2*np.pi/nfp,num=nzeta,endpoint=False) + # number of modes + nmodes = len(xn) + iradius = ns-1 + R = np.zeros((ns-1,fit_modes,nzeta)) + Z = np.zeros((ns-1,fit_modes,nzeta)) + for iradius in range(ns-1): + for itheta in range(fit_modes): + for izeta in range(nzeta): + for imode in range(nmodes): + angle = xm[imode]*theta[itheta] - xn[imode]*zeta[izeta] + R[iradius,itheta,izeta] = R[iradius,itheta,izeta] + rmnc[iradius,imode]*math.cos(angle) + Z[iradius,itheta,izeta] = Z[iradius,itheta,izeta] + zmns[iradius,imode]*math.sin(angle) + return R, Z + + +# return (R,Z) coordinates of Miller cross section +def miller_RZ(R0,Z0,r,kappa,delta,zeta,fit_modes=2**8): + theta = np.linspace(0,2*np.pi,num=fit_modes) + Z = Z0 + r * kappa * np.sin(theta + zeta * np.sin(2*theta)) + R = R0 + r * np.cos(theta + np.arcsin(delta) * np.sin(theta)) + return R,Z + + +# return the fourier coefficients of the cross section in (R,Z) coordinates +def RZ_fourier(R,Z): + # create complex (R,Z) + C = R + 1j*Z + # now do fft + C_hat = np.fft.fft(C,norm='forward') + return C_hat + + +# define fitting function which returns the difference between the Fourier coefficients +def fit_func(params,R,Z,fit_modes): + # first, get the Fourier coefficients of the VMEC cross section + C_hat = RZ_fourier(R,Z) + # now, get the Fourier coefficients of the Miller cross section + R0 = params[0] + Z0 = params[1] + r = params[2] + kappa = params[3] + delta = params[4] + zeta = params[5] + Rm,Zm = miller_RZ(R0,Z0,r,kappa,delta,zeta,fit_modes=fit_modes) + Cm_hat = RZ_fourier(Rm,Zm) + # return the difference between the two + diff = np.abs(C_hat - Cm_hat) + return np.sum(diff**2) + + +# finally, do the fitting for all the cross sections. To this end, loop over all radii +def fit_all_cross_sections(R,Z,fit_modes): + # get number of radii + nr = R.shape[0] + # initialize arrays + R0 = np.zeros(nr) + Z0 = np.zeros(nr) + r = np.zeros(nr) + kappa = np.zeros(nr) + delta = np.zeros(nr) + zeta = np.zeros(nr) + # define initial guess + params0 = np.array([1.0,0.0,0.1,1.0,0.0,0.0]) + # loop over all radii + for ir in range(nr): + # get the (R,Z) coordinates + Rr = (R[ir,:,:]).flatten() + Zr = (Z[ir,:,:]).flatten() + # do the minimization, bounded. + # R0 should be positive + # Z0 can be anything + # r should be positive + # kappa should be positive + # delta should be between -1 and 1 + # zeta can be anything + bnds = ((0.0,None),(None,None),(0.0,None),(0.0,None),(-1.0,1.0),(None,None)) + res = minimize(fit_func,params0,args=(Rr,Zr,fit_modes),bounds=bnds,tol=1e-16) + # print the minimization result + # get the result + R0[ir] = res.x[0] + Z0[ir] = res.x[1] + r[ir] = res.x[2] + kappa[ir] = res.x[3] + delta[ir] = res.x[4] + zeta[ir] = res.x[5] + # update the initial guess + # params0 = res.x + + return R0, Z0, r, kappa, delta, zeta + + +# now make a wrapper function which takes as input the filename and outputs the Miller parameters +def vmec2miller(filename,fit_modes=2**6): + # first, get the (R,Z) coordinates of the VMEC cross section + R, Z = vmec_RZ(filename,fit_modes=fit_modes) + # now, do the fitting + R0, Z0, r, kappa, delta, zeta = fit_all_cross_sections(R,Z,fit_modes) + # we construct the r derivatives of the Miller parameters + # this is done with splines + # let us first define the spline parameters + k_int = 3 + ext_int = 0 + s_val = 0.0 + # also retrieve pres, iota, Bref + rmnc, zmns, xn, xm, ns, nfp, pres, iota, edge_toroidal_flux_over_2pi, L_reference = read_file(filename) + Bref = 2 * np.abs(edge_toroidal_flux_over_2pi) / L_reference**2 + mu0 = 4*np.pi*1e-7 + mag_pres= Bref**2 / 2 / mu0 + # now construct the spline + pres_spline = sp(r,pres,k=k_int,ext=ext_int,s=s_val) + iota_spline = sp(r,iota,k=k_int,ext=ext_int,s=s_val) + R0_spline = sp(r,R0,k=k_int,ext=ext_int,s=s_val) + Z0_spline = sp(r,Z0,k=k_int,ext=ext_int,s=s_val) + kappa_spline = sp(r,kappa,k=k_int,ext=ext_int,s=s_val) + delta_spline = sp(r,delta,k=k_int,ext=ext_int,s=s_val) + zeta_spline = sp(r,zeta,k=k_int,ext=ext_int,s=s_val) + # now, define the derivatives + presp = pres_spline.derivative() + iotap = iota_spline.derivative() + R0p = R0_spline.derivative() + Z0p = Z0_spline.derivative() + kappap = kappa_spline.derivative() + deltap = delta_spline.derivative() + zetap = zeta_spline.derivative() + # finally, define the Miller parameters + alpha = - (1/iota_spline(r))**2 * R0_spline(r) * presp(r) / mag_pres + shear = - r * iotap(r)/iota_spline(r) + dR0dr = R0p(r) + dZ0dr = Z0p(r) + s_kappa = r * kappap(r)/kappa_spline(r) + s_delta = r * deltap(r)/np.sqrt(1 - delta_spline(r)**2) + s_zeta = r * zetap(r) + # return the Miller parameters (R0,Z0,r,kappa,delta,zeta,dR0dr,dZ0dr,s_kappa,s_delta,s_zeta,1/iota,shear,alpha,Bref) + return R0,Z0,r,kappa,delta,zeta,dR0dr,dZ0dr,s_kappa,s_delta,s_zeta,1/iota,shear,alpha,Bref + + +test = True +# now, let us test the code +if test==True: + # run test on wout_Lmode.nc + name = 'wout_output.nc' + R0,Z0,r,kappa,delta,zeta,dR0dr,dZ0dr,s_kappa,s_delta,s_zeta,q,shear,alpha,Bref = vmec2miller(name) + + # plot the results as a function of radius in several subplots. print Bref + fig, ax = plt.subplots(4,4,figsize=(8,8),tight_layout=True,sharex=True) + ax[0,0].plot(r,R0) + ax[0,0].set_title('R0') + ax[0,1].plot(r,Z0) + ax[0,1].set_title('Z0') + ax[0,2].plot(r,dR0dr) + ax[0,2].set_title('dR0dr') + ax[0,3].plot(r,dZ0dr) + ax[0,3].set_title('dZ0dr') + ax[1,0].plot(r,r) + ax[1,0].set_title('r') + ax[1,1].plot(r,kappa) + ax[1,1].set_title('kappa') + ax[1,2].plot(r,delta) + ax[1,2].set_title('delta') + ax[1,3].plot(r,zeta) + ax[1,3].set_title('zeta') + ax[2,0].plot(r,s_kappa) + ax[2,0].set_title('s_kappa') + ax[2,1].plot(r,s_delta) + ax[2,1].set_title('s_delta') + ax[2,2].plot(r,s_zeta) + ax[2,2].set_title('s_zeta') + ax[2,3].plot(r,q) + ax[2,3].set_title('q') + ax[3,0].plot(r,shear) + ax[3,0].set_title('shear') + ax[3,1].plot(r,alpha) + ax[3,1].set_title('alpha') + + # print Bref + print('Bref = ',Bref) + + # show the plot + plt.show() + + # also plot 3 cross sections to compare fit with original + # do for 3 radii: inner, middle, outer + # import R,Z from file + R, Z = vmec_RZ(name,fit_modes=2**8) + fig, ax = plt.subplots(1,3,figsize=(12,4),tight_layout=True) + # inner + ax[0].plot(R[1,:,0],Z[1,:,0],label='VMEC') + Rm,Zm = miller_RZ(R0[1],Z0[1],r[1],kappa[1],delta[1],zeta[1],fit_modes=2**8) + ax[0].plot(Rm,Zm,label='Miller') + ax[0].set_title('inner') + ax[0].legend() + # middle, R.shape[0]//2 + ax[1].plot(R[R.shape[0]//2,:,0],Z[R.shape[0]//2,:,0],label='VMEC') + Rm,Zm = miller_RZ(R0[R.shape[0]//2],Z0[R.shape[0]//2],r[R.shape[0]//2],kappa[R.shape[0]//2],delta[R.shape[0]//2],zeta[R.shape[0]//2],fit_modes=2**8) + ax[1].plot(Rm,Zm,label='Miller') + ax[1].set_title('middle') + ax[1].legend() + # outer + ax[2].plot(R[-1,:,0],Z[-1,:,0],label='VMEC') + Rm,Zm = miller_RZ(R0[-1],Z0[-1],r[-1],kappa[-1],delta[-1],zeta[-1],fit_modes=2**8) + ax[2].plot(Rm,Zm,label='Miller') + ax[2].set_title('outer') + ax[2].legend() + # show the plot + plt.show() diff --git a/megpy/vmec_funcs/wout_output.nc b/megpy/vmec_funcs/wout_output.nc new file mode 100644 index 0000000..1b8976a Binary files /dev/null and b/megpy/vmec_funcs/wout_output.nc differ