AirSeaFluxCode.py 28.7 KB
Newer Older
sbiri's avatar
sbiri committed
1
import numpy as np
2
import pandas as pd
sbiri's avatar
sbiri committed
3
import logging
sbiri's avatar
sbiri committed
4
from get_init import get_init
5
from hum_subs import (get_hum, gamma)
sbiri's avatar
sbiri committed
6
from util_subs import (kappa, CtoK, get_heights)
7 8
from flux_subs import (cs_C35, cs_Beljaars, cs_ecmwf, wl_ecmwf,
                       get_gust, get_L, get_strs, psim_calc,
sbiri's avatar
sbiri committed
9
                       psit_calc, cdn_calc, cd_calc, ctcq_calc, ctcqn_calc)
sbiri's avatar
sbiri committed
10 11


12 13 14

def AirSeaFluxCode(spd, T, SST, lat=None, hum=None, P=None, hin=18, hout=10,
                   Rl=None, Rs=None, cskin=None, skin="C35", wl=0, gust=None,
15
                   meth="S88", qmeth="Buck2", tol=None, n=10, out=0, L=None):
16 17 18
    """
    Calculates turbulent surface fluxes using different parameterizations
    Calculates height adjusted values for spd, T, q
sbiri's avatar
sbiri committed
19 20 21 22 23 24 25 26 27 28 29

    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
30 31 32
            latitude (deg), default 45deg
        hum : float
            humidity input switch 2x1 [x, values] default is relative humidity
33 34 35
            x='rh' : relative humidity in %
            x='q' : specific humidity (g/kg)
            x='Td' : dew point temperature (K)
sbiri's avatar
sbiri committed
36
        P : float
37
            air pressure (hPa), default 1013hPa
sbiri's avatar
sbiri committed
38
        hin : float
39
            sensor heights in m (array 3x1 or 3xn), default 18m
sbiri's avatar
sbiri committed
40 41 42 43 44 45
        hout : float
            output height, default is 10m
        Rl : float
            downward longwave radiation (W/m^2)
        Rs : float
            downward shortwave radiation (W/m^2)
sbiri's avatar
sbiri committed
46
        cskin : int
47 48
            0 switch cool skin adjustment off, else 1
            default is 1
49 50 51 52
        skin : str
            cool skin method option "C35", "ecmwf" or "Beljaars"
        wl : int
            warm layer correction default is 0, to switch on set to 1
53
        gust : int
54 55
            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
56
            zi PBL height (m) 600 for COARE, 1000 for UA and ecmwf, 800 default
57
            default for COARE [1, 1.2, 600]
58
            default for UA, ecmwf [1, 1, 1000]
59 60
            default else [1, 1.2, 800]
        meth : str
sbiri's avatar
sbiri committed
61
            "S80", "S88", "LP82", "YT96", "UA", "LY04", "C30", "C35",
62
            "ecmwf", "Beljaars"
63 64
        qmeth : str
            is the saturation evaporation method to use amongst
65 66
            "HylandWexler","Hardy","Preining","Wexler","GoffGratch","WMO",
            "MagnusTetens","Buck","Buck2","WMO2018","Sonntag","Bolton",
67 68 69 70 71 72 73
            "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
sbiri's avatar
sbiri committed
74 75
                    adjustment lim1-6
           default is tol=['all', 0.01, 0.01, 1e-05, 1e-3, 0.1, 0.1]
sbiri's avatar
sbiri committed
76
        n : int
77 78 79 80
            number of iterations (defautl = 10)
        out : int
            set 0 to set points that have not converged to missing (default)
            set 1 to keep points
81
        L : str
sbiri's avatar
sbiri committed
82
           Monin-Obukhov length definition options
sbiri's avatar
sbiri committed
83 84 85 86
           "tsrv"  : default for "S80", "S88", "LP82", "YT96", "UA", "LY04",
                     "C30", "C35"
           "Rb" : following ecmwf (IFS Documentation cy46r1), default for
                  "ecmwf", "Beljaars"
sbiri's avatar
sbiri committed
87 88 89
    Returns
    -------
        res : array that contains
sbiri's avatar
sbiri committed
90 91 92 93
                       1. momentum flux       (N/m^2)
                       2. sensible heat       (W/m^2)
                       3. latent heat         (W/m^2)
                       4. Monin-Obhukov length (m)
sbiri's avatar
sbiri committed
94 95
                       5. drag coefficient (cd)
                       6. neutral drag coefficient (cdn)
sbiri's avatar
sbiri committed
96 97
                       7. heat exchange coefficient (ct)
                       8. neutral heat exchange coefficient (ctn)
sbiri's avatar
sbiri committed
98
                       9. moisture exhange coefficient (cq)
sbiri's avatar
sbiri committed
99 100
                       10. neutral moisture exchange coefficient (cqn)
                       11. star virtual temperatcure (tsrv)
sbiri's avatar
sbiri committed
101
                       12. star temperature (tsr)
102 103
                       13. star specific humidity (qsr)
                       14. star wind speed (usr)
sbiri's avatar
sbiri committed
104
                       15. momentum stability function (psim)
105 106
                       16. heat stability function (psit)
                       17. moisture stability function (psiq)
107
                       18. 10m neutral wind speed (u10n)
108 109 110 111 112 113
                       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)
sbiri's avatar
sbiri committed
114
                       25. wind speed at reference height (uref)
115 116
                       26. temperature at reference height (tref)
                       27. specific humidity at reference height (qref)
117
                       28. number of iterations until convergence
118 119
                       29. cool-skin temperature depression (dter)
                       30. cool-skin humidity depression (dqer)
120 121 122 123 124 125
                       31. warm layer correction (dtwl)
                       32. specific humidity of air (qair)
                       33. specific humidity at sea surface (qsea)
                       34. downward longwave radiation (Rl)
                       35. downward shortwave radiation (Rs)
                       36. downward net longwave radiation (Rnl)
126 127 128 129
                       37. gust wind speed (ug)
                       38. Bulk Richardson number (Rib)
                       39. relative humidity (rh)
                       40. flag ("n": normal, "o": out of nominal range,
130
                                 "u": u10n<0, "q":q10n<0
131
                                 "m": missing, "l": Rib<-0.5 or Rib>0.2,
sbiri's avatar
sbiri committed
132
                                 "r" : rh>100%,
133
                                 "i": convergence fail at n)
134

135
    2021 / Author S. Biri
sbiri's avatar
sbiri committed
136
    """
