# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# # Source contributions

# +
import pickle
import numpy as np
import netCDF4 as nc
import pandas as pd
from calendar import monthrange

# %matplotlib inline
# -

# ###### Parameters:

# +
# domain dimensions:
imin, imax = 1479, 2179
jmin, jmax = 159, 799
isize = imax-imin
jsize = jmax-jmin

# Mn model result folders:
folder_ref      = '/data/brogalla/run_storage/Mn-reference-202110/'
folder_cleanice = '/data/brogalla/run_storage/Mn-clean-ice-202110/'
folder_spm      = '/data/brogalla/run_storage/Mn-spm-202110/'
folder_bio      = '/data/brogalla/run_storage/Mn-bio-202110/'

colors = ['#ccb598', '#448d90', '#739f78', '#CC8741', '#cee7fd', '#b9c1c7']
# -

years = [2002, 2003, 2004, 2005, 2006, 2007, 2008, 2009, 2010, 2011, 2012, 2013, 2014, 2015, 2016, \
         2017, 2018, 2019]

# ###### Load files:

# +
# ANHA12 grid:
mask       = nc.Dataset('/ocean/brogalla/GEOTRACES/data/ANHA12/ANHA12_mesh1.nc')
tmask      = np.array(mask.variables['tmask'])[0,:,imin:imax,jmin:jmax]
land_mask  = np.ma.masked_where((tmask[:,:,:] > 0.1), tmask[:,:,:]) 
e1t_base   = np.array(mask.variables['e1t'])[0,imin:imax,jmin:jmax]
e2t_base   = np.array(mask.variables['e2t'])[0,imin:imax,jmin:jmax]
e3t        = np.array(mask.variables['e3t_0'])[0,:,imin:imax,jmin:jmax]
e3t_masked = np.ma.masked_where((tmask[:,:,:] < 0.1), e3t)
nav_lev    = np.array(mask.variables['nav_lev'])
lons       = np.array(mask.variables['nav_lon'])

e1t           = np.tile(e1t_base, (50,1,1))
e2t           = np.tile(e2t_base, (50,1,1))
volume        = e1t*e2t*e3t
area_base     = e1t_base*e2t_base
volume_masked = np.ma.masked_where((tmask[:,:,:] < 0.1), volume)
area_masked   = np.ma.masked_where((tmask[0,:,:] < 0.1), area_base)


# -

# ##### Functions:

def load_results(folder_ref, year, experiment):
    
    months  = ['01', '02', '03', '04', '05', '06', '07', '08', '09', '10', '11', '12']
    dmn_riv = np.empty((12,isize,jsize))   ; dmn_sed = np.empty((12,isize,jsize)); 
    dmn_atm = np.empty((12,isize,jsize))   ; dmn_ice = np.empty((12,isize,jsize));
    dmn_sedice = np.empty((12,isize,jsize)); dmn_priv = np.empty((12,isize,jsize));
    dmn_red = np.empty((12,50,isize,jsize)); dmn_oxi = np.empty((12,50,isize,jsize));
    dmn_bio = np.empty((12,50,isize,jsize));
    
    for i, month in enumerate(months):
        file        = f'ANHA12_EXH006_1m_{year}0101_{year}1231_comp_{year}{month}-{year}{month}.nc'
        ref_monthly = nc.Dataset(folder_ref+file)
        
        dmn_riv[i,:,:]    = np.array(ref_monthly.variables['dmnriv'])[0,:,:]
        
        if experiment=='spm':
            dmn_priv[i,:,:] = np.array(ref_monthly.variables['pmnriv'])[0,:,:]
            dmn_bio[i,:,:]  = 0
        elif experiment=='bio':
            dmn_priv[i,:,:] = 0
            dmn_bio[i,:,:]  = np.array(ref_monthly.variables['dmnbio'])[0,:,:,:]
        else:
            dmn_priv[i,:,:] = 0
            dmn_bio[i,:,:]  = 0
        
        dmn_sed[i,:,:]    = np.array(ref_monthly.variables['dmnsed'])[0,:,:]
        dmn_sedice[i,:,:] = np.array(ref_monthly.variables['dmnsedice'])[0,:,:]
        dmn_atm[i,:,:]    = np.array(ref_monthly.variables['dmnatm'])[0,:,:]
        dmn_ice[i,:,:]    = np.array(ref_monthly.variables['dmnice'])[0,:,:]
        dmn_red[i,:,:,:]  = np.array(ref_monthly.variables['dmnred'])[0,:,:,:]
        dmn_oxi[i,:,:,:]  = np.array(ref_monthly.variables['dmnoxi'])[0,:,:,:]
    
    tmask_surf    = np.empty(dmn_riv.shape)
    tmask_surf[:] = tmask[0,:,:]
    tmask_full    = np.empty(dmn_red.shape)
    tmask_full[:] = tmask
    
    # Mask points on land:
    dmn_riv    = np.ma.masked_where((tmask_surf < 0.1), dmn_riv); 
    dmn_priv   = np.ma.masked_where((tmask_surf < 0.1), dmn_priv); 
    dmn_sed    = np.ma.masked_where((tmask_surf < 0.1), dmn_sed)
    dmn_sedice = np.ma.masked_where((tmask_surf < 0.1), dmn_sedice)
    dmn_atm    = np.ma.masked_where((tmask_surf < 0.1), dmn_atm)
    dmn_ice    = np.ma.masked_where((tmask_surf < 0.1), dmn_ice)
    dmn_red    = np.ma.masked_where((tmask_full < 0.1), dmn_red)
    dmn_oxi    = np.ma.masked_where((tmask_full < 0.1), dmn_oxi)
    dmn_bio    = np.ma.masked_where((tmask_full < 0.1), dmn_bio)

    return dmn_riv, dmn_priv, dmn_sed, dmn_sedice, dmn_atm, dmn_ice, dmn_red, dmn_oxi, dmn_bio


