Diagnostic plots for tritium surface fluxes

Created by Ivan Lima on Wed Jun 28 2017 08:39:08 -0400

import xarray as xr
import pandas as pd
import numpy as np
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import matplotlib.pyplot as plt
import os, cmocean
from matplotlib import ticker, colors
from mpl_utils import spd, center_cmap
from cesm_utils import *
%matplotlib notebook

Read data

case = 'test_GIAF_tr3he.gx1.015'
datadir = '/home/ivan/Data/Postproc/Tritium-3H'
ds = xr.open_dataset(os.path.join(datadir,'%s_surf.nc'%case),decode_times=False)
ds = fixtime(ds,1746)
ds = ds.where(ds.REGION_MASK>0) # mask marginal seas for plotting
lon = ds.TLONG.values
loncp = np.where(np.greater_equal(lon,lon[:,0].min()),lon-360,lon) # monotonic lon for contour plots
# cyclic lon & lat for pcolormesh
lonc = np.hstack((ds.TLONG.values,ds.TLONG.values[:,0,np.newaxis]))
latc = np.hstack((ds.TLAT.values,ds.TLAT.values[:,0,np.newaxis]))
figsize=(9.75,7.5)

Maps of annual means of last year of simulation (year 50)

Surface tritium

data = ds.TRITIUM[-12:,0,...].mean(dim='time')
fig, ax = plt.subplots(subplot_kw={'projection':ccrs.Robinson(central_longitude=210)},figsize=(9,6))
ax.set_global()
ax.add_feature(cfeature.LAND, facecolor='#b0b0b0')
ax.set_title('mean surface tritium (%d)'%(ds.time[-1]['time.year']),fontsize=12)
pm = ax.pcolormesh(lonc, latc, data,transform=ccrs.PlateCarree(),vmin=data.min(),vmax=data.max())#,cmap=plt.cm.Reds)
cb = fig.colorbar(pm,ax=ax,orientation='horizontal',pad=0.05)
cb.set_label(r'pmol m$^{-3}$')
fig.savefig('plots/surf_tritium.ps',bbox_inches='tight')

Total surface tritium flux

data = ds.STF_TRITIUM[-12:,...].mean(dim='time') * spd # pmol/m^/s -> pmol/m^2/day
fig, ax = plt.subplots(subplot_kw={'projection':ccrs.Robinson(central_longitude=210)},figsize=(9,6))
ax.set_global()
ax.add_feature(cfeature.LAND, facecolor='#b0b0b0')
ax.set_title('mean total tritium flux (%d)'%(ds.time[-1]['time.year']),fontsize=12)
pm = ax.pcolormesh(lonc, latc, data,transform=ccrs.PlateCarree(),vmin=data.min(),vmax=data.max(),cmap=plt.cm.RdBu_r)
center_cmap(pm)
cb = fig.colorbar(pm,ax=ax,orientation='horizontal',pad=0.05)
cb.set_label(r'pmol m$^{-2}$ d$^{-1}$')

Surface tritium flux due to evaporation

# data = ds.STF_TRITIUM_EVAP[-12:,...].mean(dim='time') * spd # pmol/m^/s -> pmol/m^2/day
# fig, ax = plt.subplots(subplot_kw={'projection':ccrs.Robinson(central_longitude=210)},figsize=(9,6))
# ax.set_global()
# ax.add_feature(cfeature.LAND, facecolor='#b0b0b0')
# ax.set_title('mean tritium flux due to evaporation (%d)'%(ds.time[-1]['time.year']),fontsize=12)
# pm = ax.pcolormesh(lonc, latc, data,transform=ccrs.PlateCarree(),vmin=data.min(),vmax=data.max(),cmap=plt.cm.RdBu_r)
# center_cmap(pm)
# cb = fig.colorbar(pm,ax=ax,orientation='horizontal',pad=0.05)
# cb.set_label(r'pmol m$^{-2}$ d$^{-1}$')
# fig.savefig('plots/stf_tritium_evap.ps',bbox_inches='tight')

Surface tritium flux due to precipitation

data = ds.STF_TRITIUM_PREC[-12:,...].mean(dim='time') * spd # pmol/m^/s -> pmol/m^2/day
fig, ax = plt.subplots(subplot_kw={'projection':ccrs.Robinson(central_longitude=210)},figsize=(9,6))
ax.set_global()
ax.add_feature(cfeature.LAND, facecolor='#b0b0b0')
ax.set_title('mean tritium flux due to precipitation (%d)'%(ds.time[-1]['time.year']),fontsize=12)
pm = ax.pcolormesh(lonc, latc, data,transform=ccrs.PlateCarree(),vmin=data.min(),vmax=data.max(),
                   cmap=plt.cm.RdBu_r)
center_cmap(pm)
cb = fig.colorbar(pm,ax=ax,orientation='horizontal',pad=0.05)
cb.set_label(r'pmol m$^{-2}$ d$^{-1}$')
fig.savefig('plots/stf_tritium_prec.ps',bbox_inches='tight')

Monthly maps of last year of simulation (year 50)

Surface tritium

fig, axs = plt.subplots(nrows=4,ncols=3,subplot_kw={'projection':ccrs.Robinson(central_longitude=210)},figsize=figsize)
fig.subplots_adjust(left=0.1,top=0.95,hspace=0.1,wspace=0.025)
axs = axs.ravel()
data, dates, = da2ma(ds['TRITIUM'][-12:,0,...]), ds.time[-12:] # surface field
for n in range(data.shape[0]):
    axs[n].set_global()
    axs[n].add_feature(cfeature.LAND, facecolor='#b0b0b0')
    axs[n].set_title('%d-%02d'%(dates[n]['time.year'],dates[n]['time.month']),fontsize=10)
    pm = axs[n].pcolormesh(lonc, latc, data[n,...],transform=ccrs.PlateCarree(),
                           vmin=data.min(),vmax=data.max())#,cmap=plt.cm.Reds)

l, b, w, h = axs[n].get_position().bounds
cax = fig.add_axes([0.125,0.07,0.775,0.025])
cb = fig.colorbar(pm,cax=cax,orientation='horizontal')
cb.set_label('surface tritium (pmol m$^{-3}$)')

Total surface tritium flux

fig, axs = plt.subplots(nrows=4,ncols=3,subplot_kw={'projection':ccrs.Robinson(central_longitude=210)},figsize=figsize)
fig.subplots_adjust(left=0.1,top=0.95,hspace=0.1,wspace=0.025)
axs = axs.ravel()
data, dates, = da2ma(ds['STF_TRITIUM'][-12:,...]) * spd, ds.time[-12:]
for n in range(data.shape[0]):
    axs[n].set_global()
    axs[n].add_feature(cfeature.LAND, facecolor='#b0b0b0')
    axs[n].set_title('%d-%02d'%(dates[n]['time.year'],dates[n]['time.month']),fontsize=10)
    pm = axs[n].pcolormesh(lonc, latc, data[n,...],transform=ccrs.PlateCarree(),
                           vmin=data.min(),vmax=data.max(),cmap=plt.cm.RdBu_r)
    center_cmap(pm)

l, b, w, h = axs[n].get_position().bounds
cax = fig.add_axes([0.125,0.07,0.775,0.025])
cb = fig.colorbar(pm,cax=cax,orientation='horizontal')
cb.set_label('total tritium flux (pmol m$^{-2}$ d$^{-1}$)')

Surface tritium flux due to evaporation

# fig, axs = plt.subplots(nrows=4,ncols=3,subplot_kw={'projection':ccrs.Robinson(central_longitude=210)},figsize=figsize)
# fig.subplots_adjust(left=0.1,top=0.95,hspace=0.1,wspace=0.025)
# axs = axs.ravel()
# data, dates, = da2ma(ds['STF_TRITIUM_EVAP'][-12:,...]) * spd, ds.time[-12:]
# for n in range(data.shape[0]):
#     axs[n].set_global()
#     axs[n].add_feature(cfeature.LAND, facecolor='#b0b0b0')
#     axs[n].set_title('%d-%02d'%(dates[n]['time.year'],dates[n]['time.month']),fontsize=10)
#     pm = axs[n].pcolormesh(lonc, latc, data[n,...],transform=ccrs.PlateCarree(),
#                            vmin=data.min(),vmax=data.max(),cmap=plt.cm.RdBu_r)
#     center_cmap(pm)

