icon_smtnatl_dissipation#

[1]:

# Jupyter Notebook with widget matplotlib plots
#%matplotlib notebook
# Jupyter Lab with widget matplotlib plots
#%matplotlib widget
# with Jupyter and Jupyter Lab but without widget matplotlib plots
%matplotlib inline
%load_ext autoreload
%autoreload 2

[2]:

import sys
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from netCDF4 import Dataset, num2date
from ipdb import set_trace as mybreak
import pyicon as pyic
import cartopy.crs as ccrs
import glob
import pickle
import maps_icon_smt_temp as smt
import datetime

from icon_smt_levels import dzw, dzt, depthc, depthi

-----calc
sys glob os
numpy
netcdf
Done modules calc.
-----calc_xr
sys glob os
numpy
netcdf
xarray
Done modules calc.
-----tb
sys
json
numpy
scipy
netcdf datetime
matplotlib
mybreak
pnadas
xarray
done xarray
-----IconData
-----plotting
-----view
-----calc
-----calc_xr
-----tb
-----IconData
-----plotting
-----view
-----quickplots
-----quickplots

[3]:

import xarray as xr
import dask
import dask.array as da
from dask.distributed import Client, LocalCluster
import multiprocessing

[4]:

%%javascript
IPython.notebook.kernel.execute('nb_name = "' + IPython.notebook.notebook_name + '"')

[5]:

#nb_name

[6]:

def savefig(txt):
    fbase = nb_name.split('/')[-1][:-6]
    fpath = f'{path_fig}{fbase}_{txt}.pdf'
    print('save figure: %s' % (fpath))
    plt.savefig(fpath)
    return

Setup cluster#

[7]:

cluster = LocalCluster(#ip="0.0.0.0",
                       silence_logs=50,
                       n_workers=12,  # half of workers than cores
                       threads_per_worker=1, # 1 is often fastern than 2
                       memory_limit='8G', # per worker number of tot. memory / by worker (a bit lower)
                      )
client = Client(cluster)

[8]:

client

[8]:

Client

Client-3518ef4f-0838-11ed-9b42-080038c049f9

Connection method: Cluster object Cluster type: distributed.LocalCluster
Dashboard: http://127.0.0.1:8787/status

Cluster Info

Define simulation#

[86]:

run = 'ngSMT_tke'
savefig = False
path_fig = '../pics/'
nnf=0

gname = 'smt'
lev = 'L128'

path_data = f'/work/mh0287/from_Mistral/mh0287/users/leonidas/icon/ngSMT/results/????-??/'
#fpath_tgrid = '/mnt/lustre01/work/mh0287/users/leonidas/icon/submeso/grid/cell_grid-OceanOnly_SubmesoNA_2500m_srtm30-icon.nc'
fpath_tgrid = '/home/m/m300602/work/icon/grids/smt/smt_tgrid.nc'
fpath_Tri = '/mnt/lustre01/work/mh0033/m300602/tmp/Tri.pkl'

path_grid     = f'/work/mh0033/m300602/icon/grids/{gname}/'
path_ckdtree  = f'{path_grid}ckdtree/'
fpath_ckdtree = f'{path_grid}ckdtree/rectgrids/{gname}_res0.02_180W-180E_90S-90N.npz'

[30]:

time0 = np.datetime64('2010-02-01T15:00:00')

Load the grid#

[31]:

!ls /mnt/lustre01/work/mh0033/m300602/icon/grids/smt/

ls: cannot access '/mnt/lustre01/work/mh0033/m300602/icon/grids/smt/': No such file or directory

[32]:

# %%time
# IcD = pyic.IconData(
#                     fname        = f'{run}_T_S_sp_001-016_*.nc',
#                     path_data    = '/mnt/lustre01/work/mh0287/users/leonidas/icon/ngSMT/results/2010-03/',
#                     path_grid    = path_grid,
#                     gname        = gname,
#                     lev          = lev,
#                     do_triangulation      = False,
#                     omit_last_file        = True,
#                     calc_coeff            = False,
#                     load_vertical_grid    = False,
#                     load_triangular_grid  = True,
#                     load_rectangular_grid = False,
#                     load_variable_info    = False,
#                     verbose               = True,
#                     time_mode             ='float2date',
#                     )

