cutCoilsFromRegcoil

#!/usr/bin/env python3

import sys
import numpy as np
from scipy.io import netcdf_file
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

print("Usage: cutCoilsFromRegcoil <regcoil_out.XXX> <nescin.XXX> <# of coils per half period> <thetaShift> <ilambda>")
print("This script assumes the contours do not zig-zag back and forth across the theta=0 line,")
print("after shifting the current potential by thetaShift grid points.")
print("The nescin file is used to provide the coil winding surface, so make sure this is consistent with the regcoil run.")
print("ilambda is the index in the lambda scan which you want to select.")

if len(sys.argv) != 6:
   print("Error! Wrong number of arguments")
   exit(1)

filename = sys.argv[1]
print(filename[:12])
if filename[:12] != 'regcoil_out.':
   print("Error! First argument should be regcoil_out.XXX")
   exit(1)
coilsFilename='coils.'+filename[12:-3]
print("coilsFilename:",coilsFilename)
nescinFilename = sys.argv[2]
coilsPerHalfPeriod = int(sys.argv[3])
thetaShift = int(sys.argv[4])
ilambda = int(sys.argv[5])
print("coilsPerHalfPeriod:",coilsPerHalfPeriod)
print("thetaShift:",thetaShift)
print("ilambda:",ilambda)

f = netcdf_file(filename,'r',mmap=False)
theta = f.variables['theta_coil'][()]
zeta = f.variables['zeta_coil'][()]
nfp = f.variables['nfp'][()]
net_poloidal_current_Amperes = f.variables['net_poloidal_current_Amperes'][()]
current_potential = f.variables['current_potential'][()]
# = f.variables[''][()]
# = f.variables[''][()]
f.close()

print("current_potential.shape",current_potential.shape)
if abs(net_poloidal_current_Amperes) > np.finfo(float).eps:
   data = current_potential[ilambda,:,:] / net_poloidal_current_Amperes * nfp
else:
   data = current_potential[ilambda,:,:] / np.max(current_potential[ilambda,:,:])
#data = current_potential[ilambda,:,:] / np.max(abs(current_potential[ilambda,:,:] ))

print("Theta before shift:")
print(theta)
# First apply 'roll' to be sure I use the same convention as numpy:
theta = np.roll(theta,thetaShift)
# Now just generate a new monotonic array with the correct first value:
theta = theta[0] + np.linspace(0,2*np.pi,len(theta),endpoint=False)
#theta = np.mod(theta,2*np.pi)
print("Theta after shift:")
print(theta)

data = np.roll(data,thetaShift,axis=1)

d = 2*np.pi/nfp
zeta_3 = np.concatenate((zeta-d, zeta, zeta+d))
data_3 = np.concatenate((data-1,data,data+1))
print("data_3.shape",data_3.shape)

#d=2*np.pi
#theta_3 = np.concatenate((theta-d, theta, theta+d))
#data_3x3 = np.concatenate((data_3,data_3,data_3),1)

fig = plt.figure()

contours = np.linspace(-1,2,coilsPerHalfPeriod*2*3+1)
d = contours[1]-contours[0]
contours = contours + d/2
#cdata = plt.contour(zeta_3,theta_3,np.transpose(data_3x3),contours)
cdata = plt.contour(zeta_3,theta,np.transpose(data_3),contours)
plt.colorbar()
plt.xlabel('zeta')
plt.ylabel('theta')

# Repeat with just the contours we care about:
contours = np.linspace(0,1,coilsPerHalfPeriod*2,endpoint=False)
d = contours[1]-contours[0]
contours = contours + d/2
#cdata = plt.contour(zeta_3,theta_3,np.transpose(data_3x3),contours)
cdata = plt.contour(zeta_3,theta,np.transpose(data_3),contours,colors='k')

#print(cdata.collections[0].get_paths())
print("cdata.collections:")
print(cdata.collections)
numCoilsFound = len(cdata.collections)
print("len(cdata.collections):",len(cdata.collections))
if numCoilsFound != 2*coilsPerHalfPeriod:
   print("WARNING!!! The expected number of coils was not the number found.")

contour_zeta=[]
contour_theta=[]
numCoils = 0
for j in range(numCoilsFound):
   p = cdata.collections[j].get_paths()[0]
   v = p.vertices
   # Make sure the contours have increasing theta:
   if v[1,1]<v[0,1]:
      v = np.flipud(v)

   # close the contours by adding a copy of the first point to the end
   #v = np.append(v, [v[0,:]], axis=0)
   for jfp in range(nfp):
      d = 2*np.pi/nfp*jfp
      contour_zeta.append(v[:,0]+d)
      contour_theta.append(v[:,1])
      numCoils += 1
      plt.plot(contour_zeta[-1],contour_theta[-1],'.-r',linewidth=1)
      plt.plot(contour_zeta[-1][0],contour_theta[-1][0],'sk')
      