# l, b, w, h = axs[n].get_position().bounds
# cax = fig.add_axes([0.125,0.07,0.775,0.025])
# cb = fig.colorbar(pm,cax=cax,orientation='horizontal')
# cb.set_label('tritium flux due to evaporation (pmol m$^{-2}$ d$^{-1}$)')

Surface tritium flux due to precipitation

fig, axs = plt.subplots(nrows=4,ncols=3,subplot_kw={'projection':ccrs.Robinson(central_longitude=210)},figsize=figsize)
fig.subplots_adjust(left=0.1,top=0.95,hspace=0.1,wspace=0.025)
axs = axs.ravel()
data, dates, = da2ma(ds['STF_TRITIUM_PREC'][-12:,...]) * spd, ds.time[-12:]
for n in range(data.shape[0]):
    axs[n].set_global()
    axs[n].add_feature(cfeature.LAND, facecolor='#b0b0b0')
    axs[n].set_title('%d-%02d'%(dates[n]['time.year'],dates[n]['time.month']),fontsize=10)
    pm = axs[n].pcolormesh(lonc, latc, data[n,...],transform=ccrs.PlateCarree(),
                           vmin=data.min(),vmax=data.max(),cmap=plt.cm.RdBu_r)
    center_cmap(pm)

l, b, w, h = axs[n].get_position().bounds
cax = fig.add_axes([0.125,0.07,0.775,0.025])
cb = fig.colorbar(pm,cax=cax,orientation='horizontal')
cb.set_label('tritium flux due to precipitation (pmol m$^{-2}$ d$^{-1}$)')

$h = \frac{q_{atm}}{q_{sat}}$

fig, axs = plt.subplots(nrows=4,ncols=3,subplot_kw={'projection':ccrs.Robinson(central_longitude=210)},figsize=figsize)
fig.subplots_adjust(left=0.1,top=0.95,hspace=0.1,wspace=0.025)
axs = axs.ravel()
data, dates, = da2ma(ds['HQ'][-12:,...]), ds.time[-12:]
for n in range(data.shape[0]):
    axs[n].set_global()
    axs[n].add_feature(cfeature.LAND, facecolor='#b0b0b0')
    axs[n].set_title('%d-%02d'%(dates[n]['time.year'],dates[n]['time.month']),fontsize=10)
    pm = axs[n].pcolormesh(lonc, latc, data[n,...],transform=ccrs.PlateCarree(),
                           vmin=data.min(),vmax=0.85)

l, b, w, h = axs[n].get_position().bounds
cax = fig.add_axes([0.125,0.07,0.775,0.025])
cb = fig.colorbar(pm,cax=cax,orientation='horizontal')
cb.set_label(r'$h$ ($\frac{q_{atm}}{q_{sat}}$)')

Precipitation

fig, axs = plt.subplots(nrows=4,ncols=3,subplot_kw={'projection':ccrs.Robinson(central_longitude=210)},figsize=figsize)
fig.subplots_adjust(left=0.1,top=0.95,hspace=0.1,wspace=0.025)
axs = axs.ravel()
data, dates, = da2ma(ds['PREC_F'][-12:,...]) * spd, ds.time[-12:]
for n in range(data.shape[0]):
    axs[n].set_global()
    axs[n].add_feature(cfeature.LAND, facecolor='#b0b0b0')
    axs[n].set_title('%d-%02d'%(dates[n]['time.year'],dates[n]['time.month']),fontsize=10)
    pm = axs[n].pcolormesh(lonc, latc, data[n,...],transform=ccrs.PlateCarree(),
                           vmin=data.min(),vmax=data.max(),cmap=plt.cm.RdBu_r)
    center_cmap(pm)

l, b, w, h = axs[n].get_position().bounds
cax = fig.add_axes([0.125,0.07,0.775,0.025])
cb = fig.colorbar(pm,cax=cax,orientation='horizontal')
cb.set_label('precipitation (kg m$^{-2}$ d$^{-1}$)')

Evaporation

