Meridional Overturning

mom6_tools.moc collection of functions for computing and plotting meridional overturning circulation.

The goal of this notebook is the following:

  1. server as an example on to compute a meridional overturning streamfunction (global and Atalntic) from CESM/MOM output;

  2. evaluate model experiments by comparing transports against observed estimates and other model results.

[1]:

%load_ext autoreload
%autoreload 2

[1]:

%matplotlib inline
import warnings
warnings.filterwarnings("ignore")
import matplotlib
import numpy as np
import xarray as xr
# mom6_tools
from mom6_tools.MOM6grid import MOM6grid
from mom6_tools.moc import  *
from ncar_jobqueue import NCARCluster
from dask.distributed import Client
from mom6_tools.m6toolbox import genBasinMasks, add_global_attrs
from mom6_tools.m6toolbox import cime_xmlquery
import matplotlib.pyplot as plt
import yaml, os, intake

Basemap module not found. Some regional plots may not function properly

ERROR 1: PROJ: proj_create_from_database: Open of /glade/work/gmarques/conda-envs/mom6-tools/share/proj failed

[3]:

# Read in the yaml file
diag_config_yml_path = "diag_config.yml"
diag_config_yml = yaml.load(open(diag_config_yml_path,'r'), Loader=yaml.Loader)

[4]:

caseroot = diag_config_yml['Case']['CASEROOT']
casename = cime_xmlquery(caseroot, 'CASE')
DOUT_S = cime_xmlquery(caseroot, 'DOUT_S')
if DOUT_S:
  OUTDIR = cime_xmlquery(caseroot, 'DOUT_S_ROOT')+'/ocn/hist/'
else:
  OUTDIR = cime_xmlquery(caseroot, 'RUNDIR')

print('Output directory is:', OUTDIR)
print('Casename is:', casename)

Output directory is: /glade/derecho/scratch/gmarques/archive/g.e30_a03c.GJRAv4.TL319_t232_wgx3_hycom1_N75.2024.079/ocn/hist/
Casename is: g.e30_a03c.GJRAv4.TL319_t232_wgx3_hycom1_N75.2024.079

[5]:

# The following parameters must be set accordingly
######################################################
# add your name and email address below
author = 'Gustavo Marques (gmarques@ucar.edu)'
######################################################
# create an empty class object
class args:
  pass

# load avg dates
avg = diag_config_yml['Avg']

args.infile = OUTDIR
args.monthly = casename+diag_config_yml['Fnames']['z']
args.sigma2 = casename+diag_config_yml['Fnames']['rho2']
args.static = casename+diag_config_yml['Fnames']['static']
args.geom =   casename+diag_config_yml['Fnames']['geom']
args.start_date = avg['start_date']
args.end_date = avg['end_date']
args.case_name = casename
args.label = ''
args.savefigs = False

[6]:

# read grid info
geom_file = OUTDIR+'/'+args.geom
if os.path.exists(geom_file):
  grd = MOM6grid(OUTDIR+'/'+args.static, geom_file)
else:
  grd = MOM6grid(OUTDIR+'/'+args.static)

try:
  depth = grd.depth_ocean
except:
  depth = grd.deptho

MOM6 grid successfully loaded...


[7]:

# remove Nan's, otherwise genBasinMasks won't work
depth[np.isnan(depth)] = 0.0
basin_code = genBasinMasks(grd.geolon, grd.geolat, depth, verbose=False)
basin_code_xr = genBasinMasks(grd.geolon, grd.geolat, depth, verbose=False, xda=True)

[8]:

cluster = NCARCluster()
cluster.scale(6)
client = Client(cluster)
client

[8]:

Client

Client-57b39531-9e42-11ef-9c8d-3cecef1b11de

Connection method: Cluster object Cluster type: dask_jobqueue.PBSCluster
Dashboard: https://jupyterhub.hpc.ucar.edu/stable/user/gmarques/High-mem/proxy/8787/status

Cluster Info

 
[9]:

def preprocess(ds):
    variables = ['vmo','vhml','vhGM']
    for v in variables:
      if v not in ds.variables:
        ds[v] = xr.zeros_like(ds.vo)
    return ds[variables]

[10]:

print('\n Reading dataset...')
# load data
%time ds = xr.open_mfdataset(OUTDIR+'/'+args.monthly, parallel=True, \
                             combine="nested", concat_dim="time", \
                             preprocess=preprocess).chunk({"time": 12})


 Reading dataset...
CPU times: user 4.56 s, sys: 249 ms, total: 4.81 s
Wall time: 23.8 s

[10]:

attrs =  {
         'description': 'Annual mean meridional thickness flux by components ',
         'reduction_method': 'annual mean weighted by days in each month',
         'casename': casename
         }

[12]:

ds_ann =  m6toolbox.weighted_temporal_mean_vars(ds,attrs=attrs)

[13]:

print('\n Selecting data between {} and {}...'.format(args.start_date, args.end_date))
%time ds_sel = ds_ann.sel(time=slice(args.start_date, args.end_date))


 Selecting data between 0031-01-01 and 0062-01-01...
CPU times: user 6.25 ms, sys: 0 ns, total: 6.25 ms
Wall time: 8.52 ms

Compute temporal mean for each term

[14]:

stream = True
# create a ndarray subclass
class C(np.ndarray): pass

[15]:

print('\n Computing time mean...')
%time ds_mean = ds_sel.mean('time').compute()


 Computing time mean...
CPU times: user 7.07 s, sys: 703 ms, total: 7.78 s
Wall time: 1min 30s

[16]:

# create a ndarray subclass
class C(np.ndarray): pass

if 'vmo' in ds.variables:
  varName = 'vmo'; conversion_factor = 1.e-9
elif 'vh' in ds.variables:
  varName = 'vh'; conversion_factor = 1.e-6
  if 'zw' in ds.variables: conversion_factor = 1.e-9 # Backwards compatible for when we had wrong units for 'vh'
else: raise Exception('Could not find "vh" or "vmo" in file "%s"'%(args.infile+args.monthly))

tmp = np.ma.masked_invalid(ds_sel[varName].mean('time').values)
tmp = tmp[:].filled(0.)
VHmod = tmp.view(C)
VHmod.units = ds_sel[varName].units

Zmod = m6toolbox.get_z(ds, depth, varName)

if args.case_name != '':  case_name = args.case_name + ' ' + args.label
else: case_name = rootGroup.title + ' ' + args.label

Global MOC

[17]:

%matplotlib inline

# Global MOC
m6plot.setFigureSize([16,9],576,debug=False)
axis = plt.gca()
cmap = plt.get_cmap('dunnePM')
zg = Zmod.min(axis=-1); psiPlot = MOCpsi(VHmod)*conversion_factor
psiPlot = 0.5 * (psiPlot[0:-1,:]+psiPlot[1::,:])
yyg = grd.geolat_c[:,:].max(axis=-1)+0*zg
ci=m6plot.pmCI(0.,40.,5.)
plotPsi(yyg, zg, psiPlot, ci, 'Global MOC [Sv]')
plt.xlabel(r'Latitude [$\degree$N]')
plt.suptitle(case_name)
findExtrema(yyg, zg, psiPlot, max_lat=-30.)
findExtrema(yyg, zg, psiPlot, min_lat=25., min_depth=250.)
findExtrema(yyg, zg, psiPlot, min_depth=2000., mult=-1.)
plt.gca().invert_yaxis()
../_images/examples_meridional_overturning_19_0.png 
[18]:

# create dataset to store results
moc = xr.Dataset(data_vars={ 'moc' :    (('z_l','yq'), psiPlot),
                            'amoc' :   (('z_l','yq'), np.zeros((psiPlot.shape))),
                            'moc_FFM' :   (('z_l','yq'), np.zeros((psiPlot.shape))),
                            'moc_GM' : (('z_l','yq'), np.zeros((psiPlot.shape))),
                            'amoc_45' : (('time'), np.zeros((ds_ann.time.shape))),
                            'moc_GM_ACC' : (('time'), np.zeros((ds_ann.time.shape))),
                            'amoc_26' : (('time'), np.zeros((ds_ann.time.shape))) },
                            coords={'z_l': ds.z_l, 'yq':ds.yq, 'time':ds_ann.time})
