
%matplotlib inline
import xarray as xr
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import os, datetime, warnings, cmocean
from matplotlib import colors
from cartopy import crs as ccrs
from cartopy import feature as cfeature
from cesm_utils import fixtime, spd
print('Last updated on {}'.format(datetime.datetime.now().ctime()))

Last updated on Wed Sep 15 12:12:13 2021


warnings.filterwarnings('ignore')
plt.rcParams['figure.dpi'] = 100


from pencil_cases import *
case = 'cesm2.B1850.f09_g17.1dpop.006'
datadir = '/home/ivan/Data/Postproc/1DPOP_CESM2/'
infile = datadir + '{}.abs.max_2D.nc'.format(case)
ds_max_2D = xr.open_dataset(infile, decode_times=False, chunks='auto',
                         drop_variables=['HMXL','TEMP','UVEL','VVEL','ULAT','ULONG','z_w','dz','dzw'])
ds_max_2D = fixtime(ds_max_2D, 1851)
print('{}: {}'.format(case, case_info[case]))

stacked = ds_max_2D.SALT.stack(z=('nlat','nlon'))
maxind = stacked.argmax(axis=1)
maxpos = stacked['z'][maxind]
ntime = ds_max_2D.dims['time']
imax = [maxpos.values[t][0] for t in range(ntime)]
jmax = [maxpos.values[t][1] for t in range(ntime)]

cesm2.B1850.f09_g17.1dpop.006: 1D + G-terms + SSS restoring (tau = 365 days) + EBM off


tlat, tlon = ds_max_2D.TLAT, ds_max_2D.TLONG
tlat, tlon = np.c_[tlat,tlat[:,0]], np.c_[tlon,tlon[:,0]]

smax = ds_max_2D.SALT.max(dim='time')
data = []

fig, ax = plt.subplots(figsize=(10,10),subplot_kw={'projection':ccrs.Robinson(central_longitude=210)})
ax.set_global()
ax.add_feature(cfeature.LAND, facecolor='#b0b0b0')
ax.coastlines(linewidth=0.5)
for i, j in sorted(set([p for p in zip(imax,jmax)])):
#     print('SALT = {:.2f} g/kg : lat = {:.2f}N, lon = {:.2f}E (i={:d}, j={:d})'.format(smax.values[i,j], tlat[i,j], tlon[i,j], i, j))
    data.append([smax.values[i,j], tlat[i,j], tlon[i,j], i, j])
    _ = ax.plot(tlon[i,j], tlat[i,j], 'C3x', transform=ccrs.PlateCarree())
    _ = ax.set_title('position of salinity maxima')

pd.DataFrame(data, columns=['SALT', 'lat', 'lon', 'i', 'j'])#.style.format('{:.2f}')


# maxind2 = stacked.fillna(-99).values.argsort(axis=1)[:,-5]
# maxpos2 = stacked['z'][maxind2]
# imax2 = [maxpos2.values[t][0] for t in range(ntime)]
# jmax2 = [maxpos2.values[t][1] for t in range(ntime)]
# tups = sorted(set([p for p in zip(imax2,jmax2)]))

# data = []
# for i, j in tups:
#     data.append([smax.values[i,j], tlat[i,j], tlon[i,j], i, j])

# pd.DataFrame(data, columns=['SALT', 'lat', 'lon', 'i', 'j'])


import seaborn as sns
sns.set_theme(style='whitegrid', context='paper', palette='tab10')

ds_max = xr.open_dataset(os.path.join(datadir, '{}.abs.max.nc'.format(case)), decode_times=False)
ds_max = fixtime(ds_max, 1851)
ds_max['time'] = ds_max.time.to_index().to_datetimeindex()

data_dict = {vname: ds_max[vname].to_series() for vname in ['HMXL','TEMP','SALT','UVEL','VVEL']}
df = pd.DataFrame(data_dict)
df = df.stack().reset_index().rename(columns={'level_1':'variable', 0:'value'})

with sns.axes_style('whitegrid'):
    g = sns.relplot(x='time', y='value', data=df, kind='line', col='variable', col_wrap=3,
                    height=4, facet_kws={'sharey': False, 'sharex': True})
    _ = g.set_xticklabels(rotation=-45, ha='left')
    _ = g.set_titles(col_template='absolute max {col_name}')
    for ax, vname in zip(g.axes.flat, ['HMXL','TEMP','SALT','UVEL','VVEL']):
        ax.set_ylabel('{}'.format(ds_max[vname].units))
        ax.autoscale(axis='x', tight=True)


infile = datadir + '{}.weddell_sea.nc'.format(case)
ds_tseries = xr.open_dataset(infile, decode_times=False)
ds_tseries = fixtime(ds_tseries, 1851)

ds_tseries['FRAZIL'] = - ds_tseries.QFLUX/ds_tseries.latent_heat_fusion_mks
ds_tseries['FRAZIL'].attrs = {'long_name':'frazil ice formation', 'units':'kg/m^2/s'}

# - [Kg of salt / m^2 / s] x (1/0.0347) [ Kg of sea water / Kg of salt] x (1/1027) [m^3/kg of sea water] x 1000 [kg of freshwater / m^3]
ds_tseries['SALT_F'] = - ds_tseries.SALT_F * 1/0.0347 * 1/1027 * 1000
ds_tseries['SALT_F'].attrs = {'long_name':'fresh water flux', 'units':'kg/m^2/s'}

zz = 37
fig, ax = plt.subplots(1,2,figsize=(12,6))
pm = ax[0].pcolormesh(ds_tseries.time.to_index().to_datetimeindex(), -ds_tseries.z_t[:zz].values/100, ds_tseries.SALT[:,:zz,0,0].transpose().values,
                      cmap=cmocean.cm.haline)
#                   cmap=cmocean.cm.balance, norm=colors.TwoSlopeNorm(0))
_ = ax[0].set(title='salinity profile at Weddell Sea', ylabel='depth (m)')
ax[0].tick_params(axis='x', labelrotation = 45)
cb = fig.colorbar(pm, ax=ax[0], orientation='horizontal')
cb.set_label('g/kg')

_ = sns.lineplot(ds_tseries.time.to_index().to_datetimeindex(), ds_tseries.SALT[:,0,0,0].values, ax=ax[1])
_ = ax[1].set(xlabel='time', ylabel=ds_tseries['SALT'].units, title='surface salinity')


# varlist = ['SFWF','EVAP_F','PREC_F','ROFF_F','MELT_F']
varlist = ['SFWF','EVAP_F','PREC_F','ROFF_F','MELT_F','IOFF_F','SALT_F','FRAZIL']
data_dict = {vname: ds_tseries[vname][:,0,0].values for vname in varlist}
# for vname in ['S_INTERIOR']:
#     data_dict[vname] = ds_tseries[vname][:,0,0,0].values
df_fluxes = pd.DataFrame(data_dict, index=ds_tseries.time.to_index().to_datetimeindex())
df_fluxes.index.name = 'time'

df = df_fluxes.stack().reset_index().rename(columns={'level_1':'source', 0:'flux'})
g = sns.relplot(x='time', y='flux', data=df, kind='line', col='source', col_wrap=4,
                height=4, facet_kws={'sharey': False, 'sharex': True})
_ = g.set_xticklabels(rotation=-45, ha='left')
_ = g.set_titles(col_template='{col_name}')
for ax, vname in zip(g.axes.flat, varlist + ['S_INTERIOR','SALT']):
    ax.set_ylabel('{}'.format(ds_tseries[vname].units))
    ax.autoscale(axis='x', tight=True)


# area = (ds_tseries.TAREA[0,0] * 1.e-4).values.item()
# df_fluxes['sum'] = df_fluxes.EVAP_F + df_fluxes.PREC_F + df_fluxes.ROFF_F + df_fluxes.MELT_F + df_fluxes.IOFF_F + df_fluxes.S_INTERIOR * 1000 / area
df_fluxes['sum'] = df_fluxes.EVAP_F + df_fluxes.PREC_F + df_fluxes.ROFF_F + df_fluxes.MELT_F + df_fluxes.IOFF_F + df_fluxes.SALT_F + df_fluxes.FRAZIL
df = df_fluxes[['SFWF','sum']].stack().reset_index().rename(columns={'level_1':'source', 0:'flux'})
g = sns.relplot(x='time', y='flux', data=df, kind='line', hue='source', alpha=0.75, aspect=1.1)
_ = g.set_xticklabels(rotation=-45, ha='left')
g.ax.set_ylabel('{}'.format(ds_tseries['SFWF'].units))
g.ax.autoscale(axis='x', tight=True)


df_fluxes_monthly = df_fluxes.resample('M').mean()
df_fluxes_monthly['sum'] = (df_fluxes_monthly.EVAP_F + df_fluxes_monthly.PREC_F + df_fluxes_monthly.ROFF_F + df_fluxes_monthly.MELT_F + 
                            df_fluxes_monthly.IOFF_F + df_fluxes_monthly.SALT_F + df_fluxes_monthly.FRAZIL)
