Full Variance Budget for Busecke and Abernathey 2017

In this notebook I evaluate the full tracer variance budget for all runs assembled for the 2017 ENSO diffusivity paper.

The full budget (from Abernathey and Marshall 2013) can be written as:

$$ \underbrace{\frac{\partial}{\partial t}\frac{\overline{{q^\prime}^2}}{2}} _\text{tendency} + \underbrace{\nabla \cdot \overline{v \frac{{q^\prime}^2}{2}}} _\text{advection}+ \underbrace{\overline{v^\prime q^\prime} \cdot \nabla \overline{q}} _\text{production}= \underbrace{\nabla^2 \left({\kappa \overline{ \frac{{q^\prime}^2}{2}}}\right)} _\text{diffusion} - \underbrace{\kappa \overline{|{\nabla q^\prime}|^2}} _\text{destruction}$$

I do have the following

% For the paper: \usepackage{amsmath}

The following terms are put out as diagnostics from the model:

  • $$\overline{q}$$
  • $$\overline{q^2}$$
  • $$

The terms are calculated as follows:

Production: $$ \overline{v^\prime q^\prime} \cdot \nabla \overline{q} = \overline{u^\prime q^\prime} \frac{\partial \overline{q}}{\partial x} + \overline{v^\prime q^\prime} \frac{\partial \overline{q}}{\partial y} = (\overline{uq} - \overline{u} \, \overline{q}) \frac{\partial \overline{q}}{\partial x} + (\overline{vq} - \overline{v} \, \overline{q}) \frac{\partial \overline{q}}{\partial y} $$

Advection: $$ \nabla = \overline{ \vec{u} \frac{{q^\prime}^2}{2}} = \frac{\partial}{\partial x} \overline{u \frac{{q^\prime}^2}{2}} $$


In [1]:
from dask.distributed import Client
client=Client(scheduler_file='/rigel/home/jb3210/scheduler.json')
client


Out[1]:

Client

Cluster

  • Workers: 41
  • Cores: 164
  • Memory: 492.00 GB

In [2]:
import xarray as xr
import numpy as np
from xmitgcm import open_mdsdataset
from xgcm import Grid
from os.path import join as pjoin
%matplotlib inline

In [3]:
# Basic vector algebra (This should update the aviso tracer repo)
def _get_matching(da, da_list):
    """from a list of xarray data arrays (da_list) return the one with dims 
    matching obj"""
    matched = [a for a in da_list if set(a.dims).issubset(set(da.dims))]
    if len(matched) == 0:
        raise RuntimeError('no matching coordinates found. Try interpolating them using xgcm')
    elif len(matched) >1 :
        raise RuntimeError('More than one matching coord found, check list of data arrays')
    else:
        matched = matched[0]
    return matched

def _get_dx(da, grid_ds):
    """get horizontal distance"""
    dx_list = [grid_ds[a] for a in ['dxC', 'dxG']]
    return _get_matching(da,dx_list)

def _get_dy(da, grid_ds):
    """get horizontal distance"""
    dy_list = [grid_ds[a] for a in ['dyC', 'dyG']]
    return _get_matching(da,dy_list)

def _get_hfac(da, grid_ds):
    """get horizontal distance"""
    hfac_list = [grid_ds[a] for a in ['hFacC', 'hFacW', 'hFacS']]
    return _get_matching(da,hfac_list)
    
def _get_area(da, grid_ds):
    area_list = [grid_ds[a] for a in ['rA', 'rAw', 'rAs']]
    return _get_matching(da,area_list)

def gradient(da, grid):
    delta_x = grid.diff(da, 'X')
    delta_y = grid.diff(da, 'Y')
    dx = _get_dx(delta_x,grid._ds)
    dy = _get_dy(delta_y,grid._ds)
    da_dx = delta_x/dx
    da_dy = delta_y/dy
    return da_dx, da_dy

