# By default, have the iter indices range over full file range
index_first, index_last = 0, nfiles - 1
for arg in args:
if arg in ['-range', '-centerrange', '-leftrange', '-rightrange', '-n',\
'-f', '-all', '-iter']:
index_first, index_last = get_desired_range(int_file_list, args)
savename = dirname_stripped + '_partial_wreath_trace_' +\
file_list[index_first] + '_' + file_list[index_last] + '.pkl'
savefile = datadir + savename
print('Your data will be saved in the file %s' %savename)
# Read in first Shell_Slices/AZ_Avgs file
a0 = Shell_Slices(radatadir + file_list[index_first], '')
az0 = AZ_Avgs(dirname + '/AZ_Avgs/' + file_list[index_first], '')
# Read in grid info from AZ_Avgs slice
rr = az0.radius
ri, ro = np.min(rr), np.max(rr)
sint, cost = az0.sintheta, az0.costheta
tt = np.arccos(cost)
tt_lat = (np.pi/2 - tt)*180./np.pi
nt, nr = len(sint), len(rr)
nphi = 2*nt
lons = np.arange(0., 360., 360/nphi)
# Desired latitude range
lat1 = -25.
lat2 = 25.
ith1 = np.argmin(np.abs(tt_lat - lat1))
#
# In the example that follows, we demonstrate how to plot azimuthal averages, including how to generate streamlines of mass flux. Note that since the benchmark is Boussinesq, our velocity and mass flux fields are identical. This is not the case when running an anelastic simulation.
#
# We begin with the usual preamble and also import two helper routines used for displaying azimuthal averages.
#
# Examining the data structure, we see that the vals array is dimensioned to account for latitudinal variation, and that we have new attributes costheta and sintheta used for referencing locations in the theta direction.
# In[13]:
from rayleigh_diagnostics import AZ_Avgs, build_file_list, plot_azav, streamfunction
import matplotlib.pyplot as plt
import pylab
import numpy
#from azavg_util import *
files = build_file_list(30000, 40000, path='AZ_Avgs')
az = AZ_Avgs(files[0], path='')
# ***
# Before creating our plots, let's time-average over the last two files that were output (thus sampling the equilibrated phase).
# In[14]:
nfiles = len(files)
tcount = 0
for i in range(nfiles):
az = AZ_Avgs(files[i], path='')
if (i == 0):
nr = az.nr
ntheta = az.ntheta
nq = az.nq
T = eq.temperature
dlnT = eq.dlnT
dsdr = eq.dsdr
rhoT = rho * T
grav = eq.gravity
rho_2d = rho.reshape((1, nr))
T_2d = T.reshape((1, nr))
dsdr_2d = dsdr.reshape((1, nr))
rhoT_2d = rhoT.reshape((1, nr))
dlnT_2d = dlnT.reshape((1, nr))
grav_2d = grav.reshape((1, nr))
# loop over data and make plots
for i in range(index_first, index_last + 1, nskip):
a = AZ_Avgs(radatadir + file_list[i], '')
nrec_tot = a.niter
nstep = nrec_tot // nrec
if nstep == 0:
nstep = 1
for j in range(0, nrec_tot, nstep):
vals = a.vals[:, :, :, j]
lut = a.lut
iter_loc = a.iters[j]
t_loc = a.time[j]
#========================
# KINETIC ENERGY EQUATION
#========================
# Create the save directory if it doesn't already exist
plotdir = dirname + '/plots/moll_azav_tl_movies/' + varname + '/' +\
('rval%0.3f' %rval) + '/'
if not os.path.isdir(plotdir):
os.makedirs(plotdir)
# Name of logfile to catalog progress, if desired
logfile = dirname + '/plots/zz_logfile_mollazavtlmovie_' + varname +\
('_rval%0.3f' %rval) + '.txt'
if log_progress:
append_logfile(logfile, '==================================\n')
append_logfile(logfile, '========== Begin Log File ==========\n')
# Get grid information from last AZ_Avgs file
az0 = AZ_Avgs(azdatadir + file_list[index_last], '')
rr = az0.radius
ri, ro = np.min(rr), np.max(rr)
d = ro - ri
rr_depth = (ro - rr) / d
rr_height = (rr - ri) / d
sint = az0.sintheta
cost = az0.costheta
tt = np.arccos(cost)
tt_lat = (np.pi / 2 - tt) * 180 / np.pi
nr = az0.nr
nt = az0.ntheta
# compute some derivative quantities for the grid
tt_2d, rr_2d = np.meshgrid(tt, rr, indexing='ij')
sint_2d = np.sin(tt_2d)
else:
time_unit = compute_tdt(dirname)
time_label = r'$\rm{TDT}$'
if plotdir is None:
plotdir = dirname + '/plots/'
if not os.path.isdir(plotdir):
os.makedirs(plotdir)
iter_val = int_file_list[iiter]
fname = file_list[iiter]
# Check if there are corresponding Shell_Avgs or AZ_Avgs files if needed
if use_az and varname[-5:] == 'prime':
try:
az = AZ_Avgs(dirname + '/AZ_Avgs/' + fname, '')
print("read AZ_Avgs/" + fname)
except:
print("No file AZ_Avgs/" + fname)
print("setting az = None")
az = None
else:
az = None
if use_sh and 'prime_sph' in varname:
try:
sh = Shell_Avgs(dirname + '/Shell_Avgs/' + fname, '')
print("read Shell_Avgs/" + fname)
except:
print("No file Shell_Avgs/" + fname)
print("setting sh = None")
# just plot the last file
iter1, iter2 = int_file_list[index_first], int_file_list[index_last]
for i in range(nargs):
arg = args[i]
if arg == '-minmax':
mins = float(args[i + 1]), float(args[i + 3]), float(args[i + 5])
maxes = float(args[i + 2]), float(args[i + 4]), float(args[i + 6])
elif arg == '-rbcz':
rbcz = float(args[i + 1])
elif arg == '-start':
count = int(args[i + 1])
# Get grid information from first AZ_Avgs file
az0 = AZ_Avgs(radatadir + file_list[index_first], '')
rr = az0.radius
ri, ro = np.min(rr), np.max(rr)
d = ro - ri
rr_depth = (ro - rr) / d
rr_height = (rr - ri) / d
sint = az0.sintheta
cost = az0.costheta
tt = np.arccos(cost)
tt_lat = (np.pi / 2 - tt) * 180 / np.pi
nr = az0.nr
nt = az0.ntheta
# compute some derivative quantities for the grid
tt_2d, rr_2d = np.meshgrid(tt, rr, indexing='ij')
sint_2d = np.sin(tt_2d)