<jupyter_start><jupyter_text>First, Let's try just applying the aguadv_omz_busecke_2021 framework globally<jupyter_code>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)<1.0).all() #failed 0.5, change to 1.

fname = ofolder.joinpath('woa13.nc')
woa_sigma = write(woa_sigma, fname, overwrite=True)
woa_pacific = mask_basin(woa_sigma, region='Pacific')
woa_atlantic = mask_basin(woa_sigma, region='Atlantic')
woa_indian = mask_basin(woa_sigma, region='Indian')
#look at the preset masks
#woa_sigma.isel(sigma_0=2).o2.plot(robust = True)
#woa_pacific.isel(sigma_0=2).o2.plot(robust = True)
#woa_atlantic.isel(sigma_0=2).o2.plot(robust = True)
#woa_indian.isel(sigma_0=2).o2.plot(x = 'lon', y = 'lat', robust = True)
#want 30S (y=121) to 30N (y = 362), o2 = 80 \mu mol/kg
woa_pac = woa_pacific.isel(o2_bin = 5, y = list(range(121,362)))
omz_thick_pac = np.nanmean(woa_pac.omz_thickness, axis = 1)
woa_atl = woa_atlantic.isel(o2_bin = 5, y = list(range(121,362)))
omz_thick_atl = np.nanmean(woa_atl.omz_thickness, axis = 1)
woa_AS = woa_indian.isel(o2_bin = 5, y = list(range(121,362)), x = list(range(0, 222)))
omz_thick_AS = np.nanmean(woa_AS.omz_thickness, axis = 1)
woa_BB = woa_indian.isel(o2_bin = 5, y = list(range(121,362)), x = list(range(222,400)))
omz_thick_BB = np.nanmean(woa_BB.omz_thickness, axis = 1)
basins = ['Pacific', 'Atlantic', 'Arabian Sea', 'Bay of Bengal']
labels = ['surface', 'thermocline', 'Intermed.', 'Deep']
lats = [woa_pac.lat[0], woa_atl.lat[0], woa_AS.lat[0], woa_BB.lat[0]]
omz_thick = [omz_thick_pac, omz_thick_atl, omz_thick_AS, omz_thick_BB]
fig, axes = plt.subplots(2,2,figsize = (20,20))
for j, ax in enumerate(axes.flat):
    for i in range(len(labels)):
        ax.set_title(basins[j], fontsize = 20)
        ax.plot(lats[j], omz_thick[j][:,i], label = labels[i], linewidth = 2)
        ax.set_xlabel('LATITUDE', fontsize = 20)
        ax.set_ylabel('OMZ thickness (m)', fontsize = 20)
        #ax.set_ylim(.1, 5000)
        #ax.set_yscale('log')
        ax.legend(fontsize = 15)<jupyter_output><empty_output><jupyter_text>Now, let's try applying our idea for a global framework using stratification<jupyter_code>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_glob = transform_wrapper_again(woa, intensive_vars=[va for va in woa.data_vars if va not in ['omz_thickness']], 
                                   new_sigma = True)

# confirm that the volume is conserved
with ProgressBar():
    vol_sigma = (woa_sigma_glob.omz_thickness * woa_sigma_glob.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)<1.0).all() #failed 0.5, change to 1.

fname = ofolder.joinpath('woa13_global.nc')
woa_sigma_glob = write(woa_sigma_glob, fname, overwrite=True)
#woa_sigma_global = xr.open_dataset('../../data/density_remapping_obs/woa13_global.nc', decode_times = False)
woa_sigma_glob
woa_pacific2 = mask_basin(woa_sigma_glob, region='Pacific')
woa_atlantic2 = mask_basin(woa_sigma_glob, region='Atlantic')
woa_indian2 = mask_basin(woa_sigma_glob, region='Indian')
#want 30S (y=121) to 30N (y = 362)
woa_pac = woa_pacific2.isel(o2_bin = 5, y = list(range(121,362)))
omz_thick_pac = np.nanmean(woa_pac.omz_thickness, axis = 1)
woa_atl = woa_atlantic2.isel(o2_bin = 5, y = list(range(121,362)))
omz_thick_atl = np.nanmean(woa_atl.omz_thickness, axis = 1)
woa_AS = woa_indian2.isel(o2_bin = 5, y = list(range(121,362)), x = list(range(0, 222)))
omz_thick_AS = np.nanmean(woa_AS.omz_thickness, axis = 1)
woa_BB = woa_indian2.isel(o2_bin = 5, y = list(range(121,362)), x = list(range(222,400)))
omz_thick_BB = np.nanmean(woa_BB.omz_thickness, axis = 1)
#[0, 24.5, 25.5, 26.5, 26.65, 26.7, 27.25, 27.65, 27.8, 100]
sigma_transform_inds = {
    'Pacific':['0:2', '2:4', '4:8', '8:'],
    'Atlantic': ['0:3', '3:5', '5:7', '7:'],
    'Arabian Sea': ['0:2', '2:6', '6:9', '9'],
    'Bay of Bengal': ['0', '1:6', '6:8', '8:']
}
%matplotlib inline
basins = ['Pacific', 'Atlantic', 'Arabian Sea', 'Bay of Bengal']
labels = ['surface', 'thermocline', 'Intermed.', 'Deep']
lats = [woa_pac.lat[0], woa_atl.lat[0], woa_AS.lat[0], woa_BB.lat[0]]
omz_thick = [omz_thick_pac, omz_thick_atl, omz_thick_AS, omz_thick_BB]
fig, axes = plt.subplots(2,2,figsize = (20,20))
for j, ax in enumerate(axes.flat):
    for i in range(len(labels)):
        inds = sigma_transform_inds[basins[j]][i]
        if ':' in inds:
            curve = eval(f'np.nansum(omz_thick[j][:,{inds}], axis = 1)')
        else:
            curve = eval(f'omz_thick[j][:,{inds}]')
        ax.set_title(basins[j], fontsize = 20)
        ax.plot(lats[j], curve, label = labels[i], linewidth = 2)
        ax.set_xlabel('LATITUDE', fontsize = 20)
        ax.set_ylabel('OMZ thickness (m)', fontsize = 20)
        #ax.set_ylim(.1, 5000)
        #ax.set_yscale('log')
        ax.legend(fontsize = 15)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.18257188216616688, 'fraction_numerical': 0.05309578408890457, 'mean_word_length': 4.981012658227848, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '627463', 'license': 'MIT', 'n_tokens_mistral': 2735, 'n_tokens_neox': 2432, 'n_words': 435, 'path': 'notebooks/sam_tests/global_density_concept_observations.ipynb', 'repo_name': 'jbusecke/cmip6_omz', 'size': '5669'}