Commit ef5ba2a2 authored by Nils Brüggemann's avatar Nils Brüggemann
Browse files

notebooks/examp_plotting_dual_and_primal_grid.ipynb: Deleted functions which...

notebooks/examp_plotting_dual_and_primal_grid.ipynb: Deleted functions which are now part of the pyicon library.
parent deb42472
Pipeline #10617 passed with stages
in 8 seconds
%% Cell type:code id:4952e91c tags:
 
``` python
# 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
```
 
%% Cell type:code id:071b629f tags:
 
``` python
import matplotlib.pyplot as plt
import numpy as np
import xarray as xr
import matplotlib
from matplotlib.patches import Polygon
from matplotlib.collections import PatchCollection
from matplotlib.collections import PolyCollection
import pyicon as pyic
import cartopy.crs as ccrs
```
 
%%%% Output: stream
 
-----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
xarray
done xarray
-----IconData
-----plotting
-----view
-----calc
-----calc_xr
-----tb
-----IconData
-----plotting
-----view
-----quickplots
-----quickplots
 
%% Cell type:markdown id:2569cd5d tags:
 
# Create edges of primal and dual grid
 
%% Cell type:code id:c7a3f2a0 tags:
``` python
def patch_plot_derive_bnds(ds_tgrid, lon_reg=[-180, 180], lat_reg=[-90, 90]):
# --- load tgrid
vertex_of_cell = ds_tgrid['vertex_of_cell'][:].transpose()-1
clon_bnds = ds_tgrid.vlon[vertex_of_cell] * 180./np.pi
clat_bnds = ds_tgrid.vlat[vertex_of_cell] * 180./np.pi
cells_of_vertex = ds_tgrid['cells_of_vertex'][:].transpose()-1
vlon_bnds = ds_tgrid.clon[cells_of_vertex] * 180./np.pi
vlat_bnds = ds_tgrid.clat[cells_of_vertex] * 180./np.pi
vlon_bnds = vlon_bnds.to_masked_array()
vlat_bnds = vlat_bnds.to_masked_array()
cells_of_vertex = cells_of_vertex.to_masked_array()
vlon_bnds[cells_of_vertex==-1] = np.ma.masked
vlat_bnds[cells_of_vertex==-1] = np.ma.masked
clon = ds_tgrid.clon.data * 180./np.pi
clat = ds_tgrid.clat.data * 180./np.pi
vlon = ds_tgrid.vlon.data * 180./np.pi
vlat = ds_tgrid.vlat.data * 180./np.pi
# --- sort edges of dual grid
x = (vlon_bnds-vlon[:,np.newaxis])
y = (vlat_bnds-vlat[:,np.newaxis])
ex, ey = 1., 0.
ab_scalar = (x*ex+y*ey)/np.sqrt(x**2+y**2)
ab_cross = (x*ey-y*ex)/np.sqrt(x**2+y**2)
sort_crit = np.ma.arccos(ab_scalar) * np.sign(ab_cross) * 180./np.pi
isort = np.ma.argsort(sort_crit, axis=1)
sorted_sort_crit = np.take_along_axis(sort_crit, isort, axis=1)
cells_of_vertex_sorted = np.take_along_axis(cells_of_vertex, isort, axis=1)
vlon_bnds_sorted = np.take_along_axis(vlon_bnds, isort, axis=1)
vlat_bnds_sorted = np.take_along_axis(vlat_bnds, isort, axis=1)
vlon_bnds_sorted = xr.DataArray(vlon_bnds_sorted, dims=['vertex', 'nc'])
vlat_bnds_sorted = xr.DataArray(vlat_bnds_sorted, dims=['vertex', 'nc'])
# --- cut region
ireg_v = (vlon>lon_reg[0]) & (vlon<=lon_reg[1]) & (vlat>lat_reg[0]) & (vlat<=lat_reg[1])
ireg_c = (clon>lon_reg[0]) & (clon<=lon_reg[1]) & (clat>lat_reg[0]) & (clat<=lat_reg[1])
vlon_bnds_sorted = vlon_bnds_sorted[ireg_v]
vlat_bnds_sorted = vlat_bnds_sorted[ireg_v]
cells_of_vertex_sorted = cells_of_vertex_sorted[ireg_v]
clon_bnds = clon_bnds[ireg_c]
clat_bnds = clat_bnds[ireg_c]
return clon_bnds, clat_bnds, vlon_bnds_sorted, vlat_bnds_sorted, cells_of_vertex_sorted
```
%% Cell type:code id:acff2ac8 tags:
``` python
def patch_plot_patches_from_bnds(clon_bnds, clat_bnds, vlon_bnds, vlat_bnds, cells_of_vertex):
xy = np.concatenate((clon_bnds.data[:,:, np.newaxis],clat_bnds.data[:,:, np.newaxis]), axis=2)
patches_c = []
nc = clon_bnds.shape[0]
for nn in range(nc):
if nn%1000==0:
print(f'nn = {nn}/{nc})', end='\r')
polygon = Polygon(xy[nn,:,:], True, edgecolor='k', facecolor='b')
patches_c.append(polygon)
patches_c = np.array(patches_c)
patches_v = []
nv = vlon_bnds.shape[0]
for nn in range(nv):
if nn%1000==0:
print(f'nn = {nn}/{nv})', end='\r')
ivalid = cells_of_vertex[nn,:]!=-1
xy = np.concatenate((vlon_bnds.data[nn,ivalid, np.newaxis],vlat_bnds.data[nn,ivalid, np.newaxis]), axis=1)
polygon = Polygon(xy, True, edgecolor='k', facecolor='b')
patches_v.append(polygon)
patches_v = np.array(patches_v)
return patches_c, patches_v
```
%% Cell type:code id:43a9d26b tags:
``` python
def patch_plot_shade(patches, datai, clim='auto', cmap='auto', ax='auto', cax='auto', edgecolor='none', logplot=False):
# --- mask 0 and negative values in case of log plot
data = 1.*datai
if logplot and isinstance(data, np.ma.MaskedArray):
data[data<=0.0] = np.ma.masked
data = np.ma.log10(data)
elif logplot and not isinstance(data, np.ma.MaskedArray):
data[data<=0.0] = np.nan
data = np.log10(data)
# --- clim
if isinstance(clim, str) and clim=='auto':
clim = [None, None]
elif isinstance(clim, str) and clim=='sym':
clim = np.abs(data).max()
clim=np.array(clim)
if clim.size==1:
clim = np.array([-1, 1])*clim
if clim[0] is None:
clim[0] = data.min()
if clim[1] is None:
clim[1] = data.max()
# --- cmap
if (clim[0]==-clim[1]) and cmap=='auto':
cmap = 'RdBu_r'
elif cmap=='auto':
#cmap = 'viridis'
cmap = 'RdYlBu_r'
if isinstance(cmap, str):
cmap = getattr(plt.cm, cmap)
p = PatchCollection(patches, cmap='RdBu_r', edgecolor=edgecolor)
p.set_array(data)
p.set_clim(clim)
ax.add_collection(p)
plt.colorbar(p, cax=cax)
return p
```
%% Cell type:code id:c018d524 tags:
 