attrs = {'description': 'MOC time-mean sections and time-series', 'unit': 'Sv', 'start_date': avg['start_date'],
       'end_date': avg['end_date']}
add_global_attrs(moc,attrs)

[19]:

print('Saving netCDF files...')
if not os.path.isdir('ncfiles'):
      os.system('mkdir -p ncfiles')
moc.to_netcdf('ncfiles/'+str(casename)+'_MOC.nc')

Saving netCDF files...

Atlantic MOC

[20]:

m6plot.setFigureSize([16,9],576,debug=False)
cmap = plt.get_cmap('dunnePM')
m = 0*basin_code; m[(basin_code==2) | (basin_code==4) | (basin_code==6) | (basin_code==7) | (basin_code==8)]=1
ci=m6plot.pmCI(0.,22.,2.)
z = (m*Zmod).min(axis=-1); psiPlot = MOCpsi(VHmod, vmsk=m*np.roll(m,-1,axis=-2))*conversion_factor
psiPlot = 0.5 * (psiPlot[0:-1,:]+psiPlot[1::,:])
yy = grd.geolat_c[:,:].max(axis=-1)+0*z
plotPsi(yy, z, psiPlot, ci, 'Atlantic MOC [Sv]')
plt.xlabel(r'Latitude [$\degree$N]')
plt.suptitle(case_name)
findExtrema(yy, z, psiPlot, min_lat=26.5, max_lat=27., min_depth=250.) # RAPID
findExtrema(yy, z, psiPlot, min_lat=44, max_lat=46., min_depth=250.) # RAPID
findExtrema(yy, z, psiPlot, max_lat=-33.)
findExtrema(yy, z, psiPlot)
findExtrema(yy, z, psiPlot, min_lat=5.)
plt.gca().invert_yaxis()
moc['amoc'].data = psiPlot
../_images/examples_meridional_overturning_23_0.png

AMOC profile at 26N

[21]:

catalog = intake.open_catalog(diag_config_yml['oce_cat'])
rapid_vertical = catalog["moc-rapid"].to_dask()

[22]:

if 'zl' in ds:
  zl=ds.zl.values
elif 'z_l' in ds:
  zl=ds.z_l.values
else:
  raise ValueError("Dataset does not have vertical coordinate zl or z_l")

[23]:

fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(8, 6))
ax.plot(rapid_vertical.stream_function_mar.mean('time'), rapid_vertical.depth, 'k', label='RAPID')
ax.plot(moc['amoc'].sel(yq=26, method='nearest'), moc.z_l, label=case_name)
ax.legend()
plt.gca().invert_yaxis()
plt.grid()
ax.set_xlabel('AMOC @ 26N [Sv]')
ax.set_ylabel('Depth [m]');
../_images/examples_meridional_overturning_27_0.png

AMOC time series

[24]:

dtime = ds_ann.time.values

amoc_26 = np.zeros(len(dtime))
amoc_45 = np.zeros(len(dtime))
moc_GM_ACC = np.zeros(len(dtime))

# loop in time
for t in range(len(dtime)):
    tmp = np.ma.masked_invalid(ds_ann[varName].sel(time=dtime[t]).values)
    tmp = tmp[:].filled(0.)
    psi = MOCpsi(tmp, vmsk=m*np.roll(m,-1,axis=-2))*conversion_factor
    psi = 0.5 * (psi[0:-1,:]+psi[1::,:])
    amoc_26[t] = findExtrema(yy, z, psi, min_lat=26.5, max_lat=27., plot=False)
    amoc_45[t] = findExtrema(yy, z, psi, min_lat=44., max_lat=46., plot=False)
    tmp_GM = np.ma.masked_invalid(ds_ann['vhGM'][t,:].values)
    tmp_GM = tmp_GM[:].filled(0.)
    psiGM = MOCpsi(tmp_GM)*conversion_factor
    psiGM = 0.5 * (psiGM[0:-1,:]+psiGM[1::,:])
    moc_GM_ACC[t] = findExtrema(yyg, zg, psiGM, min_lat=-65., max_lat=-30, mult=-1., plot=False)