def dot_prod(u1, v1, u2, v2, grid, ref):
    u = u1*u2
    v = v1*v2
    
    dyu = _get_dy(u, grid._ds)
    dxv = _get_dx(v, grid._ds)
    hfacu = _get_hfac(u, grid._ds)
    hfacv = _get_hfac(v, grid._ds)
    
    u_int = u*dyu*hfacu
    v_int = v*dxv*hfacv
    dot_prod_int = xr.DataArray(u_int.data+v_int.data,
                                coords=ref.coords, dims=ref.dims)
    area = _get_area(dot_prod_int, grid._ds)
    hfac = _get_hfac(dot_prod_int, grid._ds)
    return dot_prod_int/area/hfac
    

def divergence(u, v, grid):
    dyu = _get_dy(u, grid._ds)
    dxv = _get_dx(v, grid._ds)
    hfacu = _get_hfac(u, grid._ds)
    hfacv = _get_hfac(v, grid._ds)
    delta_u = grid.diff(u*dyu*hfacu, 'X')
    delta_v = grid.diff(v*dxv*hfacv, 'Y')
    area = _get_area(delta_u, grid._ds)
    hfac = _get_hfac(delta_u, grid._ds)
    return (delta_u+delta_v)/area/hfac

def laplacian(da, grid):
    gradx, grady = gradient(da, grid)
    return divergence(gradx, grady, grid)

In [4]:
# import matplotlib.pyplot as plt
# # What am I doing wrong with the divergence?
# test_u = ds_vel['UVEL'].isel(time=45)
# test_v = ds_vel['VVEL'].isel(time=45)

# test_dx_u = grid.interp(grid.diff(test_u, 'X'), 'X')
# test_dy_v = grid.interp(grid.diff(test_v, 'Y'), 'Y')
# test_div_alt = grid.interp(test_dx_u/_get_dx(test_dx_u, grid._ds), 'X') + \
#                     grid.interp(test_dy_v/_get_dy(test_dy_v, grid._ds), 'Y')
# test_div = divergence(test_u,test_v, grid)

In [5]:
# plt.figure()
# test_div.plot(vmin=-5e-8, vmax=5e-8, cmap=plt.cm.RdBu_r)
# plt.figure()
# test_div_alt.plot(vmin=-5e-8, vmax=5e-8, cmap=plt.cm.RdBu_r)

# # plt.figure()
# # test_div/test_div_alt.plot(robust=True)

In [6]:
# This could potentially be integrated in xmitgcm
def convert_trnum2dimension(ds):
    def drop_nonmatching_vars(ds, var_list):
        data_vars = list(ds.data_vars)
        drop_vars = [a for a in data_vars if a not in var_list]
        return ds.drop(drop_vars)
    
    def rename_vars(ds):
        for vv in list(ds.data_vars):
            ds = ds.rename({vv:vv[0:-2]})
        return ds
        
    tr_num = list(set([a[-2:] for a in list(ds.data_vars)]))
    tr_num.sort()
    tr_vars = list(set([a[:-2] for a in list(ds.data_vars)]))
    
    datasets = [drop_nonmatching_vars(ds, [a+n for a in tr_vars]) for n in tr_num]
    datasets = [rename_vars(b) for b in datasets]
    
    tr_num_int = [int(a) for a in tr_num]
    
    tr_dim = xr.DataArray(tr_num_int, 
                          dims='tracer_no',
                          coords={'tracer_no': (['tracer_no', ], tr_num_int)})
    
    ds_combined = xr.concat(datasets,dim='tracer_no')
    ds_combined['tracer_no'] = tr_dim
    
    return ds_combined

# These should all go to the final version of the github repo

In [7]:
ddir = '/rigel/ocp/users/jb3210/projects/aviso_surface_tracer/runs'

run = 'run_KOC_PSI_variance_budget'
rundir = pjoin(ddir,run)
timestep = 900 # in seconds
readin_dict = dict(delta_t=timestep, swap_dims=False)
kappa = 63

In [8]:
ds_mean = convert_trnum2dimension(open_mdsdataset(rundir,prefix=['tracer_diags'],
                                                  **readin_dict))
variance_mean = (ds_mean['TRACSQ']-(ds_mean['TRAC']**2))/2
# ds_mean


