AirSeaFluxCode.py 26.1 KB
Newer Older
sbiri's avatar
sbiri committed
1 2 3
import numpy as np
import sys
import logging
4 5 6 7
from flux_subs import (kappa, CtoK, get_heights, get_skin, get_gust, get_L,
                       psim_calc, psit_calc, psit_26, psiu_26,
                       cdn_calc, cd_calc, ctcq_calc, ctcqn_calc,
                       gc, qsat_sea, qsat_air)
sbiri's avatar
sbiri committed
8 9


10
def AirSeaFluxCode(spd, T, SST, lat=None, hum=None, P=None,
11 12 13
                   hin=18, hout=10, Rl=None, Rs=None, cskin=None,
                   gust=None, meth="S80", qmeth="Buck2", tol=None, n=10,
                   out=0, L=None):
sbiri's avatar
sbiri committed
14 15 16 17 18 19 20 21 22 23 24 25
    """ Calculates momentum and heat fluxes using different parameterizations

    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
26 27 28
            latitude (deg), default 45deg
        hum : float
            humidity input switch 2x1 [x, values] default is relative humidity
29 30 31
            x='rh' : relative humidity in %
            x='q' : specific humidity (g/kg)
            x='Td' : dew point temperature (K)
sbiri's avatar
sbiri committed
32
        P : float
33
            air pressure (hPa), default 1013hPa
sbiri's avatar
sbiri committed
34
        hin : float
35
            sensor heights in m (array 3x1 or 3xn), default 18m
sbiri's avatar
sbiri committed
36 37 38 39 40 41
        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
42
        cskin : int
43 44
            0 switch cool skin adjustment off, else 1
            default is 1
45
        gust : int
46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66
            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
            zi PBL height (m) 600 for COARE, 1000 for UA and ERA5, 800 default
            default for COARE [1, 1.2, 600]
            default for UA, ERA5 [1, 1, 1000]
            default else [1, 1.2, 800]
        meth : str
            "S80","S88","LP82","YT96","UA","LY04","C30","C35","C40","ERA5"
        qmeth : str
            is the saturation evaporation method to use amongst
            "HylandWexler","Hardy","Preining","Wexler","GoffGratch","CIMO",
            "MagnusTetens","Buck","Buck2","WMO","WMO2000","Sonntag","Bolton",
            "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
                    adjustment lim1-6 ['all', 0.01, 0.01, 5e-05, 0.01, 1, 1]
           default is tol=['flux', 0.01, 1, 1]
sbiri's avatar
sbiri committed
67
        n : int
68 69 70 71
            number of iterations (defautl = 10)
        out : int
            set 0 to set points that have not converged to missing (default)
            set 1 to keep points
sbiri's avatar
sbiri committed
72 73 74 75 76 77
        L : int
           Monin-Obukhov length definition options
           0 : constant, default for S80, S88, LP82, YT96 and LY04
           1 : following UA (Zeng et al., 1998), default for UA
           2 : following ERA5 (IFS Documentation cy46r1), default for ERA5
           3 : COARE3.5 (Edson et al., 2013), default for C30, C35 and C40
sbiri's avatar
sbiri committed
78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95
    Returns
    -------
        res : array that contains
                       1. momentum flux (W/m^2)
                       2. sensible heat (W/m^2)
                       3. latent heat (W/m^2)
                       4. Monin-Obhukov length (mb)
                       5. drag coefficient (cd)
                       6. neutral drag coefficient (cdn)
                       7. heat exhange coefficient (ct)
                       8. neutral heat exhange coefficient (ctn)
                       9. moisture exhange coefficient (cq)
                       10. neutral moisture exhange coefficient (cqn)
                       11. star virtual temperature (tsrv)
                       12. star temperature (tsr)
                       13. star humidity (qsr)
                       14. star velocity (usr)
                       15. momentum stability function (psim)
96 97 98 99 100 101 102 103 104
                       16. heat stability function (psit)
                       17. moisture stability function (psiq)
                       18. 10m neutral velocity (u10n)
                       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)
105 106 107
                       25. velocity at reference height (uref)
                       26. temperature at reference height (tref)
                       27. specific humidity at reference height (qref)
108
                       28. number of iterations until convergence
sbiri's avatar
sbiri committed
109 110 111 112 113 114 115
        ind : int
            the indices in the matrix for the points that did not converge
            after the maximum number of iterations
    The code is based on bform.f and flux_calc.R modified by S. Biri
    """
    logging.basicConfig(filename='flux_calc.log',
                        format='%(asctime)s %(message)s',level=logging.INFO)