sbiri's avatar
sbiri committed
137
    logging.basicConfig(filename='flux_calc.log', filemode="w",
sbiri's avatar
sbiri committed
138
                        format='%(asctime)s %(message)s',level=logging.INFO)
139 140
    logging.captureWarnings(True)
    #  check input values and set defaults where appropriate
sbiri's avatar
sbiri committed
141 142 143 144 145 146 147 148 149 150
    lat, hum, P, Rl, Rs, cskin, skin, wl, gust, tol, L, n = get_init(spd, T,
                                                                     SST, lat,
                                                                     hum, P,
                                                                     Rl, Rs,
                                                                     cskin,
                                                                     skin,
                                                                     wl, gust,
                                                                     L, tol, n,
                                                                     meth,
                                                                     qmeth)
151
    ref_ht = 10        # reference height
152 153
    h_in = get_heights(hin, len(spd))  # heights of input measurements/fields
    h_out = get_heights(hout, 1)       # desired height of output variables
154
    logging.info('method %s, inputs: lat: %s | P: %s | Rl: %s |'
155
                 ' Rs: %s | gust: %s | cskin: %s | L : %s', meth,
sbiri's avatar
sbiri committed
156 157 158
                 np.nanmedian(lat), np.round(np.nanmedian(P), 2),
                 np.round(np.nanmedian(Rl),2 ), np.round(np.nanmedian(Rs), 2),
                 gust, cskin, L)
159
    #  set up/calculate temperatures and specific humidities
160
    th = np.where(T < 200, (np.copy(T)+CtoK) *
sbiri's avatar
sbiri committed
161 162
                  np.power(1000/P,287.1/1004.67),
                  np.copy(T)*np.power(1000/P,287.1/1004.67))  # potential T