/rigel/home/jb3210/src/xmitgcm/xmitgcm/utils.py:314: UserWarning: Not sure what to do with rlev = L
  warnings.warn("Not sure what to do with rlev = " + rlev)
/rigel/home/jb3210/src/xmitgcm/xmitgcm/mds_store.py:220: FutureWarning: iteration over an xarray.Dataset will change in xarray v0.11 to only include data variables, not coordinates. Iterate over the Dataset.variables property instead to preserve existing behavior in a forwards compatible manner.
  for vname in ds:

In [9]:
readin_dict = dict(delta_t=timestep, swap_dims=False)
ds_snap = convert_trnum2dimension(open_mdsdataset(rundir,prefix=['tracer_snapshots'],
                                                  **readin_dict))
variance_snap = (ds_snap['TRACSQ']-(ds_snap['TRAC']**2))/2
# ds_snap


/rigel/home/jb3210/src/xmitgcm/xmitgcm/utils.py:314: UserWarning: Not sure what to do with rlev = L
  warnings.warn("Not sure what to do with rlev = " + rlev)
/rigel/home/jb3210/src/xmitgcm/xmitgcm/mds_store.py:220: FutureWarning: iteration over an xarray.Dataset will change in xarray v0.11 to only include data variables, not coordinates. Iterate over the Dataset.variables property instead to preserve existing behavior in a forwards compatible manner.
  for vname in ds:

In [10]:
readin_dict = dict(delta_t=timestep, swap_dims=False)
ds_vel = open_mdsdataset(rundir,prefix=['vel_diags'], **readin_dict)
# ds_vel


/rigel/home/jb3210/src/xmitgcm/xmitgcm/utils.py:314: UserWarning: Not sure what to do with rlev = L
  warnings.warn("Not sure what to do with rlev = " + rlev)
/rigel/home/jb3210/src/xmitgcm/xmitgcm/mds_store.py:220: FutureWarning: iteration over an xarray.Dataset will change in xarray v0.11 to only include data variables, not coordinates. Iterate over the Dataset.variables property instead to preserve existing behavior in a forwards compatible manner.
  for vname in ds:

In [11]:
# xgcm grid
grid = Grid(ds_mean)

In [12]:
# # Variance Destruction
q_mean_gradx, q_mean_grady = gradient(ds_mean['TRAC'], grid)
destruction_x = ds_mean['DXSqTr']-(q_mean_gradx**2)
destruction_y = ds_mean['DYSqTr']-(q_mean_grady**2)
destruction = -kappa*dot_prod(destruction_x, destruction_y, 1, 1, grid, ds_mean['TRAC'])
destruction


Out[12]:
<xarray.DataArray (tracer_no: 2, time: 292, j: 1600, i: 3600)>
dask.array<shape=(2, 292, 1600, 3600), dtype=float32, chunksize=(1, 1, 1600, 3600)>
Coordinates:
  * tracer_no  (tracer_no) int64 1 2
  * i          (i) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 ...
  * j          (j) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 ...
    XC         (j, i) float32 0.05 0.15 0.25 0.35 0.45 0.55 0.65 0.75 0.85 ...
    YC         (j, i) float32 -79.95 -79.95 -79.95 -79.95 -79.95 -79.95 ...
    rA         (j, i) float32 2.15699e+07 2.15699e+07 2.15699e+07 ...
    Depth      (j, i) float32 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ...
    hFacC      (j, i) float32 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ...
    iter       (time) int64 2880 5760 8640 11520 14400 17280 20160 23040 ...
  * time       (time) int64 2592000 5184000 7776000 10368000 12960000 ...

In [13]:
#Variance advection
q_on_u = grid.interp(ds_mean['TRAC'], 'X')
q_on_v = grid.interp(ds_mean['TRAC'], 'Y')
u = ds_vel['UVEL']
v = ds_vel['VVEL']
uq = ds_mean['UTRAC']
vq = ds_mean['VTRAC']
uq_sq = ds_mean['UTrSq']
vq_sq = ds_mean['VTrSq']

