scripts/03P_WOA_zarr_data_extraction.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a9a0b1d6-bb0d-4030-bc0c-7ed2534c4b00",
   "metadata": {},
   "source": [
    "# Extracting WOA23 data using FishMIP regional model boundaries\n",
    "**Author:** Denisse Fierro Arcos  \n",
    "**Date:** 2024-09-04  \n",
    "  \n",
    "This script uses the cloud optimised WOA23 data (i.e., `zarr` files produced in the [02P_WOA_netcdf_to_zarr.ipynb](02P_WOA_netcdf_to_zarr.ipynb) script) to extract data for all FishMIP regional models.\n",
    "\n",
    "We use the FishMIP regional model shapefile, which is available via our THREDDS server. You can refer to this [notebook](https://github.com/Fish-MIP/FishMIP_regions/blob/main/scripts/02_Mapping_Regional_Models.md) for instructions on how to download this shapefile.\n",
    "\n",
    "Additionally, you will need a mask containing all FishMIP regional models. Instructions on how to create this mask are available [here](https://github.com/Fish-MIP/FishMIP_regions/blob/main/scripts/03a_Regional_Models_2DMasks.md)\n",
    "\n",
    "We recommend that you store both the shapefile and mask in the same folder."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65ced475-7b83-44dd-a06f-fb235f8e0ba7",
   "metadata": {},
   "source": [
    "## Setting working directory\n",
    "Remember to change the working directory below to the location of the scripts in your own local machine. Update the folder path below before continuing with the next chunk."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "3944e9dc-62ad-43b3-8f93-f5a5fbddf403",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.chdir('/g/data/vf71/la6889/processing_WOA_data/scripts')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6bd1c29-91ea-4637-958b-0d3e15e474b7",
   "metadata": {},
   "source": [
    "## Loading libraries\n",
    "We will load published Python libraries as well as our custom-made `useful_functions` library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6d343d26-90b1-4f14-9dbb-1cc4cda6459c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.distributed import Client\n",
    "from glob import glob\n",
    "import xarray as xr\n",
    "import geopandas as gpd\n",
    "from glob import glob\n",
    "import useful_functions as uf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a6d647c-567c-4d45-b736-f2263a1c8cc7",
   "metadata": {},
   "source": [
    "## Starting a cluster\n",
    "This will allow us to automatically parallelising tasks on large datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d56e7f29-4751-45cf-abbc-d4fc512ac2f8",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/g/data/hh5/public/apps/miniconda3/envs/analysis3-24.07/lib/python3.10/site-packages/distributed/node.py:187: UserWarning: Port 8787 is already in use.\n",
      "Perhaps you already have a cluster running?\n",
      "Hosting the HTTP server on port 38019 instead\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "client = Client(threads_per_worker = 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f9bbac4-d89b-4a27-ab52-bb125c9f4c80",
   "metadata": {},
   "source": [
    "## Defining basic variables\n",
    "Before continuing with the script, remember to update the `regions_folder` variable below with the location of the folder where you are storing the FishMIP regional model shapefile and gridded mask."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1ca5544b-4162-499e-8b5e-55a0dd32aa9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "regions_folder = '/g/data/vf71/shared_resources/FishMIP_regional_models/'\n",
    "mask_folder = '/g/data/vf71/shared_resources/FishMIPMasks/merged_regional_fishmip'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a535c493-c54f-43ad-b5d4-d2b33db8a3c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Loading FishMIP regional models shapefile\n",
    "rmes = gpd.read_file(os.path.join(regions_folder, 'FishMIP_regional_models.shp'))\n",
    "#Loading FishMIP regional models gridded mask\n",
    "mask_ras = xr.open_dataset(os.path.join(mask_folder, \n",
    "                                        'gfdl-mom6-cobalt2_areacello_15arcmin_fishMIP_regional_merged.nc')).region\n",
    "#Renaming coordinate dimensions\n",
    "mask_ras = mask_ras.rename({'latitude': 'lat', 'longitude': 'lon'})\n",
    "#Rechunking data to make it more manageable\n",
    "mask_ras = mask_ras.chunk({'lat': 144, 'lon': 288})\n",
    "\n",
    "# Getting a list of all WOA zarr files available \n",
    "WOA_zarr = glob('/g/data/vf71/WOA_data/global/woa*.zarr')\n",
    "\n",
    "#Define (or create) folders where outputs will be stored\n",
    "base_out_clim = '/g/data/vf71/WOA_data/regional/climatology'\n",
    "os.makedirs(base_out_clim, exist_ok = True)\n",
    "base_out_month = '/g/data/vf71/WOA_data/regional/monthly'\n",
    "os.makedirs(base_out_month, exist_ok = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebfac74d-9cf3-4c36-9dd8-abb8b8f893f8",
   "metadata": {},
   "source": [
    "## Extracting WOA data for each region"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "b334e242-4ee6-40da-a677-3751232d6ade",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Applying functions to WOA files\n",
    "for f in WOA_zarr:\n",
    "    #Open data array as ARD\n",
    "    da = uf.mask_boolean_ard_data(f, mask_ras)   \n",
    "    base_name = os.path.basename(f).replace('zarr', 'parquet')\n",
    "    \n",
    "    #Create full file path\n",
    "    if 'month' in f:\n",
    "         #Adding output folder to create full file path\n",
    "        full_file_out = os.path.join(base_out_month, base_name)\n",
    "    else:\n",
    "        full_file_out = os.path.join(base_out_clim, base_name)\n",
    "\n",
    "    #Extract data for each region included in the regional mask\n",
    "    for i in rmes.region:\n",
    "        #Get polygon for each region\n",
    "        mask = rmes[rmes.region == i]\n",
    "        #Get name of region and clean it for use in output file\n",
    "        reg_name = mask['region'].values[0].lower().replace(\" \", \"-\").replace(\"'\", \"\")\n",
    "        #File name out - Replacing \"global\" for region name\n",
    "        file_out = full_file_out.replace('woa23_', f'woa23_{reg_name}_')\n",
    "        #Extract data and save masked data - but only if file does not already exist\n",
    "        if os.path.isdir(file_out) | os.path.isfile(file_out):\n",
    "            continue\n",
    "        uf.mask_ard_data(da, mask, file_out)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "886e4051-e014-46ad-b768-bd3e034abeb3",
   "metadata": {},
   "source": [
    "## Extracting WOA data from regridded files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "42098d29-01f1-4d24-b0af-a159f2cd1a15",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Getting a list of all WOA zarr files available \n",
    "WOA_zarr = glob('/g/data/vf71/WOA_data/global/regridded_woa*.zarr')\n",
    "\n",
    "#Define (or create) folders where outputs will be stored\n",
    "base_out_clim = '/g/data/vf71/WOA_data/regional/climatology/comp_clim/'\n",
    "os.makedirs(base_out_clim, exist_ok = True)\n",
    "base_out_month = '/g/data/vf71/WOA_data/regional/monthly/regridded/'\n",
    "os.makedirs(base_out_month, exist_ok = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1cd12d1-3163-4237-ba1f-1827c12ace9c",
   "metadata": {},
   "source": [
    "## Extracting monthly WOA data for each region"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "94081b38-b382-48f9-909f-6f8443897ffe",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Applying functions to WOA files\n",
    "for f in WOA_zarr:\n",
    "    #Open data array as ARD\n",
    "    da = uf.mask_boolean_ard_data(f, mask_ras)   \n",
    "    base_name = os.path.basename(f).replace('zarr', 'parquet')\n",
    "    \n",
    "    #Create full file path\n",
    "    if 'month' in f:\n",
    "         #Adding output folder to create full file path\n",
    "        full_file_out = os.path.join(base_out_month, \n",
    "                                     base_name.replace('woa_month', \n",
    "                                                       'woa_global_month'))\n",
    "    else:\n",
    "        full_file_out = os.path.join(base_out_clim, \n",
    "                                     base_name.replace('woa_clim', \n",
    "                                                       'woa23_global_clim'))\n",
    "        \n",
    "    #Extract data for each region included in the regional mask\n",
    "    for i in rmes.region:\n",
    "        #Get polygon for each region\n",
    "        mask = rmes[rmes.region == i]\n",
    "        #Get name of region and clean it for use in output file\n",
    "        reg_name = mask['region'].values[0].lower().replace(\" \", \"-\").replace(\"'\", \"\")\n",
    "        file_out = full_file_out.replace('_global_', f'_{reg_name}_')\n",
    "        #Extract data and save masked data - but only if file does not already exist\n",
    "        if os.path.isdir(file_out) | os.path.isfile(file_out):\n",
    "            continue\n",
    "        uf.mask_ard_data(da, mask, file_out)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:analysis3-24.07]",
   "language": "python",
   "name": "conda-env-analysis3-24.07-py"
  },
  "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.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}