Adding a derived variable

import os
import numpy as np
import xarray as xr
import matplotlib.pyplot as plt
import nc_time_axis  # registers the cftime axis converter for matplotlib
import x4c

x4c.set_style('journal')

# after set_style: it resets rcParams from matplotlibrc defaults, which includes the
# backend, so assert the inline backend last or figures are never captured
%matplotlib inline

# The tutorial runs against a reduced copy of a real CESM case. It is published as a
# GitHub Release asset rather than committed, so the first call downloads it into
# ~/.cache/x4c (override with $X4C_CACHE_DIR) and later calls reuse it. Set
# $X4C_SAMPLE_DIR to point at a copy you already have.
case_dir = x4c.fetch_sample_data(case='cesm1', verbose=False)
casename = os.path.basename(case_dir)
print('x4c', x4c.__version__)

from x4c import diags
from x4c.diags import F

case = x4c.Timeseries(case_dir, grid_dict={'atm': 'ne16np4', 'ocn': 'g16'},
                      cesm_ver=1)

import inspect
print(inspect.getsource(diags.DiagCalc.get_PRECT))

from x4c import utils


@F
def get_PRECT_MM(case, **kws):
    """Total precipitation rate in mm/day."""
    case.load('PRECC', **kws)
    case.load('PRECL', **kws)

    precip = case.ds['PRECC'].x.da + case.ds['PRECL'].x.da   # m/s
    da = precip * 1000 * 86400                               # -> mm/day

    # Adding two DataArrays drops the DataArray-valued attrs, `gw` among them, even
    # with keep_attrs=True: a binary op keeps only attrs that are *identical* in both
    # operands, and `gw` compares as conflicting. Without this the next `.x.gm` would
    # fail with KeyError: 'gw'.
    utils.copy_grid_attrs(da, case.ds['PRECC'].x.da)

    da.name = 'PRECT_MM'
    da.attrs['long_name'] = 'Total precipitation rate'
    da.attrs['units'] = 'mm/day'
    return da


print('registered:', 'PRECT_MM' in diags.Registry.funcs)

case.load('PRECT_MM', timespan=(1, 10))
da = case.ds['PRECT_MM']
print(dict(da.sizes), '|', da.attrs.get('units'))

gm = case.calc('PRECT_MM:ann:gm', timespan=(1, 10))
print('global-mean precipitation (mm/day):')
print(np.round(gm.values, 4))

ssh_path = os.path.join(case_dir, 'ocn', 'proc', 'tseries', 'month_1',
                        f'{casename}.pop.h.SSH.000101-000512.nc')
ssv = x4c.open_dataset(ssh_path, comp='ocn', grid='g16', vn='SSH',
                       shift_time=True).x.regrid().x.da.mean('time')

field = case.calc('PRECT_MM|regrid(2,2):ann', timespan=(1, 10))
fig, ax = field.mean('time').x.plot(
    ssv=ssv, cmap='BrBG', levels=np.linspace(0, 12, 13),
    title='Annual-mean precipitation')
x4c.showfig(fig)

@F
def get_TS_ANOM(case, **kws):
    """Surface-temperature anomaly about the record mean."""
    case.load('TS', **kws)
    ts = case.ds['TS'].x.da
    da = ts - ts.mean('time')
    da.name = 'TS_ANOM'
    da.attrs['long_name'] = 'Surface temperature anomaly'
    da.attrs['units'] = 'degC'    # NOT 'K': this is a difference, not a temperature
    return da


anom = case.calc('TS_ANOM:ann:gm', timespan=(1, 10))
print('anomaly (degC):', np.round(anom.values, 4))