``` python
fpath_tgrid = '/home/mpim/m300602/work/icon/grids/r2b8_oce_r0004/r2b8_oce_r0004_tgrid.nc'
fpath_data = '/work/bm1102/m211054/smtwave/zstar2/experiments/exp.ocean_era51h_zstar_r2b8_21070-SMT/outdata/exp.ocean_era51h_zstar_r2b8_21070-SMT_P1D_kin_20160416T000000Z.nc'
 
# lon_reg = [4,10]
# lat_reg = [53,57]
 
lon_reg = [0,12]
lat_reg = [52,58]
```
 
%% Cell type:code id:e60d45b6 tags:
 
``` python
ds_tgrid = xr.open_dataset(fpath_tgrid)
clon = ds_tgrid.clon.data * 180./np.pi
clat = ds_tgrid.clat.data * 180./np.pi
vlon = ds_tgrid.vlon.data * 180./np.pi
vlat = ds_tgrid.vlat.data * 180./np.pi
 
ireg_v = (vlon>lon_reg[0]) & (vlon<=lon_reg[1]) & (vlat>lat_reg[0]) & (vlat<=lat_reg[1])
ireg_c = (clon>lon_reg[0]) & (clon<=lon_reg[1]) & (clat>lat_reg[0]) & (clat<=lat_reg[1])
```
 
%% Cell type:code id:0d6e7d50 tags:
 
``` python
%%time
clon_bnds, clat_bnds, vlon_bnds, vlat_bnds, cells_of_vertex = pyic.patch_plot_derive_bnds(ds_tgrid)
```
 
%%%% Output: stream
 
CPU times: user 4.63 s, sys: 3.14 s, total: 7.76 s
Wall time: 7.79 s
 
%% Cell type:code id:bc7b3162 tags:
 
``` python
vlon_bnds_reg = vlon_bnds[ireg_v]
vlat_bnds_reg = vlat_bnds[ireg_v]
cells_of_vertex_reg = cells_of_vertex[ireg_v]
clon_bnds_reg = clon_bnds[ireg_c]
clat_bnds_reg = clat_bnds[ireg_c]
```
 
%% Cell type:code id:4ca4684a tags:
 
``` python
%%time
patches_c, patches_v = pyic.patch_plot_patches_from_bnds(clon_bnds_reg, clat_bnds_reg, vlon_bnds_reg, vlat_bnds_reg, cells_of_vertex_reg)
```
 
