CMIP6_growth_potential.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "Script to calculate growth rates of fish during winter and summer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "source": [
    "import xarray as xr\n",
    "import seaborn as sns\n",
    "import warnings\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "plt.style.use('default')\n",
    "import os\n",
    "warnings.simplefilter('ignore') # filter some warning messages\n",
    "sns.axes_style(\"whitegrid\")\n",
    "\n",
    "import datetime\n",
    "import logging\n",
    "\n",
    "import xarray as xr\n",
    "import matplotlib as mpl\n",
    "import numpy as np\n",
    "import warnings\n",
    "import matplotlib.pyplot as plt\n",
    "import pandas as pd\n",
    "import os\n",
    "import seaborn as sns\n",
    "import cartopy.crs as ccrs\n",
    "import cartopy.feature as cfeature\n",
    "from shapely.geometry import box, mapping\n",
    "import geopandas as gpd\n",
    "import global_land_mask\n",
    "import xesmf as xe\n",
    "import iris\n",
    "import CMIP6_downscale_iris\n",
    "\n",
    "import sys\n",
    "from dateutil.relativedelta import relativedelta\n",
    "from cartopy.util import add_cyclic_point\n",
    "from CMIP6_model import CMIP6_MODEL\n",
    "import CMIP6_IO\n",
    "import pyproj\n",
    "pyproj.datadir.set_data_dir('/Users/trondkr/anaconda3/envs/gcs2/share/proj')\n",
    "pyproj.datadir.get_data_dir()"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "source": [
    "# Global variables\n",
    "par_threshold = 0.0368\n",
    "depth = 5 # meter\n",
    "attenuation = 0.18"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "source": [
    "def convert_to_180(ds):\n",
    "    return (ds.assign_coords(lon=(((ds.lon + 180) % 360) - 180))).sortby('lon')\n",
    "\n",
    "def convert_time(ds):\n",
    "    if not ds.indexes[\"time\"].dtype in [\"datetime64[ns]\"]:\n",
    "\n",
    "        time_objects = ds.indexes['time'].to_datetimeindex()\n",
    "        ds=ds.assign_coords({\"time\": time_objects})\n",
    "        ds = xr.decode_cf(ds)\n",
    "\n",
    "    return ds\n",
    "\n",
    "def xr_add_cyclic_point(da):\n",
    "    \"\"\"\n",
    "    Inputs\n",
    "    da: xr.DataArray with dimensions (time,lat,lon)\n",
    "    \"\"\"\n",
    "\n",
    "    # Use add_cyclic_point to interpolate input data\n",
    "    lon_idx = 2 #da.dims.index('lon')\n",
    "\n",
    "    wrap_data, wrap_lon = add_cyclic_point(da.values, coord=da.lon, axis=lon_idx)\n",
    "\n",
    "    # Generate output DataArray with new data but same structure as input\n",
    "    outp_da = xr.DataArray(data=wrap_data,\n",
    "                           coords = {'time': da.time, 'lat': da.lat, 'lon': wrap_lon},\n",
    "                           dims=da.dims,\n",
    "                           attrs=da.attrs)\n",
    "\n",
    "    return outp_da\n",
    "\n",
    "def create_land_ocean_mask(cmip6_grid: xr.Dataset) -> xr.DataArray:\n",
    "    print(\"[create_land_ocean_mask] Running create_land_ocean_mask\")\n",
    "    lon = cmip6_grid.lon.values\n",
    "    lat = cmip6_grid.lat.values\n",
    "    lon_180 = xr.where(lon > 180, lon - 360, lon)\n",
    "\n",
    "    lon_grid, lat_grid = np.meshgrid(lon_180, lat)\n",
    "\n",
    "    mask_data = global_land_mask.globe.is_ocean(lat_grid, lon_grid).astype(int)\n",
    "\n",
    "    return xr.DataArray(mask_data, coords={'lat': lat, 'lon': lon},\n",
    "                        dims=['lat', 'lon'])\n",
    "\n",
    "def get_LME_records():\n",
    "    lme_file = \"gs://actea-shared/Shapefiles/LME66/LMEs66.shp\"\n",
    "    return gpd.read_file(lme_file)\n",
    "\n",
    "\n",
    "def get_LME_records_plot():\n",
    "    lme_file = \"gs://actea-shared/Shapefiles/LME66_180/LME66_180.shp\"\n",
    "    return gpd.read_file(lme_file)\n",
    "\n",
    "\n",
    "def create_colors(N):\n",
    "    color = iter(cm.tab20b(np.linspace(0, 1, N)))\n",
    "    return [next(color) for c in range(N)]\n",
    "\n",
    "def get_data_within_LME(ds, var_name, LME, create_maps):\n",
    "    # Extract the polygon defining the boundaries of the LME\n",
    "    print(f\"Working on LME: {LME}\")\n",
    "          \n",
    "    if LME=='Northern Bering - Chukchi Seas':              \n",
    "        shdf = get_LME_records_plot()\n",
    "    elif LME=='Barents Sea':              \n",
    "        shdf = get_LME_records()\n",
    "    else:\n",
    "        raise Exception(f\"Unable to parse LME {LME}\")   \n",
    "    \n",
    "    shdf_sel = shdf[shdf[\"LME_NAME\"] == LME]\n",
    "    \n",
    "    # Rioxarray requires x and y dimensions - we convert these back to lon and lat later.\n",
    "    # We also add the projection (lat-lon) so that rioxarray can do the clipping of the data according to the\n",
    "    # shapefile.\n",
    "    ds = ds.rename({\"x\": \"lon\", \"y\": \"lat\"})\n",
    "    ds = ds.rio.write_crs(4326)\n",
    "    \n",
    "    ds.rio.set_spatial_dims(y_dim=\"lat\", x_dim=\"lon\", inplace=True)\n",
    "    ds.rio.write_crs(4326,  inplace=True)\n",
    "    \n",
    "    clipped = ds.rio.clip(geometries=shdf_sel.geometry.apply(mapping), crs=ds.rio.crs)\n",
    "    if \"x\" in ds.dims:\n",
    "        clipped = clipped.rename({\"x\": \"lon\", \"y\": \"lat\"})\n",
    "    return clipped\n",
    "\n",
    "def create_LME_figure(ax, LMES, projection, show, extent,data_to_contour=None):\n",
    "\n",
    "    ax.add_feature(cfeature.LAND)\n",
    "    ax.add_feature(cfeature.COASTLINE)\n",
    "\n",
    "    ax.set_extent(extent)\n",
    "\n",
    "    # Get the -180-180 projected shapefile containing LMEs to make it\n",
    "    # easy to plot across the Pacific Ocean\n",
    "    shdf = get_LME_records_plot()\n",
    "    colors_rgb=create_colors(len(LMES))\n",
    "    counter=0\n",
    "    for LME_NAME,LME_NUMBER in zip(shdf['LME_NAME'],shdf['LME_NUMBER']):\n",
    "\n",
    "        shdf_sel = shdf[ shdf['LME_NAME']==LME_NAME ]\n",
    "\n",
    "        if (LME_NAME in LMES):\n",
    "           # print(\"Adding geometry for LME {}\".format(LME_NAME))\n",
    "            # Add the geometry and fill it with color\n",
    "            if len(LMES)==1:\n",
    "                color=\"red\"\n",
    "            else:\n",
    "                color=colors_rgb[counter]\n",
    "            ax.add_geometries(shdf_sel['geometry'],\n",
    "                              projection,\n",
    "                              facecolor=color,\n",
    "                              edgecolor='k',\n",
    "                              zorder=8)\n",
    "            if data_to_contour is not None:\n",
    "\n",
    "                ax.contourf(data_to_contour.lon,\n",
    "                            data_to_contour.lat,\n",
    "                            data_to_contour,\n",
    "                            zorder=10,\n",
    "                   cmap=sns.color_palette(\"Spectral_r\", as_cmap=True),\n",
    "                   transform=ccrs.PlateCarree())\n",
    "\n",
    "            # Add the label LME_NUMBER of the selected LME at the center of the LME\n",
    "          #  ax.annotate(s=LME_NUMBER,\n",
    "          #              xy=(shdf_sel.centroid.x,shdf_sel.centroid.y),\n",
    "          #              color=\"white\",\n",
    "          #              fontsize=13)\n",
    "            counter+=1\n",
    "        else:\n",
    "            ax.add_geometries(shdf_sel['geometry'],\n",
    "                              projection,\n",
    "                              facecolor='LightGray',\n",
    "                              edgecolor='k')\n",
    "\n",
    "    if show:\n",
    "        plotfile=\"Figures/CMIP6_lightpaper_map_{}.png\".format(LMES[0])\n",
    "        print(\"Created figure {}\".format(plotfile))\n",
    "        plt.savefig(plotfile, dpi=200,\n",
    "                        bbox_inches = 'tight')\n",
    "        plt.show()\n",
    "\n",
    "def get_data_within_LME(ds,var_name,LME,create_maps, scenario):\n",
    "\n",
    "    print(\"Extracting data within LME {} for {} scenario {}\".format(LME, var_name, scenario))\n",
    "\n",
    "    # Extract the polygon defining the boundaries of the LME\n",
    "    shdf = get_LME_records()\n",
    "    shdf_sel = shdf[ shdf['LME_NAME']==LME ]\n",
    "\n",
    "    # Create the map of the LME boundaries and color it.\n",
    "    # The active LME has color while the others are grey.\n",
    "    if create_maps:\n",
    "        # Setup the figure panels\n",
    "        fig = plt.figure(figsize=(13, 8))\n",
    "        if LME in [\"Barents Sea\",\"Arctic Ocean\"]:\n",
    "            projection=ccrs.NorthPolarStereo() #ccrs.PlateCarree(central_longitude=0)\n",
    "            extent = [-8, 80, 67, 90]\n",
    "        else:\n",
    "            projection=ccrs.PlateCarree(central_longitude=-180)\n",
    "            extent = [-252, -100, 10, 65]\n",
    "            extent = [-220, -135, 30, 85]\n",
    "        ax1 = fig.add_subplot(111, projection=projection)\n",
    "\n",
    "        create_LME_figure(ax1, [LME], ccrs.PlateCarree(central_longitude=-180),True,extent)\n",
    "\n",
    "    # Rioxarray requires x and y dimensions - we convert these back to lon and lat later.\n",
    "    # We also add the projection (lat-lon) so that rioxarray can do the clipping of the data according to the\n",
    "    # shapefile.\n",
    "\n",
    "    tos=ds.rename({'lon': 'x','lat': 'y'})\n",
    "    tos=tos.rio.write_crs(4326)\n",
    "\n",
    "    # Clip the data within the LME. We have to convert the polygon geometry to a geodataframe using\n",
    "    # `shapely.geometry`. The clipping of data within the polygon is done using rioxarray.clip function\n",
    "\n",
    "    clipped = tos.rio.clip(geometries=shdf_sel.geometry.apply(mapping), crs=tos.rio.crs)\n",
    "    clipped=clipped.rename({'x': 'lon','y': 'lat'}) #.to_dataset()\n",
    "\n",
    "    return clipped\n",
    "\n",
    "\n",
    "def get_formatted_ds_within_LME(ds, var_name, LME, scenario):\n",
    "    if var_name in [\"tos_mean\",\"siconc_mean\"]:\n",
    "        ds=ds.assign_coords(lat=ds.y, lon=ds.x)\n",
    "        ds=ds.rename({\"x\":\"lon\",\"y\":\"lat\"})\n",
    "    ds = convert_to_180(ds)\n",
    "    ds = convert_time(ds)\n",
    "\n",
    "    ds_lme = get_data_within_LME(ds, var_name, LME, False, scenario)\n",
    "    ds_lme = xr_add_cyclic_point(ds_lme[var_name])\n",
    "    ds_lme = ds_lme.to_dataset(name=var_name)\n",
    "\n",
    "    # Add land mask\n",
    "    ds_lme[\"mask\"] = create_land_ocean_mask(ds_lme.isel(time=0))\n",
    "    return ds_lme.where(ds_lme.mask == 1)\n",
    "\n",
    "def calculate_daily_uvb_exposure_and_combine_datasets(ds, ds_uvb, ds_par):\n",
    "    # Sum up the four UVB values per day to get daily exposure (kJm-2)\n",
    "    ds[uvb_var] = (ds_uvb[uvb_var] * 3600*4.0)/1000.\n",
    "    ds[par_var] = ds_par[par_var]\n",
    "\n",
    "    return ds.resample(time=\"MS\", skipna=True).sum()\n",
    "\n",
    "def calculate_egg_survival(T, specie):\n",
    "    if specie == \"pacific_cod\":\n",
    "        # Laurel et al., 2020\n",
    "        H = (0.453 / (1 + ((T - 4.192) / 2.125)**2)) * 100\n",
    "    elif specie == \"atlantic_cod\":\n",
    "        # Dahlke et al., 2018\n",
    "        H = -318027.1388 + (318119.8195 / (1 + ((T-4.337) /  467.4799)**2))\n",
    "    elif specie == \"polar_cod\":\n",
    "        # Laurel et al., 2018\n",
    "        H = 87.926 + 1.266*T - 3.6582*T**2\n",
    "    elif specie==\"walleye_pollock\":\n",
    "        # Laurel et al., 2018\n",
    "        H=72.127+5.775*T - 0.801*T**2\n",
    "    H =xr.where(H < 0, 0.0, H)\n",
    "    return H\n",
    "\n",
    "def egg_survival(ds:xr.Dataset, var_name:str, species:str):\n",
    "    \"\"\"Calculate egg survival in surface for different species in percentage given\n",
    "    temperature of water column \"\"\"\n",
    "\n",
    "    ds[\"par_mean\"]=ds[\"par_mean\"]*np.exp(-depth*attenuation)\n",
    "    ds[\"areacello_par_masked\"]= xr.where(ds[\"par_mean\"] >= par_threshold, ds[\"areacello\"], np.nan)\n",
    "    ds[\"tos_mean_par_masked\"]= xr.where(ds[\"par_mean\"] >= par_threshold, ds[\"tos_mean\"], np.nan)\n",
    "\n",
    "    for specie in species:\n",
    "        print(\"Calculating survival for {}\".format(specie))\n",
    "        T = ds[\"tos_mean\"]\n",
    "        H = calculate_egg_survival(T, specie)\n",
    "\n",
    "        ds[\"{}_egg_par_masked\".format(specie)]=(('time', 'lat', 'lon'),  H.data)\n",
    "    return ds\n",
    "\n",
    "def calculate_areacello(ds, model_obj, var_name, project_name):\n",
    "\n",
    "        # Calculate the area based on the longitude - latitude\n",
    "        if ds.lon.ndim == 2:\n",
    "            lon = ds.lon.values[0, :]\n",
    "            lat = ds.lat.values[:, 0]\n",
    "        else:\n",
    "            lon = ds.lon.values\n",
    "            lat = ds.lat.values\n",
    "\n",
    "        ds_singletime = ds.isel(time=0)\n",
    "\n",
    "        # Convert the dataset to a cube as this adds correct units required by iris\n",
    "        sdiris = CMIP6_downscale_iris.CMIP6_downscale_iris()\n",
    "        cube = sdiris.ds_to_iris(ds_singletime,\n",
    "                                 var_name,\n",
    "                                 model_obj,\n",
    "                                 project_name, prefix=\"cube\")\n",
    "\n",
    "        # Calculate the areacello for the grid and convert the result to km2\n",
    "        # Uses iris area_weights function.\n",
    "        # https://scitools.org.uk/iris/docs/v2.4.0/iris/iris/analysis/cartography.html#iris.analysis.cartography.area_weights\n",
    "        m2_to_km2 = 1.0e-6\n",
    "        area_ends = (iris.analysis.cartography.area_weights(cube, normalize=False)) * m2_to_km2\n",
    "\n",
    "        # Now convert the numpy array of areas to a dataset with the same dimension as the siconc\n",
    "        area_ds = xr.DataArray(name=\"areacello\",\n",
    "                               data=area_ends,\n",
    "                               coords={\"lat\": lat,\n",
    "                                       \"lon\": lon},\n",
    "                               dims=[\"lat\", \"lon\"]).to_dataset()\n",
    "\n",
    "        # Convert the resulting dataset to an iris cube\n",
    "        area_cube = sdiris.ds_to_iris(area_ds,\n",
    "                                      \"areacello\",\n",
    "                                      model_obj,\n",
    "                                      project_name,\n",
    "                                      prefix=\"areacello_\")\n",
    "\n",
    "        # Fix the coordinates so that we add geographic information to the cube,\n",
    "        # before saving the cube to the siconc dataset\n",
    "        area_cube = sdiris.fix_coordinates_cube(area_cube/1.e6)\n",
    "\n",
    "        return xr.DataArray.from_iris(area_cube)\n",
    "\n",
    "\n",
    "def calculate_juvenile_growth(T, specie):\n",
    "    if specie == \"atlantic_cod\":\n",
    "        # Bjørnsson et al., 2007\n",
    "        G = -0.2425 + 0.1519 * T + 0.0552*T**2 - 0.002931*T**3\n",
    "       \n",
    "    elif specie == \"polar_cod\":\n",
    "        # Laurel et al., 2016\n",
    "        G = 0.8290 + 0.1638 * T - 0.0054* T**2 - 0.0005 * T**3\n",
    "       \n",
    "    elif specie == \"pacific_cod\":\n",
    "        G = 0.2494 + 0.3216 * T - 0.0069 * T ** 2 - 0.0004 * T ** 3\n",
    "       \n",
    "    elif specie == \"walleye_pollock\":\n",
    "        G = 0.2023 + 0.0992 * T + 0.0335 * T ** 2 - 0.0019 * T ** 3\n",
    "    return G\n",
    "\n",
    "def cumulative_juvenile_growth_potential(ds:xr.Dataset, var_name:str, species:str):\n",
    "    \"\"\"Calculate larval survival in surface for different species in percentage given\n",
    "    temperature of water column\"\"\"\n",
    "\n",
    "    ds_daily = ds.resample(time=\"1D\", skipna=True).interpolate('linear')\n",
    "    ds_daily[\"par_mean\"]=ds_daily[\"par_mean\"]*np.exp(-depth*attenuation)\n",
    "\n",
    "    ds_daily[\"tos_mean\"]= xr.where(ds_daily[\"par_mean\"] >= par_threshold, ds_daily[\"tos_mean\"], np.nan)\n",
    "    ds_daily[\"areacello_par_masked\"]= xr.where(ds_daily[\"par_mean\"] >= par_threshold, ds_daily[\"areacello\"], np.nan)\n",
    "    ds_daily[\"tos_mean_par_masked\"]= xr.where(ds_daily[\"par_mean\"] >= par_threshold, ds_daily[\"tos_mean\"], np.nan)\n",
    "\n",
    "    ds_daily=ds_daily.transpose(\"time\",\"lat\",\"lon\")\n",
    "\n",
    "    print(\"total sum 1 : {}\".format(ds_daily[\"tos_mean\"].count(dim=[\"lat\",\"lon\"]).sum().values))\n",
    "    print(\"total sum 2 : {}\".format(ds_daily[\"tos_mean_par_masked\"].count(dim=[\"lat\",\"lon\"]).sum().values))\n",
    "\n",
    "    for specie in species:\n",
    "        print(\"Calculating growth potential for {}\".format(specie))\n",
    "        T = ds_daily[var_name]\n",
    "        T_par = ds_daily[\"tos_mean_par_masked\"]\n",
    "\n",
    "        G = calculate_juvenile_growth(T, specie)\n",
    "        G_par = calculate_juvenile_growth(T_par, specie)\n",
    "\n",
    "        ds_daily[\"{}_juvenile\".format(specie)]=(('time', 'lat', 'lon'),  G.data)\n",
    "        ds_daily[\"{}_juvenile_par_masked\".format(specie)]=(('time', 'lat', 'lon'),  G_par.data)\n",
    "\n",
    "    return ds_daily.resample(time=\"MS\", skipna=True).mean()\n",
    "\n",
    "\n",
    "\n",
    "def calculate_combined_uvb_scaled_survival_and_combine_datasets(ds:xr.Dataset,\n",
    "                                                                ds_uvb_lme:xr.Dataset,\n",
    "                                                                ds_sic_lme:xr.Dataset,\n",
    "                                                                ds_par_lme:xr.Dataset,\n",
    "                                                                species:str,\n",
    "                                                                seaice_minimum:int):\n",
    "\n",
    "    regridder = xe.Regridder(ds_uvb_lme,ds, \"bilinear\",\n",
    "                                             periodic=False,\n",
    "                                             extrap_method='inverse_dist',\n",
    "                                             extrap_num_src_pnts=2,\n",
    "                                             extrap_dist_exponent=1,\n",
    "                                             ignore_degenerate=True)\n",
    "\n",
    "    regridder_sic = xe.Regridder(ds_sic_lme,ds, \"bilinear\",\n",
    "                                             periodic=False,\n",
    "                                             extrap_method='inverse_dist',\n",
    "                                             extrap_num_src_pnts=2,\n",
    "                                             extrap_dist_exponent=1,\n",
    "                                             ignore_degenerate=True)\n",
    "\n",
    "    ds[uvb_var]=(('time', 'lat', 'lon'), regridder(ds_uvb_lme[uvb_var]).data)\n",
    "    ds[par_var]=(('time', 'lat', 'lon'), regridder(ds_par_lme[par_var]).data)\n",
    "    ds[\"siconc_mean\"]=(('time', 'lat', 'lon'), regridder_sic(ds_sic_lme[\"siconc_mean\"]).data)\n",
    "\n",
    "    \"\"\"\n",
    "    Filter out only areas where minimum light threshold is adequate for feeding. any other regions will\n",
    "    not satisfy survival for larval stages.\n",
    "    \"\"\"\n",
    "    ds[uvb_var]= xr.where(ds[\"par_mean\"] >= par_threshold, ds[uvb_var], np.nan)\n",
    "    ds[par_var]= xr.where(ds[\"par_mean\"] >= par_threshold, ds[par_var], np.nan)\n",
    "\n",
    "    \"\"\"\n",
    "    Calculate the monthly relative change in uvb light for the LMe and provide as timeseries\n",
    "    \"\"\"\n",
    "    clim_uvb = ds[uvb_var].sel(time=(slice(\"1980-01-01\",\"2000-12-31\"))).groupby(\"time.month\").mean(dim=\"time\")\n",
    "    clim_par = ds[par_var].sel(time=(slice(\"1980-01-01\",\"2000-12-31\"))).groupby(\"time.month\").mean(dim=\"time\")\n",
    "\n",
    "    ds[\"uvb_ice\"] = (ds[uvb_var].groupby(\"time.month\") - clim_uvb)/clim_uvb.max(dim=\"month\")\n",
    "    ds[\"par_ice\"] = (ds[par_var].groupby(\"time.month\") - clim_par)/clim_par.max(dim=\"month\")\n",
    "\n",
    "    #xr.where(((ds[\"siconc_mean\"] > seaice_minimum) ), ds[uvb_var], np.nan)\n",
    "   # ds[\"uvb_ice\"]=xr.where(((np.isnan(ds[\"siconc_mean\"]) | (ds[\"siconc_mean\"] == 0.0))), ds[\"uvb_ice\"], np.nan)\n",
    "   # ds[\"par_ice\"]=xr.where(((np.isnan(ds[\"siconc_mean\"]) | (ds[\"siconc_mean\"] == 0.0))), ds[\"par_ice\"], np.nan)\n",
    "\n",
    "    return ds\n",
    "\n",
    "def q5(x):\n",
    "    return x.quantile(0.05)\n",
    "def q95(x):\n",
    "    return x.quantile(0.95)\n",
    "\n",
    "def plot_averaged_growth_or_survival(df: pd.DataFrame,\n",
    "                                     df_survival: pd.DataFrame,\n",
    "                                     specie: str,\n",
    "                                     rolling_mean_years: int,\n",
    "                                     egg_or_larva: str, LME: str, dfs_all_seasons, season):\n",
    "\n",
    "    # This plotting method takes the SSP245 and SSP585 dataframe for survival/growth, temperature,\n",
    "    # and uvb light and plots each of these variables. Dataframe is in monthly values but an aggregator\n",
    "    # function calculates the annual average mean and quantiles within each year fulfilling the months requirement.\n",
    "    # For larval and juvenile stages we consider only the summer months e.g. months=[5,6,7,8]\n",
    "\n",
    "    sns.set(style=\"white\", rc={\"axes.facecolor\": (0, 0, 0, 0)}, font_scale=1.5)\n",
    "    f=plt.figure(figsize=(12, 6))\n",
    "    f.clf()\n",
    "\n",
    "    gs = f.add_gridspec(1, 3)\n",
    "    ax = f.add_subplot(gs[0, 0])\n",
    "    ax2 = f.add_subplot(gs[0, 1])\n",
    "    ax3 = f.add_subplot(gs[0, 2])\n",
    "\n",
    "    sns.set_palette(\"tab10\")\n",
    "\n",
    "    # Organize the data for dataframe survival --------\n",
    "    df_survival[\"time_bin\"] = pd.to_datetime(df_survival.time)\n",
    "    bins=np.arange(1979,2100,1)\n",
    "\n",
    "    dfs245=df_survival[df_survival[\"scenario\"]==\"ssp245\"]\n",
    "    dfs585=df_survival[df_survival[\"scenario\"]==\"ssp585\"]\n",
    "\n",
    "    # Growth or survival in areas adequate for feeding\n",
    "    dfs245['time_bin'] = pd.cut(dfs245.time_bin.dt.year, bins, right=False)\n",
    "    if egg_or_larva== \"egg\":\n",
    "        specie_var=\"{}_egg_par_masked\".format(specie)\n",
    "    else:\n",
    "        specie_var=\"{}_juvenile_par_masked\".format(specie)\n",
    "      \n",
    "    dfgs245 = dfs245.groupby('time_bin', as_index=False)[specie_var].agg([q5,q95,\"mean\"])\n",
    "    dfgs245[\"proper_time\"]=[datetime.datetime(1980,6,15)+relativedelta(years=i) for i in range(0,120)]\n",
    "\n",
    "    dfs585['time_bin'] = pd.cut(dfs585.time_bin.dt.year, bins, right=False)\n",
    "    dfgs585 = dfs585.groupby('time_bin', as_index=False)[specie_var].agg([q5,q95,\"mean\"])\n",
    "    dfgs585[\"proper_time\"]=[datetime.datetime(1980,6,15)+relativedelta(years=i) for i in range(0,120)]\n",
    "\n",
    "    # Areacello where PAR above threshold for feeding\n",
    "    if egg_or_larva== \"egg\":\n",
    "        sel_var = \"siconc_mean\"\n",
    "    else:\n",
    "        sel_var = \"areacello_par_masked\"\n",
    "\n",
    "    dftos245 = dfs245.groupby('time_bin', as_index=False)[\"tos_mean\"].agg([q5,q95,\"mean\"])\n",
    "    dftos245[\"proper_time\"]=[datetime.datetime(1980,6,15)+relativedelta(years=i) for i in range(0,120)]\n",
    "\n",
    "    dftos585 = dfs585.groupby('time_bin', as_index=False)[\"tos_mean\"].agg([q5,q95,\"mean\"])\n",
    "    dftos585[\"proper_time\"]=[datetime.datetime(1980,6,15)+relativedelta(years=i) for i in range(0,120)]\n",
    "\n",
    "    # Ocean UVB under ice - not used as very small values\n",
    "    dfarea_or_ice245 = dfs245.groupby('time_bin', as_index=False)[sel_var].agg([q5,q95,\"mean\"])\n",
    "    dfarea_or_ice245[\"proper_time\"]=[datetime.datetime(1980,6,15)+relativedelta(years=i) for i in range(0,120)]\n",
    "\n",
    "    dfarea_or_ice585 = dfs585.groupby('time_bin', as_index=False)[sel_var].agg([q5,q95,\"mean\"])\n",
    "    dfarea_or_ice585[\"proper_time\"]=[datetime.datetime(1980,6,15)+relativedelta(years=i) for i in range(0,120)]\n",
    "\n",
    "    alpha=0.2\n",
    "    tickfontsize=22\n",
    "\n",
    "    # SSP245 growth/survival\n",
    "    ax3.fill_between(dfgs245['proper_time'],\n",
    "                    dfgs245['q5'].rolling(rolling_mean_years).mean(),\n",
    "                    dfgs245['q95'].rolling(rolling_mean_years).mean(),\n",
    "                    alpha=alpha,color = 'tab:blue')\n",
    "\n",
    "    ax3.plot(dfgs245[\"proper_time\"],\n",
    "            dfgs245['mean'].rolling(rolling_mean_years).mean(),\n",
    "            color = 'tab:blue', alpha = 1.0, linewidth = 4)\n",
    "\n",
    "    # SSP585 growth/survival\n",
    "    ax3.fill_between(dfgs585['proper_time'],\n",
    "                    dfgs585['q5'].rolling(rolling_mean_years).mean(),\n",
    "                    dfgs585['q95'].rolling(rolling_mean_years).mean(),\n",
    "                    alpha=alpha,color = 'tab:orange')\n",
    "\n",
    "    ax3.plot(dfgs585[\"proper_time\"],\n",
    "            dfgs585['mean'].rolling(rolling_mean_years).mean(),\n",
    "            color = 'tab:orange', alpha = 1.0, linewidth = 4)\n",
    "\n",
    "    # SSP245 ocean temperature\n",
    "    ax.fill_between(dftos245['proper_time'],\n",
    "                    dftos245['q5'].rolling(rolling_mean_years).mean(),\n",
    "                    dftos245['q95'].rolling(rolling_mean_years).mean(),\n",
    "                    alpha=alpha,\n",
    "                    color = 'tab:blue')\n",
    "\n",
    "    ax.plot(dftos245[\"proper_time\"],\n",
    "             dftos245['mean'].rolling(rolling_mean_years).mean(),\n",
    "         #    linestyle=\"--\",\n",
    "             color = 'tab:blue',\n",
    "             alpha = 1.0, linewidth = 4)\n",
    "\n",
    "    # SSP585 TOS\n",
    "\n",
    "    ax.fill_between(dftos585['proper_time'],\n",
    "                    dftos585['q5'].rolling(rolling_mean_years).mean(),\n",
    "                    dftos585['q95'].rolling(rolling_mean_years).mean(),\n",
    "                    alpha=alpha,color = 'tab:orange')\n",
    "\n",
    "\n",
    "    ax.plot(dfgs585[\"proper_time\"],\n",
    "             dftos585['mean'].rolling(rolling_mean_years).mean(),\n",
    "             #linestyle=\"--\",\n",
    "             color = 'tab:orange',\n",
    "             alpha = 1.0, linewidth = 4)\n",
    "   \n",
    "    ylabels = ['{:,.2f}{}'.format(y,u\"$^{o}C$\") for y in ax.get_yticks()]\n",
    "    ax.set_yticklabels(ylabels)\n",
    "\n",
    "#    ylabels = ['{:.0e}{}'.format(y,u\"\\N{DEGREE SIGN}C\") for y in ax3.get_yticks()]\n",
    "  #  if egg_or_larva== \"egg\":\n",
    "#    ylabels = ['{:,.2f}{}'.format(y,u\"\\N{DEGREE SIGN}C\") for y in ax3.get_yticks()]\n",
    "   # else:\n",
    "   #     ylabels = ['{:,.0f}{}'.format(y,u\"km$^{2}$\") for y in ax3.get_yticks()]\n",
    "   # ax3.set_yticklabels(ylabels)\n",
    "    plt.legend(loc=\"upper right\", frameon=False, fontsize=24)\n",
    "\n",
    "    # https://reckoningrisk.com/watercolor-regression/\n",
    "    df[\"time_bin\"] = pd.to_datetime(df.time)\n",
    "    df245=df[df[\"scenario\"]==\"ssp245\"]\n",
    "    df585=df[df[\"scenario\"]==\"ssp585\"]\n",
    "\n",
    "    # Plot PAR for the selected months\n",
    "    df245['time_bin'] = pd.cut(df245.time_bin.dt.year, bins, right=False)\n",
    "    dfg245 = df245.groupby('time_bin', as_index=False)[\"par_ice\"].agg([q5,q95,\"mean\"])\n",
    "    dfg245[\"proper_time\"]=[datetime.datetime(1980,6,15)+relativedelta(years=i) for i in range(0,120)]\n",
    "\n",
    "    df585['time_bin'] = pd.cut(df585.time_bin.dt.year, bins, right=False)\n",
    "    dfg585 = df585.groupby('time_bin', as_index=False)[\"par_ice\"].agg([q5,q95,\"mean\"])\n",
    "    dfg585[\"proper_time\"]=[datetime.datetime(1980,6,15)+relativedelta(years=i) for i in range(0,120)]\n",
    "\n",
    "    ax2.fill_between(dfarea_or_ice245['proper_time'],\n",
    "                     dfarea_or_ice245['q5'].rolling(rolling_mean_years).mean(),\n",
    "                     dfarea_or_ice245['q95'].rolling(rolling_mean_years).mean(),\n",
    "                     alpha=alpha,color = 'tab:blue')\n",
    "\n",
    "    ax2.plot(dfarea_or_ice245[\"proper_time\"],\n",
    "                  dfarea_or_ice245['mean'].rolling(rolling_mean_years).mean(),\n",
    "                  color = 'tab:blue', alpha = 1.0, linewidth = 4)\n",
    "\n",
    "    ax2.fill_between(dfarea_or_ice585['proper_time'],\n",
    "                     dfarea_or_ice585['q5'].rolling(rolling_mean_years).mean(),\n",
    "                     dfarea_or_ice585['q95'].rolling(rolling_mean_years).mean(),\n",
    "                     alpha=alpha,color = 'tab:orange')\n",
    "\n",
    "    ax2.plot(dfarea_or_ice585['proper_time'],\n",
    "             dfarea_or_ice585['mean'].rolling(rolling_mean_years).mean(),\n",
    "             color = 'tab:orange', alpha = 1.0, linewidth = 4)\n",
    "\n",
    "    ax2.tick_params(labelsize=tickfontsize)\n",
    "    ax2.set_ylabel(\"\") #\"kJm$^{-2}$\", fontsize=24)\n",
    "    if sel_var == \"siconc_mean\":\n",
    "    \n",
    "        ylabels =  ['{:,.1f}{}'.format(y,u\"%\") for y in ax2.get_yticks()]\n",
    "    else:\n",
    "        ylabels =  ['{:,.2f}{}'.format(y,u\"\") for y in ax2.get_yticks()]\n",
    "    \n",
    "    ax2.set_yticklabels(ylabels)\n",
    "    f.canvas.draw()\n",
    "\n",
    "    # Before you can edit and rotate the labels you have to draw them\n",
    "    plt.draw()\n",
    "    ylabels =  ['{:,.1f}{}'.format(y,u\"%\") for y in ax3.get_yticks()]\n",
    "    ax3.set_yticklabels(ylabels)\n",
    "    \n",
    "    ax.set_xticklabels(ax.get_xticklabels(), rotation=-90, ha='right', fontsize=16)\n",
    "    ax2.set_xticklabels(ax2.get_xticklabels(), rotation=-90, ha='right',fontsize=16)\n",
    "    ax3.set_xticklabels(ax3.get_xticklabels(), rotation=-90, ha='right',fontsize=16)\n",
    "    ax.tick_params(labelsize=tickfontsize)\n",
    "    ax3.tick_params(labelsize=tickfontsize)\n",
    "    ax2.tick_params(labelsize=tickfontsize)\n",
    "\n",
    "    if not os.path.exists(\"Figures\"):\n",
    "        os.makedirs(\"Figures\")\n",
    "    plotfile=\"Figures/{}_uvb_{}_{}.png\".format(egg_or_larva, LME, specie)\n",
    "    print(\"Created figure {}\".format(plotfile))\n",
    "    gs.tight_layout(f)\n",
    "\n",
    "    plt.savefig(plotfile, dpi=300,\n",
    "                bbox_inches = 'tight')\n",
    "    plt.show()\n",
    "\n",
    "    # Store the growth and survival data\n",
    "    print(f\"Saving data to dictionary: {egg_or_larva}_{season}_{specie}_ssp245_{LME}\")\n",
    "    dfs_all_seasons[f\"{egg_or_larva}_{season}_{specie}_ssp245_{LME}\"]=dfgs245\n",
    "    dfs_all_seasons[f\"{egg_or_larva}_{season}_{specie}_ssp585_{LME}\"]=dfgs585\n",
    "\n",
    "    return dfs_all_seasons"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "source": [
    "LMES = ['Barents Sea','Northern Bering - Chukchi Seas'] #,'Central Arctic']\n",
    "\n",
    "SPECIES_PER_LME = [[\"polar_cod\",\"atlantic_cod\"],[\"polar_cod\",\"walleye_pollock\"]]\n",
    "\n",
    "create_ensemble = True\n",
    "scenarios=[\"ssp245\",\"ssp585\"]\n",
    "var_names=[\"par_\"] #\"uv_\",\"uvi_\",\"par_\", \"uvb_\",\"_tos.nc\",\"_siconc.nc\"]\n",
    "\n",
    "start_time = \"1979-01-01\"\n",
    "end_time = \"2099-12-16\"\n",
    "spring_months=[4,5,6]\n",
    "winter_months=[1,2,3]\n",
    "\n",
    "rolling_mean_years=5\n",
    "seaice_minimum=15\n",
    "dfs_all_seasons = {}\n",
    "\n",
    "for months in [winter_months,spring_months]:\n",
    "    for LME, species in zip(LMES, SPECIES_PER_LME):\n",
    "        if months==winter_months:\n",
    "            season=\"winter\"\n",
    "        else:\n",
    "            season=\"spring\"\n",
    "            \n",
    "        if LME in [\"Barents Sea\",\"Arctic Ocean\"]:\n",
    "            projection=ccrs.NorthPolarStereo()\n",
    "            extent = [-10, 80, 67, 85]\n",
    "        else:\n",
    "            projection=ccrs.PlateCarree(central_longitude=-180)\n",
    "            extent = [-200, -145, 50, 85]\n",
    "\n",
    "        dfs=[]\n",
    "        dfs_survival=[]\n",
    "        uvb_var=\"uvb_mean\"\n",
    "        par_var=\"par_mean\"\n",
    "\n",
    "        for scenario in scenarios:\n",
    "\n",
    "            ensemble_path = \"light/ncfiles/ensemble\"   \n",
    "            lightpath = f\"light/{scenario}/ensemble\"\n",
    "                    \n",
    "            io = CMIP6_IO.CMIP6_IO()\n",
    "            tos = io.open_dataset_on_gs(f\"{lightpath}/CMIP6_ensemble_stats_{scenario}_tos_weighted.nc\")\n",
    "            sic = io.open_dataset_on_gs(f\"{lightpath}/CMIP6_ensemble_stats_{scenario}_siconc_weighted.nc\")\n",
    "            uvb = io.open_dataset_on_gs(f\"{ensemble_path}/CMIP6_ensemble_stats_{scenario}_uvb_weighted.nc\")\n",
    "            par = io.open_dataset_on_gs(f\"{ensemble_path}/CMIP6_ensemble_stats_{scenario}_par_weighted.nc\")\n",
    "\n",
    "            par = par.sel(time=slice(start_time,end_time))\n",
    "            tos = tos.sel(time=slice(start_time,end_time))\n",
    "            uvb = uvb.sel(time=slice(start_time,end_time))\n",
    "            sic = sic.sel(time=slice(start_time,end_time))\n",
    "\n",
    "            ds_uvb_lme = get_formatted_ds_within_LME(uvb, uvb_var, LME, scenario)\n",
    "            ds_par_lme = get_formatted_ds_within_LME(par, par_var, LME, scenario)\n",
    "            ds_sic_lme = get_formatted_ds_within_LME(sic, \"siconc_mean\", LME, scenario).resample(time=\"MS\", skipna=True).mean()\n",
    "            ds_tos_lme = get_formatted_ds_within_LME(tos, \"tos_mean\", LME, scenario).resample(time=\"MS\", skipna=True).mean()\n",
    "\n",
    "            ds_sic_lme = xr.where(ds_sic_lme <= 0, np.nan, ds_sic_lme)\n",
    "            ds_tos_lme = xr.where(ds_tos_lme <= -2.0, np.nan, ds_tos_lme)\n",
    "\n",
    "            ds_tos_lme = xr.where(((ds_tos_lme < 1.e-20) | (ds_tos_lme > 1e20)), np.nan, ds_tos_lme)\n",
    "            ds_sic_lme = xr.where(((ds_sic_lme < 1.e-20) | (ds_sic_lme > 1e20)), np.nan, ds_sic_lme)\n",
    "            ds_par_lme = xr.where(((ds_par_lme < 1.e-20) | (ds_par_lme > 1e20)), np.nan, ds_par_lme)\n",
    "            ds_uvb_lme = xr.where(((ds_uvb_lme < 1.e-20) | (ds_uvb_lme > 1e20)), np.nan, ds_uvb_lme)\n",
    "\n",
    "            # Select the time period\n",
    "            ds_uvb_lme = ds_uvb_lme.sel(lat=slice(50,91))\n",
    "            ds_par_lme = ds_par_lme.sel(lat=slice(50,91))\n",
    "            ds_tos_lme = ds_tos_lme.sel(lat=slice(50,91))\n",
    "            ds_sic_lme = ds_sic_lme.sel(lat=slice(50,91))\n",
    "\n",
    "            ds_tos_lme[\"mask\"] = create_land_ocean_mask(ds_tos_lme.isel(time=0))\n",
    "            ds_tos_lme = ds_tos_lme.where(ds_tos_lme.mask == 1)\n",
    "            model_obj = CMIP6_MODEL(\"ensemble\")\n",
    "\n",
    "            ds_tos_lme[\"areacello\"] = calculate_areacello(ds_tos_lme, model_obj, \"tos_mean\", project_name=\"light\")\n",
    "            # Clip the area to the polygon - this also clips to the time varying var_name# which results in\n",
    "            # time varying areacello variable that we can use to sum up the size of open water by\n",
    "            # looking at the annual mean.\n",
    "            # First we set all areas outside of the polygon to nan - which also now includes\n",
    "            # areacello which as part of teh calculations actually covered the whole region (and\n",
    "            # not just the polygon).\n",
    "            ds_tos_lme = xr.where(np.isnan(ds_tos_lme[\"tos_mean\"]), np.nan, ds_tos_lme)\n",
    "            total_area_lme = np.nansum(ds_tos_lme[\"areacello\"].mean(dim=\"time\").values)\n",
    "            print(\"[calculate_areacello] Average total area {:,.2f} km2\".format(total_area_lme))\n",
    "\n",
    "            ds_uvb_lme = ds_uvb_lme.resample(time=\"MS\", skipna=True).mean()\n",
    "            ds_sic_lme = ds_sic_lme.resample(time=\"MS\", skipna=True).mean()\n",
    "            ds_par_lme = ds_par_lme.resample(time=\"MS\", skipna=True).mean()\n",
    "            ds_uvb_lme[\"mask\"]=(('lat', 'lon'), ds_uvb_lme[\"mask\"].isel(time=0).data)\n",
    "            ds_sic_lme[\"mask\"]=(('lat', 'lon'), ds_sic_lme[\"mask\"].isel(time=0).data)\n",
    "            ds_par_lme[\"mask\"]=(('lat', 'lon'), ds_par_lme[\"mask\"].isel(time=0).data)\n",
    "\n",
    "            ds = calculate_daily_uvb_exposure_and_combine_datasets(ds_tos_lme, ds_uvb_lme, ds_par_lme)\n",
    "\n",
    "            # Light (par, uvb) and forcing (tos, sic) are on slightly different grids. We combine all\n",
    "            # datasets to a common grid.\n",
    "            ds[\"mask\"] = create_land_ocean_mask(ds.isel(time=0))\n",
    "\n",
    "            ds = calculate_combined_uvb_scaled_survival_and_combine_datasets(ds, ds_uvb_lme, ds_sic_lme, ds_par_lme, species, seaice_minimum)\n",
    "\n",
    "            # Calculate the survival potential for all species\n",
    "            if months[0]>=4:\n",
    "                larva_or_egg= \"juvenile\"\n",
    "                ds = cumulative_juvenile_growth_potential(ds, \"tos_mean\", species)\n",
    "            else:\n",
    "                larva_or_egg= \"egg\"\n",
    "                ds = egg_survival(ds, \"tos_mean\", species)\n",
    "             \n",
    "\n",
    "            # Important to set zeros to nan - the zeroes come after performing resample\n",
    "           # ds = xr.where(ds <= 0, np.nan, ds)\n",
    "\n",
    "            ds_avg=ds.sel(time=ds.time.dt.month.isin(months)).mean({\"lat\",\"lon\"})\n",
    "\n",
    "            ds_avg = ds_avg.assign(areacello_par_masked = ds[\"areacello_par_masked\"].sum({\"lat\",\"lon\"}))\n",
    "\n",
    "\n",
    "            df = ds_avg.to_dataframe() #.resample(\"A\").mean()\n",
    "            df_survival = ds_avg.to_dataframe() #.resample(\"A\").mean()\n",
    "\n",
    "            df = df.reset_index(level=[0])\n",
    "            df_survival = df_survival.reset_index(level=[0])\n",
    "\n",
    "            df[\"scenario\"]=scenario\n",
    "            df_survival[\"scenario\"]=scenario\n",
    "\n",
    "            dfs.append(df)\n",
    "            dfs_survival.append(df_survival)\n",
    "\n",
    "        # Concatenate all scenarios into one dataframe for this LME\n",
    "        df_all = pd.concat(dfs)\n",
    "        df_all_survival = pd.concat(dfs_survival)\n",
    "\n",
    "        # Plot the results for this LME\n",
    "        for specie in species:\n",
    "            dfs_all_seasons = plot_averaged_growth_or_survival(df_all, df_all_survival, specie, rolling_mean_years,\n",
    "                                             larva_or_egg, LME, dfs_all_seasons, season)\n",
    "print(dfs_all_seasons)"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "source": [
    "def plot_habitat_suitability(df_habitat_245, df_habitat_585, specie):\n",
    "        sns.set(style=\"white\", rc={\"axes.facecolor\": (0, 0, 0, 0)}, font_scale=1.5)\n",
    "        fig, ax = plt.subplots()\n",
    "        sns.set_palette(\"tab10\")\n",
    "\n",
    "        alpha=0.2\n",
    "        tickfontsize=22\n",
    "\n",
    "        # SSP245 habitat suitability\n",
    "        ax.fill_between(df_habitat_245['proper_time'],\n",
    "                        df_habitat_245['q5'].rolling(rolling_mean_years).mean(),\n",
    "                        df_habitat_245['q95'].rolling(rolling_mean_years).mean(),\n",
    "                        alpha=alpha,color = 'tab:blue')\n",
    "\n",
    "        ax.plot(df_habitat_245[\"proper_time\"],\n",
    "                df_habitat_245['mean'].rolling(rolling_mean_years).mean(),\n",
    "                color = 'tab:blue', alpha = 1.0, linewidth = 4)\n",
    "\n",
    "        # SSP585 growth/survival\n",
    "        ax.fill_between(df_habitat_585['proper_time'],\n",
    "                        df_habitat_585['q5'].rolling(rolling_mean_years).mean(),\n",
    "                        df_habitat_585['q95'].rolling(rolling_mean_years).mean(),\n",
    "                        alpha=alpha,color = 'tab:orange')\n",
    "\n",
    "        ax.plot(df_habitat_585[\"proper_time\"],\n",
    "                df_habitat_585['mean'].rolling(rolling_mean_years).mean(),\n",
    "                color = 'tab:orange', alpha = 1.0, linewidth = 4)\n",
    "\n",
    "        ax.tick_params(labelsize=tickfontsize)\n",
    "        ax.set_ylabel(\"\") #\"kJm$^{-2}$\", fontsize=24)\n",
    "        ylabels = ['{:,.1%}'.format(y) for y in ax.get_yticks()]\n",
    "        ax.set_yticklabels(ylabels)\n",
    "        fig.canvas.draw()\n",
    "\n",
    "        # Before you can edit and rotate the labels you have to draw them\n",
    "        plt.draw()\n",
    "        ax.set_xticklabels(ax.get_xticklabels(), rotation=-90, ha='right', fontsize=18)\n",
    "        ax.tick_params(labelsize=tickfontsize)\n",
    "        \n",
    "        if not os.path.exists(\"Figures\"):\n",
    "                os.makedirs(\"Figures\")\n",
    "        plotfile=\"Figures/habitat_suitability_{}_{}.png\".format(LME, specie)\n",
    "        print(f\"Created figure {plotfile}\")\n",
    "       \n",
    "        plt.savefig(plotfile, dpi=300,\n",
    "                        bbox_inches = 'tight')\n",
    "        plt.show()\n",
    "\n",
    "\n",
    "\n",
    "import math\n",
    "from operator import itemgetter\n",
    "\n",
    "for LME, species in zip(LMES, SPECIES_PER_LME):\n",
    "\n",
    "        for specie in species:\n",
    "                # Find maximum of calculate_juvenile_growth at all temperatures. This is used to normalize the data \n",
    "                # found earlier. For egg survival, the maximum survival will always be 100 so we use that.\n",
    "                temperature_range = list(np.arange(-2,30,0.1))\n",
    "        \n",
    "                result = max(((e, calculate_juvenile_growth(e, specie)) for e in temperature_range), key=itemgetter(1))\n",
    "                print(f\"Maximum for function calculate_juvenile_growth {specie} - {result[0]:3.3f} growth rate:  {result[1]:3.3f}\")\n",
    "        \n",
    "        #   print(dfs_all_seasons.keys())\n",
    "                \n",
    "                df_winter_245 = dfs_all_seasons[f\"egg_winter_{specie}_ssp245_{LME}\"]\n",
    "                df_winter_585 = dfs_all_seasons[f\"egg_winter_{specie}_ssp585_{LME}\"]\n",
    "                df_spring_245 = dfs_all_seasons[f\"juvenile_spring_{specie}_ssp245_{LME}\"]\n",
    "                df_spring_585 = dfs_all_seasons[f\"juvenile_spring_{specie}_ssp585_{LME}\"]\n",
    "                df_habitat_245 = df_spring_245\n",
    "                df_habitat_585 = df_spring_585\n",
    "\n",
    "                for k in ['q5','q95','mean']:\n",
    "                        \n",
    "                        df_winter_245[k] = dfs_all_seasons[f\"egg_winter_{specie}_ssp245_{LME}\"][k]/100.\n",
    "                        df_winter_585[k] = dfs_all_seasons[f\"egg_winter_{specie}_ssp585_{LME}\"][k]/100.\n",
    "                        \n",
    "                        df_spring_245[k] = dfs_all_seasons[f\"juvenile_spring_{specie}_ssp245_{LME}\"][k]/float(result[1])\n",
    "                        df_spring_585[k] = dfs_all_seasons[f\"juvenile_spring_{specie}_ssp585_{LME}\"][k]/float(result[1])\n",
    "                        \n",
    "                        df_habitat_245[k] = df_spring_245[k] * df_winter_245[k]\n",
    "                        df_habitat_585[k] = df_spring_585[k] * df_winter_585[k]\n",
    "                \n",
    "                print(f\"{df_habitat_245.head(4)}\")\n",
    "\n",
    "\n",
    "                plot_habitat_suitability(df_habitat_245, df_habitat_585, specie)"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "source": [
    "SPECIES_PER_LME"
   ],
   "outputs": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "source": [],
   "outputs": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gcs2",
   "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.9.7"
  },
  "vscode": {
   "interpreter": {
    "hash": "ce65ac3c58bbfe9a8d7f0a7a880a285b16a99b928790773caf2c61080a97596a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}