sigma_to_z_interpolation_tests.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import ESMF\n",
    "import os\n",
    "%matplotlib inline\n",
    "import xarray as xr\n",
    "import matplotlib.pyplot as plt\n",
    "import cartopy.crs as ccrs\n",
    "import cartopy.feature as cpf\n",
    "plt.style.use('default')\n",
    "import IOverticalGrid\n",
    "import seaborn as sns\n",
    "import grd\n",
    "import configM2R\n",
    "import xarray as xr\n",
    "import glob"
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "This is a test script to enable interpolation directly from sigma coordinates to z coordinates.\n",
    "The goal is to be able to use ROMS results files as forcing using model2roms.\n",
    "\n",
    "Trond Kristiansen, 12.04.2021, 13.04.2021"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "outputs": [],
   "source": [
    "def remove_all_small_variables(ds):\n",
    "    \"\"\" remove all the variables that have less than 3 dimensions\n",
    "\n",
    "        Parameters:\n",
    "            ds (xarray.Dataset): ROMS dataset\n",
    "    \"\"\"\n",
    "    for v in ds.variables:\n",
    "        if v not in ['Cs_r', 'Cs_w', 'hc', 'Vtransform', 'ocean_time']:\n",
    "            ds = ds.drop_vars(v) if len(ds[v].dims) < 3 else ds\n",
    "    return ds\n",
    "\n",
    "def select_interior(ds):\n",
    "    \"\"\"\n",
    "    discard \"exterior\" u,v,rho-points to build a symetric grid\n",
    "\n",
    "        Parameters:\n",
    "            ds (xarray.Dataset): ROMS dataset\n",
    "    \"\"\"\n",
    "    ds = ds.isel(xi_rho=slice(1,-1), eta_rho=slice(1,-1))\n",
    "    if 'xi_v' in ds.dims:\n",
    "        ds = ds.isel(xi_v=slice(1,-1))\n",
    "    if 'eta_u' in ds.dims:\n",
    "        ds = ds.isel(eta_u=slice(1,-1))\n",
    "    return ds\n",
    "\n",
    "def rename_dims(ds):\n",
    "    \"\"\" rename dimensions\n",
    "\n",
    "        Parameters:\n",
    "            ds (xarray.Dataset): ROMS dataset\n",
    "\n",
    "    \"\"\"\n",
    "    ds = ds.rename({'xi_rho': 'xh', 'xi_v': 'xh', 'xi_u': 'xq',\n",
    "                    'eta_rho': 'yh', 'eta_v': 'yq', 'eta_u': 'yh',\n",
    "                    'ocean_time': 'time'\n",
    "                    })\n",
    "    return ds\n",
    "\n",
    "def compute_depth_layers(ds, grid, hmin=0.1):\n",
    "    \"\"\" compute depths of ROMS vertical levels (Vtransform = 2) \"\"\"\n",
    "\n",
    "    # compute vertical transformation functional\n",
    "    S_rho = (ds.hc * ds.s_rho + ds.Cs_r * ds.h) / (ds.hc + ds.h)\n",
    "    S_w = (ds.hc * ds.s_w + ds.Cs_w * ds.h) / (ds.hc + ds.h)\n",
    "\n",
    "    # compute depth of rho (layers) and w (interfaces) points\n",
    "    z_rho = ds.zeta + (ds.zeta + ds.h) * S_rho\n",
    "    z_w = ds.zeta + (ds.zeta + ds.h) * S_w\n",
    "\n",
    "    # transpose arrays and fill NaNs with a minimal depth\n",
    "    ds['z_rho'] = z_rho.transpose(*('time', 's_rho','yh','xh'),\n",
    "                                  transpose_coords=False).fillna(hmin)\n",
    "\n",
    "    ds['z_w'] = z_w.transpose(*('time', 's_w','yh','xh'),\n",
    "                                  transpose_coords=False).fillna(hmin)\n",
    "\n",
    "    # interpolate depth of levels at U and V points\n",
    "    ds['z_u'] = grid.interp(ds['z_rho'], 'X', boundary='fill')\n",
    "    ds['z_v'] = grid.interp(ds['z_rho'], 'Y', boundary='fill')\n",
    "\n",
    "    # compute layer thickness as difference between interfaces\n",
    "    ds['dz'] = grid.diff(ds['z_w'], 'Z')\n",
    "\n",
    "    # add z_rho and z_w to xarray coordinates\n",
    "    ds = ds.set_coords(['z_rho', 'z_w', 'z_v', 'z_u'])\n",
    "\n",
    "    return ds\n",
    "\n",
    "def add_coords(ds):\n",
    "    \"\"\" set coordinate variables as xarray coordinates\n",
    "\n",
    "        Parameters:\n",
    "            ds (xarray.Dataset): ROMS dataset\n",
    "    \"\"\"\n",
    "    ds = ds.set_coords(['Cs_r', 'Cs_w', 'hc', 'h', 'Vtransform', 'time',\n",
    "                        'lon_rho', 'lon_v', 'lon_u', 'lon_psi',\n",
    "                        'lat_rho', 'lat_v', 'lat_u', 'lat_psi'])\n",
    "    return ds"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Iteration for file  Examples/Grids/norkyst_800m_his.nc4_2019121501-2019121600\n",
      "Ite 1\n",
      "Ite 2\n",
      "Ite 3\n",
      "Ite 4\n",
      "Ite 5\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "One or more of the specified variables cannot be found in this dataset",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mValueError\u001B[0m                                Traceback (most recent call last)",
      "\u001B[0;32m<ipython-input-6-1ada490ca10f>\u001B[0m in \u001B[0;36m<module>\u001B[0;34m\u001B[0m\n\u001B[1;32m     28\u001B[0m     \u001B[0mds\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mrename_dims\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mds\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m     29\u001B[0m     \u001B[0mprint\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m\"Ite 5\"\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m---> 30\u001B[0;31m     \u001B[0mds\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0madd_coords\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mds\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m     31\u001B[0m     \u001B[0mprint\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m\"Ite 6\"\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m     32\u001B[0m     \u001B[0mds\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mds\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mchunk\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m{\u001B[0m\u001B[0;34m'time'\u001B[0m\u001B[0;34m:\u001B[0m \u001B[0;36m1\u001B[0m\u001B[0;34m}\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m<ipython-input-5-20ed2e6df181>\u001B[0m in \u001B[0;36madd_coords\u001B[0;34m(ds)\u001B[0m\n\u001B[1;32m     73\u001B[0m             \u001B[0mds\u001B[0m \u001B[0;34m(\u001B[0m\u001B[0mxarray\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mDataset\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m:\u001B[0m \u001B[0mROMS\u001B[0m \u001B[0mdataset\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m     74\u001B[0m     \"\"\"\n\u001B[0;32m---> 75\u001B[0;31m     ds = ds.set_coords(['Cs_r', 'Cs_w', 'hc', 'h', 'Vtransform', 'time',\n\u001B[0m\u001B[1;32m     76\u001B[0m                         \u001B[0;34m'lon_rho'\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m'lon_v'\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m'lon_u'\u001B[0m\u001B[0;34m,\u001B[0m \u001B[0;34m'lon_psi'\u001B[0m\u001B[0;34m,\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m     77\u001B[0m                         'lat_rho', 'lat_v', 'lat_u', 'lat_psi'])\n",
      "\u001B[0;32m/opt/anaconda3/envs/xesmf2/lib/python3.8/site-packages/xarray/core/dataset.py\u001B[0m in \u001B[0;36mset_coords\u001B[0;34m(self, names)\u001B[0m\n\u001B[1;32m   1551\u001B[0m         \u001B[0;32melse\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m   1552\u001B[0m             \u001B[0mnames\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mlist\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mnames\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m-> 1553\u001B[0;31m         \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0m_assert_all_in_dataset\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mnames\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0m\u001B[1;32m   1554\u001B[0m         \u001B[0mobj\u001B[0m \u001B[0;34m=\u001B[0m \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mcopy\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m   1555\u001B[0m         \u001B[0mobj\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0m_coord_names\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mupdate\u001B[0m\u001B[0;34m(\u001B[0m\u001B[0mnames\u001B[0m\u001B[0;34m)\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;32m/opt/anaconda3/envs/xesmf2/lib/python3.8/site-packages/xarray/core/dataset.py\u001B[0m in \u001B[0;36m_assert_all_in_dataset\u001B[0;34m(self, names, virtual_okay)\u001B[0m\n\u001B[1;32m   3976\u001B[0m             \u001B[0mbad_names\u001B[0m \u001B[0;34m-=\u001B[0m \u001B[0mself\u001B[0m\u001B[0;34m.\u001B[0m\u001B[0mvirtual_variables\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m   3977\u001B[0m         \u001B[0;32mif\u001B[0m \u001B[0mbad_names\u001B[0m\u001B[0;34m:\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[0;32m-> 3978\u001B[0;31m             raise ValueError(\n\u001B[0m\u001B[1;32m   3979\u001B[0m                 \u001B[0;34m\"One or more of the specified variables \"\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n\u001B[1;32m   3980\u001B[0m                 \u001B[0;34m\"cannot be found in this dataset\"\u001B[0m\u001B[0;34m\u001B[0m\u001B[0;34m\u001B[0m\u001B[0m\n",
      "\u001B[0;31mValueError\u001B[0m: One or more of the specified variables cannot be found in this dataset"
     ]
    }
   ],
   "source": [
    "# https://raphaeldussin.medium.com/modern-python-tools-for-the-roms-ocean-model-bfca8642db01\n",
    "# use glob to create a list of files to convert\n",
    "filelist = glob.glob('Examples/Grids/norkyst_800m*.nc*')\n",
    "# path to grid file\n",
    "gridfile = 'Examples/Grids/norkyst_800m_his.nc4_2019121501-2019121600'\n",
    "# output store name\n",
    "store_name = 'CCS2_store'\n",
    "\n",
    "# load the grid\n",
    "ds_grid = xr.open_dataset(gridfile)\n",
    "# drop some unnecessary variables\n",
    "ds_grid = ds_grid.drop_vars(['Cs_r', 'Cs_w', 'hc'])\n",
    "# select interior points for grid variables\n",
    "ds_grid = select_interior(ds_grid)\n",
    "\n",
    "first=True\n",
    "# iterate on files\n",
    "for ncfile in filelist:\n",
    "    print(\"Iteration for file \",ncfile)\n",
    "    ds = xr.open_dataset(ncfile)\n",
    "    print(\"Ite 1\")\n",
    "    ds = remove_all_small_variables(ds)\n",
    "    print(\"Ite 2\")\n",
    "    ds = select_interior(ds)\n",
    "    print(\"Ite 3\")\n",
    "    ds = xr.merge([ds, ds_grid])\n",
    "    print(\"Ite 4\")\n",
    "    ds = rename_dims(ds)\n",
    "    print(\"Ite 5\")\n",
    "    ds = add_coords(ds)\n",
    "    print(\"Ite 6\")\n",
    "    ds = ds.chunk({'time': 1})\n",
    "    print(\"Ite 7\")\n",
    "    print(f\"working on {ncfile}\")\n",
    "    if first:\n",
    "        ds.to_zarr(store_name, consolidated=True, mode='w')\n",
    "    else:\n",
    "        ds.to_zarr(store_name, consolidated=True, mode='a', append_dim='time')\n",
    "    first=False"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "from xgcm import Grid\n",
    "\n",
    "CCS2 = xr.open_zarr(CCS2_store, consolidated=True)\n",
    "\n",
    "# Create xgcm grid object\n",
    "grid_ccs2 = Grid(CCS2, coords={'X': {'center': 'xh', 'outer': 'xq'},\n",
    "                               'Y': {'center': 'yh', 'outer': 'yq'},\n",
    "                               'Z': {'center': 's_rho', 'outer': 's_w'}},\n",
    "                 periodic=False)\n",
    "\n",
    "# Add depths of layers and interfaces to dataset\n",
    "CCS2 = compute_depth_layers(CCS2, grid_ccs2)"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "to_grid_file=\"Examples/Grids/norfjords_160m_grid.nc_A04\"\n",
    "from_grid_file=\"Examples/Grids/norkyst_800m_his.nc4_2019121501-2019121600\"\n",
    "\n",
    "ds=xr.open_dataset(from_grid_file)\n",
    "confM2R = configM2R.Model2romsConfig()\n",
    "\n",
    "print(confM2R.hc,confM2R.vtransform)\n",
    "grd_obj = grd.Grd(\"ROMS\",confM2R=confM2R)\n",
    "\n",
    "grd_obj.create_object(confM2R, from_grid_file)\n",
    "\n",
    "IOverticalGrid.calculateVgrid(grd_obj)\n",
    "\n",
    "ax=plt.figure(figsize=(8,8)).gca(projection=ccrs.PlateCarree())\n",
    "print(grd_obj.lon_rho)\n",
    "print(grd_obj.lat_rho)\n",
    "print(np.shape(grd_obj.z_r))\n",
    "cs= ax.pcolormesh(grd_obj.lon_rho.values,\n",
    "                  grd_obj.lat_rho.values,\n",
    "                  grd_obj.z_r[0],\n",
    "                  cmap=sns.color_palette(\"Spectral_r\", as_cmap=True),\n",
    "                          transform=ccrs.PlateCarree())\n",
    "plt.colorbar(cs, shrink=.4)\n",
    "ax.coastlines()\n",
    "#lons,lats =np.meshgrid(grd_obj.lon_rho,grd_obj.lat_rho)\n",
    "#step=12\n",
    "#ax.scatter(lons[::step,::step],lats[::step,::step],3,c=\"w\")\n",
    "#plt.title(\"test\")\n",
    "plt.show()"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "lons = grd_obj.lon_rho.values.ravel()\n",
    "lats = grd_obj.lat_rho.values.ravel()\n",
    "print(len(lats),np.shape(grd_obj.z_r)[0])\n",
    "z = grd_obj.z_r.ravel().reshape(len(lats),np.shape(grd_obj.z_r)[0])\n",
    "\n",
    "print(np.shape(lons),np.shape(lats),np.shape(z))\n",
    "grid_create_from_coordinates_3d(lons,lats,z)\n",
    "\n",
    "#for lon, lat, in zip(ds.lon_rho, ds.lat_rho):\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}