116 117 118
    if ((type(spd) != np.ndarray) or (type(T) != np.ndarray) or
        (type(SST) != np.ndarray)):
        sys.exit("input type of spd, T and SST should be numpy.ndarray")
119 120 121
    ref_ht, tlapse = 10, 0.0098        # reference height, lapse rate
    h_in = get_heights(hin, len(spd))  # heights of input measurements/fields
    h_out = get_heights(hout, 1)       # desired height of output variables
sbiri's avatar
sbiri committed
122
    # if input values are nan break
123 124 125 126 127 128 129
    if meth not in ["S80", "S88", "LP82", "YT96", "UA", "LY04", "C30", "C35",
                    "C40","ERA5"]:
        sys.exit("unknown method")
    if qmeth not in ["HylandWexler", "Hardy", "Preining", "Wexler", "CIMO",
                     "GoffGratch", "MagnusTetens", "Buck", "Buck2", "WMO",
                     "WMO2000", "Sonntag", "Bolton", "IAPWS", "MurphyKoop"]:
        sys.exit("unknown q-method")
sbiri's avatar
sbiri committed
130 131
    if (np.all(np.isnan(spd)) or np.all(np.isnan(T)) or np.all(np.isnan(SST))):
        sys.exit("input wind, T or SST is empty")
132
        logging.debug('all spd or T or SST input is nan')
133
    if (np.all(lat == None)):  # set latitude to 45deg if empty
134
        lat = 45*np.ones(spd.shape)
135
    elif ((np.all(lat != None)) and (np.size(lat) == 1)):
136 137
        lat = np.ones(spd.shape)*np.copy(lat)
    g = gc(lat, None)             # acceleration due to gravity
138
    if ((np.all(P == None)) and (meth == "C30" or meth == "C40")):
sbiri's avatar
sbiri committed
139
        P = np.ones(spd.shape)*1015  # if empty set to default for COARE3.0
140
    elif ((np.all(P == None)) or np.all(np.isnan(P))):
141 142
        P = np.ones(spd.shape)*1013
        logging.debug('input P is empty and set to 1013hPa')
143
    elif (((np.all(P != None)) or np.all(~np.isnan(P))) and np.size(P) == 1):
144
        P = np.ones(spd.shape)*np.copy(P)
145
    if ((np.all(Rl == None) or np.all(np.isnan(Rl))) and meth == "C30"):
sbiri's avatar
sbiri committed
146
        Rl = np.ones(spd.shape)*150    # set to default for COARE3.0
147 148
    elif (((np.all(Rl == None) or np.all(np.isnan(Rl))) and meth == "C35") or
          ((np.all(Rl == None) or np.all(np.isnan(Rl))) and meth == "C40")):
sbiri's avatar
sbiri committed
149
        Rl = np.ones(spd.shape)*370    # set to default for COARE3.5
150
    elif (np.all(Rl == None) or np.all(np.isnan(Rl))):
sbiri's avatar
sbiri committed
151
        Rl = np.ones(spd.shape)*370    # set to default for COARE3.5
152
    if ((np.all(Rs == None) or np.all(np.isnan(Rs))) and meth == "C30"):
sbiri's avatar
sbiri committed
153
        Rs = np.ones(spd.shape)*370  # set to default for COARE3.0
154
    elif (np.all(Rs == None) or np.all(np.isnan(Rs))):
sbiri's avatar
sbiri committed
155
        Rs = np.ones(spd.shape)*150  # set to default for COARE3.5
156 157 158 159 160 161
    if ((gust == None) and (meth == "C30" or meth == "C35" or meth == "C40")):
        gust = [1, 1.2, 600]
    elif ((gust == None) and (meth == "UA" or meth == "ERA5")):
        gust = [1, 1, 1000]
    elif (gust == None):
        gust = [1, 1.2, 800]
162 163
    elif (np.size(gust) < 3):
        sys.exit("gust input must be a 3x1 array")
164 165
    if (tol == None):
        tol = ['flux', 0.01, 1, 1]
166 167
    elif (tol[0] not in ['flux', 'ref', 'all']):
        sys.exit("unknown tolerance input")
