process_freshwater_final.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\n",
    "from functools import partial\n",
    "import dask"
   ]
  },
  {
   "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": "d47c6799",
   "metadata": {},
   "outputs": [],
   "source": [
    "# change this to decrease resolution but decrease run time\n",
    "sub = 1\n",
    "\n",
    "now = dt.datetime.now()\n",
    "\n",
    "flist_s = sorted(glob.glob(fn_nemo_dat_s1))\n",
    "flist_s.extend(sorted(glob.glob(fn_nemo_dat_s2)))\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']][:, :, 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']][:] # t, z, y, x\n",
    "\n",
    "\n",
    "lon_bnds, lat_bnds = (-180, 180), (60, 90)\n",
    "y1 = np.min(np.nonzero((lat >= lat_bnds[0]))[0])\n",
    "y2 = np.max(np.nonzero((lat <= lat_bnds[1]))[0])\n",
    "x1 = np.min(np.nonzero((lon >= lon_bnds[0]))[1])\n",
    "x2 = np.max(np.nonzero((lon <= lon_bnds[1]))[1])\n",
    "\n",
    "lat = lat[y1:y2:sub, x1:x2:sub]\n",
    "lon = lon[y1:y2:sub, x1:x2:sub]\n",
    "tmask = (sal[:, :, y1:y2:sub, x1:x2:sub] == 1e20) | (tmask[:, :, y1:y2:sub, x1:x2:sub] == 0)\n",
    "e1t = e1t[y1:y2:sub, x1:x2:sub]\n",
    "e2t = e2t[y1:y2:sub, x1:x2:sub]\n",
    "e3t = e3t[:, y1:y2:sub, x1:x2:sub]\n",
    "\n",
    "def dask_load_ts(fns, tmask, y1, y2, x1, x2, sub):\n",
    "\n",
    "    lazy_arrays = [dask.delayed(nc.Dataset)(fn, 'r') for fn in fns]\n",
    "    lazy_arrays = [dask.array.from_delayed(\n",
    "                        x.variables[v_map['sal']][:, :, y1:y2:sub, x1:x2:sub], shape=sal[:, :, y1:y2:sub, x1:x2:sub].shape, dtype=np.float64) \n",
    "                        for x in lazy_arrays]\n",
    "    ds_s = dask.array.concatenate(lazy_arrays[:], axis=0)\n",
    "    ds_s = ds_s.rechunk(chunks={0:1, 1:-1, 2:100, 3:100})#, balance=True)\n",
    "    tmask = dask.array.from_array(tmask)\n",
    "    tmask = tmask.rechunk(chunks={0:1, 1:-1, 2:100, 3:100})#, balance=True)\n",
    "    tmask = tmask.repeat(ds_s.shape[0], axis=0)\n",
    "    ds_s = dask.array.ma.masked_array(ds_s, mask=tmask)\n",
    "    return ds_s\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "481b58fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "\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, pres)\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",
    "def fc(fresh, volume):\n",
    "    return fresh * volume\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",
    "depth_g = np.tile(depth, (lon.shape[1], lon.shape[0], 1)).T\n",
    "mask_south = np.zeros((lat.shape[0], lat.shape[1], 3), dtype=bool)\n",
    "mask_south[:, :, 0] = lat < 70 # mask where less\n",
    "mask_south[:, :, 1] = lat < 75\n",
    "mask_south[:, :, 2] = lat < 80\n",
    "\n",
    "\n",
    "ref = 'days since 1950-01-01'\n",
    "date = np.zeros((len(flist_s)), dtype=object)\n",
    "for i in range(len(flist_s)):\n",
    "    with nc.Dataset(flist_s[i], 'r') as nc_fid:\n",
    "        time = nc_fid.variables[v_map['time']][:]\n",
    "    date[i] = cftime.num2date(time, ref, calendar='360_day')[0]\n",
    "    \n",
    "lat_g = dask.array.from_array(lat[np.newaxis, np.newaxis, :, :], chunks={0:1, 0:-1, 1:100, 2:100})\n",
    "lon_g = dask.array.from_array(lon[np.newaxis, np.newaxis, :, :], chunks={0:1, 0:-1, 1:100, 2:100})\n",
    "depth_g = dask.array.from_array(depth_g[np.newaxis, :, :, :], chunks={0:1, 0:-1, 1:100, 2:100})\n",
    "volume = dask.array.from_array(volume[np.newaxis, :, :, :], chunks={0:1, 0:-1, 1:100, 2:100})\n",
    "mask_south = dask.array.from_array(mask_south[np.newaxis, :, :, :], chunks={0:1, 0:100, 1:100, 2:-1})\n",
    "lat_g = lat_g.repeat(depth_g.shape[1], axis=1).repeat(12, axis=0)\n",
    "lon_g = lon_g.repeat(depth_g.shape[1], axis=1).repeat(12, axis=0)\n",
    "depth_g = depth_g.repeat(12, axis=0)\n",
    "volume = volume.repeat(12, axis=0)\n",
    "mask_south = mask_south.repeat(12, axis=0)\n",
    "\n",
    "fresh_time = np.ma.zeros((len(flist_s), 3))\n",
    "for i in range(len(flist_s) // 12):\n",
    "    i1 = int(i * 12)\n",
    "    i2 = int((i + 1) * 12)\n",
    "    \n",
    "    ds_s = dask_load_ts(flist_s[i1:i2], tmask, y1, y2, x1, x2, sub)\n",
    "    \n",
    "    fresh_part = (\n",
    "        dask.array.ma.masked_where(mask_south, \n",
    "        fc(\n",
    "        calc_fresh(ds_s), \n",
    "        volume)        # t, z, y, x\n",
    "        .sum(axis=1)   # t, y, x\n",
    "        [:, :, :, np.newaxis].repeat(3, axis=3)) # t, y, x, mask\n",
    "        .sum(axis=(1, 2)) # t, mask\n",
    "        )\n",
    "    print(date[i1])\n",
    "    print(fresh_part)\n",
    "    fresh_time[i1:i2, :] = fresh_part.compute()\n",
    "print(fresh_time.shape)\n",
    "    \n",
    "                          "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee167fa5-d9b1-46e5-831a-418c4666e2b7",
   "metadata": {},
   "source": [
    "Save data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f8a29af-8460-4317-ad53-5679a2433bfb",
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Runtime:', dt.datetime.now() - now)\n",
    "print(fresh_time)\n",
    "fresh_time = fresh_time.filled(-1e20)\n",
    "np.savez(out_file + 'freshwater_content.npz', fresh_time=fresh_time, date=date)"
   ]
  }
 ],
 "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
}