AirSeaFluxCode.py 29.2 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
def AirSeaFluxCode(spd, T, SST, lat=None, hum=None, P=None,
10 11 12
                   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
            latitude (deg), default 45deg
        hum : float
            humidity input switch 2x1 [x, values] default is relative humidity
28 29 30
            x='rh' : relative humidity in %
            x='q' : specific humidity (g/kg)
            x='Td' : 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
    # if input values are nan break
113 114 115 116 117 118 119
    if meth not in ["S80", "S88", "LP82", "YT96", "UA", "LY04", "C30", "C35",
                    "C40","ERA5"]:
        sys.exit("unknown method")
    if qmeth not in ["HylandWexler", "Hardy", "Preining", "Wexler", "CIMO",
                     "GoffGratch", "MagnusTetens", "Buck", "Buck2", "WMO",
                     "WMO2000", "Sonntag", "Bolton", "IAPWS", "MurphyKoop"]:
        sys.exit("unknown q-method")
sbiri's avatar
sbiri committed
120 121
    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")
122
        logging.debug('all spd or T or SST input is nan')
123
    if (np.all(lat == None)):  # set latitude to 45deg if empty
124
        lat = 45*np.ones(spd.shape)
125
    elif ((np.all(lat != None)) and (np.size(lat) == 1)):
126 127
        lat = np.ones(spd.shape)*np.copy(lat)
    g = gc(lat, None)             # acceleration due to gravity
128
    if ((np.all(P == None)) and (meth == "C30" or meth == "C40")):
sbiri's avatar
sbiri committed
129
        P = np.ones(spd.shape)*1015  # if empty set to default for COARE3.0
130
    elif ((np.all(P == None)) or np.all(np.isnan(P))):
131 132
        P = np.ones(spd.shape)*1013
        logging.debug('input P is empty and set to 1013hPa')
133
    elif (((np.all(P != None)) or np.all(~np.isnan(P))) and np.size(P) == 1):
134
        P = np.ones(spd.shape)*np.copy(P)
135
    if ((np.all(Rl == None) or np.all(np.isnan(Rl))) and meth == "C30"):
sbiri's avatar
sbiri committed
136
        Rl = np.ones(spd.shape)*150    # set to default for COARE3.0
137 138
    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
139
        Rl = np.ones(spd.shape)*370    # set to default for COARE3.5
140
    elif (np.all(Rl == None) or np.all(np.isnan(Rl))):
sbiri's avatar
sbiri committed
141
        Rl = np.ones(spd.shape)*370    # set to default for COARE3.5
142
    if ((np.all(Rs == None) or np.all(np.isnan(Rs))) and meth == "C30"):
sbiri's avatar
sbiri committed
143
        Rs = np.ones(spd.shape)*370  # set to default for COARE3.0
144
    elif (np.all(Rs == None) or np.all(np.isnan(Rs))):
sbiri's avatar
sbiri committed
145
        Rs = np.ones(spd.shape)*150  # set to default for COARE3.5
146 147 148 149 150 151
    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]
152 153
    elif (np.size(gust) < 3):
        sys.exit("gust input must be a 3x1 array")
154 155
    if (tol == None):
        tol = ['flux', 0.01, 1, 1]
156 157
    elif (tol[0] not in ['flux', 'ref', 'all']):
        sys.exit("unknown tolerance input")
158 159 160 161
    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"
162 163
                               or meth == "C35" or meth == "C40"
                               or meth == "ERA5")):
164
       jcool = 1
165 166
    logging.info('method %s, inputs: lat: %s | P: %s | Rl: %s |'
                 ' Rs: %s | gust: %s | jcool: %s', meth,
167 168
                 np.nanmedian(lat), np.nanmedian(P), np.nanmedian(Rl),
                 np.nanmedian(Rs), gust, jcool)
sbiri's avatar
sbiri committed
169
    ####
170
    th = np.where(T < 200, (np.copy(T)+CtoK) *
sbiri's avatar
sbiri committed
171 172
                  np.power(1000/P,287.1/1004.67),
                  np.copy(T)*np.power(1000/P,287.1/1004.67))  # potential T
173 174 175
    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))
176 177 178 179
    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)
180 181 182
    elif (hum[0] not in ['rh', 'q', 'Td']):
        sys.exit("unknown humidity input")
    elif (hum[0] == 'rh'):
183
        RH = hum[1]
184
        if (np.all(RH < 1)):
185 186 187
            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)
188
    elif (hum[0] == 'q'):
189
        qair = hum[1]
190
        qsea = qsat_sea(sst, P, qmeth)/1000  # surface water q (g/kg)
191
    elif (hum[0] == 'Td'):
192 193 194 195 196 197 198 199 200 201
        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
202 203 204 205 206 207 208 209 210 211 212
    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
    # ------------
213
    rho = P*100/(287.1*tv10n)
sbiri's avatar
sbiri committed
214 215 216
    lv = (2.501-0.00237*SST)*1e6
    cp = 1004.67*(1 + 0.00084*qsea)
    u10n = np.copy(spd)
217
    monob = -100*np.ones(spd.shape)  # Monin-Obukhov length
sbiri's avatar
sbiri committed
218
    cdn = cdn_calc(u10n, Ta, None, lat, meth)
219 220 221 222 223 224 225 226
    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
227 228
    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
229
    dter = np.ones(T.shape)*0.3
230
    dqer = dter*0.622*lv*qsea/(287.1*np.power(sst, 2))
231
    if (gust[0] == 1 and meth == "UA"):
232 233
        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)))
234
    elif (gust[0] == 1):
235
        wind = np.sqrt(np.power(np.copy(spd), 2)+np.power(0.5, 2))
236
    elif (gust[0] == 0):
237 238 239
        wind = np.copy(spd)

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