FESOM2.1-REcoM3 model analysis

import os

import numpy as np
import xarray as xr

#from scipy.interpolate import griddata
#from scipy.stats import linregress

import matplotlib.pyplot as plt
#import matplotlib.colors as colors
from matplotlib import ticker

import cartopy
import cartopy.crs as ccrs
import cartopy.feature as cfeature

# pyfesom2 package as published in pip or manual installation (recommended)
if False:
    # conda/pip
    import pyfesom2 as pf

else:
    # manually 
    # from https://github.com/FESOM/pyfesom2
    # import path to module
    import sys
    
    #sys.path.append('../../../py_f2recom/modules/') # in py_f2recom
    sys.path.append('../../../pyfesom2/') # official pyfesom2
    import pyfesom2 as pf

## paths to data
resultpath = '/Users/mozeisin/data/Zeising_2026_GMD/'
meshpath = '/Users/mozeisin/data/Zeising_2026_GMD/mesh/'
savepath = '/Users/mozeisin/data/Zeising_2026_GMD/results/'

## script behavior
silent = False
savefig = False
dpi = 90

## load FESOM2 mesh
mesh = pf.load_mesh(meshpath)#, usepickle=False, usejoblib=False)

# for debugging, try
#mesh = pf.load_mesh(meshpath, usepickle=False, usejoblib=False)

#meshdiag = pf.get_meshdiag(mesh, meshpath)
meshdiag = xr.open_dataset(meshpath + 'fesom.mesh.diag.nc')

# open ds_3D_Monthly_1990_2019 from Zenodo
TEP_3D_Monthly_1990_2019 = xr.open_dataset(
    resultpath + 'ds_3D_Monthly_1990_2019.nc')
TEP_3D_Monthly_1990_2019

# pyfesom2 own Arctic masks
mask_AO = pf.get_mask(mesh, "Arctic_Basin") 
mask_Global = pf.get_mask(mesh, "Global Ocean")

# subset dataset to 2000-2005
TEP_3D_Monthly_2000_2005 = TEP_3D_Monthly_1990_2019.sel(time = slice('2000-01','2005-12'))

# subset TEP to Arctic Ocean 
TEP_AO = TEP_3D_Monthly_2000_2005.TEP[:,mask_AO]

# select surface layer
TEP_AO_Surface = TEP_AO.sel(nz1=2.5)

# average over Arctic Ocean
TEP_AO_Surface_mean = TEP_AO_Surface.mean(dim='nod2')

if not silent:
    print(f'Shape of TEP Arctic Ocean surface mean array: {np.shape(TEP_AO_Surface_mean)}')

# TEP unit conversion
# [mmol C/m3] -> [mug /L] 

print(TEP_3D_Monthly_1990_2019.TEP.units)

# conversion factor
conv = 12.01 # [mmol C/m3] -> [mug /L] 
TEP_label = 'TEP [$\mu$g L$^{-1}$]'

# simple figure of timeseries
fig, ax = plt.subplots()
ax.plot(TEP_AO_Surface_mean.time, TEP_AO_Surface_mean * conv)
ax.set_ylabel(TEP_label)
ax.set_xlabel('Year')
ax.set_title('Surface mean of Arctic Ocean')
ax.grid()

if savefig:
    plt.savefig(savepath + 'TEP_timeseries_Arctic_Ocean_surface_2000_2005.png',
                dpi=dpi, bbox_inches='tight')

plt.show(block=False)                

# calculate climatology
# subset to upper 30 m
# group by months and calculate mean
TEP_AO_clim = TEP_AO.sel(nz1=slice(0, 30)).groupby("time.month").mean("time")

if not silent:
    display(TEP_AO_clim)

# calculate mean TEP concentration across Arctic Ocean region
TEP_AO_clim_mean = TEP_AO_clim.mean(dim='nod2').compute()

# calculate standard deviation
TEP_AO_clim_std = TEP_AO_clim.std(dim='nod2').compute()

# name variables within Datsets
TEP_AO_clim_mean.name = 'TEP_mean'
TEP_AO_clim_std.name = 'TEP_std'

# plotting climatology

# Paul Tol's bright colors
# blue, cyan, green, yellow, red, pruple, grey
CB_Tol_bright = ['#4477AA', '#66CCEE', '#228833',
                 '#CCBB44', '#EE6677', '#AA3377', '#BBBBBB']

months = TEP_AO_clim.month

lwd = 2
alpha = 0.5
fontsize = 13

ylim = (400)

