AirSeaFluxCode.py 28.6 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_moist)
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
83
           "S80"  : default for S80, S88, LP82, YT96 and LY04
84 85
           "ecmwf" : following ecmwf (IFS Documentation cy46r1), default for
           ecmwf
sbiri's avatar
sbiri committed
86 87 88
    Returns
    -------
        res : array that contains
sbiri's avatar
sbiri committed
89 90 91 92
                       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
93 94
                       5. drag coefficient (cd)
                       6. neutral drag coefficient (cdn)
sbiri's avatar
sbiri committed
95 96
                       7. heat exchange coefficient (ct)
                       8. neutral heat exchange coefficient (ctn)
sbiri's avatar
sbiri committed
97
                       9. moisture exhange coefficient (cq)
sbiri's avatar
sbiri committed
98 99
                       10. neutral moisture exchange coefficient (cqn)
                       11. star virtual temperatcure (tsrv)
sbiri's avatar
sbiri committed
100
                       12. star temperature (tsr)
101 102
                       13. star specific humidity (qsr)
                       14. star wind speed (usr)
sbiri's avatar
sbiri committed
103
                       15. momentum stability function (psim)
104 105
                       16. heat stability function (psit)
                       17. moisture stability function (psiq)
106
                       18. 10m neutral wind speed (u10n)
107 108 109 110 111 112
                       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
113
                       25. wind speed at reference height (uref)
114 115
                       26. temperature at reference height (tref)
                       27. specific humidity at reference height (qref)
116
                       28. number of iterations until convergence
117 118
                       29. cool-skin temperature depression (dter)
                       30. cool-skin humidity depression (dqer)
119 120 121 122 123 124
                       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)
125 126 127 128
                       37. gust wind speed (ug)
                       38. Bulk Richardson number (Rib)
                       39. relative humidity (rh)
                       40. flag ("n": normal, "o": out of nominal range,
129
                                 "u": u10n<0, "q":q10n<0
130
                                 "m": missing, "l": Rib<-0.5 or Rib>0.2,
sbiri's avatar
sbiri committed
131
                                 "r" : rh>100%,
132
                                 "i": convergence fail at n)
133

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

sbiri's avatar
sbiri committed
192 193
    dt = Ta - sst
    dq = qair - qsea
194

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

522 523

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

564 565
    if (out == 0):
        res[:, ind] = np.nan
sbiri's avatar
sbiri committed
566 567 568 569 570
        # 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)])
571 572
    # output with pandas
    resAll = pd.DataFrame(data=res.T, index=range(len(spd)),
sbiri's avatar
sbiri committed
573 574 575 576 577 578 579
                          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"])
580 581
    resAll["flag"] = flag
    return resAll
582