%%%% Output: stream
 
CPU times: user 639 ms, sys: 19 ms, total: 658 ms
Wall time: 645 ms
 
%% Cell type:markdown id:b2e04e04 tags:
 
### Save grid
 
%% Cell type:code id:99bf80ca tags:
 
``` python
fpath_out = '/mnt/lustre01/work/mh0033/m300602/for_others/for_niklas/examp_plotting_dual_cell.nc'
```
 
%% Cell type:code id:2dab721c tags:
 
``` python
ds_out = xr.Dataset()
ds_out['clon_bnds'] = clon_bnds
ds_out['clat_bnds'] = clat_bnds
ds_out['vlon_bnds'] = vlon_bnds
ds_out['vlat_bnds'] = vlat_bnds
```
 
%% Cell type:code id:93319c54 tags:
 
``` python
ds_out
```
 
%%%% Output: execute_result
 
<xarray.Dataset>
Dimensions: (cell: 3729001, nc: 6, nv: 3, vertex: 1882979)
Coordinates:
vlon (cell, nv) float64 1.278 1.274 1.27 1.278 ... 1.335 1.337 1.333
vlat (cell, nv) float64 1.258 1.26 1.258 ... -0.7877 -0.7898 -0.7899
Dimensions without coordinates: cell, nc, nv, vertex
Data variables:
clon_bnds (cell, nv) float64 73.23 73.0 72.77 73.23 ... 76.49 76.59 76.4
clat_bnds (cell, nv) float64 72.1 72.21 72.1 72.1 ... -45.13 -45.25 -45.26
vlon_bnds (vertex, nc) float64 73.0 73.23 73.46 73.46 ... 76.5 76.41 76.31
vlat_bnds (vertex, nc) float64 72.13 72.18 72.13 ... -45.3 -45.34 -45.3
 
%% Cell type:code id:23d21db5 tags:
 
``` python
print(f'Writing netcdf file {fpath_out}')
ds_out.to_netcdf(fpath_out)
```
 
%%%% Output: stream
 
Writing netcdf file /mnt/lustre01/work/mh0033/m300602/for_others/for_niklas/examp_plotting_dual_cell.nc
 
%% Cell type:markdown id:4bbdf1b9 tags:
 
### Plot the grid
 
%% Cell type:code id:e457f73a tags:
 
``` python
%%time
 
proj = ccrs.PlateCarree()
hca, hcb = pyic.arrange_axes(1,1, plot_cb=False, asp=0.5, fig_size_fac=2., sharex=False, projection=proj)
ii=-1
 
 
ii+=1; ax=hca[ii]; cax=hcb[ii]
 
pc = PatchCollection(patches_c, edgecolor='k', facecolor='none')
ax.add_collection(pc)
 
pv = PatchCollection(patches_v, edgecolor='r', facecolor='none')
ax.add_collection(pv)
 
ax.set_title('Primal and dual grid')
 
 
for ax in hca:
pyic.plot_settings(ax, xlim=(6, 10), ylim=(53, 55))
```
 
%%%% Output: display_data
 
 
%%%% Output: display_data
 
 
%%%% Output: stream
 
CPU times: user 125 ms, sys: 10 ms, total: 135 ms
Wall time: 129 ms
 
%% Cell type:markdown id:a1f84b7c tags:
 
## Load data
 
%% Cell type:code id:9f002fcc tags:
 
``` python
ds = xr.open_dataset(fpath_data)
 
vort = ds.vort[0,0,ireg_v]
kin = ds.kin[0,0,ireg_c]
```
 
%% Cell type:code id:c4d8842f tags:
 
``` python
%%time
proj = ccrs.PlateCarree()
hca, hcb = pyic.arrange_axes(1,2, plot_cb=True, asp=0.5, fig_size_fac=2., sharex=False, projection=proj)
ii=-1
 
ii+=1; ax=hca[ii]; cax=hcb[ii]
p = pyic.patch_plot_shade(patches_c, kin, ax=ax, cax=cax, clim='auto', logplot=True)
ax.set_title('Kin. energy on dual grid')
 
ii+=1; ax=hca[ii]; cax=hcb[ii]
p = pyic.patch_plot_shade(patches_v, vort, ax=ax, cax=cax, clim=1e-5)
ax.set_title('Vorticity on dual grid')
 
for ax in hca:
pyic.plot_settings(ax, xlim=lon_reg, ylim=lat_reg)
```
 
%%%% Output: display_data
 
 
%%%% Output: display_data
 
 
%%%% Output: stream
 
CPU times: user 685 ms, sys: 39 ms, total: 724 ms
Wall time: 838 ms
 
%% Cell type:code id:acae041d tags:
 
``` python
```
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment