AirSeaFluxCode.py 30 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", "NCAR", "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
            number of iterations (defautl = 10)
        out : int
79 80
            set 0 to set points that have not converged, negative values of
                  u10n and q10n to missing (default)
81
            set 1 to keep points
82
        L : str
sbiri's avatar
sbiri committed
83
           Monin-Obukhov length definition options
sbiri's avatar
sbiri committed
84
           "tsrv"  : default for "S80", "S88", "LP82", "YT96", "UA", "NCAR",
sbiri's avatar
sbiri committed
85 86 87
                     "C30", "C35"
           "Rb" : following ecmwf (IFS Documentation cy46r1), default for
                  "ecmwf", "Beljaars"
sbiri's avatar
sbiri committed
88 89 90
    Returns
    -------
        res : array that contains
sbiri's avatar
sbiri committed
91 92 93 94
                       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
95 96
                       5. drag coefficient (cd)
                       6. neutral drag coefficient (cdn)
sbiri's avatar
sbiri committed
97 98
                       7. heat exchange coefficient (ct)
                       8. neutral heat exchange coefficient (ctn)
sbiri's avatar
sbiri committed
99
                       9. moisture exhange coefficient (cq)
sbiri's avatar
sbiri committed
100 101
                       10. neutral moisture exchange coefficient (cqn)
                       11. star virtual temperatcure (tsrv)
sbiri's avatar
sbiri committed
102
                       12. star temperature (tsr)
103 104
                       13. star specific humidity (qsr)
                       14. star wind speed (usr)
sbiri's avatar
sbiri committed
105
                       15. momentum stability function (psim)
106 107
                       16. heat stability function (psit)
                       17. moisture stability function (psiq)
108
                       18. 10m neutral wind speed (u10n)
109 110 111 112 113 114
                       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
115
                       25. wind speed at reference height (uref)
116 117
                       26. temperature at reference height (tref)
                       27. specific humidity at reference height (qref)
118
                       28. number of iterations until convergence
119 120
                       29. cool-skin temperature depression (dter)
                       30. cool-skin humidity depression (dqer)
121 122 123 124 125 126
                       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)
127 128 129
                       37. gust wind speed (ug)
                       38. Bulk Richardson number (Rib)
                       39. relative humidity (rh)
130 131
                       40. thickness of the viscous layer (delta)
                       41. lv latent heat of vaporization (Jkg−1)
132
                       42. flag ("n": normal, "o": out of nominal range,
133
                                 "u": u10n<0, "q":q10n<0
134
                                 "m": missing,
135
                                 "l": Rib<-0.5 or Rib>0.2 or z/L>1000,
sbiri's avatar
sbiri committed
136
                                 "r" : rh>100%,
137
                                 "i": convergence fail at n)
138

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

sbiri's avatar
sbiri committed
207 208
    dt = Ta - sst
    dq = qair - qsea
209

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

545
    res = np.zeros((41, len(spd)))
sbiri's avatar
sbiri committed
546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561
    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
562 563 564 565 566 567 568 569
    res[16][:] = psiq
    res[17][:] = u10n
    res[18][:] = t10n
    res[19][:] = tv10n
    res[20][:] = q10n
    res[21][:] = zo
    res[22][:] = zot
    res[23][:] = zoq
570 571 572
    res[24][:] = uref
    res[25][:] = tref
    res[26][:] = qref
573
    res[27][:] = itera
574 575
    res[28][:] = dter
    res[29][:] = dqer
576 577 578 579 580 581
    res[30][:] = dtwl
    res[31][:] = qair
    res[32][:] = qsea
    res[33][:] = Rl
    res[34][:] = Rs
    res[35][:] = Rnl
582 583 584
    res[36][:] = np.sqrt(np.power(wind, 2)-np.power(spd, 2))
    res[37][:] = Rb
    res[38][:] = rh
585
    res[39][:] = tkt
586
    res[40][:] = lv
587

588 589
    if (out == 0):
        res[:, ind] = np.nan
sbiri's avatar
sbiri committed
590
        # set missing values where data have non acceptable values
591 592 593
        if (hum[0] != 'no'):
            res = np.asarray([np.where(q10n < 0, np.nan,
                                        res[i][:]) for i in range(41)])
sbiri's avatar
sbiri committed
594
        res = np.asarray([np.where(u10n < 0, np.nan,
595
                                   res[i][:]) for i in range(41)])
596 597 598
    elif (out == 1):
        print("Warning: the output will contain values for points that have"
              " not converged and negative values (if any) for u10n/q10n")
599 600
    # output with pandas
    resAll = pd.DataFrame(data=res.T, index=range(len(spd)),
sbiri's avatar
sbiri committed
601 602 603 604 605 606
                          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",
607
                                   "Rnl", "ug", "Rib", "rh", "delta", "lv"])
608 609
    resAll["flag"] = flag
    return resAll
610