plot_inte = True
elif arg == '-tote':
plot_tote = True
elif arg == '-name':
savename = args[i + 1] + '.png'
elif arg == '-tol':
tol = float(args[i + 1])
elif arg == '-equil':
plot_equil_time = True
elif arg == '-eqtime':
chosen_eqtime = float(args[i + 1])
# get density and grid info
eq = get_eq(dirname)
rho = eq.rho
gi = GridInfo(dirname + '/grid_info', '')
rr = gi.radius
rw = gi.rweights
nr = len(rr)
# Get the frame rate energy
Om0 = 2 * np.pi / compute_Prot(dirname)
# Get the radial integration weights and density
rw = gi.rweights
w0 = Om0**2.0 * np.sum(rho * rr**2.0 * rw) / 3.0
# Compute the equilibrated DRKE
# Read in the KE data (dictionary form)
if 'trace' in the_file:
def get_slice(a, varname, dirname=None, j=0):
# first test if varname is valid and if it's a basic variable
basic = is_basic(varname)
# first get the appropriate time slice
vals = a.vals[..., j]
lut = a.lut
rr = a.radius
if basic:
# gets a basic field associated with Rayleigh object a
if vals.ndim == 4 and not hasattr(a, 'lpower'):
# Shell_Slice or Meridional_Slice
cost = a.costheta
nt = len(cost)
cost = cost.reshape((1, nt, 1))
else:
# note...don't do cylindrical projections for Shell_Spectra!
# they don't multiply well
cost = 0.
# first get root variable name and store any modifiers
varname, deriv, primevar, sphvar = get_varprops(varname)
# get sine/cotangent from cosine
sint = np.sin(np.arccos(cost))
cott = cost / sint
# shape to make geometric fields
zero = np.zeros(np.array(np.shape(vals[..., 0])))
# return the basic field based on the variable name
if varname in var_indices:
# this is really only option for Shell_Spectra...
the_slice = vals[..., lut[var_indices[varname]]]
elif varname[-1] in ['l', 'z']: # cylindrical variable
the_slice_r = vals[..., lut[var_indices[varname[:-1] + 'r']]]
if deriv:
the_slice_t = vals[..., lut[var_indices[varname[:-1] + 'T']]]
else:
the_slice_t = vals[..., lut[var_indices[varname[:-1] + 't']]]
if varname[-1] == 'l':
the_slice = sint * the_slice_r + cost * the_slice_t
elif varname[-1] == 'z':
the_slice = cost * the_slice_r - sint * the_slice_t
elif is_an_int(varname):
the_slice = vals[..., lut[int(varname)]]
if primevar:
the_slice = prime(the_slice)
elif sphvar:
gi = GridInfo(dirname + '/grid_info')
tw = gi.tweights
the_slice = prime_sph(the_slice, tw)
del vals # free up memory
return the_slice
else:
if '+' in varname or '=' in varname:
signature = [1]
for char in varname:
if char == '+':
signature.append(1)
elif char == '=':
signature.append(-1)
the_slice = np.zeros_like(vals[..., 0])
count = 0
for subvar in varname.split('+'):
for subsubvar in subvar.split('='):
the_slice += get_slice(a, subsubvar, dirname,
j) * signature[count]
count += 1
elif '*' in varname or '/' in varname:
signature = [1]
for char in varname:
if char == '*':
signature.append(1)
elif char == '/':
signature.append(-1)
the_slice = np.ones_like(vals[..., 0])
count = 0
for subvar in varname.split('*'):
for subsubvar in subvar.split('/'):
the_slice *= (get_slice(a, subsubvar, dirname,
j))**signature[count]
count += 1
return the_slice
def get_sslice(a, varname, dirname=None, old=False, j=0):
# Given a shell_slice object (nphi, ntheta, nr, nq, nt),
# return the field (nphi, ntheta, nr) associated with [varname]
# (take the jth time slice (default j=0))
# for "primed" variables, the script will use radial/latitudinal
# integration weights (from GridInfo) to subtract off various
# averages
# for thermal variables, will need to provide dirname to get
# reference-state parameters (get_eq)
# _prime refers to az-avg subtracted
# _prime_sph refers to sph-avg subtracted
smooth_desired = False
if 'smooth' in varname:
smooth_desired = True
varname = varname.replace('smooth', '')
nphi_av = int(varname[:3])
varname = varname[3:]
# Get integration weights if needed
if '_sph' in varname:
gi = GridInfo(dirname + '/grid_info')
tw = gi.tweights
# Get sine and cosine if needed
if 'l' in varname or 'z' in varname:
sint = (a.sintheta).reshape((1, a.ntheta, 1))
cost = (a.costheta).reshape((1, a.ntheta, 1))
# get reference-state stuff if needed
if not dirname is None:
eq = get_eq(dirname)
ref_rho = (eq.density)[a.inds].reshape((1, 1, a.nr))
ref_T = (eq.temperature)[a.inds].reshape((1, 1, a.nr))
ref_P = (eq.pressure)[a.inds].reshape((1, 1, a.nr))
# Get raw values associated with jth time slice
vals = a.vals[:, :, :, :, j]
# Get quantity indexes (new or old)---do I actually ever need the
# old ones again?
qind_vr = var_indices['vr']
qind_vt = var_indices['vt']
qind_vp = var_indices['vp']
qind_s = var_indices['s']
qind_p = var_indices['p']
qind_omr = var_indices['omr']
qind_omt = var_indices['omt']
qind_omp = var_indices['omp']
qind_br = var_indices['br']
qind_bt = var_indices['bt']
qind_bp = var_indices['bp']
if old:
qind_vr = var_indices_old['vr']
qind_vt = var_indices_old['vt']
qind_vp = var_indices_old['vp']
qind_omr = var_indices_old['omr']
qind_omt = var_indices_old['omt']
qind_omp = var_indices_old['omp']
qind_s = var_indices_old['s']
qind_p = var_indices_old['p']
qind_omr = var_indices_old['br']
qind_omt = var_indices_old['bt']
qind_omp = var_indices_old['bp']
# Spherical velocities
if varname == 'vr':
ind_vr = a.lut[qind_vr]
sslice = vals[:, :, :, ind_vr] / 100. # measure velocities in m/s
elif varname == 'vt':
ind_vt = a.lut[qind_vt]
sslice = vals[:, :, :, ind_vt] / 100.
elif varname == 'vp':
ind_vp = a.lut[qind_vp]
sslice = vals[:, :, :, ind_vp] / 100.
elif varname == 'vr_prime':
ind_vr = a.lut[qind_vr]
vr_slice = vals[:, :, :, ind_vr] / 100.
sslice = prime(vr_slice)
elif varname == 'vt_prime':
ind_vt = a.lut[qind_vt]
vt_slice = vals[:, :, :, ind_vt] / 100.
sslice = prime(vt_slice)
elif varname == 'vp_prime':
ind_vp = a.lut[qind_vp]
vp_slice = vals[:, :, :, ind_vp] / 100.
sslice = prime(vp_slice)
# Cylindrical-coordinate velocities
elif varname == 'vl':
ind_vr, ind_vt = a.lut[qind_vr], a.lut[qind_vt]
vr_slice = vals[:, :, :, ind_vr] / 100.
vt_slice = vals[:, :, :, ind_vt] / 100.
sslice = vr_slice * sint + vt_slice * cost
elif varname == 'vz':
ind_vr, ind_vt = a.lut[qind_vr], a.lut[qind_vt]
vr_slice = vals[:, :, :, ind_vr] / 100.
vt_slice = vals[:, :, :, ind_vt] / 100.
sslice = vr_slice * cost - vt_slice * sint
elif varname == 'vl_prime':
sint = (a.sintheta).reshape((1, a.ntheta, 1))
cost = (a.costheta).reshape((1, a.ntheta, 1))
ind_vr, ind_vt = a.lut[qind_vr], a.lut[qind_vt]
vr_slice = vals[:, :, :, ind_vr] / 100.
vt_slice = vals[:, :, :, ind_vt] / 100.
vl_slice = vr_slice * sint + vt_slice * cost
sslice = prime(vl_slice)
elif varname == 'vz_prime':
ind_vr, ind_vt = a.lut[qind_vr], a.lut[qind_vt]
vr_slice = vals[:, :, :, ind_vr] / 100.
vt_slice = vals[:, :, :, ind_vt] / 100.
vz_slice = vr_slice * cost - vt_slice * sint
sslice = prime(vz_slice)
# Spherical vorticities
elif varname == 'omr':
ind_omr = a.lut[qind_omr]
sslice = vals[:, :, :, ind_omr]
elif varname == 'omt':
ind_omt = a.lut[qind_omt]
sslice = vals[:, :, :, ind_omt]
elif varname == 'omp':
ind_omp = a.lut[qind_omp]
sslice = vals[:, :, :, ind_omp]
elif varname == 'omr_prime':
ind_omr = a.lut[qind_omr]
omr_slice = vals[:, :, :, ind_omr]
sslice = prime(omr_slice)
elif varname == 'omt_prime':
ind_omt = a.lut[qind_omt]
omt_slice = vals[:, :, :, ind_omt]
sslice = prime(omt_slice)
elif varname == 'omp_prime':
ind_omp = a.lut[qind_omp]
omp_slice = vals[:, :, :, ind_omp]
sslice = prime(omp_slice)
# Cylindrical vorticities
elif varname == 'oml':
ind_omr, ind_omt = a.lut[qind_omr], a.lut[qind_omt]
omr_slice = vals[:, :, :, ind_omr]
omt_slice = vals[:, :, :, ind_omt]
sslice = omr_slice * sint + omt_slice * cost
elif varname == 'omz':
ind_omr, ind_omt = a.lut[qind_omr], a.lut[qind_omt]
omr_slice = vals[:, :, :, ind_omr]
omt_slice = vals[:, :, :, ind_omt]
sslice = omr_slice * cost - omt_slice * sint
elif varname == 'oml_prime':
sint = (a.sintheta).reshape((1, a.ntheta, 1))
cost = (a.costheta).reshape((1, a.ntheta, 1))
ind_omr, ind_omt = a.lut[qind_omr], a.lut[qind_omt]
omr_slice = vals[:, :, :, ind_omr]
omt_slice = vals[:, :, :, ind_omt]
oml_slice = omr_slice * sint + omt_slice * cost
sslice = prime(oml_slice)
elif varname == 'omz_prime':
ind_omr, ind_omt = a.lut[qind_omr], a.lut[qind_omt]
omr_slice = vals[:, :, :, ind_omr]
omt_slice = vals[:, :, :, ind_omt]
omz_slice = omr_slice * cost - omt_slice * sint
sslice = prime(omz_slice)
# Enstrophy (total and fluctuating)
elif varname == 'omsq':
ind_omr, ind_omt, ind_omp = a.lut[qind_omr], a.lut[qind_omt],\
a.lut[qind_omp]
sslice = vals[:, :, :, ind_omr]**2 + vals[:, :, :, ind_omt]**2 +\
vals[:, :, :, ind_omp]**2
elif varname == 'omsq_prime':
ind_omr, ind_omt, ind_omp = a.lut[qind_omr], a.lut[qind_omt],\
a.lut[qind_omp]
omr_slice = vals[:, :, :, ind_omr]
omt_slice = vals[:, :, :, ind_omt]
omp_slice = vals[:, :, :, ind_omp]
omr_prime_slice = prime(omr_slice)
omt_prime_slice = prime(omt_slice)
omp_prime_slice = prime(omp_slice)
sslice = omr_prime_slice**2 + omt_prime_slice**2 +\
omp_prime_slice**2
# Thermodynamic variables: deviations from reference state
elif varname == 's':
ind_s = a.lut[qind_s]
sslice = vals[:, :, :, ind_s]
elif varname == 'p':
ind_p = a.lut[qind_p]
sslice = vals[:, :, :, ind_p]
elif varname == 'rho':
ind_s, ind_p = a.lut[qind_s], a.lut[qind_p]
s_slice, p_slice = vals[:, :, :, ind_s], vals[:, :, :, ind_p]
sslice = ref_rho * (p_slice / ref_P / thermo_gamma - s_slice / c_P)
elif varname == 't':
ind_s, ind_p = a.lut[qind_s], a.lut[qind_p]
s_slice, p_slice = vals[:, :, :, ind_s], vals[:, :, :, ind_p]
sslice = ref_T*(p_slice/ref_P*(1. - 1./thermo_gamma) +\
s_slice/c_P)
# Thermodynamic variables: deviations from zonal mean
elif varname == 's_prime':
ind_s = a.lut[qind_s]
s_slice = vals[:, :, :, ind_s]
sslice = prime(s_slice)
elif varname == 'p_prime':
ind_p = a.lut[qind_p]
p_slice = vals[:, :, :, ind_p]
sslice = prime(p_slice)
elif varname == 'rho_prime':
ind_s, ind_p = a.lut[qind_s], a.lut[qind_p]
s_slice, p_slice = vals[:, :, :, ind_s], vals[:, :, :, ind_p]
rho_slice = ref_rho * (p_slice / ref_P / thermo_gamma - s_slice / c_P)
sslice = prime(rho_slice)
elif varname == 't_prime':
ind_s, ind_p = a.lut[qind_s], a.lut[qind_p]
s_slice, p_slice = vals[:, :, :, ind_s], vals[:, :, :, ind_p]
t_slice = ref_T*(p_slice/ref_P*(1. - 1./thermo_gamma) +\
s_slice/c_P)
sslice = prime(t_slice)
# Thermodynamic variables: deviations from spherical mean
elif varname == 's_prime_sph':
ind_s = a.lut[qind_s]
s_slice = vals[:, :, :, ind_s]
sslice = prime_sph(s_slice, tw)
elif varname == 'p_prime_sph':
ind_p = a.lut[qind_p]
p_slice = vals[:, :, :, ind_p]
sslice = prime_sph(p_slice, tw)
elif varname == 'rho_prime_sph':
ind_s, ind_p = a.lut[qind_s], a.lut[qind_p]
s_slice, p_slice = vals[:, :, :, ind_s], vals[:, :, :, ind_p]
rho_slice = ref_rho * (p_slice / ref_P / thermo_gamma - s_slice / c_P)
sslice = prime_sph(rho_slice, tw)
elif varname == 't_prime_sph':
ind_s, ind_p = a.lut[qind_s], a.lut[qind_p]
s_slice, p_slice = vals[:, :, :, ind_s], vals[:, :, :, ind_p]
t_slice = ref_T * (p_slice / ref_P *
(1. - 1. / thermo_gamma) + s_slice / c_P)
sslice = prime_sph(t_slice, tw)
# Spherical magnetic fields
elif varname == 'br':
ind_br = a.lut[qind_br]
sslice = vals[:, :, :, ind_br]
elif varname == 'bt':
ind_bt = a.lut[qind_bt]
sslice = vals[:, :, :, ind_bt]
elif varname == 'bp':
ind_bp = a.lut[qind_bp]
sslice = vals[:, :, :, ind_bp]
elif varname == 'br_prime':
ind_br = a.lut[qind_br]
br_slice = vals[:, :, :, ind_br]
sslice = prime(br_slice)
elif varname == 'bt_prime':
ind_bt = a.lut[qind_bt]
bt_slice = vals[:, :, :, ind_bt]
sslice = prime(bt_slice)
elif varname == 'bp_prime':
ind_bp = a.lut[qind_bp]
bp_slice = vals[:, :, :, ind_bp]
sslice = prime(bp_slice)
# Cylindrical magnetic fields
elif varname == 'bl':
ind_br, ind_bt = a.lut[qind_br], a.lut[qind_bt]
br_slice = vals[:, :, :, ind_br]
bt_slice = vals[:, :, :, ind_bt]
sslice = br_slice * sint + bt_slice * cost
elif varname == 'bz':
ind_br, ind_bt = a.lut[qind_br], a.lut[qind_bt]
br_slice = vals[:, :, :, ind_br]
bt_slice = vals[:, :, :, ind_bt]
sslice = br_slice * cost - bt_slice * sint
elif varname == 'bl_prime':
ind_br, ind_bt = a.lut[qind_br], a.lut[qind_bt]
br_slice = vals[:, :, :, ind_br]
bt_slice = vals[:, :, :, ind_bt]
bl_slice = br_slice * sint + bt_slice * cost
sslice = prime(bl_slice)
elif varname == 'bz_prime':
ind_br, ind_bt = a.lut[qind_br], a.lut[qind_bt]
br_slice = vals[:, :, :, ind_br]
bt_slice = vals[:, :, :, ind_bt]
bz_slice = br_slice * cost - bt_slice * sint
sslice = prime(bz_slice)
# Spherical fluctuating velocity products
# (multiply by rho first to get units of pressure)
elif varname == 'vsq':
ind_vr, ind_vt, ind_vp = a.lut[qind_vr], a.lut[qind_vt],\
a.lut[qind_vp]
vr_slice = vals[:, :, :, ind_vr]
vt_slice = vals[:, :, :, ind_vt]
vp_slice = vals[:, :, :, ind_vp]
vr_prime_slice = prime(vr_slice)
vt_prime_slice = prime(vt_slice)
vp_prime_slice = prime(vp_slice)
sslice = ref_rho*(vr_prime_slice**2 + vt_prime_slice**2 +\
vp_prime_slice**2)
elif varname == 'vrsq':
ind_vr = a.lut[qind_vr]
vr_slice = vals[:, :, :, ind_vr]
vr_prime_slice = vr_slice - np.mean(vr_slice, axis=0)
sslice = ref_rho * vr_prime_slice**2
elif varname == 'vtsq':
ind_vt = a.lut[qind_vt]
vt_slice = vals[:, :, :, ind_vt]
vt_prime_slice = prime(vt_slice)
sslice = ref_rho * vt_prime_slice**2
elif varname == 'vpsq':
ind_vp = a.lut[qind_vp]
vp_slice = vals[:, :, :, ind_vp]
vp_prime_slice = prime(vp_slice)
sslice = ref_rho * vp_prime_slice**2
elif varname == 'vhsq':
ind_vt, ind_vp = a.lut[qind_vt], a.lut[qind_vp]
vt_slice = vals[:, :, :, ind_vt]
vp_slice = vals[:, :, :, ind_vp]
vt_prime_slice = prime(vt_slice)
vp_prime_slice = prime(vp_slice)
sslice = ref_rho * (vt_prime_slice**2 + vp_prime_slice**2)
elif varname == 'vrvp':
ind_vr, ind_vp = a.lut[qind_vr], a.lut[qind_vp]
vr_slice = vals[:, :, :, ind_vr]
vp_slice = vals[:, :, :, ind_vp]
vr_prime_slice = prime(vr_slice)
vp_prime_slice = prime(vp_slice)
sslice = ref_rho * vr_prime_slice * vp_prime_slice
elif varname == 'vrvt':
ind_vr, ind_vt = a.lut[qind_vr], a.lut[qind_vt]
vr_slice = vals[:, :, :, ind_vr]
vt_slice = vals[:, :, :, ind_vt]
vr_prime_slice = prime(vr_slice)
vt_prime_slice = prime(vt_slice)
sslice = ref_rho * vr_prime_slice * vt_prime_slice
elif varname == 'vtvp':
ind_vt, ind_vp = a.lut[qind_vt], a.lut[qind_vp]
vt_slice = vals[:, :, :, ind_vt]
vp_slice = vals[:, :, :, ind_vp]
vt_prime_slice = prime(vt_slice)
vp_prime_slice = prime(vp_slice)
sslice = ref_rho * vt_prime_slice * vp_prime_slice
# Cylindrical fluctuating velocity products velocity products
# Multiplied by rho to get units of pressure
elif varname == 'vlsq':
ind_vr, ind_vt = a.lut[qind_vr], a.lut[qind_vt]
vr_slice = vals[:, :, :, ind_vr]
vt_slice = vals[:, :, :, ind_vt]
vl_slice = vr_slice * sint + vt_slice * cost
vl_prime_slice = prime(vl_slice)
sslice = ref_rho * vl_prime_slice**2
elif varname == 'vzsq':
ind_vr, ind_vt = a.lut[qind_vr], a.lut[qind_vt]
vr_slice = vals[:, :, :, ind_vr]
vt_slice = vals[:, :, :, ind_vt]
vz_slice = vr_slice * cost - vt_slice * sint
vz_prime_slice = prime(vz_slice)
sslice = ref_rho * vz_prime_slice**2
elif varname == 'vmsq':
ind_vr, ind_vt = a.lut[qind_vr], a.lut[qind_vt]
vr_slice = vals[:, :, :, ind_vr]
vt_slice = vals[:, :, :, ind_vt]
vr_prime_slice = prime(vr_slice)
vt_prime_slice = prime(vt_slice)
sslice = ref_rho * (vr_prime_slice**2 + vt_prime_slice**2)
elif varname == 'vlvp':
ind_vr, ind_vt, ind_vp = a.lut[qind_vr], a.lut[qind_vt],\
a.lut[qind_vp]
vr_slice = vals[:, :, :, ind_vr]
vt_slice = vals[:, :, :, ind_vt]
vp_slice = vals[:, :, :, ind_vp]
vr_prime_slice = prime(vr_slice)
vt_prime_slice = prime(vt_slice)
vp_prime_slice = prime(vp_slice)
vl_prime_slice = vr_prime_slice * sint + vt_prime_slice * cost
sslice = ref_rho * vl_prime_slice * vp_prime_slice
elif varname == 'vlvz':
ind_vr, ind_vt, ind_vp = a.lut[qind_vr], a.lut[qind_vt],\
a.lut[qind_vp]
vr_slice = vals[:, :, :, ind_vr]
vt_slice = vals[:, :, :, ind_vt]
vp_slice = vals[:, :, :, ind_vp]
vr_prime_slice = prime(vr_slice)
vt_prime_slice = prime(vt_slice)
vp_prime_slice = prime(vp_slice)
vl_prime_slice = vr_prime_slice * sint + vt_prime_slice * cost
vz_prime_slice = vr_prime_slice * cost - vt_prime_slice * sint
sslice = ref_rho * vl_prime_slice * vz_prime_slice
elif varname == 'vpvz':
ind_vr, ind_vt, ind_vp = a.lut[qind_vr], a.lut[qind_vt],\
a.lut[qind_vp]
vr_slice = vals[:, :, :, ind_vr]
vt_slice = vals[:, :, :, ind_vt]
vp_slice = vals[:, :, :, ind_vp]
vr_prime_slice = prime(vr_slice)
vt_prime_slice = prime(vt_slice)
vp_prime_slice = prime(vp_slice)
vz_prime_slice = vr_prime_slice * cost - vt_prime_slice * sint
sslice = ref_rho * vp_prime_slice * vz_prime_slice
# Spherical fluctuating magnetic field products
# (multiply by 1/4*pi to get units of magnetic pressure)
elif varname == 'bsq':
ind_br, ind_bt, ind_bp = a.lut[qind_br], a.lut[qind_bt],\
a.lut[qind_bp]
br_slice = vals[:, :, :, ind_br]
bt_slice = vals[:, :, :, ind_bt]
bp_slice = vals[:, :, :, ind_bp]
br_prime_slice = prime(br_slice)
bt_prime_slice = prime(bt_slice)
bp_prime_slice = prime(bp_slice)
sslice = (br_prime_slice**2 + bt_prime_slice**2 +\
bp_prime_slice**2)/(4.*np.pi)
elif varname == 'brsq':
ind_br = a.lut[qind_br]
br_slice = vals[:, :, :, ind_br]
br_prime_slice = prime(br_slice)
sslice = br_prime_slice**2 / (4. * np.pi)
elif varname == 'btsq':
ind_bt = a.lut[qind_bt]
bt_slice = vals[:, :, :, ind_bt]
bt_prime_slice = prime(bt_slice)
sslice = bt_prime_slice**2 / (4. * np.pi)
elif varname == 'bpsq':
ind_bp = a.lut[qind_bp]
bp_slice = vals[:, :, :, ind_bp]
bp_prime_slice = prime(bp_slice)
sslice = bp_prime_slice**2 / (4. * np.pi)
elif varname == 'bhsq':
ind_bt, ind_bp = a.lut[qind_bt], a.lut[qind_bp]
bt_slice = vals[:, :, :, ind_bt]
bp_slice = vals[:, :, :, ind_bp]
bt_prime_slice = prime(bt_slice)
bp_prime_slice = prime(bp_slice)
sslice = (bt_prime_slice**2 + bp_prime_slice**2) / (4. * np.pi)
elif varname == 'brbp':
ind_br, ind_bp = a.lut[qind_br], a.lut[qind_bp]
br_slice = vals[:, :, :, ind_br]
bp_slice = vals[:, :, :, ind_bp]
#br_prime_slice = prime(br_slice)
#bp_prime_slice = prime(bp_slice)
#sslice = br_prime_slice*bp_prime_slice/(4.*np.pi)
sslice = br_slice * bp_slice / (4 * np.pi)
elif varname == 'brbt':
ind_br, ind_bt = a.lut[qind_br], a.lut[qind_bt]
br_slice = vals[:, :, :, ind_br]
bt_slice = vals[:, :, :, ind_bt]
#br_prime_slice = prime(br_slice)
#bt_prime_slice = prime(bt_slice)
#sslice = br_prime_slice*bt_prime_slice/(4.*np.pi)
sslice = br_slice * bt_slice / (4 * np.pi)
elif varname == 'btbp':
ind_bt, ind_bp = a.lut[qind_bt], a.lut[qind_bp]
bt_slice = vals[:, :, :, ind_bt]
bp_slice = vals[:, :, :, ind_bp]
#bt_prime_slice = prime(bt_slice)
#bp_prime_slice = prime(bp_slice)
#sslice = bt_prime_slice*bp_prime_slice/(4.*np.pi)
sslice = bt_slice * bp_slice / (4 * np.pi)
# Cylindrical fluctuating magnetic field products
# Multiplied by 1/4*pi to get units of magnetic pressure
elif varname == 'blsq':
sint = (a.sintheta).reshape((1, a.ntheta, 1))
cost = (a.costheta).reshape((1, a.ntheta, 1))
ind_br, ind_bt = a.lut[qind_br], a.lut[qind_bt]
br_slice = vals[:, :, :, ind_br]
bt_slice = vals[:, :, :, ind_bt]
bl_slice = br_slice * sint + bt_slice * cost
bl_prime_slice = prime(bl_slice)
sslice = bl_prime_slice**2 / (4. * np.pi)
elif varname == 'bzsq':
ind_br, ind_bt = a.lut[qind_br], a.lut[qind_bt]
br_slice = vals[:, :, :, ind_br]
bt_slice = vals[:, :, :, ind_bt]
bz_slice = br_slice * cost - bt_slice * sint
bz_prime_slice = prime(bz_slice)
sslice = bz_prime_slice**2 / (4. * np.pi)
elif varname == 'bmsq':
ind_br, ind_bt = a.lut[qind_br], a.lut[qind_bt]
br_slice = vals[:, :, :, ind_br]
bt_slice = vals[:, :, :, ind_bt]
br_prime_slice = prime(br_slice)
bt_prime_slice = prime(bt_slice)
sslice = (br_prime_slice**2 + bt_prime_slice**2) / (4. * np.pi)
elif varname == 'blbp':
ind_br, ind_bt, ind_bp = a.lut[qind_br], a.lut[qind_bt],\
a.lut[qind_bp]
br_slice = vals[:, :, :, ind_br]
bt_slice = vals[:, :, :, ind_bt]
bp_slice = vals[:, :, :, ind_bp]
br_prime_slice = prime(br_slice)
bt_prime_slice = prime(bt_slice)
bp_prime_slice = prime(bp_slice)
bl_prime_slice = br_prime_slice * sint + bt_prime_slice * cost
sslice = bl_prime_slice * bp_prime_slice / (4. * np.pi)
elif varname == 'blbz':
ind_br, ind_bt, ind_bp = a.lut[qind_br], a.lut[qind_bt],\
a.lut[qind_bp]
br_slice = vals[:, :, :, ind_br]
bt_slice = vals[:, :, :, ind_bt]
bp_slice = vals[:, :, :, ind_bp]
br_prime_slice = prime(br_slice)
bt_prime_slice = prime(bt_slice)
bp_prime_slice = prime(bp_slice)
bl_prime_slice = br_prime_slice * sint + bt_prime_slice * cost
bz_prime_slice = br_prime_slice * cost - bt_prime_slice * sint
sslice = bl_prime_slice * bz_prime_slice / (4. * np.pi)
elif varname == 'bpbz':
ind_br, ind_bt, ind_bp = a.lut[qind_br], a.lut[qind_bt],\
a.lut[qind_bp]
br_slice = vals[:, :, :, ind_br]
bt_slice = vals[:, :, :, ind_bt]
bp_slice = vals[:, :, :, ind_bp]
br_prime_slice = prime(br_slice)
bt_prime_slice = prime(bt_slice)
bp_prime_slice = prime(bp_slice)
bz_prime_slice = br_prime_slice * cost - bt_prime_slice * sint
sslice = bp_prime_slice * bz_prime_slice / (4. * np.pi)
# Correlations between vertical flow (vr) and temperature/entropy
elif varname == 'vrt_sph':
ind_vr, ind_s, ind_p = a.lut[qind_vr], a.lut[qind_s], a.lut[qind_p]
vr_slice = vals[:, :, :, ind_vr]
s_slice = vals[:, :, :, ind_s]
p_slice = vals[:, :, :, ind_p]
t_slice = ref_T*(p_slice/ref_P*(1. - 1./thermo_gamma)\
+ s_slice/c_P)
vr_prime_sph = prime_sph(vr_slice, tw)
t_prime_sph = prime_sph(t_slice, tw)
sslice = ref_rho * c_P * vr_prime_sph * t_prime_sph
elif varname == 'vrs_sph':
ind_vr, ind_s, ind_p = a.lut[qind_vr], a.lut[qind_s], a.lut[qind_p]
vr_slice = vals[:, :, :, ind_vr]
s_slice = vals[:, :, :, ind_s]
p_slice = vals[:, :, :, ind_p]
vr_prime_sph = prime_sph(vr_slice, tw)
s_prime_sph = prime_sph(s_slice, tw)
sslice = ref_rho * ref_T * vr_prime_sph * s_prime_sph
elif varname == 'vrp_sph':
ind_vr, ind_s, ind_p = a.lut[qind_vr], a.lut[qind_s], a.lut[qind_p]
vr_slice = vals[:, :, :, ind_vr]
s_slice = vals[:, :, :, ind_s]
p_slice = vals[:, :, :, ind_p]
vr_prime_sph = prime_sph(vr_slice, tw)
p_prime_sph = prime_sph(p_slice, tw)
sslice = vr_prime_sph * p_prime_sph
elif varname == 'vtt':
ind_vt, ind_s, ind_p = a.lut[qind_vt], a.lut[qind_s], a.lut[qind_p]
vt_slice = vals[:, :, :, ind_vt]
s_slice = vals[:, :, :, ind_s]
p_slice = vals[:, :, :, ind_p]
vt_prime_slice = prime(vt_slice)
s_prime_slice = prime(s_slice)
p_prime_slice = prime(p_slice)
t_prime_slice = ref_T*(p_prime_slice/ref_P*(1. - 1./thermo_gamma)\
+ s_prime_slice/c_P)
sslice = vt_prime_slice * t_prime_slice
elif varname == 'dvtdr':
ind = a.lut[var_indices[varname]]
sslice = prime(vals[:, :, :, ind])
elif varname == 'brdvtdr':
slice_br = vals[:, :, :, a.lut[var_indices['br']]]
slice_dv = vals[:, :, :, a.lut[var_indices['dvtdr']]]
sslice = prime(slice_br) * prime(slice_dv)
elif is_an_int(varname):
sslice = vals[:, :, :, a.lut[int(varname)]]
elif 'plus' in varname and 'times' in varname: # order of operations!
sslice = np.zeros_like(vals[:, :, :, 0])
for st1 in varname.split('plus'):
vals_loc = np.ones_like(vals[:, :, :, 0])
for st2 in st1.split('times'):
vals_loc *= vals[:, :, :, a.lut[int(st2)]]
sslice += vals_loc
elif 'plus' in varname:
sslice = np.zeros_like(vals[:, :, :, 0])
for st in varname.split('plus'):
sslice += vals[:, :, :, a.lut[int(st)]]
elif 'times' in varname:
sslice = np.ones_like(vals[:, :, :, 0])
for st in varname.split('times'):
sslice *= vals[:, :, :, a.lut[int(st)]]
else:
print("get_sslice(): unknown variable name %s" % varname)
print("exiting")
sys.exit()
# possibly smooth the variable
if smooth_desired:
sslice = smooth(sslice, nphi_av)
del vals # free up memory
return sslice
# Get the run directory on which to perform the analysis
dirname = sys.argv[1]
dirname_stripped = strip_dirname(dirname)
# Data and plot directories
datadir = dirname + '/data/'
plotdir = dirname + '/plots/'
if not os.path.isdir(plotdir):
os.makedirs(plotdir)
# more defaults
minmax1 = None
minmax2 = None
dr = 0.03 # By default average energies over 10% the depth of the shell
# (shell_depth ~ 0.27 Rsun)
gi = GridInfo('/altair/loma3853/dyn_nkeom3.0-1/grid_info', '')
rw = gi.rweights
saveplot = True
symlog = False
varname = 'bp' # by default, plot B_phi and KE_phi
nstd = None
linthresh = None
desired_rvals = [0.83] # by default, plot time-radius diagram for fields
# mid-CZ (units of solar radius)
# Get command-line arguments
args = sys.argv[2:]
nargs = len(args)
for i in range(nargs):
arg = args[i]