Ocean macronutrients

import xarray as xr
import glob
import numpy as np
import matplotlib.pyplot as plt
import cartopy
import cartopy.crs as ccrs
import pop_tools
from dask.distributed import LocalCluster
import s3fs

from module import adjust_pop_grid

cluster = LocalCluster()
client = cluster.get_client()

cluster.scale(20)

client

ds_grid = pop_tools.get_grid('POP_gx1v7')
lons = ds_grid.TLONG
lats = ds_grid.TLAT
depths = ds_grid.z_t * 0.01

jetstream_url = 'https://js2.jetstream-cloud.org:8001/'

s3 = s3fs.S3FileSystem(anon=True, client_kwargs=dict(endpoint_url=jetstream_url))

# Generate a list of all files in CESM folder
s3path = 's3://pythia/ocean-bgc/cesm/g.e22.GOMIPECOIAF_JRA-1p4-2018.TL319_g17.4p2z.002branch/ocn/proc/tseries/month_1/*'
remote_files = s3.glob(s3path)
s3.invalidate_cache()
# Open all files from folder
fileset = [s3.open(file) for file in remote_files]

# Open with xarray
ds = xr.open_mfdataset(fileset, data_vars="minimal", coords='minimal', compat="override", parallel=True,
                       drop_variables=["transport_components", "transport_regions", 'moc_components'], decode_times=True)

variables =['PO4','NO3','SiO3']

keep_vars=['z_t','z_t_150m','dz','time_bound','time','TAREA','TLAT','TLONG'] + variables
ds = ds.drop_vars([v for v in ds.variables if v not in keep_vars])

ds.NO3.isel(time=0,z_t=0).plot(cmap="viridis")

def year_mean(ds):
    """
    Properly convert monthly data to annual means, taking into account month lengths.
    Source: https://ncar.github.io/esds/posts/2021/yearly-averages-xarray/
    """
    
    # Make a DataArray with the number of days in each month, size = len(time)
    month_length = ds.time.dt.days_in_month

    # Calculate the weights by grouping by 'time.year'
    weights = (
        month_length.groupby("time.year") / month_length.groupby("time.year").sum()
    )

    # Test that the sum of the year for each season is 1.0
    np.testing.assert_allclose(weights.groupby("time.year").sum().values, np.ones((len(ds.groupby("time.year")), )))

    # Calculate the weighted average
    return (ds * weights).groupby("time.year").sum(dim="time")

ds_annual = year_mean(ds)
ds_annual

ds_annual['NO3'].sel(year=2010).isel(z_t=0).plot()

fig = plt.figure(figsize=(8,6))

ax = fig.add_subplot(1,1,1, projection=ccrs.Robinson(central_longitude=305.0))
land = cartopy.feature.NaturalEarthFeature('physical', 'land', scale='110m', edgecolor='k', facecolor='white', linewidth=0.5)
ax.add_feature(land)

ax.set_title('CESM surface NO$_3$', fontsize=10)
lon, lat, field = adjust_pop_grid(lons, lats, ds_annual.NO3.sel(year=2010).isel(z_t=0))
pc1=ax.pcolormesh(lon, lat,field, vmin=0, vmax=20, cmap='Greens',
                 transform=ccrs.PlateCarree())

cbar1 = fig.colorbar(pc1, ax=ax,extend='max',label='NO$_3$ (mmol m$^{-3}$)')

woa_file_path = 's3://pythia/ocean-bgc/obs/WOA2018_POPgrid.nc'

woa_file = s3.open(woa_file_path)

ds_woa = xr.load_dataset(woa_file, decode_times=False, decode_coords=False)
ds_woa['z_t'] = ds.z_t
ds_woa

ds_annual_subset = ds_annual.sel(year=[2010,2011])
ds_mean = ds_annual_subset.mean("year").compute()

NO3_diff = ds_mean.NO3 - ds_woa.NO3
PO4_diff = ds_mean.PO4 - ds_woa.PO4
SiO3_diff = ds_mean.SiO3 - ds_woa.SiO3

ds_dict_surf = {'CESMNO3': {'title': 'CESM surface NO$_3$',
                       'label': 'NO$_3$ (mmol m$^{-3}$)',
                       'cmap': 'Greens',
                       'vmin': 0, 'vmax': 20,
                       'ds': ds_mean.NO3},
           'WOANO3':  {'title': 'WOA surface NO$_3$',
                       'label': 'NO$_3$ (mmol m$^{-3}$)',
                       'cmap': 'Greens',
                       'vmin': 0, 'vmax': 20,
                       'ds': ds_woa.NO3},
           'DIFFNO3': {'title': 'Surface NO$_3$ model bias',
                       'label': 'NO$_3$ bias (mmol m$^{-3}$)',
                       'cmap': 'bwr',
                       'vmin': -10, 'vmax': 10,
                       'ds': ds_mean.NO3 - ds_woa.NO3},
           'CESMPO4': {'title': 'CESM surface PO$_4$',
                       'label': 'PO$_4$ (mmol m$^{-3}$)',
                       'cmap': 'Oranges',
                       'vmin': 0, 'vmax': 2,
                       'ds': ds_mean.PO4},
           'WOAPO4':  {'title': 'WOA surface PO$_4$',
                       'label': 'PO$_4$ (mmol m$^{-3}$)',
                       'cmap': 'Oranges',
                       'vmin': 0, 'vmax': 2,
                       'ds': ds_woa.PO4},
           'DIFFPO4': {'title': 'Surface PO$_4$ model bias',
                       'label': 'PO$_4$ bias (mmol m$^{-3}$)',
                       'cmap': 'bwr',
                       'vmin': -1, 'vmax': 1,
                       'ds': ds_mean.PO4 - ds_woa.PO4},
           'CESMSiO3': {'title': 'CESM surface SiO$_3$',
                       'label': 'SiO$_3$ (mmol m$^{-3}$)',
                       'cmap': 'Blues',
                       'vmin': 0, 'vmax': 30,
                       'ds': ds_mean.SiO3},
           'WOASiO3':  {'title': 'WOA surface SiO$_3$',
                       'label': 'SiO$_3$ (mmol m$^{-3}$)',
                       'cmap': 'Blues',
                       'vmin': 0, 'vmax': 30,
                       'ds': ds_woa.SiO3},
           'DIFFSiO3': {'title': 'Surface SiO$_3$ model bias',
                       'label': 'SiO$_3$ bias (mmol m$^{-3}$)',
                       'cmap': 'bwr',
                       'vmin': -15, 'vmax': 15,
                       'ds': ds_mean.SiO3 - ds_woa.SiO3}
          }
                        

fig = plt.figure(figsize=(18,10))


plot_count = 1
for key, item in ds_dict_surf.items():
    ax = fig.add_subplot(3,3,plot_count, projection=ccrs.Robinson(central_longitude=305.0))
    land = cartopy.feature.NaturalEarthFeature('physical', 'land', scale='110m', edgecolor='k', facecolor='white', linewidth=0.5)
    ax.add_feature(land)
    ax.set_title(item['title'], fontsize=10)
    lon, lat, field = adjust_pop_grid(lons, lats, item['ds'].isel(z_t=0))
    pc=ax.pcolormesh(lon, lat,field, vmin=item['vmin'], vmax=item['vmax'], cmap=item['cmap'],
                     transform=ccrs.PlateCarree())
    cbar1 = fig.colorbar(pc, ax=ax,label=item['label'])

    plot_count += 1
    

ds_dict_100m = {'CESMNO3': {'title': 'CESM 100m NO$_3$',
                       'label': 'NO$_3$ (mmol m$^{-3}$)',
                       'cmap': 'Greens',
                       'vmin': 0, 'vmax': 20,
                       'ds': ds_mean.NO3},
           'WOANO3':  {'title': 'WOA 100m NO$_3$',
                       'label': 'NO$_3$ (mmol m$^{-3}$)',
                       'cmap': 'Greens',
                       'vmin': 0, 'vmax': 20,
                       'ds': ds_woa.NO3},
           'DIFFNO3': {'title': '100m NO$_3$ model bias',
                       'label': 'NO$_3$ bias (mmol m$^{-3}$)',
                       'cmap': 'bwr',
                       'vmin': -10, 'vmax': 10,
                       'ds': ds_mean.NO3 - ds_woa.NO3},
           'CESMPO4': {'title': 'CESM 100m PO$_4$',
                       'label': 'PO$_4$ (mmol m$^{-3}$)',
                       'cmap': 'Oranges',
                       'vmin': 0, 'vmax': 2,
                       'ds': ds_mean.PO4},
           'WOAPO4':  {'title': 'WOA 100m PO$_4$',
                       'label': 'PO$_4$ (mmol m$^{-3}$)',
                       'cmap': 'Oranges',
                       'vmin': 0, 'vmax': 2,
                       'ds': ds_woa.PO4},
           'DIFFPO4': {'title': '100m PO$_4$ model bias',
                       'label': 'PO$_4$ bias (mmol m$^{-3}$)',
                       'cmap': 'bwr',
                       'vmin': -1, 'vmax': 1,
                       'ds': ds_mean.PO4 - ds_woa.PO4},
           'CESMSiO3': {'title': 'CESM 100m SiO$_3$',
                       'label': 'SiO$_3$ (mmol m$^{-3}$)',
                       'cmap': 'Blues',
                       'vmin': 0, 'vmax': 30,
                       'ds': ds_mean.SiO3},
           'WOASiO3':  {'title': 'WOA 100m SiO$_3$',
                       'label': 'SiO$_3$ (mmol m$^{-3}$)',
                       'cmap': 'Blues',
                       'vmin': 0, 'vmax': 30,
                       'ds': ds_woa.SiO3},
           'DIFFSiO3': {'title': '100m SiO$_3$ model bias',
                       'label': 'SiO$_3$ bias (mmol m$^{-3}$)',
                       'cmap': 'bwr',
                       'vmin': -15, 'vmax': 15,
                       'ds': ds_mean.SiO3 - ds_woa.SiO3}
          }
                        

fig = plt.figure(figsize=(18,10))


plot_count = 1
for key, item in ds_dict_100m.items():
    ax = fig.add_subplot(3,3,plot_count, projection=ccrs.Robinson(central_longitude=305.0))
    land = cartopy.feature.NaturalEarthFeature('physical', 'land', scale='110m', edgecolor='k', facecolor='white', linewidth=0.5)
    ax.add_feature(land)
    ax.set_title(item['title'], fontsize=10)
    lon, lat, field = adjust_pop_grid(lons, lats, item['ds'].isel(z_t=10))
    pc=ax.pcolormesh(lon, lat,field, vmin=item['vmin'], vmax=item['vmax'], cmap=item['cmap'],
                     transform=ccrs.PlateCarree())
    cbar1 = fig.colorbar(pc, ax=ax,label=item['label'])

    plot_count += 1

def global_mean(ds, ds_grid, compute_vars, include_ms=False):
    """
    Compute the global mean on a POP dataset. 
    Return computed quantity in conventional units.
    """
    dso = xr.Dataset({v: ds_grid[v] for v in ['z_t']})
    
    for var in compute_vars:
        
        area_depth = np.full([384,320,60],np.nan)
        var_profile = np.full([60],np.nan)

        for z in np.arange(0,60,1):
            
            if include_ms: # marginal seas
                area_depth[:,:,z] = ds_grid.TAREA.where(ds_grid.KMT > 0).where(ds[var].isel(z_t=z) > 0)
            
            else: 
                area_depth[:,:,z] = ds_grid.TAREA.where(ds_grid.REGION_MASK > 0).where(ds[var].isel(z_t=z) > 0)  
            
        area_depth = xr.DataArray(area_depth,dims=('nlat','nlon','z_t'))
            
        for z in np.arange(0,60,1):
            
            var_profile[z] = (ds[var].isel(z_t=z) * area_depth.isel(z_t=z)).sum(dim=('nlon','nlat')) / area_depth.isel(z_t=z).sum(dim=('nlon','nlat'))
    
        dso[var] = var_profile
                
    return dso

ds_glb = global_mean(ds_mean, ds_grid, ['NO3','PO4','SiO3']).compute()

ds_glb_woa = global_mean(ds_woa, ds_grid, ['NO3','PO4','SiO3']).compute()

fig = plt.figure(figsize=(6,10))

plt.suptitle('Global mean macronutrient profiles', fontsize=14)

### Row 1  - NO3

ax = fig.add_subplot(3,2,1)
ax.set_title('Global mean NO$_3$')
ax.plot(ds_glb_woa['NO3'].values, depths, label='WOA', linewidth=3, color='lightgreen')
ax.plot(ds_glb['NO3'].values, depths, label='CESM', linewidth=3, color='green')
ax.legend()
ax.set(ylabel='depth (m)',xlabel='NO$_3$ (mmol m$^{-3}$)')
plt.gca().invert_yaxis()

# Bias plot

ax = fig.add_subplot(3,2,2)
ax.plot(ds_glb['NO3'].values - ds_glb_woa['NO3'].values, depths, label='bias', linewidth=3, color='black')
ax.legend()
ax.set(ylabel='depth (m)',xlabel='NO$_3$ bias (mmol m$^{-3}$)')
ax.axvline(x=0,color='black',linestyle='--',linewidth=0.5)
plt.gca().invert_yaxis()

### Row 2  - PO4

ax = fig.add_subplot(3,2,3)
ax.set_title('Global mean PO$_4$')
ax.plot(ds_glb_woa['PO4'].values, depths, label='WOA', linewidth=3, color='peachpuff')
ax.plot(ds_glb['PO4'].values, depths, label='CESM', linewidth=3, color='orange')
ax.legend()
ax.set(ylabel='depth (m)',xlabel='PO$_4$ (mmol m$^{-3}$)')
plt.gca().invert_yaxis()

# Bias plot

ax = fig.add_subplot(3,2,4)
ax.plot(ds_glb['PO4'].values - ds_glb_woa['PO4'].values, depths, label='bias', linewidth=3, color='black')
ax.legend()
ax.set(ylabel='depth (m)',xlabel='PO$_4$ bias (mmol m$^{-3}$)')
ax.axvline(x=0,color='black',linestyle='--',linewidth=0.5)
plt.gca().invert_yaxis()

### Row 3  - SiO3

ax = fig.add_subplot(3,2,5)
ax.set_title('Global mean SiO$_3$')
ax.plot(ds_glb_woa['SiO3'].values, depths, label='WOA', linewidth=3, color='lightblue')
ax.plot(ds_glb['SiO3'].values, depths, label='CESM', linewidth=3, color='blue')
ax.legend()
ax.set(ylabel='depth (m)',xlabel='SiO$_3$ (mmol m$^{-3}$)')
plt.gca().invert_yaxis()

# Bias plot

ax = fig.add_subplot(3,2,6)
ax.plot(ds_glb['SiO3'].values - ds_glb_woa['SiO3'].values, depths, label='bias', linewidth=3, color='black')
ax.legend()
ax.set(ylabel='depth (m)',xlabel='SiO$_3$ bias (mmol m$^{-3}$)')
ax.axvline(x=0,color='black',linestyle='--',linewidth=0.5)
plt.gca().invert_yaxis()

fig.tight_layout()

cluster.close()