163
    sst = np.where(SST < 200, np.copy(SST)+CtoK, np.copy(SST))
164
    qair, qsea = get_hum(hum, T, sst, P, qmeth)
sbiri's avatar
sbiri committed
165 166 167 168
    # mask to preserve missing values when initialising variables
    msk = np.empty(sst.shape)
    msk = np.where(np.isnan(spd+T+SST), np.nan, 1)
    Rb = np.empty(sst.shape)*msk
169
    # specific heat
sbiri's avatar
sbiri committed
170
    cp = 1004.67*(1+0.00084*qsea)
171
    #lapse rate
172
    tlapse = gamma("dry", SST, T, qair/1000, cp)
173 174
    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
175
    logging.info('method %s and q method %s | qsea:%s, qair:%s', meth, qmeth,
sbiri's avatar
sbiri committed
176 177
                 np.round(np.nanmedian(qsea), 7),
                 np.round(np.nanmedian(qair), 7))
sbiri's avatar
sbiri committed
178 179
    if (np.all(np.isnan(qsea)) or np.all(np.isnan(qair))):
        print("qsea and qair cannot be nan")
sbiri's avatar
sbiri committed
180 181 182 183 184 185 186 187 188 189 190 191 192 193
    if (hum == None):
        rh = np.ones(sst.shape)*80
    elif (hum[0] == 'rh'):
        rh = hum[1]
    elif (hum[0] == 'Td'):
        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-CtoK)/(Td-32.19)))
        es = 611.21*np.exp(17.502*((T-CtoK)/(T-32.19)))
        rh = 100*esd/es
    flag = np.empty(spd.shape, dtype="object")
    flag[:] = "n"
    flag = np.where(np.isnan(spd+T+SST+hum[1]+P+Rs+Rl), "m", flag)
sbiri's avatar
sbiri committed
194
    flag = np.where(rh > 100, "r", flag)
195

sbiri's avatar
sbiri committed
196 197
    dt = Ta - sst
    dq = qair - qsea
198

199
    #  first guesses
sbiri's avatar
sbiri committed
200
    t10n, q10n = np.copy(Ta), np.copy(qair)
201
    tv10n = t10n*(1+0.6077*q10n)
sbiri's avatar
sbiri committed
202
    #  Zeng et al. 1998
203 204
    tv=th*(1+0.6077*qair)   # virtual potential T
    dtv=dt*(1+0.6077*qair)+0.6077*th*dq
sbiri's avatar
sbiri committed
205
    # ------------
206
    rho = P*100/(287.1*tv10n)
207
    lv = (2.501-0.00237*(sst-CtoK))*1e6
sbiri's avatar
sbiri committed
208
    u10n = np.copy(spd)
209 210
    usr = 0.035*u10n
    cd10n = cdn_calc(u10n, usr, Ta, lat, meth)
sbiri's avatar
sbiri committed
211 212
    psim, psit, psiq = (np.zeros(spd.shape)*msk, np.zeros(spd.shape)*msk,
                        np.zeros(spd.shape)*msk)
213
    cd = cd_calc(cd10n, h_in[0], ref_ht, psim)
sbiri's avatar
sbiri committed
214 215
    tsr, tsrv, qsr = (np.zeros(spd.shape)*msk, np.zeros(spd.shape)*msk,
                      np.zeros(spd.shape)*msk)
sbiri's avatar
sbiri committed
216
    # cskin parameters
sbiri's avatar
sbiri committed
217 218
    tkt = 0.001*np.ones(T.shape)*msk
    dter = np.ones(T.shape)*0.3*msk
219
    dqer = dter*0.622*lv*qsea/(287.1*np.power(sst, 2))
220 221
    Rnl = 0.97*(5.67e-8*np.power(sst-0.3*cskin, 4)-Rl)
    Qs = 0.945*Rs
sbiri's avatar
sbiri committed
222
    dtwl = np.ones(T.shape)*0.3*msk
223
    skt = np.copy(sst)
224
    # gustiness adjustment
225
    if (gust[0] == 1 and meth == "UA"):
226 227
        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)))
228
    elif (gust[0] == 1):