[55]:

# --- load grid variables
f = Dataset(fpath_tgrid, 'r')
clon = f.variables['clon'][:] * 180./np.pi
clat = f.variables['clat'][:] * 180./np.pi
vlon = f.variables['vlon'][:] * 180./np.pi
vlat = f.variables['vlat'][:] * 180./np.pi
elon = f.variables['elon'][:] * 180./np.pi
elat = f.variables['elat'][:] * 180./np.pi
vertex_of_cell = f.variables['vertex_of_cell'][:].transpose()-1
edge_of_cell = f.variables['edge_of_cell'][:].transpose()-1
edges_of_vertex = f.variables['edges_of_vertex'][:].transpose()-1
dual_edge_length = f.variables['dual_edge_length'][:]
edge_orientation = f.variables['edge_orientation'][:].transpose()
edge_length = f.variables['edge_length'][:]
cell_area_p = f.variables['cell_area_p'][:]
dual_area = f.variables['dual_area'][:]
edge_length = f.variables['edge_length'][:]
adjacent_cell_of_edge = f.variables['adjacent_cell_of_edge'][:].transpose()-1
orientation_of_normal = f.variables['orientation_of_normal'][:].transpose()
edge_vertices = f.variables['edge_vertices'][:].transpose()-1

dtype = 'float32'
cell_cart_vec = np.ma.zeros((clon.size,3), dtype=dtype)
cell_cart_vec[:,0] = f.variables['cell_circumcenter_cartesian_x'][:]
cell_cart_vec[:,1] = f.variables['cell_circumcenter_cartesian_y'][:]
cell_cart_vec[:,2] = f.variables['cell_circumcenter_cartesian_z'][:]
# vert_cart_vec = np.ma.zeros((vlon.size,3), dtype=dtype)
# vert_cart_vec[:,0] = f.variables['cartesian_x_vertices'][:]
# vert_cart_vec[:,1] = f.variables['cartesian_y_vertices'][:]
# vert_cart_vec[:,2] = f.variables['cartesian_z_vertices'][:]
edge_cart_vec = np.ma.zeros((elon.size,3), dtype=dtype)
edge_cart_vec[:,0] = f.variables['edge_middle_cartesian_x'][:]
edge_cart_vec[:,1] = f.variables['edge_middle_cartesian_y'][:]
edge_cart_vec[:,2] = f.variables['edge_middle_cartesian_z'][:]
# edge_prim_norm = np.ma.zeros((elon.size,3), dtype=dtype)
# edge_prim_norm[:,0] = f.variables['edge_primal_normal_cartesian_x'][:]
# edge_prim_norm[:,1] = f.variables['edge_primal_normal_cartesian_y'][:]
# edge_prim_norm[:,2] = f.variables['edge_primal_normal_cartesian_z'][:]
f.close()

[51]:

grid_sphere_radius = 6.371229e6
earth_angular_velocity = 7.29212e-05
fv = 2.* earth_angular_velocity * np.sin(vlat*np.pi/180.)

[53]:

%%time
# --- derive velocity conversion coeffs
sinLon = np.sin(clon*np.pi/180.)
cosLon = np.cos(clon*np.pi/180.)
sinLat = np.sin(clat*np.pi/180.)
cosLat = np.cos(clat*np.pi/180.)

dist_vector = edge_cart_vec[edge_of_cell,:] - cell_cart_vec[:,np.newaxis,:]
norm = np.sqrt(pyic.scalar_product(dist_vector,dist_vector,dim=2))
prime_edge_length = edge_length/grid_sphere_radius
fixed_vol_norm = (  0.5 * norm
                * (prime_edge_length[edge_of_cell]))
fixed_vol_norm = fixed_vol_norm.sum(axis=1)
edge2cell_coeff_cc = (  dist_vector
                    * (edge_length[edge_of_cell,np.newaxis] / grid_sphere_radius)
                    * orientation_of_normal[:,:,np.newaxis] )

[56]:

%%time
# --- derive dissipation coefficients
rot_coeff = (  dual_edge_length[edges_of_vertex]*edge_orientation ) / dual_area[:,np.newaxis]
div_coeff = (  edge_length[edge_of_cell]*orientation_of_normal ) / cell_area_p[:,np.newaxis]
grad_coeff = (1./dual_edge_length)
rottr_coeff = (1./edge_length)