sbiri's avatar
sbiri committed
168
    if ((cskin == None) and (meth == "S80" or meth == "S88" or meth == "LP82"
169
                             or meth == "YT96")):
sbiri's avatar
sbiri committed
170 171
       cskin = 0
    elif ((cskin == None) and (meth == "UA" or meth == "LY04" or meth == "C30"
172 173
                               or meth == "C35" or meth == "C40"
                               or meth == "ERA5")):
sbiri's avatar
sbiri committed
174
       cskin = 1
175
    logging.info('method %s, inputs: lat: %s | P: %s | Rl: %s |'
sbiri's avatar
sbiri committed
176
                 ' Rs: %s | gust: %s | cskin: %s', meth,
177
                 np.nanmedian(lat), np.nanmedian(P), np.nanmedian(Rl),
sbiri's avatar
sbiri committed
178 179 180 181 182 183 184 185 186 187 188 189
                 np.nanmedian(Rs), gust, cskin)
    if (L not in [None, 0, 1, 2, 3]):
        sys.exit("L input must be either None, 0, 1, 2 or 3")
    if ((L == None) and (meth == "S80" or meth == "S88" or meth == "LP82"
                             or meth == "YT96" or meth == "LY04")):
       L = 0
    elif ((L == None) and (meth == "UA")):
       L = 1
    elif ((L == None) and (meth == "ERA5")):
       L = 2
    elif ((L == None) and (meth == "C30" or meth == "C35" or meth == "C40")):
       L = 3
sbiri's avatar
sbiri committed
190
    ####
191
    th = np.where(T < 200, (np.copy(T)+CtoK) *
sbiri's avatar
sbiri committed
192 193
                  np.power(1000/P,287.1/1004.67),
                  np.copy(T)*np.power(1000/P,287.1/1004.67))  # potential T
194 195 196
    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
    sst = np.where(SST < 200, np.copy(SST)+CtoK, np.copy(SST))
197 198 199 200
    if (hum == None):
        RH = np.ones(spd.shape)*80
        qsea = qsat_sea(sst, P, qmeth)/1000     # surface water q (g/kg)
        qair = qsat_air(T, P, RH, qmeth)/1000   # q of air (g/kg)
201 202 203
    elif (hum[0] not in ['rh', 'q', 'Td']):
        sys.exit("unknown humidity input")
    elif (hum[0] == 'rh'):
204
        RH = hum[1]
205
        if (np.all(RH < 1)):
206 207 208
            sys.exit("input relative humidity units should be \%")
        qsea = qsat_sea(sst, P, qmeth)/1000    # surface water q (g/kg)
        qair = qsat_air(T, P, RH, qmeth)/1000  # q of air (g/kg)
209
    elif (hum[0] == 'q'):
210
        qair = hum[1]
211
        qsea = qsat_sea(sst, P, qmeth)/1000  # surface water q (g/kg)
212
    elif (hum[0] == 'Td'):
213 214 215 216 217 218 219 220 221 222
        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-273.16)/(Td-32.19)))
        es = 611.21*np.exp(17.502*((T-273.16)/(T-32.19)))
        RH = 100*esd/es
        qair = qsat_air(T, P, RH, qmeth)/1000  # q of air (g/kg)
        qsea = qsat_sea(sst, P, qmeth)/1000    # surface water q (g/kg)
    logging.info('method %s and q method %s | qsea:%s, qair:%s', meth, qmeth,
                  np.nanmedian(qsea), np.nanmedian(qair))
sbiri's avatar
sbiri committed
223 224 225 226 227 228 229 230 231 232 233
    if (np.all(np.isnan(qsea)) or np.all(np.isnan(qair))):
        print("qsea and qair cannot be nan")
    # first guesses
    dt = Ta - sst
    dq = qair - qsea
    t10n, q10n = np.copy(Ta), np.copy(qair)
    tv10n = t10n*(1 + 0.61*q10n)
    #  Zeng et al. 1998
    tv=th*(1.+0.61*qair)   # virtual potential T
    dtv=dt*(1.+0.61*qair)+0.61*th*dq
    # ------------
234
    rho = P*100/(287.1*tv10n)
235
    lv = (2.501-0.00237*(sst-CtoK))*1e6
sbiri's avatar
sbiri committed
236 237
    cp = 1004.67*(1 + 0.00084*qsea)
    u10n = np.copy(spd)
