Xorca_autogenerate_vars_dicts.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "typical_ds_T = xr.open_dataset(\n",
    "    (\n",
    "        \"https://vesg.ipsl.upmc.fr/thredds/dodsC/store/p86mart/\"\n",
    "        \"IPSLCM6/PROD/piControl/SWSTART2/OCE/Analyse/SE/\"\n",
    "        \"SWSTART2_SE_1950_1959_1M_grid_T.nc\"\n",
    "    )\n",
    ")\n",
    "\n",
    "typical_ds_U = xr.open_dataset(\n",
    "    (\n",
    "        \"https://vesg.ipsl.upmc.fr/thredds/dodsC/store/p86mart/\"\n",
    "        \"IPSLCM6/PROD/piControl/SWSTART2/OCE/Analyse/SE/\"\n",
    "        \"SWSTART2_SE_1950_1959_1M_grid_U.nc\"\n",
    "    )\n",
    ")\n",
    "\n",
    "typical_ds_V = xr.open_dataset(\n",
    "    (\n",
    "        \"https://vesg.ipsl.upmc.fr/thredds/dodsC/store/p86mart/\"\n",
    "        \"IPSLCM6/PROD/piControl/SWSTART2/OCE/Analyse/SE/\"\n",
    "        \"SWSTART2_SE_1950_1959_1M_grid_V.nc\"\n",
    "    )\n",
    ")\n",
    "\n",
    "typical_ds_W = xr.open_dataset(\n",
    "    (\n",
    "        \"https://vesg.ipsl.upmc.fr/thredds/dodsC/store/p86mart/\"\n",
    "        \"IPSLCM6/PROD/piControl/SWSTART2/OCE/Analyse/SE/\"\n",
    "        \"SWSTART2_SE_1950_1959_1M_grid_W.nc\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "replace_names_T = {\n",
    "    \"time_counter\": \"t\",\n",
    "    \"olevel\": \"z_c\",\n",
    "    \"y\": \"y_c\",\n",
    "    \"x\": \"x_c\"\n",
    "}\n",
    "\n",
    "replace_names_U = {\n",
    "    \"time_counter\": \"t\",\n",
    "    \"olevel\": \"z_c\",\n",
    "    \"y\": \"y_c\",\n",
    "    \"x\": \"x_r\"\n",
    "}\n",
    "\n",
    "replace_names_V = {\n",
    "    \"time_counter\": \"t\",\n",
    "    \"olevel\": \"z_c\",\n",
    "    \"y\": \"y_r\",\n",
    "    \"x\": \"x_c\"\n",
    "}\n",
    "\n",
    "replace_names_W = {\n",
    "    \"time_counter\": \"t\",\n",
    "    \"olevel\": \"z_l\",\n",
    "    \"y\": \"y_c\",\n",
    "    \"x\": \"x_c\"\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _generate_dims_list(variable, replace_names=None):\n",
    "    return [replace_names[dim_name] for dim_name in variable.dims]\n",
    "\n",
    "\n",
    "# A unit test\n",
    "assert all(\n",
    "    generated == truth\n",
    "    for generated, truth in zip(\n",
    "        [\"t\", \"z_c\", \"y_c\", \"x_c\"],\n",
    "        _generate_dims_list(typical_ds_T[\"thetao\"], replace_names_T)\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# will drop vars with non-t,z,y,x dims (like the bounds dims)\n",
    "typical_variables_T = filter(\n",
    "    lambda v: all(d in replace_names_T for d in typical_ds_T[v].dims),\n",
    "    typical_ds_T.data_vars\n",
    ")\n",
    "\n",
    "# will drop vars with non-t,z,y,x dims (like the bounds dims)\n",
    "typical_variables_U = filter(\n",
    "    lambda v: all(d in replace_names_U for d in typical_ds_U[v].dims),\n",
    "    typical_ds_U.data_vars\n",
    ")\n",
    "\n",
    "# will drop vars with non-t,z,y,x dims (like the bounds dims)\n",
    "typical_variables_V = filter(\n",
    "    lambda v: all(d in replace_names_V for d in typical_ds_V[v].dims),\n",
    "    typical_ds_V.data_vars\n",
    ")\n",
    "\n",
    "# will drop vars with non-t,z,y,x dims (like the bounds dims)\n",
    "typical_variables_W = filter(\n",
    "    lambda v: all(d in replace_names_W for d in typical_ds_W[v].dims),\n",
    "    typical_ds_W.data_vars\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# generate updated vars dict\n",
    "update_orca_variables = {}\n",
    "\n",
    "update_orca_variables.update(\n",
    "    {\n",
    "        v: {\"dims\": _generate_dims_list(typical_ds_T[v], replace_names_T)}\n",
    "        for v in typical_variables_T\n",
    "    }\n",
    ")\n",
    "\n",
    "update_orca_variables.update(\n",
    "    {\n",
    "        v: {\"dims\": _generate_dims_list(typical_ds_U[v], replace_names_U)}\n",
    "        for v in typical_variables_U\n",
    "    }\n",
    ")\n",
    "\n",
    "update_orca_variables.update(\n",
    "    {\n",
    "        v: {\"dims\": _generate_dims_list(typical_ds_V[v], replace_names_V)}\n",
    "        for v in typical_variables_V\n",
    "    }\n",
    ")\n",
    "\n",
    "update_orca_variables.update(\n",
    "    {\n",
    "        v: {\"dims\": _generate_dims_list(typical_ds_W[v], replace_names_W)}\n",
    "        for v in typical_variables_W\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'zos': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'thetao': {'dims': ['t', 'z_c', 'y_c', 'x_c']},\n",
       " 'tos': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'so': {'dims': ['t', 'z_c', 'y_c', 'x_c']},\n",
       " 'sos': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'mldr10_1': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'wfcorr': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'nshfls': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'rsntds': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'hfcorr': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'BLT': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'tinv': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'depti': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'hc300': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'heatc': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'qt_oce': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'qemp_oce': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'qt_ice': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'qemp_ice': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'hflx_rain_cea': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'hflx_evap_cea': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'hflx_snow_cea': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'hflx_cal_cea': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'wfo': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'emp_oce': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'emp_ice': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'runoffs': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'friver': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'calving': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'iceberg': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'iceshelf': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'vfxice': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'vfxsnw': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'vfxsub': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'vfxspr': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'rain': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'snow_ao_cea': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'snow_ai_cea': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'evap_ao_cea': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'subl_ai_cea': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'sosflxdo': {'dims': ['t', 'y_c', 'x_c']},\n",
       " 'e3u': {'dims': ['t', 'z_c', 'y_c', 'x_r']},\n",
       " 'ahu_bbl': {'dims': ['t', 'y_c', 'x_r']},\n",
       " 'e3v': {'dims': ['t', 'z_c', 'y_r', 'x_c']},\n",
       " 'ahv_bbl': {'dims': ['t', 'y_r', 'x_c']},\n",
       " 'NEMO_difvho': {'dims': ['t', 'z_l', 'y_c', 'x_c']},\n",
       " 'NEMO_difvho_noevd': {'dims': ['t', 'z_l', 'y_c', 'x_c']},\n",
       " 'NEMO_difvso': {'dims': ['t', 'z_l', 'y_c', 'x_c']},\n",
       " 'NEMO_difvso_noevd': {'dims': ['t', 'z_l', 'y_c', 'x_c']}}"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "update_orca_variables"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:nemo_cf]",
   "language": "python",
   "name": "conda-env-nemo_cf-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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}