AirSeaFluxCode.py 28.4 KB
Newer Older
sbiri's avatar
sbiri committed
1 2 3 4 5
import numpy as np
import sys
import logging
from flux_subs import (kappa, CtoK, get_heights, cdn_calc, cd_calc, get_skin,
                       psim_calc, psit_calc, ctcq_calc, ctcqn_calc, get_gust,
6
                       gc, qsat_sea, qsat_air, visc_air, psit_26, psiu_26)
sbiri's avatar
sbiri committed
7 8


9 10 11 12
def AirSeaFluxCode(spd, T, SST, lat=None, hum=None, P=None,
                       hin=18, hout=10, Rl=None, Rs=None, jcool=None,
                       gust=None, meth="S80", qmeth="Buck2", tol=None, n=10,
                       out=0):
sbiri's avatar
sbiri committed
13 14 15 16 17 18 19 20 21 22 23 24
    """ Calculates momentum and heat fluxes using different parameterizations

    Parameters
    ----------
        spd : float
            relative wind speed in m/s (is assumed as magnitude difference
            between wind and surface current vectors)
        T : float
            air temperature in K (will convert if < 200)
        SST : float
            sea surface temperature in K (will convert if < 200)
        lat : float
25 26 27 28 29 30
            latitude (deg), default 45deg
        hum : float
            humidity input switch 2x1 [x, values] default is relative humidity
            x=0 : relative humidity in %
            x=1 : specific humidity (g/kg)
            x=2 : dew point temperature (K)
sbiri's avatar
sbiri committed
31
        P : float
32
            air pressure (hPa), default 1013hPa
sbiri's avatar
sbiri committed
33
        hin : float
34
            sensor heights in m (array 3x1 or 3xn), default 18m
sbiri's avatar
sbiri committed
35 36 37 38 39 40
        hout : float
            output height, default is 10m
        Rl : float
            downward longwave radiation (W/m^2)
        Rs : float
            downward shortwave radiation (W/m^2)
41
        jcool : int
42 43
            0 switch cool skin adjustment off, else 1
            default is 1
44
        gust : int
45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65
            3x1 [x, beta, zi] x=1 to include the effect of gustiness, else 0
            beta gustiness parameter, beta=1 for UA, beta=1.2 for COARE
            zi PBL height (m) 600 for COARE, 1000 for UA and ERA5, 800 default
            default for COARE [1, 1.2, 600]
            default for UA, ERA5 [1, 1, 1000]
            default else [1, 1.2, 800]
        meth : str
            "S80","S88","LP82","YT96","UA","LY04","C30","C35","C40","ERA5"
        qmeth : str
            is the saturation evaporation method to use amongst
            "HylandWexler","Hardy","Preining","Wexler","GoffGratch","CIMO",
            "MagnusTetens","Buck","Buck2","WMO","WMO2000","Sonntag","Bolton",
            "IAPWS","MurphyKoop"]
            default is Buck2
        tol : float
           4x1 or 7x1 [option, lim1-3 or lim1-6]
           option : 'flux' to set tolerance limits for fluxes only lim1-3
           option : 'ref' to set tolerance limits for height adjustment lim-1-3
           option : 'all' to set tolerance limits for both fluxes and height
                    adjustment lim1-6 ['all', 0.01, 0.01, 5e-05, 0.01, 1, 1]
           default is tol=['flux', 0.01, 1, 1]
sbiri's avatar
sbiri committed
66
        n : int
67 68 69 70
            number of iterations (defautl = 10)
        out : int
            set 0 to set points that have not converged to missing (default)
            set 1 to keep points
sbiri's avatar
sbiri committed
71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88
    Returns
    -------
        res : array that contains
                       1. momentum flux (W/m^2)
                       2. sensible heat (W/m^2)
                       3. latent heat (W/m^2)
                       4. Monin-Obhukov length (mb)
                       5. drag coefficient (cd)
                       6. neutral drag coefficient (cdn)
                       7. heat exhange coefficient (ct)
                       8. neutral heat exhange coefficient (ctn)
                       9. moisture exhange coefficient (cq)
                       10. neutral moisture exhange coefficient (cqn)
                       11. star virtual temperature (tsrv)
                       12. star temperature (tsr)
                       13. star humidity (qsr)
                       14. star velocity (usr)
                       15. momentum stability function (psim)
89 90 91 92 93 94 95 96 97
                       16. heat stability function (psit)
                       17. moisture stability function (psiq)
                       18. 10m neutral velocity (u10n)
                       19. 10m neutral temperature (t10n)
                       20. 10m neutral virtual temperature (tv10n)
                       21. 10m neutral specific humidity (q10n)
                       22. surface roughness length (zo)
                       23. heat roughness length (zot)
                       24. moisture roughness length (zoq)
98 99 100
                       25. velocity at reference height (uref)
                       26. temperature at reference height (tref)
                       27. specific humidity at reference height (qref)
101
                       28. number of iterations until convergence
sbiri's avatar
sbiri committed
102 103 104 105 106 107 108
        ind : int
            the indices in the matrix for the points that did not converge
            after the maximum number of iterations
    The code is based on bform.f and flux_calc.R modified by S. Biri
    """
    logging.basicConfig(filename='flux_calc.log',
                        format='%(asctime)s %(message)s',level=logging.INFO)
