<jupyter_start><jupyter_text>Calculate surface-forced watermass transformationPresently, information about the products, e.g. names of variables is hard-coded. Could improve robustness by specifying generic names.<jupyter_code>%load_ext autoreload
%autoreload 2
from dask.diagnostics import ProgressBar
import so_decadal_variability as so
import numpy as np
import xarray as xr
import os
fluxname='era5'
oceanname='en4'
density = 'gamman'
timeslice = slice('1980','2018')
dg = 0.1
density_edges = np.arange(1026,1028.5,dg)
save = False
ds = so.process.loaddata(fluxname,oceanname,timeslice)
ds, xgrid = so.grid.get_xgcm(ds)

# Hack to get around grid errors in IAP for now
# if oceanname=='iap':
#     ds['vol4d'] = ds['dx']*ds['dy']*ds['dz']*xr.ones_like(ds['vol4d'])

b_cap = True
cap = 9999
if b_cap:
    condition = (ds['b']>cap) & (~np.isnan(ds['b']))
    ds['b'] = ds['b'].where(~condition,cap)
    
G = so.wmt.calc_watermasstransformation(ds,xgrid,density_edges,density=density,b_ones=False)

if save:
    with ProgressBar():
        localdir = '../data/processed/'+fluxname+'_'+oceanname
        if ~os.path.isdir(localdir):
            !mkdir $localdir
        filename = 'G.'+density+'.'+str(dg)
        if (density=='gamman') & b_cap:
            filename = filename+'.b_cap-'+str(cap)
        G.to_netcdf(localdir+'/'+filename+'.nc')
with ProgressBar():
    G.load()
G['heat'].mean('time').plot()
G['fw'].mean('time').plot()
G['total'].mean('time').plot()
G['heat'].mean('time').plot()
G['fw'].mean('time').plot()
G['total'].mean('time').plot()
G['heat'].sel(time='2002-10-01',gamman_bin=1027.85,method='nearest')
dsnow = ds.sel(time=slice('2002-09-01','2002-10-01'))
# from xhistogram.xarray import histogram

bins=np.arange(1026,1028.5,0.1)
F = so.wmt.calc_densityflux(dsnow,xgrid)
gbFV = ds['gamman']*ds['b']*F['heat']*ds['vol4d']
# dsnow['gamman'].sel(depth=0,method='nearest').plot()
condition = (dsnow['gamman']>1027.7)&(dsnow['gamman']<=1027.8)
gbFV.where(condition,0).isel(time=0).sum(xr.ALL_DIMS).values
dsnow['b'].max().values
G['heat']['gamman_bin'].diff('gamman_bin')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.20132013201320131, 'fraction_numerical': 0.039132484677039135, 'mean_word_length': 5.549382716049383, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 7, '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': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '593569', 'license': 'MIT', 'n_tokens_mistral': 892, 'n_tokens_neox': 822, 'n_words': 131, 'path': 'notebooks/reanalyses/calc_wmt-surfacefluxes_reanalyses_oceanTS.ipynb', 'repo_name': 'gmacgilchrist/so_decadal_variability', 'size': '2121'}