<jupyter_start><jupyter_text>prepare yearly forcing files based on eartH2Observe data<jupyter_code>import xarray as xr
import pandas as pd
import os, sys
import numpy as np
from os.path import join, dirname, realpath, basename
from datetime import datetime
import pdb
import glob

# multiprocessing libraries
import dask
ddir = r'/home/dirk/datasets/E2O'
name = 'metfr'
var_name="Runoff"
options = dict(scale = -86400, xdim='lon', ydim='lat', tdim='time')
fns = glob.glob(join(ddir, name, 'org', f'e2o_{name}_wrr2_glob15_day_{var_name}_*-*.nc'))
fns
import matplotlib.pyplot as plt
cmap = plt.cm.Blues
cmap.set_under('red')
@dask.delayed
def main(fn, options, name, ddir, var_name):
# for fn in fns:
    tdim, xdim, ydim = options['tdim'], options['xdim'], options['ydim']
    scale = options['scale'] # to 'mm.day-1'
    units = 'mm day-1'

    with xr.open_dataset(fn, chunks={tdim: 30, xdim:-1, ydim:-1}) as ds:
        ds = ds.rename({xdim: 'lon', ydim: 'lat', tdim:'time'})
        yrs = np.unique(ds.time.dt.year)
        for yr in yrs:
            fn_out = f'e2o_{name}_wrr2_glob15_day_{var_name}_{yr}.nc'
            print(fn_out)

            ds_yr = ds.sel(time = slice(f'01-01-{yr:04d}', f'31-12-{yr:04d}')).load()
            # fix data
            # runoff is positive
            xvar = ds_yr[var_name] * scale 

            # make figure
#             fig = plt.figure()
            xvar_seas = xvar.groupby('time.season').mean('time')
            xvar_seas.plot(x='lon', y='lat', col='season', vmin=0, vmax=10, col_wrap=2, cmap=cmap, figsize=(14,8))
            fig = plt.gcf()
            fig.suptitle(f'{name} - {yr}', y=1.01, fontsize='large')
            plt.savefig(join(ddir, name, 'fig', fn_out.replace('.nc','.png')), bbox_inches='tight', dpi = 320)
            plt.close('all')

            # remove negative values
#             xvar = xr.ufuncs.maximum(xvar, 0)
            # write output
            xvar.name = var_name
            xvar.attrs.update(ds_yr[var_name].attrs)
            xvar.attrs.update(units=units) # overwrite unit.
            ds_out = xvar.to_dataset()
            ds_out.attrs.update(ds_yr.attrs)
            encoding= {var_name: {'zlib': True}}

            ds_out.to_netcdf(join(ddir, name, fn_out), encoding=encoding)
tasks = [main(fn, options, name, ddir, var_name) for fn in fns]
dask.compute(*tasks, scheduler='processes', num_workers=3)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1522905489063145, 'fraction_numerical': 0.01981015270326042, 'mean_word_length': 2.897106109324759, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '262342', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 920, 'n_tokens_neox': 890, 'n_words': 190, 'path': 'src/1-prepare/s05-e2o_forcing.ipynb', 'repo_name': 'DirkEilander/compound_hotspots', 'size': '2423'}