229
        wind = np.sqrt(np.power(np.copy(spd), 2)+np.power(0.5, 2))
230
    elif (gust[0] == 0):
231
        wind = np.copy(spd)
232
    # stars and roughness lengths
233
    usr = np.sqrt(cd*np.power(wind, 2))
sbiri's avatar
sbiri committed
234 235
    zo, zot, zoq = (1e-4*np.ones(spd.shape)*msk, 1e-4*np.ones(spd.shape)*msk,
                    1e-4*np.ones(spd.shape)*msk)
sbiri's avatar
sbiri committed
236 237 238 239 240 241
    ct10n = np.power(kappa, 2)/(np.log(h_in[0]/zo)*np.log(h_in[1]/zot))
    cq10n = np.power(kappa, 2)/(np.log(h_in[0]/zo)*np.log(h_in[2]/zoq))
    ct = np.power(kappa, 2)/((np.log(h_in[0]/zo)-psim) *
                             (np.log(h_in[1]/zot)-psit))
    cq = np.power(kappa, 2)/((np.log(h_in[0]/zo)-psim) *
                             (np.log(h_in[2]/zoq)-psiq))
sbiri's avatar
sbiri committed
242
    monob = -100*np.ones(spd.shape)*msk  # Monin-Obukhov length
243 244
    tsr = (dt+dter*cskin-dtwl*wl)*kappa/(np.log(h_in[1]/zot) -
                                         psit_calc(h_in[1]/monob, meth))
sbiri's avatar
sbiri committed
245
    qsr = (dq+dqer*cskin)*kappa/(np.log(h_in[2]/zoq) -
246
                                 psit_calc(h_in[2]/monob, meth))
247
    # set-up to feed into iteration loop
248
    it, ind = 0, np.where(spd > 0)
sbiri's avatar
sbiri committed
249 250 251 252
    ii, itera = True, -1*np.ones(spd.shape)*msk
    tau = 0.05*np.ones(spd.shape)*msk
    sensible = 5*np.ones(spd.shape)*msk
    latent = 65*np.ones(spd.shape)*msk
253
    #  iteration loop
sbiri's avatar
sbiri committed
254 255 256 257
    while np.any(ii):
        it += 1
        if it > n:
            break
258 259 260 261 262 263 264
        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)])
265
        cd10n[ind] = cdn_calc(u10n[ind], usr[ind], Ta[ind], lat[ind], meth)
266
        if (np.all(np.isnan(cd10n))):
267
            break
268 269
            logging.info('break %s at iteration %s cd10n<0', meth, it)
        zo[ind] = ref_ht/np.exp(kappa/np.sqrt(cd10n[ind]))
270
        psim[ind] = psim_calc(h_in[0, ind]/monob[ind], meth)
271 272
        cd[ind] = cd_calc(cd10n[ind], h_in[0, ind], ref_ht, psim[ind])
        ct10n[ind], cq10n[ind] = ctcqn_calc(h_in[1, ind]/monob[ind],
273
                                            cd10n[ind], usr[ind], zo[ind],
274
                                            Ta[ind], meth)
275
        zot[ind] = ref_ht/(np.exp(np.power(kappa, 2) /
276
                           (ct10n[ind]*np.log(ref_ht/zo[ind]))))
277
        zoq[ind] = ref_ht/(np.exp(np.power(kappa, 2) /
278
                           (cq10n[ind]*np.log(ref_ht/zo[ind]))))
279 280
        psit[ind] = psit_calc(h_in[1, ind]/monob[ind], meth)
        psiq[ind] = psit_calc(h_in[2, ind]/monob[ind], meth)
281
        ct[ind], cq[ind] = ctcq_calc(cd10n[ind], cd[ind], ct10n[ind], cq10n[ind],
282
                                      h_in[:, ind], [ref_ht, ref_ht, ref_ht],
283
                                      psit[ind], psiq[ind])
sbiri's avatar
sbiri committed
284 285 286 287 288
        if (meth == "LY04"):
            cd = np.maximum(np.copy(cd), 1e-4)
            ct = np.maximum(np.copy(ct), 1e-4)
            cq = np.maximum(np.copy(cq), 1e-4)
            zo = np.minimum(np.copy(zo), 0.0025)
sbiri's avatar
sbiri committed
289 290 291
        usr[ind], tsr[ind], qsr[ind] = get_strs(h_in[:, ind], monob[ind],
                                                wind[ind], zo[ind], zot[ind],
                                                zoq[ind], dt[ind], dq[ind],
292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351
                                                dter[ind], dqer[ind], dtwl[ind],
                                                ct[ind], cq[ind], cskin, wl,
                                                meth)
        if ((cskin == 1) and (wl == 0)):
            if (skin == "C35"):
                dter[ind], dqer[ind], tkt[ind] = cs_C35(sst[ind], qsea[ind],
                                                        rho[ind], Rs[ind],
                                                        Rnl[ind],
                                                        cp[ind], lv[ind],
                                                        np.copy(tkt[ind]),
                                                        usr[ind], tsr[ind],
                                                        qsr[ind], lat[ind])
            elif (skin == "ecmwf"):
                dter[ind] = cs_ecmwf(rho[ind], Rs[ind], Rnl[ind], cp[ind],
                                     lv[ind], usr[ind], tsr[ind], qsr[ind],
                                     sst[ind], lat[ind])
                dqer[ind] = (dter[ind]*0.622*lv[ind]*qsea[ind] /
                             (287.1*np.power(sst[ind], 2)))
            elif (skin == "Beljaars"):
                Qs[ind], dter[ind] = cs_Beljaars(rho[ind], Rs[ind], Rnl[ind],
                                                 cp[ind], lv[ind], usr[ind],
                                                 tsr[ind], qsr[ind], lat[ind],
                                                 np.copy(Qs[ind]))
                dqer = dter*0.622*lv*qsea/(287.1*np.power(sst, 2))
        elif ((cskin == 1) and (wl == 1)):
            if (skin == "C35"):
                dter[ind], dqer[ind], tkt[ind] = cs_C35(sst[ind], qsea[ind],
                                                        rho[ind], Rs[ind],
                                                        Rnl[ind],
                                                        cp[ind], lv[ind],
                                                        np.copy(tkt[ind]),
                                                        usr[ind], tsr[ind],
                                                        qsr[ind], lat[ind])
                dtwl[ind] = wl_ecmwf(rho[ind], Rs[ind], Rnl[ind], cp[ind],
                                     lv[ind], usr[ind], tsr[ind], qsr[ind],
                                     np.copy(sst[ind]), np.copy(skt[ind]),
                                     np.copy(dter[ind]), lat[ind])
                skt = np.copy(sst)-dter+dtwl
            elif (skin == "ecmwf"):
                dter[ind] = cs_ecmwf(rho[ind], Rs[ind], Rnl[ind], cp[ind],
                                     lv[ind], usr[ind], tsr[ind], qsr[ind],
                                     sst[ind], lat[ind])
                dtwl[ind] = wl_ecmwf(rho[ind], Rs[ind], Rnl[ind], cp[ind],
                                     lv[ind], usr[ind], tsr[ind], qsr[ind],
                                     np.copy(sst[ind]), np.copy(skt[ind]),
                                     np.copy(dter[ind]), lat[ind])
                skt = np.copy(sst)-dter+dtwl
                dqer[ind] = (dter[ind]*0.622*lv[ind]*qsea[ind] /
                             (287.1*np.power(skt[ind], 2)))
            elif (skin == "Beljaars"):
                Qs[ind], dter[ind] = cs_Beljaars(rho[ind], Rs[ind], Rnl[ind],
                                                 cp[ind], lv[ind], usr[ind],
                                                 tsr[ind], qsr[ind], lat[ind],
                                                 np.copy(Qs[ind]))
                dtwl[ind] = wl_ecmwf(rho[ind], Rs[ind], Rnl[ind], cp[ind],
                                     lv[ind], usr[ind], tsr[ind], qsr[ind],
                                     np.copy(sst[ind]), np.copy(skt[ind]),
                                     np.copy(dter[ind]), lat[ind])
                skt = np.copy(sst)-dter+dtwl
                dqer = dter*0.622*lv*qsea/(287.1*np.power(skt, 2))
352 353 354
        else:
           dter[ind] = np.zeros(sst[ind].shape)
           dqer[ind] = np.zeros(sst[ind].shape)
355
           tkt[ind] = 0.001*np.ones(T[ind].shape)
sbiri's avatar
sbiri committed
356 357
        logging.info('method %s | dter = %s | dqer = %s | tkt = %s | Rnl = %s '
                     '| usr = %s | tsr = %s | qsr = %s', meth,
sbiri's avatar
sbiri committed
358 359 360 361 362 363 364
                     np.round(np.nanmedian(dter), 2),
                     np.round(np.nanmedian(dqer), 7),
                     np.round(np.nanmedian(tkt), 2),
                     np.round(np.nanmedian(Rnl), 2),
                     np.round(np.nanmedian(usr), 3),
                     np.round(np.nanmedian(tsr), 4),
                     np.round(np.nanmedian(qsr), 7))
365 366
        Rnl[ind] = 0.97*(5.67e-8*np.power(sst[ind] -
                          dter[ind]*cskin, 4)-Rl[ind])
367 368 369 370
        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]))
371 372 373 374 375
        tv10n[ind] = t10n[ind]*(1+0.6077*q10n[ind])
        tsrv[ind], monob[ind], Rb[ind] = get_L(L, lat[ind], usr[ind], tsr[ind],
                                               qsr[ind], h_in[:, ind], Ta[ind],
                                               sst[ind]-dter[ind]*cskin+dtwl[ind]*wl,
                                               qair[ind], qsea[ind], wind[ind],
sbiri's avatar
sbiri committed
376 377
                                               np.copy(monob[ind]), zo[ind],
                                               zot[ind], psim[ind], meth)
378
        # sst[ind]-dter[ind]*cskin+dtwl[ind]*wl
379 380 381
        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)
382
        if (gust[0] == 1 and meth == "UA"):
sbiri's avatar
sbiri committed
383
            wind[ind] = np.where(dtv[ind] >= 0, np.where(spd[ind] > 0.1,
384 385 386 387 388
                                  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)
sbiri's avatar
sbiri committed
389
        elif (gust[0] == 1 and (meth == "C30" or meth == "C35")): # or meth == "C40"
sbiri's avatar
sbiri committed
390
            wind[ind] = np.sqrt(np.power(np.copy(spd[ind]), 2) +
391 392 393
                                np.power(get_gust(gust[1], Ta[ind], usr[ind],
                                tsrv[ind], gust[2], lat[ind]), 2))
        elif (gust[0] == 1):
394
            wind[ind] = np.sqrt(np.power(np.copy(spd[ind]), 2) +
395 396 397
                                np.power(get_gust(gust[1], Ta[ind], usr[ind],
                                tsrv[ind], gust[2], lat[ind]), 2))
        elif (gust[0] == 0):
398
            wind[ind] = np.copy(spd[ind])
399 400
        u10n[ind] = wind[ind]-usr[ind]/kappa*(np.log(h_in[0, ind]/10) -
                                              psim[ind])
sbiri's avatar
sbiri committed
401 402
        if (it < 4): # make sure you allow small negative values convergence
            u10n = np.where(u10n < 0, 0.5, u10n)
sbiri's avatar
sbiri committed
403 404
        utmp = np.copy(u10n)
        utmp = np.where(utmp < 0, np.nan, utmp)
sbiri's avatar
sbiri committed
405
        itera[ind] = np.ones(1)*it
406 407 408 409 410 411 412 413 414
        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'):
sbiri's avatar
sbiri committed
415
            new = np.array([np.copy(utmp), np.copy(t10n), np.copy(q10n)])
416
        elif (tol[0] == 'all'):
sbiri's avatar
sbiri committed
417
            new = np.array([np.copy(utmp), np.copy(t10n), np.copy(q10n),
418
                            np.copy(tau), np.copy(sensible), np.copy(latent)])
sbiri's avatar
sbiri committed
419 420 421 422 423 424 425 426 427 428 429 430
        if (it > 2): # force at least two iterations
            d = np.abs(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]))
431 432
        if (ind[0].size == 0):
            ii = False
