AirSeaFluxCode.py 28.8 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)
403
        flag = np.where((u10n < 0) & (flag == "n"), "u",
sbiri's avatar
sbiri committed
404 405
                        np.where((u10n < 0) &
                                 (np.char.find(flag.astype(str), 'u') == -1),
sbiri's avatar
sbiri committed
406
                                 flag+[","]+["u"], flag))
sbiri's avatar
sbiri committed
407 408
        utmp = np.copy(u10n)
        utmp = np.where(utmp < 0, np.nan, utmp)
sbiri's avatar
sbiri committed
409
        itera[ind] = np.ones(1)*it
sbiri's avatar
sbiri committed
410
        rho = (0.34838*P)/(tv10n)
411 412 413 414 415 416 417 418 419
        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
420
            new = np.array([np.copy(utmp), np.copy(t10n), np.copy(q10n)])
421
        elif (tol[0] == 'all'):
sbiri's avatar
sbiri committed
422
            new = np.array([np.copy(utmp), np.copy(t10n), np.copy(q10n),
423
                            np.copy(tau), np.copy(sensible), np.copy(latent)])
sbiri's avatar
sbiri committed
424 425 426 427 428 429 430 431 432 433 434 435
        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]))
436 437
        if (ind[0].size == 0):
            ii = False
sbiri's avatar
sbiri committed
438
        else:
439
            ii = True
440
    itera[ind] = -1
sbiri's avatar
sbiri committed
441
    itera = np.where(itera > n, -1, itera)
442
    logging.info('method %s | # of iterations:%s', meth, it)
sbiri's avatar
sbiri committed
443
    logging.info('method %s | # of points that did not converge :%s \n', meth,
444
                  ind[0].size)
sbiri's avatar
sbiri committed
445
    # calculate output parameters
446
    rho = (0.34838*P)/(tv10n)
sbiri's avatar
sbiri committed
447
    t10n = t10n-(273.16+tlapse*ref_ht)
448 449 450 451
    # 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
452 453
    cd = cd_calc(cdn, h_out[0], h_out[0], psim)
    # solve for zot, zoq from ct10n, cq10n
454 455 456
    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
457 458
    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))
459
    ct, cq = ctcq_calc(cdn, cd, ctn, cqn, h_out, h_out, psit, psiq)
460 461 462 463
    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)))
464
    tref = tref-(CtoK+tlapse*h_out[1])
465 466
    qref = (qair-qsr/kappa*(np.log(h_in[2]/h_out[2]) -
            psit+psit_calc(h_out[2]/monob, meth)))
467
    if (wl == 0):
sbiri's avatar
sbiri committed
468
        dtwl = np.zeros(T.shape)*msk # reset to zero if not used
469 470 471
    flag = np.where((q10n < 0) & (flag == "n"), "q",
                    np.where((q10n < 0) & (flag != "n"), flag+[","]+["q"],
                             flag))
sbiri's avatar
sbiri committed
472
    flag = np.where(((Rb < -0.5) | (Rb > 0.2)) & (flag == "n"), "l",
sbiri's avatar
sbiri committed
473
                    np.where(((Rb < -0.5) | (Rb > 0.2)) &
sbiri's avatar
sbiri committed
474 475 476 477 478 479 480 481 482 483 484 485 486 487
                             (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))
488
    if (meth == "S80"):
sbiri's avatar
sbiri committed
489 490 491 492 493
        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)),
494 495
                                 flag+[","]+["o"], flag))
    elif (meth == "LP82"):
sbiri's avatar
sbiri committed
496 497 498 499 500
        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)),
501 502
                                 flag+[","]+["o"], flag))
    elif (meth == "YT96"):
sbiri's avatar
sbiri committed
503 504 505 506 507
        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)),
508 509
                                 flag+[","]+["o"], flag))
    elif (meth == "UA"):
sbiri's avatar
sbiri committed
510 511 512 513 514
        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)),
515 516
                                 flag+[","]+["o"], flag))
    elif (meth == "LY04"):
sbiri's avatar
sbiri committed
517 518 519 520 521
        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)),
522
                                 flag+[","]+["o"], flag))
sbiri's avatar
sbiri committed
523

524 525

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

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