<jupyter_start><jupyter_text>Setup<jupyter_code>from dask.distributed import Client, progress

from dask_kubernetes import KubeCluster
cluster = KubeCluster(n_workers=4)
cluster.adapt(minimum=1, maximum=40)
cluster
client = Client(cluster)
client<jupyter_output><empty_output><jupyter_text>Load data<jupyter_code>import intake
import util

if util.is_ncar_host():
    print("On NCAR")
    col = intake.open_esm_datastore("../catalogs/glade-cmip6.json")
else:
    print("On GCP")
    col = intake.open_esm_datastore("../catalogs/pangeo-cmip6.json")
from collections import namedtuple

model_id = namedtuple('model_id', ['institution_id', 'source_id'])
models = [
    model_id('MIROC',  'MIROC6'),
    model_id('NCAR',  'CESM2-WACCM'),
    model_id('CCCma',  'CanESM5'),
    model_id('CAMS',  'CAMS-CSM1-0'),
    model_id('MRI',  'MRI-ESM2-0'),
    model_id('MOHC',  'UKESM1-0-LL'),
    model_id('MIROC',  'MIROC-ES2L'),
    model_id('NCAR',  'CESM2'),
    model_id('BCC', 'BCC-CSM2-MR')
]
def get_data(variable_id, model):
    
    institution_id = model.institution_id
    source_id = model.source_id
    
    cat = col.search(
        experiment_id=['piControl', 'historical', 'ssp585'],
        variable_id=variable_id,
        source_id=[source_id, ]
    )
    
    dset_dict = cat.to_dataset_dict(
        zarr_kwargs={'consolidated': True, 
                     'decode_times': True},
        cdf_kwargs={},
    )
    pi = f'CMIP.{institution_id}.{source_id}.piControl.Amon.gn'
    hist =  f'CMIP.{institution_id}.{source_id}.historical.Amon.gn'
    fut = f'ScenarioMIP.{institution_id}.{source_id}.ssp585.Amon.gn'

    pi_ds = dset_dict[pi]
    hist_ds = dset_dict[hist]
    fut_ds = dset_dict[fut]
    
    return pi_ds, hist_ds, fut_ds

def fit_hs(pi_ds, hist_ds, fut_ds, model):
    model_id = f'{model.institution_id}.{model.source_id}'
    hs = HawkinsSutton2012(
        pi_ds, hist_ds, fut_ds, variable_id, model_id
    )
    # hs.fit()
    return hs
toe_dict = {}
for m in models:
    print(m)

    pi, hist, fut = get_data('tas', m)
    hs = fit_hs(pi, hist, fut, m)
    hs.fit()
    toe = hs.time_of_emergence(2)
    toe_dict[f'{m.institution_id}.{m.source_id}'] = toe
print("done")
len(toe_dict)
_t.lon
from tools import shift_lons, shift_roll

nrows, ncols = 3, 3
size, aspect = 4., 1.8
width, height = size*aspect*ncols, size*nrows
figsize = (width, height)

fig, axs = plt.subplots(nrows, ncols, figsize=figsize,
                        subplot_kw=dict(
                            projection=ccrs.PlateCarree(),
                            aspect='auto'
                        ))

cmap_kws = dict(
    vmin=2005, vmax=2065, extend='max', levels=13,
    cmap="RdYlBu"
)

for ax, m in zip(axs.flatten(), models):
    model_name = f'{m.institution_id}.{m.source_id}'
    print(model_name)
    
    ax.set_title(f'{m.institution_id} - {m.source_id}', loc='left')
    _t = toe_dict[model_name]['TOE'].copy()
#     _t = shift_lons(_t)
#     nlon = len(_t.lon)
#     _t = _t.roll(lon=)
    cf = _t.plot.pcolormesh('lon', 'lat', 
                            ax=ax, transform=ccrs.PlateCarree(),
                            add_colorbar=False, 
                            **cmap_kws)
    ax.coastlines()
    ax.set_aspect('auto', 'box')
#     break
    
cb = plt.colorbar(cf, ax=axs, label='Year of Emergence')
# plt.tight_layout()

plt.savefig("toe.composite.png", dpi=200, transparent=True, bbox_inches='tight')<jupyter_output>MIROC.MIROC6
NCAR.CESM2-WACCM
CCCma.CanESM5
CAMS.CAMS-CSM1-0
MRI.MRI-ESM2-0
MOHC.UKESM1-0-LL
MIROC.MIROC-ES2L
NCAR.CESM2
BCC.BCC-CSM2-MR
{'fraction_non_alphanumeric': 0.14864487286951664, 'fraction_numerical': 0.01648505169041632, 'mean_word_length': 3.0497737556561084, 'pattern_counts': {'":': 0, '<': 8, '<?xml version=': 0, '>': 8, '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': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '62312', 'license': 'MIT', 'n_tokens_mistral': 1494, 'n_tokens_neox': 1441, 'n_words': 222, 'path': 'notebooks/hawkins_automate_example.ipynb', 'repo_name': 'kaedonkers/cmip6hack-toe', 'size': '3579'}