#!/usr/bin/env python import matplotlib.pyplot as plt from scipy.io import * import numpy as np import sys import xarray as xr plt.ion() # for xarray, use flag_scipy = 0, for scipy.io use flag_scipy = 1 flag_scipy = 0 t_chuncks = 1 # time average dir0 = "/summer/plume/convection_jeanzay/outdir_" #dir0 = "/ccc/cont003/home/ige/deremblb/storedir/convection/outdir_" if len(sys.argv) > 1: dir0 = dir0 + str(format(sys.argv[1])).zfill(4) + '/' Re = None exec(open(dir0 + "params.in").read()) if Re is not None: nu = 1/Re file1 = 'vars.nc' if flag_scipy: f = netcdf_file(dir0 + file1, 'r') t = f.variables['time'][:].copy() x = f.variables['x'][:].copy() y = f.variables['y'][:].copy() z = f.variables['z'][:].copy() else: ds = xr.open_dataset(dir0 + file1, engine='netcdf4') t = ds['time'][:].values.squeeze() x = ds['x'][:].values.squeeze() y = ds['y'][:].values.squeeze() z = ds['z'][:].values.squeeze() si_t = len(t) si_t2 = int(np.floor(si_t/t_chuncks)) si_x = len(x) si_y = len(y) si_z = len(z) dx = x[1] - x[0] dy = y[1] - y[0] dz = z[1] - z[0] ke_prof = np.zeros((si_t2,si_z)) pe_prof = np.zeros((si_t2,si_z)) wb_prof = np.zeros((si_t2,si_z)) wpbp_prof = np.zeros((si_t2,si_z)) wpup_prof = np.zeros((si_t2,si_z)) bpbp_prof = np.zeros((si_t2,si_z)) b_prof = np.zeros((si_t2,si_z)) u_prof = np.zeros((si_t2,si_z)) v_prof = np.zeros((si_t2,si_z)) w_prof = np.zeros((si_t2,si_z)) #epsilon_prof = np.zeros((si_t2,si_z)) eke_prof = np.zeros((si_t2,si_z)) up2_prof = np.zeros((si_t2,si_z)) vp2_prof = np.zeros((si_t2,si_z)) wp2_prof = np.zeros((si_t2,si_z)) up3_prof = np.zeros((si_t2,si_z)) vp3_prof = np.zeros((si_t2,si_z)) wp3_prof = np.zeros((si_t2,si_z)) up4_prof = np.zeros((si_t2,si_z)) vp4_prof = np.zeros((si_t2,si_z)) wp4_prof = np.zeros((si_t2,si_z)) b_top = np.zeros((si_t2)) b_bot = np.zeros((si_t2)) for i in range(0,si_t2): print("iter ", i, "/", si_t2-1, end="\r") if flag_scipy: if (t_chuncks >1): print("option not coded yet") b = f.variables['b'][i,:,:].copy().squeeze() u = f.variables['u.x'][i,:,:].copy().squeeze() v = f.variables['u.y'][i,:,:].copy().squeeze() w = f.variables['u.z'][i,:,:].copy().squeeze() else: b = ds['b'][i*t_chuncks:(i+1)*t_chuncks,:,:,:] u = ds['u.x'][i*t_chuncks:(i+1)*t_chuncks,:,:,:] v = ds['u.y'][i*t_chuncks:(i+1)*t_chuncks,:,:,:] w = ds['u.z'][i*t_chuncks:(i+1)*t_chuncks,:,:,:] # dudz, dudy, dudx = np.gradient(u,dx) # dwdz, dwdy, dwdx = np.gradient(w,dz) if flag_scipy: xy_axis = (1,2) z_slice = np.s_[:,None,None] else: xy_axis = (0,2,3) z_slice = np.s_[None,:,None,None] b_prof[i] = b.mean(axis=xy_axis) u_prof[i] = u.mean(axis=xy_axis) v_prof[i] = v.mean(axis=xy_axis) w_prof[i] = w.mean(axis=xy_axis) ke_prof[i] = 0.5*np.mean(u**2 + v**2 + w**2,axis=xy_axis) pe_prof[i] = -np.mean(b*z[z_slice],axis=xy_axis) wb_prof[i] = np.mean(w*b,axis=xy_axis) wpbp_prof[i] = np.mean((w-w_prof[i][z_slice])*(b-b_prof[i][z_slice]),axis=xy_axis) wpup_prof[i] = np.mean((w-w_prof[i][z_slice])*(u-u_prof[i][z_slice]),axis=xy_axis) bpbp_prof[i] = np.mean((b-b_prof[i][z_slice])*(b-b_prof[i][z_slice]),axis=xy_axis) eke_prof[i] = 0.5*np.mean((u-u_prof[i][z_slice])**2 + (v-v_prof[i][z_slice])**2 + (w-w_prof[i][z_slice])**2,axis=xy_axis) up2_prof[i] = np.mean((u-u_prof[i][z_slice])**2,axis=xy_axis) vp2_prof[i] = np.mean((v-v_prof[i][z_slice])**2,axis=xy_axis) wp2_prof[i] = np.mean((w-w_prof[i][z_slice])**2,axis=xy_axis) up3_prof[i] = np.mean((u-u_prof[i][z_slice])**3,axis=xy_axis) vp3_prof[i] = np.mean((v-v_prof[i][z_slice])**3,axis=xy_axis) wp3_prof[i] = np.mean((w-w_prof[i][z_slice])**3,axis=xy_axis) up4_prof[i] = np.mean((u-u_prof[i][z_slice])**4,axis=xy_axis) vp4_prof[i] = np.mean((v-v_prof[i][z_slice])**4,axis=xy_axis) wp4_prof[i] = np.mean((w-w_prof[i][z_slice])**4,axis=xy_axis) # epsilon_prof[i] = nu*np.mean(dudx**2 + dudz**2 + dwdx**2 + dwdz**2,axis = 1) b_top[i] = np.mean(b[-1,:]) b_bot[i] = np.mean(b[0,:]) if flag_scipy: f.close() else: ds.close() # create netcdf file fileout = dir0 + 'profiles.nc' fo = netcdf_file(fileout,'w') fo.createDimension('t',None) fo.createDimension('z',si_z) to = fo.createVariable('t', 'f', ('t',)) zo = fo.createVariable('z', 'f', ('z',)) bo = fo.createVariable('b' , 'f', ('t','z',)) uo = fo.createVariable('u' , 'f', ('t','z',)) vo = fo.createVariable('v' , 'f', ('t','z',)) wo = fo.createVariable('w' , 'f', ('t','z',)) keo = fo.createVariable('ke' , 'f', ('t','z',)) peo = fo.createVariable('pe' , 'f', ('t','z',)) ekeo = fo.createVariable('eke' , 'f', ('t','z',)) up2o = fo.createVariable('up2' , 'f', ('t','z',)) vp2o = fo.createVariable('vp2' , 'f', ('t','z',)) wp2o = fo.createVariable('wp2' , 'f', ('t','z',)) up3o = fo.createVariable('up3' , 'f', ('t','z',)) vp3o = fo.createVariable('vp3' , 'f', ('t','z',)) wp3o = fo.createVariable('wp3' , 'f', ('t','z',)) up4o = fo.createVariable('up4' , 'f', ('t','z',)) vp4o = fo.createVariable('vp4' , 'f', ('t','z',)) wp4o = fo.createVariable('wp4' , 'f', ('t','z',)) wbo = fo.createVariable('wb' , 'f', ('t','z',)) wpbpo= fo.createVariable('wpbp' , 'f', ('t','z',)) wpupo = fo.createVariable('wpup' , 'f', ('t','z',)) bpbpo = fo.createVariable('bpbp' , 'f', ('t','z',)) #epso = fo.createVariable('epsilon' , 'f', ('t','z',)) to[:] = t[:si_t2] zo[:] = z[:] bo[:,:] = b_prof uo[:,:] = u_prof vo[:,:] = v_prof wo[:,:] = w_prof keo[:,:] = ke_prof ekeo[:,:] = eke_prof up2o[:,:] = up2_prof vp2o[:,:] = vp2_prof wp2o[:,:] = wp2_prof up3o[:,:] = up3_prof vp3o[:,:] = vp3_prof wp3o[:,:] = wp3_prof up4o[:,:] = up4_prof vp4o[:,:] = vp4_prof wp4o[:,:] = wp4_prof peo[:,:] = pe_prof wbo[:,:] = wb_prof wpbpo[:,:] = wpbp_prof wpupo[:,:] = wpup_prof bpbpo[:,:] = bpbp_prof #epso[:,:] = epsilon_prof fo.close() print("Done\n") # # WARNING: CHANGED SUM CHANGED TO MEAN ABOVE # ke_int = np.sum(ke_prof,axis=1)*dz # pe_int = np.sum(pe_prof,axis=1)*dz # wb_int = np.sum(wb_prof,axis=1)*dz # b_int = np.sum(b_prof,axis=1)*dz # epsilon_int = np.sum(epsilon_prof,axis=1)*dz # i0 = 10 # plt.figure() # plt.plot(t[i0:],pe_int[i0:]-pe_int[i0],label='pe') # plt.plot(t[i0:],ke_int[i0:],label='ke') # plt.plot(t[i0:],dtout*np.cumsum(wb_int[i0:]),label='int_wb') # plt.plot(t[i0:],dtout*np.cumsum(epsilon_int[i0:]),label='eps') # plt.plot(t[i0:],dtout*nu*np.cumsum(b_top[i0:]-b_bot[i0:]),label='t-b') # plt.legend() #plt.plot(b_prof,z)