From f03c38f99d39bb57d4dcf69f8b581fafa9ed858f Mon Sep 17 00:00:00 2001
From: aaronspring <aaron.spring@mpimet.mpg.de>
Date: Tue, 29 Jun 2021 14:30:32 +0200
Subject: [PATCH] xoak example
---
ci/pymistral.yml | 24 ++-
notebooks/xarray_xoak_ICON.ipynb | 356 +++++++++++++++++++++++++++++++
2 files changed, 370 insertions(+), 10 deletions(-)
create mode 100644 notebooks/xarray_xoak_ICON.ipynb
diff --git a/ci/pymistral.yml b/ci/pymistral.yml
index c396676..9d5c278 100644
--- a/ci/pymistral.yml
+++ b/ci/pymistral.yml
@@ -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
diff --git a/notebooks/xarray_xoak_ICON.ipynb b/notebooks/xarray_xoak_ICON.ipynb
new file mode 100644
index 0000000..736c8ef
--- /dev/null
+++ b/notebooks/xarray_xoak_ICON.ipynb
@@ -0,0 +1,356 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# `xarray` and `ICON`"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import xarray as xr\n",
+ "xr.set_options(display_style='text')\n",
+ "import numpy as np\n",
+ "import matplotlib.pyplot as plt\n",
+ "import matplotlib\n",
+ "import cartopy.crs as ccrs\n",
+ "import cartopy as cp"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# `xoak`"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "- https://xoak.readthedocs.io/en/latest/examples/introduction.html\n",
+ "- https://xoak.readthedocs.io/en/latest/examples/dask_support.html"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import xoak"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def rad_to_deg(ds):\n",
+ " \"\"\"Convert radian units to deg.\"\"\"\n",
+ " #ds.coords.compute()\n",
+ " for c in ds.coords:\n",
+ " if 'units' in ds[c].attrs:\n",
+ " if ds[c].attrs['units'] == 'radian':\n",
+ " print(f'convert {c} from rad to deg')\n",
+ " ds[c] = ds[c]* 180./np.pi\n",
+ " ds[c].attrs['units'] = 'degrees'\n",
+ " elif 'bnds' in c:\n",
+ " print(f'convert {c} from rad to deg')\n",
+ " ds[c] = ds[c]* 180./np.pi\n",
+ " ds[c].attrs['units'] = 'degrees'\n",
+ " return ds"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "#igrid = xr.open_dataset('/work/mh0727/m300732/icon-oes-hamocc/experiments/ler0956/icon_grid_0036_R02B04_O.nc')\n",
+ "#igrid = rad_to_deg(igrid)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": []
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# ICON output I want to work with\n",
+ "ds = xr.open_dataset('/work/mh0727/m300732/icon-oes-hamocc/experiments/ler0956/ler0956_hamocc_34500101T000000Z.nc',\n",
+ " chunks={'time': 1})"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# move from data_vars to coords\n",
+ "ds = ds.set_coords(['clon_bnds','clat_bnds'])\n",
+ "\n",
+ "ds = rad_to_deg(ds)\n",
+ "# \n",
+ "ds['phyp'] # ncells is dimension"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# new grid I want to work on\n",
+ "mpiom_grid = xr.open_dataset('/home/mpim/m300524/mpiom_fx.nc')"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## remap with xoak"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# set xoak index\n",
+ "ds.xoak.set_index(['clat', 'clon'], 'sklearn_geo_balltree')"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from dask.diagnostics import ProgressBar\n",
+ "import dask\n",
+ "with ProgressBar(), dask.config.set(scheduler='processes'):\n",
+ " ds_selection = ds.xoak.sel(clat=mpiom_grid.lat, clon=mpiom_grid.lon)\n",
+ "\n",
+ "# %time ds_selection = ds.xoak.sel(clat=mpiom_grid.lat, clon=mpiom_grid.lon) # also runs without dask scheduler\n",
+ "print('remapping',ds.nbytes/1e6,'MB to ',ds_selection.nbytes/1e6,'MB')"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# all lazy dask arrays\n",
+ "ds_selection['phyp'] # ncells is no dimension anymore, but x and y"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## plot remapped with `xarray.plot()`"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "ds_selection['phyp'].isel(depth=0).plot(col='time', col_wrap=5, robust=True, yincrease=False)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "## plot remapped with `pymistral.plot()` wrapping xarray and cartopy"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import pymistral\n",
+ "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",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": []
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": []
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "# Triangular plotting\n",
+ "\n",
+ "plotting the triangles with `plt.tripcolor` from `xarray`"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "- could be handle by accessor: http://xarray.pydata.org/en/stable/internals/extending-xarray.html\n",
+ "- rewrite nicer with xarray instead of np commands where possible"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "@xr.register_dataset_accessor('icon')\n",
+ "class CartopyMap(object):\n",
+ " def __init__(self, xarray_obj):\n",
+ " self._obj = xarray_obj\n",
+ "\n",
+ " def plot(self,\n",
+ " v=None,\n",
+ " ax=None,\n",
+ " proj=ccrs.PlateCarree(),\n",
+ " robust=False,\n",
+ " **tripcolor_kwargs\n",
+ " ):\n",
+ " \"\"\"\n",
+ " Plot the variable v from an xr.Dataset.\n",
+ " \n",
+ " Note: xr.DataArray.icon.plot() would be nicer,\n",
+ " but the xr.DataArray doesnt carry the neccessary clon_bnds and clat_bnds coords.\n",
+ " \n",
+ " Also having some issues how to get the plot nice without buggy triangle,\n",
+ " but that happened also in matplotlib.\n",
+ " \n",
+ " Ideally, this plot function could work like xr.DataArray.plot(col, row, vmin, vmax, robust, ...)\n",
+ " \"\"\"\n",
+ " da = self._obj\n",
+ " # --- Triangulation\n",
+ " ntr = da.clon.size\n",
+ " clon_bnds_rs = da['clon_bnds'].values.reshape(ntr*3)\n",
+ " clat_bnds_rs = da['clat_bnds'].values.reshape(ntr*3)\n",
+ " triangles = np.arange(ntr*3).reshape(ntr,3)\n",
+ " Tri = matplotlib.tri.Triangulation(clon_bnds_rs, clat_bnds_rs, triangles=triangles)\n",
+ "\n",
+ " if isinstance(da,xr.Dataset):\n",
+ " if v is None:\n",
+ " v=list(da.data_vars)\n",
+ " if len(v)>0:\n",
+ " v=v[0] # take first variable\n",
+ " da=da[v]\n",
+ " else:\n",
+ " da=da[v]\n",
+ "\n",
+ " vmin = tripcolor_kwargs.pop('vmin',None)\n",
+ " vmax = tripcolor_kwargs.pop('vmax',None)\n",
+ " if vmin is None and vmax is None and robust:\n",
+ " vmin = da.quantile(0.02).values\n",
+ " vmax = da.quantile(.98).values\n",
+ " \n",
+ " # what if more than two dims: see xr.facetgrid\n",
+ " fig, axes = plt.subplots(subplot_kw=dict(projection=proj),figsize=(8,4))\n",
+ " # best would be to add tripcolor to xarray\n",
+ " hm = axes.tripcolor(Tri,\n",
+ " da,\n",
+ " vmin=vmin,\n",
+ " vmax=vmax,\n",
+ " #edgecolors='k',\n",
+ " #lw=0.01,\n",
+ " transform=ccrs.PlateCarree(),\n",
+ " alpha=.5,\n",
+ " **tripcolor_kwargs\n",
+ " )\n",
+ " # now add features, labels etc\n",
+ " axes.add_feature(cp.feature.LAND, zorder=0)\n",
+ " #axes.set_extent((-178, 178, -88, 88), crs=ccrs.PlateCarree())\n",
+ " cb = plt.colorbar(mappable=hm, ax=axes, extend='both')\n",
+ " ## use metadata\n",
+ " title_str = ''\n",
+ " title_labels = [c for c in da.coords if da.coords[c].size == 1]\n",
+ " for l in title_labels:\n",
+ " title_str += f', {l} = {str(da[l].astype(str).values)}'\n",
+ " axes.set_title(title_str[1:])\n",
+ " cb.set_label(f'{da.attrs[\"long_name\"]}\\n[{da.attrs[\"units\"]}]', rotation = 90)\n",
+ " # add x,y ticks and labels\n",
+ " return axes"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "ds.isel(depth=0,time=2).icon.plot('phyp', cmap='RdBu')"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "ds.isel(depth=0,time=2).icon.plot('po4', robust=True)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": []
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "Python [conda env:pymistral]",
+ "language": "python",
+ "name": "conda-env-pymistral-py"
+ },
+ "language_info": {
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "file_extension": ".py",
+ "mimetype": "text/x-python",
+ "name": "python",
+ "nbconvert_exporter": "python",
+ "pygments_lexer": "ipython3",
+ "version": "3.8.6"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}
--
GitLab