# figure
fig, ax = plt.subplots()

# surface layer, nz1 = 2.5 m
for i, nz1 in enumerate([2.5, 7.5, 15]):
    m, = ax.plot(months, TEP_AO_clim_mean.sel(nz1=nz1) * conv, '-',
                 color=CB_Tol_bright[i], label=f'{nz1:2.1f} m', linewidth=lwd)
    ax.fill_between(months,
                    (TEP_AO_clim_mean.sel(nz1=nz1) +
                     TEP_AO_clim_std.sel(nz1=nz1)) * conv,
                    (TEP_AO_clim_mean.sel(nz1=nz1) -
                     TEP_AO_clim_std.sel(nz1=nz1)) * conv,
                    color=CB_Tol_bright[i], alpha=alpha)

# ax outline
ax.axhline(y=0, linewidth=.70, linestyle='-', color='k')
ax.set_xlabel('Month', fontsize=fontsize)
ax.set_ylabel(TEP_label, fontsize=fontsize)
ax.set_ylim(0, ylim)
ax.grid()
ax.set_yticks([y for y in range(0, ylim, 25)], minor=True)
ax.yaxis.grid(True, which='minor', alpha=0.3)
ax.set_xticks([x for x in months])  # , minor=True)
ax.legend()
ax.set_title('Climatology for the Arctic Ocean, 2000-2005, layered')

if savefig:
    plt.savefig(savepath + 'TEP_climatology_Arctic_Ocean_layered_2000_2005.png',
                dpi=dpi, bbox_inches='tight')

plt.show(block=False)

# same plot for climatology
# now figure with errorbars


# Paul Tol bright colors
# blue, cyan, green, yellow, red, pruple, grey
CB_Tol_bright = ['#4477AA', '#66CCEE', '#228833',
                 '#CCBB44', '#EE6677', '#AA3377', '#BBBBBB']

months = TEP_AO_clim.month

lwd = 2
alpha = 0.5
fontsize = 13

ylim = (400)

# figure
fig, ax = plt.subplots()

# surface layer, nz1 = 2.5 m
for i, nz1 in enumerate([2.5, 7.5, 15]):
    m, = ax.plot(months, TEP_AO_clim_mean.sel(nz1=nz1) * conv, '-',
                 color=CB_Tol_bright[i], label=f'{nz1:2.1f} m', linewidth=lwd)
    ax.errorbar(months, TEP_AO_clim_mean.sel(nz1=nz1) * conv,
                yerr=TEP_AO_clim_std.sel(nz1=nz1) * conv,
                capsize=4, linestyle='', color=CB_Tol_bright[i],
                )

# ax outline
ax.axhline(y=0, linewidth=.70, linestyle='-', color='k')
ax.set_xlabel('Month', fontsize=fontsize)
ax.set_ylabel(TEP_label, fontsize=fontsize)
ax.set_ylim(0, ylim)
ax.grid()
ax.set_yticks([y for y in range(0, ylim, 25)], minor=True)
ax.yaxis.grid(True, which='minor', alpha=0.3)
ax.set_xticks([x for x in months])  # , minor=True)
ax.legend()
ax.set_title('Climatology for the Arctic Ocean, 2000-2005, layered')

plt.show(block=False)

# subset TEP to surface data for 2000-2005
TEP_Surface = TEP_3D_Monthly_1990_2019.TEP.sel(
    time=slice('2000-01', '2005-12')).sel(nz1=2.5)

# subset to contain only May
TEP_Surface_May = TEP_Surface.sel(time=TEP_Surface.time.dt.month.isin([5]))

# linear regression per grid cell
from scipy.stats import linregress

res_May = TEP_Surface_May.polyfit(dim='time', deg=1, skipna=True)

# specification for interpolation to regular
box = [-180,180,-80,90]
left, right, down, up = box

# resolution
res_factor = 2
lonNumber, latNumber = [360 * res_factor, 180 * res_factor] 

# create meshgrid
lonreg = np.linspace(left, right, lonNumber)
latreg = np.linspace(down, up, latNumber)
lons, lats = np.meshgrid(lonreg, latreg)

# create land-ocean mask based on natural Earth
mask_NaturalEarth = pf.mask_ne(lons, lats)

# control on number of grid cells
mesh.n2d == len(res_May.polyfit_coefficients[0,:])

# interpolation based on pyfesom2 
reg = pf.fesom2regular(res_May.polyfit_coefficients[0,:].values, mesh, lons, lats)

# land masking
aux = np.ma.masked_where(mask_NaturalEarth, reg)
res_May_slope = np.ma.masked_equal(aux, 0)

# cyclic points for plotting
res_May_slope_cyc, lon_cyc = pf.add_cyclic_point(res_May_slope, coord=lonreg)

# new meshgrid for plotting
lons_cyc, lats_cyc = np.meshgrid(lon_cyc, latreg)

def mygrid(m, box=[-180, 180, 65, 90], draw_labels=True,
           gridcolor='lightgray', fontsize=14, labelcolor='black'):
    '''
    Add natural earth fetures, add grid to panel, and set extent to specified lat/lon box.
    '''

    # natural earth feature
    m.add_feature(cfeature.LAND, zorder=1,
                  edgecolor='black', facecolor='darkgray')
    # m.add_feature(cfeature.COASTLINE.with_scale('110m'), linewidth=0.1, color='black')

    # box
    m.set_extent(box, crs=ccrs.PlateCarree())

    # grid
    gl = m.gridlines(draw_labels=draw_labels, dms=True, y_inline=True, alpha=0.8, color=gridcolor,
                     xlabel_style=dict(rotation=0, fontsize=fontsize),
                     ylabel_style=dict(color=labelcolor, fontsize=fontsize))

    if draw_labels:
        # gl.xformatter = ticker.FuncFormatter(lambda x, _: f"{abs(x):.0f}° {'E' if x >= 0 else 'W'}")
        gl.yformatter = ticker.FuncFormatter(
            lambda y, _: f"{abs(y):.0f}° {'N' if y >= 0 else 'S'}")

# plot TEP with specific colorbar
# for months of main blooming, i.e., here only for May
# other months can easily be added as panels and included in subplot_mosaic

# map outline
mapproj = ccrs.NorthPolarStereo()
box = [-180, 180, 68, 90]

# colormap
cmap = 'RdBu_r'

# norm
vmin = -10e-17
vmax = -vmin

# levels = np.arange(-0.02,0.0205,0.0025)
# ticks = np.arange(-0.2,0.21,0.125)
# norm = colors.BoundaryNorm(boundaries=levels, ncolors=256)
norm = None

# formatting
sfmt = ticker.ScalarFormatter(useMathText=True)
sfmt.set_powerlimits((-3, 4))

# figure
fig = plt.figure(figsize=(8, 6))
axes = fig.subplot_mosaic(
    """
            A
            """,
    gridspec_kw={'hspace': 0.02, 'wspace': 0.05},
    subplot_kw=dict(projection=mapproj))

# -----------------------------------------------------
# May
ax1 = axes['A']
m1 = ax1.pcolormesh(lons_cyc, lats_cyc, res_May_slope_cyc,
                    norm=norm,
                    vmin=vmin, vmax=vmax,
                    cmap=cmap,
                    transform=ccrs.PlateCarree())
mygrid(ax1)
ax1.set_title('May', fontsize=20)

# -----------------------------------------------------
# -----------------------------------------------------
# colorbar
cbar = fig.colorbar(m1, ax=[ax1],
                    orientation='horizontal',
                    # ticks = ticks,
                    # format = sfmt,
                    extend='both',
                    fraction=0.1, pad=0.05, shrink=0.75)
cbar.set_label('Slope TEP [$\mu$g C L$^{{-1}}$ yr$^{-1}$]', fontsize=16)
cbar.ax.tick_params(labelsize=12)


# -----------------------------------------------------
if savefig:
    plt.savefig(savepath+'TEP_LinReg_May_2000-2005.png',
                dpi=dpi, bbox_inches='tight')

plt.show(block=False)

def vw_mean_data(data, mesh, uplow=None, meshdiag=None, runid="fesom", mask=None,
                  silent = True, compare = False):
    """Calculate volume weighted mean over the range of depths.

    Parameters
     ----------
     data: xarray.DataArray, numpy.array
        Input data, that can be ether xarray data, or just numpy array,
        with values of 3d tracer variable (on nodes).
        Input should be 3d (time, nodes, levels(nz1))
        The 2D input (to calculate value for only one time step)
        in form of (nodes, levels(nz1)) is possible, but not recomended :)
    mesh: mesh object
        FESOM2 mesh object.
    uplow: list
        if None, all depths will be selected
        if e.g. [2000, 'depth'] all from model depth closest to 200 down to depth will be selected
        if e.g. [0, 700] all between 0 and closest depth to 700 will be selected.
        if e.g. [500, 500] only model level closest to 500 will be selected.
    meshdiag: str
        path to *mesh.diag.nc file, that is created during fesom cold start.
    runid: str
        name of the run. Usually just `fesom`.
    mask: array of bool
        array of boolian values of the same shape as 2D variable.
        True where data are selected.
    silent: bool
        Print diagnostics while processing.
    compare: bool
        Plot comparison of weighted and unweighted mean.

    Returns
    -------
    time series: xarray.DataArray
        time series in dimension (time,nodes) of volume weighted scalar.
    """
    
    if len(data.shape) == 2:
        data = pf.add_timedim(data)
        
    # select corresponding indices
    if uplow is None:
        uplow = [0,100]
    
    indexes = pf.select_depths(uplow, mesh,verbose=np.invert(silent))
    print(f'Selected indexes: {indexes}\nNumber of indexes: {len(indexes)-1}')
    
    # Thickness of layers
    delta_z = np.abs(np.diff(mesh.zlev))
    
    # Mask data if necessary
    if mask is not None:
        data = data[:, mask, :]
        
    # Load mesh area
    #diag = pf.get_meshdiag(mesh, meshdiag, runid)
    diag = xr.open_dataset(mesh.path+'/fesom.mesh.diag.nc')
    
    nod_area = diag.rename_dims({"nl": "nz1", "nod_n": "nod2"}).nod_area
    nod_area.load()

    # Mask mesh area if necessary
    if mask is not None:
        nod_area = nod_area[:, mask]

    # Storage arrays
    total_tracer = xr.zeros_like(data)
    aux_volume = xr.zeros_like(nod_area[0,:])
    #print(np.shape(aux_volume))

    # Calculation per layer, index also contains lower boundary!!!
    for i in indexes[:-1]: 
        if not silent:
            print('Processing depth index {0}'.format(i))
            
        #nod_area_at_level = np.ma.masked_equal(nod_area[i, :].data, 0)
        nod_vol_at_level = np.ma.masked_equal(nod_area[i, :].data, 0) * delta_z[i]
        total_tracer[:,:,i] = data[:, :, i] * nod_vol_at_level[:] 
        aux_volume = aux_volume + nod_vol_at_level
        
        if False:
            print(f'    delta_z: {delta_z[i]}')
            print(f'    sum data: {np.nansum(data[:, :, i])}')
            print(f'    nod_vol_at_level: {np.nansum(nod_vol_at_level)}')

    output = total_tracer.sum(axis=2) / aux_volume #total_volume
    
       
    if not silent:
        print('Output shape: {0}\nmin = {1:3.6f}\nmax = {2:3.6f}'.format(
            np.shape(output),np.nanmin(output), np.nanmax(output)))
        print('mean = {0:3.6f}'.format(np.nanmean(output)))
    
    return output

def get_depthIintegral_perNode(data, mesh, uplow=None, meshdiag=None, runid="fesom", mask=None,
                  silent = True):
    """Calculate integral of tracer over the range of depth.

    Parameters
     ----------
    data: xarray.DataArray, numpy.array
    mesh: mesh object
    uplow: list
    meshdiag: str
    runid: str
    mask: array of bool
    silent: bool

    Returns
    -------
    time series: xarray.DataArray
        time series in dimension (time,nodes) of integrated scalar.
    """
    
    import pyfesom2 as pf
    import numpy as np
    import xarray as xr
    
    if len(data.shape) == 2:
        data = pf.add_timedim(data)
        
    # select corresponding indices
    if uplow is None:
        uplow = [0,100]
    
    indexes = pf.select_depths(uplow, mesh,verbose=np.invert(silent))
    if not silent:
        print('Total number of selected layers: {0}'.format(len(indexes[:-1])))
        
    # Thickness of layers
    delta_z = np.abs(np.diff(mesh.zlev))
    
    # Mask data if necessary
    if mask is not None:
        data = data[:, mask, :]

    # Storage arrays
    weighted_tracer = xr.zeros_like(data)

    # Calculation per layer, index also contains lower boundary!!!
    for i in indexes[:-1]: 
        if not silent:
            print('Processing depth index {0}'.format(i))
        delta_z_at_level = delta_z[i]
        weighted_tracer[:,:,i] = data[:,:,i] * delta_z_at_level
        
    #integrated_tracer = np.nansum(weighted_tracer, axis = 2)
    integrated_tracer = weighted_tracer.sum(axis = 2)
    
    if not silent:
        print('Output shape: {0}'.format(np.shape(integrated_tracer)))
        
    return integrated_tracer