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)
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
    # specific heat
    cp = 1004.67*(1+0.00084*qsea) #1005*np.ones(spd.shape)#
167
    #lapse rate
168
    tlapse = gamma("dry", SST, T, qair/1000, cp)
169 170
    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
171
    logging.info('method %s and q method %s | qsea:%s, qair:%s', meth, qmeth,
sbiri's avatar
sbiri committed
172 173
                 np.round(np.nanmedian(qsea), 7),
                 np.round(np.nanmedian(qair), 7))
sbiri's avatar
sbiri committed
174 175
    if (np.all(np.isnan(qsea)) or np.all(np.isnan(qair))):
        print("qsea and qair cannot be nan")
sbiri's avatar
sbiri committed
176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191
    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
192
    flag = np.where(rh > 100, "r", flag)
193

sbiri's avatar
sbiri committed
194 195
    dt = Ta - sst
    dq = qair - qsea
196

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