Created by Ivan Lima on Wed Sep 16 2020 15:04:26 -0400
%matplotlib inline
import xarray as xr
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import os, datetime, warnings
from cesm_utils import fixtime
print('Last updated on {}'.format(datetime.datetime.now().ctime()))
Last updated on Mon Oct 12 12:18:16 2020
warnings.filterwarnings('ignore')
plt.rcParams['figure.dpi'] = 100
plt.rcParams['xtick.minor.bottom'] = False # fix strange minor ticks in time series plots
datadir = '/media/ivan/WD/Data/Postproc/1DPOP'
caselist = ['cesm.1dpop.03','cesm.1dpop.07','cesm.1dpop.08','cesm.1dpop.09']
cases, cases_upper = {}, {}
droplist = ['UVEL','VVEL','TEMP','SALT','SFWF','EVAP_F','PREC_F','ROFF_F','ULAT','ULONG']
for run in sorted(caselist):
ds = xr.open_dataset(os.path.join(datadir, '{}.surf.nc'.format(run)), decode_times=False,
drop_variables=droplist)
ds = fixtime(ds, 2051)
ds['time'] = ds.time.to_index().to_datetimeindex()
cases[run] = ds.squeeze('z_t', drop=True)
# convert lon from 0:360 to -180:180
tlon = np.where(ds.TLONG.values>180, ds.TLONG.values-360, ds.TLONG.values)
ds_3D = xr.open_dataset(os.path.join(datadir, 'cesm.3d.gterms.test.01.surf.nc'), decode_times=False,
drop_variables=droplist)
ds_3D = fixtime(ds_3D, 2051)
ds_3D['time'] = ds_3D.time.to_index().to_datetimeindex()
ds_3D = ds_3D.squeeze('z_t', drop=True)
case_info = {
'cesm.1dpop.03': 'U & V restoring',
'cesm.1dpop.07': 'G-term monthly',
'cesm.1dpop.08': 'G-term daily',
'cesm.1dpop.09': 'SSS restoring',
}
# Atlantic sector (65W - 20E)
def socean_mean_atlantic(ds, vname):
"""Compute time series of area weighted means for Atlantic sector of Southern Ocean"""
data = ds[vname].where(ds.TLAT<=-60)
weights = ds.TAREA.where(ds.TLAT<=-60)
data = data.where((tlon>=-65) & (tlon<=20))
weights = weights.where((tlon>=-65) & (tlon<=20))
return ((data * weights).sum(dim=['nlat','nlon']) / weights.sum()).to_series()
so_atlantic, so_atlantic_3D, units = {}, {}, {}
for vname in ['HMXL','HBLT','IFRAC']:
datadict = {case_info[k]:socean_mean_atlantic(cases[k], vname) for k in sorted(cases.keys())}
so_atlantic[vname] = pd.DataFrame(datadict)
so_atlantic_3D[vname] = pd.DataFrame(socean_mean_atlantic(ds_3D, vname), columns=['3-D'])
units[vname] = cases['cesm.1dpop.03'][vname].units
# Indian sector (20E - 145E)
def socean_mean_indian(ds, vname):
"""Compute time series of area weighted means for Indian sector of Southern Ocean"""
data = ds[vname].where(ds.TLAT<=-60)
data = data.where((tlon>=20) & (tlon<=145))
weights = ds.TAREA.where(ds.TLAT<=-60)
weights = weights.where((tlon>=20) & (tlon<=145))
return ((data * weights).sum(dim=['nlat','nlon']) / weights.sum()).to_series()
so_indian, so_indian_3D, units = {}, {}, {}
for vname in ['HMXL','HBLT','IFRAC']:
datadict = {case_info[k]:socean_mean_indian(cases[k], vname) for k in sorted(cases.keys())}
so_indian[vname] = pd.DataFrame(datadict)
so_indian_3D[vname] = pd.DataFrame(socean_mean_indian(ds_3D, vname), columns=['3-D'])
units[vname] = cases['cesm.1dpop.03'][vname].units
# Pacific sector (145E - 295E)
def socean_mean_pacific(ds, vname):
"""Compute time series of area weighted means for Pacific sector of Southern Ocean"""
data = ds[vname].where(ds.TLAT<=-60)
data = data.where((ds.TLONG.values>=145) & (ds.TLONG.values<=295))
weights = ds.TAREA.where(ds.TLAT<=-60)
weights = weights.where((ds.TLONG.values>=145) & (ds.TLONG.values<=295))
return ((data * weights).sum(dim=['nlat','nlon']) / weights.sum()).to_series()
so_pacific, so_pacific_3D, units = {}, {}, {}
for vname in ['HMXL','HBLT','IFRAC']:
datadict = {case_info[k]:socean_mean_pacific(cases[k], vname) for k in sorted(cases.keys())}
so_pacific[vname] = pd.DataFrame(datadict)
so_pacific_3D[vname] = pd.DataFrame(socean_mean_pacific(ds_3D, vname), columns=['3-D'])
units[vname] = cases['cesm.1dpop.03'][vname].units
fig, axs = plt.subplots(2, 2, sharex=True, figsize=(11,9))
fig.delaxes(axs[-1,-1])
for ax, vname in zip(axs.ravel(), ['HMXL','HBLT','IFRAC']):
_ = so_atlantic[vname].plot(ax=ax, alpha=0.75)
_ = so_atlantic_3D[vname].plot(ax=ax, alpha=0.75)
_ = ax.set(title='Atlantic sector', ylabel='{} ({})'.format(vname, units[vname]))
fig, axs = plt.subplots(2, 2, sharex=True, figsize=(11,9))
fig.delaxes(axs[-1,-1])
for ax, vname in zip(axs.ravel(), ['HMXL','HBLT','IFRAC']):
_ = so_indian[vname].plot(ax=ax, alpha=0.75)
_ = so_indian_3D[vname].plot(ax=ax, alpha=0.75)
_ = ax.set(title='Indian sector', ylabel='{} ({})'.format(vname, units[vname]))
fig, axs = plt.subplots(2, 2, sharex=True, figsize=(11,9))
fig.delaxes(axs[-1,-1])
for ax, vname in zip(axs.ravel(), ['HMXL','HBLT','IFRAC']):
_ = so_pacific[vname].plot(ax=ax, alpha=0.75)
_ = so_pacific_3D[vname].plot(ax=ax, alpha=0.75)
_ = ax.set(title='Pacific sector', ylabel='{} ({})'.format(vname, units[vname]))
# datadir = '/media/ivan/WD/Data/Postproc/1DPOP'
# ds_surf = xr.open_dataset(os.path.join(datadir, 'cesm.1dpop.08.surf.nc'), decode_times=False,
# drop_variables=['UVEL','VVEL','TEMP','SALT','SFWF','EVAP_F','PREC_F','ROFF_F'])
# ds_surf = fixtime(ds_surf, 2051)
# ds_rmask = xr.open_dataset('/home/ivan/Python/data/BEC_REGION_MASK_gx1v6.nc')
# # ds_surf = ds_surf.where(ds_rmask.REGION_MASK==1, drop=False)
# fig, ax = plt.subplots()
# im = ax.imshow(ds_surf.HMXL[0,...], origin='lower')
# # im = ax.imshow(ds_surf.TLONG, origin='lower')
# cb = fig.colorbar(im,ax=ax)
# fig, ax = plt.subplots()
# # im = ax.imshow(tlon, origin='lower')
# im = ax.imshow(ds_surf.TLONG, origin='lower')
# cb = fig.colorbar(im,ax=ax)
# foo = socean_mean_atlantic(ds_surf, 'HMXL')
# fig, ax = plt.subplots()
# im = ax.imshow(foo[0,...], origin='lower')
# cb = fig.colorbar(im,ax=ax)