{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 11,
   "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": 12,
   "id": "6130b8b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Set paths\n",
    "dir_data = '/Users/jzhang/Documents/Data/'\n",
    "dir_proj = '/Users/jzhang/Documents/Projects/Albedo_Symmetry_WE/codes/open_code/'\n",
    "\n",
    "#Load data\n",
    "ebaf = xr.open_mfdataset(sorted(glob(dir_data+'CERES/EBAFed421/CERES_EBAF_Ed4.2.1_Subset_*.nc'))).sel(time=slice(np.datetime64('2001-01-01'),np.datetime64('2024-12-31')))\n",
    "ocean = xr.open_dataset(dir_data+'CERES/SYN1deg/CERES_SYN1deg-Month_Terra-Aqua-NOAA20_Ed4.2_Subset_200003-202512.nc').sel(time=slice(np.datetime64('2001-01-01'),np.datetime64('2024-12-31')))\n",
    "oce_frac = ocean['aux_ocean_mon']/100.\n",
    "snow_frac = ocean['aux_snow_mon']/100."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e673245",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Calculating:\n",
      "...all...\n",
      "...clr...\n",
      "...clr...\n",
      "...aer...\n",
      "...sfc...\n",
      "Done!\n",
      "\n",
      "Saving files:\n",
      "...all...\n",
      "...cld...\n",
      "...ocean and snow...\n",
      "...clr...\n",
      "Done!\n"
     ]
    }
   ],
   "source": [
    "#Set up master array to store decompositions\n",
    "xR = xr.Dataset() \n",
    "xR['time'] = ebaf.time\n",
    "xR['lat'] = ebaf.lat\n",
    "xR['lon'] = ebaf.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 = ebaf['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 = ebaf['solar_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 = ebaf['toa_sw_clr_t_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 for total region'}\n",
    "\n",
    "###Cloudy-sky, i.e. Cloud Radiative Effect (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 for total region'}\n",
    "\n",
    "### Out going LW\n",
    "\n",
    "OLR = ebaf['toa_lw_all_mon']\n",
    "OLR = OLR.where(OLR>0,0)\n",
    "xR['OLR'] = OLR\n",
    "xR['OLR'].attrs = {'units' : 'W m-2', 'long_name' : 'Outgoing longwave radiation'}\n",
    "\n",
    "### Earth Energy Imbalance (EEI)\n",
    "EEI = S - R - OLR\n",
    "xR['EEI'] = EEI\n",
    "xR['EEI'].attrs = {'units' : 'W m-2', 'long_name' : 'Earth Energy Imbalance'}\n",
    "\n",
    "#Use EBAF values to help with clear-sky decomposition\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",
    "###Ocean-snow-fraction\n",
    "\n",
    "O_frac = oce_frac\n",
    "xR['O_frac'] = O_frac\n",
    "xR['O_frac'].attrs = {'units' : '1', 'long_name' : 'Ocean fraction'}\n",
    "\n",
    "S_frac = snow_frac\n",
    "xR['S_frac'] = S_frac\n",
    "xR['S_frac'].attrs = {'units' : '1', 'long_name' : 'Snow fraction'}\n",
    "\n",
    "#Calculate climatological values and deseasonalized anomalies for all-sky, cloudy-sky, clear-sky, OLR, EEI, and ocean-snow-fraction values\n",
    "\n",
    "#Feb_wts = 5*[28,28,28,29] #leap years between 2005-2024\n",
    "#Feb_wts = [28]+4*[29,28,28,28]+[29,28,28] #leap years between 2003-2022\n",
    "Feb_wts = 6*[28,28,28,29] #leap years between 2001-2024\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 (2001-01 to 2024-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 (2001-01 to 2024-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 (2001-01 to 2024-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['OLRc'] = (['month','lat','lon'],np.nan*np.ones((12,180,360)))\n",
    "xR['OLRc'].attrs = {'units' : 'W m-2', 'long_name' : 'Climatological (2001-01 to 2024-12) value of OLR'}\n",
    "xR['OLRa'] = (['time','lat','lon'],np.nan*np.ones(xR['R'].shape))\n",
    "xR['OLRa'].attrs = {'units' : 'W m-2', 'long_name' : 'Deseasonalized anomaly of OLR'}\n",
    "\n",
    "xR['EEIc'] = (['month','lat','lon'],np.nan*np.ones((12,180,360)))\n",
    "xR['EEIc'].attrs = {'units' : 'W m-2', 'long_name' : 'Climatological (2001-01 to 2024-12) value of EEI'}\n",
    "xR['EEIa'] = (['time','lat','lon'],np.nan*np.ones(xR['R'].shape))\n",
    "xR['EEIa'].attrs = {'units' : 'W m-2', 'long_name' : 'Deseasonalized anomaly of EEI'}\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 (2001-01 to 2024-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",
    "xR['Sc_frac'] = (['month','lat','lon'],np.nan*np.ones((12,180,360)))\n",
    "xR['Sc_frac'].attrs = {'units' : '1', 'long_name' : 'Climatological (2001-01 to 2024-12) value of S_frac'}\n",
    "xR['Sa_frac'] = (['time','lat','lon'],np.nan*np.ones(xR['R'].shape))\n",
    "xR['Sa_frac'].attrs = {'units' : '1', 'long_name' : 'Deseasonalized anomaly of S_frac'}\n",
    "\n",
    "for m in range(1,13):\n",
    "\n",
    "    R = xR['R'].sel(time=slice(np.datetime64('2001-01'),np.datetime64('2025-01')))\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('2001-01'),np.datetime64('2025-01')))\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('2001-01'),np.datetime64('2025-01')))\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",
    "    EE = xR['EEI'].sel(time=slice(np.datetime64('2001-01'),np.datetime64('2025-01')))\n",
    "    if m ==2: xR['EEIc'][m-1] = np.average(EE[EE.time.dt.month==m],weights=Feb_wts,axis=0)\n",
    "    else: xR['EEIc'][m-1] = EE[EE.time.dt.month==m].mean(axis=0)\n",
    "    xR['EEIa'][xR.time.dt.month==m] = xR['EEI'][xR.time.dt.month==m].values-xR['EEIc'][m-1].values\n",
    "\n",
    "    LL = xR['OLR'].sel(time=slice(np.datetime64('2001-01'),np.datetime64('2025-01')))\n",
    "    if m ==2: xR['OLRc'][m-1] = np.average(LL[LL.time.dt.month==m],weights=Feb_wts,axis=0)\n",
    "    else: xR['OLRc'][m-1] = LL[LL.time.dt.month==m].mean(axis=0)\n",
    "    xR['OLRa'][xR.time.dt.month==m] = xR['OLR'][xR.time.dt.month==m].values-xR['OLRc'][m-1].values\n",
    "\n",
    "    OO = xR['O_frac'].sel(time=slice(np.datetime64('2001-01'),np.datetime64('2025-01')))\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",
    "    WW = xR['S_frac'].sel(time=slice(np.datetime64('2001-01'),np.datetime64('2025-01')))\n",
    "    if m ==2: xR['Sc_frac'][m-1] = np.average(WW[WW.time.dt.month==m],weights=Feb_wts,axis=0)\n",
    "    else: xR['Sc_frac'][m-1] = WW[WW.time.dt.month==m].mean(axis=0)\n",
    "    xR['Sa_frac'][xR.time.dt.month==m] = xR['S_frac'][xR.time.dt.month==m].values-xR['Sc_frac'][m-1].values\n",
    "\n",
    "print('...clr...')\n",
    "\n",
    "for loc in ['clr','aer','sfc']:\n",
    "    print('...'+loc+'...')\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 (2001-01 to 2024-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('2001-01'),np.datetime64('2025-01')))\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",
    "print('Done!')\n",
    "    \n",
    "print('Saving files:')\n",
    "\n",
    "#All-sky values\n",
    "print('...all...')\n",
    "tmp = xr.Dataset()\n",
    "tmp['time'] = ebaf.time\n",
    "tmp['lat'] = ebaf.lat\n",
    "tmp['lon'] = ebaf.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','OLR','OLRc','OLRa','EEI','EEIc','EEIa']:\n",
    "    tmp[var] = xR[var]\n",
    "\n",
    "filename = dir_data+'CERES/EBAFed421/EBAF_decomposition_all_01-24.nc'\n",
    "os.system('rm %s' % filename)\n",
    "tmp.to_netcdf(path=filename,mode='w')\n",
    "\n",
    "#CRE values\n",
    "print('...cld...')\n",
    "tmp = xr.Dataset()\n",
    "tmp['time'] = ebaf.time\n",
    "tmp['lat'] = ebaf.lat\n",
    "tmp['lon'] = ebaf.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/EBAFed421/EBAF_decomposition_cld_01-24.nc'\n",
    "os.system('rm %s' % filename)\n",
    "tmp.to_netcdf(path=filename,mode='w')\n",
    "\n",
    "\n",
    "#Ocean-snow-fraction values\n",
    "print('...ocean and snow...')\n",
    "tmp = xr.Dataset()\n",
    "tmp['time'] = ebaf.time\n",
    "tmp['lat'] = ebaf.lat\n",
    "tmp['lon'] = ebaf.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",
    "for var in ['S_frac','Sc_frac','Sa_frac']:\n",
    "    tmp[var] = xR[var]\n",
    "\n",
    "filename = dir_data+'CERES/EBAFed421/EBAF_decomposition_ocean_snow_01-24.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'] = ebaf.time\n",
    "tmp['lat'] = ebaf.lat\n",
    "tmp['lon'] = ebaf.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/EBAFed421/EBAF_decomposition_clr_01-24.nc'\n",
    "os.system('rm %s' % filename)\n",
    "tmp.to_netcdf(path=filename,mode='w')    \n",
    "    \n",
    "print('Done!')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b3e56b8",
   "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
}