u_q_pr_sq_mean = uq_sq - (uq*q_on_u) - (2*u*q_on_u**2)
v_q_pr_sq_mean = vq_sq - (vq*q_on_v) - (2*v*q_on_v**2)
advection = -divergence(u_q_pr_sq_mean, v_q_pr_sq_mean, grid)
advection


Out[13]:
<xarray.DataArray (tracer_no: 2, time: 292, j: 1600, i: 3600)>
dask.array<shape=(2, 292, 1600, 3600), dtype=float32, chunksize=(1, 1, 1600, 3600)>
Coordinates:
  * tracer_no  (tracer_no) int64 1 2
  * time       (time) int64 2592000 5184000 7776000 10368000 12960000 ...
  * j          (j) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 ...
  * i          (i) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 ...
    XC         (j, i) float32 0.05 0.15 0.25 0.35 0.45 0.55 0.65 0.75 0.85 ...
    YC         (j, i) float32 -79.95 -79.95 -79.95 -79.95 -79.95 -79.95 ...
    rA         (j, i) float32 2.15699e+07 2.15699e+07 2.15699e+07 ...
    Depth      (j, i) float32 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ...
    hFacC      (j, i) float32 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ...

In [14]:
# Variance production
q_mean_gradx, q_mean_grady = gradient(ds_mean['TRAC'], grid)
q_pr_u_pr = (ds_mean['UTRAC']-(grid.interp(ds_mean['TRAC'],'X')*ds_vel['UVEL']))
q_pr_v_pr = (ds_mean['VTRAC']-(grid.interp(ds_mean['TRAC'],'Y')*ds_vel['VVEL']))
production = -dot_prod(q_pr_u_pr, q_pr_v_pr, q_mean_gradx, q_mean_grady, grid, ds_mean['TRAC'])
production

# production_x = (ds_mean['UTRAC']-(grid.interp(ds_mean['TRAC'],'X')*ds_vel['UVEL']))*q_mean_gradx
# production_y = (ds_mean['VTRAC']-(grid.interp(ds_mean['TRAC'],'Y')*ds_vel['VVEL']))*q_mean_grady
# production_x_int = production_x*_get_dy(production_x, grid._ds)
# production_y_int = production_y*_get_dx(production_y, grid._ds)
# production_int = xr.DataArray(production_x_int.data+production_y_int.data,
#                              coords=ds_mean['TRAC'].coords, dims=ds_mean['TRAC'].dims)
# production = production_int/_get_area(production_int, grid._ds)


Out[14]:
<xarray.DataArray (tracer_no: 2, time: 292, j: 1600, i: 3600)>
dask.array<shape=(2, 292, 1600, 3600), dtype=float32, chunksize=(1, 1, 1600, 3600)>
Coordinates:
  * tracer_no  (tracer_no) int64 1 2
  * i          (i) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 ...
  * j          (j) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 ...
    XC         (j, i) float32 0.05 0.15 0.25 0.35 0.45 0.55 0.65 0.75 0.85 ...
    YC         (j, i) float32 -79.95 -79.95 -79.95 -79.95 -79.95 -79.95 ...
    rA         (j, i) float32 2.15699e+07 2.15699e+07 2.15699e+07 ...
    Depth      (j, i) float32 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ...
    hFacC      (j, i) float32 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ...
    iter       (time) int64 2880 5760 8640 11520 14400 17280 20160 23040 ...
  * time       (time) int64 2592000 5184000 7776000 10368000 12960000 ...

In [15]:
# Variance tendency
q_pr_sq_snap = (ds_snap['TRACSQ']-(ds_snap['TRAC']**2))/2
dt = ds_snap['TRACSQ'].time[0:2].diff('time').data
tendency = (q_pr_sq_snap/2).diff('time')/dt
tendency


