Poleward Heat Transport
mom6_tools.polar_heat_transport collection of functions for computing and plotting poleward heat transport.
The goal of this notebook is the following:
server as an example on to compute polar heat transport using CESM/MOM6 output;
evaluate model experiments by comparing transports against observed and other model estimates;
[1]:
%load_ext autoreload
%autoreload 2
[2]:
import warnings
warnings.filterwarnings("ignore")
from mom6_tools.poleward_heat_transport import *
from mom6_tools.m6toolbox import cime_xmlquery, add_global_attrs, genBasinMasks
from mom6_tools.m6toolbox import weighted_temporal_mean_vars
from mom6_tools.jobqueue import get_cluster
from datetime import datetime, date
import yaml, os
import matplotlib.pyplot as plt
import matplotlib
import numpy as np
import xarray as xr
Basemap module not found. Some regional plots may not function properly
[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)
---------------------------------------------------------------------------
FileNotFoundError Traceback (most recent call last)
Cell In[4], line 2
1 caseroot = diag_config_yml['Case']['CASEROOT']
----> 2 casename = cime_xmlquery(caseroot, 'CASE')
3 DOUT_S = cime_xmlquery(caseroot, 'DOUT_S')
4 if DOUT_S:
File ~/checkouts/readthedocs.org/user_builds/mom6-tools/envs/latest/lib/python3.10/site-packages/mom6_tools/m6toolbox.py:47, in cime_xmlquery(caseroot, varname)
45 """run CIME's xmlquery for varname in the directory caseroot, return the value"""
46 try:
---> 47 value = subprocess.check_output(
48 ["./xmlquery", "-N", "--value", varname],
49 stderr=subprocess.STDOUT,
50 cwd=caseroot,
51 )
52 except subprocess.CalledProcessError:
53 value = subprocess.check_output(
54 ["./xmlquery", "--value", varname], stderr=subprocess.STDOUT, cwd=caseroot
55 )
File ~/.asdf/installs/python/3.10.20/lib/python3.10/subprocess.py:421, in check_output(timeout, *popenargs, **kwargs)
418 empty = b''
419 kwargs['input'] = empty
--> 421 return run(*popenargs, stdout=PIPE, timeout=timeout, check=True,
422 **kwargs).stdout
File ~/.asdf/installs/python/3.10.20/lib/python3.10/subprocess.py:503, in run(input, capture_output, timeout, check, *popenargs, **kwargs)
500 kwargs['stdout'] = PIPE
501 kwargs['stderr'] = PIPE
--> 503 with Popen(*popenargs, **kwargs) as process:
504 try:
505 stdout, stderr = process.communicate(input, timeout=timeout)
File ~/.asdf/installs/python/3.10.20/lib/python3.10/subprocess.py:971, in Popen.__init__(self, args, bufsize, executable, stdin, stdout, stderr, preexec_fn, close_fds, shell, cwd, env, universal_newlines, startupinfo, creationflags, restore_signals, start_new_session, pass_fds, user, group, extra_groups, encoding, errors, text, umask, pipesize)
967 if self.text_mode:
968 self.stderr = io.TextIOWrapper(self.stderr,
969 encoding=encoding, errors=errors)
--> 971 self._execute_child(args, executable, preexec_fn, close_fds,
972 pass_fds, cwd, env,
973 startupinfo, creationflags, shell,
974 p2cread, p2cwrite,
975 c2pread, c2pwrite,
976 errread, errwrite,
977 restore_signals,
978 gid, gids, uid, umask,
979 start_new_session)
980 except:
981 # Cleanup if the child failed starting.
982 for f in filter(None, (self.stdin, self.stdout, self.stderr)):
File ~/.asdf/installs/python/3.10.20/lib/python3.10/subprocess.py:1863, in Popen._execute_child(self, args, executable, preexec_fn, close_fds, pass_fds, cwd, env, startupinfo, creationflags, shell, p2cread, p2cwrite, c2pread, c2pwrite, errread, errwrite, restore_signals, gid, gids, uid, umask, start_new_session)
1861 if errno_num != 0:
1862 err_msg = os.strerror(errno_num)
-> 1863 raise child_exception_type(errno_num, err_msg, err_filename)
1864 raise child_exception_type(err_msg)
FileNotFoundError: [Errno 2] No such file or directory: '/glade/work/gmarques/cesm.cases/G/g.e30_a07c_cesm.GJRAv4.TL319_t232_wgx3_hycom1_N75.2025.130/'
[5]:
# create an empty class object
class args:
pass
args.casename = casename
# set avg dates
avg = diag_config_yml['Avg']
args.start_date = avg['start_date']
args.end_date = avg['end_date']
args.native = casename+diag_config_yml['Fnames']['native']
args.static = casename+diag_config_yml['Fnames']['static']
args.geom = casename+diag_config_yml['Fnames']['geom']
args.savefigs = False
args.nw = 6 # requesting 6 workers
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[5], line 5
2 class args:
3 pass
----> 5 args.casename = casename
6 # set avg dates
7 avg = diag_config_yml['Avg']
NameError: name 'casename' is not defined
[6]:
# read grid info
geom_file = OUTDIR+'/'+args.geom
if os.path.exists(geom_file):
grd = MOM6grid(OUTDIR+'/'+args.static, geom_file, xrformat=True)
else:
grd = MOM6grid(OUTDIR+'/'+args.static, xrformat=True)
try:
depth = grd.depth_ocean.values
except:
depth = grd.deptho.values
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[6], line 2
1 # read grid info
----> 2 geom_file = OUTDIR+'/'+args.geom
3 if os.path.exists(geom_file):
4 grd = MOM6grid(OUTDIR+'/'+args.static, geom_file, xrformat=True)
NameError: name 'OUTDIR' is not defined
[7]:
# basin masks - remove Nan's, otherwise genBasinMasks won't work
depth[np.isnan(depth)] = 0.0
basin_code = genBasinMasks(grd.geolon.values, grd.geolat.values, depth, verbose=False)
basin_code_xr = genBasinMasks(grd.geolon.values, grd.geolat.values, depth, verbose=False, xda=True)
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[7], line 2
1 # basin masks - remove Nan's, otherwise genBasinMasks won't work
----> 2 depth[np.isnan(depth)] = 0.0
3 basin_code = genBasinMasks(grd.geolon.values, grd.geolat.values, depth, verbose=False)
4 basin_code_xr = genBasinMasks(grd.geolon.values, grd.geolat.values, depth, verbose=False, xda=True)
NameError: name 'depth' is not defined
[8]:
parallel, cluster, client = get_cluster(args.nw, cluster_class='PBSCluster')
client
---------------------------------------------------------------------------
AttributeError Traceback (most recent call last)
Cell In[8], line 1
----> 1 parallel, cluster, client = get_cluster(args.nw, cluster_class='PBSCluster')
2 client
AttributeError: type object 'args' has no attribute 'nw'
[9]:
def preprocess(ds):
''' Compute montly averages and return the dataset with variables'''
variables = ['T_ady_2d', 'T_diffy_2d', 'T_hbd_diffy_2d']
for var in variables:
print('Processing {}'.format(var))
if var not in ds.variables:
print('WARNING: ds does not have variable {}. Creating dataarray with zeros'.format(var))
jm, im = grd.geolat.shape
tm = len(ds.time)
da = xr.DataArray(np.zeros((tm, jm, im)), dims=['time', 'yq','xh'], \
coords={'yq' : grd.yq, 'xh' : grd.xh, 'time' : ds.time}).rename(var)
ds = xr.merge([ds, da])
return ds[variables]
[10]:
print('\n Reading monthly native history files...')
# load data
%time ds = xr.open_mfdataset(OUTDIR+'/'+args.native, \
parallel=True, data_vars='minimal', chunks={'time': 12},\
coords='minimal', compat='override', preprocess=preprocess)
Reading monthly native history files...
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
File <timed exec>:1
NameError: name 'OUTDIR' is not defined
[11]:
print('\n Selecting data between {} and {}...'.format(args.start_date, args.end_date))
%time ds_sel = ds.sel(time=slice(args.start_date, args.end_date)).load()
---------------------------------------------------------------------------
AttributeError Traceback (most recent call last)
Cell In[11], line 1
----> 1 print('\n Selecting data between {} and {}...'.format(args.start_date, args.end_date))
2 get_ipython().run_line_magic('time', 'ds_sel = ds.sel(time=slice(args.start_date, args.end_date)).load()')
AttributeError: type object 'args' has no attribute 'start_date'
[12]:
attrs = {
'description': 'Annual mean of poleward heat transport by components ',
'start_date': args.start_date,
'end_date': args.end_date,
'reduction_method': 'annual mean weighted by days in each month',
'casename': casename
}
---------------------------------------------------------------------------
AttributeError Traceback (most recent call last)
Cell In[12], line 3
1 attrs = {
2 'description': 'Annual mean of poleward heat transport by components ',
----> 3 'start_date': args.start_date,
4 'end_date': args.end_date,
5 'reduction_method': 'annual mean weighted by days in each month',
6 'casename': casename
7 }
AttributeError: type object 'args' has no attribute 'start_date'
[13]:
ds_ann = weighted_temporal_mean_vars(ds_sel,attrs=attrs)
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[13], line 1
----> 1 ds_ann = weighted_temporal_mean_vars(ds_sel,attrs=attrs)
NameError: name 'ds_sel' is not defined
[14]:
# Heat Transport Time Series at 26.5°N (Atlantic)
ds_atl_ts = (ds*basin_code_xr.sel(region='AtlanticOcean').rename({'yh':'yq'})).sel(yq=26.5,
method='nearest').sum('xh').drop('yq')
#ds_atl_ts
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[14], line 2
1 # Heat Transport Time Series at 26.5°N (Atlantic)
----> 2 ds_atl_ts = (ds*basin_code_xr.sel(region='AtlanticOcean').rename({'yh':'yq'})).sel(yq=26.5,
3 method='nearest').sum('xh').drop('yq')
4 #ds_atl_ts
NameError: name 'ds' is not defined
[15]:
# Heat Transport Time Series at the Equator (Global)
ds_eq_ts = ds.sel(yq=0.0, method='nearest').sum('xh').drop('yq')
# Build a rename mapping
rename_dict = {var: f"{var}_26.5" for var in ds_eq_ts.data_vars}
# Apply renaming
ds_eq_ts = ds_eq_ts.rename(rename_dict)
ds_eq_ts
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[15], line 2
1 # Heat Transport Time Series at the Equator (Global)
----> 2 ds_eq_ts = ds.sel(yq=0.0, method='nearest').sum('xh').drop('yq')
3 # Build a rename mapping
4 rename_dict = {var: f"{var}_26.5" for var in ds_eq_ts.data_vars}
NameError: name 'ds' is not defined
[16]:
%matplotlib inline
ds_eq_ts.T_ady_2d[:].plot()
plt.show()
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[16], line 2
1 get_ipython().run_line_magic('matplotlib', 'inline')
----> 2 ds_eq_ts.T_ady_2d[:].plot()
3 plt.show()
NameError: name 'ds_eq_ts' is not defined
[17]:
ds_mean = ds_ann.mean('time').load()
ds_mean
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[17], line 1
----> 1 ds_mean = ds_ann.mean('time').load()
2 ds_mean
NameError: name 'ds_ann' is not defined
[18]:
varName = 'T_ady_2d'
print('Saving netCDF files...')
os.makedirs('ncfiles', exist_ok=True)
ds_mean = ds_ann.mean('time').load()
attrs = {'description': 'Time-mean poleward heat transport by components ', 'units': ds[varName].units,
'start_date': args.start_date, 'end_date': args.end_date, 'casename': casename}
add_global_attrs(ds_mean,attrs)
ds_mean.to_netcdf('ncfiles/'+casename+'_heat_transport.nc')
Saving netCDF files...
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[18], line 6
2 print('Saving netCDF files...')
4 os.makedirs('ncfiles', exist_ok=True)
----> 6 ds_mean = ds_ann.mean('time').load()
7 attrs = {'description': 'Time-mean poleward heat transport by components ', 'units': ds[varName].units,
8 'start_date': args.start_date, 'end_date': args.end_date, 'casename': casename}
9 add_global_attrs(ds_mean,attrs)
NameError: name 'ds_ann' is not defined
[ ]:
[19]:
# fix coords
basin_code_xr['xh'] = ds_sel.xh
basin_code_xr = basin_code_xr.rename({'yh':'yq'})
basin_code_xr['yq'] = ds_sel.yq
basin_code_xr.to_netcdf('ncfiles/'+casename+'_region_masks.nc')
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[19], line 2
1 # fix coords
----> 2 basin_code_xr['xh'] = ds_sel.xh
3 basin_code_xr = basin_code_xr.rename({'yh':'yq'})
4 basin_code_xr['yq'] = ds_sel.yq
NameError: name 'ds_sel' is not defined
Hovmoller plots
Global
[20]:
%matplotlib inline
f, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(10, 6))
plt.suptitle('Global poleward heat transport by components')
#T_ady_2d
(ds_ann.T_ady_2d*1.0e-15).sum(dim='xh').plot(ax=ax1,cbar_kwargs={"label": "TW"});
ax1.set_title("T_ady_2d")
ax1.set_xlabel("")
# T_diffy_2d
(ds_ann.T_diffy_2d*1.0e-15).sum(dim='xh').plot(ax=ax2,cbar_kwargs={"label": "TW"});
ax2.set_title("T_diffy_2d")
ax2.set_xlabel("")
ax2.set_ylabel("")
# T_hbd_diffy_2d
(ds_ann.T_hbd_diffy_2d*1.0e-15).sum(dim='xh').plot(ax=ax3,cbar_kwargs={"label": "TW"});
ax3.set_title("T_hbd_diffy_2d")
# T_hbd_diffy_2d
total = (ds_ann.T_hbd_diffy_2d + ds_ann.T_diffy_2d + ds_ann.T_ady_2d).rename('total')
(total*1.0e-15).sum(dim='xh').plot(ax=ax4,cbar_kwargs={"label": "TW"});
ax4.set_title("total")
ax4.set_ylabel("")
# Make it nice
plt.tight_layout()
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[20], line 7
4 plt.suptitle('Global poleward heat transport by components')
6 #T_ady_2d
----> 7 (ds_ann.T_ady_2d*1.0e-15).sum(dim='xh').plot(ax=ax1,cbar_kwargs={"label": "TW"});
8 ax1.set_title("T_ady_2d")
9 ax1.set_xlabel("")
NameError: name 'ds_ann' is not defined
Atlantic
[21]:
atl = (basin_code_xr.sel(region='MedSea') + basin_code_xr.sel(region='HudsonBay') + \
basin_code_xr.sel(region='Arctic') + basin_code_xr.sel(region='AtlanticOcean') + \
basin_code_xr.sel(region='BlackSea'))
atl['region'] = 'Altantic mask'
atl.plot();
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[21], line 1
----> 1 atl = (basin_code_xr.sel(region='MedSea') + basin_code_xr.sel(region='HudsonBay') + \
2 basin_code_xr.sel(region='Arctic') + basin_code_xr.sel(region='AtlanticOcean') + \
3 basin_code_xr.sel(region='BlackSea'))
4 atl['region'] = 'Altantic mask'
5 atl.plot();
NameError: name 'basin_code_xr' is not defined
[22]:
%matplotlib inline
f, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(10, 6))
plt.suptitle('Atlantic poleward heat transport by components')
#T_ady_2d
(ds_ann.T_ady_2d*1.0e-15*atl).sum(dim='xh').plot(ax=ax1,cbar_kwargs={"label": "TW"});
ax1.set_title("T_ady_2d")
ax1.set_xlabel("")
# T_diffy_2d
(ds_ann.T_diffy_2d*1.0e-15*atl).sum(dim='xh').plot(ax=ax2,cbar_kwargs={"label": "TW"});
ax2.set_title("T_diffy_2d")
ax2.set_xlabel("")
ax2.set_ylabel("")
# T_hbd_diffy_2d
(ds_ann.T_hbd_diffy_2d*1.0e-15*atl).sum(dim='xh').plot(ax=ax3,cbar_kwargs={"label": "TW"});
ax3.set_title("T_hbd_diffy_2d")
# T_hbd_diffy_2d
total = (ds_ann.T_hbd_diffy_2d + ds_ann.T_diffy_2d + ds_ann.T_ady_2d).rename('total')
(total*1.0e-15*atl).sum(dim='xh').plot(ax=ax4,cbar_kwargs={"label": "TW"});
ax4.set_title("total")
ax4.set_ylabel("")
# Make it nice
plt.tight_layout()
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[22], line 7
4 plt.suptitle('Atlantic poleward heat transport by components')
6 #T_ady_2d
----> 7 (ds_ann.T_ady_2d*1.0e-15*atl).sum(dim='xh').plot(ax=ax1,cbar_kwargs={"label": "TW"});
8 ax1.set_title("T_ady_2d")
9 ax1.set_xlabel("")
NameError: name 'ds_ann' is not defined
Compute temporal mean for each term
[23]:
stream = True
# create a ndarray subclass
class C(np.ndarray): pass
[24]:
# advection
varName = 'T_ady_2d'
if varName in ds_sel.variables:
tmp = np.ma.masked_invalid(ds_ann[varName].mean('time').values)
tmp = tmp[:].filled(0.)
advective = tmp.view(C)
advective.units = ds_ann[varName].units
else:
raise Exception('Could not find "T_ady_2d" in ds')
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[24], line 3
1 # advection
2 varName = 'T_ady_2d'
----> 3 if varName in ds_sel.variables:
4 tmp = np.ma.masked_invalid(ds_ann[varName].mean('time').values)
5 tmp = tmp[:].filled(0.)
NameError: name 'ds_sel' is not defined
[25]:
# neutral diffusion
varName = 'T_diffy_2d'
if varName in ds.variables:
tmp = np.ma.masked_invalid(ds_ann[varName].mean('time').values)
tmp = tmp[:].filled(0.)
diffusive = tmp.view(C)
diffusive.units = ds_ann[varName].units
else:
diffusive = None
warnings.warn('Neutrally-diffusive temperature term not found. This will result in an underestimation of the heat transport.')
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[25], line 3
1 # neutral diffusion
2 varName = 'T_diffy_2d'
----> 3 if varName in ds.variables:
4 tmp = np.ma.masked_invalid(ds_ann[varName].mean('time').values)
5 tmp = tmp[:].filled(0.)
NameError: name 'ds' is not defined
[26]:
# horizontal diffusion
varName = 'T_hbd_diffy_2d'
if varName in ds.variables:
tmp = np.ma.masked_invalid(ds_ann[varName].mean('time').values)
tmp = tmp[:].filled(0.)
hbd = tmp.view(C)
hbd.units = ds_ann[varName].units
else:
hbd = None
warnings.warn('Horizontal diffusion term not found. This will result in an underestimation of the heat transport.')
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[26], line 3
1 # horizontal diffusion
2 varName = 'T_hbd_diffy_2d'
----> 3 if varName in ds.variables:
4 tmp = np.ma.masked_invalid(ds_ann[varName].mean('time').values)
5 tmp = tmp[:].filled(0.)
NameError: name 'ds' is not defined
[27]:
# release workers
client.close(); cluster.close()
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[27], line 2
1 # release workers
----> 2 client.close(); cluster.close()
NameError: name 'client' is not defined
Plotting
[28]:
%matplotlib inline
plt_heat_transport_model_vs_obs(advective, diffusive, hbd, basin_code, grd, args)
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
Cell In[28], line 2
1 get_ipython().run_line_magic('matplotlib', 'inline')
----> 2 plt_heat_transport_model_vs_obs(advective, diffusive, hbd, basin_code, grd, args)
NameError: name 'advective' is not defined