# ---
# 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
# ---

# ### Define fxn to calculate regress and correl coeff maps

# - Calculate correl and regress coeff map btwn two (time,lat,lon) xr dataarrays
def gettemprcmap_loop(x3dnow,y3dnow,namebasenow):

    # - Initialize nan DataArray same size as x3dnow
    # --> new DataArray takes name of old though, so will need to rename
    # (Apparently you can't initialize empty or full DataArrays yet still: https://github.com/pydata/xarray/issues/277)
    rc2dnow = x3dnow.mean(dim='time') * np.nan
    int2dnow = x3dnow.mean(dim='time') * np.nan
    rcpval2dnow = x3dnow.mean(dim='time') * np.nan
    cc2dnow = x3dnow.mean(dim='time') * np.nan
    ccpval2dnow = x3dnow.mean(dim='time') * np.nan
    stderr2dnow = x3dnow.mean(dim='time') * np.nan

    # - Loop over each lat/lon combo
    for ilat in range(x3dnow.lat.size):
        for ilon in range(x3dnow.lon.size):
            xnow = x3dnow[:,ilat,ilon]
            ynow = y3dnow[:,ilat,ilon]
            # - Get rid of any times where xnow or ynow is nan or inf
            valsnow = (~xr.ufuncs.isnan(xnow))&(~xr.ufuncs.isnan(ynow))&(~xr.ufuncs.isinf(xnow))&(~xr.ufuncs.isinf(ynow))
            if valsnow.sum()>=2:
                #print([ilat,ilon])
                ccnow = stats.pearsonr(xnow[valsnow],ynow[valsnow]) 
                cc2dnow[ilat,ilon] = ccnow[0]
                ccpval2dnow[ilat,ilon] = ccnow[1]
                #[rc2dnow[ilat,ilon], int2dnow[ilat,ilon], cc2dnow[ilat,ilon], pval2dnow[ilat,ilon], 
                # stderr2dnow[ilat,ilon]] = stats.linregress(xnow[valsnow],ynow[valsnow]) 
                linregnow = stats.linregress(xnow[valsnow],ynow[valsnow]) 
                rc2dnow[ilat,ilon] = linregnow[0]
                int2dnow[ilat,ilon] = linregnow[1]
                rcpval2dnow[ilat,ilon] = linregnow[3]
                stderr2dnow[ilat,ilon] = linregnow[4]

    # - Rename
    rc2dnow.name = namebasenow+'_rc'
    int2dnow.name = namebasenow+'_int'
    rcpval2dnow.name = namebasenow+'_rc_pval'
    cc2dnow.name = namebasenow+'_cc'
    ccpval2dnow.name = namebasenow+'_cc_pval'
    stderr2dnow.name = namebasenow+'_stderr'
    
    return [rc2dnow,int2dnow,rcpval2dnow,cc2dnow,ccpval2dnow,stderr2dnow]


# ### Define fxn to generate Wilcoxon ranksum p-value maps

# - Calculate wilcoxon rank sum p-val map btwn two (time,lat,lon) xr dataarrays
def wrspvalmap_loop(x3dnow,y3dnow,namebasenow):

    # - Initialize nan DataArray same size as x3dnow
    # --> new DataArray takes name of old though, so will need to rename
    pval2dnow = xr.full_like(x3dnow.mean(dim='time'), np.nan)

    # - Loop over each lat/lon combo
    for ilat in range(x3dnow.lat.size):
        for ilon in range(x3dnow.lon.size):
            xnow = x3dnow[:,ilat,ilon]
            ynow = y3dnow[:,ilat,ilon]
            # - Get rid of any times where xnow or ynow is nan or inf
            xvalsnow = (~xr.ufuncs.isnan(xnow))&(~xr.ufuncs.isinf(xnow))
            yvalsnow = (~xr.ufuncs.isnan(ynow))&(~xr.ufuncs.isinf(ynow))
            if (xvalsnow.sum()>=1)&(yvalsnow.sum()>=1):
                #print([ilat,ilon])
                wrsnow = stats.ranksums(xnow[xvalsnow],ynow[yvalsnow]) 
                pval2dnow[ilat,ilon] = wrsnow[1]
    
    # - Rename
    pval2dnow.name = namebasenow+'_wrs_pval'
    
    return pval2dnow


