{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2e336d6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Import libraries\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "import scipy\n",
    "from scipy import stats\n",
    "import scipy.special as sp\n",
    "import matplotlib.pyplot as plt\n",
    "from glob import glob\n",
    "import os\n",
    "import warnings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6130b8b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Set paths\n",
    "dir_data = '/Users/jzhang/Documents/Data/'\n",
    "\n",
    "#Load data\n",
    "cFlx = xr.open_mfdataset(sorted(glob(dir_data+'CERES/FluxByCldTyp/CERES_FluxByCldTyp-MON_Terra-Aqua-MODIS_Ed4.1_*nc'))).sel(time=slice(np.datetime64('2003-01-01'),np.datetime64('2022-12-31')))\n",
    "ebaf = xr.open_mfdataset(sorted(glob(dir_data+'CERES/EBAFed42/CERES_EBAF_Ed4.2_Subset_200301-202212.nc'))).sel(time=slice(np.datetime64('2003-01-01'),np.datetime64('2022-12-31')))\n",
    "ocean = xr.open_dataset(dir_data+'CERES/SYN1deg/CERES_SYN1deg-Month_Terra-Aqua-NOAA20_Ed4.2_Subset_200003-202507.nc').sel(time=slice(np.datetime64('2003-01-01'),np.datetime64('2022-12-31')))\n",
    "oce_frac = ocean['aux_ocean_mon']/100.\n",
    "\n",
    "cldtypes = ['Cu','Sc','St','Ac','As','Ns','Ci','Cs','Cb']\n",
    "dFlx = {} #Each cloud type aggregated separately\n",
    "for cldtype in cldtypes:\n",
    "    dFlx[cldtype] = xr.open_dataset(dir_data+'CERES/FluxByCldTyp/FBCT_cldtype_%s_mon.nc' % cldtype).sel(time=slice(np.datetime64('2003-01-01'),np.datetime64('2022-12-31')))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e673245",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Calculating:\n",
      "...all...\n",
      "...clr...\n",
      "...Cu...\n",
      "...Sc...\n",
      "...St...\n",
      "...Ac...\n",
      "...As...\n",
      "...Ns...\n",
      "...Ci...\n",
      "...Cs...\n",
      "...Cb...\n",
      "Done!\n",
      "\n",
      "Saving files:\n",
      "...all...\n",
      "...cld...\n",
      "...ocean...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "rm: /Users/jzhang/Documents/Data/CERES/FluxByCldTyp/FBCT_decomposition_ocean.nc: No such file or directory\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "...clr...\n",
      "...Cu...\n",
      "...Sc...\n",
      "...St...\n",
      "...Ac...\n",
      "...As...\n",
      "...Ns...\n",
      "...Ci...\n",
      "...Cs...\n",
      "...Cb...\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "#Set up master array to store decompositions\n",
    "xR = xr.Dataset() \n",
    "xR['time'] = cFlx.time\n",
    "xR['lat'] = cFlx.lat\n",
    "xR['lon'] = cFlx.lon\n",
    "xR['month'] = np.arange(1,13)\n",
    "xR['month'].attrs = {'units' : '1', 'long_name' : 'Month of the year'}\n",
    "\n",
    "###All-sky total\n",
    "\n",
    "R = cFlx['toa_sw_all_mon']\n",
    "R = R.where(R>0,0)\n",
    "\n",
    "xR['R'] = R\n",
    "xR['R'].attrs = {'units' : 'W m-2', 'long_name' : 'All-sky reflection'}\n",
    "\n",
    "S = cFlx['toa_solar_all_mon']\n",
    "\n",
    "xR['S'] = S\n",
    "xR['S'].attrs = {'units' : 'W m-2', 'long_name' : 'Insolation'}\n",
    "\n",
    "###Clear-sky\n",
    "\n",
    "R_clr = cFlx['toa_sw_clr_mon']\n",
    "R_clr = R_clr.where(R_clr>0,0)\n",
    "\n",
    "xR['R_clr'] = R_clr\n",
    "xR['R_clr'].attrs = {'units' : 'W m-2', 'long_name' : 'Clear-sky reflection'}\n",
    "\n",
    "### CRE\n",
    "\n",
    "R_cld = R - R_clr\n",
    "\n",
    "xR['R_cld'] = R_cld\n",
    "xR['R_cld'].attrs = {'units' : 'W m-2', 'long_name' : 'Cloudy-sky reflection'}\n",
    "\n",
    "###Ocean-fraction\n",
    "\n",
    "xR['O_frac'] = oce_frac\n",
    "xR['O_frac'].attrs = {'units' : '1', 'long_name' : 'Ocean fraction'}\n",
    "\n",
    "#Use EBAF values to help with decomposition\n",
    "#Note S is from FBCT, F_tow,down is from EBAF to be consistent with the use of surface fluxes from EBAF when calculate relative measures\n",
    "\n",
    "rsdt = ebaf['solar_mon']\n",
    "rsut = ebaf['toa_sw_clr_t_mon']\n",
    "rsds = ebaf['sfc_sw_down_clr_t_mon']\n",
    "rsus = ebaf['sfc_sw_up_clr_t_mon']\n",
    "\n",
    "R_aer = S*(rsdt*rsut-rsds*rsus)/(rsdt**2-rsus**2)\n",
    "R_aer = R_aer.where(R_aer>0,0)\n",
    "R_aer = R_aer.where(R_aer<R_clr,R_clr)\n",
    "\n",
    "xR['R_aer'] = R_aer\n",
    "xR['R_aer'].attrs = {'units' : 'W m-2', 'long_name' : 'Clear-sky atmospheric reflection'}\n",
    "\n",
    "xR['R_sfc'] = xR['R_clr']-xR['R_aer']\n",
    "xR['R_sfc'].attrs = {'units' : 'W m-2', 'long_name' : 'Clear-sky surface reflection'}\n",
    "\n",
    "###Cloud types\n",
    "\n",
    "for i in range(len(cldtypes)):\n",
    "    ct = cldtypes[i]\n",
    "\n",
    "    R_cld = dFlx[ct]['toa_sw_cldtyp_mon']\n",
    "    C = dFlx[ct]['cldarea_cldtyp_mon']/100\n",
    "\n",
    "    xR['R_%s' % ct] = C*(R_cld-R_clr) #Positive up\n",
    "    xR['R_%s' % ct].attrs = {'units' : 'W m-2', 'long_name' : 'Cloud radiative effect for %s scenes' % ct}\n",
    "\n",
    "\n",
    "\n",
    "# Calculate climatologies\n",
    "\n",
    "Feb_wts = [28]+4*[29,28,28,28]+[29,28,28] #leap years\n",
    "\n",
    "print('Calculating:')\n",
    "\n",
    "#\n",
    "###All-sky reflection and insolation\n",
    "#\n",
    "print('...all...')\n",
    "xR['Rc'] = (['month','lat','lon'],np.nan*np.ones((12,180,360)))\n",
    "xR['Rc'].attrs = {'units' : 'W m-2', 'long_name' : 'Climatological (2003-01 to 2022-12) value of R'}\n",
    "xR['Ra'] = (['time','lat','lon'],np.nan*np.ones(xR['R'].shape))\n",
    "xR['Ra'].attrs = {'units' : 'W m-2', 'long_name' : 'Deseasonalized anomaly of R'}\n",
    "\n",
    "xR['Sc'] = (['month','lat','lon'],np.nan*np.ones((12,180,360)))\n",
    "xR['Sc'].attrs = {'units' : 'W m-2', 'long_name' : 'Climatological (2003-01 to 2022-12) value of S'}\n",
    "xR['Sa'] = (['time','lat','lon'],np.nan*np.ones(xR['R'].shape))\n",
    "xR['Sa'].attrs = {'units' : 'W m-2', 'long_name' : 'Deseasonalized anomaly of S'}\n",
    "\n",
    "xR['Rc_cld'] = (['month','lat','lon'],np.nan*np.ones((12,180,360)))\n",
    "xR['Rc_cld'].attrs = {'units' : 'W m-2', 'long_name' : 'Climatological (2003-01 to 2022-12) value of R_cld'}\n",
    "xR['Ra_cld'] = (['time','lat','lon'],np.nan*np.ones(xR['R'].shape))\n",
    "xR['Ra_cld'].attrs = {'units' : 'W m-2', 'long_name' : 'Deseasonalized anomaly of R_cld'}\n",
    "\n",
    "xR['Oc_frac'] = (['month','lat','lon'],np.nan*np.ones((12,180,360)))\n",
    "xR['Oc_frac'].attrs = {'units' : '1', 'long_name' : 'Climatological (2003-01 to 2022-12) value of O_frac'}\n",
    "xR['Oa_frac'] = (['time','lat','lon'],np.nan*np.ones(xR['R'].shape))\n",
    "xR['Oa_frac'].attrs = {'units' : '1', 'long_name' : 'Deseasonalized anomaly of O_frac'}\n",
    "\n",
    "for m in range(1,13):\n",
    "\n",
    "    R = xR['R'].sel(time=slice(np.datetime64('2003-01'),np.datetime64('2022-12')))\n",
    "    if m ==2: xR['Rc'][m-1] = np.average(R[R.time.dt.month==m],weights=Feb_wts,axis=0)\n",
    "    else: xR['Rc'][m-1] = R[R.time.dt.month==m].mean(axis=0)\n",
    "    xR['Ra'][xR.time.dt.month==m] = xR['R'][xR.time.dt.month==m].values-xR['Rc'][m-1].values\n",
    "\n",
    "    SS = xR['S'].sel(time=slice(np.datetime64('2003-01'),np.datetime64('2022-12')))\n",
    "    if m ==2: xR['Sc'][m-1] = np.average(SS[SS.time.dt.month==m],weights=Feb_wts,axis=0)\n",
    "    else: xR['Sc'][m-1] = SS[SS.time.dt.month==m].mean(axis=0)\n",
    "    xR['Sa'][xR.time.dt.month==m] = xR['S'][xR.time.dt.month==m].values-xR['Sc'][m-1].values\n",
    "\n",
    "    RR = xR['R_cld'].sel(time=slice(np.datetime64('2003-01'),np.datetime64('2022-12')))\n",
    "    if m ==2: xR['Rc_cld'][m-1] = np.average(RR[RR.time.dt.month==m],weights=Feb_wts,axis=0)\n",
    "    else: xR['Rc_cld'][m-1] = RR[RR.time.dt.month==m].mean(axis=0)\n",
    "    xR['Ra_cld'][xR.time.dt.month==m] = xR['R_cld'][xR.time.dt.month==m].values-xR['Rc_cld'][m-1].values\n",
    "\n",
    "    OO = xR['O_frac'].sel(time=slice(np.datetime64('2003-01'),np.datetime64('2022-12')))\n",
    "    if m ==2: xR['Oc_frac'][m-1] = np.average(OO[OO.time.dt.month==m],weights=Feb_wts,axis=0)\n",
    "    else: xR['Oc_frac'][m-1] = OO[OO.time.dt.month==m].mean(axis=0)\n",
    "    xR['Oa_frac'][xR.time.dt.month==m] = xR['O_frac'][xR.time.dt.month==m].values-xR['Oc_frac'][m-1].values\n",
    "\n",
    "print('...clr...')\n",
    "\n",
    "#Climatology for S*Aaer and S*Asfc\n",
    "for loc in ['clr','aer','sfc']:\n",
    "    \n",
    "    with warnings.catch_warnings():\n",
    "        warnings.simplefilter(\"ignore\")\n",
    "    \n",
    "        xR['Rc_%s' % loc] = (['month','lat','lon'],np.nan*np.ones((12,180,360)))\n",
    "        xR['Rc_%s' % loc].attrs = {'units' : 'W m-2', 'long_name' : 'Climatological (2003-01 to 2022-12) value of R%s' % loc}\n",
    "        xR['Ra_%s' % loc] = (['time','lat','lon'],np.nan*np.ones(xR['R'].shape))\n",
    "        xR['Ra_%s' % loc].attrs = {'units' : 'W m-2', 'long_name' : 'Deseasonalized anomaly of R%s' % loc}\n",
    "\n",
    "        for m in range(1,13):\n",
    "\n",
    "            R = xR['R_%s' % (loc)].sel(time=slice(np.datetime64('2003-01'),np.datetime64('2022-12')))\n",
    "            if m ==2: xR['Rc_%s' % (loc)][m-1] = np.average(R[R.time.dt.month==m],weights=Feb_wts,axis=0)\n",
    "            else: xR['Rc_%s' % (loc)][m-1] = R[R.time.dt.month==m].mean(axis=0)\n",
    "            xR['Ra_%s' % (loc)][xR.time.dt.month==m] = xR['R_%s' % (loc)][xR.time.dt.month==m].values-xR['Rc_%s' % (loc)][m-1].values\n",
    "\n",
    "#Climatology for R by cloud scene\n",
    "for ct in cldtypes:\n",
    "    print('...'+ct+'...')\n",
    "\n",
    "    xR['Rc_%s' % ct] = (['month','lat','lon'],np.nan*np.ones((12,180,360)))\n",
    "    xR['Rc_%s' % ct].attrs = {'units' : 'W m-2', 'long_name' : 'Climatological (2003-01 to 2022-12) value of R for %s scenes' % ct}\n",
    "    xR['Ra_%s' % ct] = (['time','lat','lon'],np.nan*np.ones(xR['R'].shape))\n",
    "    xR['Ra_%s' % ct].attrs = {'units' : 'W m-2', 'long_name' : 'Deseasonalized anomaly of reflection for %s scenes' % ct}\n",
    "\n",
    "    for m in range(1,13):\n",
    "\n",
    "        Ratm = xR['R_%s' % ct].sel(time=slice(np.datetime64('2003-01'),np.datetime64('2022-12')))\n",
    "        if m ==2: xR['Rc_%s' % ct][m-1] = np.average(Ratm[Ratm.time.dt.month==m],weights=Feb_wts,axis=0)\n",
    "        else: xR['Rc_%s' % ct][m-1] = Ratm[Ratm.time.dt.month==m].mean(axis=0)\n",
    "        xR['Ra_%s' % ct][xR.time.dt.month==m] = xR['R_%s' % ct][xR.time.dt.month==m].values-xR['Rc_%s' % ct][m-1].values\n",
    "\n",
    "print('Done!')\n",
    "print()\n",
    "\n",
    "\n",
    "print('Saving files:')\n",
    "\n",
    "#All-sky values\n",
    "print('...all...')\n",
    "tmp = xr.Dataset()\n",
    "tmp['time'] = cFlx.time\n",
    "tmp['lat'] = cFlx.lat\n",
    "tmp['lon'] = cFlx.lon\n",
    "tmp['month'] = np.arange(1,13)\n",
    "tmp['month'].attrs = {'units' : '1', 'long_name' : 'month of the year'}\n",
    "\n",
    "for var in ['R','Rc','Ra','S','Sc','Sa']:\n",
    "    tmp[var] = xR[var]\n",
    "\n",
    "filename = dir_data+'CERES/FluxByCldTyp/FBCT_decomposition_all.nc'\n",
    "os.system('rm %s' % filename)\n",
    "tmp.to_netcdf(path=filename,mode='w')\n",
    "\n",
    "#Cloudy-sky values\n",
    "print('...cld...')\n",
    "tmp = xr.Dataset()\n",
    "tmp['time'] = cFlx.time\n",
    "tmp['lat'] = cFlx.lat\n",
    "tmp['lon'] = cFlx.lon\n",
    "tmp['month'] = np.arange(1,13)\n",
    "tmp['month'].attrs = {'units' : '1', 'long_name' : 'month of the year'}\n",
    "\n",
    "for var in ['R_cld','Rc_cld','Ra_cld']:\n",
    "    tmp[var] = xR[var]\n",
    "\n",
    "filename = dir_data+'CERES/FluxByCldTyp/FBCT_decomposition_cld.nc'\n",
    "os.system('rm %s' % filename)\n",
    "tmp.to_netcdf(path=filename,mode='w')\n",
    "\n",
    "#Ocean-fraction values\n",
    "print('...ocean...')\n",
    "tmp = xr.Dataset()\n",
    "tmp['time'] = cFlx.time\n",
    "tmp['lat'] = cFlx.lat\n",
    "tmp['lon'] = cFlx.lon\n",
    "tmp['month'] = np.arange(1,13)\n",
    "tmp['month'].attrs = {'units' : '1', 'long_name' : 'month of the year'}\n",
    "\n",
    "for var in ['O_frac','Oc_frac','Oa_frac']:\n",
    "    tmp[var] = xR[var]\n",
    "\n",
    "filename = dir_data+'CERES/FluxByCldTyp/FBCT_decomposition_ocean.nc'\n",
    "os.system('rm %s' % filename)\n",
    "tmp.to_netcdf(path=filename,mode='w')\n",
    "    \n",
    "#Clear-sky values\n",
    "print('...clr...')\n",
    "tmp = xr.Dataset()\n",
    "tmp['time'] = cFlx.time\n",
    "tmp['lat'] = cFlx.lat\n",
    "tmp['lon'] = cFlx.lon\n",
    "tmp['month'] = np.arange(1,13)\n",
    "tmp['month'].attrs = {'units' : '1', 'long_name' : 'month of the year'}\n",
    "\n",
    "for var in ['R_clr','Rc_clr','Ra_clr',\n",
    "            'R_aer','Rc_aer','Ra_aer',\n",
    "            'R_sfc','Rc_sfc','Ra_sfc']:\n",
    "    tmp[var] = xR[var]\n",
    "\n",
    "filename = dir_data+'CERES/FluxByCldTyp/FBCT_decomposition_clr.nc'\n",
    "os.system('rm %s' % filename)\n",
    "tmp.to_netcdf(path=filename,mode='w')\n",
    "    \n",
    "#Cloud-type values\n",
    "for cld in cldtypes:\n",
    "    print('...'+cld+'...')\n",
    "    \n",
    "    tmp = xr.Dataset()\n",
    "    tmp['time'] = cFlx.time\n",
    "    tmp['lat'] = cFlx.lat\n",
    "    tmp['lon'] = cFlx.lon\n",
    "    tmp['month'] = np.arange(1,13)\n",
    "    tmp['month'].attrs = {'units' : '1', 'long_name' : 'month of the year'}\n",
    "\n",
    "    for var in ['R_','Rc_','Ra_']:\n",
    "        tmp[var+cld] = xR[var+cld]\n",
    "\n",
    "    filename = dir_data+'CERES/FluxByCldTyp/FBCT_decomposition_%s.nc' % cld\n",
    "    os.system('rm %s' % filename)\n",
    "    tmp.to_netcdf(path=filename,mode='w')\n",
    "    \n",
    "    \n",
    "print('Done!')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3205aa6",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tf215env",
   "language": "python",
   "name": "python3"
  },
  "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.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}