Out[15]:
<xarray.DataArray (tracer_no: 2, time: 292, j: 1600, i: 3600)>
dask.array<shape=(2, 292, 1600, 3600), dtype=float64, chunksize=(1, 1, 1600, 3600)>
Coordinates:
  * tracer_no  (tracer_no) int64 1 2
  * i          (i) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 ...
  * j          (j) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 ...
    XC         (j, i) float32 0.05 0.15 0.25 0.35 0.45 0.55 0.65 0.75 0.85 ...
    YC         (j, i) float32 -79.95 -79.95 -79.95 -79.95 -79.95 -79.95 ...
    rA         (j, i) float32 2.15699e+07 2.15699e+07 2.15699e+07 ...
    Depth      (j, i) float32 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ...
    hFacC      (j, i) float32 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ...
    iter       (time) int64 2880 5760 8640 11520 14400 17280 20160 23040 ...
  * time       (time) int64 2592000 5184000 7776000 10368000 12960000 ...

In [16]:
# Variance diffusion
diffusion = laplacian(variance_mean*kappa,grid)
diffusion


Out[16]:
<xarray.DataArray (tracer_no: 2, time: 292, j: 1600, i: 3600)>
dask.array<shape=(2, 292, 1600, 3600), dtype=float32, chunksize=(1, 1, 1600, 3600)>
Coordinates:
  * tracer_no  (tracer_no) int64 1 2
  * time       (time) int64 2592000 5184000 7776000 10368000 12960000 ...
  * j          (j) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 ...
  * i          (i) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 ...
    XC         (j, i) float32 0.05 0.15 0.25 0.35 0.45 0.55 0.65 0.75 0.85 ...
    YC         (j, i) float32 -79.95 -79.95 -79.95 -79.95 -79.95 -79.95 ...
    rA         (j, i) float32 2.15699e+07 2.15699e+07 2.15699e+07 ...
    Depth      (j, i) float32 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ...
    hFacC      (j, i) float32 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 ...

In [17]:
# alternative way to compute diffusion
variance_gradx, variance_grady = gradient(variance_mean, grid)
variance_gradxx,_ = gradient(grid.interp(variance_gradx, 'X'), grid)
_,variance_gradyy = gradient(grid.interp(variance_grady, 'Y'), grid)
diffusion_alt = kappa*(grid.interp(variance_gradxx, 'X')+
                       grid.interp(variance_gradyy, 'Y'))
diffusion_alt


Out[17]:
<xarray.DataArray (tracer_no: 2, time: 292, j: 1600, i: 3600)>
dask.array<shape=(2, 292, 1600, 3600), dtype=float32, chunksize=(1, 1, 1600, 3600)>
Coordinates:
  * tracer_no  (tracer_no) int64 1 2
  * time       (time) int64 2592000 5184000 7776000 10368000 12960000 ...
  * j          (j) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 ...
  * i          (i) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 ...

In [18]:
# sum of all terms
all_sum = diffusion+advection+destruction+production
part_sum = destruction+production

In [19]:
# quick order of magnitude sanity check.
combo = xr.concat([tendency.drop(tendency.coords),
                   diffusion.drop(diffusion.coords),
                   advection.drop(advection.coords),
                   destruction.drop(destruction.coords),
                   production.drop(production.coords),
                   all_sum.drop(all_sum.coords),
                   part_sum.drop(part_sum.coords)],
                dim=['tendency',
                     'diffusion',
                     'advection',
                     'destruction',
                     'production',
                     'sum',
                     'part_sum'])
combo


Out[19]:
<xarray.DataArray (concat_dim: 7, tracer_no: 2, time: 292, j: 1600, i: 3600)>
dask.array<shape=(7, 2, 292, 1600, 3600), dtype=float64, chunksize=(1, 1, 1, 1600, 3600)>
Coordinates:
  * concat_dim  (concat_dim) <U11 'tendency' 'diffusion' 'advection' ...
Dimensions without coordinates: tracer_no, time, j, i

In [20]:
# import matplotlib as mpl

In [21]:
# combo.isel(time=slice(30,40)).mean('time').plot(col='concat_dim',row='tracer_no',robust=True, norm=mpl.colors.SymLogNorm(1e-16))

In [22]:
# combo.mean('time').plot(col='concat_dim',row='tracer_no',robust=True, norm=mpl.colors.SymLogNorm(1e-18))

In [ ]:
combo.isel(concat_dim=[0, 3, 4, 6]).mean('time').plot(col='concat_dim',
                                                      row='tracer_no',robust=True)

In [ ]: