Postprocessing: history files to timeseries

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__)

import shutil, subprocess

for tool in ['ncks', 'ncrcat']:
    path = shutil.which(tool)
    print(f'{tool:8s} {path or "NOT FOUND -- run `module load nco`"}')

try:
    from mpi4py import MPI
    print(f'mpi4py   ok (rank {MPI.COMM_WORLD.Get_rank()} of '
          f'{MPI.COMM_WORLD.Get_size()})')
except ImportError as e:
    print('mpi4py   NOT FOUND:', e)

hist_dir = os.path.join(case_dir, 'atm', 'hist')
for f in sorted(os.listdir(hist_dir))[:4]:
    mb = os.path.getsize(os.path.join(hist_dir, f)) / 1e6
    print(f'  {mb:5.1f} MB  {f}')
print(f'  ... {len(os.listdir(hist_dir))} files total')

hcase = x4c.History(case_dir, comps=['atm'], casename=casename)

print('streams found:', hcase.comps_info['atm'])
print()
vns = hcase.vns['atm']['cam.h0']
print(f'{len(vns)} time-varying variables in cam.h0:')
print(', '.join(vns))

import tempfile

work = tempfile.mkdtemp(prefix='x4c-gen_ts-')
out_dir = os.path.join(work, 'timeseries')
staging = os.path.join(work, 'staging')
print('output :', out_dir)
print('staging:', staging)

import contextlib, io, time

# gen_ts writes tqdm progress bars to stderr. They are useful interactively but add
# hundreds of carriage-return fragments to a rendered notebook, so capture them here
# and report just the summary. Drop the redirect to watch it live.
buf = io.StringIO()
t0 = time.time()
with contextlib.redirect_stderr(buf):
    hcase.gen_ts(
        output_dirpath=out_dir,
        staging_dirpath=staging,
        comps=['atm'],
        timespan=(1, 1),          # year 1 only
        timestep=1,
        timestep_unit='year',
        nproc=1,
        overwrite=True,
    )
print(f'gen_ts finished in {time.time() - t0:.0f} s')

# last line of the captured progress output, as a sanity check
tail = [ln for ln in buf.getvalue().replace(chr(13), chr(10)).splitlines() if ln.strip()]
print('last progress line:', tail[-1][:100] if tail else '(none)')

produced = []
for root, _, files in os.walk(out_dir):
    for f in sorted(files):
        if f.endswith('.nc'):
            produced.append(os.path.join(root, f))

print(f'{len(produced)} timeseries files produced\n')
print('relative to the output root:')
for p in produced[:6]:
    mb = os.path.getsize(p) / 1e6
    print(f'  {mb:5.1f} MB  {os.path.relpath(p, out_dir)}')
if len(produced) > 6:
    print(f'  ... and {len(produced) - 6} more')

ts_new = [p for p in produced if '.TS.' in os.path.basename(p)][0]
print('generated:', os.path.basename(ts_new))

new = x4c.open_dataset(ts_new, comp='atm', grid='ne16np4', vn='TS')

ref_path = os.path.join(case_dir, 'atm', 'proc', 'tseries', 'month_1',
                        f'{casename}.cam.h0.TS.000101-000512.nc')
ref = x4c.open_dataset(ref_path, comp='atm', grid='ne16np4', vn='TS')
ref12 = ref.x.da.isel(time=slice(0, 12))

print()
print('generated:', dict(new.x.da.sizes))
print('reference:', dict(ref12.sizes))
print()
diff = float(np.abs(new.x.da.values - ref12.values).max())
print(f'max absolute difference: {diff:.3e} K')
print('round-trip exact:' , diff == 0.0)

paths = hcase.get_paths('atm', 'cam.h0', timespan=('0001-01', '0001-03'))
print(f'{len(paths)} history files in 0001-01..0001-03:')
for p in paths:
    print('  ', os.path.basename(p))

print()
print('which stream holds TS?', hcase.get_hstr_based_on_vn('TS'))

found = hcase.find_timespan_files((1, 1), comps=['atm'])
print(f'{len(found)} files would be matched for year 1')

# rehearsal=True (the default) prints the list and deletes nothing
_ = hcase.rm_timespan((9000, 9000), comps=['atm'])

import shutil as _sh
_sh.rmtree(work)
print('removed the temporary output tree')