Skip to content

Commit

Permalink
Merge pull request #14 from PrincetonUniversity/pySTEL
Browse files Browse the repository at this point in the history
Updates to pySTEL
  • Loading branch information
lazersos authored Aug 30, 2019
2 parents b3b204c + a1e4c3a commit bd4dca8
Show file tree
Hide file tree
Showing 3 changed files with 212 additions and 25 deletions.
54 changes: 34 additions & 20 deletions pySTEL/STELLOPT.py
Original file line number Diff line number Diff line change
Expand Up @@ -823,15 +823,17 @@ def LoadSTELLOPT(self):
self.ui.ComboBoxOPTplot_type.addItem('Iota')
self.ui.ComboBoxOPTplot_type.addItem('q-prof')
self.ui.ComboBoxOPTplot_type.addItem('<j*B>')
self.wout_files = sorted([k for k in files if 'wout' in k])
wout_files = sorted([k for k in files if 'wout' in k])
self.wout_files = sorted([k for k in wout_files if '_opt' not in k])
# Handle Kinetic Profiles
if any('tprof.' in mystring for mystring in files):
self.ui.ComboBoxOPTplot_type.addItem('----- Kinetics -----')
self.ui.ComboBoxOPTplot_type.addItem('Electron Temperature')
self.ui.ComboBoxOPTplot_type.addItem('Electron Density')
self.ui.ComboBoxOPTplot_type.addItem('Ion Temperature')
self.ui.ComboBoxOPTplot_type.addItem('Z Effective')
self.tprof_files = sorted([k for k in files if 'tprof.' in k])
tprof_files = sorted([k for k in files if 'tprof.' in k])
self.tprof_files = sorted([k for k in tprof_files if '_opt' not in k])
# Handle Diagnostic Profiles
if any('dprof.' in mystring for mystring in files):
self.ui.ComboBoxOPTplot_type.addItem('----- Diagnostic -----')
Expand All @@ -844,7 +846,8 @@ def LoadSTELLOPT(self):
self.ui.ComboBoxOPTplot_type.addItem('Bootstrap Profile')
self.ui.ComboBoxOPTplot_type.addItem('Beam Profile')
self.ui.ComboBoxOPTplot_type.addItem('Total Current Profile')
self.jprof_files = sorted([k for k in files if 'jprof.' in k])
jprof_files = sorted([k for k in files if 'tprof.' in k])
self.jprof_files = sorted([k for k in jprof_files if '_opt' not in k])