[25]:

# add dataarays to the moc dataset
moc['amoc_26'].data = amoc_26
moc['amoc_45'].data = amoc_45
moc['moc_GM_ACC'].data = moc_GM_ACC

[27]:

# load datasets from oce catalog
amoc_core_26 = catalog["moc-core2-26p5"].to_dask()
amoc_pop_26  = catalog["moc-pop-jra-26"].to_dask()
rapid = m6toolbox.weighted_temporal_mean_vars(catalog["transports-rapid"].to_dask())

amoc_core_45 = catalog["moc-core2-45"].to_dask()

amoc_pop_45 = catalog["moc-pop-jra-45"].to_dask()

#list(catalog)

AMOC @ 26 \(^o\) N

[28]:

# plot
fig = plt.figure(figsize=(12, 6))
plt.plot(np.arange(len(moc.time))+1958.5 ,moc['amoc_26'].values, color='k', label=case_name, lw=2)
# core data
core_mean = amoc_core_26['MOC'].mean(axis=0).data
core_std = amoc_core_26['MOC'].std(axis=0).data
plt.plot(amoc_core_26.time,core_mean, 'k', label='CORE II (group mean)', color='#1B2ACC', lw=1)
plt.fill_between(amoc_core_26.time, core_mean-core_std, core_mean+core_std,
  alpha=0.25, edgecolor='#1B2ACC', facecolor='#089FFF')
# pop data
plt.plot(np.arange(len(amoc_pop_26.time))+1958.5 ,amoc_pop_26.AMOC_26n.values, color='r', label='POP', lw=1)
# rapid
plt.plot(np.arange(len(rapid.time))+2004.5 ,rapid.moc_mar_hc10.values, color='green', label='RAPID', lw=1)
#plt.plot(np.arange(len(rapid_filtered.time))+2004.5 ,rapid_filtered.values, color='green', label='RAPID', lw=1)

plt.title('AMOC @ 26 $^o$ N', fontsize=16)
plt.ylim(5,20)
plt.xlim(1948,1958.5+len(moc.time))
plt.xlabel('Time [years]', fontsize=16); plt.ylabel('Sv', fontsize=16)
plt.legend(fontsize=13, ncol=2)

[28]:

<matplotlib.legend.Legend at 0x1492608e2910>
../_images/examples_meridional_overturning_33_1.png

AMOC @ 45 \(^o\) N

[29]:

# plot
fig = plt.figure(figsize=(12, 6))
plt.plot(np.arange(len(moc.time))+1958.5 ,moc['amoc_45'].values, color='k', label=case_name, lw=2)
# core data
core_mean = amoc_core_45['MOC'].mean(axis=0).data
core_std = amoc_core_45['MOC'].std(axis=0).data
plt.plot(amoc_core_45.time,core_mean, 'k', label='CORE II (group mean)', color='#1B2ACC', lw=1)
plt.fill_between(amoc_core_45.time, core_mean-core_std, core_mean+core_std,
  alpha=0.25, edgecolor='#1B2ACC', facecolor='#089FFF')
# pop data
plt.plot(np.arange(len(amoc_pop_45.time))+1958. ,amoc_pop_45.AMOC_45n.values, color='r', label='POP', lw=1)

plt.title('AMOC @ 45 $^o$ N', fontsize=16)
plt.ylim(5,20)
plt.xlim(1948,1958+len(moc.time))
plt.xlabel('Time [years]', fontsize=16); plt.ylabel('Sv', fontsize=16)
plt.legend(fontsize=13, ncol=3)

[29]:

<matplotlib.legend.Legend at 0x149256dcf090>
../_images/examples_meridional_overturning_35_1.png

Submesoscale-induced Global MOC

[30]:

# create a ndarray subclass
class C(np.ndarray): pass

if 'vhml' in ds.variables:
  varName = 'vhml'; conversion_factor = 1.e-9
else: raise Exception('Could not find "vhml" in file "%s"'%(args.infile+args.monthly))

tmp = np.ma.masked_invalid(ds_mean[varName].values)
tmp = tmp[:].filled(0.)
VHmod = tmp.view(C)
VHmod.units = ds[varName].units

# Global MOC
m6plot.setFigureSize([16,9],576,debug=False)
axis = plt.gca()
cmap = plt.get_cmap('dunnePM')
z = Zmod.min(axis=-1); psiPlot = MOCpsi(VHmod)*conversion_factor
psiPlot = 0.5 * (psiPlot[0:-1,:]+psiPlot[1::,:])
#yy = y[1:,:].max(axis=-1)+0*z
yy = grd.geolat_c[:,:].max(axis=-1)+0*z
ci=m6plot.pmCI(0.,20.,2.)
plotPsi(yy, z, psiPlot, ci, 'Global MOC [Sv] due to vhML', zval=[0.,-400.,-6500.])
plt.xlabel(r'Latitude [$\degree$N]')
plt.suptitle(case_name)
plt.gca().invert_yaxis()
moc['moc_FFM'].data = psiPlot
../_images/examples_meridional_overturning_37_0.png

Eddy(GM)-induced Global MOC

[31]:

# create a ndarray subclass
class C(np.ndarray): pass

if 'vhGM' in ds.variables:
  varName = 'vhGM'; conversion_factor = 1.e-9
else: raise Exception('Could not find "vhGM" in file "%s"'%(args.infile+args.monthly))

tmp = np.ma.masked_invalid(ds_mean[varName].values)
tmp = tmp[:].filled(0.)
VHmod = tmp.view(C)
VHmod.units = ds[varName].units

# Global MOC
m6plot.setFigureSize([16,9],576,debug=False)
axis = plt.gca()
cmap = plt.get_cmap('dunnePM')
z = Zmod.min(axis=-1); psiPlot = MOCpsi(VHmod)*conversion_factor
psiPlot = 0.5 * (psiPlot[0:-1,:]+psiPlot[1::,:])
yy = grd.geolat_c[:,:].max(axis=-1)+0*z
ci=m6plot.pmCI(0.,20.,1.)
plotPsi(yy, z, psiPlot, ci, 'Global MOC [Sv] due to GM')
plt.xlabel(r'Latitude [$\degree$N]')
plt.suptitle(case_name)
findExtrema(yy, z, psiPlot, min_lat=-65., max_lat=-30, mult=-1.)
plt.gca().invert_yaxis()
moc['moc_GM'].data = psiPlot
../_images/examples_meridional_overturning_39_0.png 
[32]:

print('Saving netCDF files...')
moc.to_netcdf('ncfiles/'+str(casename)+'_MOC.nc')

Saving netCDF files...

Sigma-2 space (to be implemented…)

[46]:

def calc_moc_rho(vmo):
    # Sum over the zonal direction and integrate along density
    integ_layers = (
        vmo.sum("xh").cumsum("rho2_l") - vmo.sum("xh").sum("rho2_l")
    ) / rho0 / 1.0e6 + 0.1
    # The result of the integration over layers is evaluated at the interfaces
    # with psi = 0 as the bottom boundary condition for the integration
    bottom_condition = xr.zeros_like(integ_layers.isel({"rho2_l": 0}))
    # combine bottom condition with data array
    # psi_raw = xr.concat([integ_layers, bottom_condition], dim='rho2_l')
    psi_raw = xr.concat([bottom_condition, integ_layers], dim="rho2_l")
    # rename to correct dimension and add correct vertical coordinate
    psi = psi_raw.rename({"rho2_l": "rho2_i"}).transpose("rho2_i", "yq")
    psi["rho2_i"] = xr.concat([vmo.rho2_l[0]*0, vmo.rho2_l], dim="rho2_l").rename({"rho2_l": "rho2_i"})
    #psi = psi.assign_coords(rho2_i=rho2_i)
    psi.name = "psi"
    return psi.load()

