ptracer_calc.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculation for PTRACER\n",
    "import xarray as xr\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import glob"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(69, 175)"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hs = 1\n",
    "# /Volumes/Extreme SSD/MITgcm outputs/ & ../Model_outputs\n",
    "tracer = xr.open_dataset('../Model_outputs/TracerAge/ptracers_' + str(format(hs,'03d')) + '.nc')\n",
    "#tracer2 = xr.open_dataset('../Model_outputs/Tracer01/ptracers_' + str(format(hs,'03d')) + '.nc')\n",
    "state = xr.open_dataset('../Model_outputs/TracerAge/state_' + str(format(hs,'03d')) + '.nc')\n",
    "Grid = xr.open_dataset('../Model_outputs/TracerAge/grid_' + str(format(hs,'03d')) + '.nc')\n",
    "\n",
    "HFacC = Grid.HFacC.data\n",
    "drF = Grid.drF.data\n",
    "rA = Grid.rA.data\n",
    "cv = np.broadcast_to(rA[np.newaxis, :, :], HFacC.shape) * np.broadcast_to(drF[:, np.newaxis, np.newaxis], HFacC.shape) * HFacC\n",
    "\n",
    "time = state.T.data\n",
    "x = state.X.data\n",
    "z = state.Z.data\n",
    "\n",
    "tage1 = tracer.tr1_age.data\n",
    "#tr1 = tracer.tr1.data\n",
    "#tr1_rlx1 = tracer.tr1_rlx.data\n",
    "#tr1_rlx2 = tracer2.tr1_rlx.data\n",
    "rA.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Volume weighted concentration\n",
    "\n",
    "Tr1 = np.empty(len(time))\n",
    "Tr1_rlx1 = np.empty(len(time))\n",
    "Tr1_rlx2 = np.empty(len(time))\n",
    "\n",
    "for i in range(len(time)):\n",
    "    Tr1[i] = np.sum(tr1[i,40:,:,120]*cv[40:,:,120]) / cv[40:,:,120].sum()\n",
    "    Tr1_rlx1[i] = np.sum(tr1_rlx1[i,40:,:,120]*cv[40:,:,120]) / cv[40:,:,120].sum()\n",
    "    Tr1_rlx2[i] = np.sum(tr1_rlx2[i,40:,:,120]*cv[40:,:,120]) / cv[40:,:,120].sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig1,ax1 = plt.subplots(figsize=(10,7))\n",
    "\n",
    "ax1.plot(time,Tr1,'-',label='Initial(no relax)')\n",
    "ax1.plot(time,Tr1_rlx1,'-',label='Relax-100s')\n",
    "ax1.plot(time,Tr1_rlx2,'-',label='Relax-1800s')\n",
    "\n",
    "ax1.legend(loc='best', fontsize=15)\n",
    "ax1.set_ylabel(r'Concentration',size = 20)\n",
    "ax1.set_xlabel(r'Time (s)', size = 20)\n",
    "ax1.grid(linestyle='--', linewidth='0.5')\n",
    "ax1.set(xlim=(0, time[-1]), ylim=(0.1, 1))\n",
    "#plt.title(r'Relaxing Timescale = 1800s',size=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Estimate the age of water parcel (Rayson et al., 2016)\n",
    "\n",
    "# Average in y direction\n",
    "tr1_zx = tr1.mean(2)\n",
    "tr1rlx_zx = tr1_rlx1.mean(2)\n",
    "dt = time[1] - time[0]\n",
    "\n",
    "NT, NZ, NX = tr1_zx.shape\n",
    "alpha = np.zeros((NT, NZ, NX)) # age concentration\n",
    "C = tr1_zx # total concentration\n",
    "\n",
    "for j in range(1,len(time)-1):\n",
    "    alpha[1,:,:] = tr1rlx_zx[1,:,:] * dt\n",
    "    alpha[j+1,:,:] = alpha[j] + tr1rlx_zx[j,:,:] * dt\n",
    "\n",
    "#for j in range(len(time)):\n",
    "#    alpha[j,:,:] = tr1rlx_zx[j,:,:] * time[j]    \n",
    "    \n",
    "    \n",
    "# Mean age of the water parcel\n",
    "a = alpha / C\n",
    "#a[np.isnan(a)] == 0\n",
    "A = np.nan_to_num(a)\n",
    "#np.nanmax(a), np.nanmin(a), np.nanmean(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "A[50,:,118]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Contour plot of the age\n",
    "cm = np.linspace(0,1e4,20) # Contour range\n",
    "\n",
    "cmap = 'RdBu_r'\n",
    "\n",
    "\n",
    "t1 = -1\n",
    "xl1, xr1 = 100, 123\n",
    "\n",
    "fig2 = plt.figure(figsize=(20,5))\n",
    "fig2.tight_layout()  \n",
    "       \n",
    "ax1 = fig2.add_subplot(111)\n",
    "im1 = ax1.contourf(x[xl1:xr1], z, A[t1,:,xl1:xr1], cm, cmap=cmap,extend='max')\n",
    "ax1.axvline(x=x[120], color='g',lw=2)\n",
    "\n",
    "\n",
    "cbar1 = fig2.colorbar(im1)\n",
    "cbar1.set_ticks([0, .5e4, 1e4])\n",
    "cbar1.ax.set_title(r'Age (s)', size=12)\n",
    "\n",
    "# Set Colorbar\n",
    "#cbaxes = fig2.add_axes([0.75, .9, 0.15, 0.02])\n",
    "#cbar = fig2.colorbar(im1,cax=cbaxes, orientation='horizontal')\n",
    "#cbar.set_ticks([-5, 0, 5])\n",
    "#cbar.ax.set_title(r'Diffusive Heat $(\\times{10^{-5}}^\\circ$C/s)', size=12)\n",
    "\n",
    "ax1.set_ylabel('Depth (m)', size=15)\n",
    "#ax1.set_title(r'(b) Initial', loc='left', size=15, weight='bold')\n",
    "ax1.set_xlabel('Distance from glacier (m)', size=15)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Contour plot of the tracer concentration\n",
    "cm1 = np.linspace(0,1,20) # Contour range\n",
    "cm2 = np.linspace(0,.2,20) # Contour range\n",
    "cmap = 'RdBu_r'\n",
    "\n",
    "t2 = -1\n",
    "xl2, xr2 = 100, 123\n",
    "\n",
    "fig3 = plt.figure(figsize=(18,9))\n",
    "fig3.tight_layout()  \n",
    "       \n",
    "ax1 = fig3.add_subplot(211)\n",
    "im1 = ax1.contourf(x[xl2:xr2], z, tr1rlx_zx[t2,:,xl2:xr2], cm1, cmap=cmap)\n",
    "ax1.axvline(x=x[120], color='g',lw=2)\n",
    "\n",
    "ax1.set_ylabel('Depth (m)', size=15)\n",
    "ax1.xaxis.set_visible(False)\n",
    "#ax1.set_xlim([x[Xl], x[Xr-1]])\n",
    "ax1.set_facecolor('k')\n",
    "ax1.set_title(r'(a) Source', loc='left', size=15, weight='bold')\n",
    "\n",
    "# Set Colorbar\n",
    "\n",
    "#plt.clabel(im1, fmt='%2.1f', fontsize=14)\n",
    "#plt.colorbar(im1)\n",
    "\n",
    "#cbaxes1 = fig3.add_axes([0.75, .91, 0.15, 0.02])\n",
    "cbar1 = fig3.colorbar(im1)\n",
    "cbar1.set_ticks([0, .5, 1])\n",
    "cbar1.ax.set_title(r'Concentration', size=12)\n",
    "\n",
    "       \n",
    "ax2 = fig3.add_subplot(212)\n",
    "im2 = ax2.contourf(x[xl2:xr2], z, tr1_zx[t2,:,xl2:xr2], cm2, cmap=cmap)\n",
    "ax2.axvline(x=x[120], color='g',lw=2)\n",
    "\n",
    "cbar2 = fig3.colorbar(im2)\n",
    "cbar2.set_ticks([0, .1, .2])\n",
    "cbar2.ax.set_title(r'Concentration', size=12)\n",
    "#cbaxes2 = fig3.add_axes([0.75, .49, 0.15, 0.02])\n",
    "#cbar2 = fig3.colorbar(im2,cax=cbaxes2, orientation='horizontal')\n",
    "#cbar2.set_ticks([0,.04,.08])\n",
    "#cbar2.ax.set_title(r'Concentration', size=12)\n",
    "\n",
    "ax2.set_ylabel('Depth (m)', size=15)\n",
    "ax2.set_title(r'(b) Initial', loc='left', size=15, weight='bold')\n",
    "ax2.set_xlabel('Distance from glacier (m)', size=15)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/weiyangbao/miniconda3/envs/py37/lib/python3.7/site-packages/ipykernel_launcher.py:16: RuntimeWarning: divide by zero encountered in log10\n",
      "  app.launch_new_instance()\n"
     ]
    }
   ],
   "source": [
    "cm4 = np.linspace(0,2,20) # Contour range\n",
    "cmap4 = 'RdBu_r'\n",
    "\n",
    "tage = tage1.mean(2)\n",
    "xl2, xr2 = 105, 125\n",
    "\n",
    "fig_path = \"/Users/weiyangbao/Documents/Results_Figs/PTRACERS/Age/\"\n",
    "if not os.path.exists(fig_path):\n",
    "    os.makedirs(fig_path)\n",
    "\n",
    "for t in range(len(time)):\n",
    "    fig4 = plt.figure(figsize=(24,4))\n",
    "    fig4.tight_layout()  \n",
    "       \n",
    "    ax4 = fig4.add_subplot(111) \n",
    "    im4 = ax4.contourf(x[xl2:xr2], z, np.log10(tage[t,:,xl2:xr2]), cm4, cmap=cmap4)\n",
    "    ax4.axvline(x=x[120], color='g',lw=2)\n",
    "\n",
    "    ax4.set_ylabel('Depth (m)', size=15)\n",
    "    ax4.set_xlabel('Distance from glacier (m)', size=15)\n",
    "\n",
    "    #ax4.xaxis.set_visible(False)\n",
    "    #ax1.set_xlim([x[Xl], x[Xr-1]])\n",
    "    ax4.set_facecolor('k')\n",
    "    ax4.set_title(r'Time: '+str(t*100)+' s', loc='center', size=15, weight='bold')\n",
    "    \n",
    "    # Set Colorbar\n",
    "\n",
    "    #plt.clabel(im1, fmt='%2.1f', fontsize=14)\n",
    "    #plt.colorbar(im1)\n",
    "\n",
    "    #cbaxes1 = fig3.add_axes([0.75, .91, 0.15, 0.02])\n",
    "    cbar4 = fig4.colorbar(im4)\n",
    "    cbar4.set_ticks([0, .5, 1, 1.5, 2])\n",
    "    cbar4.ax.set_title(r'Concentration ($\\log_{10}$)', size=12)\n",
    " \n",
    "    \n",
    "    plt.close()\n",
    "    fig4.savefig(fig_path + '/time_' + format(t, '03d') + '.png', dpi=150, bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate GIF from saved output figures\n",
    "import imageio\n",
    "import glob\n",
    "\n",
    "images = []\n",
    "imagenames = sorted(glob.glob(fig_path + '/time_*.png'))\n",
    "\n",
    "for filename in imagenames:\n",
    "    images.append(imageio.imread(filename))\n",
    "imageio.mimsave(fig_path + 'TracerAge.gif', images, duration=.2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig_path = \"/Users/weiyangbao/Documents/Results_Figs/PTRACERS/\"\n",
    "if not os.path.exists(fig_path):\n",
    "    os.makedirs(fig_path)\n",
    "\n",
    "#fig1.savefig(fig_path + '/Rlx_scale.png', dpi=200, bbox_inches='tight')\n",
    "fig3.savefig(fig_path + '/Concentration.png', dpi=200, bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([  0,   2,   4,   6,   8,  10,  12,  14,  16,  18,  20,  22,  24,\n",
       "        26,  28,  30,  32,  34,  36,  38,  40,  42,  44,  46,  48,  50,\n",
       "        52,  54,  56,  58,  60,  62,  64,  66,  68,  70,  72,  74,  76,\n",
       "        78,  80,  82,  84,  86,  88,  90,  92,  94,  96,  98, 100])"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.arange(0,101,2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}