def make1EFITobject(self, shot, t, gfile=None):
"""
Creates an equilParams_class object for a SINGLE timesteps. equilParams
is a class from the ORNL_Fusion github repo, and was developed by
A. Wingen
"""
if self.shotPath[-1] != '/':
shotPath = self.shotPath + '/'
else:
shotPath = self.shotPath
if gfile is None:
gfile = shotPath + '{:06d}/g{:06d}.{:05d}'.format(t, shot, t)
self.ep = EP.equilParams(gfile)
#Correct for weird helicity and field directions that are occasionally
#in gfile. Here we assume that Ip is the CCW direction as viewed from
#above tokamak.
# self.dsign = np.sign(self.ep.g['Ip'])
# self.gsign = np.sign( (self.ep.g['psiAxis'] - self.ep.g['psiSep']) )
# self.qsign = np.sign(self.ep.g['Fpol'][-1]) #F_edge sign
# print('dsign = {:f}'.format(self.dsign))
# print('gsign = {:f}'.format(self.gsign))
# print('qsign = {:f}'.format(self.qsign))
# self.ep.g['FFprime'] *= self.dsign*self.qsign*self.gsign
# self.ep.g['Pprime'] *= self.dsign*self.qsign*self.gsign
# self.ep.g['psiRZ'] *= self.dsign*self.qsign*self.gsign
# self.ep.g['qpsi'] *= -self.qsign*self.dsign
# self.ep.g['psiSep'] *= self.dsign*self.qsign*self.gsign
# self.ep.g['psiAxis'] *= self.dsign*self.qsign*self.gsign
# self.ep.g['Fpol'] *= self.dsign*self.qsign
return
def makeEFITobjects(self):
"""
Creates an equilParams_class object for MULTIPLE timesteps. equilParams
is a class from the ORNL_Fusion github repo, and was developed by
A. Wingen
gfiles should be placed in the dataPath before running this function
"""
self.ep = ['None' for i in range(len(self.timesteps))]
for idx, t in enumerate(self.timesteps):
if self.shotPath[-1] != '/':
shotPath = self.shotPath + '/'
else:
shotPath = self.shotPath
gfile = shotPath + '{:06d}/g{:06d}.{:05d}'.format(t, self.shot, t)
self.ep[idx] = EP.equilParams(gfile)
return
def setupForTerminalUse(gFile=None):
"""
Sets up an MHD object so that it can be used from python console
without running HEAT. This is convenient when a user wants to load
MHD EQ directly from the python console.
To use, user needs to set pythonpath to include path of this file and EFIT:
import sys
EFITPath = '/home/tom/source'
HEATPath = '/home/tom/source/HEAT/github/source'
sys.path.append(EFITPath)
sys.path.append(HEATPath)
import MHDClass
Then user needs to run this function to set up the MHD object:
MHD = MHDClass.setupForTerminalUse()
if user wants to load MHD EQ, they need to set gFile
arguments when calling this function, otherwise it just returns an
empty MHD object:
MHD = MHDClass.setupForTerminalUse(gFile=gFilePath)
if you want access some gFile parameters, use the ep.g:
Example to see available parameters:
MHD.ep.g.keys()
Example to see Rlim, Zlim:
MHD.ep.g['wall']
returns an MHD class object ie: MHD = MHDClass.MHD()
"""
rootDir = ''
dataPath = ''
import MHDClass
MHD = MHDClass.MHD(rootDir, dataPath)
if gFile != None:
print("Making MHD() object with ep included")
MHD.ep = EP.equilParams(gFile)
else:
print("Not including ep in MHD() object")
return MHD
def copyGfile2tree(self, gFileName, idx, clobberflag=True):
"""
Copies gfile to HEAT tree
gFileName is name of gFile that is already located in self.tmpDir
"""
oldgfile = self.tmpDir + gFileName
#try to make EP object if naming follows d3d gFile naming convention
try:
ep = EP.equilParams(oldgfile)
shot = ep.g['shot']
time = ep.g['time']
#if gfile doesn't follow naming convention define manually
except:
print("Couldn't open gFile with equilParams_class")
log.info("Couldn't open gFile with equilParams_class")
if self.shot is None:
shot = 1
else:
shot = self.shot
time = idx
name = 'g{:06d}.{:05d}'.format(shot, time)
#make tree for this shot
tools.makeDir(self.shotPath,
clobberFlag=False,
mode=self.chmod,
UID=self.UID,
GID=self.GID)
#make tree for this timestep
if self.shotPath[-1] != '/': self.shotPath += '/'
timeDir = self.shotPath + '{:06d}/'.format(time)
newgfile = timeDir + name
#clobber and make time directory
tools.makeDir(timeDir,
clobberFlag=clobberflag,
mode=self.chmod,
UID=self.UID,
GID=self.GID)
shutil.copyfile(oldgfile, newgfile)
return time
def makePlotlyEQDiv(shot,
time,
MachFlag,
ep,
height=None,
gfile=None,
logFile=False):
"""
returns a DASH object for use directly in dash app
"""
if logFile is True:
log = logging.getLogger(__name__)
#Use Equilparamsclass to create EQ object
if ep is None:
print('Note: no EP object')
if gfile is None:
print("Error generating EQ plot: no EQ object or gfile")
sys.exit()
else:
ep = EP.equilParams(gfile)
r = ep.g['R']
z = ep.g['Z']
psi = ep.g['psiRZn']
R, Z = np.meshgrid(r, z)
rbdry = ep.g['lcfs'][:, 0]
zbdry = ep.g['lcfs'][:, 1]
if MachFlag == 'nstx':
rlim, zlim = nstxu_wall(oldwall=False) #FOR NSTXU
else:
rlim = ep.g['wall'][:, 0]
zlim = ep.g['wall'][:, 1]
# for aspect ratio
# if height==None:
# height=1000
# aspect = (z.max()-z.min()) / (r.max()-r.min())
# width = (1.0/aspect)*height
levels = sorted(
np.append([0.0, 0.05, 0.1, 0.25, 0.5, 0.75, 0.95, 1.0],
np.linspace(0.99, psi.max(), 15)))
# aspect = (z.max()-z.min()) / (r.max()-r.min())
# width = (1.0/aspect)*height
import plotly
import plotly.graph_objects as go
import plotly.express as px
#psi data
fig = go.Figure(data=go.Contour(
z=psi,
x=r, # horizontal axis
y=z, # vertical axis
colorscale='cividis',
contours_coloring='heatmap',
name='psiN',
showscale=False,
ncontours=20))
#Wall in green
fig.add_trace(
go.Scatter(x=rlim,
y=zlim,
mode="markers+lines",
name="Wall",
line=dict(color="#19fa1d")))
#white lines around separatrix
levelsAtLCFS = np.linspace(0.95, 1.05, 15)
CS = plt.contourf(R, Z, psi, levelsAtLCFS, cmap=plt.cm.cividis)
for i in range(len(levelsAtLCFS)):
levelsCS = plt.contour(R, Z, psi, levels=[levelsAtLCFS[i]])
for j in range(len(levelsCS.allsegs[0])):
r = levelsCS.allsegs[0][j][:, 0]
z = levelsCS.allsegs[0][j][:, 1]
fig.add_trace(
go.Scatter(x=r,
y=z,
mode="lines",
line=dict(
color="white",
width=1,
dash='dot',
)))
#Seperatrix in red. Sometimes this fails if psi is negative
#so we try and except.
#if try fails, just plot using rbdry,zbdry from gfile
try:
CS = plt.contourf(R, Z, psi, levels, cmap=plt.cm.cividis)
lcfsCS = plt.contour(CS, levels=[1.0])
for i in range(len(lcfsCS.allsegs[0])):
rlcfs = lcfsCS.allsegs[0][i][:, 0]
zlcfs = lcfsCS.allsegs[0][i][:, 1]
fig.add_trace(
go.Scatter(x=rlcfs,
y=zlcfs,
mode="lines",
name="LCFS",
line=dict(
color="red",
width=4,
)))
except:
print("Could not create contour plot. Psi levels must be increasing.")
print("Try flipping psi sign and replotting.")
print("plotting rbdry, zbdry from gfile (not contour)")
if logFile is True:
log.info(
"Could not create contour plot. Psi levels must be increasing."
)
log.info("Try flipping psi sign and replotting.")
log.info("plotting rbdry, zbdry from gfile (not contour)")
fig.add_trace(
go.Scatter(x=rbdry,
y=zbdry,
mode="lines",
name="LCFS",
line=dict(
color="red",
width=4,
)))
fig.update_layout(
title="{:06d}@{:05d}ms".format(shot, time),
xaxis_title="R [m]",
yaxis_title="Z [m]",
autosize=True,
#for aspect ratio
#autosize=False,
#width=width*1.1,
#height=height,
paper_bgcolor='rgba(0,0,0,0)',
plot_bgcolor='rgba(0,0,0,0)',
showlegend=False,
font=dict(
# family="Courier New",
size=18,
color="#dcdce3"))
return fig
def writePlotlyEQ(shot,
time,
outFile,
MachFlag,
ep=None,
gfile=None,
logFile=False):
"""
saves a plotly webpage to a file that can be imported to html via iframe
"""
if logFile is True:
log = logging.getLogger(__name__)
#Use Equilparamsclass to create EQ object
if ep is None:
if gfile is None:
print("Error generating EQ plot: no EQ object or gfile")
sys.exit()
else:
ep = EP.equilParams(gfile)
r = ep.g['R']
z = ep.g['Z']
psi = ep.g['psiRZn']
R, Z = np.meshgrid(r, z)
rbdry = ep.g['lcfs'][:, 0]
zbdry = ep.g['lcfs'][:, 1]
if MachFlag == 'nstx':
rlim, zlim = nstxu_wall(oldwall=False) #FOR NSTXU
else:
rlim = ep.g['wall'][:, 0]
zlim = ep.g['wall'][:, 1]
levels = sorted(
np.append([0.0, 0.05, 0.1, 0.25, 0.5, 0.75, 1.0],
np.linspace(1.01, psi.max(), 20)))
import plotly
import plotly.graph_objects as go
import plotly.express as px
#psi data
fig = go.Figure(data=go.Contour(
z=psi,
x=r, # horizontal axis
y=z, # vertical axis
colorscale='cividis',
contours_coloring='heatmap',
name='psiN',
showscale=False,
))
#Wall in green
fig.add_trace(
go.Scatter(x=rlim,
y=zlim,
mode="markers+lines",
name="Wall",
line=dict(color="#19fa1d")))
#Seperatrix in red. Sometimes this fails if psi is negative
#so we try and except.
#if try fails, just plot using rbdry,zbdry from gfile
try:
CS = plt.contourf(R, Z, psi, levels, cmap=plt.cm.cividis)
lcfsCS = plt.contour(CS, levels=[1.0])
for i in range(len(lcfsCS.allsegs[0])):
rlcfs = lcfsCS.allsegs[0][i][:, 0]
zlcfs = lcfsCS.allsegs[0][i][:, 1]
fig.add_trace(
go.Scatter(x=rlcfs,
y=zlcfs,
mode="lines",
name="LCFS",
line=dict(
color="red",
width=4,
)))
except:
print("Could not create contour plot. Psi levels must be increasing.")
print("Try flipping psi sign and replotting.")
print("plotting rbdry, zbdry from gfile (not contour)")
if logFile is True:
log.info(
"Could not create contour plot. Psi levels must be increasing."
)
log.info("Try flipping psi sign and replotting.")
log.info("plotting rbdry, zbdry from gfile (not contour)")
fig.add_trace(
go.Scatter(x=rbdry,
y=zbdry,
mode="lines",
name="LCFS",
line=dict(
color="red",
width=4,
)))
fig.update_layout(title="{:06d}@{:05d}ms".format(shot, time),
xaxis_title="R [m]",
yaxis_title="Z [m]",
paper_bgcolor='rgba(0,0,0,0)',
plot_bgcolor='rgba(0,0,0,0)',
showlegend=False,
font=dict(family="Anurati", size=18, color="#dcdce3"))
print("writing EQ output to: " + outFile)
fig.write_html(outFile)
def __init__(self,
gfile,
fluxLimit,
Terms,
dic=None,
path='.',
Time=None,
tag=None):
# --- input variables ---
idx = gfile[::-1].find(
'/') # returns location of last '/' in gfile or -1 if not found
if (idx == -1):
gpath = '.'
else:
idx *= -1
gpath = gfile[0:idx - 1] # path without a final '/'
gfile = gfile[idx::]
idx = gfile.find('.')
fmtstr = '0' + str(idx - 1) + 'd'
shot, time = int(gfile[1:idx]), int(gfile[idx + 1::])
gfile = 'g' + format(shot, fmtstr) + '.' + format(time, '05d')
fluxLimit = float(fluxLimit)
Terms = format(int(Terms), '06d')
if (Terms == format(int(000000), '06d')): fitFLAG = 'akima'
else: fitFLAG = 'vmec'
# --- set member variables ---
self.gpath = gpath # path to g-file and a-file
self.shot = shot # shot number
self.TIME = time # time of g-file
self.gfile = gfile # g-file name
self.fluxLimit = fluxLimit # percentage of normalized poloidal flux to mark VMEC boundary
self.Terms = Terms # number of terms in polynomial fit for AC,AI,AM; 000000 switches to akima splines
self.fitFLAG = fitFLAG # flag for poly fit or akima splines
self.path = path # path to write VMEC input file
self.filename = 'input.' + fitFLAG + '.' + gfile[
1::] # raw name (no tags) for VMEC input file
self.dic = {} # dictionary of input file data
if not (dic == None): self.dic = dic
else:
try:
import g2vmi_defaults
reload(g2vmi_defaults)
self.dic = g2vmi_defaults.set_defaults(self)
except:
self.dic = self.set_defaults()
if not (Time == None): self.time = Time
else:
self.time = time # time to get I- & C-coil currents from, if different from gfile time
if not (tag == None): self.tag = tag
else: self.tag = None # tag to append on input file name
# --- load g-file and get toroidal flux ---
print 'Loading g-file: ', gpath + '/' + gfile
self.ep = EPC.equilParams(gpath + '/' + gfile)
self.psitordic = self.ep.getTorPsi()
# --- read a-file & locate currents ---
afile = 'a' + format(shot, fmtstr) + '.' + format(time, '05d')
self.afile = {'name': afile}
if os.path.isfile(gpath + '/' + afile):
print 'Loading a-file: ', gpath + '/' + afile
with open(gpath + '/' + afile, 'r') as f:
a_input = []
for line in f: # read to end of file
a_input.append(line)
self.afile['data'] = a_input
# locate F-coil & E-coil currents
for i, line in enumerate(a_input):
line = line.strip().split()
if (len(line) == 4):
line = [float(line[k]) for k in xrange(4)]
if (line[0] == int(line[0])) & (line[1] == int(
line[1])) & (line[2] == int(
line[2])) & (line[3] == int(line[3])):
F_coil_idx = i + int(ceil(
(line[0] + line[1]) / 4.0)) + 1
E_coil_idx = i + int(
ceil((line[0] + line[1] + line[2]) / 4.0)) + 1
break
self.afile['F_coil_idx'] = F_coil_idx
self.afile['E_coil_idx'] = E_coil_idx
elif (not self.dic['mgrid']) | (self.dic['mgrid'] == 'None') | (
self.dic['mgrid'] == 'none'):
print 'no a-file found, but no mgrid file either, so set to fixed boundary mode'
self.dic['mgrid'] = None
self.dic['LFREEB'] = 'F'
else:
raise RuntimeError('a-file not found')