sbiri's avatar
sbiri committed
433
        else:
434
            ii = True
435
    itera[ind] = -1
sbiri's avatar
sbiri committed
436
    itera = np.where(itera > n, -1, itera)
437
    logging.info('method %s | # of iterations:%s', meth, it)
sbiri's avatar
sbiri committed
438
    logging.info('method %s | # of points that did not converge :%s \n', meth,
439
                  ind[0].size)
sbiri's avatar
sbiri committed
440
    # calculate output parameters
441
    rho = (0.34838*P)/(tv10n)
sbiri's avatar
sbiri committed
442
    t10n = t10n-(273.16+tlapse*ref_ht)
443 444 445 446
    # solve for zo from cd10n
    zo = ref_ht/np.exp(kappa/np.sqrt(cd10n))
    # adjust neutral cdn at any output height
    cdn = np.power(kappa/np.log(hout/zo), 2)
sbiri's avatar
sbiri committed
447 448
    cd = cd_calc(cdn, h_out[0], h_out[0], psim)
    # solve for zot, zoq from ct10n, cq10n
449 450 451
    zot = ref_ht/(np.exp(kappa**2/(ct10n*np.log(ref_ht/zo))))
    zoq = ref_ht/(np.exp(kappa**2/(cq10n*np.log(ref_ht/zo))))
    # adjust neutral ctn, cqn at any output height
sbiri's avatar
sbiri committed
452 453
    ctn =np.power(kappa, 2)/(np.log(h_out[0]/zo)*np.log(h_out[1]/zot))
    cqn =np.power(kappa, 2)/(np.log(h_out[0]/zo)*np.log(h_out[2]/zoq))
454
    ct, cq = ctcq_calc(cdn, cd, ctn, cqn, h_out, h_out, psit, psiq)
455 456 457 458
    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)))
459
    tref = tref-(CtoK+tlapse*h_out[1])
460 461
    qref = (qair-qsr/kappa*(np.log(h_in[2]/h_out[2]) -
            psit+psit_calc(h_out[2]/monob, meth)))
462
    if (wl == 0):
sbiri's avatar
sbiri committed
463
        dtwl = np.zeros(T.shape)*msk # reset to zero if not used
464 465 466 467
    flag = np.where((u10n < 0) & (flag == "n"), "u",
                        np.where((u10n < 0) &
                                 (np.char.find(flag.astype(str), 'u') == -1),
                                 flag+[","]+["u"], flag))
468 469 470
    flag = np.where((q10n < 0) & (flag == "n"), "q",
                    np.where((q10n < 0) & (flag != "n"), flag+[","]+["q"],
                             flag))
sbiri's avatar
sbiri committed
471
    flag = np.where(((Rb < -0.5) | (Rb > 0.2)) & (flag == "n"), "l",
sbiri's avatar
sbiri committed
472
                    np.where(((Rb < -0.5) | (Rb > 0.2)) &
sbiri's avatar
sbiri committed
473 474 475 476 477 478 479 480 481 482 483 484 485 486
                             (flag != "n"), flag+[","]+["l"], flag))
    if (out == 1):
        flag = np.where((itera == -1) & (flag == "n"), "i",
                        np.where((itera == -1) &
                                 ((flag != "n") &
                                  (np.char.find(flag.astype(str), 'm') == -1)),
                                 flag+[","]+["i"], flag))
    else:
        flag = np.where((itera == -1) & (flag == "n"), "i",
                        np.where((itera == -1) &
                                 ((flag != "n") &
                                  (np.char.find(flag.astype(str), 'm') == -1) &
                                  (np.char.find(flag.astype(str), 'u') == -1)),
                                 flag+[","]+["i"], flag))
487
    if (meth == "S80"):
sbiri's avatar
sbiri committed
488 489 490 491 492
        flag = np.where(((utmp < 6) | (utmp > 22)) & (flag == "n"), "o",
                        np.where(((utmp < 6) | (utmp > 22)) &
                                 ((flag != "n") &
                                  (np.char.find(flag.astype(str), 'u') == -1) &
                                  (np.char.find(flag.astype(str), 'q') == -1)),
493 494
                                 flag+[","]+["o"], flag))
    elif (meth == "LP82"):