sbiri's avatar
sbiri committed
238
    monob = -100*np.ones(spd.shape)
sbiri's avatar
sbiri committed
239
    cdn = cdn_calc(u10n, Ta, None, lat, meth)
240 241 242 243 244 245 246
    ctn, ct, cqn, cq = (np.zeros(spd.shape)*np.nan, np.zeros(spd.shape)*np.nan,
                        np.zeros(spd.shape)*np.nan, np.zeros(spd.shape)*np.nan)
    psim, psit, psiq = (np.zeros(spd.shape), np.zeros(spd.shape),
                        np.zeros(spd.shape))
    cd = cd_calc(cdn, h_in[0], ref_ht, psim)
    tsr, tsrv = np.zeros(spd.shape), np.zeros(spd.shape)
    qsr = np.zeros(spd.shape)
sbiri's avatar
sbiri committed
247
    # cskin parameters
248
    tkt = 0.001*np.ones(T.shape)
sbiri's avatar
sbiri committed
249
    Rnl = 0.97*(5.67e-8*np.power(sst-0.3*cskin+CtoK, 4)-Rl)
sbiri's avatar
sbiri committed
250
    dter = np.ones(T.shape)*0.3
251
    dqer = dter*0.622*lv*qsea/(287.1*np.power(sst, 2))
252
    if (gust[0] == 1 and meth == "UA"):
253 254
        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)))
255
    elif (gust[0] == 1):
256
        wind = np.sqrt(np.power(np.copy(spd), 2)+np.power(0.5, 2))
257
    elif (gust[0] == 0):
258
        wind = np.copy(spd)
259 260 261 262 263 264
    usr = np.sqrt(cd*np.power(wind, 2))
    Rb = g*h_in[0]*dtv/((T+CtoK)*np.power(wind, 2))
    zo = 0.0001*np.ones(spd.shape)
    zot, zoq = 0.0001*np.ones(spd.shape), 0.0001*np.ones(spd.shape)
    monob = -100*np.ones(spd.shape)  # Monin-Obukhov length
    zol = h_in[0]/monob
sbiri's avatar
sbiri committed
265
    tsr = (dt+dter*cskin)*kappa/(np.log(h_in[1]/zot) -
266
                                 psit_calc(h_in[1]/monob, meth))
sbiri's avatar
sbiri committed
267
    qsr = (dq+dqer*cskin)*kappa/(np.log(h_in[2]/zoq) -
268
                                 psit_calc(h_in[2]/monob, meth))
269 270
    it, ind = 0, np.where(spd > 0)
    ii, itera = True, np.zeros(spd.shape)*np.nan
271 272 273
    tau = 0.05*np.ones(spd.shape)
    sensible = 5*np.ones(spd.shape)
    latent = 65*np.ones(spd.shape)
sbiri's avatar
sbiri committed
274 275 276 277
    while np.any(ii):
        it += 1
        if it > n:
            break
278 279 280 281 282 283 284
        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)])
sbiri's avatar
sbiri committed
285 286
        cdn[ind] = cdn_calc(u10n[ind], Ta[ind], None, lat[ind], meth)
        if (np.all(np.isnan(cdn))):
287
            break
sbiri's avatar
sbiri committed
288 289
            logging.info('break %s at iteration %s cdn<0', meth, it)
        zo[ind] = ref_ht/np.exp(kappa/np.sqrt(cdn[ind]))
290 291 292
        psim[ind] = psim_calc(h_in[0, ind]/monob[ind], meth)
        cd[ind] = cd_calc(cdn[ind], h_in[0, ind], ref_ht, psim[ind])
        ctn[ind], cqn[ind] = ctcqn_calc(h_in[1, ind]/monob[ind], cdn[ind],
sbiri's avatar
sbiri committed
293
                                        u10n[ind], zo[ind], Ta[ind], meth)
294 295 296 297
        zot[ind] = ref_ht/(np.exp(np.power(kappa, 2) /
                           (ctn[ind]*np.log(ref_ht/zo[ind]))))
        zoq[ind] = ref_ht/(np.exp(np.power(kappa, 2) /
                           (cqn[ind]*np.log(ref_ht/zo[ind]))))
298 299
        psit[ind] = psit_calc(h_in[1, ind]/monob[ind], meth)
        psiq[ind] = psit_calc(h_in[2, ind]/monob[ind], meth)