# ### Define fxn to generate Kruskal-Wallis p-value maps (compares seasons)

# - Calculate kruskal-wallis p-val map btwn four (time,lat,lon) xr dataarrays
def kwpvalmap_loop(var1_3dnow,var2_3dnow,var3_3dnow,var4_3dnow,namebasenow):

    # - Initialize nan DataArray same size as var1_3dnow
    # --> new DataArray takes name of old though, so will need to rename
    pval2dnow = xr.full_like(var1_3dnow.mean(dim='time'), np.nan)

    # - Loop over each lat/lon combo
    for ilat in range(var1_3dnow.lat.size):
        for ilon in range(var1_3dnow.lon.size):
            x1now = var1_3dnow[:,ilat,ilon]
            x2now = var2_3dnow[:,ilat,ilon]
            x3now = var3_3dnow[:,ilat,ilon]
            x4now = var4_3dnow[:,ilat,ilon]
            # - Get rid of any times where xnow is nan or inf
            x1valsnow = (~xr.ufuncs.isnan(x1now))&(~xr.ufuncs.isinf(x1now))
            x2valsnow = (~xr.ufuncs.isnan(x2now))&(~xr.ufuncs.isinf(x2now))
            x3valsnow = (~xr.ufuncs.isnan(x3now))&(~xr.ufuncs.isinf(x3now))
            x4valsnow = (~xr.ufuncs.isnan(x4now))&(~xr.ufuncs.isinf(x4now))
            if (x1now[x1valsnow].sum()>=1)&(x2now[x2valsnow].sum()>=1)& \
               (x3now[x3valsnow].sum()>=1)&(x4now[x4valsnow].sum()>=1):
                #print([ilat,ilon])
                kwnow = stats.kruskal(x1now[x1valsnow],x2now[x2valsnow],
                                      x3now[x3valsnow],x4now[x4valsnow]) 
                pval2dnow[ilat,ilon] = kwnow[1]
    
    # - Rename
    pval2dnow.name = namebasenow+'_kw_pval'
    
    return pval2dnow


# ### Define fxn to control false discovery rate in 2D maps

# - Calculate minimum significant p-val threshold from false discovery rate (fdr)
def controlfdr2d(pval2dnow,alphafdr):
    pvalstack=pval2dnow.stack(x=['lat','lon'])
    pvalstack=pvalstack[pvalstack.notnull()]
    sortedpvalstack = pvalstack.sortby(pvalstack).values
    N = sortedpvalstack.size
    pfdrarr = alphafdr*np.arange(1,N+1)/N
    if np.sum((sortedpvalstack-pfdrarr)<=0)>0:
        pthreshfdr = sortedpvalstack[(pfdrarr-sortedpvalstack)>=0].max()
    else:
        pthreshfdr = 0
    return pthreshfdr


# ### Define fxn to find spots where p-value is below certain value

# - Find lat and lon where all p-values are small
def find_where_pval_small(pvalmap,alpha):
    # - Find spatial pts where all p-val are smaller than alpha
    pvalmap_small_nonnan = pvalmap.where(pvalmap<alpha)
    # - Get lat,lon where small p-vals are (https://stackoverflow.com/questions/40592630/get-coordinates-of-non-nan-values-of-xarray-dataset)
    pvalmap_small_nonnan_stacked = pvalmap_small_nonnan.stack(x=['lat','lon'])
    pvalmap_small = pvalmap_small_nonnan_stacked[pvalmap_small_nonnan_stacked.notnull()]
    return [pvalmap_small.lon,pvalmap_small.lat]


# ### Define fxn for boxplotting values over EEZs

def get_allenln_eezmask_enlnwrspval(varnow,eeznames,eezmask,dfeeznamesmask):
    varallnow, varennow, varlnnow = [], [], []
    enlnpvalnow = np.full(len(eeznames), np.nan) 
    for ieez in range(0,len(eeznames)):
        eeznumnow = dfeeznamesmask['numbers'][dfeeznamesmask['names']==eeznames[ieez]].values
        allnow = varnow.where(np.isin(eezmask,eeznumnow)).values
        allnow = allnow[~np.isnan(allnow)]
        ennow = varnow[onienln==1].where(np.isin(eezmask,eeznumnow)).values
        ennow = ennow[~np.isnan(ennow)]
        lnnow = varnow[onienln==-1].where(np.isin(eezmask,eeznumnow)).values
        lnnow = lnnow[~np.isnan(lnnow)]
        if len(ennow)>1 and len(lnnow)>1:
            _,enlnpvalnow[ieez] = stats.ranksums(ennow,lnnow)
        varallnow.append(allnow)
        varennow.append(ennow)
        varlnnow.append(lnnow)

    return varallnow, varennow, varlnnow, enlnpvalnow


def get_seas_eezmask_seaskwpval(varnow,eeznames,eezmask,dfeeznamesmask):
    varwinnow, varsprnow, varsumnow, varautnow = [], [], [], []
    seaspvalnow = np.full(len(eeznames), np.nan) 
    for ieez in range(0,len(eeznames)):
        eeznumnow = dfeeznamesmask['numbers'][dfeeznamesmask['names']==eeznames[ieez]].values
        winnow = varnow.sel(time=varnow['time.season']=='DJF'
                  ).where(np.isin(eezmask,eeznumnow)).values
        winnow = winnow[~np.isnan(winnow)]
        sprnow = varnow.sel(time=varnow['time.season']=='MAM'
                  ).where(np.isin(eezmask,eeznumnow)).values
        sprnow = sprnow[~np.isnan(sprnow)]
        sumnow = varnow.sel(time=varnow['time.season']=='JJA'
                  ).where(np.isin(eezmask,eeznumnow)).values
        sumnow = sumnow[~np.isnan(sumnow)]
        autnow = varnow.sel(time=varnow['time.season']=='SON'
                  ).where(np.isin(eezmask,eeznumnow)).values
        autnow = autnow[~np.isnan(autnow)]
        if len(winnow)>1 and len(sprnow)>1 and len(sumnow)>1 and len(autnow)>1:
            _,seaspvalnow[ieez] = stats.kruskal(winnow, sprnow, sumnow, autnow)
        varwinnow.append(winnow)
        varsprnow.append(sprnow)
        varsumnow.append(sumnow)
        varautnow.append(autnow)
    return varwinnow, varsprnow, varsumnow, varautnow, seaspvalnow


def set_box_color(bp, color):
    plt.setp(bp['boxes'], color=color)
    plt.setp(bp['whiskers'], color=color)
    plt.setp(bp['caps'], color=color)
    plt.setp(bp['medians'], color=color)


# ### Define fxn to control false discovery rate in 1D array

def controlfdr1d(pvals1d,alphafdr):
    pvals1d = pvals1d[~np.isnan(pvals1d)]
    sortedpvals = np.sort(pvals1d)
    N = sortedpvals.size
    pfdrarr = alphafdr*np.arange(1,N+1)/N
    if np.sum((sortedpvals-pfdrarr)<=0)>0:
        pthreshfdr = sortedpvals[(sortedpvals-pfdrarr)<=0].max()
    else:
        pthreshfdr = 0
    return pthreshfdr


# ### Define fxn for calculating pO2

# +
# function from <NAME>'s github:
# https://github.com/kallisons/pO2_conversion/function_pO2.py
# with some modifications for style and units

# pO2 Conversion:
# Oxygen concentration is converted to percent oxygen saturation using the equations from Garcia and Gordon (1992).
# The percent oxygen saturation is divided by 0.21 (the fractional atmospheric concentration of oxygen) to get pO2 in atmospheres (atm).
# pO2 is then corrected for the hydrostatic pressure at depth (Enns et al., 1965).
# The units for pO2 are converted to kilopascals (kPa), the SI Units for pressure.
# References:
# - <NAME>, <NAME> (1992) Oxygen solubility in seawater: Better fitting equations. Limnology and Oceanography, 37, 1307–1312.
# - Enns T, Scholander PF, Bradstreet ED (1965) Effect of hydrostatic pressure on gases dissolved in water. The Journal of Physical Chemistry, 69, 389–391.
 
# UNITS:
# o2 in umol/Kg
# temp in Celsius (= potential temperature, NOT in situ)
# sal in psu
# depth in m
# po2 returned in kPa

def calc_po2(o2, temp, sal, depth):
    """Computes po2 from o2 [umol/kg], potential temperature [Celsius], salinity [psu], depth [m]."""
    a_0 = 5.80871
    a_1 = 3.20291
    a_2 = 4.17887
    a_3 = 5.10006
    a_4 = -9.86643e-2
    a_5 = 3.80369
    b_0 = -7.01577e-3
    b_1 = -7.70028e-3
    b_2 =  -1.13864e-2
    b_3 = -9.51519e-3
    c_0 = -2.75915E-7

    tt = 298.15 - temp
    tk = 273.15 + temp
    ts = np.log(tt / tk)

    #correct for pressure at depth
    V = 32e-6 #partial molar volume of O2 (m3/mol)
    R = 8.31 #Gas constant [J/mol/K]
    db2Pa = 1e4 #convert pressure: decibar to Pascal
    atm2Pa = 1.01325e5 #convert pressure: atm to Pascal

    #calculate pressure in dB from depth in m
    pres = depth*(1.0076+depth*(2.3487e-6 - depth*1.2887e-11));

    #convert pressure from decibar to pascal
    dp = pres*db2Pa
    pCor = np.exp((V*dp)/(R*(temp+273.15)))

    o2_sat = np.exp(a_0 + a_1*ts + a_2*ts**2 + a_3*ts**3 + a_4*ts**4 + a_5*ts**5 + sal*(b_0 + b_1*ts + b_2*ts**2 + b_3*ts**3) + c_0*sal**2)
    
    o2_alpha = (o2_sat / 0.21)  #0.21 is atmospheric composition of O2
    kh = o2_alpha*pCor
    po2 = (o2 / kh)*101.32501  #convert po2 from atm to kPa

    return po2


# -

# ### Define fxns for performing + plotting quotient analysis

def quotient_analysis_full(ivfull,dvfull,ivname,dvname,nbins,nruns,plothists):
    # IV for indy var, DV for dep var
    iv = ivfull.where(dvfull.notnull())
    dv = dvfull.where(ivfull.notnull())
    
    ivminf = np.floor(10*iv.min())/10
    ivmaxc = np.ceil(10*iv.max())/10
    binedges = np.linspace(ivminf,ivmaxc,nbins+1)
    
    if plothists==1:
        fig,axes = plt.subplots(figsize=(7,5),nrows=2,ncols=2)
        ivfull.plot.hist(ax=axes[0][0],bins=20); axes[0][0].set_title('IV - all values');
        iv.plot.hist(ax=axes[0][1],bins=20); axes[0][1].set_title('IV - assoc w/ DV values');
        dvfull.plot.hist(ax=axes[1][0],bins=20); axes[1][0].set_title('DV - all values');
        dv.plot.hist(ax=axes[1][1],bins=20); axes[1][1].set_title('DV - assoc w/ IV values');
        fig.tight_layout()
    
    dsqa = xr.merge([iv, dv])
    dfqa = dsqa.to_dataframe()
    dfqa.rename(columns={ivname: 'IV', dvname: 'DV'}, inplace=True)
    dfqa.reset_index(inplace=True)
    dfqa = dfqa.dropna(subset=['IV','DV'], how='any')

    ivcounts, dvcounts, dvquot = quotient_analysis(dfqa,binedges)

    # Bernal et al. 2007: replace=FALSE in the code, but in the paper it says replace=TRUE
    # --> I think I'll go w/ what the paper says? though difference
    # between the two methods is small
    dfsimreplaceF = pd.DataFrame(); dfsimreplaceT = pd.DataFrame()
    for i in range(nruns):
        dfsimreplaceF=pd.concat([dfsimreplaceF,dfqa['DV'].sample(
            n=len(dfqa['DV']), replace=False).reset_index(
            drop=True)], axis=1).rename(columns={'DV':i})
        dfsimreplaceT=pd.concat([dfsimreplaceT,dfqa['DV'].sample(
            n=len(dfqa['DV']), replace=True).reset_index(
            drop=True)], axis=1).rename(columns={'DV':i})

    dfsimreplaceF = dfsimreplaceF.assign(IV=dfqa['IV'].values)
    dfsimreplaceT = dfsimreplaceT.assign(IV=dfqa['IV'].values)

    quotsimreplaceF = pd.DataFrame(); quotsimreplaceT = pd.DataFrame()
    for i in range(nruns):
        _,_,quotsimreplaceFtemp = quotient_analysis(
            dfsimreplaceF[[i,'IV']].rename(columns={i:'DV'}),binedges)
        quotsimreplaceF = pd.concat([quotsimreplaceF,quotsimreplaceFtemp], axis=1)
        _,_,quotsimreplaceTtemp = quotient_analysis(
            dfsimreplaceT[[i,'IV']].rename(columns={i:'DV'}),binedges)
        quotsimreplaceT = pd.concat([quotsimreplaceT,quotsimreplaceTtemp], axis=1)

    qlimsreplaceT = quotsimreplaceT.quantile([0.025, 0.975], axis=1)
    qlimsreplaceF = quotsimreplaceF.quantile([0.025, 0.975], axis=1)
    bincenters = (binedges[1:] + binedges[:-1])/2

    return dfqa['IV'], binedges, bincenters, ivcounts, dvcounts, dvquot, qlimsreplaceT, qlimsreplaceF


def quotient_analysis(dfqa,binedges):
    # IV for indy var, DV for dep var
    dfqa['IV_bin'] = pd.cut(dfqa['IV'],binedges)
    ivcounts = dfqa.groupby('IV_bin')['IV'].count()
    dvcounts = dfqa.groupby('IV_bin')['DV'].sum()
    ivfreq = ivcounts*100/ivcounts.sum()
    dvfreq = dvcounts*100/dvcounts.sum()
    dvquot = dvfreq/ivfreq
    return ivcounts, dvcounts, dvquot


def define_atp(row):
    if row['quot']<row['2pt5p']:
        return -1 # avoidance
    if row['quot']>=row['2pt5p'] and row['quot']<=row['97pt5p']:
        return 0 # tolerance
    if row['quot']>row['97pt5p']:
        return 1 # preference


def get_dfatp(qlimsreplaceT,dvquot,binedges):
    dfatp = pd.concat([qlimsreplaceT.transpose(),dvquot], axis=1)
    dfatp = dfatp.rename(columns={0:'quot', 0.025:'2pt5p', 0.975:'97pt5p'})
    dfatp['ATP'] = dfatp.apply(define_atp, axis=1)
    dfatp.reset_index(inplace=True)
    dfatp = pd.concat([dfatp,pd.Series(binedges[0:-1],name='lbinedges'),
                      pd.Series(binedges[1:],name='rbinedges')], axis=1)    
    return dfatp

StarcoderdataJupyterScriptsDedupFiltered