fig5.ipynb

Kernel: Python 3 (system-wide)

import numpy as np
import matplotlib.pyplot as plt
import cmocean 
from matplotlib import cm
from matplotlib.colors import ListedColormap, LinearSegmentedColormap
import scipy as sp
import scipy.special 
import seaborn as sns
import scipy.ndimage
from matplotlib.gridspec import GridSpec
%matplotlib inline
from matplotlib import rc
rc('font',**{'family':'sans-serif','sans-serif':['Helvetica'], 'size':20})
rc('text', usetex=True)
rc('text.latex', preamble=r'\usepackage{mathrsfs}')
from matplotlib.gridspec import GridSpec
def remove_ticks(ax):
    for axis in ax:
        axis.set_xticks([])
        axis.set_yticks([])

RC = {'figure.figsize':(10,5),
      'axes.facecolor':'white',
      'axes.edgecolor':' .15',
      'axes.linewidth':' 1.25',
      'axes.grid' : True,
      'grid.color': '.8',
      'font.size' : 15,
     "xtick.major.size": 4, "ytick.major.size": 4}

plt.rcParams.update(RC)
sns.set_palette(sns.color_palette('dark'), n_colors=None, desat=None, color_codes=False)

#Define non-dimensional parameters for the simulation 
Re = 2.478883e+03
Fr = 1.105914e+00
Ri = 1/((Fr/(2*np.pi)))**2
Pr = 7
ch = 50 #input column height

#Define dimensionless functions f and g from equations (2.14) and (2.15) in the manuscript
#These can be easily modified by the user
a = 1
b = 0.8
c = 0.9
d = 0.9 
def f(x):
    return 19/8 + 11/8*np.tanh(a*np.log(x)-b)

def g(x):
    return 2 + np.tanh(c*np.log(x)-d)

#Model name
model = 'shearr_mixed_all_ch50'
data_loc = '../processed_data/test/Pr7_decaying/'
output_loc = '../PCNN_eps/'

t = ['T=0.5','T=1', 'T=2', 'T=4', 'T=6', 'T=7.7']
xdim, zdim = 500, 500
dudz, drdz, eps, isotropic_eps, empirical_eps, pcnn_eps, pcnn_mean_eps, reb, drtdz, T = [], [], [],[],[], [],[], [],[], []
count=0
for i in range(0,6): #Choose timestep(s) to load here (from the list 't' above)
    #This is a bit clunky because basic Cocalc struggles to load the data from all 6 timesteps at once.
    T.append(t[i])
    dudz.append(np.load(data_loc+ T[count]+'/dudz_surr.npy').reshape((xdim,zdim)))
    drtdz.append(np.load(data_loc+ T[count]+'/drtdz_surr.npy').reshape((xdim,zdim)))
    eps.append(np.load(data_loc + T[count]+'/test_eps.npy').reshape((xdim,zdim)))
    isotropic_eps.append(15/4*(dudz[count]))
    empirical_eps.append(f(sp.ndimage.gaussian_filter(dudz[count], sigma=(0,50))*Re/Ri)*dudz[count])
    reb.append(np.mean(eps[count])*Re/Ri)
    pcnn_eps.append(np.load(output_loc+model+'_ensemble_pred_'+T[count]+'.npy'))
    pcnn_mean_eps.append(np.load(output_loc+model+'_means_'+T[count]+'.npy').reshape((xdim,zdim)))
    count+=1

print(T)

['T=0.5', 'T=1', 'T=2', 'T=4', 'T=6', 'T=7.7']

red = sns.color_palette('deep')[2]
fig=plt.figure(figsize=(18,3))
gs=GridSpec(1,3)
ax1 = fig.add_subplot(gs[0,0], rasterized=True)
ax2 = fig.add_subplot(gs[0,1], rasterized=True)
axs = [ax1,ax2]
count=0
for J in [1,5]:
    ax = axs[count]
    rc('font',**{'family':'sans-serif','sans-serif':['Helvetica'], 'size':15})
    hist=np.histogram(pcnn_mean_eps[J].flatten(), bins=50, range=(-5,2), density=True, )
    hist2=np.histogram(pcnn_eps[J][0].flatten(), bins=50, range=(-5,2), density=True)
    hist3=np.histogram(np.log10(isotropic_eps[J].flatten()), bins=50, range=(-5,2), density=True)
    hist4=np.histogram(np.log10(empirical_eps[J].flatten()), bins=50, range=(-5,2), density=True)
    hist5=np.histogram(np.log10(eps[J].flatten()), bins=50, range=(-5,2), density=True)
    ax.plot(hist[1][1:], hist[0], label='$\\mathrm{Model\ ensemble}$')
    ax.plot(hist2[1][1:], hist2[0], label='$\\mathrm{Model\ sample}$')
    ax.plot(hist3[1][1:], hist3[0],color=red,  label='$\\mathrm{Iso\ surrogate}$')
    ax.plot(hist4[1][1:], hist4[0],color=red, linestyle='--', label='$\\mathrm{Empirical\ surrogate}$')
    ax.plot(hist5[1][1:], hist5[0], color='k', label='$\\mathrm{DNS}$', marker='x', linestyle='')
    ax.set_yscale('log')
#     plt.legend()
    ax.set_ylim(0.001,2)
    ax.set_title('$'+str(T[count])+'$', size=17)
    ax.set_xlabel('$\\log_{10}\\varepsilon_0$')
    count+=1
ax1.set_ylabel('$\\mathrm{p.d.f}$')
ax1.annotate('$(a)$', (-6,3), annotation_clip=False)
ax1.annotate('$(b)$', (3,3), annotation_clip=False)

ax2.set_yticklabels([])
ax2.set_title('$T=7.7$',size=17)
ax1.legend(loc='upper left', fontsize=10)
fig.show()