Skip to content
Snippets Groups Projects
Commit f03c38f9 authored by Aaron Spring's avatar Aaron Spring
Browse files

xoak example

parent bbdb1533
No related branches found
No related tags found
No related merge requests found
Hide whitespace changes
Inline Side-by-side
ci/pymistral.yml
+ 14
− 10
View file @ f03c38f9
......@@ -3,7 +3,8 @@ channels:
- conda-forge
dependencies:
- bokeh
- cartopy
- cartopy>=0.18
- matplotlib>=3.3
- cdo
- dask
- dask-jobqueue
......@@ -20,9 +21,6 @@ dependencies:
- intake
- intake-xarray
- hvplot
- geoviews
- holoviews
- pytest<5.4.0
- pytest-cov
- pytest-sugar
- nb_conda_kernels
......@@ -34,15 +32,21 @@ dependencies:
- pynio
- flake8
- nc-time-axis
- xrviz
- xhistogram
- nodejs>=10.0.0
- nb_black
- xoak
- climpred
- regionmask
- pip:
- git+https://github.com/jbusecke/cmip6_preprocessing.git
- git+https://github.com/xgcm/xgcm.git
- git+https://github.com/xgcm/xrft.git
- git+https://github.com/mathause/regionmask.git
- cmip6_preprocessing
- xgcm
- git+https://github.com/intake/intake-esm.git
- pre-commit
- pytest-lazy-fixture
- pytest-tldr
- git+https://github.com/dask/dask-labextension.git
- git+https://github.com/intake/intake_geopandas.git
- git+https://github.com/NCAR/intake-thredds.git
#- git+https://github.com/intake/intake.git
- git+https://github.com/intake/intake-xarray.git
- git+https://github.com/intake/filesystem_spec.git
notebooks/xarray_xoak_ICON.ipynb 0 → 100644
+ 356
− 0
View file @ f03c38f9
%% Cell type:markdown id: tags:
# `xarray` and `ICON`
%% Cell type:code id: tags:
``` python
import xarray as xr
xr.set_options(display_style='text')
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
import cartopy.crs as ccrs
import cartopy as cp
```
%% Cell type:markdown id: tags:
# `xoak`
%% Cell type:markdown id: tags:
- https://xoak.readthedocs.io/en/latest/examples/introduction.html
- https://xoak.readthedocs.io/en/latest/examples/dask_support.html
%% Cell type:code id: tags:
``` python
import xoak
```
%% Cell type:code id: tags:
``` python
def rad_to_deg(ds):
"""Convert radian units to deg."""
#ds.coords.compute()
for c in ds.coords:
if 'units' in ds[c].attrs:
if ds[c].attrs['units'] == 'radian':
print(f'convert {c} from rad to deg')
ds[c] = ds[c]* 180./np.pi
ds[c].attrs['units'] = 'degrees'
elif 'bnds' in c:
print(f'convert {c} from rad to deg')
ds[c] = ds[c]* 180./np.pi
ds[c].attrs['units'] = 'degrees'
return ds
```
%% Cell type:code id: tags:
``` python
#igrid = xr.open_dataset('/work/mh0727/m300732/icon-oes-hamocc/experiments/ler0956/icon_grid_0036_R02B04_O.nc')
#igrid = rad_to_deg(igrid)
```
%% Cell type:code id: tags:
``` python
```
%% Cell type:code id: tags:
``` python
# ICON output I want to work with
ds = xr.open_dataset('/work/mh0727/m300732/icon-oes-hamocc/experiments/ler0956/ler0956_hamocc_34500101T000000Z.nc',
chunks={'time': 1})
```
%% Cell type:code id: tags:
``` python
# move from data_vars to coords
ds = ds.set_coords(['clon_bnds','clat_bnds'])
ds = rad_to_deg(ds)
#
ds['phyp'] # ncells is dimension
```
%% Cell type:code id: tags:
``` python
# new grid I want to work on
mpiom_grid = xr.open_dataset('/home/mpim/m300524/mpiom_fx.nc')
```
%% Cell type:markdown id: tags:
## remap with xoak
%% Cell type:code id: tags:
``` python
# set xoak index
ds.xoak.set_index(['clat', 'clon'], 'sklearn_geo_balltree')
```
%% Cell type:code id: tags:
``` python
from dask.diagnostics import ProgressBar
import dask
with ProgressBar(), dask.config.set(scheduler='processes'):
ds_selection = ds.xoak.sel(clat=mpiom_grid.lat, clon=mpiom_grid.lon)
# %time ds_selection = ds.xoak.sel(clat=mpiom_grid.lat, clon=mpiom_grid.lon) # also runs without dask scheduler
print('remapping',ds.nbytes/1e6,'MB to ',ds_selection.nbytes/1e6,'MB')
```
%% Cell type:code id: tags:
``` python
# all lazy dask arrays
ds_selection['phyp'] # ncells is no dimension anymore, but x and y
```
%% Cell type:markdown id: tags:
## plot remapped with `xarray.plot()`
%% Cell type:code id: tags:
``` python
ds_selection['phyp'].isel(depth=0).plot(col='time', col_wrap=5, robust=True, yincrease=False)
```
%% Cell type:markdown id: tags:
## plot remapped with `pymistral.plot()` wrapping xarray and cartopy
%% Cell type:code id: tags:
``` python
import pymistral
ds_selection['phyp'].isel(depth=0).rename({'clon':'lon','clat':'lat'}).compute().plot_map(col='time',col_wrap=3,robust=True, aspect=2, feature='land')
```
%% Cell type:code id: tags:
``` python
```
%% Cell type:code id: tags:
``` python
```
%% Cell type:markdown id: tags:
# Triangular plotting
plotting the triangles with `plt.tripcolor` from `xarray`
%% Cell type:markdown id: tags:
- could be handle by accessor: http://xarray.pydata.org/en/stable/internals/extending-xarray.html
- rewrite nicer with xarray instead of np commands where possible
%% Cell type:code id: tags:
``` python
@xr.register_dataset_accessor('icon')
class CartopyMap(object):
def __init__(self, xarray_obj):
self._obj = xarray_obj
def plot(self,
v=None,
ax=None,
proj=ccrs.PlateCarree(),
robust=False,
**tripcolor_kwargs
):
"""
Plot the variable v from an xr.Dataset.
Note: xr.DataArray.icon.plot() would be nicer,
but the xr.DataArray doesnt carry the neccessary clon_bnds and clat_bnds coords.
Also having some issues how to get the plot nice without buggy triangle,
but that happened also in matplotlib.
Ideally, this plot function could work like xr.DataArray.plot(col, row, vmin, vmax, robust, ...)
"""
da = self._obj
# --- Triangulation
ntr = da.clon.size
clon_bnds_rs = da['clon_bnds'].values.reshape(ntr*3)
clat_bnds_rs = da['clat_bnds'].values.reshape(ntr*3)
triangles = np.arange(ntr*3).reshape(ntr,3)
Tri = matplotlib.tri.Triangulation(clon_bnds_rs, clat_bnds_rs, triangles=triangles)
if isinstance(da,xr.Dataset):
if v is None:
v=list(da.data_vars)
if len(v)>0:
v=v[0] # take first variable
da=da[v]
else:
da=da[v]
vmin = tripcolor_kwargs.pop('vmin',None)
vmax = tripcolor_kwargs.pop('vmax',None)
if vmin is None and vmax is None and robust:
vmin = da.quantile(0.02).values
vmax = da.quantile(.98).values
# what if more than two dims: see xr.facetgrid
fig, axes = plt.subplots(subplot_kw=dict(projection=proj),figsize=(8,4))
# best would be to add tripcolor to xarray
hm = axes.tripcolor(Tri,
da,
vmin=vmin,
vmax=vmax,
#edgecolors='k',
#lw=0.01,
transform=ccrs.PlateCarree(),
alpha=.5,
**tripcolor_kwargs
)
# now add features, labels etc
axes.add_feature(cp.feature.LAND, zorder=0)
#axes.set_extent((-178, 178, -88, 88), crs=ccrs.PlateCarree())
cb = plt.colorbar(mappable=hm, ax=axes, extend='both')
## use metadata
title_str = ''
title_labels = [c for c in da.coords if da.coords[c].size == 1]
for l in title_labels:
title_str += f', {l} = {str(da[l].astype(str).values)}'
axes.set_title(title_str[1:])
cb.set_label(f'{da.attrs["long_name"]}\n[{da.attrs["units"]}]', rotation = 90)
# add x,y ticks and labels
return axes
```
%% Cell type:code id: tags:
``` python
ds.isel(depth=0,time=2).icon.plot('phyp', cmap='RdBu')
```
%% Cell type:code id: tags:
``` python
ds.isel(depth=0,time=2).icon.plot('po4', robust=True)
```
%% Cell type:code id: tags:
``` python
```
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

Imprint | Privacy Policy