def calculate_contribution(year, mask, experiment='ref'):
    # Calculate the contribution of the model components by region:
    if experiment=='spm':
        folder_year = folder_spm + f'ANHA12_spm-{year}_20211026/'
    elif experiment=='bio':
        folder_year = folder_bio + f'ANHA12_bio-{year}_20211022/'
    else:
        folder_year = folder_ref + f'ANHA12_ref-{year}_20211012/'
    
    if experiment=='bio' and year>2015:
        total_yearly_contribution = np.empty((9))
        ave_yearly_contribution   = np.empty((9)) 
        total_yearly_contribution[:] = np.NaN
        ave_yearly_contribution[:]   = np.NaN
        return total_yearly_contribution, ave_yearly_contribution
    
    dmn_mriv, dmn_mpriv, dmn_msed, dmn_msedice, dmn_matm, dmn_mice, dmn_mred, dmn_moxi, dmn_mbio = \
                       load_results(folder_year, year, experiment)

    # Calculate contributions to the upper water column, so mask locations where sediment resuspension is added in the deep.
    indexes_bottom = np.tile(index_bottom, (12,1,1))
    
    # Momentary contribution: moles / second (from moles / L / s --> moles / m3 / s --> moles / s)
    # Mask is for the points associated with the specified region. 
    priv   = np.ma.masked_where(mask==0, dmn_mpriv)   * volume_masked[0,:,:] * 1e3
    riv    = np.ma.masked_where(mask==0, dmn_mriv)    * volume_masked[0,:,:] * 1e3
    sed    = np.ma.masked_where(mask==0, dmn_msed)    * volume_bottom        * 1e3
    sedice = np.ma.masked_where(mask==0, dmn_msedice) * volume_masked[0,:,:] * 1e3
    atm    = np.ma.masked_where(mask==0, dmn_matm)    * volume_masked[0,:,:] * 1e3
    ice    = np.ma.masked_where(mask==0, dmn_mice)    * volume_masked[0,:,:] * 1e3
    # 3D fields only calculated over polar mixed layer:
    mask_with_depths = np.tile(mask, (50,1,1,1)).reshape((12,50,700,640))
    bio    = np.ma.masked_where(mask_with_depths==0, dmn_mbio)[:,0:17,:,:] * volume_masked[0:17,:,:] * 1e3
    red    = np.ma.masked_where(mask_with_depths==0, dmn_mred)[:,0:17,:,:] * volume_masked[0:17,:,:] * 1e3
    oxi    = np.ma.masked_where(mask_with_depths==0, dmn_moxi)[:,0:17,:,:] * volume_masked[0:17,:,:] * 1e3

    # Yearly contribution to domain: moles / year (from moles / s / month --> moles / month --> moles / year)
    atm_year = 0; riv_year = 0; priv_year = 0; sed_year = 0; sedice_year = 0;
    ice_year = 0; bio_year = 0; red_year = 0; oxi_year = 0;
    for month in range(1,13):   
        index = month-1
        days_in_month = monthrange(year, month)[1]
        if days_in_month == 29: # ignore leap year
            days_in_month = 28
        
        atm_year    = atm_year    + np.ma.sum(atm[index,:,:],axis=(0,1))   *3600*24*days_in_month
        riv_year    = riv_year    + np.ma.sum(riv[index,:,:],axis=(0,1))   *3600*24*days_in_month
        priv_year   = priv_year   + np.ma.sum(priv[index,:,:],axis=(0,1))  *3600*24*days_in_month
        sed_year    = sed_year    + np.ma.sum(sed[index,:,:],axis=(0,1))   *3600*24*days_in_month
        sedice_year = sedice_year + np.ma.sum(sedice[index,:,:],axis=(0,1))*3600*24*days_in_month
        ice_year    = ice_year    + np.ma.sum(ice[index,:,:],axis=(0,1))   *3600*24*days_in_month
        
        bio_year    = bio_year    + np.ma.sum(bio[index,:,:,:],axis=(0,1,2)) *3600*24*days_in_month
        red_year    = red_year    + np.ma.sum(red[index,:,:,:],axis=(0,1,2)) *3600*24*days_in_month
        oxi_year    = oxi_year    + np.ma.sum(oxi[index,:,:,:],axis=(0,1,2)) *3600*24*days_in_month
    
    # Average yearly contribution over domain: moles / m2 / year
    total_area     = np.ma.sum(np.ma.masked_where(mask[0,:,:]==0, area_masked[:,:]))
    atmm2_year     =  atm_year    / total_area
    rivm2_year     =  riv_year    / total_area
    privm2_year    =  priv_year   / total_area
    sedm2_year     =  sed_year    / total_area 
    sedicem2_year  =  sedice_year / total_area
    icem2_year     =  ice_year    / total_area   
    # Polar mixed layer average contribution:
    PML_depth      = nav_lev[17]
    biom2_year     =  bio_year    / (total_area*PML_depth)
    redm2_year     =  red_year    / (total_area*PML_depth)
    oxim2_year     =  oxi_year    / (total_area*PML_depth)
    
    total_yearly_contribution = np.array([priv_year, riv_year, sed_year, sedice_year, ice_year, atm_year, \
                                          bio_year, red_year, oxi_year])           # mol/yr
    ave_yearly_contribution   = np.array([privm2_year, rivm2_year, sedm2_year, sedicem2_year, icem2_year, \
                                          atmm2_year, biom2_year, redm2_year, oxim2_year]) # mol/m2/yr
 
    return total_yearly_contribution, ave_yearly_contribution


def calculate_regional_contributions(mask):
    
    totals_ref   = np.empty((len(years),9)); totals_spm   = np.empty((len(years),9)); totals_bio   = np.empty((len(years),9));
    averages_ref = np.empty((len(years),9)); averages_spm = np.empty((len(years),9)); averages_bio = np.empty((len(years),9));

    totals_ref[:]   = np.NaN; totals_spm[:]    = np.NaN; totals_bio[:]   = np.NaN;
    averages_ref[:] = np.NaN; averages_spm[:]  = np.NaN; averages_spm[:] = np.NaN;

    for i, year in enumerate(years):
        total_ref_year, ave_ref_year  = calculate_contribution(year, mask, experiment='ref')
        total_spm_year, ave_spm_year  = calculate_contribution(year, mask, experiment='spm')
        total_bio_year, ave_bio_year  = calculate_contribution(year, mask, experiment='bio')

        totals_ref[i,:]        = total_ref_year 
        totals_spm[i,:]        = total_spm_year 
        totals_bio[i,:]        = total_bio_year 
        averages_ref[i,:]      = ave_ref_year  
        averages_spm[i,:]      = ave_spm_year 
        averages_bio[i,:]      = ave_bio_year 
    
    return totals_ref, totals_spm, totals_bio, averages_ref, averages_spm, averages_bio


def reshape_arrays(array_in):
    fixed      = np.zeros((len(years),8))
    
    for i in range(0,8):
        if i==0: # Sum the particulate and dissolved river contributions
            fixed[:,i] = array_in[:,i]+array_in[:,i+1]
        else:
            fixed[:,i] = array_in[:,i+1]
    
    return fixed


def pipeline(totals_input_ref, totals_input_spm, totals_input_bio, 
             averages_input_ref, averages_input_spm, averages_input_bio):
    
    regions = ['full', 'CB', 'CAA', 'wCAA', 'eCAA']
    
    totals_fixed_ref = totals_input_ref.copy(); averages_fixed_ref = averages_input_ref.copy();
    totals_fixed_spm = totals_input_spm.copy(); averages_fixed_spm = averages_input_spm.copy();
    totals_fixed_bio = totals_input_bio.copy(); averages_fixed_bio = averages_input_bio.copy();
    
    # Reshape arrays so that particulate and dissolved river contributions are combined into one.
    for region in regions:

        totals_fixed_ref[region]   = reshape_arrays(totals_input_ref[region])
        totals_fixed_spm[region]   = reshape_arrays(totals_input_spm[region])
        totals_fixed_bio[region]   = reshape_arrays(totals_input_bio[region])
        averages_fixed_ref[region] = reshape_arrays(averages_input_ref[region])
        averages_fixed_spm[region] = reshape_arrays(averages_input_spm[region])
        averages_fixed_bio[region] = reshape_arrays(averages_input_bio[region])
    
    return totals_fixed_ref, totals_fixed_spm, totals_fixed_bio, averages_fixed_ref, averages_fixed_spm, averages_fixed_bio


def print_interannual_averages(totals_ref, totals_spm, totals_bio, 
                               averages_ref, averages_spm, averages_bio,
                               location='Full'):

    print(f'Interannual average values in {location}')
    print('----------------------------------------------------------------')
    total_ref_annual = np.sum(totals_ref[location][:,0:5], axis=1)
    total_spm_annual = np.sum(totals_spm[location][:,0:5], axis=1)
    total_bio_annual = np.sum(totals_bio[location][:,0:5], axis=1)
    
    # Total annual contribution is averaged over the 18-year time series and multiplied by 1e-6 to convert to Mmol
    print(f"{'Total annual contribution of Mn [Mmol/yr]:':<55}" +
          f"{np.average(total_ref_annual)*1e-6:<5.0f} ({np.average(total_spm_annual)*1e-6:<5.0f})")
    print(f"{'-- River discharge ---':<40}" +
          f" {np.average(totals_ref[location][:,0])*1e-6:<8.2f} ({np.average(totals_spm[location][:,0])*1e-6:<5.2f})")
    print(f"{'-- Sediment resuspension ---':<40}" +
          f" {np.average(totals_ref[location][:,1])*1e-6:<8.2f} ({np.average(totals_spm[location][:,1])*1e-6:<5.2f})")
    print(f"{'-- Sediment released by sea ice ---':<40}" +
          f" {np.average(totals_ref[location][:,2])*1e-6:<8.2f} ({np.average(totals_spm[location][:,2])*1e-6:<5.2f})")
    print(f"{'-- Dust released by sea ice ---':<40}" +
          f" {np.average(totals_ref[location][:,3])*1e-6:<8.2f} ({np.average(totals_spm[location][:,3])*1e-6:<5.2f})")
    print(f"{'-- Dust deposition ---':<40}" +
          f" {np.average(totals_ref[location][:,4])*1e-6:<8.2f} ({np.average(totals_spm[location][:,4])*1e-6:<5.2f})") 

    tref = np.sum(totals_ref[location][:,0:5], axis=(0,1))
    tspm = np.sum(totals_spm[location][:,0:5], axis=(0,1))
    tbio = np.sum(totals_bio[location][:,0:5], axis=(0,1))

    # Percent is calculated from sum of component contribution over all months / total contributions over full time period
    print('----------------------------------------------------------------')
    print('Total annual contribution of Mn [%]:') 
    print(f"{'-- River discharge ---':<40}" +
          f"{np.sum(totals_ref[location][:,0])*100/tref:<5.2f} ({np.sum(totals_spm[location][:,0])*100/tspm:<5.2f})")
    print(f"{'-- Sediment resuspension ---':<40}" +
          f"{np.sum(totals_ref[location][:,1])*100/tref:<5.2f} ({np.sum(totals_spm[location][:,1])*100/tspm:<5.2f})")
    print(f"{'-- Sediment released by sea ice ---':<40}" +
          f"{np.sum(totals_ref[location][:,2])*100/tref:<5.2f} ({np.sum(totals_spm[location][:,2])*100/tspm:<5.2f})")
    print(f"{'-- Dust released by sea ice ---':<40}" +
          f"{np.sum(totals_ref[location][:,3])*100/tref:<5.2f} ({np.sum(totals_spm[location][:,3])*100/tspm:<5.2f})")
    print(f"{'-- Dust deposition ---':<40}" +
          f"{np.sum(totals_ref[location][:,4])*100/tref:<5.2f} ({np.sum(totals_spm[location][:,4])*100/tspm:<5.2f})")

    # Average annual contribution is calculated as an average over the full time series and converted to millimol
    print('----------------------------------------------------------------')
    print(f"{'Average annual contribution of Mn [micromol/m2/yr]:':<55}" +
          f"{np.average(np.sum(averages_ref[location][:,0:5], axis=1)*1e6):<5.0f}" +
          f"({np.average(np.sum(averages_spm[location][:,0:5], axis=1)*1e6):<5.0f})")
    print(f"{'-- River discharge ---':<40}" +
          f"{np.average(averages_ref[location][:,0])*1e6:<8.2f} ({np.average(averages_spm[location][:,0])*1e6:<5.2f})")
    print(f"{'-- Sediment resuspension ---':<40}" +
          f"{np.average(averages_ref[location][:,1])*1e6:<8.2f} ({np.average(averages_spm[location][:,1])*1e6:<5.2f})")
    print(f"{'-- Sediment released by sea ice ---':<40}" +
          f"{np.average(averages_ref[location][:,2])*1e6:<8.2f} ({np.average(averages_spm[location][:,2])*1e6:<5.2f})")
    print(f"{'-- Dust released by sea ice ---':<40}" +
          f"{np.average(averages_ref[location][:,3])*1e6:<8.2f} ({np.average(averages_spm[location][:,3])*1e6:<5.2f})")
    print(f"{'-- Dust deposition ---':<40}" +
          f"{np.average(averages_ref[location][:,4])*1e6:<8.2f} ({np.average(averages_spm[location][:,4])*1e6:<5.2f})")

    return


