Salinity in the Baltic Sea

Created by Ivan Lima on Tue Sep 1 2020 09:54:24 -0400

%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 cesm_utils import fixtime, da2ma, spd
from mpl_utils import center_cmap
warnings.filterwarnings('ignore')
plt.rcParams['figure.dpi'] = 100
print('Last updated on {}'.format(datetime.datetime.now().ctime()))

Last updated on Tue Sep  1 10:43:29 2020

from cartopy import crs as ccrs
from cartopy import feature as cfeature

datadir = '/media/ivan/WD/Data/Postproc/1DPOP'
droplist = ['HMXL','TEMP','UVEL','VVEL','ULAT','ULONG','z_w','dz','dzw','SFWF','EVAP_F','PREC_F','ROFF_F','IFRAC']

ds_rest = xr.open_dataset(os.path.join(datadir, 'cesm.1dpop.03.0005.monthly.nc'),decode_times=False,
                           drop_variables=droplist)
ds_rest = fixtime(ds_rest, 2051)
ds_gterm = xr.open_dataset(os.path.join(datadir, 'cesm.1dpop.07.0005.monthly.nc'),decode_times=False,
                           drop_variables=droplist)
ds_gterm = fixtime(ds_gterm, 2051)
ds_3D = xr.open_dataset(os.path.join(datadir, 'cesm.3d.gterms.test.01.0001.monthly.nc'),decode_times=False,
                           drop_variables=droplist)
ds_3D = fixtime(ds_3D, 2051)

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

minlon, maxlon = 5, 30
minlat, maxlat =  52, 66
lonc = (minlon + maxlon)/2.
latc = (minlat + maxlat)/2.
proj = ccrs.EquidistantConic(central_longitude=lonc,central_latitude=latc)

def surf_maps(ds, vname, cmap=plt.cm.viridis, log=False):
    fig, axs = plt.subplots(nrows=3, ncols=4, sharex=True, sharey=True, subplot_kw={'projection':proj},
                            figsize=(9.5,7.75))
    fig.subplots_adjust(top=0.95,left=0.05,right=0.95,hspace=0.175,wspace=0.025)
    axs = axs.ravel()
    for ax in axs:
        ax.set_extent([minlon,maxlon,minlat,maxlat])
        ax.coastlines(linewidth=0.5)
        gl = ax.gridlines(xlocs=np.arange(-180,180,10),ylocs=np.arange(0,90,5),draw_labels=True)
        gl.xlabels_top, gl.xlabels_bottom, gl.ylabels_right, gl.ylabels_left = False, False, False, False
    
    cax = fig.add_axes([0.1,0.07,0.8,0.025])
    if 'z_t' in ds[vname].dims:
        data = ds[vname].isel(z_t=0).where(ds.REGION_MASK<0) # mask non-marginal seas
    else:
        data = ds[vname].where(ds.REGION_MASK<0) # mask non-marginal seas

    if log:
        norm = norm=colors.LogNorm()
    else:
        norm = None
    if vname == 'SALT':
        vmin, vmax = 5, 10
    else:
        vmin, vmax = data.min().item(), data.max().item()
    for t in range(ds.dims['time']):
        pm = axs[t].pcolormesh(tlon, tlat, data[t,...], transform=ccrs.PlateCarree(),
                               vmin=vmin, vmax=vmax, norm=norm, cmap=cmap)
        if vname in ['UVEL','VVEL','EVAP_F']:
            center_cmap(pm)
            
        _ = axs[t].set_title('%d-%02d'%(ds.time[t]['time.year'],ds.time[t]['time.month']),fontsize=10)

    label = '{} ({})'.format(vname, ds[vname].units)
    cb = fig.colorbar(pm, cax=cax, orientation='horizontal')
    cb.set_label(label)
    return fig, axs, cb

fig, axs, cb = surf_maps(ds_gterm, 'SALT', cmocean.cm.haline)
_ = fig.suptitle('G-term run', x=0.5, y=1.02)

fig, axs, cb = surf_maps(ds_rest, 'SALT', cmocean.cm.haline)
_ = fig.suptitle('U & V restoring run', x=0.5, y=1.02)

fig, axs, cb = surf_maps(ds_3D, 'SALT', cmocean.cm.haline)
_ = fig.suptitle('3-D run', x=0.5, y=1.02)