fig, axs = plt.subplots(nrows=4,ncols=3,subplot_kw={'projection':ccrs.Robinson(central_longitude=210)},figsize=figsize)
fig.subplots_adjust(left=0.1,top=0.95,hspace=0.1,wspace=0.025)
axs = axs.ravel()
data, dates, = da2ma(ds['EVAP_F'][-12:,...]) * spd, ds.time[-12:]
for n in range(data.shape[0]):
    axs[n].set_global()
    axs[n].add_feature(cfeature.LAND, facecolor='#b0b0b0')
    axs[n].set_title('%d-%02d'%(dates[n]['time.year'],dates[n]['time.month']),fontsize=10)
    pm = axs[n].pcolormesh(lonc, latc, data[n,...],transform=ccrs.PlateCarree(),
                           vmin=data.min(),vmax=data.max(),cmap=plt.cm.RdBu_r)
    center_cmap(pm)

l, b, w, h = axs[n].get_position().bounds
cax = fig.add_axes([0.125,0.07,0.775,0.025])
cb = fig.colorbar(pm,cax=cax,orientation='horizontal')
cb.set_label('evaporation (kg m$^{-2}$ d$^{-1}$)')

Time series

store = pd.HDFStore(os.path.join(datadir,'%s_tseries.h5'%case))
df_tseries = store['df']
store.close()

Tritium inventory

fig, ax = plt.subplots()
_ = df_tseries.TRITIUM.plot(ax=ax,title='tritium inventory')
_ = ax.set_ylabel(r'$10^6$ mols')
fig.savefig('plots/tritium_inv_tseries.ps',bbox_inches='tight')

Tritium surface fluxes

#df_stf = df_tseries[['STF_TRITIUM','STF_TRITIUM_EVAP','STF_TRITIUM_PREC']]
df_stf = df_tseries[['STF_TRITIUM','STF_TRITIUM_PREC']]
fig, ax = plt.subplots()
_ = df_stf.plot(ax=ax,title='tritium surface fluxes')
_ = ax.set_ylabel(r'mol d$^{-1}$')
fig.savefig('plots/tritium_stf_tseries.ps',bbox_inches='tight')

# fig, ax = plt.subplots()
# _ = df_stf.STF_TRITIUM_PREC.plot(title='tritium flux due to precipitation')
# _ = ax.set_ylabel(r'mol d$^{-1}$')
# fig.savefig('plots/tritium_stf_prec_tseries.ps',bbox_inches='tight')

Tritium decay

fig, ax = plt.subplots()
_ = df_tseries.TRITIUM_DECAY.plot(ax=ax,title='tritium decay')
_ = ax.set_ylabel(r'mol d$^{-1}$')
fig.savefig('plots/tritium_decay_tseries.ps',bbox_inches='tight')

Helium 3 inventory

fig, ax = plt.subplots()
_ = df_tseries[['HELIUM3','INERT_HELIUM3']].plot(ax=ax,title='helium 3 inventory')
_ = ax.set_ylabel(r'$10^6$ mols')

Vertical Sections

ds_sec_j12  = xr.open_dataset(os.path.join(datadir,'%s_sec_j12.nc'%case),decode_times=False)
ds_sec_j225 = xr.open_dataset(os.path.join(datadir,'%s_sec_j225.nc'%case),decode_times=False)
ds_sec_j12  = fixtime(ds_sec_j12)
ds_sec_j225 = fixtime(ds_sec_j225)

def secfig_mon():
    """monthly sections"""
    fig, axs = plt.subplots(nrows=4,ncols=3,sharex=True,sharey=True,figsize=(9.5,9.5))
    fig.subplots_adjust(top=0.95,wspace=0.05,hspace=0.15)
    axs = axs.ravel()
    for n in range(data.shape[0]):
        axs[n].set_title('%d-%02d'%(dates[n]['time.year'],dates[n]['time.month']),fontsize=10)
        axs[n].set_axis_bgcolor('#b0b0b0')

    cax = fig.add_axes([0.125,0.05,0.775,0.025])
    fig.text(0.06,0.51,'depth (m)',ha='center',va='center',rotation=90)
    return fig, axs, cax

Position of sections

fig, ax = plt.subplots(subplot_kw={'projection':ccrs.Robinson(central_longitude=210)})
ax.set_global()
ax.add_feature(cfeature.LAND,facecolor='#b0b0b0')
l1, = ax.plot(ds_sec_j12.TLONG, ds_sec_j12.TLAT,'C0--',transform=ccrs.Geodetic())
l2, = ax.plot(ds_sec_j225.TLONG, ds_sec_j225.TLAT,'C1--',transform=ccrs.Geodetic())