def UpdateOptplot(self):
Expand Down Expand Up @@ -916,8 +919,8 @@ def UpdateOptplot(self):
self.ax2.set_xlabel('Radial Grid')
self.ax2.set_ylabel('Epsilon Effective')
self.ax2.set_title('Neoclassical Helical Ripple (NEO)')
elif (plot_name == 'HELICITY_evolution'):
self.ax2.plot(self.stel_data['HELICITY_equil'].T,'o',fillstyle='none')
elif (plot_name == 'HELICITY_FULL_evolution'):
self.ax2.plot(self.stel_data['HELICITY_FULL_equil'].T,'o',fillstyle='none')
self.ax2.set_ylabel('Helicity')
self.ax2.set_title('Boozer Spectrum Helicity')
elif (plot_name == 'B_PROBE_evolution'):
Expand Down Expand Up @@ -1487,7 +1490,8 @@ def UpdateOptplot(self):
self.ax2.set_title('XICS Velocity Reconstruction')
elif (plot_name == 'Pressure'):
l=0
dl = len(self.wout_files)
dl = len(self.wout_files)-1
if dl == 0 : dl = 1
for string in self.wout_files:
if 'wout' in string:
vmec_data=read_vmec(self.workdir+string)
Expand All @@ -1502,7 +1506,8 @@ def UpdateOptplot(self):
self.ax2.set_xlim((0,1))
elif (plot_name == 'I-prime'):
l=0
dl = len(self.wout_files)
dl = len(self.wout_files)-1
if dl == 0 : dl = 1
for string in self.wout_files:
if 'wout' in string:
vmec_data=read_vmec(self.workdir+string)
Expand All @@ -1517,7 +1522,7 @@ def UpdateOptplot(self):
self.ax2.set_xlim((0,1))
elif (plot_name == 'Iota'):
l=0
dl = len(self.wout_files)
dl = len(self.wout_files)-1
for string in self.wout_files:
if 'wout' in string:
vmec_data=read_vmec(self.workdir+string)
Expand Down Expand Up @@ -1547,7 +1552,7 @@ def UpdateOptplot(self):
self.ax2.set_xlim((0,1))
elif (plot_name == 'Current'):
l=0
dl = len(self.wout_files)
dl = len(self.wout_files)-1
for string in self.wout_files:
if 'wout' in string:
vmec_data=read_vmec(self.workdir+string)
Expand All @@ -1562,7 +1567,7 @@ def UpdateOptplot(self):
self.ax2.set_xlim((0,1))
elif (plot_name == '<j*B>'):
l=0
dl = len(self.wout_files)
dl = len(self.wout_files)-1
for string in self.wout_files:
if 'wout' in string:
vmec_data=read_vmec(self.workdir+string)
Expand All @@ -1577,7 +1582,7 @@ def UpdateOptplot(self):
self.ax2.set_xlim((0,1))
elif (plot_name == 'Flux0'):
l=0
dl = len(self.wout_files)
dl = len(self.wout_files)-1
for string in self.wout_files:
if 'wout' in string:
vmec_data=read_vmec(self.workdir+string)
Expand All @@ -1601,7 +1606,8 @@ def UpdateOptplot(self):
self.ax2.set_aspect('equal')
elif (plot_name == 'FluxPI'):
l=0
dl = len(self.wout_files)
dl = len(self.wout_files)-1
if dl == 0 : dl = 1
for string in self.wout_files:
if 'wout' in string:
vmec_data=read_vmec(self.workdir+string)
Expand All @@ -1625,7 +1631,8 @@ def UpdateOptplot(self):
self.ax2.set_aspect('equal')
elif (plot_name == 'Electron Temperature'):
l=0
dl = len(self.tprof_files)
dl = len(self.tprof_files)-1
if dl == 0 : dl = 1
for string in self.tprof_files:
if 'tprof' in string:
tprof = np.loadtxt(self.workdir+string,skiprows=1)
Expand All @@ -1637,7 +1644,8 @@ def UpdateOptplot(self):
self.ax2.set_xlim((0,1))
elif (plot_name == 'Electron Density'):
l=0
dl = len(self.tprof_files)
dl = len(self.tprof_files)-1
if dl == 0 : dl = 1
for string in self.tprof_files:
if 'tprof' in string:
tprof = np.loadtxt(self.workdir+string,skiprows=1)
Expand All @@ -1649,7 +1657,8 @@ def UpdateOptplot(self):
self.ax2.set_xlim((0,1))
elif (plot_name == 'Ion Temperature'):
l=0
dl = len(self.tprof_files)
dl = len(self.tprof_files)-1
if dl == 0 : dl = 1
for string in self.tprof_files:
if 'tprof' in string:
tprof = np.loadtxt(self.workdir+string,skiprows=1)
Expand All @@ -1661,7 +1670,8 @@ def UpdateOptplot(self):
self.ax2.set_xlim((0,1))
elif (plot_name == 'Z Effective'):
l=0
dl = len(self.tprof_files)
dl = len(self.tprof_files)-1
if dl == 0 : dl = 1
for string in self.tprof_files:
if 'tprof' in string:
tprof = np.loadtxt(self.workdir+string,skiprows=1)
Expand All @@ -1673,7 +1683,8 @@ def UpdateOptplot(self):
self.ax2.set_xlim((0,1))
elif (plot_name == 'XICS Emissivity'):
l=0
dl = len(self.dprof_files)
dl = len(self.dprof_files)-1
if dl == 0 : dl = 1
for string in self.dprof_files:
if 'dprof' in string:
dprof = np.loadtxt(self.workdir+string,skiprows=1)
Expand All @@ -1697,7 +1708,8 @@ def UpdateOptplot(self):
self.ax2.set_xlim((0,1))
elif (plot_name == 'Bootstrap Profile'):
l=0
dl = len(self.jprof_files)
dl = len(self.jprof_files)-1
if dl == 0 : dl = 1
for string in self.jprof_files:
if 'jprof' in string:
jprof = np.loadtxt(self.workdir+string,skiprows=1)
Expand All @@ -1709,7 +1721,8 @@ def UpdateOptplot(self):
self.ax2.set_xlim((0,1))
elif (plot_name == 'Beam Profile'):
l=0
dl = len(self.jprof_files)
dl = len(self.jprof_files)-1
if dl == 0 : dl = 1
for string in self.jprof_files:
if 'jprof' in string:
jprof = np.loadtxt(self.workdir+string,skiprows=1)
Expand All @@ -1721,7 +1734,8 @@ def UpdateOptplot(self):
self.ax2.set_xlim((0,1))
elif (plot_name == 'Total Current Profile'):
l=0
dl = len(self.jprof_files)
dl = len(self.jprof_files)-1
if dl == 0 : dl = 1
jprof = np.loadtxt(self.workdir+self.jprof_files[0],skiprows=1)
self.ax2.plot(jprof[:,0],jprof[:,2],'--b')
self.ax2.plot(jprof[:,0],jprof[:,1],':b')
Expand Down
51 changes: 46 additions & 5 deletions pySTEL/STELLOPTrefit.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,46 @@
#!/usr/bin/env python3
import sys, os, globsys.path.insert(0, '../../../pySTEL/')
import sys, os, glob
sys.path.insert(0, '../../../pySTEL/')
import ctypes as ct
import numpy as np #For Arrays
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
from math import pi
from libstell.stellopt import read_stellopt


