{ "cells": [ { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [ "## import self defined functions\n", "import sys\n", "import os\n", "from os.path import expanduser\n", "home = expanduser(\"~\")\n", "# insert at 1, 0 is the script path (or '' in REPL)\n", "sys.path.insert(1, '/tigress/cw55/local/python_lib')\n" ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "## import most commonly used packages\n", "import numpy as np\n", "import subprocess\n", "import xarray as xr\n", "%matplotlib widget\n", "import matplotlib.pyplot as plt\n", "from mpl_toolkits.mplot3d import axes3d\n", "import time\n", "sstart_time0 = time.time()\n", "\n", "## other packages\n", "import xesmf as xe\n", "# import cftime\n", "# import json\n", "import numba\n", "from numba import njit" ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [], "source": [ "def error_info(string):\n", " \"\"\"By C.Wang\n", " print any input info if it is not empty\n", " and raise exception \"\"\"\n", " if string!=\"\" :\n", " raise Exception(\"Error : \"+string)" ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [], "source": [ "\n", "def satur_specific_humidity_xarray(T):\n", "# es_611 = np.exp(17.67*(T-273.16)/(T-29.65))\n", "# q_s = (0.622*611)*es_611/(T.plev - es_611)\n", " q_s = (0.622*611)*np.exp(17.67*(T-273.16)/(T-29.65))/T.plev\n", " return q_s\n", "\n", "\n", "@njit(parallel=True)\n", "def satur_specific_humidity_TPLL_fast(T_TPLL,plev):\n", " \"\"\"\n", " ALL inputs have to be numpy array\n", " \"\"\"\n", " qs = np.empty_like(T_TPLL,dtype = np.float32)\n", " for pi in numba.prange(T_TPLL.shape[1]):\n", " Kp = 0.622*611/plev[pi]\n", " for ti in numba.prange(T_TPLL.shape[0]):\n", " for li in numba.prange(T_TPLL.shape[2]):\n", " for lj in numba.prange(T_TPLL.shape[3]):\n", " qs[ti,pi,li,lj] = Kp*np.exp(17.67*(T_TPLL[ti,pi,li,lj]-273.16)/(T_TPLL[ti,pi,li,lj]-29.65))\n", " return qs\n", "\n", "def global_mean_weight_lat(ds):\n", " data_ori = ds.copy()\n", " lat = data_ori.coords['lat'] # readin lat\n", " # global mean\n", " # compute cos(lat) as a weight function\n", " weight_lat = np.cos(np.deg2rad(lat))/np.mean(np.cos(np.deg2rad(lat)))\n", " data_gm = data_ori.mean(dim=['lon'])*weight_lat\n", " data_gm = data_gm.mean(dim=['lat'])\n", " return data_gm\n", "\n", "@njit(parallel=True)\n", "def alb_diff_pert_mon_cont_12mon_TLL(pert_rsus,pert_rsds,cont_rsus,cont_rsds):\n", " \"\"\"\n", " ALL inputs have to be numpy array\n", " \"\"\"\n", " alb_diff = np.empty_like(pert_rsus,dtype = np.float32)\n", " alb_cont_mon = np.empty_like(cont_rsus[0,:,:])\n", " for mon in numba.prange(cont_rsus.shape[0]):\n", " alb_cont_mon = cont_rsus[mon,:,:] / cont_rsds[mon,:,:] \n", " for i in numba.prange(pert_rsus.shape[0]/12):\n", " alb_diff[i*12+mon,:,:] = pert_rsus[i*12+mon,:,:] / pert_rsds[i*12+mon,:,:] - alb_cont_mon\n", " return alb_diff\n", "\n", "@njit(parallel=True)\n", "def diff_pert_mon_cont_12mon_TLL(pert_TLL,cont_TLL):\n", " \"\"\"\n", " ALL inputs have to be numpy array\n", " \"\"\"\n", " diff = np.empty_like(pert_TLL,dtype = np.float32)\n", " for mon in numba.prange(cont_TLL.shape[0]):\n", " for i in numba.prange(pert_TLL.shape[0]/12):\n", " diff[i*12+mon,:,:] = pert_TLL[i*12+mon,:,:] - cont_TLL[mon,:,:]\n", " return diff\n", "\n", "@njit(parallel=True)\n", "def diff_pert_mon_cont_12mon_TPLL(pert_TPLL,cont_TPLL):\n", " \"\"\"\n", " ALL inputs have to be numpy array\n", " \"\"\"\n", " diff = np.empty_like(pert_TPLL,dtype = np.float32)\n", " for mon in numba.prange(cont_TPLL.shape[0]):\n", " for i in numba.prange(pert_TPLL.shape[0]/12):\n", " diff[i*12+mon,:,:,:] = pert_TPLL[i*12+mon,:,:,:] - cont_TPLL[mon,:,:,:]\n", " return diff\n", "\n", "@njit(parallel=True)\n", "def omega_wv_fast_compile(qs_pert_TPLL,qs_cont_TPLL,hus_pert_TPLL,hus_cont_TPLL,ta_anom_TPLL):\n", " \"\"\"\n", " ALL inputs have to be numpy array\n", " \"\"\"\n", " omega_wv = np.empty_like(qs_pert_TPLL,dtype = np.float32)\n", " for mon in numba.prange(qs_cont_TPLL.shape[0]):\n", " for i in numba.prange(qs_pert_TPLL.shape[0]/12):\n", " omega_wv[i*12+mon,:,:,:] = (qs_cont_TPLL[mon,:,:,:] / hus_pert_TPLL[i*12+mon,:,:,:]) \\\n", " * ta_anom_TPLL[i*12+mon,:,:,:] / (qs_pert_TPLL[i*12+mon,:,:,:]-qs_cont_TPLL[mon,:,:,:]) \\\n", " * (hus_pert_TPLL[i*12+mon,:,:,:] - hus_cont_TPLL[mon,:,:,:])\n", " return omega_wv\n", "\n", "def omega_wv_xarray(qs_pert_TPLL,qs_cont_TPLL,hus_pert_TPLL,hus_cont_TPLL,ta_anom_TPLL):\n", " \"\"\"\n", " ALL inputs have to be xarray array\n", " \"\"\"\n", "# print(hus_pert_TPLL.groupby('time.month'))\n", "# print(hus_cont_TPLL.dims)\n", " omega_wv = (qs_cont_TPLL / hus_pert_TPLL.groupby('time.month')) \\\n", " / (qs_pert_TPLL.groupby('time.month') - qs_cont_TPLL)*ta_anom_TPLL \\\n", " * (hus_pert_TPLL.groupby('time.month') - hus_cont_TPLL)\n", " return omega_wv\n", "@njit(parallel=True)\n", "def RK_compute_TLL_fast(var_mon, rk_mon_cli):\n", " \"\"\"\n", " ALL inputs have to be numpy array\n", " \"\"\"\n", " dR_TLL = np.empty_like(var_mon,dtype = np.float32)\n", " for mon in numba.prange(rk_mon_cli.shape[0]):\n", " for i in numba.prange(var_mon.shape[0]/12):\n", " dR_TLL[i*12+mon,:,:] = rk_mon_cli[mon,:,:] * var_mon[i*12+mon,:,:]\n", " return dR_TLL\n", "\n", "@njit(parallel=True)\n", "def RK_compute_TPLL_fast(var_mon, rk_mon_cli):\n", " \"\"\"\n", " ALL inputs have to be numpy array\n", " \"\"\"\n", " plev_weight = [0.425, 0.75 , 1.125, 1.25 , 1., 1. , 1.,\n", " 0.75 , 0.5 ,0.5 , 0.5 , 0.4 , 0.25 ,\n", " 0.2 , 0.15 , 0.1 , 0.075, 0.045,0.005]\n", " dR_TLL = np.zeros_like(var_mon[:,0,:,:],dtype = np.float32)\n", " for li in numba.prange(var_mon.shape[2]):\n", " for lj in numba.prange(var_mon.shape[3]):\n", " for mon in numba.prange(rk_mon_cli.shape[0]):\n", " for i in range(var_mon.shape[0]/12):\n", " for pi in range(var_mon.shape[1]):\n", " if np.isnan(rk_mon_cli[mon,pi,li,lj]) or np.isnan(var_mon[i*12+mon,pi,li,lj]): \n", " continue\n", " dR_TLL[i*12+mon,li,lj] += rk_mon_cli[mon,pi,li,lj] \\\n", " * var_mon[i*12+mon,pi,li,lj] \\\n", " * plev_weight[pi]\n", " return dR_TLL\n", "\n", "def RK_compute_suf(var, rk):\n", " dR = var.groupby('time.month') *rk\n", " return dR\n", "def RK_compute_TPLL(var, rk):\n", " tmp = var.groupby('time.month')*rk\n", " \n", " # weighted by pressure and sum\n", "# var = ta_anom\n", "# plev_weight = var.plev.copy().values\n", "# plev_weight[0] = (1.01e5 - var.plev[1])/20000\n", "# plev_weight[-1] = (var.plev[-1].values-0)/20000\n", "# plev_weight[1:-1] = (var.plev[:-2].values - var.plev[2:].values)/20000\n", "\n", " # only correct for cmip6 data\n", " plev_weight = [0.425, 0.75 , 1.125, 1.25 , 1., 1. , 1., \\\n", " 0.75 , 0.5 ,0.5 , 0.5 , 0.4 , 0.25 ,\\\n", " 0.2 , 0.15 , 0.1 , 0.075, 0.045,0.005]\n", " plev_weight = xr.DataArray(plev_weight,var.plev.coords,dims=['plev'])\n", " dR = tmp *plev_weight\n", " dR = dR.sum(dim='plev',skipna=True)\n", " return dR\n", "\n", "# run dummy data to compile the jit functions\n", "# dummy_TPLL = np.random.rand(120,19,100,2)\n", "# dummy_TPLL12 = np.random.rand(12,19,100,2)\n", "# dummy_TLL = np.random.rand(120,100,2)\n", "# dummy_TLL12 = np.random.rand(12,100,2)\n", "# dummy_plev = np.random.rand(19)\n", "# tmp_compile = satur_specific_humidity_TPLL_fast(dummy_TPLL, dummy_plev)\n", "# tmp_compile = diff_pert_mon_cont_12mon_TPLL(dummy_TPLL, dummy_TPLL12)\n", "# tmp_compile = diff_pert_mon_cont_12mon_TLL(dummy_TLL, dummy_TLL12)\n", "# tmp_compile = alb_diff_pert_mon_cont_12mon_TLL(dummy_TLL, dummy_TLL, dummy_TLL12,dummy_TLL12)\n", "# tmp_compile = omega_wv_fast_compile(dummy_TPLL,dummy_TPLL12,dummy_TPLL,dummy_TPLL12,dummy_TPLL)\n", "# tmp_compile = RK_compute_TLL_fast (dummy_TLL, dummy_TLL12)\n", "# tmp_compile = RK_compute_TPLL_fast(dummy_TPLL, dummy_TPLL12)\n", "# print(\"Functions finished compile!\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "# Calculate the feedbacks in single model abrupt-4xCO2 experiment by using the radiatvie kernel \n", "ref: Soden, et.al., (2008)\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Prepare the data \n", " - __ta(3D), hus(3D), ts(2D)__: \n", "> in both perturbation (abrupt-4xCO2) and control (piControl) \\\n", "> compute the perturbation for temperature, water vapor \\\n", "> $dt = T_{per} - T_{con} $\n", " - __rlut rsdt rsut__ : \n", "> in perturbation \\\n", "> for the dR\n", " - __rsutcs rlutcs__: \n", "> clear-sky data needed in both perturbation (abrupt-4xCO2) and control (piControl) \\\n", "> for the dR^0, as well as the direct forcing in clear sky (and therefore we can infer the direct forcing in all sky ) \\\n", "> $D = D^0/1.16$" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Compute the cloud feedback\n", "$\n", "\\begin{align}\n", "dR &= D + dR_wv + dR_a + dR_T + dR_c \\\\\n", "dR^0 &= D^0 + dR^0_{wv} + dR^0_a + dR^0_T \\\\\n", "D^0 &= dR^0 - (dR^0_{wv} + dR^0_a + dR^0_T)\\\\\n", "\\Rightarrow dR_c &= dR - D - (dR_{wv} + dR_a + dR_T) \\\\\n", "with\\quad D &= D^0/1.16 \n", "\\end{align}\n", "$" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## convert the hus for RK and calculate the $dR_{wv}$\n", "$\n", "\\begin{align}\n", "dR_{wv} &= K^w dq \\\\\n", " &= \\underline{(K^w * \\zeta)} \\quad \\underline{(dq / \\zeta)}\\\\\n", " &= K^{\\omega} \\omega \\\\\n", "\\zeta &= \\frac{q}{q_s} \\frac{dq_s}{dT} \\\\\n", "\\Rightarrow \\omega &= dq / \\zeta \\\\\n", " &= \\frac{q_s}{q} \\frac{dT}{dq_s} dq\n", "\\end{align}\n", "$\n", "\n", "Since the water vapor RK is store with factor $\\zeta$ and the unit $W/m^{-2} K^{-1}$, the convertion for specific humidity is needed for the $dR_{wv}$ calculation." ] }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [], "source": [ "experiments = [\"piControl\",\"abrupt-4xCO2\"]\n", "\n", "var_list = \"ta hus ts rlut rsdt rsut rlutcs rsutcs rsus rsds\".split()\n", "res_hori = '2x2.5'\n", "## run the regrid scripts before this one!\n", "dvc_info = '2020051522'" ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ ">>> start GFDL-CM4\n", ">>> reading | ta\n", ">>> reading | hus\n", ">>> reading | ts\n", ">>> reading | rlut\n", ">>> reading | rsdt\n", ">>> reading | rsut\n", ">>> reading | rlutcs\n", ">>> reading | rsutcs\n", ">>> reading | rsus\n", ">>> reading | rsds\n" ] } ], "source": [ " # def decompose_RK(model)\n", " # process for each model\n", " model = 'GFDL-CM4'\n", " # model = 'CESM2'\n", " var_cont = {}\n", " var_pert = {}\n", " model_ava = []\n", " # read in all var needed\n", " print(\">>> start \"+model)\n", " for var in var_list:\n", " print(\">>> reading |{0:>7s}\".format(var) )\n", " try:\n", " dirpath = \"/tigress/cw55/data/CMIP6_post/CMIP6_post_large_scratch/piControl/{3:}/{0:}/{1:}.*0001-0450.ltm.nc.{2:}.2x2.5\".format(var,var,dvc_info,model)\n", " tmp = subprocess.run([\"source ~/.bash_env; ls \"+dirpath], shell=True, capture_output=True)\n", " error_info(tmp.stderr.decode(\"utf-8\"))\n", " except:\n", " print(\"model {0:} does't have variable {1:} in control experiment\".format(model,var))\n", "# return\n", " filelist = tmp.stdout.decode(\"utf-8\")[:-1]\n", " # read variables\n", " with xr.open_dataset(filelist) as ds:\n", "# print(ds)\n", " var_cont[var] = ds.isel(model=0).load()\n", "\n", " try:\n", " dirpath = \"/tigress/cw55/data/CMIP6_post/CMIP6_post_large_scratch/abrupt-4xCO2/{3:}/{0:}/{1:}.*.nc.{2:}.2x2.5\".format(var,var,dvc_info,model)\n", " tmp = subprocess.run([\"source ~/.bash_env; ls \"+dirpath], shell=True, capture_output=True)\n", " error_info(tmp.stderr.decode(\"utf-8\"))\n", " except:\n", " print(\"model {0:} does't have variable {1:} in perturbe experiment\".format(model,var))\n", "# return\n", " filelist = tmp.stdout.decode(\"utf-8\")[:-1]\n", " # read variables\n", " with xr.open_dataset(filelist) as ds:\n", " var_pert[var] = ds.isel(model=0).load()" ] }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "/tigress/cw55/data/Radiative_kernel/kernels_TOA_GFDL_CMIP6-standard.nc\n", "\n" ] } ], "source": [ "## import kernel file and compute: dR_wv, dR_T, dR_Ts, dR_alb + clear sky\n", "rk_source = 'GFDL'\n", "# kernel file is from Brian Soden\n", "dirpath = \"/tigress/cw55/data/Radiative_kernel/kernels_TOA_\"+rk_source+\"_CMIP6-standard.nc\"\n", "tmp = subprocess.run([\"source ~/.bash_env; ls \"+dirpath], shell=True, capture_output=True)\n", "file_rk = tmp.stdout.decode(\"utf-8\")\n", "print(file_rk)\n", "f_RK = xr.open_dataset(file_rk[:-1],decode_times=False) \n", "f_RK=f_RK.rename({'time': 'month'})\n", "f_RK.coords['month']= np.arange(1,13,1)\n", "f_RK.coords['plev']= f_RK.coords['plev']*100" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## time benchmark (test with GFDL-CM4)\n", "- compile and parallel with numba ~ 10s\n", "- pure xarray ~ 100s" ] }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "CPU times: user 2min 2s, sys: 5.13 s, total: 2min 7s\n", "Wall time: 17.1 s\n" ] }, { "data": { "text/plain": [ "1.3000169" ] }, "execution_count": 8, "metadata": {}, "output_type": "execute_result" } ], "source": [ "%%time\n", "# numba version (fast, will be faster 2x if run it after successfully compiled)\n", "## convert the hus for RK and calculate the dR_wv\n", "qs_cont = satur_specific_humidity_TPLL_fast(var_cont['ta'].ta.values,var_cont['ta'].plev.values)\n", "qs_pert = satur_specific_humidity_TPLL_fast(var_pert['ta'].ta.values,var_pert['ta'].plev.values)\n", "ta_anom = diff_pert_mon_cont_12mon_TPLL(var_pert['ta'].ta.values,var_cont['ta'].ta.values)\n", "omega_wv = omega_wv_fast_compile (qs_pert, qs_cont,\n", " var_pert['hus'].hus.values,var_cont['hus'].hus.values,\n", " ta_anom)\n", "ts_anom = diff_pert_mon_cont_12mon_TLL(var_pert['ts'].ts.values, var_cont['ts'].ts.values)\n", "alb_anom = alb_diff_pert_mon_cont_12mon_TLL(var_pert['rsus'].rsus.values, var_pert['rsds'].rsds.values, \\\n", " var_cont['rsus'].rsus.values, var_cont['rsds'].rsds.values)\n", "dR_sw = diff_pert_mon_cont_12mon_TLL((var_pert['rsdt'].rsdt.values-var_pert['rsut'].rsut.values),\n", " (var_cont['rsdt'].rsdt.values-var_cont['rsut'].rsut.values) )\n", "dR_lw = diff_pert_mon_cont_12mon_TLL((-var_pert['rlut'].rlut.values),\n", " (-var_cont['rlut'].rlut.values) )\n", "dRcs_sw = diff_pert_mon_cont_12mon_TLL((var_pert['rsdt'].rsdt.values-var_pert['rsutcs'].rsutcs.values),\n", " (var_cont['rsdt'].rsdt.values-var_cont['rsutcs'].rsutcs.values) )\n", "dRcs_lw = diff_pert_mon_cont_12mon_TLL((-var_pert['rlutcs'].rlutcs.values),\n", " (-var_cont['rlutcs'].rlutcs.values) )\n", "dR_wv_lw = RK_compute_TPLL_fast(omega_wv,f_RK.lw_q.values)\n", "dR_wv_sw = RK_compute_TPLL_fast(omega_wv,f_RK.sw_q.values)\n", "dR_wvcs_lw = RK_compute_TPLL_fast(omega_wv,f_RK.lwclr_q.values)\n", "dR_wvcs_sw = RK_compute_TPLL_fast(omega_wv,f_RK.swclr_q.values)\n", "dR_Ta = RK_compute_TPLL_fast(ta_anom,f_RK.lw_ta.values)\n", "dR_Tacs = RK_compute_TPLL_fast(ta_anom,f_RK.lwclr_ta.values)\n", "dR_Ts = RK_compute_TLL_fast (ts_anom,f_RK.lw_ts.values)\n", "dR_Tscs = RK_compute_TLL_fast (ts_anom,f_RK.lwclr_ts.values)\n", "dR_alb = RK_compute_TLL_fast (alb_anom*100,f_RK.sw_alb.values)\n", "dR_albcs = RK_compute_TLL_fast (alb_anom*100,f_RK.swclr_alb.values)\n", "\n", "## dR due to cloud change\n", "Dcs_lw = dRcs_lw - (dR_Tacs - dR_Tscs - dR_wvcs_lw)\n", "Dcs_sw = dRcs_sw - (- dR_albcs - dR_wvcs_sw)\n", "D_lw = Dcs_lw / 1.16\n", "D_sw = Dcs_sw / 1.16\n", "dR_c_lw = dR_lw - D_lw - (dR_Ta - dR_Ts - dR_wv_lw)\n", "dR_c_sw = dR_sw - D_sw - (- dR_alb - dR_wv_sw)\n", "np.nanmean(dR_c_lw+dR_c_sw)" ] }, { "cell_type": "code", "execution_count": 9, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "\n", "array(1.30001656)\n", "Coordinates:\n", " model