df = df_fluxes_monthly[['SFWF','sum']].stack().reset_index().rename(columns={'level_1':'source', 0:'flux'})
g = sns.relplot(x='time', y='flux', data=df, kind='line', hue='source', alpha=0.75, aspect=1.1)
_ = g.set_xticklabels(rotation=-45, ha='left')
g.ax.set_ylabel('{}'.format(ds_tseries['SFWF'].units))
g.ax.autoscale(axis='x', tight=True)

df_fluxes_monthly


df_fluxes['sum'] = df_fluxes.EVAP_F + df_fluxes.PREC_F  + df_fluxes.MELT_F + df_fluxes.IOFF_F + df_fluxes.SALT_F + df_fluxes.FRAZIL
df = df_fluxes[['SFWF','sum']].stack().reset_index().rename(columns={'level_1':'source', 0:'flux'})
g = sns.relplot(x='time', y='flux', data=df, kind='line', hue='source', alpha=0.75, aspect=1.1)
_ = g.set_xticklabels(rotation=-45, ha='left')
g.ax.set_ylabel('{}'.format(ds_tseries['SFWF'].units))
g.ax.autoscale(axis='x', tight=True)


df_fluxes_monthly = df_fluxes.resample('M').mean()
df_fluxes_monthly['sum'] = (df_fluxes_monthly.EVAP_F + df_fluxes_monthly.PREC_F + df_fluxes_monthly.MELT_F + 
                            df_fluxes_monthly.IOFF_F + df_fluxes_monthly.SALT_F + df_fluxes_monthly.FRAZIL)
df = df_fluxes_monthly[['SFWF','sum']].stack().reset_index().rename(columns={'level_1':'source', 0:'flux'})
g = sns.relplot(x='time', y='flux', data=df, kind='line', hue='source', alpha=0.75, aspect=1.1)
_ = g.set_xticklabels(rotation=-45, ha='left')
g.ax.set_ylabel('{}'.format(ds_tseries['SFWF'].units))
g.ax.autoscale(axis='x', tight=True)

df_fluxes_monthly


# def fixtime_ice(ds,yr0=1851):
#     ds2 = ds.assign(time = ds.time.values - 1)
#     ds2['time'] = ds2.time.assign_attrs(**ds.time.attrs)
#     ds2['time'].attrs['units'] = ds.time.units.replace('0001',str(yr0))
#     if 'time_bound' in ds2:
#         ds = ds.drop('time_bound')
#     return xr.decode_cf(ds2)

# infile = datadir + 'cesm2.B1850.f09_g17.1dpop.003.ice.hudson_bay.nc'
# ds_tseries_ice = xr.open_dataset(infile, decode_times=False)
# ds_tseries_ice = fixtime_ice(ds_tseries_ice, 1851)

# varlist = ['frazil','fresh','fsalt']
# data_dict = {vname: ds_tseries_ice[vname][:,0,0].values for vname in varlist}
# df_fluxes_ice = pd.DataFrame(data_dict, index=ds_tseries_ice.time.to_index().to_datetimeindex())
# df_fluxes_ice.index.name = 'time'

# df = df_fluxes_ice.stack().reset_index().rename(columns={'level_1':'source', 0:'flux'})
# g = sns.relplot(x='time', y='flux', data=df, kind='line', col='source', col_wrap=3,
#                 facet_kws={'sharey': False, 'sharex': True})
# _ = g.set_xticklabels(rotation=-45, ha='left')
# _ = g.set_titles(col_template='{col_name}')
# for ax, vname in zip(g.axes.flat, varlist):
#     ax.set_ylabel('{}'.format(ds_tseries_ice[vname].units))
#     ax.autoscale(axis='x', tight=True)


# ds_clim = xr.open_dataset(datadir + 'forcing/SALT_12_month.CMIP6-piControl.001.0400-2000.hudson_bay.nc', decode_times=False)
# ds_clim = fixtime(ds_clim,653)
# tau = 365.
# ds_monthly = ds_tseries.resample(time='M').mean()
# df_restore = pd.DataFrame({'climatology': np.tile(ds_clim.SALT[:,0,0,0].values, 10),
#                            'model': ds_tseries.SALT[:,0,0,0].resample(time='M').mean()},
#                           index=ds_monthly.time.to_index().to_datetimeindex())
# df_restore['flux'] = 1/(spd*tau) * (df_restore['climatology'] - df_restore['model'])