109 110 111
    ref_ht, tlapse = 10, 0.0098        # reference height, lapse rate
    h_in = get_heights(hin, len(spd))  # heights of input measurements/fields
    h_out = get_heights(hout, 1)       # desired height of output variables
sbiri's avatar
sbiri committed
112 113 114
    # if input values are nan break
    if (np.all(np.isnan(spd)) or np.all(np.isnan(T)) or np.all(np.isnan(SST))):
        sys.exit("input wind, T or SST is empty")
115 116 117 118 119 120 121
        logging.debug('all spd or T or SST input is nan')
    if (np.all(lat) == None):  # set latitude to 45deg if empty
        lat = 45*np.ones(spd.shape)
    elif ((np.all(lat) != None) and (np.size(lat) == 1)):
        lat = np.ones(spd.shape)*np.copy(lat)
    g = gc(lat, None)             # acceleration due to gravity
    if ((np.all(P) == None) and (meth == "C30" or meth == "C40")):
sbiri's avatar
sbiri committed
122
        P = np.ones(spd.shape)*1015  # if empty set to default for COARE3.0
123 124 125 126 127 128
    elif ((np.all(P) == None) or np.all(np.isnan(P))):
        P = np.ones(spd.shape)*1013
        logging.debug('input P is empty and set to 1013hPa')
    elif (((np.all(P) != None) or np.all(~np.isnan(P))) and np.size(P) == 1):
        P = np.ones(spd.shape)*np.copy(P)
    if ((np.all(Rl) == None or np.all(np.isnan(Rl))) and meth == "C30"):
sbiri's avatar
sbiri committed
129
        Rl = np.ones(spd.shape)*150    # set to default for COARE3.0
130 131
    elif (((np.all(Rl) == None or np.all(np.isnan(Rl))) and meth == "C35") or
          ((np.all(Rl) == None or np.all(np.isnan(Rl))) and meth == "C40")):
sbiri's avatar
sbiri committed
132
        Rl = np.ones(spd.shape)*370    # set to default for COARE3.5
133
    elif (np.all(Rl) == None or np.all(np.isnan(Rl))):
sbiri's avatar
sbiri committed
134
        Rl = np.ones(spd.shape)*370    # set to default for COARE3.5
135
    if ((np.all(Rs) == None or np.all(np.isnan(Rs))) and meth == "C30"):
sbiri's avatar
sbiri committed
136
        Rs = np.ones(spd.shape)*370  # set to default for COARE3.0
137
    elif (np.all(Rs) == None or np.all(np.isnan(Rs))):
sbiri's avatar
sbiri committed
138
        Rs = np.ones(spd.shape)*150  # set to default for COARE3.5
139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157
    if ((gust == None) and (meth == "C30" or meth == "C35" or meth == "C40")):
        gust = [1, 1.2, 600]
    elif ((gust == None) and (meth == "UA" or meth == "ERA5")):
        gust = [1, 1, 1000]
    elif (gust == None):
        gust = [1, 1.2, 800]
    if (tol == None):
        tol = ['flux', 0.01, 1, 1]
    if ((jcool == None) and (meth == "S80" or meth == "S88" or meth == "LP82"
                             or meth == "YT96")):
       jcool = 0
    elif ((jcool == None) and (meth == "UA" or meth == "LY04" or meth == "C30"
                              or meth == "C35" or meth == "C40"
                              or meth == "ERA5")):
       jcool = 1
    logging.info('method %s, inputs: h_in: %s | lat: %s | P: %s | Rl: %s |'
                 ' Rs: %s | gust: %s | jcool: %s', meth, h_in[:, 1],
                 np.nanmedian(lat), np.nanmedian(P), np.nanmedian(Rl),
                 np.nanmedian(Rs), gust, jcool)
sbiri's avatar
sbiri committed
158
    ####
159
    th = np.where(T < 200, (np.copy(T)+CtoK) *
sbiri's avatar
sbiri committed
160 161
                  np.power(1000/P,287.1/1004.67),
                  np.copy(T)*np.power(1000/P,287.1/1004.67))  # potential T
162 163 164
    Ta = np.where(T < 200, np.copy(T)+CtoK+tlapse*h_in[1],
                  np.copy(T)+tlapse*h_in[1])  # convert to Kelvin if needed
    sst = np.where(SST < 200, np.copy(SST)+CtoK, np.copy(SST))