[12]:

print('\n Reading dataset...')
# load data
%time ds_sigma2 = xr.open_mfdataset(OUTDIR+'/'+args.sigma2, parallel=True, \
                             combine="nested", concat_dim="time", \
                             preprocess=preprocess).chunk({"time": 12})


 Reading dataset...
CPU times: user 4.59 s, sys: 345 ms, total: 4.94 s
Wall time: 24.9 s

[13]:

ds_ann_sigma2 =  m6toolbox.weighted_temporal_mean_vars(ds_sigma2,attrs=attrs)

[14]:

print('\n Selecting data between {} and {}...'.format(args.start_date, args.end_date))
%time ds_sel_sigma2 = ds_ann_sigma2.sel(time=slice(args.start_date, args.end_date))


 Selecting data between 0031-01-01 and 0062-01-01...
CPU times: user 2.21 ms, sys: 3.55 ms, total: 5.77 ms
Wall time: 5.95 ms

[15]:

print('\n Computing time mean...')
%time ds_mean_sigma2 = ds_sel_sigma2.mean('time').compute()


 Computing time mean...
CPU times: user 9.89 s, sys: 1.9 s, total: 11.8 s
Wall time: 3min 28s

[49]:

# The following is from John Krasting's notebook. Thanks, John!
# https://github.com/jkrasting/mar/blob/main/src/gfdlnb/notebooks/ocean/Global_Meridional_Overturning.ipynb

rho0 = 1035.
ax = plt.subplot(1, 1, 1)
vmo = ds_mean_sigma2.vmo.where(ds_mean_sigma2.vmo < 1e14)
psi = calc_moc_rho(vmo)

levels = np.arange(-40, 42, 2)

cb = ax.contourf(psi.yq, psi.rho2_i, psi, levels=levels, cmap="RdBu_r")
cs = ax.contour(psi.yq, psi.rho2_i, psi, levels=levels, colors="k", linewidths=0.3)

stats = {}

y = psi.yq
r = psi.rho2_i
Y, R = np.meshgrid(y, r)
mask = (Y >= -75) & (Y <= -50) & (R >= 1036.8) & (R <= 1037.4)
_psi = psi.values
min_psi = np.min(_psi[mask])
min_index = np.unravel_index(np.argmin(_psi[mask]), _psi[mask].shape)
min_y = Y[mask][min_index]
min_r = R[mask][min_index]

stats = {
    "min_psi": round(float(min_psi),2),
    "min_lat": round(float(min_y),2),
    "min_rho": round(float(min_r),2),
}

square_size = 10
ax.plot(
    min_y,
    min_r,
    marker="s",
    color="magenta",
    markersize=square_size + 5,
    markeredgewidth=2,
    markeredgecolor="magenta",
    markerfacecolor="none",
)
ax.annotate(
    f"{min_psi:.1f} Sv",
    xy=(min_y, min_r),
    xytext=(10, 10),
    textcoords="offset points",
    arrowprops=dict(arrowstyle="->", color="black"),
    bbox=dict(
        facecolor="white",
        edgecolor="black",
        boxstyle="round,pad=0.2",
        linewidth=0.5,
        alpha=0.7,
    ),
)

ax.set_yscale("splitscale", zval=[1037.25, 1036.5, 1028.8])
plt.colorbar(cb)

date_range = 'Years '+args.start_date+' to '+ args.end_date
ax.text(.99,1.03,date_range,ha="right", transform=ax.transAxes)
ax.text(.01,1.03,f"Global",ha="left", transform=ax.transAxes)

[49]:

Text(0.01, 1.03, 'Global')
../_images/examples_meridional_overturning_47_1.png 
[50]:

# release workers
client.close(); cluster.close()