sbiri's avatar
sbiri committed
300
        ct[ind], cq[ind] = ctcq_calc(cdn[ind], cd[ind], ctn[ind], cqn[ind],
301 302
                                      h_in[1, ind], h_in[2, ind], ref_ht,
                                      psit[ind], psiq[ind])
sbiri's avatar
sbiri committed
303
        if (meth == "UA"):
304 305
            usr[ind] = np.where(h_in[0, ind]/monob[ind] < -1.574,
                                kappa*wind[ind] /
sbiri's avatar
sbiri committed
306 307 308
                                (np.log(-1.574*monob[ind]/zo[ind]) -
                                psim_calc(-1.574, meth) +
                                psim_calc(zo[ind]/monob[ind], meth) +
309
                                1.14*(np.power(-h_in[0, ind]/monob[ind], 1/3) -
310
                                np.power(1.574, 1/3))),
311 312 313 314
                                np.where((h_in[0, ind]/monob[ind] > -1.574) &
                                (h_in[0, ind]/monob[ind] < 0),
                                kappa*wind[ind]/(np.log(h_in[0, ind]/zo[ind]) -
                                psim_calc(h_in[0, ind]/monob[ind], meth) +
sbiri's avatar
sbiri committed
315
                                psim_calc(zo[ind]/monob[ind], meth)),
316 317 318 319 320 321
                                np.where((h_in[0, ind]/monob[ind] > 0) &
                                (h_in[0, ind]/monob[ind] < 1),
                                kappa*wind[ind]/(np.log(h_in[0, ind]/zo[ind]) +
                                5*h_in[0, ind]/monob[ind]-5*zo[ind] /
                                monob[ind]), kappa*wind[ind] /
                                (np.log(monob[ind]/zo[ind]) +
sbiri's avatar
sbiri committed
322
                                5-5*zo[ind]/monob[ind] +
323 324
                                5*np.log(h_in[0, ind]/monob[ind]) +
                                h_in[0, ind]/monob[ind]-1))))
sbiri's avatar
sbiri committed
325
                                # Zeng et al. 1998 (7-10)
326 327
            tsr[ind] = np.where(h_in[1, ind]/monob[ind] < -0.465,
                                kappa*(dt[ind] +
sbiri's avatar
sbiri committed
328
                                dter[ind]*cskin) /
sbiri's avatar
sbiri committed
329 330
                                (np.log((-0.465*monob[ind])/zot[ind]) -
                                psit_calc(-0.465, meth)+0.8 *
331
                                (np.power(0.465, -1/3) -
332 333 334
                                np.power(-h_in[1, ind]/monob[ind], -1/3))),
                                np.where((h_in[1, ind]/monob[ind] > -0.465) &
                                (h_in[1, ind]/monob[ind] < 0),
sbiri's avatar
sbiri committed
335
                                kappa*(dt[ind]+dter[ind]*cskin) /
336 337
                                (np.log(h_in[1, ind]/zot[ind]) -
                                psit_calc(h_in[1, ind]/monob[ind], meth) +
sbiri's avatar
sbiri committed
338
                                psit_calc(zot[ind]/monob[ind], meth)),
339 340
                                np.where((h_in[1, ind]/monob[ind] > 0) &
                                (h_in[1, ind]/monob[ind] < 1),
sbiri's avatar
sbiri committed
341
                                kappa*(dt[ind]+dter[ind]*cskin) /
342 343 344
                                (np.log(h_in[1, ind]/zot[ind]) +
                                5*h_in[1, ind]/monob[ind]-5*zot[ind] /
                                monob[ind]),
sbiri's avatar
sbiri committed
345
                                kappa*(dt[ind]+dter[ind]*cskin) /
346 347
                                (np.log(monob[ind]/zot[ind])+5 -
                                5*zot[ind]/monob[ind] +
348 349
                                5*np.log(h_in[1, ind]/monob[ind]) +
                                h_in[1, ind]/monob[ind]-1))))
sbiri's avatar
sbiri committed
350
                                # Zeng et al. 1998 (11-14)