165 166 167 168 169 170 171 172 173 174 175 176
    if (hum == None):
        RH = np.ones(spd.shape)*80
        qsea = qsat_sea(sst, P, qmeth)/1000     # surface water q (g/kg)
        qair = qsat_air(T, P, RH, qmeth)/1000   # q of air (g/kg)
    elif (hum[0] == 0):
        RH = hum[1]
        if (np.all(RH) < 1):
            sys.exit("input relative humidity units should be \%")
        qsea = qsat_sea(sst, P, qmeth)/1000    # surface water q (g/kg)
        qair = qsat_air(T, P, RH, qmeth)/1000  # q of air (g/kg)
    elif (hum[0] == 1):
        qair = hum[1]
177
        qsea = qsat_sea(sst, P, qmeth)/1000  # surface water q (g/kg)
178 179 180 181 182 183 184 185 186 187 188
    elif (hum[0] == 2):
        Td = hum[1] # dew point temperature (K)
        Td = np.where(Td < 200, np.copy(Td)+CtoK, np.copy(Td))
        T = np.where(T < 200, np.copy(T)+CtoK, np.copy(T))
        esd = 611.21*np.exp(17.502*((Td-273.16)/(Td-32.19)))
        es = 611.21*np.exp(17.502*((T-273.16)/(T-32.19)))
        RH = 100*esd/es
        qair = qsat_air(T, P, RH, qmeth)/1000  # q of air (g/kg)
        qsea = qsat_sea(sst, P, qmeth)/1000    # surface water q (g/kg)
    logging.info('method %s and q method %s | qsea:%s, qair:%s', meth, qmeth,
                  np.nanmedian(qsea), np.nanmedian(qair))
sbiri's avatar
sbiri committed
189 190 191 192 193 194 195 196 197 198 199
    if (np.all(np.isnan(qsea)) or np.all(np.isnan(qair))):
        print("qsea and qair cannot be nan")
    # first guesses
    dt = Ta - sst
    dq = qair - qsea
    t10n, q10n = np.copy(Ta), np.copy(qair)
    tv10n = t10n*(1 + 0.61*q10n)
    #  Zeng et al. 1998
    tv=th*(1.+0.61*qair)   # virtual potential T
    dtv=dt*(1.+0.61*qair)+0.61*th*dq
    # ------------
200
    rho = P*100/(287.1*tv10n)
sbiri's avatar
sbiri committed
201 202 203
    lv = (2.501-0.00237*SST)*1e6
    cp = 1004.67*(1 + 0.00084*qsea)
    u10n = np.copy(spd)
204
    monob = -100*np.ones(spd.shape)  # Monin-Obukhov length
sbiri's avatar
sbiri committed
205
    cdn = cdn_calc(u10n, Ta, None, lat, meth)
206 207 208 209 210 211 212 213
    ctn, ct, cqn, cq = (np.zeros(spd.shape)*np.nan, np.zeros(spd.shape)*np.nan,
                        np.zeros(spd.shape)*np.nan, np.zeros(spd.shape)*np.nan)
    psim, psit, psiq = (np.zeros(spd.shape), np.zeros(spd.shape),
                        np.zeros(spd.shape))
    cd = cd_calc(cdn, h_in[0], ref_ht, psim)
    tsr, tsrv = np.zeros(spd.shape), np.zeros(spd.shape)
    qsr = np.zeros(spd.shape)
    # jcool parameters
214 215
    tkt = 0.001*np.ones(T.shape)
    Rnl = 0.97*(5.67e-8*np.power(sst-0.3*jcool+CtoK, 4)-Rl)
sbiri's avatar
sbiri committed
216
    dter = np.ones(T.shape)*0.3
217
    dqer = dter*0.622*lv*qsea/(287.1*np.power(sst, 2))
218
    if (gust[0] == 1 and meth == "UA"):
219 220
        wind = np.where(dtv >= 0, np.where(spd > 0.1, spd, 0.1),
                        np.sqrt(np.power(np.copy(spd), 2)+np.power(0.5, 2)))
221
    elif (gust[0] == 1):
222
        wind = np.sqrt(np.power(np.copy(spd), 2)+np.power(0.5, 2))
223
    elif (gust[0] == 0):
224 225 226
        wind = np.copy(spd)

    if (meth == "UA"):
sbiri's avatar
sbiri committed
227 228 229
        usr = 0.06
        for i in range(5):
            zo = 0.013*np.power(usr,2)/g+0.11*visc_air(Ta)/usr
230 231 232
            usr=kappa*wind/np.log(h_in[0]/zo)
        Rb = g*h_in[0]*dtv/(tv*wind**2)
        zol = np.where(Rb >= 0, Rb*np.log(h_in[0]/zo) /
233 234
                        (1-5*np.where(Rb < 0.19, Rb, 0.19)),
                        Rb*np.log(h_in[0]/zo))
235
        monob = h_in[0]/zol
sbiri's avatar
sbiri committed
236 237 238
        zo = 0.013*np.power(usr, 2)/g + 0.11*visc_air(Ta)/usr
        zot = zo/np.exp(2.67*np.power(usr*zo/visc_air(Ta), 0.25)-2.57)
        zoq = zot
239
        logging.info('method %s | wind:%s, usr:%s, '
240 241 242 243
                      'zo:%s, zot:%s, zoq:%s, Rb:%s, monob:%s', meth,
                      np.nanmedian(wind), np.nanmedian(usr), np.nanmedian(zo),
                      np.nanmedian(zot), np.nanmedian(zoq), np.nanmedian(Rb),
                      np.nanmedian(monob))
sbiri's avatar
sbiri committed
244
    elif (meth == "ERA5"):
245
        usr = np.sqrt(cd*np.power(wind, 2))
246
        Rb = ((g*h_in[0]*((2*dt)/(Ta+sst-g*h_in[0]) +
sbiri's avatar
sbiri committed
247 248 249
                0.61*dq))/np.power(wind, 2))
        zo = 0.11*visc_air(Ta)/usr+0.018*np.power(usr, 2)/g
        zot = 0.40*visc_air(Ta)/usr
250
        zoq = 0.62*visc_air(Ta)/usr
251
        zol = (Rb*((np.log((h_in[0]+zo)/zo)-psim_calc((h_in[0]+zo) /
252 253 254 255
                monob, meth)+psim_calc(zo/monob, meth)) /
                (np.log((h_in[0]+zo)/zot) -
                psit_calc((h_in[0]+zo)/monob, meth) +
                psit_calc(zot/monob, meth))))
256
        monob = h_in[0]/zol
257
        logging.info('method %s | wind:%s, usr:%s, '
258 259 260
                      'zo:%s, zot:%s, Rb:%s, monob:%s', meth,
                      np.nanmedian(wind), np.nanmedian(usr), np.nanmedian(zo),
                      np.nanmedian(zot), np.nanmedian(Rb), np.nanmedian(monob))
sbiri's avatar
sbiri committed
261
    elif (meth == "C30" or meth == "C35" or meth == "C40"):
262
        usr = np.sqrt(cd*np.power(wind, 2))
sbiri's avatar
sbiri committed
263 264
        a = 0.011*np.ones(T.shape)
        a = np.where(wind > 10, 0.011+(wind-10)/(18-10)*(0.018-0.011),
265
                          np.where(wind > 18, 0.018, a))
sbiri's avatar
sbiri committed
266 267 268
        zo = a*np.power(usr, 2)/g+0.11*visc_air(T)/usr
        rr = zo*usr/visc_air(T)
        zoq = np.minimum(5e-5/np.power(rr, 0.6), 1.15e-4)
269
        zot = np.copy(zoq)
270 271
        Rb = g*h_in[0]*dtv/((T+CtoK)*np.power(wind, 2))
        zol =  (Rb*((np.log((h_in[0]+zo)/zo)-psim_calc((h_in[0]+zo) /
272 273 274 275
                monob, meth)+psim_calc(zo/monob, meth)) /
                (np.log((h_in[0]+zo)/zot) -
                psit_calc((h_in[0]+zo)/monob, meth) +
                psit_calc(zot/monob, meth))))
276
        monob = h_in[0]/zol
277 278
        logging.info('method %s | wind:%s, usr:%s, '
                     'zo:%s, zot:%s, Rb:%s, monob:%s', meth,
279 280
                     np.nanmedian(wind), np.nanmedian(usr), np.nanmedian(zo),
                     np.nanmedian(zot), np.nanmedian(Rb), np.nanmedian(monob))
sbiri's avatar
sbiri committed
281
    else:
282 283
        zo = 0.0001*np.ones(spd.shape)
        zot, zoq = 0.0001*np.ones(spd.shape), 0.0001*np.ones(spd.shape)
284 285
        usr = np.sqrt(cd*np.power(wind, 2))
        logging.info('method %s | wind:%s, usr:%s, '
286 287 288
                     'zo:%s, zot:%s, monob:%s', meth,
                     np.nanmedian(wind), np.nanmedian(usr), np.nanmedian(zo),
                     np.nanmedian(zot), np.nanmedian(monob))
289
    tsr = (dt+dter*jcool)*kappa/(np.log(hin[1]/zot) -
290
                                 psit_calc(h_in[1]/monob, meth))
291
    qsr = (dq+dqer*jcool)*kappa/(np.log(hin[2]/zoq) -
292
                                 psit_calc(h_in[2]/monob, meth))
293 294
    it, ind = 0, np.where(spd > 0)
    ii, itera = True, np.zeros(spd.shape)*np.nan
295 296 297
    tau = 0.05*np.ones(spd.shape)
    sensible = 5*np.ones(spd.shape)
    latent = 65*np.ones(spd.shape)
sbiri's avatar
sbiri committed
298 299 300 301
    while np.any(ii):
        it += 1
        if it > n:
            break
302 303 304 305 306 307 308
        if (tol[0] == 'flux'):
            old = np.array([np.copy(tau), np.copy(sensible), np.copy(latent)])
        elif (tol[0] == 'ref'):
            old = np.array([np.copy(u10n), np.copy(t10n), np.copy(q10n)])
        elif (tol[0] == 'all'):
            old = np.array([np.copy(u10n), np.copy(t10n), np.copy(q10n),
                            np.copy(tau), np.copy(sensible), np.copy(latent)])
sbiri's avatar
sbiri committed
309 310
        cdn[ind] = cdn_calc(u10n[ind], Ta[ind], None, lat[ind], meth)
        if (np.all(np.isnan(cdn))):
311
            break
sbiri's avatar
sbiri committed
312 313
            logging.info('break %s at iteration %s cdn<0', meth, it)
        zo[ind] = ref_ht/np.exp(kappa/np.sqrt(cdn[ind]))
314 315 316
        psim[ind] = psim_calc(h_in[0, ind]/monob[ind], meth)
        cd[ind] = cd_calc(cdn[ind], h_in[0, ind], ref_ht, psim[ind])
        ctn[ind], cqn[ind] = ctcqn_calc(h_in[1, ind]/monob[ind], cdn[ind],
sbiri's avatar
sbiri committed
317
                                        u10n[ind], zo[ind], Ta[ind], meth)
318 319 320 321
        zot[ind] = ref_ht/(np.exp(np.power(kappa, 2) /
                           (ctn[ind]*np.log(ref_ht/zo[ind]))))
        zoq[ind] = ref_ht/(np.exp(np.power(kappa, 2) /
                           (cqn[ind]*np.log(ref_ht/zo[ind]))))
322 323
        psit[ind] = psit_calc(h_in[1, ind]/monob[ind], meth)
        psiq[ind] = psit_calc(h_in[2, ind]/monob[ind], meth)
sbiri's avatar
sbiri committed
324
        ct[ind], cq[ind] = ctcq_calc(cdn[ind], cd[ind], ctn[ind], cqn[ind],
325 326
                                      h_in[1, ind], h_in[2, ind], ref_ht,
                                      psit[ind], psiq[ind])
sbiri's avatar
sbiri committed
327
        if (meth == "UA"):
328 329
            usr[ind] = np.where(h_in[0, ind]/monob[ind] < -1.574,
                                kappa*wind[ind] /
sbiri's avatar
sbiri committed
330 331 332
                                (np.log(-1.574*monob[ind]/zo[ind]) -
                                psim_calc(-1.574, meth) +
                                psim_calc(zo[ind]/monob[ind], meth) +
333
                                1.14*(np.power(-h_in[0, ind]/monob[ind], 1/3) -
334
                                np.power(1.574, 1/3))),
335 336 337 338
                                np.where((h_in[0, ind]/monob[ind] > -1.574) &
                                (h_in[0, ind]/monob[ind] < 0),
                                kappa*wind[ind]/(np.log(h_in[0, ind]/zo[ind]) -
                                psim_calc(h_in[0, ind]/monob[ind], meth) +
sbiri's avatar
sbiri committed
339
                                psim_calc(zo[ind]/monob[ind], meth)),
340 341 342 343 344 345
                                np.where((h_in[0, ind]/monob[ind] > 0) &
                                (h_in[0, ind]/monob[ind] < 1),
                                kappa*wind[ind]/(np.log(h_in[0, ind]/zo[ind]) +
                                5*h_in[0, ind]/monob[ind]-5*zo[ind] /
                                monob[ind]), kappa*wind[ind] /
                                (np.log(monob[ind]/zo[ind]) +
sbiri's avatar
sbiri committed
346
                                5-5*zo[ind]/monob[ind] +
347 348
                                5*np.log(h_in[0, ind]/monob[ind]) +
                                h_in[0, ind]/monob[ind]-1))))
sbiri's avatar
sbiri committed
349
                                # Zeng et al. 1998 (7-10)
350 351
            tsr[ind] = np.where(h_in[1, ind]/monob[ind] < -0.465,
                                kappa*(dt[ind] +
352
                                dter[ind]*jcool) /
sbiri's avatar
sbiri committed
353 354
                                (np.log((-0.465*monob[ind])/zot[ind]) -
                                psit_calc(-0.465, meth)+0.8 *
355
                                (np.power(0.465, -1/3) -
356 357 358
                                np.power(-h_in[1, ind]/monob[ind], -1/3))),
                                np.where((h_in[1, ind]/monob[ind] > -0.465) &
                                (h_in[1, ind]/monob[ind] < 0),
359
                                kappa*(dt[ind]+dter[ind]*jcool) /
360 361
                                (np.log(h_in[1, ind]/zot[ind]) -
                                psit_calc(h_in[1, ind]/monob[ind], meth) +
sbiri's avatar
sbiri committed
362
                                psit_calc(zot[ind]/monob[ind], meth)),
363 364
                                np.where((h_in[1, ind]/monob[ind] > 0) &
                                (h_in[1, ind]/monob[ind] < 1),
365
                                kappa*(dt[ind]+dter[ind]*jcool) /
366 367 368
                                (np.log(h_in[1, ind]/zot[ind]) +
                                5*h_in[1, ind]/monob[ind]-5*zot[ind] /
                                monob[ind]),
369 370 371
                                kappa*(dt[ind]+dter[ind]*jcool) /
                                (np.log(monob[ind]/zot[ind])+5 -
                                5*zot[ind]/monob[ind] +
372 373
                                5*np.log(h_in[1, ind]/monob[ind]) +
                                h_in[1, ind]/monob[ind]-1))))
sbiri's avatar
sbiri committed
374
                                # Zeng et al. 1998 (11-14)
375 376
            qsr[ind] = np.where(h_in[2, ind]/monob[ind] < -0.465,
                                kappa*(dq[ind] +
377
                                dqer[ind]*jcool) /
sbiri's avatar
sbiri committed
378 379 380
                                (np.log((-0.465*monob[ind])/zoq[ind]) -
                                psit_calc(-0.465, meth) +
                                psit_calc(zoq[ind]/monob[ind], meth) +
381
                                0.8*(np.power(0.465, -1/3) -
382 383 384
                                np.power(-h_in[2, ind]/monob[ind], -1/3))),
                                np.where((h_in[2, ind]/monob[ind] > -0.465) &
                                (h_in[2, ind]/monob[ind] < 0),
385
                                kappa*(dq[ind]+dqer[ind]*jcool) /
386 387
                                (np.log(h_in[1, ind]/zot[ind]) -
                                psit_calc(h_in[2, ind]/monob[ind], meth) +
sbiri's avatar
sbiri committed
388
                                psit_calc(zoq[ind]/monob[ind], meth)),
389 390
                                np.where((h_in[2, ind]/monob[ind] > 0) &
                                (h_in[2, ind]/monob[ind]<1), kappa*(dq[ind] +
391
                                dqer[ind]*jcool) /
392 393
                                (np.log(h_in[1, ind]/zoq[ind]) +
                                5*h_in[2, ind]/monob[ind]-5*zoq[ind]/monob[ind]),
394 395
                                kappa*(dq[ind]+dqer[ind]*jcool) /
                                (np.log(monob[ind]/zoq[ind])+5 -
sbiri's avatar
sbiri committed
396
                                5*zoq[ind]/monob[ind] +
397 398
                                5*np.log(h_in[2, ind]/monob[ind]) +
                                h_in[2, ind]/monob[ind]-1))))
sbiri's avatar
sbiri committed
399
        elif (meth == "C30" or meth == "C35" or meth == "C40"):
400 401
            usr[ind] = (wind[ind]*kappa/(np.log(h_in[0, ind]/zo[ind]) -
                        psiu_26(h_in[0, ind]/monob[ind], meth)))
sbiri's avatar
sbiri committed
402
            logging.info('method %s | dter = %s | Rnl = %s '
403 404 405 406
                          '| usr = %s | tsr = %s | qsr = %s', meth,
                          np.nanmedian(dter), np.nanmedian(Rnl),
                          np.nanmedian(usr), np.nanmedian(tsr),
                          np.nanmedian(qsr))
407 408 409 410
            qsr[ind] = ((dq[ind]+dqer[ind]*jcool)*(kappa/(np.log(hin[2, ind] /
                        zoq[ind])-psit_26(hin[2, ind]/monob[ind]))))
            tsr[ind] = ((dt[ind]+dter[ind]*jcool)*(kappa/(np.log(hin[1, ind] /
                        zot[ind])-psit_26(hin[1, ind]/monob[ind]))))
sbiri's avatar
sbiri committed
411
        else:
412 413
            usr[ind] = (wind[ind]*kappa/(np.log(h_in[0, ind]/zo[ind]) -
                        psim_calc(h_in[0, ind]/monob[ind], meth)))
414 415
            tsr[ind] = ct[ind]*wind[ind]*(dt[ind]+dter[ind]*jcool)/usr[ind]
            qsr[ind] = cq[ind]*wind[ind]*(dq[ind]+dqer[ind]*jcool)/usr[ind]
416 417 418 419 420 421 422 423 424 425 426 427
        if (jcool == 1):
            dter[ind], dqer[ind], tkt[ind] = get_skin(sst[ind], qsea[ind],
                                                      rho[ind], Rl[ind],
                                                      Rs[ind], Rnl[ind],
                                                      cp[ind], lv[ind],
                                                      np.copy(tkt[ind]),
                                                      usr[ind], tsr[ind],
                                                      qsr[ind], lat[ind])
        else:
           dter[ind] = np.zeros(sst[ind].shape)
           dqer[ind] = np.zeros(sst[ind].shape)
           tkt[ind] = np.zeros(sst[ind].shape)
428
        Rnl[ind] = 0.97*(5.67e-8*np.power(SST[ind] -
429 430 431 432 433
                          dter[ind]*jcool+CtoK, 4)-Rl[ind])
        t10n[ind] = (Ta[ind] -
                     tsr[ind]/kappa*(np.log(h_in[1, ind]/ref_ht)-psit[ind]))
        q10n[ind] = (qair[ind] -
                     qsr[ind]/kappa*(np.log(h_in[2, ind]/ref_ht)-psiq[ind]))
sbiri's avatar
sbiri committed
434 435 436 437 438 439
        tv10n[ind] = t10n[ind]*(1+0.61*q10n[ind])
        if (meth == "UA"):
            tsrv[ind] = tsr[ind]*(1.+0.61*qair[ind])+0.61*th[ind]*qsr[ind]
            monob[ind] = (tv[ind]*np.power(usr[ind], 2))/(kappa*g[ind]*tsrv[ind])
        elif (meth == "C30" or meth == "C35" or meth == "C40"):
            tsrv[ind] = tsr[ind]+0.61*(T[ind]+CtoK)*qsr[ind]
440
            zol[ind] = (kappa*g[ind]*h_in[0, ind]/(T[ind]+CtoK)*(tsr[ind] +
sbiri's avatar
sbiri committed
441
                        0.61*(T[ind]+CtoK)*qsr[ind])/np.power(usr[ind], 2))
442
            monob[ind] = h_in[0, ind]/zol[ind]
sbiri's avatar
sbiri committed
443 444
        elif (meth == "ERA5"):
            tsrv[ind] = tsr[ind]+0.61*t10n[ind]*qsr[ind]
445
            Rb[ind] = ((g[ind]*h_in[0, ind]*((2*dt[ind])/(Ta[ind] +
446 447
                        sst[ind]-g[ind]*h_in[0, ind])+0.61*dq[ind])) /
                        np.power(wind[ind], 2))
sbiri's avatar
sbiri committed
448 449 450
            zo[ind] = (0.11*visc_air(Ta[ind])/usr[ind]+0.018 *
                       np.power(usr[ind], 2)/g[ind])
            zot[ind] = 0.40*visc_air(Ta[ind])/usr[ind]
451
            zoq[ind] = 0.62*visc_air(Ta[ind])/usr[ind]
452 453
            zol[ind] = (Rb[ind]*((np.log((h_in[0, ind]+zo[ind])/zo[ind]) -
                        psim_calc((h_in[0, ind]+zo[ind])/monob[ind], meth) +
sbiri's avatar
sbiri committed
454
                        psim_calc(zo[ind]/monob[ind], meth)) /
455 456
                        (np.log((h_in[0, ind]+zo[ind])/zot[ind]) -
                        psit_calc((h_in[0, ind]+zo[ind])/monob[ind], meth) +
sbiri's avatar
sbiri committed
457
                        psit_calc(zot[ind]/monob[ind], meth))))
458
            monob[ind] = h_in[0, ind]/zol[ind]
sbiri's avatar
sbiri committed
459 460 461
        else:
            tsrv[ind] = tsr[ind]+0.61*t10n[ind]*qsr[ind]
            monob[ind] = (tv10n[ind]*usr[ind]**2)/(g[ind]*kappa*tsrv[ind])
462 463 464
        psim[ind] = psim_calc(h_in[0, ind]/monob[ind], meth)
        psit[ind] = psit_calc(h_in[1, ind]/monob[ind], meth)
        psiq[ind] = psit_calc(h_in[2, ind]/monob[ind], meth)
465
        if (gust[0] == 1 and meth == "UA"):
sbiri's avatar
sbiri committed
466
            wind[ind] = np.where(dtv[ind] >= 0, np.where(spd[ind] > 0.1,
467 468 469 470 471 472 473
                                  spd[ind], 0.1),
                                  np.sqrt(np.power(np.copy(spd[ind]), 2) +
                                  np.power(get_gust(gust[1], tv[ind], usr[ind],
                                  tsrv[ind], gust[2], lat[ind]), 2)))
                                  # Zeng et al. 1998 (20)
        elif (gust[0] == 1 and (meth == "C30" or meth == "C35" or
                                meth == "C40")):
sbiri's avatar
sbiri committed
474
            wind[ind] = np.sqrt(np.power(np.copy(spd[ind]), 2) +
475 476 477
                                np.power(get_gust(gust[1], Ta[ind], usr[ind],
                                tsrv[ind], gust[2], lat[ind]), 2))
        elif (gust[0] == 1):
478
            wind[ind] = np.sqrt(np.power(np.copy(spd[ind]), 2) +
479 480 481
                                np.power(get_gust(gust[1], Ta[ind], usr[ind],
                                tsrv[ind], gust[2], lat[ind]), 2))
        elif (gust[0] == 0):
482
            wind[ind] = np.copy(spd[ind])
483 484 485
        u10n[ind] = wind[ind]-usr[ind]/kappa*(np.log(h_in[0, ind]/10) -
                                              psim[ind])
        u10n[u10n < 0] = np.nan
sbiri's avatar
sbiri committed
486
        itera[ind] = np.ones(1)*it
487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510
        sensible = -rho*cp*usr*tsr
        latent = -rho*lv*usr*qsr
        if (gust[0] == 1):
            tau = rho*np.power(usr, 2)*(spd/wind)
        elif (gust[0] == 0):
            tau = rho*np.power(usr, 2)
        if (tol[0] == 'flux'):
            new = np.array([np.copy(tau), np.copy(sensible), np.copy(latent)])
        elif (tol[0] == 'ref'):
            new = np.array([np.copy(u10n), np.copy(t10n), np.copy(q10n)])
        elif (tol[0] == 'all'):
            new = np.array([np.copy(u10n), np.copy(t10n), np.copy(q10n),
                            np.copy(tau), np.copy(sensible), np.copy(latent)])
        d = np.fabs(new-old)
        if (tol[0] == 'flux'):
            ind = np.where((d[0, :] > tol[1])+(d[1, :] > tol[2]) +
                           (d[2, :] > tol[3]))
        elif (tol[0] == 'ref'):
            ind = np.where((d[0, :] > tol[1])+(d[1, :] > tol[2]) +
                           (d[2, :] > tol[3]))
        elif (tol[0] == 'all'):
            ind = np.where((d[0, :] > tol[1])+(d[1, :] > tol[2]) +
                           (d[2, :] > tol[3])+(d[3, :] > tol[4]) +
                           (d[4, :] > tol[5])+(d[5, :] > tol[6]))
511 512
        if (ind[0].size == 0):
            ii = False
sbiri's avatar
sbiri committed
513
        else:
514 515
            ii = True
    logging.info('method %s | # of iterations:%s', meth, it)
sbiri's avatar
sbiri committed
516
    logging.info('method %s | # of points that did not converge :%s', meth,
517
                  ind[0].size)
sbiri's avatar
sbiri committed
518
    # calculate output parameters
519
    rho = (0.34838*P)/(tv10n)
sbiri's avatar
sbiri committed
520
    t10n = t10n-(273.16+tlapse*ref_ht)
521 522 523 524 525 526 527 528 529 530
    # zo = ref_ht/np.exp(kappa/cdn**0.5)
    # zot = ref_ht/(np.exp(kappa**2/(ctn*np.log(ref_ht/zo))))
    # zoq = ref_ht/(np.exp(kappa**2/(cqn*np.log(ref_ht/zo))))
    uref = (spd-usr/kappa*(np.log(h_in[0]/h_out[0])-psim +
            psim_calc(h_out[0]/monob, meth)))
    tref = (Ta-tsr/kappa*(np.log(h_in[1]/h_out[1])-psit +
            psit_calc(h_out[0]/monob, meth)))
    tref = tref-(273.16+tlapse*h_out[1])
    qref = (qair-qsr/kappa*(np.log(h_in[2]/h_out[2]) -
            psit+psit_calc(h_out[2]/monob, meth)))
531
    res = np.zeros((28, len(spd)))
sbiri's avatar
sbiri committed
532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547
    res[0][:] = tau
    res[1][:] = sensible
    res[2][:] = latent
    res[3][:] = monob
    res[4][:] = cd
    res[5][:] = cdn
    res[6][:] = ct
    res[7][:] = ctn
    res[8][:] = cq
    res[9][:] = cqn
    res[10][:] = tsrv
    res[11][:] = tsr
    res[12][:] = qsr
    res[13][:] = usr
    res[14][:] = psim
    res[15][:] = psit
548 549 550 551 552 553 554 555
    res[16][:] = psiq
    res[17][:] = u10n
    res[18][:] = t10n
    res[19][:] = tv10n
    res[20][:] = q10n
    res[21][:] = zo
    res[22][:] = zot
    res[23][:] = zoq
556 557 558
    res[24][:] = uref
    res[25][:] = tref
    res[26][:] = qref
559
    res[27][:] = itera
560 561 562 563
    if (out == 0):
        res[:, ind] = np.nan
    # set missing values where data have non acceptable values
    res = np.where(spd < 0, np.nan, res)
sbiri's avatar
sbiri committed
564
    return res, ind
565