def fit_poly10(x,am0,am1,am2,am3,am4,am5,am6,am7,am8,am9,am10):
f2=np.zeros(x.shape)
for i,xx in enumerate(x):
f2[i] = (f2[i]+am10)*xx
f2[i] = (f2[i]+am9)*xx
f2[i] = (f2[i]+am8)*xx
f2[i] = (f2[i]+am7)*xx
f2[i] = (f2[i]+am6)*xx
f2[i] = (f2[i]+am5)*xx
f2[i] = (f2[i]+am4)*xx
f2[i] = (f2[i]+am3)*xx
f2[i] = (f2[i]+am2)*xx
f2[i] = (f2[i]+am1)*xx
f2[i] = f2[i]+am0
if (xx>1):
f[i]=0
return f


def fit_poly5(x,am0,am1,am2,am3,am4,am5):
f2=np.zeros(x.shape)
for i,xx in enumerate(x):
f2[i] = (f2[i]+am5)*xx
f2[i] = (f2[i]+am4)*xx
f2[i] = (f2[i]+am3)*xx
f2[i] = (f2[i]+am2)*xx
f2[i] = (f2[i]+am1)*xx
f2[i] = f2[i]+am0
if (xx>1):
f2[i]=0
return f2

try:
qtCreatorPath=os.environ["STELLOPT_PATH"]
except KeyError:
Expand All @@ -17,9 +54,13 @@
data=read_stellopt(files[0])
ndex=data['ITER'].size
print(' Iterations Detected: '+str(ndex))
print('----- Fitting ne -----')
print('----- Fitting Ne -----')
s=data['NE_s'][ndex-1,:]
f=data['NE_target'][ndex-1,:]
s_fit=[0,0.04,0.16,0.36,0.64,1.0]
f_fit=np.polyfit(s,f,npoly)
print(f_fit)
s_fit=np.array([0,0.04,0.16,0.36,0.64,1.0])
f_fit = curve_fit(fit_poly5, s, f)
f_err = np.sqrt(np.diag(f_fit[1]))
f_fit = f_fit[0]
plt.plot(s,f,'o')
plt.plot(s_fit,fit_poly5(s_fit,f_fit[0],f_fit[1],f_fit[2],f_fit[3],f_fit[4],f_fit[5]),'+')
plt.show()
132 changes: 132 additions & 0 deletions pySTEL/vmec_spectrum.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,132 @@
#!/usr/bin/env python3
import sys, os, getopt
import argparse
import numpy as np #For Arrays
from libstell.libstell import read_vmec, cfunct, sfunct, torocont, isotoro, calc_jll

