CMIP6_create_ensembles.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from xclim import ensembles\n",
    "from pathlib import Path\n",
    "import xarray as xr\n",
    "from CMIP6_IO import CMIP6_IO\n",
    "import os\n",
    "import dask\n",
    "from dask.distributed import Client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Create ensemble files from the outputs of running CMIP6_light.py and from forcing files\n",
    "\n",
    "This script will loop over all files found in a speficic folder `lightpath` and find files of the same variable and scenario to create ensemble files from. The various variables to create ensembles for include `[\"uvb\", \"uv\", \"uvi\", \"par\", \"uva\", \"tos\", \"siconc\", \"sithick\", \"tas\", \"sisnthick\", \"sisnconc\", \"uas\", \"vas\"]`.\n",
    "\n",
    "The output is stored under folder `ensemble_path` and the results are used to create the final plots of modeled lightlyfor the paper using notebooks `CMIP6_plot_light_results.ipynb` and `CMIP6_calculate_MPIESM2.ipynb`.\n",
    "\n",
    "This script may require substantial RAM depending on the region you are creating ensemble for. For datasets covering the Northern Hemisphere we used a 64CPU and 120GB RAM machine.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Define what models, scenarios, and variables to calculate ensemble files from\n",
    "\n",
    "source_ids = [\"MPI-ESM1-2-HR\",\"MPI-ESM1-2-LR\", \"CanESM5\", \"UKESM1-0-LL\"]\n",
    "\n",
    "# The weighting file is created using CMIP6_create_weights.ipynb together with \n",
    "# the ClimWIP package. The weights are calculated for variable tos using ERSSTv5 and CORA5.2 as observations\n",
    "# to compare with. The weights are calculated for the period 1993-2020.\n",
    "weights_file = \"data/calculated_weights_18092024.nc\"\n",
    "\n",
    "create_ensemble = True\n",
    "scenarios = [\"ssp245\",\"ssp585\"]\n",
    "var_names = [\"sithick\", \"tas\", \"sisnthick\", \"sisnconc\", \"uas\", \"vas\"]\n",
    "var_names = [\"tos\", \"uvb\", \"uv\", \"uva\", \"siconc\", \"sithick\", \"tas\", \"sisnthick\", \"sisnconc\", \"uas\", \"vas\"]\n",
    "var_names = [\"ghi\", \"uvb\", \"uv\", \"osa\"]\n",
    "var_names = [\"par\", \"uvb\",\"ghi\",\"osa\",\"uv\"] #, \"osa\", \"ghi\"]\n",
    "root=\"light/ncfiles_nobias\"\n",
    "percentiles = [2.5, 50.0, 97.5]\n",
    "\n",
    "create_ensemble_files = True\n",
    "io = CMIP6_IO()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false,
    "pycharm": {
     "is_executing": true,
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/tk/.conda/envs/opendrift/envs/actea/lib/python3.11/site-packages/distributed/node.py:182: UserWarning: Port 8787 is already in use.\n",
      "Perhaps you already have a cluster running?\n",
      "Hosting the HTTP server on port 33883 instead\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<Client: 'tcp://127.0.0.1:33673' processes=8 threads=64, memory=235.95 GiB>\n",
      "[CMIP6_IO] Opening file actea-shared/light/ncfiles_nobias/ssp245/par_CanESM5_r10i1p2f1_1979-01-01-2099-12-16_scenario_osa_ssp245.nc\n",
      "Found model in models CanESM5 here CanESM5_r1i1p1f1_CMIP6 weight 0.032578749610186045\n",
      "[CMIP6_IO] Opening file actea-shared/light/ncfiles_nobias/ssp245/par_CanESM5_r1i1p2f1_1979-01-01-2099-12-16_scenario_osa_ssp245.nc\n",
      "Found model in models CanESM5 here CanESM5_r1i1p1f1_CMIP6 weight 0.032578749610186045\n",
      "[CMIP6_IO] Opening file actea-shared/light/ncfiles_nobias/ssp245/par_CanESM5_r2i1p2f1_1979-01-01-2099-12-16_scenario_osa_ssp245.nc\n",
      "Found model in models CanESM5 here CanESM5_r1i1p1f1_CMIP6 weight 0.032578749610186045\n",
      "[CMIP6_IO] Opening file actea-shared/light/ncfiles_nobias/ssp245/par_CanESM5_r3i1p2f1_1979-01-01-2099-12-16_scenario_osa_ssp245.nc\n",
      "Found model in models CanESM5 here CanESM5_r1i1p1f1_CMIP6 weight 0.032578749610186045\n",
      "[CMIP6_IO] Opening file actea-shared/light/ncfiles_nobias/ssp245/par_CanESM5_r7i1p2f1_1979-01-01-2099-12-16_scenario_osa_ssp245.nc\n",
      "Found model in models CanESM5 here CanESM5_r1i1p1f1_CMIP6 weight 0.032578749610186045\n",
      "[CMIP6_IO] Opening file actea-shared/light/ncfiles_nobias/ssp245/par_CanESM5_r9i1p2f1_1979-01-01-2099-12-16_scenario_osa_ssp245.nc\n",
      "Found model in models CanESM5 here CanESM5_r1i1p1f1_CMIP6 weight 0.032578749610186045\n",
      "[CMIP6_IO] Opening file actea-shared/light/ncfiles_nobias/ssp245/par_MPI-ESM1-2-HR_r1i1p1f1_1979-01-01-2099-12-16_scenario_osa_ssp245.nc\n",
      "Found model in models MPI-ESM1-2-HR here MPI-ESM1-2-HR_r1i1p1f1_CMIP6 weight 0.07870167679146618\n",
      "[CMIP6_IO] Opening file actea-shared/light/ncfiles_nobias/ssp245/par_MPI-ESM1-2-HR_r2i1p1f1_1979-01-01-2099-12-16_scenario_osa_ssp245.nc\n",
      "Found model in models MPI-ESM1-2-HR here MPI-ESM1-2-HR_r1i1p1f1_CMIP6 weight 0.07870167679146618\n",
      "[CMIP6_IO] Opening file actea-shared/light/ncfiles_nobias/ssp245/par_MPI-ESM1-2-LR_r10i1p1f1_1979-01-01-2099-12-16_scenario_osa_ssp245.nc\n",
      "Found model in models MPI-ESM1-2-LR here MPI-ESM1-2-LR_r1i1p1f1_CMIP6 weight 0.03321524431828878\n",
      "[CMIP6_IO] Opening file actea-shared/light/ncfiles_nobias/ssp245/par_MPI-ESM1-2-LR_r1i1p1f1_1979-01-01-2099-12-16_scenario_osa_ssp245.nc\n",
      "Found model in models MPI-ESM1-2-LR here MPI-ESM1-2-LR_r1i1p1f1_CMIP6 weight 0.03321524431828878\n",
      "[CMIP6_IO] Opening file actea-shared/light/ncfiles_nobias/ssp245/par_MPI-ESM1-2-LR_r2i1p1f1_1979-01-01-2099-12-16_scenario_osa_ssp245.nc\n",
      "Found model in models MPI-ESM1-2-LR here MPI-ESM1-2-LR_r1i1p1f1_CMIP6 weight 0.03321524431828878\n",
      "[CMIP6_IO] Opening file actea-shared/light/ncfiles_nobias/ssp245/par_MPI-ESM1-2-LR_r6i1p1f1_1979-01-01-2099-12-16_scenario_osa_ssp245.nc\n",
      "Found model in models MPI-ESM1-2-LR here MPI-ESM1-2-LR_r1i1p1f1_CMIP6 weight 0.03321524431828878\n",
      "[CMIP6_IO] Opening file actea-shared/light/ncfiles_nobias/ssp245/par_UKESM1-0-LL_r1i1p1f2_1979-01-01-2099-12-16_scenario_osa_ssp245.nc\n",
      "Found model in models UKESM1-0-LL here UKESM1-0-LL_r1i1p1f2_CMIP6 weight 0.06350353688109489\n",
      "[CMIP6_IO] Opening file actea-shared/light/ncfiles_nobias/ssp245/par_UKESM1-0-LL_r2i1p1f2_1979-01-01-2099-12-16_scenario_osa_ssp245.nc\n",
      "Found model in models UKESM1-0-LL here UKESM1-0-LL_r1i1p1f2_CMIP6 weight 0.06350353688109489\n",
      "[CMIP6_IO] Opening file actea-shared/light/ncfiles_nobias/ssp245/par_UKESM1-0-LL_r3i1p1f2_1979-01-01-2099-12-16_scenario_osa_ssp245.nc\n",
      "Found model in models UKESM1-0-LL here UKESM1-0-LL_r1i1p1f2_CMIP6 weight 0.06350353688109489\n",
      "[CMIP6_IO] Opening file actea-shared/light/ncfiles_nobias/ssp245/par_UKESM1-0-LL_r4i1p1f2_1979-01-01-2099-12-16_scenario_osa_ssp245.nc\n",
      "Found model in models UKESM1-0-LL here UKESM1-0-LL_r1i1p1f2_CMIP6 weight 0.06350353688109489\n",
      "Number of datasets found 16 weights {'dims': 'realization', 'data': array([0.03257875, 0.03257875, 0.03257875, 0.03257875, 0.03257875,\n",
      "       0.03257875, 0.07870168, 0.07870168, 0.03321524, 0.03321524,\n",
      "       0.03321524, 0.03321524, 0.06350354, 0.06350354, 0.06350354,\n",
      "       0.06350354]), 'name': 'weights', 'models': ['light/ncfiles_nobias/ssp245/par_CanESM5_r10i1p2f1_1979-01-01-2099-12-16_scenario_osa_ssp245.nc', 'light/ncfiles_nobias/ssp245/par_CanESM5_r1i1p2f1_1979-01-01-2099-12-16_scenario_osa_ssp245.nc', 'light/ncfiles_nobias/ssp245/par_CanESM5_r2i1p2f1_1979-01-01-2099-12-16_scenario_osa_ssp245.nc', 'light/ncfiles_nobias/ssp245/par_CanESM5_r3i1p2f1_1979-01-01-2099-12-16_scenario_osa_ssp245.nc', 'light/ncfiles_nobias/ssp245/par_CanESM5_r7i1p2f1_1979-01-01-2099-12-16_scenario_osa_ssp245.nc', 'light/ncfiles_nobias/ssp245/par_CanESM5_r9i1p2f1_1979-01-01-2099-12-16_scenario_osa_ssp245.nc', 'light/ncfiles_nobias/ssp245/par_MPI-ESM1-2-HR_r1i1p1f1_1979-01-01-2099-12-16_scenario_osa_ssp245.nc', 'light/ncfiles_nobias/ssp245/par_MPI-ESM1-2-HR_r2i1p1f1_1979-01-01-2099-12-16_scenario_osa_ssp245.nc', 'light/ncfiles_nobias/ssp245/par_MPI-ESM1-2-LR_r10i1p1f1_1979-01-01-2099-12-16_scenario_osa_ssp245.nc', 'light/ncfiles_nobias/ssp245/par_MPI-ESM1-2-LR_r1i1p1f1_1979-01-01-2099-12-16_scenario_osa_ssp245.nc', 'light/ncfiles_nobias/ssp245/par_MPI-ESM1-2-LR_r2i1p1f1_1979-01-01-2099-12-16_scenario_osa_ssp245.nc', 'light/ncfiles_nobias/ssp245/par_MPI-ESM1-2-LR_r6i1p1f1_1979-01-01-2099-12-16_scenario_osa_ssp245.nc', 'light/ncfiles_nobias/ssp245/par_UKESM1-0-LL_r1i1p1f2_1979-01-01-2099-12-16_scenario_osa_ssp245.nc', 'light/ncfiles_nobias/ssp245/par_UKESM1-0-LL_r2i1p1f2_1979-01-01-2099-12-16_scenario_osa_ssp245.nc', 'light/ncfiles_nobias/ssp245/par_UKESM1-0-LL_r3i1p1f2_1979-01-01-2099-12-16_scenario_osa_ssp245.nc', 'light/ncfiles_nobias/ssp245/par_UKESM1-0-LL_r4i1p1f2_1979-01-01-2099-12-16_scenario_osa_ssp245.nc']}\n",
      "Creating ensemble for par and scenario ssp245\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "Alignment of dataset #00 failed: Time axis cannot be resampled to freq MS.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[3], line 148\u001b[0m\n\u001b[1;32m    146\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m Client() \u001b[38;5;28;01mas\u001b[39;00m client: \u001b[38;5;66;03m# set up local cluster on your laptop\u001b[39;00m\n\u001b[1;32m    147\u001b[0m     \u001b[38;5;28mprint\u001b[39m(client)\n\u001b[0;32m--> 148\u001b[0m     \u001b[43mcreate_ensemble_files\u001b[49m\u001b[43m(\u001b[49m\u001b[43mscenarios\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mvar_names\u001b[49m\u001b[43m)\u001b[49m\n",
      "Cell \u001b[0;32mIn[3], line 110\u001b[0m, in \u001b[0;36mcreate_ensemble_files\u001b[0;34m(scenarios, var_names)\u001b[0m\n\u001b[1;32m    108\u001b[0m     ens \u001b[38;5;241m=\u001b[39m ensembles\u001b[38;5;241m.\u001b[39mcreate_ensemble(ds_list, resample_freq\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mMS\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    109\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 110\u001b[0m     ens \u001b[38;5;241m=\u001b[39m \u001b[43mensembles\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcreate_ensemble\u001b[49m\u001b[43m(\u001b[49m\u001b[43mds_list\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mresample_freq\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mMS\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcalendar\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mnoleap\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m    112\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m use_weights:\n\u001b[1;32m    113\u001b[0m    \u001b[38;5;66;03m# if not only_perc:\u001b[39;00m\n\u001b[1;32m    114\u001b[0m     ens_stats \u001b[38;5;241m=\u001b[39m ensembles\u001b[38;5;241m.\u001b[39mensemble_mean_std_max_min(ens, weights\u001b[38;5;241m=\u001b[39mweights)\n",
      "File \u001b[0;32m~/.conda/envs/opendrift/envs/actea/lib/python3.11/site-packages/xclim/ensembles/_base.py:117\u001b[0m, in \u001b[0;36mcreate_ensemble\u001b[0;34m(datasets, multifile, resample_freq, calendar, realizations, cal_kwargs, **xr_kwargs)\u001b[0m\n\u001b[1;32m    111\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(datasets, \u001b[38;5;28mstr\u001b[39m) \u001b[38;5;129;01mand\u001b[39;00m realizations \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    112\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[1;32m    113\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mPassing `realizations` is not supported when `datasets` \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    114\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mis a glob pattern, as the final order is random.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    115\u001b[0m     )\n\u001b[0;32m--> 117\u001b[0m ds \u001b[38;5;241m=\u001b[39m \u001b[43m_ens_align_datasets\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    118\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdatasets\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    119\u001b[0m \u001b[43m    \u001b[49m\u001b[43mmultifile\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    120\u001b[0m \u001b[43m    \u001b[49m\u001b[43mresample_freq\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    121\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcalendar\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcalendar\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    122\u001b[0m \u001b[43m    \u001b[49m\u001b[43mcal_kwargs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcal_kwargs\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01mor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43m{\u001b[49m\u001b[43m}\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    123\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mxr_kwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    124\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    126\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m realizations \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    127\u001b[0m     realizations \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39marange(\u001b[38;5;28mlen\u001b[39m(ds))\n",
      "File \u001b[0;32m~/.conda/envs/opendrift/envs/actea/lib/python3.11/site-packages/xclim/ensembles/_base.py:449\u001b[0m, in \u001b[0;36m_ens_align_datasets\u001b[0;34m(datasets, multifile, resample_freq, calendar, cal_kwargs, **xr_kwargs)\u001b[0m\n\u001b[1;32m    447\u001b[0m     counts \u001b[38;5;241m=\u001b[39m time\u001b[38;5;241m.\u001b[39mastype(\u001b[38;5;28mbool\u001b[39m)\u001b[38;5;241m.\u001b[39mresample(time\u001b[38;5;241m=\u001b[39mresample_freq)\u001b[38;5;241m.\u001b[39mcount()\n\u001b[1;32m    448\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28many\u001b[39m(counts \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m1\u001b[39m):\n\u001b[0;32m--> 449\u001b[0m         \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mValueError\u001b[39;00m(\n\u001b[1;32m    450\u001b[0m             \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mAlignment of dataset #\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mi\u001b[38;5;132;01m:\u001b[39;00m\u001b[38;5;124m02d\u001b[39m\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m failed: \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    451\u001b[0m             \u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTime axis cannot be resampled to freq \u001b[39m\u001b[38;5;132;01m{\u001b[39;00mresample_freq\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m    452\u001b[0m         )\n\u001b[1;32m    453\u001b[0m     time \u001b[38;5;241m=\u001b[39m counts\u001b[38;5;241m.\u001b[39mtime\n\u001b[1;32m    455\u001b[0m ds[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtime\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m time\n",
      "\u001b[0;31mValueError\u001b[0m: Alignment of dataset #00 failed: Time axis cannot be resampled to freq MS."
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "def get_weights(weights_file: str, f: str):\n",
    "    weights = xr.open_dataset(weights_file)\n",
    "\n",
    "    weights = weights[\"weights\"]\n",
    "    models = weights[\"model_ensemble\"]\n",
    "    for source_id in source_ids:\n",
    "        if source_id in f:\n",
    "            for d, w in zip(models.values, weights.values):\n",
    "                if source_id in d.split(\"_\")[0]:\n",
    "                    print(f\"Found model in models {source_id} here {d} weight {w}\")\n",
    "                    return w\n",
    "\n",
    "\n",
    "def create_ensemble_files(scenarios, var_names):\n",
    "    max_models = 300\n",
    "    only_upload = False\n",
    "    only_perc = False\n",
    "    for use_weights in [True, False]:\n",
    "        for var_name in var_names:\n",
    "            for scenario in scenarios:\n",
    "                counter=0\n",
    "                ds_list = []\n",
    "                ensemble_stats = None\n",
    "                ensemble_perc = None\n",
    "                if var_name in [\"prw\", \"clt\", \"tos\", \"siconc\", \"sithick\", \"tas\", \"sisnthick\", \"sisnconc\", \"uas\", \"vas\"]:\n",
    "                    lightpath = f\"light/{scenario}\"\n",
    "                    ensemble_path = \"light/ensemble/{scenario}\"\n",
    "                else:\n",
    "                    lightpath = f\"{root}/{scenario}\"\n",
    "                    ensemble_path = f\"{root}/{scenario}/ensemble\"\n",
    "                if use_weights:\n",
    "                    ensemble_stats =  f\"{ensemble_path}/CMIP6_ensemble_stats_{scenario}_{var_name}_weighted.nc\"\n",
    "                    ensemble_perc =  f\"{ensemble_path}/CMIP6_ensemble_perc_{scenario}_{var_name}_weighted.nc\"\n",
    "                else:\n",
    "                    ensemble_stats =  f\"{ensemble_path}/CMIP6_ensemble_stats_{scenario}_{var_name}.nc\"\n",
    "                    ensemble_perc =  f\"{ensemble_path}/CMIP6_ensemble_perc_{scenario}_{var_name}.nc\"\n",
    "                 \n",
    "                if not os.path.exists(ensemble_path):\n",
    "                    os.makedirs(ensemble_path, exist_ok=True)\n",
    "\n",
    "                assert (\n",
    "                    ensemble_stats is not None\n",
    "                ), \"Unable to identify correct variable name to create output filename\"\n",
    "                assert (\n",
    "                    ensemble_perc is not None\n",
    "                ), \"Unable to identify correct variable name to create output filename\"\n",
    "\n",
    "                if not only_upload:\n",
    "                    if os.path.exists(ensemble_stats):\n",
    "                        os.remove(ensemble_stats)\n",
    "                    if os.path.exists(ensemble_perc):\n",
    "                        os.remove(ensemble_perc)\n",
    "\n",
    "                    current = f\"{lightpath}\"\n",
    "                \n",
    "                    file_on_gcs = io.list_dataset_on_gs(current)\n",
    "                    \n",
    "                    # Loop over all files and filter on the models defined in source_ids.\n",
    "                    # For each model read the corresponding weight from the weights file. We\n",
    "                    # will use these values to create a weights xr.Dataarray to be used to weight each \n",
    "                    # model when calculating the ensemble mean, std, max, min, and percentiles.\n",
    "                    models_weights = []\n",
    "                    model_names = []\n",
    "                    for f in file_on_gcs:\n",
    "                    \n",
    "                        var_name_mod = var_name\n",
    "                        if var_name not in [\"prw\", \"clt\", \"tos\", \"siconc\", \"sithick\", \"tas\", \"sisnthick\", \"sisnconc\", \"uas\", \"vas\"]:\n",
    "                            var_name_mod = f\"{var_name}_\"\n",
    "                            \n",
    "                        # Filter to only use the models we prefer\n",
    "                        if any(model in f for model in source_ids) and var_name_mod in f and \"uv_srf\" not in f and scenario in f and \"weight\" not in f and \"rsus\" not in f and \"rsds\" not in f:\n",
    "                            if counter >= max_models:\n",
    "                                pass\n",
    "                            else:\n",
    "                                if var_name==\"osa\" and f\"{scenario}/{var_name}_\" not in f:\n",
    "                                    pass\n",
    "                                else:\n",
    "                                    ds = io.open_dataset_on_gs(f).persist()  \n",
    "                                    models_weights.append(get_weights(weights_file, f))\n",
    "                                    model_names.append(f)\n",
    "                                    # Drop variable mask as we dont need it and it causes problems when\n",
    "                                    # calculating the ensemble mean.\n",
    "                                    if \"mask\" in ds.variables:\n",
    "                                        ds = ds.drop(\"mask\")       \n",
    "                                    ds_list.append(ds)\n",
    "                                    counter+=1\n",
    "                                    \n",
    "                    # Create the final xr.DataArray with the weights.    \n",
    "                    d = {\n",
    "                        \"dims\": \"realization\",\n",
    "                        \"data\": np.array(models_weights),\n",
    "                        \"name\": \"weights\",\n",
    "                        \"models\": model_names,\n",
    "                    }\n",
    "                    print(f\"Number of datasets found {len(ds_list)} weights {d}\")\n",
    "                    weights = xr.DataArray.from_dict(d)\n",
    "\n",
    "                    \n",
    "                    for i, ds in enumerate(ds_list):\n",
    "                        duplicates = ds['time'].to_index().duplicated()\n",
    "                        if duplicates.any():\n",
    "                            print(f\"Dataset {i} has {duplicates.sum()} duplicate time values for model {model_names[i]}\")\n",
    "            \n",
    "                    print(f\"Creating ensemble for {var_name} and scenario {scenario}\")\n",
    "                    if var_name in [\"tos\"]:\n",
    "                        ens = ensembles.create_ensemble(ds_list, resample_freq=\"MS\")\n",
    "                    else:\n",
    "                        ens = ensembles.create_ensemble(ds_list) #, resample_freq=\"MS\", calendar=\"noleap\")\n",
    "                    \n",
    "                    if use_weights:\n",
    "                       # if not only_perc:\n",
    "                        ens_stats = ensembles.ensemble_mean_std_max_min(ens, weights=weights)\n",
    "                        ens_perc = ensembles.ensemble_percentiles(\n",
    "                        ens, values=percentiles, split=False, weights=weights)   \n",
    "                    else:\n",
    "                        #if not only_perc:\n",
    "                        ens_stats = ensembles.ensemble_mean_std_max_min(ens)\n",
    "                        ens_perc = ensembles.ensemble_percentiles(\n",
    "                        ens, values=percentiles, split=False)\n",
    "                    \n",
    "                # Save to file and upload to GCS.\n",
    "                #if not only_upload:\n",
    "                #    if not only_perc:\n",
    "             #   if not only_perc:\n",
    "                ens_stats.to_netcdf(ensemble_stats)\n",
    "                ens_perc.to_netcdf(ensemble_perc)\n",
    "            \n",
    "                if Path(ensemble_stats).exists():\n",
    "                    io.upload_to_gcs(ensemble_stats)\n",
    "                print(f\"Created ensemble {ensemble_stats}\")\n",
    "\n",
    "                if Path(ensemble_perc).exists():\n",
    "                    io.upload_to_gcs(ensemble_perc)\n",
    "                print(f\"Created ensemble {ensemble_perc}\")\n",
    "\n",
    "               # if not only_perc:\n",
    "                if Path(ensemble_stats).exists():\n",
    "                    Path(ensemble_stats).unlink()\n",
    "                if Path(ensemble_perc).exists():\n",
    "                     Path(ensemble_perc).unlink()\n",
    "                \n",
    "\n",
    "dask.config.set(**{\"array.slicing.split_large_chunks\": False})  \n",
    "with Client() as client: # set up local cluster on your laptop\n",
    "    print(client)\n",
    "    create_ensemble_files(scenarios, var_names)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "actea",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}