[58]:

%%time
lon_reg_2 = [-75, -55]
lat_reg_2 = [33, 43]
lon_reg = lon_reg_2
lat_reg = lat_reg_2
clon_reg, clat_reg, vertex_of_cell_reg, edge_of_cell_reg, ind_reg = pyic.crop_tripolar_grid(lon_reg, lat_reg,
                       clon, clat, vertex_of_cell, edge_of_cell)
#Tri = matplotlib.tri.Triangulation(vlon, vlat, triangles=vertex_of_cell_reg)
nc = clon_reg.size

CPU times: user 675 ms, sys: 446 ms, total: 1.12 s
Wall time: 1.07 s

Load the data#

[42]:

# --- load times and flist
# search_str = f'{run}_T_S_sp_001-016_*.nc'
# search_str = f'{run}_vn_sp_001-016_*.nc'
search_str = f'{run}_vn_sp_081-096_*.nc'
flist = np.array(glob.glob(path_data+search_str))
flist.sort()
times, flist_ts, its = pyic.get_timesteps(flist, time_mode='float2date')

[43]:

(depthc<800).sum(), depthc[90]

[43]:

(90, 852.5)

[70]:

%%time
it = np.argmin((times-time0).astype(float)**2)
kk = 89
fpath = flist_ts[it]
f = Dataset(fpath, 'r')
# var = 'vn001_sp'
var = f'vn{kk+1:03d}_sp'
print(f'Loading {var}!')
ve = f.variables[var][its[it],:]
f.close()

Loading vn090_sp!
CPU times: user 1.12 s, sys: 422 ms, total: 1.54 s
Wall time: 1.48 s

Try with xarray#

[94]:

flist_ts[0]

[94]:

'/work/mh0287/from_Mistral/mh0287/users/leonidas/icon/ngSMT/results/2010-01/ngSMT_tke_vn_sp_081-096_20100109T010000Z.nc'

[95]:

search_str = f'{run}_vn_sp_*_20100109T010000Z.nc'
flist = glob.glob(path_data+search_str)
flist.sort()

[96]:

mfdset_kwargs = dict(
  combine='nested', concat_dim='time',
  data_vars='minimal', coords='minimal',
  compat='override', join='override',
  parallel=True,
)

[ ]:

Calculate dissipation#

[71]:

A4 = 0.04 * 500**3

[72]:

%%time
# --- derive dissipation
curl_v = (ve[edges_of_vertex] * rot_coeff).sum(axis=1)
div_v  = (ve[edge_of_cell]    * div_coeff).sum(axis=1)

grad_div_v = (div_v[adjacent_cell_of_edge[:,1]]-div_v[adjacent_cell_of_edge[:,0]])*grad_coeff
curlt_curl_v = (curl_v[edge_vertices[:,1]]-curl_v[edge_vertices[:,0]])*rottr_coeff

diss = -grad_div_v**2 - curlt_curl_v**2
diss *= A4

ro = curl_v/fv

CPU times: user 10.1 s, sys: 4.72 s, total: 14.8 s
Wall time: 14.2 s

[74]:

%%time
# --- interpolate to cell center
diss_i = diss[edge_of_cell].sum(axis=1)/3.
ro_i = ro[vertex_of_cell].sum(axis=1)/3.

# --- velocity
p_vn_c = ( edge2cell_coeff_cc[:,:,:]
       * ve[edge_of_cell,np.newaxis]
#          * prism_thick_e[iz,IcD.edge_of_cell,np.newaxis]
       * dzw[kk]
     ).sum(axis=1)
p_vn_c *= 1./(fixed_vol_norm[:,np.newaxis]*dzw[kk])

u1 = p_vn_c[:,0]
u2 = p_vn_c[:,1]
u3 = p_vn_c[:,2]

uo =   u2*cosLon - u1*sinLon
vo = -(u1*cosLon + u2*sinLon)*sinLat + u3*cosLat

CPU times: user 26.3 s, sys: 7.34 s, total: 33.6 s
Wall time: 32.5 s

Test plots#

[87]:

fpath_ckdtree

[87]:

'/work/mh0033/m300602/icon/grids/smt/ckdtree/rectgrids/smt_res0.02_180W-180E_90S-90N.npz'

[89]:

%%time
lon, lat, dissi = pyic.interp_to_rectgrid(diss_i, fpath_ckdtree,
                                          lon_reg=lon_reg, lat_reg=lat_reg,
                                         )

CPU times: user 4.02 s, sys: 2.74 s, total: 6.76 s
Wall time: 6.67 s

[93]:

ccrs_proj = ccrs.PlateCarree()
# ccrs_proj = None
hca, hcb = pyic.arrange_axes(1, 1, plot_cb=True, asp=0.5, projection=ccrs_proj, fig_size_fac=3)
ii=-1
clim = 1e10
clim = 1e-8

ii+=1; ax=hca[ii]; cax=hcb[ii]
hm = pyic.shade(lon, lat, dissi, ax=ax, cax=cax, clim=clim)
# ax.set_xlim(clon_reg.min(), clon_reg.max())
# ax.set_ylim(clat_reg.min(), clat_reg.max())

pyic.plot_settings(ax, xlim=lon_reg, ylim=lat_reg)
../_images/proj_smtnatl_icon_smtnatl_dissipation_37_0.png

Old: Calculate dissipation#

[39]:

vertex_of_cell.shape, vertex_of_cell.shape

[39]:

((59799625, 3), (59799625, 3))

[40]:

# edges_of_vertex = edges_of_vertex[vertex_of_cell,:]
# edge_orientation = edge_orientation[vertex_of_cell,:]

[41]:

%%time
rot_coeff = (  dual_edge_length[edges_of_vertex]*edge_orientation )
div_coeff = (  edge_length[edge_of_cell]*orientation_of_normal*cell_area_p[:,np.newaxis] )
grad_coeff = (1./dual_edge_length)
rottr_coeff = (1./edge_length)

CPU times: user 2.86 s, sys: 2.79 s, total: 5.65 s
Wall time: 5.46 s

[45]:

%%time
curl_v = (ve[edges_of_vertex] * rot_coeff).sum(axis=1)
div_v  = (ve[edge_of_cell]    * div_coeff).sum(axis=1)

grad_div_v = (div_v[adjacent_cell_of_edge[:,1]]-div_v[adjacent_cell_of_edge[:,0]])*grad_coeff
curlt_curl_v = (curl_v[edge_vertices[:,1]]-curl_v[edge_vertices[:,0]])*rottr_coeff

diss = -grad_div_v**2 - curlt_curl_v**2

CPU times: user 5.44 s, sys: 3.79 s, total: 9.23 s
Wall time: 8.91 s

[46]:

%%time
diss_i = diss[edge_of_cell].sum(axis=1)/3.

CPU times: user 1.71 s, sys: 912 ms, total: 2.62 s
Wall time: 2.54 s

[47]:

diss_ireg = diss_i[ind_reg]

[34]:

%%time
lon_reg = [-72, -70]
lat_reg = [33, 34.5]
clon_reg, clat_reg, vertex_of_cell_reg, edge_of_cell_reg, ind_reg = pyic.crop_tripolar_grid(lon_reg, lat_reg,
                       clon, clat, vertex_of_cell, edge_of_cell)
Tri = matplotlib.tri.Triangulation(vlon, vlat, triangles=vertex_of_cell_reg)

CPU times: user 465 ms, sys: 408 ms, total: 873 ms
Wall time: 833 ms

[48]:

# ccrs_proj = ccrs.PlateCarree()
ccrs_proj = None
hca, hcb = pyic.arrange_axes(1, 1, plot_cb=True, asp=0.5, projection=ccrs_proj, fig_size_fac=3)
ii=-1
clim = 1e10
clim = 'sym'

ii+=1; ax=hca[ii]; cax=hcb[ii]
hm = pyic.shade(Tri, diss_ireg, ax=ax, cax=cax, clim=clim)
ax.set_xlim(clon_reg.min(), clon_reg.max())
ax.set_ylim(clat_reg.min(), clat_reg.max())

[48]:

(33.00001530407567, 34.49996871173412)
../_images/proj_smtnatl_icon_smtnatl_dissipation_46_1.png 
[ ]:



[ ]: