<jupyter_start><jupyter_text>Make sure to manually install [cmip6_preprocessing](https://github.com/jbusecke/cmip6_preprocessing) by doing the following in the terminal:```shellgit clone https://github.com/jbusecke/cmip6_preprocessing.git You only need to do this once!!!cd cmip6_preprocessing  These you need to do each time after restartingpython setup.py develop```<jupyter_code>%load_ext autoreload
%autoreload 2
# Still need to recreate dzt!!!!
%matplotlib inline
import xarray as xr
import intake
import util
import warnings
from cmip6_preprocessing.preprocessing import full_preprocessing, read_data, import_data
if util.is_ncar_host():
    col = intake.open_esm_datastore("../catalogs/glade-cmip6.json")
else:
    col = intake.open_esm_datastore("../catalogs/pangeo-cmip6.json")
col
col.search(variable_id=['areacello'], grid_label='gr').df.drop_duplicates(subset='source_id')
with warnings.catch_warnings(): # these lines just make sure that the warnings dont clutter your notebook
    warnings.simplefilter("ignore")
    data_dict = read_data(col,
                    experiment_id=['piControl'],
                    grid_label='gn',
                    variable_id=['thetao', 'uo', 'vo'],#,
#                     member_id = 'r1i1p1f1', # so that this runs faster for testing
                    source_id = ['CESM2']
                    )
from xgcm import Grid
ds = data_dict['CESM2']
grid = Grid(*standardized_metrics(ds))
grid.average()
from cmip6_preprocessing.preprocessing import cmip6_renaming_dict
a = cmip6_renaming_dict()
a
from xhistogram.xarray import histogram
import numpy as np
def omz_volume(o2, vol, o2_bins, dim=['x', 'y', 'lev']):
    nanmask = np.isnan(o2)
    # Should we also check the bin data for nans?
    
    omz_vol = histogram(
        o2.where(~nanmask),
        bins=[o2_bins],
        weights=vol.where(~nanmask),
        dim=dim,
    )
    return omz_vol
o2_bins = np.arange(0,300, 20)
ds_test = data_dict['CanESM5']
omz = omz_volume(ds_test.o2*1e3, (ds_test.thkcello * ds_test.areacello), o2_bins)
# ocean volume in each oxygen bin
omz.isel(time=0).plot()
# lets look at the full volume at e.g. 80 mymol/kg
omz_80 = omz.cumsum('o2_bin').sel(o2_bin=80, method='nearest')
omz_80.isel(time=0).load()
# the time evolution of the volume of < 80 mymols/kg waters
omz_80.plot()
ds_test = data_dict['CESM2']
ds_test
ds_test = data_dict['MCM-UA-1-0']
# ok lets recreate the vertical spacing
from xgcm.autogenerate import generate_grid_ds
# Generate 'full' dataset, which includes additional coordinate `depth_left` and appropriate attributes.
ds_full = generate_grid_ds(ds_test, {'Z':'lev'}, pad=)
print(ds_full)
print(ds.depth.data)
print(ds_full.depth_left.data)
# Here I am trying to generalize the metrics. I am doing two things:
# 1) replace the reconstructed metrics with the `real` ones if they are loaded.
# 2) parse all the metrics that are found in the dataset
def standardized_metrics(ds):
    metrics = {
        ('X', 'Y'):['area_t', 'area_ne', 'area_n', 'area_e'],
        ('X',):['dxt', 'dxne', 'dxn', 'dxe'],
        ('Y',):['dyt', 'dyne', 'dyn', 'dye'],
        ('Z',):['dzt']
    }
    # parse the `real` variables over the reconstructed...
    for new, old in zip(['areacello', 'thkcello'], ['area_t', 'dzt']):
        if new in ds.variables:
            ds.coords[old] = ds[new]
            ds = ds.drop(new)
        
    # Remove metrics that are not present in dataset
    metrics = {k:[m for m in met if m in ds.variables] for k,met in metrics.items()}
    return ds, metrics
from xgcm import Grid
grids = {k: Grid(*standardized_metrics(ds)) for k,ds in data_dict.items()}
!!! z axis needs a reconstruction
ds_test = data_dict['MCM-UA-1-0']
ds_test, metrics = standardized_metrics(ds_test)
grid = Grid(ds_test, metrics=metrics)
grid
def hor_streamfunction(u,grid):
    u_int = grid.integrate(u, 'Z')
    print(u_int)
    
hor_streamfunction(ds_test.uo, grid)
import gsw
def heat_capacity():
    # for now return constant value. Should be replaced with gsw routine?
    return 3850.0
rho_ref = 1035
cp = heat_capacity()
grid.integrate(ds_test.thetao * cp * rho_ref, ['X', 'Y', 'Z']).squeeze().plot()
def moc(v, grid):
    v = v.copy()
    v_int = grid.integrate(v, 'X')
    # we need to conserve the coords!
    x_dim = grid._get_dims_from_axis(v,'X')[0]
#     temp = v[{x_dim:1}]
# #     for co in temp.coords:
# #         v_int.coords[co] = temp[co]
    print(v_int)
    moc = grid.cumsum(v_int*v_int.dyne, 'Y')
    return moc
streamf = moc(ds_test.vo, grid) 
# TODO: 
# enable the full preprocessing on just tracer set (assume right shift?)
# use area if available, also dz
# reconstruct dz if not available
# are all the models bousinesque? (and is the ref density always 1035?) - for the heat capacity

# Issue about conserving coords (at least the metrics!)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1271902700474129, 'fraction_numerical': 0.014842300556586271, 'mean_word_length': 3.852, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 5, 'https://': 2, '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': '232463', 'license': 'MIT', 'n_tokens_mistral': 1736, 'n_tokens_neox': 1673, 'n_words': 470, 'path': 'notebooks/raw_notebooks/jbusecke_testing.ipynb', 'repo_name': 'pedmondlerner/cmip6hack-ocean-bgc', 'size': '4851'}