Tritium along 26.5 W

data, dates = ds_sec_j12.TRITIUM[-12:,...], ds_sec_j12.time[-12:,...]
data = np.ma.masked_where(np.isnan(data.values),data.values)
data[data<0] = 0
clev = ticker.MaxNLocator(nbins=12).tick_values(0.005,data.max())
fig, axs, cax = secfig_mon()
for n in range(data.shape[0]):
    pm = axs[n].pcolormesh(ds_sec_j12.TLAT,-ds_sec_j12.z_t/100,data[n,...],vmin=data.min(),vmax=data.max(),
                           cmap=plt.cm.magma_r)
#     cs = axs[n].contour(ds_sec_j12.TLAT,-ds_sec_j12.z_t/100,data[n,...],clev,colors='k',linewidths=0.5)
#     cl = axs[n].clabel(cs,colors='k',fmt='%.2f',fontsize=8)

cb = fig.colorbar(pm,cax=cax,orientation='horizontal')
cb.set_label(r'tritium (pmol m$^{-3}$)')

Tritium decay along 26.5 W

data, dates = ds_sec_j12.TRITIUM_DECAY[-12:,...] * spd * 365, ds_sec_j12.time[-12:,...]
data = np.ma.masked_where(np.isnan(data.values),data.values)
clev = ticker.MaxNLocator(nbins=12).tick_values(0.0005,data.max())
fig, axs, cax = secfig_mon()
for n in range(data.shape[0]):
    pm = axs[n].pcolormesh(ds_sec_j12.TLAT,-ds_sec_j12.z_t/100,data[n,...],vmin=data.min(),vmax=data.max(),
                           cmap=plt.cm.magma_r)
#     cs = axs[n].contour(ds_sec_j12.TLAT,-ds_sec_j12.z_t/100,data[n,...],clev,colors='k',linewidths=0.5)
#     cl = axs[n].clabel(cs,colors='k',fmt='%.4f',fontsize=8)

cb = fig.colorbar(pm,cax=cax,orientation='horizontal')
cb.set_label(r'tritium decay (pmol m$^{-3}$ y$^{-1}$)')

Tritium along 145.9 W

data, dates = ds_sec_j225.TRITIUM[-12:,...], ds_sec_j225.time[-12:,...]
data = np.ma.masked_where(np.isnan(data.values),data.values)
data[data<0] = 0
clev = ticker.MaxNLocator(nbins=10).tick_values(data.min(),data.max())
fig, axs, cax = secfig_mon()
for n in range(data.shape[0]):
    pm = axs[n].pcolormesh(ds_sec_j225.TLAT,-ds_sec_j225.z_t/100,data[n,...],vmin=data.min(),vmax=data.max(),
                           cmap=plt.cm.magma_r)
#     cs = axs[n].contour(ds_sec_j225.TLAT,-ds_sec_j225.z_t/100,data[n,...],clev,colors='k',linewidths=0.5)
#     cl = axs[n].clabel(cs,colors='k',fmt='%.3f',fontsize=8)

cb = fig.colorbar(pm,cax=cax,orientation='horizontal')
cb.set_label(r'tritium (pmol m$^{-3}$)')

Tritium decay along 145.9 W

data, dates = ds_sec_j225.TRITIUM_DECAY[-12:,...] * spd * 365, ds_sec_j225.time[-12:,...]
data = np.ma.masked_where(np.isnan(data.values),data.values)
clev = ticker.MaxNLocator(nbins=10).tick_values(data.min(),data.max())
fig, axs, cax = secfig_mon()
for n in range(data.shape[0]):
    pm = axs[n].pcolormesh(ds_sec_j225.TLAT,-ds_sec_j225.z_t/100,data[n,...],vmin=data.min(),vmax=data.max(),
                           cmap=plt.cm.magma_r)
#     cs = axs[n].contour(ds_sec_j225.TLAT,-ds_sec_j225.z_t/100,data[n,...],clev,colors='k',linewidths=0.5)
#     cl = axs[n].clabel(cs,colors='k',fmt='%.3f',fontsize=8)

cb = fig.colorbar(pm,cax=cax,orientation='horizontal')
cb.set_label(r'tritium decay (pmol m$^{-3}$ y$^{-1}$)')