calculate_trends_specifictimes.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bc1c47bc-8272-4abb-b0d5-5c160e45816f",
   "metadata": {},
   "source": [
    "# Calculate pixel-level SST trends for specific year combinations\n",
    "\n",
    "A calculation, file by file, of pixel-wise linear trends of multiple specific start/end year pairs, used in Figures 4/5 and their derivatives. Set up to produce pixel-wise linear trends of SSTs and 500 hPa $\\omega$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1d517642-9106-418d-bd53-36837137157a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import os\n",
    "import glob\n",
    "import re\n",
    "from scipy import stats as sstats\n",
    "from matplotlib import pyplot as plt\n",
    "from cartopy import crs as ccrs\n",
    "from tqdm.notebook import tqdm\n",
    "import cmocean\n",
    "import warnings\n",
    "\n",
    "from funcs_support import get_filepaths,get_params\n",
    "dir_list = get_params()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "92669163-3926-47ce-a19d-fd576db84460",
   "metadata": {},
   "outputs": [],
   "source": [
    "def seasmean(ds,seas_range):\n",
    "    ''' Calculate seasonal means of one season\n",
    "\n",
    "    Note: the output year is always the year of the _first_\n",
    "    month of the season - so 2013/12 - 2014/2 is returned \n",
    "    as the \"2013\" DJF season. \n",
    "\n",
    "    Parameters\n",
    "    -------------\n",
    "    ds : :py:class:`xr.Dataset` or :py:class:`xr.DataArray`\n",
    "\n",
    "    seas_range : `py:class:`list` of :py:class:`int`\n",
    "        e.g., [3,5] for MAM. A list of the integer indices\n",
    "        of the months (inclusive) in season. THIS IS 1-,\n",
    "        NOT 0-INDEXED.\n",
    "\n",
    "    Returns\n",
    "    -------------\n",
    "    ds_out : :py:class:`xr.Dataset` or :py:class:`xr.DataArray`\n",
    "        Seasonal means of the input data, with the time \n",
    "        dimension renamed \"year\" and showing integer years\n",
    "        instead of time\n",
    "    \n",
    "    '''\n",
    "    if seas_range[0] < seas_range[1]:\n",
    "        # If the start season < end season, no wrap around, \n",
    "        # so we need months that are least the start\n",
    "        # month and at most the end month\n",
    "        tsub = ((ds.time.dt.month>=seas_range[0]) & \n",
    "                (ds.time.dt.month<=seas_range[1]))\n",
    "    else:\n",
    "        # If the start season > end season, wrap around the\n",
    "        # new year, so we need months that are least the start\n",
    "        # month OR at most the end month\n",
    "        tsub = ((ds.time.dt.month >= seas_range[0]) | \n",
    "                (ds.time.dt.month <= seas_range[1]))\n",
    "        \n",
    "    # Subset to just this season\n",
    "    ds_out = ds.isel(time=tsub)\n",
    "    \n",
    "    # Resample annually to get seasonal means (using the anchored offset\n",
    "    # to ensure wrap-around seasons are correctly averaged\n",
    "    ds_out = (ds_out.resample(time='1YS-'+\n",
    "                              pd.Timestamp('1900-'+str(seas_range[0]).zfill(2)+'-01').month_name()[0:3].upper()).\n",
    "          mean(skipna=False))\n",
    "\n",
    "    # Change time dimension to year int\n",
    "    ds_out['time'] = ds_out['time'].dt.year\n",
    "    ds_out = ds_out.rename({'time':'year'})\n",
    "\n",
    "    # Take out incomplete seasons (from wrap around, both the first\n",
    "    # and last year's seasons are incomplete means) \n",
    "    if seas_range[0] > seas_range[1]:\n",
    "        ds_out = ds_out.isel(year=slice(1,-1))\n",
    "    \n",
    "    return ds_out\n",
    "\n",
    "def seasmeans(ds,seasons):\n",
    "    ''' Calculate seasonal means of multiple seasons\n",
    "\n",
    "    A wrapper for :py:meth:`seasmean`. \n",
    "\n",
    "    Parameters\n",
    "    -------------\n",
    "    ds : :py:class:`xr.Dataset` or :py:class:`xr.DataArray`\n",
    "\n",
    "    seasons : :py:class:`dict` \n",
    "        Dictionary, with keys as season names and items as \n",
    "        lists with length 2 of the start, end months (inclusive)\n",
    "        of each season. Note that seasonal limits are 1-indexed, \n",
    "        i.e., Jan == 1. \n",
    "\n",
    "    Returns\n",
    "    -------------\n",
    "    ds_out : :py:class:`xr.Dataset` or :py:class:`xr.DataArray`\n",
    "        Seasonal means of the input data, along a new dimension \n",
    "        \"season\", with the time dimension renamed \"year\" and \n",
    "        showing integer years instead of time\n",
    "\n",
    "    '''\n",
    "\n",
    "    if type(ds) == xr.core.dataset.Dataset:\n",
    "        # Get vars with time dimension\n",
    "        vars_wtime = [v for v in ds if 'time' in ds[v].sizes]\n",
    "    \n",
    "    # Aggregate to season\n",
    "    ds_out = xr.concat([seasmean(ds[vars_wtime],seas_range) for seas,seas_range in seasons.items()],\n",
    "                      dim = pd.Index([seas for seas in seasons],name='season'))\n",
    "\n",
    "    if type(ds) == xr.core.dataset.Dataset:\n",
    "        # Add back non-time vars\n",
    "        for var in [v for v in ds if v not in vars_wtime]:\n",
    "            ds_out[var] = ds[var].copy()\n",
    "\n",
    "    # Input seasonal information\n",
    "    ds_out['season_bnds'] = xr.DataArray(np.array([seas_range for seas,seas_range in seasons.items()]),\n",
    "                                             dims = ('season','bnds'),\n",
    "                                             coords = {'season':(('season'),[seas for seas in seasons]),\n",
    "                                                       })\n",
    "    ds_out['season_bnds'].attrs['DESCRIPTION'] = 'Start, end month (inclusive) of season (1-indexed, i.e., January = 1)'\n",
    "\n",
    "    return ds_out\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7f1c445-33cd-4294-9db2-40ddc0b878d2",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "c7804c91-757d-4a3e-a522-5d762d0fb79d",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Processing model ERSST\n",
      "Processing exp hist-none, 1 runs...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fde272c75bd14ce8b98dab7f456df404",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/dx02/kschwarz/climate_proc/ERSST/tostrends_seasavg_ERSST_hist-none_obs_19820101-20221231_spectrends.nc saved!\n"
     ]
    }
   ],
   "source": [
    "#-------------- Inputs --------------\n",
    "remove_all_existing = False\n",
    "overwrite = False\n",
    "\n",
    "# Get seasonal means\n",
    "seasons = {'MAM':[3,5],\n",
    "           'AMJ':[4,6],\n",
    "           'SON':[9,11],\n",
    "           'OND':[10,12],\n",
    "           'DJF':[12,2]}\n",
    "prev_seasons = ['OND']\n",
    "\n",
    "# Get timeframes\n",
    "trends = {'rowell':[1986,2004],\n",
    "          'long0':[1982,2014],\n",
    "          'long1':[1982,2022],\n",
    "          'wet0':[2000,2014],\n",
    "          'wet1':[2000,2020],\n",
    "          'dry0':[1985,2000]}\n",
    "\n",
    "# Future experiments to extend historic runs on\n",
    "base_exp = 'historical'\n",
    "#base_exp = 'hindcastsf'\n",
    "future_exps = ['ssp245']\n",
    "# Range over which to calculate standard deviation\n",
    "sd_range = slice(1981,2014)\n",
    "\n",
    "process_var = 'tos'; freq = 'Omon'\n",
    "\n",
    "acceptable_trend_meths = ['linear']\n",
    "trend_meths = ['linear']\n",
    "\n",
    "\n",
    "# Process by variable\n",
    "for process_var,freq in zip(['tos','wap500'],['Omon','Amon']):\n",
    "    \n",
    "    # Filepaths for source files\n",
    "    df = get_filepaths()\n",
    "    df = df.query('varname == \"'+process_var+'\" and freq == \"'+freq+'\"')\n",
    "    \n",
    "    # Filepaths for (possible previously) output files\n",
    "    dfs_proc = get_filepaths(source_dir='proc')\n",
    "    \n",
    "    mods = df.model.unique()\n",
    "    for mod in tqdm(mods):\n",
    "        print('\\nProcessing model '+mod)\n",
    "        #-------------- Setup --------------\n",
    "        \n",
    "        # Get files for that model\n",
    "        df_tmp = df.query('model == \"'+mod+'\"')\n",
    "    \n",
    "        if 'historical' not in df_tmp.exp.values:\n",
    "            warnings.warn('No \"historical\" files for '+mod+', skipped!')\n",
    "            continue\n",
    "        \n",
    "        # Get runs for each experiment \n",
    "        run_dict = df_tmp.groupby('exp')['run'].apply(lambda x: np.sort(list(x))).to_dict()\n",
    "        \n",
    "        # Get for each future exp to process which runs\n",
    "        # match up with the historical runs\n",
    "        overlap_runs = dict()\n",
    "        for exp in future_exps:\n",
    "            if exp in run_dict:\n",
    "                overlap_runs = {exp:[run for run in run_dict[base_exp] if run in run_dict[exp]]}\n",
    "        \n",
    "        \n",
    "        # Set which run-exp combinations to process\n",
    "        process_list = dict()\n",
    "        for exp in future_exps:\n",
    "            if exp in overlap_runs:\n",
    "                process_list['hist-'+exp] = [df_tmp.loc[((df_tmp.exp == 'historical') | \n",
    "                                                               (df_tmp.exp == exp)) & (df_tmp.run == run)]\n",
    "                                                   for run in overlap_runs[exp]]\n",
    "        \n",
    "        # Now get any runs that are historical only to save\n",
    "        # in hist-none trends\n",
    "        hist_only_runs = [run not in np.unique(np.array([run_dict[exp] for exp in future_exps if exp in run_dict]).flatten()) \n",
    "                            for run in run_dict['historical']]\n",
    "        if np.any(hist_only_runs):\n",
    "            process_list['hist-none'] = [df_tmp.loc[(df_tmp.exp == 'historical') & (df_tmp.run == run)]\n",
    "                                                   for run in run_dict['historical'][hist_only_runs]]\n",
    "        \n",
    "        #-------------- Process by run, exp combination --------------\n",
    "        for exp in process_list:\n",
    "            print('Processing exp '+exp+', '+str(len(process_list[exp]))+' runs...')\n",
    "            for fparams in tqdm(process_list[exp]): \n",
    "                # Get save filename\n",
    "                output_fn = (dir_list['proc']+mod+'/'+\n",
    "                             process_var+'trends_seasavg_'+mod+'_'+exp+'_'+fparams.run.values[0]+'_'+\n",
    "                            str(np.min([ts for tf,ts in trends.items()]))+'0101'+'-'+\n",
    "                            str(np.max([ts for tf,ts in trends.items()]))+'1231'+'_'+\n",
    "                            'spectrends.nc')\n",
    "        \n",
    "                # Get any existing files\n",
    "                dfs_proc_tmp = (dfs_proc.query('varname == \"'+process_var+'trends\" and model == \"'+mod+\n",
    "                                               '\" and freq == \"seasavg\" and suffix == \"spectrends\"'))\n",
    "        \n",
    "                if not overwrite:\n",
    "                    if (len(dfs_proc_tmp)>0) and (dfs_proc_tmp.path.str.contains(output_fn).any()):\n",
    "                        process = False\n",
    "                    else:\n",
    "                        process = True\n",
    "                else:\n",
    "                    process = True\n",
    "        \n",
    "                if remove_all_existing:\n",
    "                    for fn in dfs_proc_tmp.path.values:\n",
    "                        os.remove(fn) \n",
    "                        print(fn+' removed.')\n",
    "        \n",
    "                if process:\n",
    "                    try:\n",
    "                        #--------- Load and cleanup\n",
    "                        # Load\n",
    "                        ds = xr.open_mfdataset(fparams.path)\n",
    "        \n",
    "                        # Subset to just years needed (with 1 year padding,\n",
    "                        # if available, to catch previous seasons) \n",
    "                        ds = ds.sel(time=slice(str(np.min([ts for tr,ts in trends.items()])-1)+'-01-01',\n",
    "                                               str(np.max([ts for tr,ts in trends.items()])+1)+'-12-'+str(np.max(ds.time.dt.daysinmonth.values))))\n",
    "        \n",
    "                        # Load into memory\n",
    "                        ds = ds.load()\n",
    "                        \n",
    "                        # Clean up\n",
    "                        # open_mfdataset sometimes concatenates along time even \n",
    "                        # if there's no time dependence\n",
    "                        for gridv in ['lon','lat']:\n",
    "                            if gridv in ds.cf.bounds:\n",
    "                                if 'time' in ds[ds.cf.bounds[gridv]].sizes:\n",
    "                                    ds[ds.cf.bounds[gridv][0]] = ds[ds.cf.bounds[gridv][0]].isel(time=0)\n",
    "        \n",
    "                \n",
    "                        #--------- Get seasonal means\n",
    "                        ds = seasmeans(ds,seasons)\n",
    "                        \n",
    "                        # Add an additional seasonal mean from the previous season\n",
    "                        for seas in prev_seasons:\n",
    "                            ds_tmp = ds.sel(season=seas)\n",
    "                            ds_tmp['year'] = ds_tmp['year'].values+1\n",
    "                            ds_tmp = ds_tmp.isel(year=slice(0,-1))\n",
    "                            ds_tmp[process_var] = ds_tmp[process_var].expand_dims({'season':['prev'+seas]})\n",
    "    \n",
    "                            # data_vars 'minimal' ensures that only the desired variable\n",
    "                            # is concatenated \n",
    "                            ds = xr.concat([ds,ds_tmp],dim='season',data_vars='minimal')\n",
    "    \n",
    "                        #--------- Get std\n",
    "                        with warnings.catch_warnings():\n",
    "                            # Catches warning about dof == 0\n",
    "                            warnings.filterwarnings('ignore')\n",
    "                            ds[process_var+'_std'] = ds[process_var].sel(year=sd_range).std('year')\n",
    "                \n",
    "                        #--------- Get trends\n",
    "                        # Get which trends to calculate\n",
    "                        trends_tmp = {trend_name:trend_years for trend_name,trend_years in trends.items()\n",
    "                                      if np.all([y in ds.year.values for y in trend_years])}\n",
    "        \n",
    "                        # Calculate by different desired trend methods\n",
    "                        for trend_meth in trend_meths:\n",
    "                            if trend_meth=='linear':\n",
    "                                # Calculate trends \n",
    "                                ds[process_var+'_lslope'] = xr.concat([(ds[process_var].\n",
    "                                                                     sel(year = slice(*trends_tmp[period])).\n",
    "                                                                     polyfit(dim='year',deg=1).sel(degree=1))\n",
    "                                                                    for period in trends_tmp],\n",
    "                                                                   dim = pd.Index([period for period in trends_tmp],name='time_period'))['polyfit_coefficients']\n",
    "                                if 'long_name' in ds[process_var].attrs:\n",
    "                                    ds[process_var+'_lslope'].attrs['long_name'] = ds[process_var].attrs['long_name']+' OLS slope'\n",
    "                                else:\n",
    "                                    ds[process_var+'_lslope'].attrs['long_name'] = process_var+' OLS slope'\n",
    "                                if 'units' in ds[process_var].attrs:\n",
    "                                    ds[process_var+'_lslope'].attrs['units'] = ds[process_var].attrs['units']+'/year'\n",
    "                                else:\n",
    "                                    ds[process_var+'_lslope'].attrs['units'] = '/year'\n",
    "                                    \n",
    "                            elif trend_meth == 'theilsen':\n",
    "                                raise NotImplementedError('theilsen slope estimator not yet implemented')\n",
    "                            else:\n",
    "                                raise KeyError('trend_meth '+trend_meth+' must be one of '+', '.join(acceptable_trend_meths))\n",
    "                            \n",
    "                        # Add start / end years for each trend period\n",
    "                        ds = xr.merge([ds,\n",
    "                                   xr.Dataset({'start_year':(('time_period'),[ys[0] for period,ys in trends_tmp.items()]),\n",
    "                                               'end_year':(('time_period'),[ys[1] for period,ys in trends_tmp.items()])},\n",
    "                                              coords = {'time_period':ds.time_period})])\n",
    "                \n",
    "                        #--------- Export\n",
    "                        # Clean up\n",
    "                        ds = (ds.set_coords(['start_year','end_year','season_bnds']).\n",
    "                                 drop_vars('degree').\n",
    "                                 drop_vars(process_var))\n",
    "            \n",
    "                        # Attributes\n",
    "                        ds.attrs['SOURCE'] = 'calculate_sst_trends_specifictimes.ipynb'\n",
    "                        ds.attrs['DESCRIPTION'] = ', '.join(trend_meths)+' trends in seasonal means calculated at each grid cell for specific year pairs.'\n",
    "            \n",
    "                        # Remove existing if needed\n",
    "                        if os.path.exists(output_fn):\n",
    "                            os.remove(output_fn)\n",
    "                            print(output_fn+' removed to allow overwrite!')\n",
    "                    \n",
    "                        # Save\n",
    "                        ds.to_netcdf(output_fn)\n",
    "                        print(output_fn+' saved!')\n",
    "        \n",
    "                        del ds\n",
    "                    except:\n",
    "                        warnings.warn('Issue with '+mod+' run '+fparams.run.values[0]+', skipped')\n",
    "                        continue\n",
    "                else:\n",
    "                    print(output_fn+' already exists, skipped!')\n",
    "\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d93c914-a170-4d23-8e08-58d957ccae0e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}