351 352
            qsr[ind] = np.where(h_in[2, ind]/monob[ind] < -0.465,
                                kappa*(dq[ind] +
sbiri's avatar
sbiri committed
353
                                dqer[ind]*cskin) /
sbiri's avatar
sbiri committed
354 355 356
                                (np.log((-0.465*monob[ind])/zoq[ind]) -
                                psit_calc(-0.465, meth) +
                                psit_calc(zoq[ind]/monob[ind], meth) +
357
                                0.8*(np.power(0.465, -1/3) -
358 359 360
                                np.power(-h_in[2, ind]/monob[ind], -1/3))),
                                np.where((h_in[2, ind]/monob[ind] > -0.465) &
                                (h_in[2, ind]/monob[ind] < 0),
sbiri's avatar
sbiri committed
361
                                kappa*(dq[ind]+dqer[ind]*cskin) /
362 363
                                (np.log(h_in[1, ind]/zot[ind]) -
                                psit_calc(h_in[2, ind]/monob[ind], meth) +
sbiri's avatar
sbiri committed
364
                                psit_calc(zoq[ind]/monob[ind], meth)),
365 366
                                np.where((h_in[2, ind]/monob[ind] > 0) &
                                (h_in[2, ind]/monob[ind]<1), kappa*(dq[ind] +
sbiri's avatar
sbiri committed
367
                                dqer[ind]*cskin) /
368 369
                                (np.log(h_in[1, ind]/zoq[ind]) +
                                5*h_in[2, ind]/monob[ind]-5*zoq[ind]/monob[ind]),
sbiri's avatar
sbiri committed
370
                                kappa*(dq[ind]+dqer[ind]*cskin) /
371
                                (np.log(monob[ind]/zoq[ind])+5 -
sbiri's avatar
sbiri committed
372
                                5*zoq[ind]/monob[ind] +
373 374
                                5*np.log(h_in[2, ind]/monob[ind]) +
                                h_in[2, ind]/monob[ind]-1))))
sbiri's avatar
sbiri committed
375
        elif (meth == "C30" or meth == "C35" or meth == "C40"):
376 377
            usr[ind] = (wind[ind]*kappa/(np.log(h_in[0, ind]/zo[ind]) -
                        psiu_26(h_in[0, ind]/monob[ind], meth)))
sbiri's avatar
sbiri committed
378 379 380 381
            qsr[ind] = ((dq[ind]+dqer[ind]*cskin)*(kappa/(np.log(h_in[2, ind] /
                        zoq[ind])-psit_26(h_in[2, ind]/monob[ind]))))
            tsr[ind] = ((dt[ind]+dter[ind]*cskin)*(kappa/(np.log(h_in[1, ind] /
                        zot[ind])-psit_26(h_in[1, ind]/monob[ind]))))
sbiri's avatar
sbiri committed
382
        else:
383 384
            usr[ind] = (wind[ind]*kappa/(np.log(h_in[0, ind]/zo[ind]) -
                        psim_calc(h_in[0, ind]/monob[ind], meth)))
sbiri's avatar
sbiri committed
385 386 387 388
            #           np.sqrt(cd[ind]*np.power(wind[ind], 2))
            qsr[ind] = cq[ind]*wind[ind]*(dq[ind]+dqer[ind]*cskin)/usr[ind]
            tsr[ind] = ct[ind]*wind[ind]*(dt[ind]+dter[ind]*cskin)/usr[ind]
        if (cskin == 1):
389 390 391 392 393 394 395 396 397 398
            dter[ind], dqer[ind], tkt[ind] = get_skin(sst[ind], qsea[ind],
                                                      rho[ind], Rl[ind],
                                                      Rs[ind], Rnl[ind],
                                                      cp[ind], lv[ind],
                                                      np.copy(tkt[ind]),
                                                      usr[ind], tsr[ind],
                                                      qsr[ind], lat[ind])
        else:
           dter[ind] = np.zeros(sst[ind].shape)
           dqer[ind] = np.zeros(sst[ind].shape)
399
           tkt[ind] = 0.001*np.ones(T[ind].shape)
sbiri's avatar
sbiri committed
400 401 402 403 404 405
        logging.info('method %s | dter = %s | dqer = %s | tkt = %s | Rnl = %s '
                     '| usr = %s | tsr = %s | qsr = %s', meth,
                     np.nanmedian(dter), np.nanmedian(dqer),
                     np.nanmedian(tkt), np.nanmedian(Rnl),
                     np.nanmedian(usr), np.nanmedian(tsr),
                     np.nanmedian(qsr))