# fig, axs = plt.subplots(1, 2, figsize=(12,6))
# _ = df_restore[['climatology','model']].plot(ax=axs[0])
# _ = df_restore.flux.plot(ax=axs[1])
# _ = axs[0].set(ylabel='salinity (g/kg)')
# _ = axs[1].set(ylabel='restoring flux (g/kg/s)')

	SFWF	EVAP_F	PREC_F	ROFF_F	MELT_F	IOFF_F	SALT_F	FRAZIL	sum
time
1851-01-31	0.000218	-4.975455e-06	1.373055e-06	4.359353e-08	0.000242	0.000006	-2.493242e-05	0.000000	0.000219
1851-02-28	0.000132	-1.212409e-05	5.998268e-06	9.348602e-08	0.000138	0.000015	-1.432395e-05	0.000000	0.000132
1851-03-31	-0.000008	-6.450993e-06	5.024930e-06	1.248645e-07	-0.000034	0.000023	4.132162e-06	0.000000	-0.000008
1851-04-30	-0.000030	-2.259952e-06	4.558785e-08	2.133147e-08	-0.000036	0.000004	4.065603e-06	0.000000	-0.000030
1851-05-31	0.000023	-1.577282e-06	2.088254e-07	1.677518e-07	-0.000002	0.000027	4.198317e-07	0.000000	0.000023
...	...	...	...	...	...	...	...	...	...
1950-08-31	-0.000066	-6.859629e-08	2.044304e-09	6.617041e-08	-0.000089	0.000013	1.005609e-05	-0.000017	-0.000083
1950-09-30	-0.000045	-3.749783e-08	1.155022e-08	2.081835e-07	-0.000078	0.000024	9.167102e-06	-0.000019	-0.000064
1950-10-31	-0.000035	-2.363167e-08	6.746582e-09	2.066991e-07	-0.000067	0.000024	7.825873e-06	-0.000022	-0.000057
1950-11-30	0.000004	-2.975349e-08	7.959450e-09	4.775293e-08	-0.000008	0.000011	1.127071e-06	-0.000009	-0.000005
1950-12-31	0.000079	-2.457978e-07	4.826351e-07	1.068390e-07	0.000075	0.000008	-3.775672e-06	0.000000	0.000080

	SFWF	EVAP_F	PREC_F	ROFF_F	MELT_F	IOFF_F	SALT_F	FRAZIL	sum
time
1851-01-31	0.000218	-4.975455e-06	1.373055e-06	4.359353e-08	0.000242	0.000006	-2.493242e-05	0.000000	0.000219
1851-02-28	0.000132	-1.212409e-05	5.998268e-06	9.348602e-08	0.000138	0.000015	-1.432395e-05	0.000000	0.000132
1851-03-31	-0.000008	-6.450993e-06	5.024930e-06	1.248645e-07	-0.000034	0.000023	4.132162e-06	0.000000	-0.000008
1851-04-30	-0.000030	-2.259952e-06	4.558785e-08	2.133147e-08	-0.000036	0.000004	4.065603e-06	0.000000	-0.000030
1851-05-31	0.000023	-1.577282e-06	2.088254e-07	1.677518e-07	-0.000002	0.000027	4.198317e-07	0.000000	0.000023
...	...	...	...	...	...	...	...	...	...
1950-08-31	-0.000066	-6.859629e-08	2.044304e-09	6.617041e-08	-0.000089	0.000013	1.005609e-05	-0.000017	-0.000084
1950-09-30	-0.000045	-3.749783e-08	1.155022e-08	2.081835e-07	-0.000078	0.000024	9.167102e-06	-0.000019	-0.000064
1950-10-31	-0.000035	-2.363167e-08	6.746582e-09	2.066991e-07	-0.000067	0.000024	7.825873e-06	-0.000022	-0.000057
1950-11-30	0.000004	-2.975349e-08	7.959450e-09	4.775293e-08	-0.000008	0.000011	1.127071e-06	-0.000009	-0.000005
1950-12-31	0.000079	-2.457978e-07	4.826351e-07	1.068390e-07	0.000075	0.000008	-3.775672e-06	0.000000	0.000079

Find location of SALT max value¶

Find position of salinity maximum absolute values¶

Time series of absolute max values¶

Time series at Weddell Sea¶

Time series of fluxes at grid point¶

Time series of SFWF and sum of fluxes at grid point¶

Converted to monthly averages¶

Time series of SFWF and sum of fluxes at grid point (excluding ROFF_F)¶

Converted to monthly averages¶

Fluxes from the ice model¶

Estimate restoring flux at Weddell Sea¶