# Now read nescin filename to map the theta-zeta coordinates of the contours to xyz coordinates.

f=open(nescinFilename,'r')

line = ''
while "np     iota_edge       phip_edge       curpol" not in line:
    line = f.readline()
line = f.readline()
nfp_nescin = int(line.split()[0])
print("nfp:",nfp_nescin)
if nfp != nfp_nescin:
   print("Error! nfp from regcoil_out does not match nfp from nescin!")
   exit(1)

#line = ''
#while "Number of fourier modes in table" not in line:
#    line = f.readline()
#line = f.readline()
#print("Number of Fourier modes in plasma surface from nescin file: ",line)
# Don't bother reading plasma surface.

contour_R = []
contour_Z = []
for j in range(numCoils):
   contour_R.append(contour_zeta[j]*0)
   contour_Z.append(contour_zeta[j]*0)

line = ''
#while "Number of fourier modes in table" not in line:
while "------ Current Surface" not in line:
    line = f.readline()
line = f.readline()
line = f.readline()
print("Number of Fourier modes in coil surface from nescin file: ",line)
nmodes = int(line)
line = f.readline()
line = f.readline()
for imode in range(nmodes):
    data = f.readline().split()
    m = int(data[0])
    #n = -int(data[1])*nfp
    n = -int(data[1]) * nfp
    # Sign flip in n because bnormal uses NESCOIL convention.
    # See bn_read_vmecf90.f line 89.
    crc = float(data[2])
    czs = float(data[3])
    crs = float(data[4])
    czc = float(data[5])
    # Skip remaining columns
    for j in range(numCoils):
       angle = m*contour_theta[j] - n*contour_zeta[j]
       contour_R[j] = contour_R[j] + crc*np.cos(angle) + crs*np.sin(angle)
       contour_Z[j] = contour_Z[j] + czs*np.sin(angle) + czc*np.cos(angle)
f.close()

contour_X = []
contour_Y = []
fig=plt.figure(figsize=(7,7))
ax = fig.add_subplot(projection='3d')
maxR=0
for j in range(numCoils):
   maxR = np.max((maxR,np.max(contour_R[j])))
   contour_X.append(contour_R[j]*np.cos(contour_zeta[j]))
   contour_Y.append(contour_R[j]*np.sin(contour_zeta[j]))
   ax.plot(contour_X[j],contour_Y[j],contour_Z[j],'.-')

ax.auto_scale_xyz([-maxR,maxR],[-maxR,maxR],[-maxR,maxR])

coilCurrent = net_poloidal_current_Amperes / numCoils

# Find the point of minimum separation
minSeparation2=1.0e+20
#for whichCoil1 in [5*nfp]:
#   for whichCoil2 in [4*nfp]:
for whichCoil1 in range(numCoils):
   for whichCoil2 in range(whichCoil1):
      for whichPoint in range(len(contour_X[whichCoil1])):
         dx = contour_X[whichCoil1][whichPoint] - contour_X[whichCoil2]
         dy = contour_Y[whichCoil1][whichPoint] - contour_Y[whichCoil2]
         dz = contour_Z[whichCoil1][whichPoint] - contour_Z[whichCoil2]
         separation2 = dx*dx+dy*dy+dz*dz
         this_minSeparation2 = np.min(separation2)
         if this_minSeparation2<minSeparation2:
            minSeparation2 = this_minSeparation2
            x1 = contour_X[whichCoil1][whichPoint]
            y1 = contour_Y[whichCoil1][whichPoint]
            z1 = contour_Z[whichCoil1][whichPoint]
            index=np.argmin(separation2)
            x2 = contour_X[whichCoil2][index]
            y2 = contour_Y[whichCoil2][index]
            z2 = contour_Z[whichCoil2][index]
            
print("Minimum coil separation:",np.sqrt(minSeparation2))
ax.plot([x1,x2],[y1,y2],[z1,z2],'k',linewidth=3)


# Write coils file
f = open(coilsFilename,'w')
f.write('periods '+str(nfp)+'\n')
f.write('begin filament\n')
f.write('mirror NIL\n')

for j in range(numCoils):
   N = len(contour_X[j])
   for k in range(N):
      f.write('{:14.22e} {:14.22e} {:14.22e} {:14.22e}\n'.format(contour_X[j][k],contour_Y[j][k],contour_Z[j][k],coilCurrent))
   # Close the loop
   k=0
   f.write('{:14.22e} {:14.22e} {:14.22e} {:14.22e} 1 Modular\n'.format(contour_X[j][k],contour_Y[j][k],contour_Z[j][k],0))

f.write('end\n')
f.close()

plt.show()