<jupyter_start><jupyter_text>Complete reconception of the dataset loading Priorities- Reduce the number of chunks. Combining members in intake-esm makes everything a lot worse.Should I completely ditch intake-esm? At this point I barely see an advantage. The methodology for now looks like this:1. Read all single datasets into a larger dict with intake-esm, but do not aggregate. `dataset` level preprocessing is applied here.2. Detrend on the `dataset` level3. Move datasets to a nested dictionary4. Combine grid_labels by using xesmf5. Parse metrics to get a full `member` dataset. Things still to fix.- [x] How is the control run making it through the detrending? Might work if it has the same time convention. But I prob have to code an exception in there...- [x] ProgressBar for the trend loading- [x] ProgressBar for the dim split saving- [x] Compare results for CanESM5 (I have some old files for that...- [x] Show non unique entries (https://stackoverflow.com/questions/9835762/how-do-i-find-the-duplicates-in-a-list-and-create-another-list-with-them)    - [x] UKESM has duplicate versions...whyyyyyy- [x] **Why does the rechunking 'setup' take so long? Raise issue?**- [x] Norwegian models dont all have lev_bounds. Investigate    - [x] Regridding needs to ignore area and thckcello! (Norwegian models)- [ ] Sort out the datasets with missing fields before- [ ] missing values alignment should also include dz_t?- [ ] All nan slope files- [ ] MPI-HR has missing detrending fields- [ ]  ddict_nested_metrics['CMIP.MOHC.UKESM1-0-LL.piControl.r1i1p1f2'] has a drift in the upper dzt cells..roi = {'x':180, 'y':130, 'lev':0}- **GFDL weird naming issue need to be fixed upstream**- [ ] reconstruct metrics if they are not there (or parse them from another run? E.g. piControl is missing the area...)- [ ] The removal of files is incredibly slow. `shutil.rmtree()` takes hours e.g. for CM4- [ ] Try to use nc files instead of zarr? Might take too long...- ~~[ ] warnings.warn('Input array is not F_CONTIGUOUS. ' 'Will affect performance.') Can I fix this somehow? its xesmf stuff./~~<jupyter_code>import dask
dask.__version__<jupyter_output><empty_output><jupyter_text>Should always have this to be able to go back with the versions<jupyter_code>!conda list
%load_ext autoreload
%autoreload 2

import numpy as np
import dask
import time
import shutil
import warnings
import xarray as xr

import matplotlib.pyplot as plt
from fastjmd95 import rho
import fastjmd95

assert fastjmd95.__version__ >= "0.2.1"
# import fastjmd95 #quick fix? https://github.com/dask/distributed/issues/3450#issuecomment-584739501
# Nope
from multiprocessing.pool import ThreadPool
from dask.diagnostics import ProgressBar
import matplotlib.pyplot as plt

from xgcm import Grid
from fastprogress.fastprogress import progress_bar

from xarrayutils.file_handling import (
    write,
    maybe_create_folder,
    file_exist_check,
    temp_write_split,
)
from xarrayutils.utils import remove_bottom_values


from aguadv_omz_busecke_2021.vertical_transformation import (
    xgcm_transform_wrapper as transform_wrapper,
)
from aguadv_omz_busecke_2021.datasets import woa13
from aguadv_omz_busecke_2021.utils import print_html
from aguadv_omz_busecke_2021.preprocessing import (
    mixed_layer_wrapper,
    save_and_reload_rechunker,
    preprocessing_wrapper,
    strip_encoding,
)
from aguadv_omz_busecke_2021.omz_tools import (
    volume_consistency_checks,
    plot_omz_results,
    omz_thickness,
)
from aguadv_omz_busecke_2021.cmip6_stash import (
    detrend,
    construct_cfdate,
    _get_calendar,
    cmip6_dataset_id,
    combined_preprocessing,
    load_single_datasets,
    nested_dataset_wrapper,
    match_and_detrend,
    combine_grid_labels,
    parse_metrics,
)

# !!! Replace with the new metric parsing...

# from euc_dynamics.file_handling import
# from euc_dynamics.cmip6_preprocessing_stash import cmip6_dataset_id

# from euc_dynamics.cmip6_preprocessing_stash import combined_preprocessing

# from euc_dynamics.o2_utils import omz_thickness

# from euc_dynamics.utils import (mask_pacific, member_treatment,load_time_averages,
#                                 force_datasets_on_same_grid, master_load, o2_models,
#                                 combine_grid_labels, pick_time)


# dask.config.set(pool=ThreadPool(4)) ### New
dask.config.set(pool=ThreadPool(32))<jupyter_output>The autoreload extension is already loaded. To reload it, use:
  %reload_ext autoreload<jupyter_text>This should in some form move to xgcm<jupyter_code>def transform_wrapper_again(
    ds_in,
    intensive_vars=[
        "thetao",
        "o2",
        "so",
        "agessc",
    ],
):
    # alternate name: transform mean wrapper
    # sigma_bins
    sigma_bins = np.array([0, 24.5, 26.5, 27.65, 100])

    # define variables to be averaged (intensive quantities)
    intensive_vars = [
        "thetao",
        "o2",
        "so",
        "agessc",
    ]  # add 'uo', 'agessc' etc?

    intensive_vars = [v for v in intensive_vars if v in ds_in.data_vars]

    for iv in intensive_vars:
        dz = (xr.ones_like(ds_in[iv]) * ds_in.dz_t).where(~np.isnan(ds_in[iv]))
        ds_in[iv] = ds_in[iv] * dz

    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        ds_out = transform_wrapper(
            ds_in,
            extensive_vars=["omz_thickness"] + intensive_vars,
            target=sigma_bins,
        )

    # reconvert the same variables
    dz = ds_out.dz_t
    for iv in intensive_vars:
        ds_out[iv] = ds_out[iv] / dz
    return ds_out<jupyter_output><empty_output><jupyter_text>Testing distributed processing. This doesnt work really well...At this point I have almost all models and its not worth it to pursue this further<jupyter_code># from dask_jobqueue import SLURMCluster
# from dask.distributed import Client
# cluster = SLURMCluster(
#     scheduler_options={"dashboard_address": 9999},
# )
# client = Client(cluster)
# client<jupyter_output><empty_output><jupyter_text>Template to fire up dask workers:```python/scratch/gpfs/jbusecke/conda_tigercpu/envs/aguadv_omz_busecke_2021/bin/python -m distributed.cli.dask_worker tcp://10.36.18.123:45961 --nthreads 2 --nprocs 20 --memory-limit 9GB --name  --nanny --death-timeout 60 --local-directory /tmp --interface op0```  Put out the observations in density data first<jupyter_code>foldername = "density_remapping_obs"
ofolder = maybe_create_folder(f"../../data/processed/{foldername}")
woa = woa13().drop('time')
woa['sigma_0'] = (fastjmd95.rho(woa.so, woa.thetao, 0) - 1000)
woa = woa.rename({'thkcello':'dz_t'})
woa = woa.chunk({'lev':-1,'x':180})
woa['omz_thickness'] = omz_thickness(woa, o2_bins=o2_bins)

woa_sigma = transform_wrapper_again(woa, intensive_vars=[va for va in woa.data_vars if va not in ['omz_thickness']])

# confirm that the volume is conserved
with ProgressBar():
    vol_sigma = (woa_sigma.omz_thickness * woa_sigma.areacello).sum(['x','y','sigma_0']).load()
    vol_z = (woa.omz_thickness * woa.areacello).sum(['x', 'y', 'lev']).load()
    
assert (((vol_z - vol_sigma)/vol_z*100)<0.5).all()

fname = ofolder.joinpath('woa13.nc')
woa_sigma = write(woa_sigma, fname, overwrite=True)
woa_sigma.isel(sigma_0=2).o2.plot(robust=True)<jupyter_output><empty_output><jupyter_text>Now the models<jupyter_code>foldername = "density_remapping_time_yearly_final"
ofolder = maybe_create_folder(f"../../data/processed/{foldername}")
tempfolder = maybe_create_folder(f"../../data/temp/scratch_temp/{foldername}")

# global parameters
o2_bins = np.array([10, 40, 60, 80, 100, 120])
import intake
from aguadv_omz_busecke_2021.utils import cmip6_collection
col = intake.open_esm_datastore(cmip6_collection(zarr=True))
col.df['source_id'].unique()
ddict = load_single_datasets(
    col,
    source_id=[
        "UKESM1-0-LL", # area not available for all members
        "MRI-ESM2-0", # 
        "NorESM2-LM",
        "NorESM2-MM",
        "GFDL-ESM4",
        "GFDL-CM4",
        "MIROC-ES2L", # done
        "CNRM-ESM2-1", # done bunch of members have missing thetao
        "CanESM5", # done
        "CanESM5-CanOE", # done
        "MPI-ESM1-2-HR", # done
        "ACCESS-ESM1-5", # done
        "IPSL-CM6A-LR", # done
        "NorESM1-F", #meh, dont need it
        "NorCPM1", #meh, dont need it
        "MPI-ESM-1-2-HAM", # done
        "MPI-ESM1-2-LR", # done
    ],
    grid_label=["gr", "gn"],
    experiment_id=["historical", "ssp585", "piControl"],
    table_id=["Omon", "Ofx"],
    variable_id=["thetao", "so", "o2", "mlotst", "agessc", "thkcello", "areacello"],
)
# Load all trends into dict

import pathlib
trendfolder = pathlib.Path('../../data/processed/linear_regression_time_zarr_multimember/')
trend_models = np.unique([ds.attrs['source_id'] for ds in ddict.values()])
flist = []
for tm in trend_models:
    flist = flist + list(trendfolder.glob(f'*{tm}*_trend.nc'))

total = len(flist)
progress = progress_bar(range(total))

trend_dict = {}
for i,path in enumerate(flist):
    key = path.stem
    ds = xr.open_mfdataset([path])
    # write the filename in the dataset
    ds.attrs.update({'filepath':str(path)})
    # exclude all nan slopes (why are these there in the first place?)
    if not np.isnan(ds.slope).all():
        trend_dict[key] = ds
    else:
        print(f"found all nan regression data for {path}")
    progress.update(i)
progress.update(total)
ddict_detrended = match_and_detrend(
    ddict,
    trend_dict,
    verbose=True,
    pass_variables=["areacello", "thkcello", "mlotst"],
)
ddict_detrended_filtered = {k:v for k,v in ddict_detrended.items() if 'rho' not in v.dims}
import warnings
with warnings.catch_warnings():
    warnings.simplefilter("ignore")
    ddict_nested = nested_dataset_wrapper(ddict_detrended_filtered)
ddict_detrended_filtered['CMIP.NOAA']<jupyter_output><empty_output><jupyter_text>List the model/variables that had to be regridded<jupyter_code>for k, ds in ddict_nested.items():
    hits = []
    if ds is not None:
        for va in ds.data_vars:
            if 'horizontally_regridded_from_grid_label' in ds[va].attrs.keys():
                hits.append(va)
        if len(hits)>0:
            print(f"regridded data for {k}:{hits}")
ddict_nested_metrics = parse_metrics(ddict_nested)
list(ddict_nested_metrics.keys())<jupyter_output><empty_output><jupyter_text>First test of the detrended data<jupyter_code># # quick detrending QC example
# roi = {'x':180, 'y':130, 'lev':30}
# plot_vars = ['thetao', 'so', 'o2']
# test_plot_dict_raw = ddict
# test_plot_dict_detrended = ddict_detrended


# for k in [kk for kk in test_plot_dict_detrended.keys() if 'historical' in kk]:
#     print(k)
#     if not any([s in k for s in plot_vars]):
#         print('k')
#         pass
#     else:
#         hist_detrended = test_plot_dict_detrended[k]
#         hist = test_plot_dict_raw[k]
        
#         var = hist.attrs['variable_id']
#         with ProgressBar():
#             hist[var].isel(**roi).resample(time='5AS').mean().plot(color='C0')
#             hist_detrended[var].isel(**roi).resample(time='5AS').mean().plot(color='C1')
            
#         try:
#             k_cc = k.replace('CMIP', 'ScenarioMIP').replace('historical', 'ssp585')
#             cc_detrended = test_plot_dict_detrended[k_cc]
#             cc = test_plot_dict_raw[k_cc]
#             cc[var].isel(**roi).resample(time='5AS').mean().plot(color='C0')
#             cc_detrended[var].isel(**roi).resample(time='5AS').mean().plot(color='C1')
#         except:
#             print(f'cc failed for {k}')
            
            
#         plt.title(k)
#         plt.show()<jupyter_output><empty_output><jupyter_text>Test the nested (and regridded) output<jupyter_code># # ds_plot = ddict_nested_metrics['CMIP.MPI-M.MPI-ESM1-2-HR.historical.r1i1p1f1']
# ds_plot = ddict_nested_metrics['CMIP.MIROC.MIROC-ES2L.historical.r3i1p1f2']
# ds_plot
# # roi = {'x':180, 'y':130, 'lev':30}

# roi = {'x':180, 'y':130, 'lev':20} # for medium res...
# # roi = {'x':300, 'y':250, 'lev':20} # for higher res...
# with ProgressBar():
#     ds_plot.isel(**roi).thetao.plot()
# with ProgressBar():
#     ds_plot.isel(**roi).agessc.plot()
# roi = {'x':180, 'y':130, 'lev':0}
# with ProgressBar():
#     ds_plot.isel(**roi).dz_t.plot()
# ds_plot.isel(**roi).dz_t<jupyter_output><empty_output><jupyter_text>Cheap copy paste of the transformation code on the quick<jupyter_code>def resample_yearly(ds_in, freq="1AS"):
    # this drops some coordinates, so i need to convert them to data_vars and then reconvert
    time_coords = [
        co
        for co in list(ds_in.coords)
        if "time" in ds_in[co].dims and co not in ["time", "time_bounds"]
    ]
    ds_out = ds_in.reset_coords(time_coords).resample(time=freq).mean()
    ds_out = ds_out.assign_coords({co: ds_out[co] for co in time_coords})
    ds_out.attrs.update({k: v for k, v in ds_in.attrs.items() if k not in ["table_id"]})
    return ds_out


def align_missing(ds_in):
    """Make sure that nans in all fields of a dataset are consistent.
    Requires"""
    # Due to the interpolation between `gr` and `gn`, we have to make sure that all data variables are masked in the same way!

    ds_mask = ds_in
    # for generalization np.logical_or.reduce((x, y, z))https://stackoverflow.com/questions/20528328/numpy-logical-or-for-more-than-two-arguments
    combo_nanmask = np.logical_or(
        np.isnan(ds_mask.o2).all("time").load(),
        np.isnan(ds_mask.thetao).all("time").load(),
    )
    try:
        plt.figure()
        combo_nanmask.isel(lev=5).plot()
        plt.show()
    except:
        pass

    ds_out = ds_in.where(~combo_nanmask)

    return ds_out
list(ddict_nested_metrics.keys())
for name, ds in ddict_nested_metrics.items():
    required_vars = ["thetao", "so", "o2", "dz_t"]#"areacello", 
    if not all(
        [va in ds.variables for va in required_vars]
    ):
        print(f"!!!\nMissing Variables {name}\n!!!")
        print([va for va in list(ds.variables) if va in required_vars])
    else:
        print(f"@@@@@@@@@@@@@@@@@\nProcessing {name}\n@@@@@@@@@@@@@@@@@")

        dataset_id = f"{cmip6_dataset_id(ds)}_{ds.attrs['variant_label']}"
        filename = ofolder.joinpath(f"{dataset_id}.zarr")

        t0 = time.time()
        if file_exist_check(filename) and not overwrite:
            print("Skipping. File exists already.")
            ds_sigma_reloaded = xr.open_zarr(
                filename, use_cftime=True, consolidated=True
            )
            try:
                plot_omz_results(ds_sigma_reloaded)
            except Exception as e:
                print(f"Plotting failed with: {e}")
        else:

            print(f"Writing to {filename}")

            tempfilelist = []

            plt.figure()
            ds.isel(time=200, lev=10).thetao.plot()
            plt.show()

            ## make sure the variables have been detrended (control run should not be)

            for va in [vv for vv in ds.data_vars if vv != "mlotst"]:
                print(va)
                if ds.attrs["experiment_id"] == "piControl":
                    assert "detrended_with_file" not in ds[va].attrs.keys()
                else:
                    assert "detrended_with_file" in ds[va].attrs.keys()

            # make sure the area has not been interpolated. I think the thickness is ok?
            for va in ["areacello"]:
                if va in ds.variables:
                    assert (
                        "horizontally_regridded_from_grid_label" not in ds[va].attrs.keys()
                    )

            if ds.attrs["experiment_id"] == "piControl":
                ds = ds.isel(time=slice(-300 * 12, None))

            ds = preprocessing_wrapper(ds)
            
            print('First temp save')
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore")

                ds, tempfilelist_var = temp_write_split(
                    ds,
                    tempfolder,
                    method='dimension',
                    verbose=True,
                    split_interval=12*10 if len(ds.x) < 700 else 12*1, # 2 years worked fine with ESM4 but for CM4 I just reduced it again..
                    file_name_pattern='raw_intermediate'
                )
                tempfilelist.extend(tempfilelist_var)

            print("Rechunking incoming dataset")
            with ProgressBar():
                temppath = tempfolder.joinpath("rechunked_intermediate.zarr")
                ds = save_and_reload_rechunker(ds.astype(np.float32), path=temppath)
                tempfilelist.append(temppath)

            # I need to align.mask the thickness aswell!
            ds = ds.reset_coords(["dz_t"])
            print("masking all fields consistently")
            with ProgressBar():
                ds = align_missing(ds)

            print(f"Remove bottom values")
            ds = remove_bottom_values(ds)

            print(f"Remove ML")
            ds = mixed_layer_wrapper(ds)

            ds = ds.set_coords("dz_t")

            ds["omz_thickness"] = omz_thickness(
                ds, o2_bins=o2_bins
            )

            ds["sigma_0"] = (fastjmd95.rho(ds.so, ds.thetao, 0) - 1000)
            
            ds_sigma_monthly = transform_wrapper_again(ds)

            perc_difference, omz_perc_difference = volume_consistency_checks(
                ds, ds_sigma_monthly
            )

            print(
                f"Relative difference ocean vol: {abs(perc_difference).data}% | OMZ vol {abs(omz_perc_difference).data}%"
            )

            if (abs(perc_difference) > 0.1).any() or (
                abs(omz_perc_difference) > 0.25# Had to increase for ESM4 before most would pass 0.01
            ).any():
                print("Volume differences exceed threshold. NOT SAVING.")
                print('\x1b[31m"red"\x1b[0m')
            else:
                ds_sigma_yearly = strip_encoding(resample_yearly(ds_sigma_monthly))
                with warnings.catch_warnings():
                    warnings.filterwarnings("ignore")

                    ds_sigma_yearly, tempfilelist_var = temp_write_split(
                        ds_sigma_yearly,
                        tempfolder,
                        verbose=True,
                        method='dimension',
                        split_interval=3,
                    )
                    tempfilelist.extend(tempfilelist_var)

                # rechunk the resulting file for convenience
                ds_sigma_yearly = strip_encoding(
                    ds_sigma_yearly.chunk({"x": -1, "y": -1, "time": 2})
                )

                #################### write out results ########################
                with ProgressBar():
                    ds_sigma_reloaded = write(
                        ds_sigma_yearly,
                        filename,
                        overwrite=overwrite,
                        force_load=False,
                        check_zarr_complete=True,
                    )
            ###### delete temps ######
            for tf in tempfilelist:
                if tf.exists():
                    shutil.rmtree(tf)
            ##################### Verification plotting ##########################
            try:
                plot_omz_results(ds_sigma_reloaded)
            except Exception as e:
                print(f"Plotting failed with: {e}")
        plt.show()
        t1 = time.time()
        print(f"Time passed: {(t1-t0)/60} minutes")<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1210801393728223, 'fraction_numerical': 0.01762656282025005, 'mean_word_length': 2.7765092879256965, 'pattern_counts': {'":': 6, '<': 42, '<?xml version=': 0, '>': 44, 'https://': 3, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 1, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '10', 'dir': 'jupyter-structured-clean-dedup', 'id': '275233', 'license': 'MIT', 'n_tokens_mistral': 6748, 'n_tokens_neox': 6356, 'n_words': 1468, 'path': 'notebooks/preprocessing/old/preprocessing_density_transform.ipynb', 'repo_name': 'jbusecke/busecke_etal_2021_aguadv', 'size': '19516'}