process_fresh_mean.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e699727f",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import glob\n",
    "import netCDF4 as nc\n",
    "import datetime as dt\n",
    "import sys\n",
    "import gsw as sw\n",
    "import cartopy.crs as ccrs\n",
    "import cartopy.feature as cfeature\n",
    "import matplotlib\n",
    "matplotlib.use('agg')\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.path as mpath\n",
    "import cftime\n",
    "import coast\n",
    "import xarray as xr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c57df360",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "\n",
    "var = 'so' # thetao, so, uo, vo, siconc, siage, sivol, sithick, siu, siv, \n",
    "i_o = 'O' # SI or O for sea ice or ocean\n",
    "freq = 'mon' # mon or day\n",
    "time_s = 'highres-future' # 'highres-future' or 'hist-1950'\n",
    "\n",
    "def make_path(var, i_o, freq, time_s):\n",
    "    if 'future' in time_s:\n",
    "        ddir = 'MOHC'\n",
    "    else:\n",
    "        ddir = 'NERC'\n",
    "    root = '/badc/cmip6/data/CMIP6/HighResMIP/' + ddir + '/HadGEM3-GC31-HH/' + time_s + '/r1i1p1f1/'\n",
    "    return root + i_o + freq + '/' + var + '/gn/latest/' + var + '_' + i_o + freq + '_HadGEM3-GC31-HH_' + time_s + '_r1i1p1f1_gn_*.nc'\n",
    "\n",
    "fn_nemo_dat_s1 = make_path('so', i_o, freq, 'hist-1950')\n",
    "fn_nemo_dat_s2 = make_path('so', i_o, freq, time_s)\n",
    "\n",
    "\n",
    "domain_root = '/gws/nopw/j04/nemo_vol5/acc/eORCA12-N512/domain/'\n",
    "fn_nemo_dom1 = domain_root + 'eORCA12_coordinates.nc'\n",
    "fn_nemo_dom = domain_root + 'mesh_mask_eORCA12_v2.4.nc'\n",
    "fn_nemo_bathy = domain_root + 'eORCA12_bathymetry_v2.4.nc'\n",
    "fn_config_t_grid = './config/gc31_nemo_grid_t.json'\n",
    "\n",
    "out_file = './Processed/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c157bf00-11e1-4319-8404-c807919cacda",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define start and end date for decade mean\n",
    "if 0:\n",
    "    st_date = dt.datetime(1990, 1, 1)\n",
    "    en_date = dt.datetime(2000, 1, 1)\n",
    "else:\n",
    "    st_date = dt.datetime(2040, 1, 1)\n",
    "    en_date = dt.datetime(2050, 1, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d47c6799",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "now = dt.datetime.now()\n",
    "\n",
    "flist_s = sorted(glob.glob(fn_nemo_dat_s1))\n",
    "\n",
    "flist_s.extend(sorted(glob.glob(fn_nemo_dat_s2)))\n",
    "               \n",
    "\n",
    "v_map = {}\n",
    "v_map['e1t'] = 'e1t'\n",
    "v_map['e2t'] = 'e2t'\n",
    "v_map['e3t_0'] = 'e3t_0'\n",
    "v_map['tmask'] = 'tmask'\n",
    "v_map['lat'] = 'latitude'\n",
    "v_map['lon'] = 'longitude'\n",
    "v_map['depth'] = 'lev'\n",
    "v_map['time'] = 'time'\n",
    "v_map['temp'] = 'thetao'\n",
    "v_map['sal'] = 'so' \n",
    "\n",
    "\n",
    "with nc.Dataset(fn_nemo_dom, 'r') as nc_fid:\n",
    "    e1t = nc_fid.variables[v_map['e1t']][0, ...] # t, y, x\n",
    "    e2t = nc_fid.variables[v_map['e2t']][0, ...]\n",
    "    e3t = nc_fid.variables[v_map['e3t_0']][0, ...] # t, z, y, x\n",
    "    tmask = nc_fid.variables[v_map['tmask']][0, :, 1:-1, 1:-1]\n",
    "    \n",
    "with nc.Dataset(flist_s[0], 'r') as nc_fid:\n",
    "    lat = nc_fid.variables[v_map['lat']][:]\n",
    "    lon = nc_fid.variables[v_map['lon']][:]\n",
    "    depth = nc_fid.variables[v_map['depth']][:]\n",
    "    sal = nc_fid.variables[v_map['sal']][0, ...]\n",
    "\n",
    "sal = np.ma.masked_where((sal==1e20) | (tmask==1), sal)\n",
    "mask = sal.mask"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "debf5187-857b-47bd-90c4-22c1d3cd7cf8",
   "metadata": {},
   "source": [
    "Subset data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8c5164b",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "yi1 = 2800\n",
    "\n",
    "lat = lat[yi1:, :]\n",
    "lon = lon[yi1:, :]\n",
    "mask = mask[:, yi1:, :]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90c543ee-0efc-48cc-86a5-7139ab51579c",
   "metadata": {},
   "source": [
    "Time slice. Notes dates are 360 day years so use cftime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "481b58fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "#nemo_t_subset.construct_density(eos='EOS10')#, pot_dens=True)\n",
    "\n",
    "def calc_rho(sp, tp, depth, lon, lat):\n",
    "    pres = sw.p_from_z(depth * -1, lat)\n",
    "    sa = sw.SA_from_SP(sp, pres, lon, lat)\n",
    "    ct = sw.CT_from_pt(sa, tp)\n",
    "    rho = sw.rho(sa, ct, 0)\n",
    "    return rho, ct\n",
    "\n",
    "def calc_fresh(sal):\n",
    "    s_ref = 35\n",
    "    fresh = ((s_ref - sal) / s_ref)\n",
    "    return fresh\n",
    "\n",
    "# change this to decrease resolution but decrease run time\n",
    "sub = 1\n",
    "\n",
    "\n",
    "e1t = e1t[yi1+1:-1:sub, 1:-1:sub] # y, x\n",
    "e2t = e2t[yi1+1:-1:sub, 1:-1:sub]\n",
    "e3t = e3t[:, yi1+1:-1:sub, 1:-1:sub] # z, y, x\n",
    "    \n",
    "e1t = np.tile(e1t, (e3t.shape[0], 1, 1))\n",
    "e2t = np.tile(e2t, (e3t.shape[0], 1, 1))\n",
    "print(e3t.shape, e1t.shape)\n",
    "volume = e1t * e2t * e3t\n",
    "\n",
    "      \n",
    "lat = lat[::sub, ::sub]\n",
    "lon = lon[::sub, ::sub]\n",
    "mask = mask[:, ::sub, ::sub]\n",
    "    \n",
    "depth_g = np.tile(depth, (lon.shape[1], lon.shape[0], 1)).T\n",
    "\n",
    "date_from = np.zeros((len(flist_s)), dtype=object)\n",
    "date_to = np.zeros((len(flist_s)), dtype=object)\n",
    "for i in range(len(flist_s)):\n",
    "    part = flist_s[i].split('_')[-1].split('.')[0].split('-')\n",
    "    date_from[i] = dt.datetime.strptime(part[0], '%Y%m')\n",
    "    date_to[i] = dt.datetime.strptime(part[1], '%Y%m')\n",
    "\n",
    "\n",
    "date_use = np.nonzero((date_from >= st_date) & (date_to < en_date))[0]\n",
    "date_from = date_from[date_use]\n",
    "str_year = str(st_date.year) + '-' + str(en_date.year)\n",
    "\n",
    "# output a decacal monthly mean\n",
    "mn = np.array([t.month for t in date_from])\n",
    "yr = np.array([t.year for t in date_from])\n",
    "yr_uni = np.unique(yr)\n",
    "\n",
    "\n",
    "ref = 'days since 1950-01-01'\n",
    "date = np.zeros((len(flist_s)), dtype=object)\n",
    "fresh_time = np.ma.zeros((12, lon.shape[0], lon.shape[1]))\n",
    "for m in range(1, 13):\n",
    "    for y in yr_uni:\n",
    "        ind = np.nonzero((mn == m) & (yr == y))[0][0] + date_use[0]\n",
    "        with nc.Dataset(flist_s[ind], 'r') as nc_fid:\n",
    "            sal = nc_fid.variables[v_map['sal']][0, :, yi1::sub, ::sub]\n",
    "            time = nc_fid.variables[v_map['time']][:]\n",
    "            \n",
    "        sal = np.ma.masked_where((sal==1e20), sal)       \n",
    "        fresh_time[m-1, :, :] = fresh_time[m-1, :, :] + np.ma.sum(calc_fresh(sal) * e3t, axis=0) # vertically integrated\n",
    "        date[i] = cftime.num2date(time, ref, calendar='360_day')[0]\n",
    "        print(date[i])\n",
    "\n",
    "fresh_time = fresh_time / len(yr_uni)\n",
    "fresh_time = fresh_time.filled(-1e20)\n",
    "np.savez(out_file + 'fresh_mn_' + str_year + '.npz', fresh_map=fresh_time, lat=lat, lon=lon)\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "encoding": "# coding: utf-8",
   "executable": "/usr/bin/env python",
   "main_language": "python",
   "notebook_metadata_filter": "-all"
  },
  "kernelspec": {
   "display_name": "Python 3 + Jaspy",
   "language": "python",
   "name": "jaspy"
  },
  "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}