Analysis of Surface Fields

mom6_tools.MOM6grid returns an object with MOM6 grid data.

mom6_tools.latlon_analysis has a collection of tools used to perform spatial analysis (e.g., time averages and spatial mean).

The goal of this notebook is the following:

  1. server as an example of how to post-process CESM/MOM6 output;

  2. create time averages of surface fields.

[1]:
%load_ext autoreload
%autoreload 2
[2]:
import warnings
warnings.filterwarnings("ignore")
import xarray as xr
import numpy as np
import matplotlib.pyplot as plt
import os, yaml, argparse
import pandas as pd
import dask, intake
from datetime import datetime, date
from mom6_tools.jobqueue import get_cluster
from mom6_tools.m6toolbox import cime_xmlquery
from mom6_tools.m6toolbox import add_global_attrs
from mom6_tools.m6plot import xycompare, xyplot
from mom6_tools.MOM6grid import MOM6grid
from mom6_tools.surface import get_SSH, get_MLD, get_BLD
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')
---------------------------------------------------------------------------
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]:
# The following parameters must be set accordingly
######################################################

# create an empty class object
class args:
  pass

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

args.start_date = avg['start_date']
args.end_date = avg['end_date']
args.casename = casename
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.label =  diag_config_yml['Case']['SNAME']
args.mld_obs = "mld-deboyer-tx2_3v2"
args.savefigs = diag_config_yml['Misc']['savefigs']
args.nw = 6 # requesting 6 workers
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[5], line 13
     11 args.start_date = avg['start_date']
     12 args.end_date = avg['end_date']
---> 13 args.casename = casename
     14 args.native = casename+diag_config_yml['Fnames']['native']
     15 args.static = casename+diag_config_yml['Fnames']['static']

NameError: name 'casename' is not defined
[6]:
## Creating directories to place figures
os.makedirs('PNG/BLD', exist_ok=True)
os.makedirs('PNG/MLD', exist_ok=True)

## Creating directory to place netcdf files
os.makedirs('ncfiles', exist_ok=True)

print("Directories successfully created.")
Directories successfully created.
[7]:
parallel, cluster, client = get_cluster(args.nw, cluster_class='PBSCluster')
client
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[7], line 1
----> 1 parallel, cluster, client = get_cluster(args.nw, cluster_class='PBSCluster')
      2 client

AttributeError: type object 'args' has no attribute 'nw'
[8]:
client
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[8], line 1
----> 1 client

NameError: name 'client' is not defined
[9]:
# 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
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[9], 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)

NameError: name 'OUTDIR' is not defined
[10]:
print('Reading native dataset...')
startTime = datetime.now()

def preprocess(ds):
    ''' Compute montly averages and return the dataset with variables'''
    variables = ['oml','mlotst','tos','SSH', 'SSU', 'SSV', 'speed']
    if 'time_bounds' in ds.variables:
      variables.append('time_bounds')
    elif 'time_bnds' in ds.variables:
      variables.append('time_bnds')
    for v in variables:
      if v not in ds.variables:
        ds[v] = xr.zeros_like(ds.SSH)
    return ds[variables]

ds1 = xr.open_mfdataset(OUTDIR+args.native, parallel=parallel)
ds = preprocess(ds1)

print('Time elasped: ', datetime.now() - startTime)
Reading native dataset...
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[10], line 16
     13         ds[v] = xr.zeros_like(ds.SSH)
     14     return ds[variables]
---> 16 ds1 = xr.open_mfdataset(OUTDIR+args.native, parallel=parallel)
     17 ds = preprocess(ds1)
     19 print('Time elasped: ', datetime.now() - startTime)

NameError: name 'OUTDIR' is not defined
[11]:
print('Selecting data between {} and {}...'.format(args.start_date, args.end_date))
ds_sel = ds.sel(time=slice(args.start_date, args.end_date))
Selecting data between 0006-01-01 and 0021-01-01...
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[11], line 2
      1 print('Selecting data between {} and {}...'.format(args.start_date, args.end_date))
----> 2 ds_sel = ds.sel(time=slice(args.start_date, args.end_date))

NameError: name 'ds' is not defined
[12]:
# load obs-based mld from oce-catalog
try:
  catalog = intake.open_catalog(diag_config_yml['oce_cat'])
  mld_obs = catalog[args.mld_obs].to_dask()
except Exception as e:
  print("WARNING: No obs available, check config file.")
  mld_obs = None
WARNING: No obs available, check config file.
[13]:
mld_clima = ds_sel['mlotst'].groupby("time.month").mean('time').compute()
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[13], line 1
----> 1 mld_clima = ds_sel['mlotst'].groupby("time.month").mean('time').compute()

NameError: name 'ds_sel' is not defined
[14]:
mld_clima = mld_clima.assign_coords({
    "latitude": (("yh", "xh"), grd.geolat),
    "longitude": (("yh", "xh"), grd.geolon)
})
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[14], line 1
----> 1 mld_clima = mld_clima.assign_coords({
      2     "latitude": (("yh", "xh"), grd.geolat),
      3     "longitude": (("yh", "xh"), grd.geolon)
      4 })

