AirSeaFluxCode.py 25.3 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
def AirSeaFluxCode(spd, T, SST, lat, RH, P, hin, hout, zi, Rl, Rs, jcool,
                   gust, meth, qmeth, n):
sbiri's avatar
sbiri committed
11 12 13 14 15 16
    """ Calculates momentum and heat fluxes using different parameterizations

    Parameters
    ----------
        meth : str
            "S80","S88","LP82","YT96","UA","LY04","C30","C35","C40","ERA5"
17 18 19 20 21 22
        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
sbiri's avatar
sbiri committed
23 24 25 26 27 28 29 30 31 32 33 34 35 36
        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
            latitude (deg)
        RH : float
            relative humidity in %
        P : float
            air pressure
        hin : float
37
            sensor heights in m (array 3x1 or 3xn) default 10m ([10, 10, 10])
sbiri's avatar
sbiri committed
38 39 40
        hout : float
            output height, default is 10m
        zi : int
sbiri's avatar
sbiri committed
41
            PBL height (m)
sbiri's avatar
sbiri committed
42 43 44 45
        Rl : float
            downward longwave radiation (W/m^2)
        Rs : float
            downward shortwave radiation (W/m^2)
46
        jcool : int
sbiri's avatar
sbiri committed
47
            0 if sst is true ocean skin temperature, else 1
48 49
        gust : int
            1 to include the effect of gustiness, else 0
sbiri's avatar
sbiri committed
50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70
        n : int
            number of iterations

    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)
71 72 73 74 75 76 77 78 79 80 81 82 83
                       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)
                       25. velocity at reference height (urefs)
                       26. temperature at reference height (trefs)
                       27. specific humidity at reference height (qrefs)
                       28. number of iterations until convergence
sbiri's avatar
sbiri committed
84 85 86 87 88 89 90
        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)
91 92 93 94
    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
    if np.all(np.isnan(lat)):          # set latitude to 45deg if empty
sbiri's avatar
sbiri committed
95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117
        lat=45*np.ones(spd.shape)
    g = gc(lat, None)             # acceleration due to gravity
    # 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")
        logging.debug('all input is nan')
    if (np.all(np.isnan(RH)) and meth == "C35"):
        RH = np.ones(spd.shape)*80  # if empty set to default for COARE3.5
    elif (np.all(np.isnan(RH))):
        sys.exit("input RH is empty")
        logging.debug('input RH is empty')
    if (np.all(np.isnan(P)) and (meth == "C30" or meth == "C40")):
        P = np.ones(spd.shape)*1015  # if empty set to default for COARE3.0
    elif ((np.all(np.isnan(P))) and meth == "C35"):
        P = np.ones(spd.shape)*1013  # if empty set to default for COARE3.5
    elif (np.all(np.isnan(P))):
        sys.exit("input P is empty")
        logging.debug('input P is empty')
    if (np.all(np.isnan(Rl)) and meth == "C30"):
        Rl = np.ones(spd.shape)*150    # set to default for COARE3.0
    elif ((np.all(np.isnan(Rl)) and meth == "C35") or
          (np.all(np.isnan(Rl)) and meth == "C40")):
        Rl = np.ones(spd.shape)*370    # set to default for COARE3.5
sbiri's avatar
sbiri committed
118 119
    elif (np.all(np.isnan(Rl))):
        Rl = np.ones(spd.shape)*370    # set to default for COARE3.5
sbiri's avatar
sbiri committed
120 121 122 123 124 125 126 127 128
    if (np.all(np.isnan(Rs)) and meth == "C30"):
        Rs = np.ones(spd.shape)*370  # set to default for COARE3.0
    elif ((np.all(np.isnan(Rs))) and (meth == "C35" or meth == "C40")):
        Rs = np.ones(spd.shape)*150  # set to default for COARE3.5
    if ((np.all(np.isnan(zi))) and (meth == "C30" or meth == "C35" or
        meth == "C40")):
        zi = 600  # set to default for COARE3.5
    elif ((np.all(np.isnan(zi))) and (meth == "ERA5" or meth == "UA")):
        zi = 1000
sbiri's avatar
sbiri committed
129 130
    elif (np.all(np.isnan(zi))):
        zi = 800
sbiri's avatar
sbiri committed
131
    ####
132
    th = np.where(T < 200, (np.copy(T)+CtoK) *
sbiri's avatar
sbiri committed
133 134
                  np.power(1000/P,287.1/1004.67),
                  np.copy(T)*np.power(1000/P,287.1/1004.67))  # potential T
135 136 137 138 139 140 141 142
    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))
    if qmeth:
        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)
        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
143
    else:
144 145
        qsea = qsat_sea(sst, P, "Buck2")/1000  # surface water q (g/kg)
        qair = qsat_air(T, P, RH, "Buck2")/1000     # q of air (g/kg)
sbiri's avatar
sbiri committed
146
        logging.info('method %s | qsea:%s, qair:%s', meth,
147 148
                     np.nanmedian(qsea[~np.isnan(qsea)]),
                     np.nanmedian(qair[~np.isnan(qair)]))
sbiri's avatar
sbiri committed
149 150 151
    if (np.all(np.isnan(qsea)) or np.all(np.isnan(qair))):
        print("qsea and qair cannot be nan")
        logging.info('method %s qsea and qair cannot be nan | sst:%s, Ta:%s,'
152 153
                      'P:%s, RH:%s', meth, np.nanmedian(sst), np.nanmedian(Ta),
                      np.nanmedian(P), np.nanmedian(RH))
sbiri's avatar
sbiri committed
154 155 156 157 158 159 160 161 162
    # 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
    # ------------
163
    rho = P*100/(287.1*tv10n)
sbiri's avatar
sbiri committed
164 165 166
    lv = (2.501-0.00237*SST)*1e6
    cp = 1004.67*(1 + 0.00084*qsea)
    u10n = np.copy(spd)
167
    monob = -100*np.ones(spd.shape)  # Monin-Obukhov length
sbiri's avatar
sbiri committed
168
    cdn = cdn_calc(u10n, Ta, None, lat, meth)
169 170 171 172 173 174 175 176
    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
177 178
    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
179
    dter = np.ones(T.shape)*0.3
180
    dqer = dter*0.622*lv*qsea/(287.1*np.power(sst, 2))
181
    if (gust == 1 and meth == "UA"):
182 183
        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)))
184 185 186 187 188 189
    elif (gust == 1):
        wind = np.sqrt(np.power(np.copy(spd), 2)+np.power(0.5, 2))
    elif (gust == 0):
        wind = np.copy(spd)

    if (meth == "UA"):
sbiri's avatar
sbiri committed
190 191 192
        usr = 0.06
        for i in range(5):
            zo = 0.013*np.power(usr,2)/g+0.11*visc_air(Ta)/usr
193 194 195
            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) /
sbiri's avatar
sbiri committed
196
                       (1-5*np.where(Rb < 0.19, Rb, 0.19)),
197 198
                       Rb*np.log(h_in[0]/zo))
        monob = h_in[0]/zol
sbiri's avatar
sbiri committed
199 200 201
        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
202
        logging.info('method %s | wind:%s, usr:%s, '
sbiri's avatar
sbiri committed
203
                     'zo:%s, zot:%s, zoq:%s, Rb:%s, monob:%s', meth,
204 205 206
                     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
207
    elif (meth == "ERA5"):
208
        usr = np.sqrt(cd*np.power(wind, 2))
209
        Rb = ((g*h_in[0]*((2*dt)/(Ta+sst-g*h_in[0]) +
sbiri's avatar
sbiri committed
210 211 212
                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
213
        zol = (Rb*((np.log((h_in[0]+zo)/zo)-psim_calc((h_in[0]+zo) /
sbiri's avatar
sbiri committed
214
               monob, meth)+psim_calc(zo/monob, meth)) /
215 216
               (np.log((h_in[0]+zo)/zot) -
               psit_calc((h_in[0]+zo)/monob, meth) +
sbiri's avatar
sbiri committed
217
               psit_calc(zot/monob, meth))))
218
        monob = h_in[0]/zol
219 220
        logging.info('method %s | wind:%s, usr:%s, '
                     'zo:%s, zot:%s, Rb:%s, monob:%s', meth,
221 222
                     np.nanmedian(wind), np.nanmedian(usr), np.nanmedian(zo),
                     np.nanmedian(zot), np.nanmedian(Rb), np.nanmedian(monob))
sbiri's avatar
sbiri committed
223
    elif (meth == "C30" or meth == "C35" or meth == "C40"):
224
        usr = np.sqrt(cd*np.power(wind, 2))
sbiri's avatar
sbiri committed
225 226 227 228 229 230
        a = 0.011*np.ones(T.shape)
        a = np.where(wind > 10, 0.011+(wind-10)/(18-10)*(0.018-0.011),
                         np.where(wind > 18, 0.018, a))
        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)
231
        zot = zoq
232 233
        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) /
sbiri's avatar
sbiri committed
234
               monob, meth)+psim_calc(zo/monob, meth)) /
235 236
               (np.log((h_in[0]+zo)/zot) -
               psit_calc((h_in[0]+zo)/monob, meth) +
sbiri's avatar
sbiri committed
237
               psit_calc(zot/monob, meth))))
238
        monob = h_in[0]/zol
239 240
        logging.info('method %s | wind:%s, usr:%s, '
                     'zo:%s, zot:%s, Rb:%s, monob:%s', meth,
241 242
                     np.nanmedian(wind), np.nanmedian(usr), np.nanmedian(zo),
                     np.nanmedian(zot), np.nanmedian(Rb), np.nanmedian(monob))
sbiri's avatar
sbiri committed
243
    else:
244
        zo, zot = 0.0001*np.ones(spd.shape), 0.0001*np.ones(spd.shape)
245 246
        usr = np.sqrt(cd*np.power(wind, 2))
        logging.info('method %s | wind:%s, usr:%s, '
247 248 249
                     'zo:%s, zot:%s, monob:%s', meth,
                     np.nanmedian(wind), np.nanmedian(usr), np.nanmedian(zo),
                     np.nanmedian(zot), np.nanmedian(monob))
250
    tsr = (dt+dter*jcool)*kappa/(np.log(hin[1]/zot) -
251
                                 psit_calc(h_in[1]/monob, meth))
252
    qsr = (dq+dqer*jcool)*kappa/(np.log(hin[2]/zot) -
253
                                 psit_calc(h_in[2]/monob, meth))
sbiri's avatar
sbiri committed
254
    # tolerance for u,t,q,usr,tsr,qsr
sbiri's avatar
sbiri committed
255
    tol = np.array([0.01, 0.01, 5e-05, 0.005, 0.001, 5e-07])
256 257
    it, ind = 0, np.where(spd > 0)
    ii, itera = True, np.zeros(spd.shape)*np.nan
sbiri's avatar
sbiri committed
258 259 260 261
    while np.any(ii):
        it += 1
        if it > n:
            break
262 263
        old = np.array([np.copy(u10n), np.copy(t10n), np.copy(q10n),
                       np.copy(usr), np.copy(tsr), np.copy(qsr)])
sbiri's avatar
sbiri committed
264 265
        cdn[ind] = cdn_calc(u10n[ind], Ta[ind], None, lat[ind], meth)
        if (np.all(np.isnan(cdn))):
266
            break
sbiri's avatar
sbiri committed
267 268
            logging.info('break %s at iteration %s cdn<0', meth, it)
        zo[ind] = ref_ht/np.exp(kappa/np.sqrt(cdn[ind]))
269 270 271
        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
272
                                        u10n[ind], zo[ind], Ta[ind], meth)
273 274
        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
275
        ct[ind], cq[ind] = ctcq_calc(cdn[ind], cd[ind], ctn[ind], cqn[ind],
276
                                     h_in[1, ind], h_in[2, ind], ref_ht,
sbiri's avatar
sbiri committed
277 278
                                     psit[ind], psiq[ind])
        if (meth == "UA"):
279 280
            usr[ind] = np.where(h_in[0, ind]/monob[ind] < -1.574,
                                kappa*wind[ind] /
sbiri's avatar
sbiri committed
281 282 283
                                (np.log(-1.574*monob[ind]/zo[ind]) -
                                psim_calc(-1.574, meth) +
                                psim_calc(zo[ind]/monob[ind], meth) +
284
                                1.14*(np.power(-h_in[0, ind]/monob[ind], 1/3) -
285
                                np.power(1.574, 1/3))),
286 287 288 289
                                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
290
                                psim_calc(zo[ind]/monob[ind], meth)),
291 292 293 294 295 296
                                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
297
                                5-5*zo[ind]/monob[ind] +
298 299
                                5*np.log(h_in[0, ind]/monob[ind]) +
                                h_in[0, ind]/monob[ind]-1))))
sbiri's avatar
sbiri committed
300
                                # Zeng et al. 1998 (7-10)
301 302
            tsr[ind] = np.where(h_in[1, ind]/monob[ind] < -0.465,
                                kappa*(dt[ind] +
303
                                dter[ind]*jcool) /
sbiri's avatar
sbiri committed
304 305
                                (np.log((-0.465*monob[ind])/zot[ind]) -
                                psit_calc(-0.465, meth)+0.8 *
306
                                (np.power(0.465, -1/3) -
307 308 309
                                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),
310
                                kappa*(dt[ind]+dter[ind]*jcool) /
311 312
                                (np.log(h_in[1, ind]/zot[ind]) -
                                psit_calc(h_in[1, ind]/monob[ind], meth) +
sbiri's avatar
sbiri committed
313
                                psit_calc(zot[ind]/monob[ind], meth)),
314 315
                                np.where((h_in[1, ind]/monob[ind] > 0) &
                                (h_in[1, ind]/monob[ind] < 1),
316
                                kappa*(dt[ind]+dter[ind]*jcool) /
317 318 319
                                (np.log(h_in[1, ind]/zot[ind]) +
                                5*h_in[1, ind]/monob[ind]-5*zot[ind] /
                                monob[ind]),
320 321 322
                                kappa*(dt[ind]+dter[ind]*jcool) /
                                (np.log(monob[ind]/zot[ind])+5 -
                                5*zot[ind]/monob[ind] +
323 324
                                5*np.log(h_in[1, ind]/monob[ind]) +
                                h_in[1, ind]/monob[ind]-1))))
sbiri's avatar
sbiri committed
325
                                # Zeng et al. 1998 (11-14)
326 327
            qsr[ind] = np.where(h_in[2, ind]/monob[ind] < -0.465,
                                kappa*(dq[ind] +
328
                                dqer[ind]*jcool) /
sbiri's avatar
sbiri committed
329 330 331
                                (np.log((-0.465*monob[ind])/zoq[ind]) -
                                psit_calc(-0.465, meth) +
                                psit_calc(zoq[ind]/monob[ind], meth) +
332
                                0.8*(np.power(0.465, -1/3) -
333 334 335
                                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),
336
                                kappa*(dq[ind]+dqer[ind]*jcool) /
337 338
                                (np.log(h_in[1, ind]/zot[ind]) -
                                psit_calc(h_in[2, ind]/monob[ind], meth) +
sbiri's avatar
sbiri committed
339
                                psit_calc(zoq[ind]/monob[ind], meth)),
340 341
                                np.where((h_in[2, ind]/monob[ind] > 0) &
                                (h_in[2, ind]/monob[ind]<1), kappa*(dq[ind] +
342
                                dqer[ind]*jcool) /
343 344
                                (np.log(h_in[1, ind]/zoq[ind]) +
                                5*h_in[2, ind]/monob[ind]-5*zoq[ind]/monob[ind]),
345 346
                                kappa*(dq[ind]+dqer[ind]*jcool) /
                                (np.log(monob[ind]/zoq[ind])+5 -
sbiri's avatar
sbiri committed
347
                                5*zoq[ind]/monob[ind] +
348 349
                                5*np.log(h_in[2, ind]/monob[ind]) +
                                h_in[2, ind]/monob[ind]-1))))
sbiri's avatar
sbiri committed
350
        elif (meth == "C30" or meth == "C35" or meth == "C40"):
351 352
            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
353 354
            logging.info('method %s | dter = %s | Rnl = %s '
                         '| usr = %s | tsr = %s | qsr = %s', meth,
355 356 357 358 359 360 361
                         np.nanmedian(dter), np.nanmedian(Rnl),
                         np.nanmedian(usr), np.nanmedian(tsr),
                         np.nanmedian(qsr))
            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
362
        else:
363 364
            usr[ind] = (wind[ind]*kappa/(np.log(h_in[0, ind]/zo[ind]) -
                        psim_calc(h_in[0, ind]/monob[ind], meth)))
365 366
            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]
367 368 369 370 371 372 373
        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])
374 375
        Rnl[ind] = 0.97*(5.67e-8*np.power(SST[ind] -
                         dter[ind]*jcool+CtoK, 4)-Rl[ind])
376 377 378 379
        t10n[ind] = (Ta[ind] +
                     (tsr[ind]*(np.log(h_in[1, ind]/ref_ht)-psit[ind])/kappa))
        q10n[ind] = (qair[ind] +
                     (qsr[ind]*(np.log(h_in[2, ind]/ref_ht)-psiq[ind])/kappa))
sbiri's avatar
sbiri committed
380 381 382 383 384 385
        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]
386
            zol[ind] = (kappa*g[ind]*h_in[0, ind]/(T[ind]+CtoK)*(tsr[ind] +
sbiri's avatar
sbiri committed
387
                        0.61*(T[ind]+CtoK)*qsr[ind])/np.power(usr[ind], 2))
388
            monob[ind] = h_in[0, ind]/zol[ind]
sbiri's avatar
sbiri committed
389 390
        elif (meth == "ERA5"):
            tsrv[ind] = tsr[ind]+0.61*t10n[ind]*qsr[ind]
391 392 393
            Rb[ind] = ((g[ind]*h_in[0, ind]*((2*dt[ind])/(Ta[ind] +
                       sst[ind]-g[ind]*h_in[0, ind])+0.61*dq[ind])) /
                       np.power(wind[ind], 2))
sbiri's avatar
sbiri committed
394 395 396
            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]
397 398
            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
399
                        psim_calc(zo[ind]/monob[ind], meth)) /
400 401
                        (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
402
                        psit_calc(zot[ind]/monob[ind], meth))))
403
            monob[ind] = h_in[0, ind]/zol[ind]
sbiri's avatar
sbiri committed
404 405 406
        else:
            tsrv[ind] = tsr[ind]+0.61*t10n[ind]*qsr[ind]
            monob[ind] = (tv10n[ind]*usr[ind]**2)/(g[ind]*kappa*tsrv[ind])
407 408 409 410
        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)
        if (gust == 1 and meth == "UA"):
sbiri's avatar
sbiri committed
411 412 413
            wind[ind] = np.where(dtv[ind] >= 0, np.where(spd[ind] > 0.1,
                                 spd[ind], 0.1),
                                 np.sqrt(np.power(np.copy(spd[ind]), 2) +
414
                                 np.power(get_gust(1, tv[ind], usr[ind],
sbiri's avatar
sbiri committed
415 416
                                 tsrv[ind], zi, lat[ind]), 2)))
                                 # Zeng et al. 1998 (20)
417
        elif (gust == 1 and (meth == "C30" or meth == "C35" or meth == "C40")):
sbiri's avatar
sbiri committed
418 419 420
            wind[ind] = np.sqrt(np.power(np.copy(spd[ind]), 2) +
                                np.power(get_gust(1.2, Ta[ind], usr[ind],
                                tsrv[ind], zi, lat[ind]), 2))
421 422 423 424 425 426 427
        elif (gust == 1):
            wind[ind] = np.sqrt(np.power(np.copy(spd[ind]), 2) +
                                np.power(get_gust(1, Ta[ind], usr[ind],
                                tsrv[ind], zi, lat[ind]), 2))
        elif (gust == 0):
            wind[ind] = np.copy(spd[ind])
        if (meth == "UA"):
sbiri's avatar
sbiri committed
428
            u10n[ind] = (wind[ind]-(usr[ind]/kappa)*(np.log(h_in[0, ind]/10) -
429
                         psim[ind]))
sbiri's avatar
sbiri committed
430
            u10n[u10n < 0] = np.nan
431
        elif (meth == "C30" or meth == "C35" or meth == "C40"):
sbiri's avatar
sbiri committed
432
            u10n[ind] = ((wind[ind]+usr[ind]/kappa*(np.log(10/h_in[0, ind]) -
sbiri's avatar
sbiri committed
433
                         psiu_26(10/monob[ind], meth) +
434
                         psiu_26(h_in[0, ind]/monob[ind], meth)) +
sbiri's avatar
sbiri committed
435 436
                         psiu_26(10/monob[ind], meth)*usr[ind]/kappa) /
                         (wind[ind]/spd[ind]))
sbiri's avatar
sbiri committed
437
            u10n[u10n < 0] = np.nan
sbiri's avatar
sbiri committed
438
        elif (meth == "ERA5"):
sbiri's avatar
sbiri committed
439
            u10n[ind] = (spd[ind]-(usr[ind]/kappa)*(np.log(h_in[0, ind] /
sbiri's avatar
sbiri committed
440
                         ref_ht)-psim[ind]))
sbiri's avatar
sbiri committed
441
            u10n[u10n < 0] = np.nan
sbiri's avatar
sbiri committed
442
        else:
sbiri's avatar
sbiri committed
443
            u10n[ind] = (wind[ind]-(usr[ind]/kappa)*(np.log(h_in[0, ind]/10) -
sbiri's avatar
sbiri committed
444
                         psim[ind]))
sbiri's avatar
sbiri committed
445
            u10n[u10n < 0] = np.nan
sbiri's avatar
sbiri committed
446
        itera[ind] = np.ones(1)*it
447 448 449
        new = np.array([np.copy(u10n), np.copy(t10n), np.copy(q10n),
                       np.copy(usr), np.copy(tsr), np.copy(qsr)])
        d = np.abs(new-old)
sbiri's avatar
sbiri committed
450
        ind = np.where((d[0, :] > tol[0])+(d[1, :] > tol[1]) +
sbiri's avatar
sbiri committed
451 452
                       (d[2, :] > tol[2])+(d[3, :] > tol[3]) +
                       (d[4, :] > tol[4])+(d[5, :] > tol[5]))
453 454
        if (ind[0].size == 0):
            ii = False
sbiri's avatar
sbiri committed
455
        else:
456 457
            ii = True
    logging.info('method %s | # of iterations:%s', meth, it)
sbiri's avatar
sbiri committed
458
    logging.info('method %s | # of points that did not converge :%s', meth,
459
                  ind[0].size)
sbiri's avatar
sbiri committed
460
    # calculate output parameters
461
    rho = (0.34838*P)/(tv10n)
sbiri's avatar
sbiri committed
462
    t10n = t10n-(273.16+tlapse*ref_ht)
463 464
    sensible = -rho*cp*usr*tsr
    latent = -rho*lv*usr*qsr
465
    if (gust == 1):
sbiri's avatar
sbiri committed
466
        tau = rho*np.power(usr, 2)*(spd/wind)
467
    elif (gust == 0):
sbiri's avatar
sbiri committed
468 469 470 471
        tau = rho*np.power(usr, 2)
    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))))
472 473 474 475 476 477 478 479
    urefs = (spd-(usr/kappa)*(np.log(h_in[0]/h_out[0])-psim +
             psim_calc(h_out[0]/monob, meth)))
    trefs = (Ta-(tsr/kappa)*(np.log(h_in[1]/h_out[1])-psit +
             psit_calc(h_out[0]/monob, meth)))
    trefs = trefs-(273.16+tlapse*h_out[1])
    qrefs = (qair-(qsr/kappa)*(np.log(h_in[2]/h_out[2]) -
             psit+psit_calc(h_out[2]/monob, meth)))
    res = np.zeros((28, len(spd)))
sbiri's avatar
sbiri committed
480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495
    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
496 497 498 499 500 501 502 503 504 505 506 507
    res[16][:] = psiq
    res[17][:] = u10n
    res[18][:] = t10n
    res[19][:] = tv10n
    res[20][:] = q10n
    res[21][:] = zo
    res[22][:] = zot
    res[23][:] = zoq
    res[24][:] = urefs
    res[25][:] = trefs
    res[26][:] = qrefs
    res[27][:] = itera
sbiri's avatar
sbiri committed
508
    return res, ind