bottom_speed_timeseries

import cmocean
import cosima_cookbook.distributed as ccd
import dask.distributed as dsk
import glob
import matplotlib.gridspec as gs
import matplotlib.pyplot as plt
import numpy as np
import xarray as xr

import warnings # ignore these warnings
warnings.filterwarnings("ignore", category = FutureWarning)
warnings.filterwarnings("ignore", category = UserWarning)
warnings.filterwarnings("ignore", category = RuntimeWarning)

clnt = dsk.Client()
clnt

exps = ['', 'wind-up10', 'wind-down10', 'mw-up50', 'mw-down50', 'drag0.5', 'drag2', 'lt-ridge']
labels = {'':'CTRL', 'wind-up10':'WIND +10%', 'wind-down10':'WIND -10%',
          'mw-up50':'MW +50%', 'mw-down50':'MW -50%', 'drag0.5':'0.5x CDRAG',
          'drag2':'2x CDRAG', 'lt-ridge':'LT-RID'}

def preprocess(ds):
    ds = ds.sel(xq = slice(-70, 80), yh = slice(None, -50),
                xh = slice(-70, 80), yq = slice(None, -50))
    return ds

bottomsp = {}
for exp in exps[:1]:
    if exp == '':
        files = np.sort(glob.glob('/home/561/jn8053/payu/mom6/mom6-panan/archive/output**/*ocean_month_z.nc'))
    else:
        files = np.sort(glob.glob('/home/561/jn8053/payu/mom6/mom6-panan-'+exp+'/archive/output**/*ocean_month_z.nc'))
    data = xr.open_mfdataset(files, parallel = True, preprocess = preprocess).groupby('time.year').mean('time')
    
    depth_array = data['uo'] * 0 + data['z_l']
    max_depth = depth_array.max(dim = 'z_l', skipna = True)
    u = data['uo'].where(depth_array['z_l'] >= max_depth).sum(dim = 'z_l').load()
    v = data['vo'].rename({'xh':'xq','yq':'yh'}).interp(xq = data['uo']['xq'], yh = data['uo']['yh'])
    v = v.where(depth_array['z_l'] >= max_depth).sum(dim = 'z_l').load()
    print(exp)
    
    bottomsp[exp] = np.sqrt(u**2 + v**2)
    bottomsp[exp] = bottomsp[exp].mean(['xq','yh'])