406
        Rnl[ind] = 0.97*(5.67e-8*np.power(sst[ind]-CtoK -
sbiri's avatar
sbiri committed
407
                          dter[ind]*cskin+CtoK, 4)-Rl[ind])
408 409 410 411
        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]))
sbiri's avatar
sbiri committed
412
        tv10n[ind] = t10n[ind]*(1+0.61*q10n[ind])
413 414 415 416 417 418
        tsrv[ind], monob[ind] = get_L(L, lat[ind], usr[ind], tsr[ind],
                                      qsr[ind], t10n[ind], tv10n[ind],
                                      qair[ind], h_in[:, ind], T[ind], Ta[ind],
                                      th[ind], tv[ind], sst[ind], dt[ind],
                                      dq[ind], wind[ind], np.copy(monob[ind]),
                                      meth)
419 420 421
        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)
422
        if (gust[0] == 1 and meth == "UA"):
sbiri's avatar
sbiri committed
423
            wind[ind] = np.where(dtv[ind] >= 0, np.where(spd[ind] > 0.1,
424 425 426 427 428 429 430
                                  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)
        elif (gust[0] == 1 and (meth == "C30" or meth == "C35" or
                                meth == "C40")):
sbiri's avatar
sbiri committed
431
            wind[ind] = np.sqrt(np.power(np.copy(spd[ind]), 2) +
432 433 434
                                np.power(get_gust(gust[1], Ta[ind], usr[ind],
                                tsrv[ind], gust[2], lat[ind]), 2))
        elif (gust[0] == 1):
435
            wind[ind] = np.sqrt(np.power(np.copy(spd[ind]), 2) +
436 437 438
                                np.power(get_gust(gust[1], Ta[ind], usr[ind],
                                tsrv[ind], gust[2], lat[ind]), 2))
        elif (gust[0] == 0):
439
            wind[ind] = np.copy(spd[ind])
440 441
        u10n[ind] = wind[ind]-usr[ind]/kappa*(np.log(h_in[0, ind]/10) -
                                              psim[ind])
sbiri's avatar
sbiri committed
442
        u10n = np.where(u10n < 0, np.nan, u10n)
sbiri's avatar
sbiri committed
443
        itera[ind] = np.ones(1)*it
444 445 446 447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467
        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'):
            new = np.array([np.copy(u10n), np.copy(t10n), np.copy(q10n)])
        elif (tol[0] == 'all'):
            new = np.array([np.copy(u10n), np.copy(t10n), np.copy(q10n),
                            np.copy(tau), np.copy(sensible), np.copy(latent)])
        d = np.fabs(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]))
468 469
        if (ind[0].size == 0):
            ii = False
sbiri's avatar
sbiri committed
470
        else:
471 472
            ii = True
    logging.info('method %s | # of iterations:%s', meth, it)
sbiri's avatar
sbiri committed
473
    logging.info('method %s | # of points that did not converge :%s', meth,
474
                  ind[0].size)
sbiri's avatar
sbiri committed
475
    # calculate output parameters
476
    rho = (0.34838*P)/(tv10n)
sbiri's avatar
sbiri committed
477
    t10n = t10n-(273.16+tlapse*ref_ht)
478 479 480
    zo = ref_ht/np.exp(kappa/cdn**0.5)
    zot = ref_ht/(np.exp(kappa**2/(ctn*np.log(ref_ht/zo))))
    zoq = ref_ht/(np.exp(kappa**2/(cqn*np.log(ref_ht/zo))))
481 482 483 484 485 486 487
    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)))
    tref = tref-(273.16+tlapse*h_out[1])
    qref = (qair-qsr/kappa*(np.log(h_in[2]/h_out[2]) -
            psit+psit_calc(h_out[2]/monob, meth)))
488
    res = np.zeros((28, len(spd)))
sbiri's avatar
sbiri committed
489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504
    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
505 506 507 508 509 510 511 512
    res[16][:] = psiq
    res[17][:] = u10n
    res[18][:] = t10n
    res[19][:] = tv10n
    res[20][:] = q10n
    res[21][:] = zo
    res[22][:] = zot
    res[23][:] = zoq
513 514 515
    res[24][:] = uref
    res[25][:] = tref
    res[26][:] = qref
516
    res[27][:] = itera
517 518 519 520 521
    if (out == 0):
        res[:, ind] = np.nan
    # set missing values where data have non acceptable values
    res = np.where(spd < 0, np.nan, res)

522
    return res, ind