<jupyter_start><jupyter_text>Detrending CMIP6 models This notebook is the main processing pipeline to calculate control run drift, which can then be easily removed from historical+scenario runs. At this point it should be a thin wrapper around the [drift removal functionality](https://cmip6-preprocessing.readthedocs.io/en/latest/drift_removal.html) of [cmip6_preprocessing](https://github.com/jbusecke/cmip6_preprocessing). Notes- [ ] THERE ARE DUPLICATE VERSIONS!- [ ] Use the netcdf files for the final version. This is sooooooo slow!- [ ] Delete temp folder- [x] Do the temp rechunking either on the full thing before selection or within the trend_match?- [x] Overwrite option for the rechunker- [x] Check if some are all nan? Delete and dont save those in the first place. Requires rerunning the notebook to reprocess<jupyter_code># %load_ext autoreload
# %autoreload 2
# starting from the raw data
import numpy as np
import xarray as xr
from cmip6_omz.utils import cmip6_collection, o2_models
from cmip6_omz.upstream_stash import rechunk_to_temp
import matplotlib.pyplot as plt
import intake
import pathlib
import warnings
import shutil
import os
from cmip6_preprocessing.preprocessing import combined_preprocessing
from cmip6_preprocessing.postprocessing import _match_datasets
from cmip6_preprocessing.drift_removal import calculate_drift
from cmip6_preprocessing.utils import cmip6_dataset_id

from xarrayutils.file_handling import write
from distributed import Client, LocalCluster
from fastprogress import progress_bar
import time
def debug_pp(ds):
    try:
        ds = combined_preprocessing(ds)
    except Exception as e:
        print(f"{cmip6_dataset_id(ds)} failed with ({e})")
    return ds
col = intake.open_esm_datastore(cmip6_collection(zarr=True)) # Takes wayyy to long for now
kwargs = {
#     "preprocess": combined_preprocessing,
    "preprocess": debug_pp,
    "cdf_kwargs": {"use_cftime": True, "chunks":{'time':3}, 'engine':'h5netcdf'},
    "zarr_kwargs": {'use_cftime': True, 'consolidated':True},
    "aggregate": False,
}
# regress_vars = ['agessc','thetao','o2', 'so', 'vo', 'epc100']
regress_vars = ['agessc','thetao','o2', 'so']
# source_id = o2_models()
source_id = ['IPSL-CM6A-LR']#, 
cat = col.search(source_id=source_id, variable_id=regress_vars)#'CanESM5',["thetao"]
ddict_historical = cat.search(experiment_id="historical").to_dataset_dict(**kwargs)
ddict_control = cat.search(experiment_id="piControl").to_dataset_dict(**kwargs)
mem_total = 160
workers = 30
threads = 2 # worked well for the rechunking
cluster = LocalCluster(
    memory_limit=f"{int(mem_total/workers)}GiB",
    dashboard_address=9999,
    threads_per_worker=threads,
    n_workers = workers,
                      )

rechunk_mem = np.ceil(mem_total/(workers* threads)/3)
print(rechunk_mem)
client = Client(cluster)
client<jupyter_output>1.0<jupyter_text>Main processing loop<jupyter_code>def clean_attrs(ds):
    ds.attrs = {k:v for k,v in ds.attrs.items() if k not in ['intake_esm_varname', ]}
    return ds

def clean_encoding(ds):
    for var in ds.variables:
        ds[var].encoding = {}
    return ds

def clean(ds):
    ds = ds.drop_vars([va for va in ds.data_vars if va != ds.attrs['variable_id']])
    return clean_attrs(clean_encoding(ds))
ofolder = pathlib.Path("../../data/external/cmip6_control_drifts/")
tempfolder = pathlib.Path("../../data/temp/cmip6_control_drifts/")<jupyter_output><empty_output><jupyter_text>I think this is much more efficient (wastes some space, but meh). I might consider tuning the chunksize<jupyter_code>ddict_control_rechunked = {}
for mi, (name, ds_control) in enumerate(ddict_control.items()):
    if not 'rho' in ds_control.dims:
        start = time.time()
        ds_id = cmip6_dataset_id(
            ds_control,
            id_attrs= [
                'activity_id',
                'institution_id',
                'source_id',
                'experiment_id',
                'variant_label',
                'table_id',
                'grid_label',
                'version',
                'variable_id', #needed until https://github.com/jbusecke/cmip6_preprocessing/pull/166 is merged 
            ]
        )
        print('$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$')
        print(f'#################{ds_id}  ({mi+1}/{len(ddict_control)})##################')
        store_target = tempfolder.joinpath(f"rechunked_{ds_id}.zarr")
        store_temp = tempfolder.joinpath("rechunker_temp.zarr")
        ds_control_rechunked = rechunk_to_temp(
            clean(ds_control).unify_chunks(),
            store_target,
            store_temp=store_temp,
            mem=f'{rechunk_mem} GB',
            overwrite=False
        )
        ddict_control_rechunked[name] = ds_control_rechunked
        stop = time.time()
        print(f"$$$$$$$$$$$ {(stop-start)/60} minutes $$$$$$$$$$")
# close the previous cluster. How?
cluster.close()
# set up dask again, with less workers (otherwise the calculation runs into trouble on tigressdata...)
mem_total = 350
workers = 2
threads = 12# more threads work a looooot better for the trend calculation (less comms)
cluster = LocalCluster(
    memory_limit=f"{int(mem_total/workers)}GB",
    dashboard_address=9999,
    threads_per_worker=threads,
    n_workers = workers,
                      )
client = Client(cluster)
client
missed_names = []
duplicate_versions = []

# This logic should make its way into cmip6_pp
trends = {}

for mi, (name, ds) in enumerate(ddict_historical.items()):
    start = time.time()
    print('$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$')
    print(f'#################{cmip6_dataset_id(ds)}  ({mi+1}/{len(ddict_historical)})##################')
    
    
    # set triggers
    skip = False
    
    nanmask = np.isnan(ds[ds.attrs['variable_id']].isel(time=slice(0,12))).all('time').load()

    # get potential matches

    rough_match = _match_datasets(
        ds,
        ddict_control_rechunked,
        match_attrs=["source_id", "grid_label", "variable_id", "table_id"],
        pop=False,
    )

    # match parent attributes
    exact_match = []
    match_attrs_parent = ["variant_label", "source_id", "experiment_id"]
    for ds_match in rough_match:
        if all(
            [
                ds.attrs["parent_" + ma] == ds_match.attrs[ma]
                for ma in match_attrs_parent
            ]
        ):
            exact_match.append(ds_match)
    
    if len(exact_match) == 0:
        print("Did not find a match for ...")
        missed_names.append(name)
    elif len(exact_match) > 1:
        print(
            "!!!!!!!!!!!!!!!!!!!!!!!!This should not happen. Found more than one matching control run."
        )
        duplicate_versions.append(name)
    else:
        
        filepath = ofolder.joinpath(f"{cmip6_dataset_id(ds)}.{'_'.join(list(ds.data_vars))}.nc")
        
        ds_parent = exact_match[0]

        if filepath.exists():
            print('File exists: Skipping rechunking')
        else:

            try:
                start 
                print('moving on')
                # catch non match before calculating anything
                _ = calculate_drift(
                    ds_parent, ds, ds.attrs["variable_id"],
                )

                # delete the encoding (I thought this was solved?) and the upsetting intake attrs 
                ds_parent = clean(ds_parent)

                trend = calculate_drift(
                    ds_parent, ds, ds.attrs["variable_id"],
                )
                # mask out nans
                trend = trend.where(~nanmask)
                
                trend_reloaded = write(
                    clean(trend),
                    filepath,
                    force_load=True,
                    serialize_load_over_dim=False,
                    overwrite=False,
                )
                trends[cmip6_dataset_id(ds)] = trend_reloaded
                
                # check attrs
                for att in ['source_id', 'variable_id', 'variant_label', 'grid_label']:
                    assert trend_reloaded.attrs[att] == ds.attrs[att]
                
            except Exception as e:
                missed_names.append(name)
                print(f"Could not write {cmip6_dataset_id(ds)} due to {e}")
                
        
        try:
            trend_reloaded = xr.open_dataset(filepath, chunks={'lev':1})
            # check attrs
            for att in ['source_id', 'variable_id', 'variant_label', 'grid_label']:
                assert trend_reloaded.attrs[att] == ds.attrs[att]
            plt.figure()
            trend_reloaded[trend_reloaded.attrs['variable_id']].isel(lev=0).plot()
            plt.show()
        except Exception as e:
            print(f"Plotting failed with {e}")
            
            
        try:
            
            # delete partially or fully empty ones (this is nested improperly...this failes if there is no file to load...s)
            nanmask = np.isnan(ds[ds.attrs['variable_id']].isel(time=slice(0,12))).all('time')
            nantrend = np.isnan(trend_reloaded[trend_reloaded.attrs['variable_id']])
            if not np.isclose(nanmask, nantrend).all():
                print('NON MATCHING MASK DETECTED. Deleting dataset')
                os.remove(filepath)
                missed_names.append(name)
                
        except Exception as e:
            print(f"Nancheck failed with {e}")
        
    stop = time.time()
    print(f'################# < {int(np.ceil((stop-start)/60))} minutes##################')<jupyter_output>$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$
#################CMIP.NOAA-GFDL.GFDL-ESM4.historical.r1i1p1f1.Omon.gn.none  (1/6)##################
File exists: Skipping rechunking<jupyter_text>Clean out temp files<jupyter_code># import shutil
# shutil.rmtree(tempfolder)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13962645128722867, 'fraction_numerical': 0.008581524482584554, 'mean_word_length': 2.7296686746987953, 'pattern_counts': {'":': 7, '<': 16, '<?xml version=': 0, '>': 16, 'https://': 3, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 1, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '382604', 'license': 'MIT', 'n_tokens_mistral': 3167, 'n_tokens_neox': 3052, 'n_words': 678, 'path': 'notebooks/processing/calculate_drift_files_old.ipynb', 'repo_name': 'jbusecke/cmip6_omz', 'size': '9905'}