try:
qtCreatorPath=os.environ["STELLOPT_PATH"]
except KeyError:
print("Please set environment variable STELLOPT_PATH")
sys.exit(1)

if __name__ == "__main__":
#print(sys.argv)
lvmec=0
lspec=0
lfocus=0
lflip=0
parser = argparse.ArgumentParser(description='Outputs VMEC harmonics to stdout.')
#print(parser)
parser.add_argument('filename',help='Filename (wout file) to read.')
parser.add_argument('-v','--vmec', action='store_true', dest='lvmec', help='Print VMEC harmonics.')
parser.add_argument('-f','--focus', action='store_true', dest='lfocus', help='Print FOCUS harmonics.')
parser.add_argument('-s','--spec', action='store_true', dest='lspec', help='Print SPEC harmonics.')
parser.add_argument('--flip', action='store_true', dest='lflip', help='Flip the jacobian.')
parser.add_argument('-k', action='store', dest='k', type=int, help='Evaluate Surface (default: ns)')
args = parser.parse_args()
lvmec = args.lvmec
lspec = args.lspec
lfocus = args.lfocus
lflip = args.lflip
vmec_data=read_vmec(args.filename)
# note that we use nu+nv in pystel
k = vmec_data['ns']
if args.k:
k = args.k
# Flip Jacobian
if lflip:
for i in range(vmec_data['mnmax']):
m = vmec_data['xm'][i]
if m > 0:
if np.remainder(m,2) == 0:
vmec_data['zmns'][k-1,i] = -vmec_data['zmns'][k-1,i]
else:
vmec_data['rmnc'][k-1,i] = -vmec_data['rmnc'][k-1,i]
vmec_data['xn'] = - vmec_data['xn']
print('!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!')
print('!!!!!!!!!!!!!!!! JACOBIAN SIGN FLIPPPED !!!!!!!!!!!!!!!')
print('!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!')
if lvmec:
ntor = vmec_data['ntor']
temp = ' RAXIS_CC = '
for i in range(ntor+1):
temp = temp + "{:20.10e}".format(vmec_data['rmnc'][0,i]) + ' '
print(temp)
temp = ' ZAXIS_CS = '
for i in range(ntor+1):
temp = temp + "{:20.10e}".format(vmec_data['zmns'][0,i]) + ' '
print(temp)
if vmec_data['lasym']:
temp = ' RAXIS_CS = '
for i in range(ntor+1):
temp = temp + "{:20.10e}".format(vmec_data['rmns'][0,i]) + ' '
print(temp)
temp = ' ZAXIS_CC = '
for i in range(ntor+1):
temp = temp + "{:20.10e}".format(vmec_data['zmnc'][0,i]) + ' '
print(temp)
for i in range(vmec_data['mnmax']):
temp = ' RBC(' + "{:3d}".format(int(vmec_data['xm'][i])) + ',' \
+ "{:3d}".format(-int(vmec_data['xn'][i]/vmec_data['nfp'])) + \
') = ' + "{:20.10e}".format(vmec_data['rmnc'][k-1,i]) + \
' ZBS(' + "{:3d}".format(int(vmec_data['xm'][i])) + ',' \
+ "{:3d}".format(-int(vmec_data['xn'][i]/vmec_data['nfp'])) + \
') = ' + "{:20.10e}".format(vmec_data['zmns'][k-1,i])
print(temp)
if vmec_data['lasym']:
for i in range(vmec_data['mnmax']):
temp = ' RBS(' + "{:3d}".format(int(vmec_data['xm'][i])) + ',' \
+ "{:3d}".format(-int(vmec_data['xn'][i]/vmec_data['nfp'])) + \
') = ' + "{:20.10e}".format(vmec_data['rmns'][k-1,i]) + \
' ZBC(' + "{:3d}".format(int(vmec_data['xm'][i])) + ',' \
+ "{:3d}".format(-int(vmec_data['xn'][i]/vmec_data['nfp'])) + \
') = ' + "{:20.10e}".format(vmec_data['zmnc'][k-1,i])
print(temp)
if lspec:
if (vmec_data['lasym']):
print('mn m n rmnc zmns rmns zmnc')
for i in range(vmec_data['mnmax']):
temp = "{:3d}".format(i) + ' ' + \
"{:3d}".format(int(vmec_data['xm'][i])) + ' ' + \
"{:3d}".format(int(vmec_data['xn'][i]/vmec_data['nfp'])) + ' ' + \
"{:20.10e}".format(vmec_data['rmnc'][k-1,i]) + ' ' + \
"{:20.10e}".format(vmec_data['zmns'][k-1,i]) + ' ' + \
"{:20.10e}".format(vmec_data['rmns'][k-1,i]) + ' ' + \
"{:20.10e}".format(vmec_data['zmnc'][k-1,i])
print(temp)
else:
print('mn m n rmnc zmns rmns zmnc')
for i in range(vmec_data['mnmax']):
temp = "{:3d}".format(i) + ' ' + \
"{:3d}".format(int(vmec_data['xm'][i])) + ' ' + \
"{:3d}".format(int(vmec_data['xn'][i]/vmec_data['nfp'])) + ' ' + \
"{:20.10e}".format(vmec_data['rmnc'][k-1,i]) + ' ' + \
"{:20.10e}".format(vmec_data['zmns'][k-1,i]) + ' ' + \
"{:20.10e}".format(0.0) + ' ' + \
"{:20.10e}".format(0.0)
print(temp)
if lfocus:
print('#bmn bNfp nbf')
print("{:3d}".format(vmec_data['mnmax']) + ' ' + "{:3d}".format(vmec_data['nfp']) + ' ' + "{:3d}".format(vmec_data['mnmax']))
print('#------plasma boundary harmonics-------')
print('# n m Rbc Rbs Zbc Zbs')
if (vmec_data['lasym']):
for i in range(vmec_data['mnmax']):
temp = "{:3d}".format(int(vmec_data['xn'][i]/vmec_data['nfp'])) + ' ' + \
"{:3d}".format(int(vmec_data['xm'][i])) + ' ' + \
"{:20.10e}".format(vmec_data['rmnc'][k-1,i]) + ' ' + \
"{:20.10e}".format(vmec_data['rmns'][k-1,i]) + ' ' + \
"{:20.10e}".format(vmec_data['zmnc'][k-1,i]) + ' ' + \
"{:20.10e}".format(vmec_data['zmns'][k-1,i])
print(temp)
else:
for i in range(vmec_data['mnmax']):
temp = "{:3d}".format(int(vmec_data['xn'][i]/vmec_data['nfp'])) + ' ' + \
"{:3d}".format(int(vmec_data['xm'][i])) + ' ' + \
"{:20.10e}".format(vmec_data['rmnc'][k-1,i]) + ' ' + \
"{:20.10e}".format(0.0) + ' ' + \
"{:20.10e}".format(0.0) + ' ' + \
"{:20.10e}".format(vmec_data['zmns'][k-1,i])
print(temp)

0 comments on commit bd4dca8

Please sign in to comment.