sbiri's avatar
sbiri committed
495 496 497 498 499
        flag = np.where(((utmp < 3) | (utmp > 25)) & (flag == "n"), "o",
                        np.where(((utmp < 3) | (utmp > 25)) &
                                 ((flag != "n") &
                                  (np.char.find(flag.astype(str), 'u') == -1) &
                                  (np.char.find(flag.astype(str), 'q') == -1)),
500 501
                                 flag+[","]+["o"], flag))
    elif (meth == "YT96"):
sbiri's avatar
sbiri committed
502 503 504 505 506
        flag = np.where(((utmp < 3) | (utmp > 26)) & (flag == "n"), "o",
                        np.where(((utmp < 3) | (utmp > 26)) &
                                 ((flag != "n") &
                                  (np.char.find(flag.astype(str), 'u') == -1) &
                                  (np.char.find(flag.astype(str), 'q') == -1)),
507 508
                                 flag+[","]+["o"], flag))
    elif (meth == "UA"):
sbiri's avatar
sbiri committed
509 510 511 512 513
        flag = np.where((utmp > 18) & (flag == "n"), "o",
                        np.where((utmp > 18) &
                                 ((flag != "n") &
                                  (np.char.find(flag.astype(str), 'u') == -1) &
                                  (np.char.find(flag.astype(str), 'q') == -1)),
514 515
                                 flag+[","]+["o"], flag))
    elif (meth == "LY04"):
sbiri's avatar
sbiri committed
516 517 518 519 520
        flag = np.where((utmp < 0.5) & (flag == "n"), "o",
                        np.where((utmp < 0.5) &
                                 ((flag != "n") &
                                  (np.char.find(flag.astype(str), 'u') == -1) &
                                  (np.char.find(flag.astype(str), 'q') == -1)),
521
                                 flag+[","]+["o"], flag))
sbiri's avatar
sbiri committed
522

523 524

    res = np.zeros((39, len(spd)))
sbiri's avatar
sbiri committed
525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540
    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
541 542 543 544 545 546 547 548
    res[16][:] = psiq
    res[17][:] = u10n
    res[18][:] = t10n
    res[19][:] = tv10n
    res[20][:] = q10n
    res[21][:] = zo
    res[22][:] = zot
    res[23][:] = zoq
549 550 551
    res[24][:] = uref
    res[25][:] = tref
    res[26][:] = qref
552
    res[27][:] = itera
553 554
    res[28][:] = dter
    res[29][:] = dqer
555 556 557 558 559 560
    res[30][:] = dtwl
    res[31][:] = qair
    res[32][:] = qsea
    res[33][:] = Rl
    res[34][:] = Rs
    res[35][:] = Rnl
561 562 563
    res[36][:] = np.sqrt(np.power(wind, 2)-np.power(spd, 2))
    res[37][:] = Rb
    res[38][:] = rh
564

565 566
    if (out == 0):
        res[:, ind] = np.nan
sbiri's avatar
sbiri committed
567 568 569 570 571
        # set missing values where data have non acceptable values
        res = np.asarray([np.where(q10n < 0, np.nan,
                                   res[i][:]) for i in range(39)])
        res = np.asarray([np.where(u10n < 0, np.nan,
                                   res[i][:]) for i in range(39)])
572 573
    # output with pandas
    resAll = pd.DataFrame(data=res.T, index=range(len(spd)),
sbiri's avatar
sbiri committed
574 575 576 577 578 579 580
                          columns=["tau", "shf", "lhf", "L", "cd", "cdn", "ct",
                                   "ctn", "cq", "cqn", "tsrv", "tsr", "qsr",
                                   "usr", "psim", "psit","psiq", "u10n",
                                   "t10n", "tv10n", "q10n", "zo", "zot", "zoq",
                                   "uref", "tref", "qref", "iteration", "dter",
                                   "dqer", "dtwl", "qair", "qsea", "Rl", "Rs",
                                   "Rnl", "ug", "Rib", "rh"])
581 582
    resAll["flag"] = flag
    return resAll
583