NameError: name 'mld_clima' is not defined
[15]:
# Create the faceted plot
fig = plt.figure(figsize=(15, 10))  # Adjust figure size
plot = mld_clima.plot(
    x="longitude",
    y="latitude",
    col="month",
    col_wrap=4,  # Arrange plots in a grid with 4 columns
    cmap="viridis",  # Choose a color map
    robust=True,  # Automatically set vmin/vmax for better scaling
    cbar_kwargs={
        "orientation": "horizontal",  # Horizontal colorbar
        "pad": 0.05,  # Space between colorbar and plots
        "aspect": 40,  # Control the width of the colorbar
        "shrink": 0.8,  # Shrink the colorbar size
        "label": "MLD monthly climatology (m)"  # Customize colorbar label
    }
)
plt.suptitle('{}, from {} to {}'.format(args.label, args.start_date,
            args.end_date), fontsize=16, fontweight='bold')
# Fine-tune layout
plt.subplots_adjust(top=0.93, bottom=0.26)  # Move the plots up to create space below
plt.show()
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[15], line 3
      1 # Create the faceted plot
      2 fig = plt.figure(figsize=(15, 10))  # Adjust figure size
----> 3 plot = mld_clima.plot(
      4     x="longitude",
      5     y="latitude",
      6     col="month",
      7     col_wrap=4,  # Arrange plots in a grid with 4 columns
      8     cmap="viridis",  # Choose a color map
      9     robust=True,  # Automatically set vmin/vmax for better scaling
     10     cbar_kwargs={
     11         "orientation": "horizontal",  # Horizontal colorbar
     12         "pad": 0.05,  # Space between colorbar and plots
     13         "aspect": 40,  # Control the width of the colorbar
     14         "shrink": 0.8,  # Shrink the colorbar size
     15         "label": "MLD monthly climatology (m)"  # Customize colorbar label
     16     }
     17 )
     18 plt.suptitle('{}, from {} to {}'.format(args.label, args.start_date,
     19             args.end_date), fontsize=16, fontweight='bold')
     20 # Fine-tune layout

NameError: name 'mld_clima' is not defined
[16]:
# Add a 'month' coordinate to 'reference'
if mld_obs is not None:
    mld_obs_with_month = mld_obs.assign_coords(month=mld_clima.month)
[17]:
if mld_obs is not None:
    mld_obs_monthly = mld_obs_with_month.groupby("month").mean(dim="time")
if mld_obs is not None:
    bias = mld_clima - mld_obs_monthly.mld
[18]:
%matplotlib inline
if savefigs and (mld_obs is not None):
    # Create the faceted plot
    fig = plt.figure(figsize=(15, 10))  # Adjust figure size
    plot = bias.plot(
        x="longitude",
        y="latitude",
        col="month",
        col_wrap=4,
        cmap="bwr",
        robust=True,
        cbar_kwargs={
            "orientation": "horizontal",  # Horizontal colorbar
            "pad": 0.05,  # Space between colorbar and plots
            "aspect": 40,  # Control the width of the colorbar
            "shrink": 0.8,  # Shrink the colorbar size
            "label": "MLD monthly climatology bias [model - {}] (m)".format(args.mld_obs)  # Customize colorbar label
        }
    )
    plt.suptitle('{}, from {} to {}'.format(args.label, args.start_date,
                args.end_date), fontsize=16, fontweight='bold')
    # Fine-tune layout
    plt.subplots_adjust(top=0.93, bottom=0.26)  # Move the plots up to create space below
    plt.show()
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[18], line 2
      1 get_ipython().run_line_magic('matplotlib', 'inline')
----> 2 if savefigs and (mld_obs is not None):
      3     # Create the faceted plot
      4     fig = plt.figure(figsize=(15, 10))  # Adjust figure size
      5     plot = bias.plot(
      6         x="longitude",
      7         y="latitude",
   (...)
     18         }
     19     )

NameError: name 'savefigs' is not defined

Mixed layer depth

[19]:
%matplotlib inline
# MLD
if mld_obs is not None:
    get_MLD(ds_sel,'mlotst', mld_obs, grd, args)

Boundary layer depth

[20]:
get_BLD(ds_sel, 'oml', grd, args)
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[20], line 1
----> 1 get_BLD(ds_sel, 'oml', grd, args)

NameError: name 'ds_sel' is not defined
[21]:
# SSH (not working)
# get_SSH(ds, 'SSH', grd, args)
[22]:
if parallel:
    print('\n Releasing workers...')
    client.close(); cluster.close()
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[22], line 1
----> 1 if parallel:
      2     print('\n Releasing workers...')
      3     client.close(); cluster.close()

NameError: name 'parallel' is not defined
[ ]: