global_mean_surface_temp.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Global Mean Surface Temperature\n",
    "\n",
    "This notebook uses similar techniques to `ECS_Gregory_method.ipynb`. Please refer to that notebook for details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib import pyplot as plt\n",
    "import xarray as xr\n",
    "import numpy as np\n",
    "import dask\n",
    "from dask.diagnostics import progress\n",
    "from tqdm.autonotebook import tqdm \n",
    "import intake\n",
    "import fsspec\n",
    "import seaborn as sns\n",
    "\n",
    "%matplotlib inline\n",
    "%config InlineBackend.figure_format = 'retina' "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "col = intake.open_esm_datastore(\"https://storage.googleapis.com/cmip6/pangeo-cmip6.json\")\n",
    "col"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "[eid for eid in col.df['experiment_id'].unique() if 'ssp' in eid]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# there is currently a significant amount of data for these runs\n",
    "expts = ['historical', 'ssp245', 'ssp585']\n",
    "\n",
    "query = dict(\n",
    "    experiment_id=expts,\n",
    "    table_id='Amon',                           \n",
    "    variable_id=['tas'],\n",
    "    member_id = 'r1i1p1f1',                     \n",
    ")\n",
    "\n",
    "col_subset = col.search(require_all_on=[\"source_id\"], **query)\n",
    "col_subset.df.groupby(\"source_id\")[\n",
    "    [\"experiment_id\", \"variable_id\", \"table_id\"]\n",
    "].nunique()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def drop_all_bounds(ds):\n",
    "    drop_vars = [vname for vname in ds.coords\n",
    "                 if (('_bounds') in vname ) or ('_bnds') in vname]\n",
    "    return ds.drop(drop_vars)\n",
    "\n",
    "def open_dset(df):\n",
    "    assert len(df) == 1\n",
    "    ds = xr.open_zarr(fsspec.get_mapper(df.zstore.values[0]), consolidated=True)\n",
    "    return drop_all_bounds(ds)\n",
    "\n",
    "def open_delayed(df):\n",
    "    return dask.delayed(open_dset)(df)\n",
    "\n",
    "from collections import defaultdict\n",
    "dsets = defaultdict(dict) \n",
    "\n",
    "for group, df in col_subset.df.groupby(by=['source_id', 'experiment_id']):\n",
    "    dsets[group[0]][group[1]] = open_delayed(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dsets_ = dask.compute(dict(dsets))[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# calculate global means\n",
    "\n",
    "def get_lat_name(ds):\n",
    "    for lat_name in ['lat', 'latitude']:\n",
    "        if lat_name in ds.coords:\n",
    "            return lat_name\n",
    "    raise RuntimeError(\"Couldn't find a latitude coordinate\")\n",
    "\n",
    "def global_mean(ds):\n",
    "    lat = ds[get_lat_name(ds)]\n",
    "    weight = np.cos(np.deg2rad(lat))\n",
    "    weight /= weight.mean()\n",
    "    other_dims = set(ds.dims) - {'time'}\n",
    "    return (ds * weight).mean(other_dims)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "expt_da = xr.DataArray(expts, dims='experiment_id', name='experiment_id',\n",
    "                       coords={'experiment_id': expts})\n",
    "\n",
    "dsets_aligned = {}\n",
    "\n",
    "for k, v in tqdm(dsets_.items()):\n",
    "    expt_dsets = v.values()\n",
    "    if any([d is None for d in expt_dsets]):\n",
    "        print(f\"Missing experiment for {k}\")\n",
    "        continue\n",
    "    \n",
    "    for ds in expt_dsets:\n",
    "        ds.coords['year'] = ds.time.dt.year\n",
    "        \n",
    "    # workaround for\n",
    "    # https://github.com/pydata/xarray/issues/2237#issuecomment-620961663\n",
    "    dsets_ann_mean = [v[expt].pipe(global_mean)\n",
    "                             .swap_dims({'time': 'year'})\n",
    "                             .drop('time')\n",
    "                             .coarsen(year=12).mean()\n",
    "                      for expt in expts]\n",
    "    \n",
    "    # align everything with the 4xCO2 experiment\n",
    "    dsets_aligned[k] = xr.concat(dsets_ann_mean, join='outer',\n",
    "                                 dim=expt_da)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with progress.ProgressBar():\n",
    "    dsets_aligned_ = dask.compute(dsets_aligned)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "source_ids = list(dsets_aligned_.keys())\n",
    "source_da = xr.DataArray(source_ids, dims='source_id', name='source_id',\n",
    "                         coords={'source_id': source_ids})\n",
    "\n",
    "big_ds = xr.concat([ds.reset_coords(drop=True)\n",
    "                    for ds in dsets_aligned_.values()],\n",
    "                    dim=source_da)\n",
    "\n",
    "big_ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_all = big_ds.sel(year=slice(1900, 2100)).to_dataframe().reset_index()\n",
    "df_all.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sns.relplot(data=df_all,\n",
    "            x=\"year\", y=\"tas\", hue='experiment_id',\n",
    "            kind=\"line\", ci=\"sd\", aspect=2);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# why are there axis legends for 9, 10 - doesn't make sense\n",
    "df_all.experiment_id.unique()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}