def print_interannual_averages_3D(totals_ref, totals_spm, totals_bio, 
                               averages_ref, averages_spm, averages_bio,
                               location='Full'):
    # include reduction, oxidation, and bio components calculated over the polar mixed layer

    print(f'Interannual average values in {location}')
    print('----------------------------------------------------------------')
    total_ref_annual = np.sum(totals_ref[location], axis=1)
    total_spm_annual = np.sum(totals_spm[location], axis=1)
    total_bio_annual = np.nansum(totals_bio[location], axis=1)
    
    print(f"{'Total annual contribution of Mn [Mmol/yr]:':<55}" +
          f"{np.average(total_ref_annual)*1e-6:<5.0f} ({np.average(total_spm_annual)*1e-6:<5.0f})" + \
          f" {np.nanmean(total_bio_annual)*1e-6:<8.2f}")
    print(f"{'-- River discharge ---':<40}" +
          f" {np.average(totals_ref[location][:,0])*1e-6:<8.2f} ({np.average(totals_spm[location][:,0])*1e-6:<5.2f})" + \
          f" {np.nanmean(totals_bio[location][:,0])*1e-6:<8.2f}")
    print(f"{'-- Sediment resuspension ---':<40}" +
          f" {np.average(totals_ref[location][:,1])*1e-6:<8.2f} ({np.average(totals_spm[location][:,1])*1e-6:<5.2f})" + \
          f" {np.nanmean(totals_bio[location][:,1])*1e-6:<8.2f}")
    print(f"{'-- Sediment released by sea ice ---':<40}" +
          f" {np.average(totals_ref[location][:,2])*1e-6:<8.2f} ({np.average(totals_spm[location][:,2])*1e-6:<5.2f})" + \
          f" {np.nanmean(totals_bio[location][:,2])*1e-6:<8.2f}")
    print(f"{'-- Dust released by sea ice ---':<40}" +
          f" {np.average(totals_ref[location][:,3])*1e-6:<8.2f} ({np.average(totals_spm[location][:,3])*1e-6:<5.2f})" + \
          f" {np.nanmean(totals_bio[location][:,3])*1e-6:<8.2f}")
    print(f"{'-- Dust deposition ---':<40}" +
          f" {np.average(totals_ref[location][:,4])*1e-6:<8.2f} ({np.average(totals_spm[location][:,4])*1e-6:<5.2f})" + \
          f" {np.nanmean(totals_bio[location][:,4])*1e-6:<8.2f}")
    print(f"{'-- Uptake/remineralization ---':<40}" +
          f" {np.average(totals_ref[location][:,5])*1e-6:<8.2f} ({np.average(totals_spm[location][:,5])*1e-6:<5.2f})" + \
          f" {np.nanmean(totals_bio[location][:,5])*1e-6:<8.2f}")
    print(f"{'-- Scavenging ---':<40}" +
          f" {(np.average(totals_ref[location][:,6])-np.average(totals_ref[location][:,7]))*1e-6:<8.2f}" + \
          f"({(np.average(totals_spm[location][:,6])-np.average(totals_spm[location][:,7]))*1e-6:<5.2f})" + \
          f" {(np.nanmean(totals_bio[location][:,6])-np.nanmean(totals_bio[location][:,7]))*1e-6:<8.2f}")

    tref = np.sum(totals_ref[location], axis=(0,1))
    tspm = np.sum(totals_spm[location], axis=(0,1))
    tbio = np.nansum(totals_bio[location], axis=(0,1))

    print('----------------------------------------------------------------')
    print('Total annual contribution of Mn [%]:') 
    print(f"{'-- River discharge ---':<40}" +
          f"{np.sum(totals_ref[location][:,0])*100/tref:<5.2f} ({np.sum(totals_spm[location][:,0])*100/tspm:<5.2f})" + \
          f"{np.nansum(totals_bio[location][:,0])*100/tref:<5.2f}")
    print(f"{'-- Sediment resuspension ---':<40}" +
          f"{np.sum(totals_ref[location][:,1])*100/tref:<5.2f} ({np.sum(totals_spm[location][:,1])*100/tspm:<5.2f})" + \
          f"{np.nansum(totals_bio[location][:,1])*100/tref:<5.2f}")
    print(f"{'-- Sediment released by sea ice ---':<40}" +
          f"{np.sum(totals_ref[location][:,2])*100/tref:<5.2f} ({np.sum(totals_spm[location][:,2])*100/tspm:<5.2f})" + \
          f"{np.nansum(totals_bio[location][:,2])*100/tref:<5.2f}")
    print(f"{'-- Dust released by sea ice ---':<40}" +
          f"{np.sum(totals_ref[location][:,3])*100/tref:<5.2f} ({np.sum(totals_spm[location][:,3])*100/tspm:<5.2f})" + \
          f"{np.nansum(totals_bio[location][:,3])*100/tref:<5.2f}")
    print(f"{'-- Dust deposition ---':<40}" +
          f"{np.sum(totals_ref[location][:,4])*100/tref:<5.2f} ({np.sum(totals_spm[location][:,4])*100/tspm:<5.2f})" + \
          f"{np.nansum(totals_bio[location][:,4])*100/tref:<5.2f}")
    print(f"{'-- Uptake/remineralization ---':<40}" +
          f"{np.sum(totals_ref[location][:,5])*100/tref:<5.2f} ({np.sum(totals_spm[location][:,5])*100/tspm:<5.2f})" + \
          f"{np.nansum(totals_bio[location][:,5])*100/tref:<5.2f}")
    print(f"{'-- Scavenging ---':<40}" +
          f"{np.sum(totals_ref[location][:,6]   - totals_ref[location][:,7])*100/tref:<5.2f} "+ \
          f"({np.sum(totals_spm[location][:,6]  - totals_spm[location][:,7])*100/tspm:<5.2f})" + \
          f"{np.nansum(totals_bio[location][:,6]- totals_bio[location][:,7])*100/tref:<5.2f}")

    print('----------------------------------------------------------------')
    print(f"{'Average annual contribution of Mn [micromol/m2/yr]:':<55}" +
          f"{np.average(np.sum(averages_ref[location], axis=1)*1e6):<5.0f}" +
          f"({np.average(np.sum(averages_spm[location], axis=1)*1e6):<5.0f})"+
          f"{np.nanmean(np.nansum(averages_bio[location], axis=1)*1e6):<5.0f}")
    print(f"{'-- River discharge ---':<40}" +
          f"{np.average(averages_ref[location][:,0])*1e6:<8.2f} ({np.average(averages_spm[location][:,0])*1e6:<5.2f})" +\
          f"{np.nanmean(averages_bio[location][:,0])*1e6:<8.2f}")
    print(f"{'-- Sediment resuspension ---':<40}" +
          f"{np.average(averages_ref[location][:,1])*1e6:<8.2f} ({np.average(averages_spm[location][:,1])*1e6:<5.2f})" +\
          f"{np.nanmean(averages_bio[location][:,1])*1e6:<8.2f}")
    print(f"{'-- Sediment released by sea ice ---':<40}" +
          f"{np.average(averages_ref[location][:,2])*1e6:<8.2f} ({np.average(averages_spm[location][:,2])*1e6:<5.2f})" +\
          f"{np.nanmean(averages_bio[location][:,2])*1e6:<8.2f}")
    print(f"{'-- Dust released by sea ice ---':<40}" +
          f"{np.average(averages_ref[location][:,3])*1e6:<8.2f} ({np.average(averages_spm[location][:,3])*1e6:<5.2f})" +\
          f"{np.nanmean(averages_bio[location][:,3])*1e6:<8.2f}")
    print(f"{'-- Dust deposition ---':<40}" +
          f"{np.average(averages_ref[location][:,4])*1e6:<8.2f} ({np.average(averages_spm[location][:,4])*1e6:<5.2f})" +\
          f"{np.nanmean(averages_bio[location][:,4])*1e6:<8.2f}")
    print(f"{'-- Uptake/remineralization ---':<40}" +
          f"{np.average(averages_ref[location][:,5])*1e6:<8.2f} ({np.average(averages_spm[location][:,5])*1e6:<5.2f})" +\
          f"{np.nanmean(averages_bio[location][:,5])*1e6:<8.2f}")
    print(f"{'-- Scavenging ---':<40}" +
          f"{(np.average(averages_ref[location][:,6]) - np.average(averages_ref[location][:,7]))*1e6:<8.2f}"+ \
          f"({(np.average(averages_spm[location][:,6])- np.average(averages_spm[location][:,7]))*1e6:<5.2f})" +\
          f"{(np.nanmean(averages_bio[location][:,6]) - np.nanmean(averages_bio[location][:,7]))*1e6:<8.2f}")

    return


# ### Calculations:

# Find grid cell volume and thickness at ocean floor:

# +
tmask_bot = np.copy(tmask)
tmask_bot[0,:,:] = 1
e3t_mask = np.ma.masked_where((tmask_bot[:,:,:] < 0.1), e3t[:,:,:])

# find bottom grid cell index:
ind = np.array(np.ma.notmasked_edges(e3t_mask[:,:,:], axis=0))

# array of shape (isize, jsize) containing bottom grid cell index
index_bottom = np.ma.masked_where((tmask[0,:,:] < 0.1), ind[1][0][:].reshape(isize,jsize))

# +
e3t_bot = np.zeros((isize,jsize))
vol_bot = np.zeros((isize,jsize))

for i in range(0, isize):
    for j in range(0,jsize):
        k = index_bottom[i,j]
        try:
            e3t_bot[i,j] = e3t[k,i,j]
            vol_bot[i,j] = volume[k,i,j]
        except:
            e3t_bot[i,j] = np.nan
            vol_bot[i,j] = np.nan
        
e3t_bottom    = np.ma.masked_where((tmask[0,:,:] < 0.1), np.ma.masked_where(np.isnan(e3t_bot), e3t_bot))
volume_bottom = np.ma.masked_where((tmask[0,:,:] < 0.1), np.ma.masked_where(np.isnan(vol_bot), vol_bot))
# -

# Create masks for specific subregions:
# - Canada Basin
# - CAA
# - Full domain
# - Western CAA (separated by Barrow Sill)
# - Eastern CAA (separated by Barrow Sill)

# +
# Find indices to separate out Canada Basin:

x_ind = np.arange(1621, 2100, 1)
y_ind = (-7/8)*x_ind + 1517 + 700

CB_indx = []
CB_indy = []

for index in range(0,len(x_ind)):
    CB_x = np.arange(x_ind[index],2179,1)
    CB_y = np.ones(CB_x.shape)*y_ind[index]
    
    CB_indx = np.append(CB_x, CB_indx)
    CB_indy = np.append(CB_y, CB_indy)

# +
# Separate Canada Basin and the CAA:
mask_ini_CB  = np.zeros((isize,jsize))
mask_ini_CAA = np.ones((isize,jsize))

for i, j in zip(CB_indx, CB_indy):
    mask_ini_CB[int(i-imin),int(j-jmin)]  = 1
    mask_ini_CAA[int(i-imin),int(j-jmin)] = 0
    
mask_ini_CB[150:-1 ,-8:-1] = 1
mask_ini_CAA[150:-1,-8:-1] = 0

# +
mask_ini_wCAA = np.zeros((isize,jsize))
mask_ini_eCAA = np.zeros((isize,jsize))

# Separate the western and eastern CAA:
mask_ini_wCAA = np.where(lons[imin:imax,jmin:jmax] < -100, 1, 0)
mask_ini_eCAA = np.where(lons[imin:imax,jmin:jmax] > -100, 1, 0)

for i, j in zip(CB_indx, CB_indy):
    mask_ini_wCAA[int(i-imin),int(j-jmin)] = 0
    mask_ini_eCAA[int(i-imin),int(j-jmin)] = 0
    
mask_ini_wCAA[150:-1,-8:-1] = 0
mask_ini_eCAA[150:-1,-8:-1] = 0
# -

mask_CB_yr    = np.tile(mask_ini_CB  , (12,1,1))
mask_CAA_yr   = np.tile(mask_ini_CAA , (12,1,1))
mask_full_yr  = np.ones((12, isize, jsize))
mask_wCAA_yr  = np.tile(mask_ini_wCAA, (12,1,1))
mask_eCAA_yr  = np.tile(mask_ini_eCAA, (12,1,1))

mask_CB_depth = np.tile(mask_CB_yr  , (50,1,1,1))
mask_CB_depth = mask_CB_depth.reshape(12,50,700,640)

# ### Calculate and pickle the results:

full_totals_ref, full_totals_spm, full_totals_bio, \
        full_averages_ref, full_averages_spm, full_averages_bio = calculate_regional_contributions(mask_full_yr)

pickle.dump((full_totals_ref  , full_totals_spm, full_totals_bio    ), open('calculations/full-domain-totals-202110.pickle','wb'))
pickle.dump((full_averages_ref, full_averages_spm, full_averages_bio), open('calculations/full-domain-averages-202110.pickle','wb'))

CB_totals_ref, CB_totals_spm, CB_totals_bio, \
        CB_averages_ref, CB_averages_spm, CB_averages_bio = calculate_regional_contributions(mask_CB_yr)

pickle.dump((CB_totals_ref  , CB_totals_spm, CB_totals_bio    ), open('calculations/Canada-Basin-totals-202110--fullcolumn.pickle','wb'))
pickle.dump((CB_averages_ref, CB_averages_spm, CB_averages_bio), open('calculations/Canada-Basin-averages-202110--fullcolumn.pickle','wb'))

CAA_totals_ref, CAA_totals_spm, CAA_totals_bio, \
        CAA_averages_ref, CAA_averages_spm, CAA_averages_bio = calculate_regional_contributions(mask_CAA_yr)

pickle.dump((CAA_totals_ref  , CAA_totals_spm, CAA_totals_bio    ), open('calculations/CAA-totals-202110--fullcolumn.pickle','wb'))
pickle.dump((CAA_averages_ref, CAA_averages_spm, CAA_averages_bio), open('calculations/CAA-averages-202110--fullcolumn.pickle','wb'))

wCAA_totals_ref, wCAA_totals_spm, wCAA_totals_bio, \
        wCAA_averages_ref, wCAA_averages_spm, wCAA_averages_bio = calculate_regional_contributions(mask_wCAA_yr)

pickle.dump((wCAA_totals_ref  , wCAA_totals_spm, wCAA_totals_bio    ), open('calculations/wCAA-totals-202110.pickle','wb'))
pickle.dump((wCAA_averages_ref, wCAA_averages_spm, wCAA_averages_bio), open('calculations/wCAA-averages-202110.pickle','wb'))

eCAA_totals_ref, eCAA_totals_spm, eCAA_totals_bio,  \
        eCAA_averages_ref, eCAA_averages_spm, eCAA_averages_bio = calculate_regional_contributions(mask_eCAA_yr)

pickle.dump((eCAA_totals_ref  , eCAA_totals_spm, eCAA_totals_bio    ), open('calculations/eCAA-totals-202110.pickle','wb'))
pickle.dump((eCAA_averages_ref, eCAA_averages_spm, eCAA_averages_bio), open('calculations/eCAA-averages-202110.pickle','wb'))

# ### Overview:

# +
# Component contribution calculations partial water column:
CB_totals_ref,   CB_totals_spm,   CB_totals_bio   = pickle.load(open('calculations/Canada-Basin-totals-202110.pickle','rb')) 
CAA_totals_ref,  CAA_totals_spm,  CAA_totals_bio  = pickle.load(open('calculations/CAA-totals-202110.pickle' ,'rb')) 
wCAA_totals_ref, wCAA_totals_spm, wCAA_totals_bio = pickle.load(open('calculations/wCAA-totals-202110.pickle','rb'))  
eCAA_totals_ref, eCAA_totals_spm, eCAA_totals_bio = pickle.load(open('calculations/eCAA-totals-202110.pickle','rb')) 
totals_ref     , totals_spm     , totals_bio      = pickle.load(open('calculations/full-domain-totals-202110.pickle','rb')) 

ref_totals   = {'full': totals_ref, 'CB': CB_totals_ref, 'CAA':CAA_totals_ref, 'wCAA':wCAA_totals_ref, 'eCAA':eCAA_totals_ref}
spm_totals   = {'full': totals_spm, 'CB': CB_totals_spm, 'CAA':CAA_totals_spm, 'wCAA':wCAA_totals_spm, 'eCAA':eCAA_totals_spm}
bio_totals   = {'full': totals_bio, 'CB': CB_totals_bio, 'CAA':CAA_totals_bio, 'wCAA':wCAA_totals_bio, 'eCAA':eCAA_totals_bio}

# +
# Component contribution calculations partial water column:
CB_averages_ref,   CB_averages_spm,   CB_averages_bio   = pickle.load(open('calculations/Canada-Basin-averages-202110.pickle','rb')) 
CAA_averages_ref,  CAA_averages_spm,  CAA_averages_bio  = pickle.load(open('calculations/CAA-averages-202110.pickle' ,'rb')) 
wCAA_averages_ref, wCAA_averages_spm, wCAA_averages_bio = pickle.load(open('calculations/wCAA-averages-202110.pickle','rb'))  
eCAA_averages_ref, eCAA_averages_spm, eCAA_averages_bio = pickle.load(open('calculations/eCAA-averages-202110.pickle','rb')) 
averages_ref     , averages_spm     , averages_bio      = pickle.load(open('calculations/full-domain-averages-202110.pickle','rb')) 

ref_averages   = {'full': averages_ref, 'CB': CB_averages_ref, 'CAA':CAA_averages_ref, 'wCAA':wCAA_averages_ref, 'eCAA':eCAA_averages_ref}
spm_averages   = {'full': averages_spm, 'CB': CB_averages_spm, 'CAA':CAA_averages_spm, 'wCAA':wCAA_averages_spm, 'eCAA':eCAA_averages_spm}
bio_averages   = {'full': averages_bio, 'CB': CB_averages_bio, 'CAA':CAA_averages_bio, 'wCAA':wCAA_averages_bio, 'eCAA':eCAA_averages_bio}

# +
# Component contribution calculations full water column:
CB_totals_ref,   CB_totals_spm,   CB_totals_bio   = pickle.load(open('calculations/Canada-Basin-totals-202110--fullcolumn.pickle','rb')) 
CAA_totals_ref,  CAA_totals_spm,  CAA_totals_bio  = pickle.load(open('calculations/CAA-totals-202110--fullcolumn.pickle' ,'rb')) 
wCAA_totals_ref, wCAA_totals_spm, wCAA_totals_bio = pickle.load(open('calculations/wCAA-totals-202110.pickle','rb'))  
eCAA_totals_ref, eCAA_totals_spm, eCAA_totals_bio = pickle.load(open('calculations/eCAA-totals-202110.pickle','rb')) 
totals_ref     , totals_spm     , totals_bio      = pickle.load(open('calculations/full-domain-totals-202110.pickle','rb')) 

ref_totals   = {'full': totals_ref, 'CB': CB_totals_ref, 'CAA':CAA_totals_ref, 'wCAA':wCAA_totals_ref, 'eCAA':eCAA_totals_ref}
spm_totals   = {'full': totals_spm, 'CB': CB_totals_spm, 'CAA':CAA_totals_spm, 'wCAA':wCAA_totals_spm, 'eCAA':eCAA_totals_spm}
bio_totals   = {'full': totals_bio, 'CB': CB_totals_bio, 'CAA':CAA_totals_bio, 'wCAA':wCAA_totals_bio, 'eCAA':eCAA_totals_bio}

# +
# Component contribution calculations full water column:
CB_averages_ref,   CB_averages_spm,   CB_averages_bio   = pickle.load(open('calculations/Canada-Basin-averages-202110--fullcolumn.pickle','rb')) 
CAA_averages_ref,  CAA_averages_spm,  CAA_averages_bio  = pickle.load(open('calculations/CAA-averages-202110--fullcolumn.pickle' ,'rb')) 
wCAA_averages_ref, wCAA_averages_spm, wCAA_averages_bio = pickle.load(open('calculations/wCAA-averages-202110.pickle','rb'))  
eCAA_averages_ref, eCAA_averages_spm, eCAA_averages_bio = pickle.load(open('calculations/eCAA-averages-202110.pickle','rb')) 
averages_ref     , averages_spm     , averages_bio      = pickle.load(open('calculations/full-domain-averages-202110.pickle','rb')) 

ref_averages   = {'full': averages_ref, 'CB': CB_averages_ref, 'CAA':CAA_averages_ref, 'wCAA':wCAA_averages_ref, 'eCAA':eCAA_averages_ref}
spm_averages   = {'full': averages_spm, 'CB': CB_averages_spm, 'CAA':CAA_averages_spm, 'wCAA':wCAA_averages_spm, 'eCAA':eCAA_averages_spm}
bio_averages   = {'full': averages_bio, 'CB': CB_averages_bio, 'CAA':CAA_averages_bio, 'wCAA':wCAA_averages_bio, 'eCAA':eCAA_averages_bio}
# -

final_ref_totals, final_spm_totals, final_bio_totals, \
    final_ref_averages, final_spm_averages, final_bio_averages = pipeline(ref_totals, spm_totals, bio_totals, \
                                                                          ref_averages, spm_averages, bio_averages)

print_interannual_averages(final_ref_totals, final_spm_totals, final_bio_totals, \
                           final_ref_averages, final_spm_averages, final_bio_averages, \
                           location='full')

print_interannual_averages(final_ref_totals, final_spm_totals, final_bio_totals, \
                           final_ref_averages, final_spm_averages, final_bio_averages, \
                           location='CB')

# full water column: 
print_interannual_averages(final_ref_totals, final_spm_totals, final_bio_totals, \
                           final_ref_averages, final_spm_averages, final_bio_averages, \
                           location='CB')

# full water column:
print_interannual_averages(final_ref_totals, final_spm_totals, final_bio_totals, \
                           final_ref_averages, final_spm_averages, final_bio_averages, \
                           location='CAA')

print_interannual_averages(final_ref_totals, final_spm_totals, final_bio_totals, \
                           final_ref_averages, final_spm_averages, final_bio_averages, \
                           location='CAA')

print_interannual_averages(final_ref_totals, final_spm_totals, final_bio_totals, \
                           final_ref_averages, final_spm_averages, final_bio_averages, \
                           location='wCAA')

print_interannual_averages(final_ref_totals, final_spm_totals, final_bio_totals, \
                           final_ref_averages, final_spm_averages, final_bio_averages, \
                           location='eCAA')

# +
# biology runs only go to 2015, so later years have NaNs.
# -

print_interannual_averages_3D(final_ref_totals, final_spm_totals, final_bio_totals, \
                           final_ref_averages, final_spm_averages, final_bio_averages, \
                           location='CB')

print_interannual_averages_3D(final_ref_totals, final_spm_totals, final_bio_totals, \
                           final_ref_averages, final_spm_averages, final_bio_averages, \
                           location='CAA')



StarcoderdataJupyterScriptsDedupFiltered