def writeToTree(inp):
shot, chan, surf1, surf2 = inp
timeout = time.time()
outstr = str(chan)
if chan < 10:
outstr = '0' + outstr
try:
Tree = MDSplus.Tree('spectroscopy', shot)
b = plasma.Tokamak(eqtools.CModEFITTree(shot)) # I HATE THIS
output = effectiveHeight(
surf1, surf2, b,
scipy.mgrid[b.eq.getTimeBase()[0]:b.eq.getTimeBase()[-1]:1e-3])
dummy = MDSplus.Data.compile(
'build_signal($1,*,$2)', output,
scipy.mgrid[b.eq.getTimeBase()[0]:b.eq.getTimeBase()[-1]:1e-3])
Tree.getNode('.BOLOMETER.RESULTS.DIODE.BPLY.AREA:CHORD_' +
outstr).putData(dummy)
print('channel ' + str(chan) + ':\t ' + str(time.time() - timeout))
except IndexError:
print('ERROR IN THIS CHANNEL: ' + str(chan))
str(TePtEdge[i]))
else:
print("Could not pull edge Thomson data for this shot, " + str(s))
#Plot profile
#plt.grid(color='k', linestyle='-', linewidth=0.5)
#plt.plot(rPt,nePt/1E20,linestyle='',marker='o')
#plt.hold('True')
#plt.plot(rPtEdge,nePtEdge/1E20,linestyle='',marker='o')
#plt.xlabel('R [m]')
#plt.ylabel('$n_e$ [$\\times10^{20} $m$^{-3}$]')
#plt.show()
#Create EFIT object - pull EFIT data from tree
print("Attempting to generate EQDSK file from EFIT tree")
myEq = eqtools.CModEFITTree(s)
#Cache useful EFIT parameters
tEfit = myEq.getTimeBase()
eqInfo = myEq.getInfo()
RGrid = myEq.getRGrid()
zGrid = myEq.getZGrid()
R0 = myEq.getMagR() #Major radius of magnetic axis
z0 = myEq.getMagZ() #Height of magnetic axis
B0 = myEq.getBtVac() #Vacuum toroidal field on axis
psi0 = myEq.getFluxAxis()
psiBndry = myEq.getFluxLCFS()
psiRz = -1.0 * myEq.getCurrentSign() * myEq.getFluxGrid()
Bcentr = myEq.getBCentr()
Ip = myEq.getIpCalc()
def hirexsr_pos(shot,
hirex_branch,
tht,
primary_line,
primary_impurity,
plot_pos=False,
plot_on_tokamak=False,
check_with_tree=False,
t0=1.25):
'''
Get the POS vector as defined in the THACO manual.
Unlike in THACO, here we use POS vectors averaged over the wavelength range of the line of interest, rather than over the wavelength range that is fit (including various satellite lines). This reduces the averaging quite significantly.
Plotting functions make use of the eqtools and TRIPPy packages.
'''
specTree = MDSplus.Tree('spectroscopy', shot)
if hirex_branch == 'B':
# pos vectors for detector modules 1-3
pos1 = specTree.getNode(
r'\SPECTROSCOPY::TOP.HIREXSR.CALIB.MOD1:POS').data()
pos2 = specTree.getNode(
r'\SPECTROSCOPY::TOP.HIREXSR.CALIB.MOD2:POS').data()
pos3 = specTree.getNode(
r'\SPECTROSCOPY::TOP.HIREXSR.CALIB.MOD3:POS').data()
# wavelengths for each module
lam1 = specTree.getNode(
r'\SPECTROSCOPY::TOP.HIREXSR.CALIB.MOD1:LAMBDA').data()
lam2 = specTree.getNode(
r'\SPECTROSCOPY::TOP.HIREXSR.CALIB.MOD2:LAMBDA').data()
lam3 = specTree.getNode(
r'\SPECTROSCOPY::TOP.HIREXSR.CALIB.MOD3:LAMBDA').data()
pos_tot = np.hstack([pos1, pos2, pos3])
lam_tot = np.hstack([lam1, lam2, lam3])
else:
# 1 detector module
pos_tot = specTree.getNode(
r'\SPECTROSCOPY::TOP.HIREXSR.CALIB.MOD4:POS').data()
# wavelength
lam_tot = specTree.getNode(
r'\SPECTROSCOPY::TOP.HIREXSR.CALIB.MOD4:LAMBDA').data()
branchNode = specTree.getNode(
r'\SPECTROSCOPY::TOP.HIREXSR.ANALYSIS{:s}.{:s}LIKE'.format(
str(tht) if tht != 0 else '',
'HE' if hirex_branch == 'B' else 'H'))
# mapping from pixels to chords (wavelength x space pixels, but wavelength axis is just padding)
chmap = branchNode.getNode('BINNING:CHMAP').data()
pixels_to_chords = chmap[0, :]
# find over which wavelengths the pos vector should be averaged at every time
# get lambda bounds for specific BSFC line for accurate impurity forward modeling:
lam_bounds, primary_line = get_hirexsr_lam_bounds(primary_impurity,
primary_line,
reduced=True)
lam_all = branchNode.getNode('SPEC:LAM').data()
# exclude empty chords
mask = lam_all[0, 0, :] != -1
lam_masked = lam_all[:, :, mask]
# lambda vector does not change over time, so just use tbin=0
tbin = 0
w0 = []
w1 = []
for chbin in np.arange(lam_masked.shape[2]):
bb = np.searchsorted(lam_masked[:, tbin, chbin], lam_bounds)
w0.append(bb[0])
w1.append(bb[1])
# form chords
pos_ave = []
for chord in np.arange(lam_masked.shape[2]):
pos_ave.append(
np.mean(pos_tot[w0[chord]:w1[chord], pixels_to_chords == chord, :],
axis=(0, 1)))
pos_ave = np.array(pos_ave)
if plot_pos:
# show each component of the pos vector separately
fig, ax = plt.subplots(2, 2)
axx = ax.flatten()
for i in [0, 1, 2, 3]:
pcm = axx[i].pcolormesh(pos_tot[:, :, i].T)
axx[i].axis('equal')
fig.colorbar(pcm, ax=axx[i])
if plot_on_tokamak:
import TRIPPy
import eqtools
# visualize chords
efit_tree = eqtools.CModEFITTree(shot)
tokamak = TRIPPy.plasma.Tokamak(efit_tree)
#pos_ave[:,0]*=1.2
# pos[:,3] indicate spacing between rays
rays = [TRIPPy.beam.pos2Ray(p, tokamak) for p in pos_ave] #pos_old]
weights = TRIPPy.invert.fluxFourierSens(rays,
efit_tree.rz2psinorm,
tokamak.center,
t0,
np.linspace(0, 1, 150),
ds=1e-5)[0]
from TRIPPy.plot.pyplot import plotTokamak, plotLine
f = plt.figure()
a = f.add_subplot(1, 1, 1)
# Only plot the tokamak if an axis was not provided:
plotTokamak(tokamak)
for r in rays:
plotLine(r, pargs='r', lw=1.0)
i_flux = np.searchsorted(efit_tree.getTimeBase(), t0)
# mask out coils, where flux is highest
flux = efit_tree.getFluxGrid()[i_flux, :, :]
#flux[flux>np.percentile(flux, 75)] = np.nan
#flux[:,efit_tree.getRGrid()>0.9] = np.nan
cset = a.contour(efit_tree.getRGrid(), efit_tree.getZGrid(), flux, 80)
#f.colorbar(cset,ax=a)
if check_with_tree:
try:
pos_on_tree = branchNode.getNode('MOMENTS.{:s}:POS'.format(
primary_line.upper())).data()
except:
pos_on_tree = branchNode.getNode('MOMENTS.LYA1:POS').data()
return pos_ave, pos_on_tree
else:
return pos_ave
matplotlib.rc('font', size=20)
import scipy
import eqtools
import profiletools
import gptools
import matplotlib.pyplot as plt
plt.ion()
plt.close('all')
shot = 1101014006
t_min = 0.965
t_max = 1.365
e = eqtools.CModEFITTree(shot)
p_CTS = profiletools.TeCTS(shot, t_min=t_min, t_max=t_max, efit_tree=e)
p_CTS.time_average(weighted=True)
p_ETS = profiletools.TeETS(shot, t_min=t_min, t_max=t_max, efit_tree=e)
p_ETS.time_average(weighted=True)
p_GPC = profiletools.TeGPC(shot, t_min=t_min, t_max=t_max, efit_tree=e)
p_GPC.time_average()
p_GPC2 = profiletools.TeGPC2(shot, t_min=t_min, t_max=t_max, efit_tree=e)
p_GPC2.time_average()
f, sl = e.plotFlux(fill=False)
sl.set_val(47)
f.axes[0].plot(p_CTS.X[:, 0], p_CTS.X[:, 1], 'gs', markersize=12, label='CTS')
f.axes[0].plot(p_ETS.X[:, 0], p_ETS.X[:, 1], 'rs', markersize=12, label='ETS')
f.axes[0].plot(p_GPC2.X[:, 0],
scipy.zeros_like(p_GPC2.y),
import scipy
import scipy.stats
import scipy.io
import numpy.random
import matplotlib
matplotlib.rcParams.update({'font.size': 22})
import matplotlib.pyplot as plt
plt.ion()
#plt.close('all')
shot = 1101014006
# Start and end times of flat top:
flat_start = 0.965 #0.5
flat_stop = 1.365 #1.5
efit_tree = eqtools.CModEFITTree(shot)
t_EFIT = efit_tree.getTimeBase()
electrons = MDSplus.Tree('electrons', shot)
# Get core TS data:
N_Te_TS = electrons.getNode(r'\electrons::top.yag_new.results.profiles:te_rz')
t_Te_TS = N_Te_TS.dim_of().data()
# Only keep points that are in the RF flat top:
ok_idxs = (t_Te_TS >= flat_start) & (t_Te_TS <= flat_stop)
t_Te_TS = t_Te_TS[ok_idxs]
Te_TS = N_Te_TS.data()[:, ok_idxs]
dev_Te_TS = electrons.getNode(