def trikap(run):
# Only execute if plotting
if run.no_plot:
return
print("Plotting scan of triangularity vs elongation...")
Nfile = len(run.fnames)
# Init arrays of data used in scan.
full_lingrowth = [dict() for ifile in range(Nfile)]
# Get lingrowth data from .dat file.
for ifile in range(Nfile):
datfile_name = run.work_dir + run.dirs[ifile] + run.out_dir + run.files[ifile] + '.lingrowth.dat'
with open(datfile_name,'rb') as datfile:
full_lingrowth[ifile] = pickle.load(datfile)
nakx = len(full_lingrowth[0]['kx'])
for ifile in range(Nfile):
if len(full_lingrowth[ifile]['kx']) != nakx:
quit("Error - number of kx values should be constant across all items in the scan - Exiting...")
tri,kap = np.zeros(Nfile),np.zeros(Nfile)
for ifile in range(Nfile):
tri[ifile] = full_lingrowth[ifile]['tri']
kap[ifile] = full_lingrowth[ifile]['kap']
tris = sorted(list(set(tri)))
kaps = sorted(list(set(kap)))
print("Triangularity values: " + str(tris))
print("Elongation values: " + str(kaps))
if len(tris) * len(kaps) != Nfile:
quit("Incorrect number of files added to populate the scan - exiting")
gammas = np.zeros((len(tris), len(kaps), nakx))
kymax = np.zeros((len(tris), len(kaps), nakx))
for ifile in range(Nfile):
gamma = full_lingrowth[ifile]['gamma']
ky = full_lingrowth[ifile]['ky']
kx = full_lingrowth[ifile]['kx']
# Limits search to ITG.
ikymax = int((np.abs(ky-1.0)).argmin())
for itri in range(len(tris)):
for ikap in range(len(kaps)):
if tri[ifile] == tris[itri] and kap[ifile] == kaps[ikap]:
for ikx in range(len(kx)):
gammas[itri,ikap,ikx] = np.nanmax(gamma[:ikymax,ikx])
kymax[itri,ikap,ikx] = ky[np.nanargmax(gamma[:ikymax,ikx])]
pdflist = []
tmp_pdf_id=0
for ikx in range(nakx):
gplot.plot_2d(gammas,tris,kaps,np.min(gammas[:,:,ikx]),np.max(gammas[:,:,ikx]),cmp='Reds',xlab='$\delta$',ylab='$\kappa$',title='$\gamma a/v_t: k_x$ = ' + str(full_lingrowth[0]['kx'][ikx]))
tmp_pdfname = 'tmp'+str(tmp_pdf_id)
gplot.save_plot(tmp_pdfname, run)
pdflist.append(tmp_pdfname)
tmp_pdf_id += 1
merged_pdfname = 'tri_kap_scan'
gplot.merge_pdfs(pdflist, merged_pdfname, run)
def plot_zonal_vs_kxt(kx, time, field, title): ### Evolution of flow in fourier space versus time and kx
z = (np.abs(field)**2)/(np.max(np.abs(field)**2))
xlab = '$k_{x}\\rho_{i}$'
ylab = '$t (v_{t}/a)$'
cmap = 'YlGnBu'
fig = gplot.plot_2d(z,kx,time,z.min(),z.max(),xlab,ylab,title,cmap)
return fig
def plot_zonal_vs_xt(xgrid, time, field, title): ### Evolution of flow in real space versus time and x
# ensure there is no divide by zero
# denom = np.max(np.abs(field))
# test = denom==0
# denom[test] = 1
z = field
z_min, z_max = -np.abs(z).max(), np.abs(z).max()
xlab = '$x/\\rho_{i}$'
ylab = '$t (v_{t}/a)$'
cmap = 'RdBu'
fig = gplot.plot_2d(z,xgrid,time,z_min,z_max,xlab,ylab,title,cmap)
return fig
def plot_flux_vs_vpth(mygrids, flx, title):
from gs2_plotting import plot_2d
xlab = '$\\theta$'
ylab = '$v_{\parallel}$'
cmap = 'RdBu'
for idx in range(flx.shape[0]):
z = flx[idx, :, :]
z_min, z_max = z.min(), z.max()
fig = plot_2d(z, mygrids.theta, mygrids.vpa, z_min, z_max, xlab, ylab,
title + ' (is= ' + str(idx + 1) + ')', cmap)
return fig
def plot_phi2_vs_kxky(kx, ky, phi2, has_flowshear):
from gs2_plotting import plot_2d
title = '$\\langle\\vert\\hat{\\varphi}_k\\vert ^2\\rangle_{t,\\theta}$'
title += ' $\\forall$ $k_y > 0$'
if has_flowshear:
xlab = '$\\bar{k}_{x}\\rho_i$'
else:
xlab = '$k_{x}\\rho_i$'
ylab = '$k_{y}\\rho_i$'
cmap = 'RdBu' # 'RdBu_r','Blues'
z = phi2[1:, :] # taking out zonal modes because they are much larger
z_min, z_max = z.min(), z.max()
use_logcolor = True
fig = plot_2d(np.transpose(z), kx, ky[1:], z_min, z_max, xlab, ylab, title,
cmap, use_logcolor)
return fig
def plot_flux_vs_kxky(ispec, spec_names, kx, ky, flx, title, has_flowshear):
from gs2_plotting import plot_2d
if has_flowshear:
xlab = '$\\bar{k}_{x}\\rho_i$'
else:
xlab = '$k_{x}\\rho_i$'
ylab = '$k_{y}\\rho_i$'
cmap = 'Blues' # 'Reds','Blues'
z = flx[ispec, :, :] # OB 140918 ~ Don't take absolute value of fluxes.
z_min, z_max = 0.0, z.max()
title = 'Contributions to ' + title
if ispec > 1:
title += ' (impurity ' + str(ispec - 1) + ')'
else:
title += ' (' + spec_names[ispec] + 's)'
fig = plot_2d(np.transpose(z), kx, ky, z_min, z_max, xlab, ylab, title,
cmap)
return fig
def my_task_single(ifile, run, myin, myout):
# User parameters
dump_at_start = 0.3 # fraction of initial time to dump when fitting
ikx_list = [-1] # choose which kx to plot, negative means plot all
# Compute and save growthrate
if not run.only_plot:
grid_option = myin['kt_grids_knobs']['grid_option']
t = myout['t']
nt = t.size
it_start = round(nt*dump_at_start)
kx = myout['kx']
nakx = kx.size
ky = myout['ky']
naky = ky.size
phi2 = myout['phi2_by_mode'] # modulus squared, avged over theta (indices: [t,ky,kx])
it_end = nt
for it in range(nt):
for ikx in range(nakx):
for iky in range(naky):
if it < it_end and not np.isfinite(phi2[it,iky,ikx]):
it_end = it - 10
break
if grid_option=='range':
# In range, plot the only kx
ikx_list = [0]
elif grid_option=='box':
if ikx_list[0]==-1:
ikx_list = [i in range(nakx)]
# Fit phi to get growthrates
gamma = np.zeros([naky,len(ikx_list)])
gamma[0,:]=float('nan') # skip zonal mode
for iky in range(1,naky):
for ikx in ikx_list:
gamma[iky,ikx] = get_growthrate(t,phi2,it_start,it_end,ikx,iky)
mydict = {'ikx_list':ikx_list,'kx':kx,'ky':ky,'gamma':gamma,'tri':myin['theta_grid_parameters']['tri'],
'kap':myin['theta_grid_parameters']['akappa']}
# Save to .dat file
# OB 170918 ~ Changed output to a dict and added kappa, tri
datfile_name = run.work_dir + run.dirs[ifile] + run.out_dir + run.files[ifile] + '.lingrowth.dat'
with open(datfile_name, 'wb') as datfile:
pickle.dump(mydict,datfile)
# or read from .dat file
else:
datfile_name = run.work_dir + run.dirs[ifile] + run.out_dir + run.files[ifile] + '.lingrowth.dat'
with open(datfile_name, 'rb') as datfile:
mydict = pickle.load(datfile)
# Plotting
if not run.no_plot:
# If we ran for many kx ky, then plot colormap
if len(ikx_list) > 3 and ky.size > 3:
title = '$\\gamma \\ [v_{thr}/r_r]$'
xlab = '$k_{x}\\rho_i$'
ylab = '$k_{y}\\rho_i$'
cmap = 'RdBu'
z = gamma[:,:]
z_min, z_max = z.min(), z.max()
fig = gplot.plot_2d(z,kx,ky,z_min,z_max,xlab,ylab,title,cmap)
# Otherwise plot curves vs ky for each kx
else:
plt.figure(figsize=(12,8))
plt.xlabel('$k_y\\rho_i$')
plt.ylabel('$\\gamma \\ [v_{thr}/r_r]$')
plt.title('Linear growthrate')
plt.grid(True)
my_legend = []
for ikx in ikx_list:
plt.plot(ky,gamma[:,ikx])
my_legend.append('$\\rho_i k_x='+str(kx[ikx])+'$')
plt.legend(my_legend)
pdfname = 'lingrowth'
gplot.save_plot(pdfname, run, ifile)
print('Maximum linear growthrate: '+str(np.nanmax(gamma)))
def trikap(run):
# Only execute if plotting
if run.no_plot:
return
print("Plotting scan of triangularity vs elongation...")
Nfile = len(run.fnames)
# Init arrays of data used in scan.
full_fluxes = [dict() for ifile in range(Nfile)]
full_time = [dict() for ifile in range(Nfile)]
# Initialize fluxes and grids from .dat files.
for ifile in range(Nfile):
datfile_name = run.work_dir + run.dirs[
ifile] + run.out_dir + run.files[ifile] + '.fluxes.dat'
with open(datfile_name, 'rb') as datfile:
full_fluxes[ifile] = pickle.load(datfile)
datfile_name = run.work_dir + run.dirs[
ifile] + run.out_dir + run.files[ifile] + '.time.dat'
with open(datfile_name, 'rb') as datfile:
full_time[ifile] = pickle.load(datfile)
# Uses nspec from first file. Will quit if not constant between files.
nspec = full_fluxes[0]['nspec']
print("Number of species " + str(nspec))
tri, kap = np.zeros(Nfile), np.zeros(Nfile)
for ifile in range(Nfile):
tri[ifile] = full_fluxes[ifile]['tri']
kap[ifile] = full_fluxes[ifile]['kap']
if full_fluxes[ifile]['nspec'] != nspec:
quit("Number of species varies between files - exiting")
tris = sorted(list(set(tri)))
kaps = sorted(list(set(kap)))
print("Triangularity values: " + str(tris))
print("Elongation values: " + str(kaps))
if len(tris) * len(kaps) != Nfile:
quit("Too few files added to populate the scan - exiting")
qflx = np.zeros((len(tris), len(kaps), nspec))
for itri in range(len(tris)):
for ikap in range(len(kaps)):
for ifile in range(Nfile):
if tri[ifile] == tris[itri] and kap[ifile] == kaps[ikap]:
for ispec in range(nspec):
qflx[itri, ikap, ispec] = full_time[ifile].timeavg(
full_fluxes[ifile]['qflx'][:, ispec])
spec_names = full_fluxes[0]['spec_names']
pdflist = []
tmp_pdf_id = 0
for ispec in range(nspec):
print("Plotting for species: " + spec_names[ispec])
gplot.plot_2d(qflx[:, :, ispec],
tris,
kaps,
np.min(qflx[:, :, ispec]),
np.max(qflx[:, :, ispec]),
cmp='Reds',
xlab='$\delta$',
ylab='$\kappa$',
title='$Q_{GS2}$: ' + spec_names[ispec])
tmp_pdfname = 'tmp' + str(tmp_pdf_id)
gplot.save_plot(tmp_pdfname, run)
pdflist.append(tmp_pdfname)
tmp_pdf_id += 1
merged_pdfname = 'tri_kap_scan'
gplot.merge_pdfs(pdflist, merged_pdfname, run)