AirSeaFluxCode.py 32.9 KB
Newer Older
sbiri's avatar
sbiri committed
1
import warnings
sbiri's avatar
sbiri committed
2
import numpy as np
3
import pandas as pd
sbiri's avatar
sbiri committed
4
import logging
5
from hum_subs import (get_hum, gamma)
sbiri's avatar
sbiri committed
6 7 8
from util_subs import *
from flux_subs import *
from cs_wl_subs import *
sbiri's avatar
sbiri committed
9 10


sbiri's avatar
sbiri committed
11 12
class S88:
    def _wind_iterate(self, ind):
sbiri's avatar
sbiri committed
13
        if self.gust[0] in range(1, 6):
sbiri's avatar
sbiri committed
14
            self.wind[ind] = np.sqrt(np.power(np.copy(self.spd[ind]), 2) +
15 16 17 18 19
                                     np.power(get_gust(
                                         self.gust[1], self.gust[2],
                                         self.gust[3], self.theta[ind],
                                         self.usr[ind], self.tsrv[ind],
                                         self.grav[ind]), 2))
sbiri's avatar
sbiri committed
20 21
            if self.gust[0] in [3, 4]:
                # self.GustFact[ind] = 1
22
                # option to not remove GustFact
23
                self.u10n[ind] = self.wind[ind]-self.usr[ind]/kappa*(
24
                    np.log(self.h_in[0, ind]/self.ref10)-self.psim[ind])
sbiri's avatar
sbiri committed
25 26 27 28 29
            else:
                self.GustFact = np.sqrt(self.wind/self.spd)
                self.u10n[ind] = self.spd[ind]-self.usr[ind]/kappa / \
                    self.GustFact[ind]*(np.log(self.h_in[0, ind]/self.ref10) -
                                        self.psim[ind])
sbiri's avatar
sbiri committed
30
        else:
sbiri's avatar
sbiri committed
31
            # initalisation of wind
sbiri's avatar
sbiri committed
32
            self.wind[ind] = np.copy(self.spd[ind])
sbiri's avatar
sbiri committed
33
            self.u10n[ind] = self.wind[ind]-self.usr[ind]/kappa*(
34
                np.log(self.h_in[0, ind]/self.ref10)-self.psim[ind])
sbiri's avatar
sbiri committed
35

sbiri's avatar
sbiri committed
36 37 38 39
    def get_heights(self, hin, hout=10):
        self.hout = hout
        self.hin = hin
        self.h_in = get_heights(hin, len(self.spd))
sbiri's avatar
sbiri committed
40
        self.h_out = get_heights(self.hout, 1)
sbiri's avatar
sbiri committed
41

sbiri's avatar
sbiri committed
42 43 44
    def get_specHumidity(self, qmeth="Buck2"):
        self.qair, self.qsea = get_hum(self.hum, self.T, self.SST, self.P,
                                       qmeth)
sbiri's avatar
sbiri committed
45 46
        if (np.all(np.isnan(self.qsea)) or np.all(np.isnan(self.qair))):
            raise ValueError("qsea and qair cannot be nan")
sbiri's avatar
sbiri committed
47 48 49
        self.dq_in = self.qair-self.qsea
        self.dq_full = self.qair-self.qsea

sbiri's avatar
sbiri committed
50 51 52 53
        # Set lapse rate and Potential Temperature (now we have humdity)
        self.cp = 1004.67*(1+0.00084*self.qsea)
        self.tlapse = gamma("dry", self.SST, self.T, self.qair/1000, self.cp)
        self.theta = np.copy(self.T)+self.tlapse*self.h_in[1]
sbiri's avatar
sbiri committed
54 55
        self.dt_in = self.theta-self.SST
        self.dt_full = self.theta-self.SST
sbiri's avatar
sbiri committed
56 57

    def _fix_coolskin_warmlayer(self, wl, cskin, skin, Rl, Rs):
sbiri's avatar
sbiri committed
58
        skin = self.skin if skin is None else skin
sbiri's avatar
sbiri committed
59 60
        assert wl in [0, 1], "wl not valid"
        assert cskin in [0, 1], "cskin not valid"
sbiri's avatar
sbiri committed
61 62
        assert skin in ["C35", "ecmwf", "Beljaars"], "Skin value not valid"

sbiri's avatar
sbiri committed
63 64
        if ((cskin == 1 or wl == 1) and
            (np.all(Rl == None) or np.all(np.isnan(Rl))) and
65
                ((np.all(Rs == None) or np.all(np.isnan(Rs))))):
sbiri's avatar
sbiri committed
66 67 68 69 70
            print("Cool skin/warm layer is switched ON; "
                  "Radiation input should not be empty")
            raise

        self.wl = wl
sbiri's avatar
sbiri committed
71 72
        self.cskin = cskin
        self.skin = skin
sbiri's avatar
sbiri committed
73 74
        self.Rs = np.full(self.spd.shape, np.nan) if Rs is None else Rs
        self.Rl = np.full(self.spd.shape, np.nan) if Rl is None else Rl
sbiri's avatar
sbiri committed
75

sbiri's avatar
sbiri committed
76
    def set_coolskin_warmlayer(self, wl=0, cskin=0, skin=None, Rl=None,
sbiri's avatar
sbiri committed
77
                               Rs=None):
sbiri's avatar
sbiri committed
78
        wl = 0 if wl is None else wl
sbiri's avatar
sbiri committed
79 80
        if hasattr(self, "skin") == False:
            self.skin = "C35"
sbiri's avatar
sbiri committed
81 82
        self._fix_coolskin_warmlayer(wl, cskin, skin, Rl, Rs)

sbiri's avatar
sbiri committed
83 84
    def _update_coolskin_warmlayer(self, ind):
        if self.cskin == 1:
85 86 87 88 89
            self.dter[ind], self.tkt[ind] = cs(np.copy(
                self.SST[ind]), np.copy(self.tkt[ind]), self.rho[ind],
                self.Rs[ind], self.Rnl[ind], self.cp[ind], self.lv[ind],
                self.usr[ind], self.tsr[ind], self.qsr[ind], self.grav[ind],
                self.skin)
sbiri's avatar
sbiri committed
90 91
            self.dqer[ind] = get_dqer(self.dter[ind], self.SST[ind],
                                      self.qsea[ind], self.lv[ind])
sbiri's avatar
sbiri committed
92 93
            self.skt[ind] = np.copy(self.SST[ind])+self.dter[ind]
            self.skq[ind] = np.copy(self.qsea[ind])+self.dqer[ind]
sbiri's avatar
sbiri committed
94 95 96 97 98 99 100 101 102 103
            if self.wl == 1:
                self.dtwl[ind] = wl_ecmwf(
                    self.rho[ind], self.Rs[ind], self.Rnl[ind], self.cp[ind],
                    self.lv[ind], self.usr[ind], self.tsr[ind], self.qsr[ind],
                    np.copy(self.SST[ind]), np.copy(self.skt[ind]),
                    np.copy(self.dter[ind]), self.grav[ind])
                self.skt[ind] = (np.copy(self.SST[ind])+self.dter[ind] +
                                 self.dtwl[ind])
                self.dqer[ind] = get_dqer(self.dter[ind], self.skt[ind],
                                          self.qsea[ind], self.lv[ind])
sbiri's avatar
sbiri committed
104 105 106 107 108 109 110 111
                self.skq[ind] = np.copy(self.qsea[ind])+self.dqer[ind]
        else:
            self.dter[ind] = np.zeros(self.SST[ind].shape)
            self.dqer[ind] = np.zeros(self.SST[ind].shape)
            self.dtwl[ind] = np.zeros(self.SST[ind].shape)
            self.tkt[ind] = np.zeros(self.SST[ind].shape)

    def _first_guess(self):
112
        # reference height
sbiri's avatar
sbiri committed
113 114 115 116 117 118 119 120 121 122 123 124
        self.ref10 = 10

        #  first guesses
        self.t10n, self.q10n = np.copy(self.theta), np.copy(self.qair)
        self.rho = self.P*100/(287.1*self.t10n*(1+0.6077*self.q10n))
        self.lv = (2.501-0.00237*(self.SST-CtoK))*1e6  # J/kg

        #  Zeng et al. 1998
        self.tv = self.theta*(1+0.6077*self.qair)   # virtual potential T
        self.dtv = self.dt_in*(1+0.6077*self.qair)+0.6077*self.theta*self.dq_in

        # Set the wind array
sbiri's avatar
sbiri committed
125
        self.wind = np.sqrt(np.power(np.copy(self.spd), 2)+0.25)
sbiri's avatar
sbiri committed
126 127 128 129
        self.GustFact = self.wind*0+1

        # Rb eq. 11 Grachev & Fairall 1997, use air temp height
        # use self.tv??  adjust wind to T-height?
sbiri's avatar
sbiri committed
130
        Rb = self.grav*self.h_in[1]*self.dtv/(self.T*np.power(self.wind, 2))
131 132
        # eq. 12 Grachev & Fairall 1997   # DO.THIS
        self.monob = self.h_in[1]/12.0/Rb
sbiri's avatar
sbiri committed
133 134 135

        # ------------

136 137
        dummy_array = lambda val : np.full(self.T.shape, val)*self.msk
        # def dummy_array(val): return np.full(self.T.shape, val)*self.msk
sbiri's avatar
sbiri committed
138
        if self.cskin + self.wl > 0:
sbiri's avatar
sbiri committed
139 140
            self.dter, self.tkt, self.dtwl = [
                dummy_array(x) for x in (-0.3, 0.001, 0.3)]
141 142
            self.dqer = get_dqer(self.dter, self.SST, self.qsea,
                                 self.lv)
sbiri's avatar
sbiri committed
143 144
            self.Rnl = 0.97*(self.Rl-5.67e-8*np.power(
                self.SST-0.3*self.cskin, 4))
sbiri's avatar
sbiri committed
145 146
            self.Qs = 0.945*self.Rs
        else:
sbiri's avatar
sbiri committed
147 148 149 150
            self.dter, self.dqer, self.dtwl = [
                dummy_array(x) for x in (0.0, 0.0, 0.0)]
            self.Rnl, self.Qs, self.tkt = [
                np.empty(self.arr_shp)*self.msk for _ in range(3)]
sbiri's avatar
sbiri committed
151 152 153 154 155
        self.skt = np.copy(self.SST)
        self.skq = np.copy(self.qsea)

        self.u10n = np.copy(self.wind)
        self.usr = 0.035*self.u10n
sbiri's avatar
sbiri committed
156 157
        self.cd10n, self.zo = cdn_calc(
            self.u10n, self.usr, self.theta, self.grav, self.meth)
sbiri's avatar
sbiri committed
158 159
        self.psim = psim_calc(self.h_in[0]/self.monob, self.meth)
        self.cd = cd_calc(self.cd10n, self.h_in[0], self.ref10, self.psim)
160
        self.usr =np.sqrt(self.cd*np.power(self.wind, 2))
sbiri's avatar
sbiri committed
161 162 163 164
        self.zot, self.zoq, self.tsr, self.qsr = [
            np.empty(self.arr_shp)*self.msk for _ in range(4)]
        self.ct10n, self.cq10n, self.ct, self.cq = [
            np.empty(self.arr_shp)*self.msk for _ in range(4)]
165
        self.tv10n = self.zot
sbiri's avatar
sbiri committed
166 167 168 169 170 171

    def iterate(self, maxiter=10, tol=None):
        if maxiter < 5:
            warnings.warn("Iteration number <5 - resetting to 5.")
            maxiter = 5

sbiri's avatar
sbiri committed
172 173 174 175
        # Decide which variables to use in tolerances based on tolerance
        # specification
        tol = ['all', 0.01, 0.01, 1e-05, 1e-3,
               0.1, 0.1] if tol is None else tol
sbiri's avatar
sbiri committed
176 177
        assert tol[0] in ['flux', 'ref', 'all'], "unknown tolerance input"

sbiri's avatar
sbiri committed
178 179
        old_vars = {"flux": ["tau", "sensible", "latent"],
                    "ref": ["u10n", "t10n", "q10n"]}
sbiri's avatar
sbiri committed
180 181 182
        old_vars["all"] = old_vars["ref"] + old_vars["flux"]
        old_vars = old_vars[tol[0]]

sbiri's avatar
sbiri committed
183 184
        new_vars = {"flux": ["tau", "sensible", "latent"],
                    "ref": ["u10n", "t10n", "q10n"]}
sbiri's avatar
sbiri committed
185 186 187 188 189 190 191 192 193 194
        new_vars["all"] = new_vars["ref"] + new_vars["flux"]
        new_vars = new_vars[tol[0]]
        # extract tolerance values by deleting flag from tol
        tvals = np.delete(np.copy(tol), 0)
        tol_vals = list([float(tt) for tt in tvals])

        ind = np.where(self.spd > 0)
        it = 0

        # Setup empty arrays
sbiri's avatar
sbiri committed
195 196
        self.tsrv, self.psim, self.psit, self.psiq = [
            np.zeros(self.arr_shp)*self.msk for _ in range(4)]
sbiri's avatar
sbiri committed
197 198

        # extreme values for first comparison
199
        dummy_array = lambda val : np.full(self.T.shape, val)*self.msk
200
        # you can use def instead of lambda
sbiri's avatar
sbiri committed
201 202 203
        # def dummy_array(val): return np.full(self.arr_shp, val)*self.msk
        self.itera, self.tau, self.sensible, self.latent = [
            dummy_array(x) for x in (-1, 1e+99, 1e+99, 1e+99)]
sbiri's avatar
sbiri committed
204 205 206 207 208 209 210 211 212 213

        # Generate the first guess values
        self._first_guess()

        #  iteration loop
        ii = True
        while ii & (it < maxiter):
            it += 1

            # Set the old variables (for comparison against "new")
sbiri's avatar
sbiri committed
214
            old = np.array([np.copy(getattr(self, i)) for i in old_vars])
sbiri's avatar
sbiri committed
215 216

            # Calculate cdn
sbiri's avatar
sbiri committed
217 218 219 220 221
            self.cd10n[ind], self.zo[ind] = cdn_calc(
                self.u10n[ind], self.usr[ind], self.theta[ind], self.grav[ind],
                self.meth)

            if np.all(np.isnan(self.cd10n)):
sbiri's avatar
sbiri committed
222
                logging.info('break %s at iteration %s cd10n<0', meth, it)
sbiri's avatar
sbiri committed
223 224 225 226 227 228
                break

            self.psim[ind] = psim_calc(
                self.h_in[0, ind]/self.monob[ind], self.meth)
            self.cd[ind] = cd_calc(
                self.cd10n[ind], self.h_in[0, ind], self.ref10, self.psim[ind])
229 230
            # Update the wind values
            self._wind_iterate(ind)
sbiri's avatar
sbiri committed
231 232

            # temperature
sbiri's avatar
sbiri committed
233 234 235 236 237 238 239 240
            self.ct10n[ind], self.zot[ind] = ctqn_calc(
                "ct", self.h_in[1, ind]/self.monob[ind], self.cd10n[ind],
                self.usr[ind], self.zo[ind], self.theta[ind], self.meth)
            self.psit[ind] = psit_calc(
                self.h_in[1, ind]/self.monob[ind], self.meth)
            self.ct[ind] = ctq_calc(
                self.cd10n[ind], self.cd[ind], self.ct10n[ind],
                self.h_in[1, ind], self.ref10, self.psit[ind])
241
            # humidity
sbiri's avatar
sbiri committed
242 243 244 245 246 247 248 249
            self.cq10n[ind], self.zoq[ind] = ctqn_calc(
                "cq", self.h_in[2, ind]/self.monob[ind], self.cd10n[ind],
                self.usr[ind], self.zo[ind], self.theta[ind], self.meth)
            self.psiq[ind] = psit_calc(
                self.h_in[2, ind]/self.monob[ind], self.meth)
            self.cq[ind] = ctq_calc(
                self.cd10n[ind], self.cd[ind], self.cq10n[ind],
                self.h_in[2, ind], self.ref10, self.psiq[ind])
250

sbiri's avatar
sbiri committed
251 252 253 254 255
            # Some parameterizations set a minimum on parameters
            try:
                self._minimum_params()
            except AttributeError:
                pass
sbiri's avatar
sbiri committed
256 257 258

            self.dt_full[ind] = self.dt_in[ind] - \
                self.dter[ind]*self.cskin - self.dtwl[ind]*self.wl
sbiri's avatar
sbiri committed
259
            self.dq_full[ind] = self.dq_in[ind] - self.dqer[ind]*self.cskin
sbiri's avatar
sbiri committed
260 261 262 263 264
            self.usr[ind], self.tsr[ind], self.qsr[ind] = get_strs(
                self.h_in[:, ind], self.monob[ind], self.wind[ind],
                self.zo[ind], self.zot[ind], self.zoq[ind], self.dt_full[ind],
                self.dq_full[ind], self.cd[ind], self.ct[ind], self.cq[ind],
                self.meth)
sbiri's avatar
sbiri committed
265 266 267

            # Update CS/WL parameters
            self._update_coolskin_warmlayer(ind)
sbiri's avatar
sbiri committed
268

sbiri's avatar
sbiri committed
269
            # Logging output
sbiri's avatar
sbiri committed
270 271 272 273 274 275 276 277
            log_vars = {"dter": 2, "dqer": 7, "tkt": 2,
                        "Rnl": 2, "usr": 3, "tsr": 4, "qsr": 7}
            log_vars = [np.round(np.nanmedian(getattr(self, V)), R)
                        for V, R in log_vars.items()]
            log_vars.insert(0, self.meth)
            logging.info(
                'method {} | dter = {} | dqer = {} | tkt = {} | Rnl = {} |'
                ' usr = {} | tsr = {} | qsr = {}'.format(*log_vars))
sbiri's avatar
sbiri committed
278 279

            if self.cskin + self.wl > 0:
sbiri's avatar
sbiri committed
280 281 282
                self.Rnl[ind] = 0.97*(self.Rl[ind]-5.67e-8 *
                                      np.power(self.SST[ind] +
                                               self.dter[ind]*self.cskin, 4))
sbiri's avatar
sbiri committed
283 284
            # not sure how to handle lapse/potemp
            # well-mixed in potential temperature ...
285
            self.t10n[ind] = self.theta[ind]-self.tlapse[ind]*self.ref10 - \
sbiri's avatar
sbiri committed
286
                self.tsr[ind]/kappa * \
287
                (np.log(self.h_in[1, ind]/self.ref10)-self.psit[ind])
288
            self.q10n[ind] = self.qair[ind]-self.qsr[ind]/kappa * \
sbiri's avatar
sbiri committed
289
                (np.log(self.h_in[2, ind]/self.ref10)-self.psiq[ind])
sbiri's avatar
sbiri committed
290 291

            # update stability info
sbiri's avatar
sbiri committed
292 293 294 295 296 297
            self.tsrv[ind] = get_tsrv(
                self.tsr[ind], self.qsr[ind], self.theta[ind], self.qair[ind])
            self.Rb[ind] = get_Rb(
                self.grav[ind], self.usr[ind], self.h_in[0, ind],
                self.h_in[1, ind], self.tv[ind], self.dtv[ind], self.wind[ind],
                self.monob[ind], self.meth)
sbiri's avatar
sbiri committed
298
            if self.L == "tsrv":
sbiri's avatar
sbiri committed
299 300 301
                self.monob[ind] = get_Ltsrv(
                    self.tsrv[ind], self.grav[ind], self.tv[ind],
                    self.usr[ind])
sbiri's avatar
sbiri committed
302
            else:
sbiri's avatar
sbiri committed
303
                self.monob[ind] = get_LRb(
304
                    self.Rb[ind], self.h_in[1, ind], self.monob[ind],
sbiri's avatar
sbiri committed
305
                    self.zo[ind], self.zot[ind], self.meth)
sbiri's avatar
sbiri committed
306 307 308

            # Update the wind values
            self._wind_iterate(ind)
sbiri's avatar
sbiri committed
309

sbiri's avatar
sbiri committed
310
            # make sure you allow small negative values convergence
sbiri's avatar
sbiri committed
311 312
            if it == 1:
                self.u10n = np.where(self.u10n < 0, 0.5, self.u10n)
sbiri's avatar
sbiri committed
313

sbiri's avatar
sbiri committed
314
            self.itera[ind] = np.full(1, it)
315 316 317
            self.tau = self.rho*np.power(self.usr, 2)
            self.sensible = self.rho*self.cp*self.usr*self.tsr
            self.latent = self.rho*self.lv*self.usr*self.qsr
sbiri's avatar
sbiri committed
318 319

            # Set the new variables (for comparison against "old")
sbiri's avatar
sbiri committed
320
            new = np.array([np.copy(getattr(self, i)) for i in new_vars])
sbiri's avatar
sbiri committed
321

sbiri's avatar
sbiri committed
322
            if it > 2:  # force at least three iterations
323
                d = np.abs(new-old)  # change over this iteration
sbiri's avatar
sbiri committed
324 325 326 327
                for ii in range(0, len(tol_vals)):
                    d[ii, ] = d[ii, ]/tol_vals[ii]  # ratio to tolerance
                # identifies non-convergence
                ind = np.where(d.max(axis=0) >= 1)
sbiri's avatar
sbiri committed
328 329 330 331 332 333 334 335

            self.ind = np.copy(ind)
            ii = False if (ind[0].size == 0) else True
            # End of iteration loop

        self.itera[ind] = -1
        self.itera = np.where(self.itera > maxiter, -1, self.itera)
        logging.info('method %s | # of iterations:%s', self.meth, it)
sbiri's avatar
sbiri committed
336 337
        logging.info('method %s | # of points that did not converge :%s \n',
                     self.meth, self.ind[0].size)
sbiri's avatar
sbiri committed
338 339

    def _get_humidity(self):
340
        """Calculate RH used for flagging purposes & output."""
sbiri's avatar
sbiri committed
341
        if self.hum[0] in ('rh', 'no'):
sbiri's avatar
sbiri committed
342
            self.rh = self.hum[1]
sbiri's avatar
sbiri committed
343
        elif self.hum[0] == 'Td':
sbiri's avatar
sbiri committed
344 345 346 347 348 349
            Td = self.hum[1]  # dew point temperature (K)
            Td = np.where(Td < 200, np.copy(Td)+CtoK, np.copy(Td))
            T = np.where(self.T < 200, np.copy(self.T)+CtoK, np.copy(self.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)))
            self.rh = 100*esd/es
350 351 352 353
        elif self.hum[0] == "q":
            es = 611.21*np.exp(17.502*((self.T-CtoK)/(self.T-32.19)))
            e = self.qair*self.P/(0.378*self.qair+0.622)
            self.rh = 100*e/es
sbiri's avatar
sbiri committed
354

355
    def _flag(self, out=0):
356 357 358 359
        miss = np.copy(self.msk)  # define missing input points
        if self.cskin == 1:
            miss = np.where(np.isnan(self.msk+self.P+self.Rs+self.Rl),
                            np.nan, 1)
sbiri's avatar
sbiri committed
360
        else:
361 362 363 364
            miss = np.where(np.isnan(self.msk+self.P), np.nan, 1)
        flag = set_flag(miss, self.rh, self.u10n, self.q10n, self.t10n,
                        self.Rb, self.hin, self.monob, self.itera, out=out)

sbiri's avatar
sbiri committed
365 366
        self.flag = flag

367
    def get_output(self, out_var=None, out=0):
sbiri's avatar
sbiri committed
368 369

        assert out in [0, 1], "out must be either 0 or 1"
sbiri's avatar
sbiri committed
370

371
        self._get_humidity()  # Get the Relative humidity
sbiri's avatar
sbiri committed
372 373
        self._flag(out=out)  # Get flags

374 375 376 377 378
        self.GustFact = apply_GF(self.gust, self.spd, self.wind, "TSF")
        self.tau = self.rho*np.power(self.usr, 2)/self.GustFact[0]
        self.sensible = self.rho*self.cp*(self.usr/self.GustFact[1])*self.tsr
        self.latent = self.rho*self.lv*(self.usr/self.GustFact[2])*self.qsr

379
        self.GustFact = apply_GF(self.gust, self.spd, self.wind, "u")
sbiri's avatar
sbiri committed
380 381 382 383 384 385 386 387 388 389 390 391 392
        if self.gust[0] in [3, 4]:
            self.u10n = self.wind-self.usr/kappa*(
                np.log(self.h_in[0]/self.ref10)-self.psim)
            self.uref = self.wind-self.usr/kappa*(
                np.log(self.h_in[0]/self.h_out[0])-self.psim +
                psim_calc(self.h_out[0]/self.monob, self.meth))
        else:
            self.u10n = self.spd-self.usr/kappa/self.GustFact*(
                np.log(self.h_in[0]/self.ref10)-self.psim) 
            self.uref = self.spd-self.usr/kappa/self.GustFact * \
                (np.log(self.h_in[0]/self.h_out[0])-self.psim +
                 psim_calc(self.h_out[0]/self.monob, self.meth))

393
        self.usrGF = self.usr/self.GustFact
sbiri's avatar
sbiri committed
394
        # include lapse rate adjustment as theta is well-mixed
sbiri's avatar
sbiri committed
395 396
        self.tref = self.theta-self.tlapse*self.h_out[1]-self.tsr/kappa * \
            (np.log(self.h_in[1]/self.h_out[1])-self.psit +
397
             psit_calc(self.h_out[1]/self.monob, self.meth))
sbiri's avatar
sbiri committed
398 399
        self.qref = self.qair-self.qsr/kappa * \
            (np.log(self.h_in[2]/self.h_out[2])-self.psiq +
400
             psit_calc(self.h_out[2]/self.monob, self.meth))
401 402 403
        self.psim_ref = psim_calc(self.h_out[0]/self.monob, self.meth)
        self.psit_ref = psit_calc(self.h_out[1]/self.monob, self.meth)
        self.psiq_ref = psit_calc(self.h_out[2]/self.monob, self.meth)
sbiri's avatar
sbiri committed
404

sbiri's avatar
sbiri committed
405 406 407
        if self.wl == 0:
            self.dtwl = np.zeros(self.T.shape)*self.msk
            # reset to zero if not used
sbiri's avatar
sbiri committed
408 409

        # Do not calculate lhf if a measure of humidity is not input
410
        # This gets filled into a pd dataframe and so no need to specify
sbiri's avatar
sbiri committed
411 412 413 414
        # dimension of array
        if self.hum[0] == 'no':
            self.latent, self.qsr, self.q10n = np.empty(3)
            self.qref, self.qair, self.rh = np.empty(3)
sbiri's avatar
sbiri committed
415 416 417 418 419

        # Set the final wind speed values
        # this seems to be gust (was wind_speed)
        self.ug = np.sqrt(np.power(self.wind, 2)-np.power(self.spd, 2))

sbiri's avatar
sbiri committed
420 421
        # Get class specific flags (will only work if self.u_hi and self.u_lo
        # have been set in the class)
sbiri's avatar
sbiri committed
422 423 424 425 426 427
        try:
            self._class_flag()
        except AttributeError:
            pass

        # Combine all output variables into a pandas array
428 429
        res_vars = get_outvars(out_var, self.cskin, self.gust)

sbiri's avatar
sbiri committed
430 431

        res = np.zeros((len(res_vars), len(self.spd)))
sbiri's avatar
sbiri committed
432 433
        for i, value in enumerate(res_vars):
            res[i][:] = getattr(self, value)
sbiri's avatar
sbiri committed
434

sbiri's avatar
sbiri committed
435
        if out == 0:
sbiri's avatar
sbiri committed
436 437
            res[:, self.ind] = np.nan
            # set missing values where data have non acceptable values
sbiri's avatar
sbiri committed
438
            if self.hum[0] != 'no':
439
                res = np.asarray([
440
                    np.where(self.q10n < 0, np.nan, res[i][:])
441
                    for i in range(len(res_vars))])
442 443
                # len(res_vars)-1 instead of len(res_vars) in order to keep
                # itera= -1 for no convergence
444
            res = np.asarray([
445
                np.where(self.u10n < 0, np.nan, res[i][:])
446 447 448 449 450
                for i in range(len(res_vars))])
            res = np.asarray([
                np.where(((self.t10n < 173) | (self.t10n > 373)), np.nan,
                          res[i][:])
                for i in range(len(res_vars))])
sbiri's avatar
sbiri committed
451
        else:
sbiri's avatar
sbiri committed
452 453 454 455
            warnings.warn("Warning: the output will contain values for points"
                          " that have not converged and negative values "
                          "(if any) for u10n/q10n")

456 457
        resAll = pd.DataFrame(data=res.T, index=range(self.nlen),
                              columns=res_vars)
458 459
        if "itera" in res_vars:
            resAll["itera"] = self.itera  # restore itera
sbiri's avatar
sbiri committed
460 461 462 463 464 465 466 467

        resAll["flag"] = self.flag

        return resAll

    def add_variables(self, spd, T, SST, lat=None, hum=None, P=None, L=None):

        # Add the mandatory variables
sbiri's avatar
sbiri committed
468 469 470 471
        assert type(spd) == type(T) == type(
            SST) == np.ndarray, "input type of spd, T and SST should be"
        " numpy.ndarray"
        self.L = "tsrv" if L is None else L
sbiri's avatar
sbiri committed
472 473 474 475
        self.arr_shp = spd.shape
        self.nlen = len(spd)
        self.spd = spd
        self.T = np.where(T < 200, np.copy(T)+CtoK, np.copy(T))
sbiri's avatar
sbiri committed
476
        self.hum = ['no', np.full(SST.shape, 80)] if hum is None else hum
sbiri's avatar
sbiri committed
477
        self.SST = np.where(SST < 200, np.copy(SST)+CtoK, np.copy(SST))
sbiri's avatar
sbiri committed
478
        self.lat = np.full(self.arr_shp, 45) if lat is None else lat
sbiri's avatar
sbiri committed
479
        self.grav = gc(self.lat)
480
        self.P = np.full(self.nlen, 1013) if P is None else P
sbiri's avatar
sbiri committed
481 482

        # mask to preserve missing values when initialising variables
sbiri's avatar
sbiri committed
483
        self.msk = np.empty(SST.shape)
sbiri's avatar
sbiri committed
484 485 486
        self.msk = np.where(np.isnan(spd+T+SST), np.nan, 1)
        self.Rb = np.empty(SST.shape)*self.msk

sbiri's avatar
sbiri committed
487 488
    def add_gust(self, gust=None):
        if np.all(gust is None):
sbiri's avatar
sbiri committed
489 490 491
            try:
                gust = self.default_gust
            except AttributeError:
492 493
                gust = [0, 0, 0, 0]  # gustiness OFF
                # gust = [1, 1.2, 800]
sbiri's avatar
sbiri committed
494
        elif ((np.size(gust) < 3) and (gust == 0)):
495
            gust = [0, 0, 0, 0]
sbiri's avatar
sbiri committed
496

497
        assert np.size(gust) == 4, "gust input must be a 4x1 array"
498
        assert gust[0] in range(6), "gust at position 0 must be 0 to 5"
sbiri's avatar
sbiri committed
499 500 501
        self.gust = gust

    def _class_flag(self):
sbiri's avatar
sbiri committed
502 503 504 505 506 507 508 509 510 511 512
        """A flag specific to this class - only used for certain classes where
         u_lo and u_hi are defined"""
        self.flag = np.where(((self.u10n < self.u_lo[0]) |
                              (self.u10n > self.u_hi[0])) &
                             (self.flag == "n"), "o",
                             np.where(((self.u10n < self.u_lo[1]) |
                                       (self.u10n > self.u_hi[1])) &
                                      ((self.flag != "n") & (np.char.find(
                                          self.flag.astype(str), 'u') == -1) &
                                       (np.char.find(
                                           self.flag.astype(str), 'q') == -1)),
sbiri's avatar
sbiri committed
513 514 515 516 517
                                      self.flag+[","]+["o"], self.flag))

    def __init__(self):
        self.meth = "S88"

518

sbiri's avatar
sbiri committed
519 520 521 522
class S80(S88):

    def __init__(self):
        self.meth = "S80"
sbiri's avatar
sbiri committed
523 524
        self.u_lo = [6, 6]
        self.u_hi = [22, 22]
sbiri's avatar
sbiri committed
525

sbiri's avatar
sbiri committed
526

sbiri's avatar
sbiri committed
527 528 529 530
class YT96(S88):

    def __init__(self):
        self.meth = "YT96"
531
        # no limits to u range as we use eq. 21 for cdn
sbiri's avatar
sbiri committed
532 533
        # self.u_lo = [0, 3]
        # self.u_hi = [26, 26]
sbiri's avatar
sbiri committed
534

535

sbiri's avatar
sbiri committed
536
class LP82(S88):
sbiri's avatar
sbiri committed
537

sbiri's avatar
sbiri committed
538 539
    def __init__(self):
        self.meth = "LP82"
sbiri's avatar
sbiri committed
540 541
        self.u_lo = [3, 3]
        self.u_hi = [25, 25]
sbiri's avatar
sbiri committed
542 543 544 545 546 547 548 549


class NCAR(S88):

    def _minimum_params(self):
        self.cd = np.maximum(np.copy(self.cd), 1e-4)
        self.ct = np.maximum(np.copy(self.ct), 1e-4)
        self.cq = np.maximum(np.copy(self.cq), 1e-4)
sbiri's avatar
sbiri committed
550
        # self.zo = np.minimum(np.copy(self.zo), 0.0025)
sbiri's avatar
sbiri committed
551

sbiri's avatar
sbiri committed
552 553
    def __init__(self):
        self.meth = "NCAR"
sbiri's avatar
sbiri committed
554 555
        self.u_lo = [0.5, 0.5]
        self.u_hi = [999, 999]
sbiri's avatar
sbiri committed
556

557

sbiri's avatar
sbiri committed
558
class UA(S88):
sbiri's avatar
sbiri committed
559

sbiri's avatar
sbiri committed
560 561
    def __init__(self):
        self.meth = "UA"
sbiri's avatar
sbiri committed
562
        self.default_gust = [1, 1.2, 600, 0.01]
sbiri's avatar
sbiri committed
563 564
        self.u_lo = [-999, -999]
        self.u_hi = [18, 18]
sbiri's avatar
sbiri committed
565 566 567


class C30(S88):
sbiri's avatar
sbiri committed
568 569
    # def set_coolskin_warmlayer(self, wl=0, cskin=1, skin="C35", Rl=None, Rs=None):
    #     self._fix_coolskin_warmlayer(wl, cskin, skin, Rl, Rs)
sbiri's avatar
sbiri committed
570 571 572

    def __init__(self):
        self.meth = "C30"
sbiri's avatar
sbiri committed
573
        self.default_gust = [1, 1.2, 600, 0.01]
sbiri's avatar
sbiri committed
574
        self.skin = "C35"
sbiri's avatar
sbiri committed
575

576

sbiri's avatar
sbiri committed
577 578 579
class C35(C30):
    def __init__(self):
        self.meth = "C35"
sbiri's avatar
sbiri committed
580
        self.default_gust = [1, 1.2, 600, 0.01]
sbiri's avatar
sbiri committed
581
        self.skin = "C35"
sbiri's avatar
sbiri committed
582

583

sbiri's avatar
sbiri committed
584
class ecmwf(C30):
sbiri's avatar
sbiri committed
585 586 587
    # def set_coolskin_warmlayer(self, wl=0, cskin=1, skin="ecmwf", Rl=None,
    #                            Rs=None):
    #     self._fix_coolskin_warmlayer(wl, cskin, skin, Rl, Rs)
588 589 590 591 592
    # def _minimum_params(self):
    #     self.cdn = np.maximum(np.copy(self.cdn), 0.1e-3)
    #     self.ctn = np.maximum(np.copy(self.ctn), 0.1e-3)
    #     self.cqn = np.maximum(np.copy(self.cqn), 0.1e-3)
    #     self.wind = np.maximum(np.copy(self.wind), 0.2)
sbiri's avatar
sbiri committed
593 594 595

    def __init__(self):
        self.meth = "ecmwf"
sbiri's avatar
sbiri committed
596
        self.default_gust = [1, 1.2, 600, 0.01]
sbiri's avatar
sbiri committed
597
        self.skin = "ecmwf"
sbiri's avatar
sbiri committed
598

599

sbiri's avatar
sbiri committed
600
class Beljaars(C30):
sbiri's avatar
sbiri committed
601 602 603
    # def set_coolskin_warmlayer(self, wl=0, cskin=1, skin="Beljaars", Rl=None,
    #                            Rs=None):
    #     self._fix_coolskin_warmlayer(wl, cskin, skin, Rl, Rs)
sbiri's avatar
sbiri committed
604

sbiri's avatar
sbiri committed
605 606
    def __init__(self):
        self.meth = "Beljaars"
sbiri's avatar
sbiri committed
607 608 609
        self.default_gust = [1, 1.2, 600, 0.01]
        self.skin = "ecmwf"
        # self.skin = "Beljaars"
sbiri's avatar
sbiri committed
610 611


612 613 614 615
def AirSeaFluxCode(spd, T, SST, meth, lat=None, hum=None, P=None, hin=18,
                   hout=10, Rl=None, Rs=None, cskin=0, skin=None, wl=0,
                   gust=None, qmeth="Buck2", tol=None, maxiter=10, out=0,
                   out_var=None, L=None):
616
    """
sbiri's avatar
sbiri committed
617 618 619
    Calculate turbulent surface fluxes using different parameterizations.

    Calculate height adjusted values for spd, T, q
sbiri's avatar
sbiri committed
620 621 622 623 624 625 626 627 628 629

    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)
630 631 632
        meth : str
            "S80", "S88", "LP82", "YT96", "UA", "NCAR", "C30", "C35",
            "ecmwf", "Beljaars"
sbiri's avatar
sbiri committed
633
        lat : float
634 635 636
            latitude (deg), default 45deg
        hum : float
            humidity input switch 2x1 [x, values] default is relative humidity
637 638 639
            x='rh' : relative humidity in %
            x='q' : specific humidity (g/kg)
            x='Td' : dew point temperature (K)
sbiri's avatar
sbiri committed
640
        P : float
641
            air pressure (hPa), default 1013hPa
sbiri's avatar
sbiri committed
642
        hin : float
643
            sensor heights in m (array 3x1 or 3xn), default 18m
sbiri's avatar
sbiri committed
644 645 646 647 648 649
        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
650
        cskin : int
651
            0 switch cool skin adjustment off, else 1
652
            default is 0
653 654 655 656
        skin : str
            cool skin method option "C35", "ecmwf" or "Beljaars"
        wl : int
            warm layer correction default is 0, to switch on set to 1
657
        gust : int
658 659
            4x1 [x, beta, zi, ustb] x=0 gustiness is OFF, x=1-5 gustiness is ON
            and use gustiness factor: 1. Fairall et al. 2003, 2. GF is removed
sbiri's avatar
sbiri committed
660 661
            from TSFs u10n, uref, 3. GF=1, 4. following Zeng et al. 1998 or
            Brodeau et al. 2006, 5. following C35 matlab code;
662 663 664 665
            beta gustiness parameter, default is 1.2,
            zi PBL height (m) default is 600,
            min is the value for gust speed in stable conditions,
            default is 0.01ms^{-1}
666 667
        qmeth : str
            is the saturation evaporation method to use amongst
668 669
            "HylandWexler","Hardy","Preining","Wexler","GoffGratch","WMO",
            "MagnusTetens","Buck","Buck2","WMO2018","Sonntag","Bolton",
670 671 672 673 674 675 676
            "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
677 678
                    adjustment lim1-6
           default is tol=['all', 0.01, 0.01, 1e-05, 1e-3, 0.1, 0.1]
sbiri's avatar
sbiri committed
679 680
        maxiter : int
            number of iterations (default = 10)
681
        out : int
682
            set 0 to set points that have not converged, negative values of
683
                  u10n, q10n or T10n out of limits to missing (default)
684
            set 1 to keep points
685 686 687 688
        out_var : str
           optional. user can define pandas array of variables to be output.
           the default full pandas array is :
               out_var = ("tau", "sensible", "latent", "monob", "cd", "cd10n",
689 690 691 692 693 694 695
                           "ct", "ct10n", "cq", "cq10n", "tsrv", "tsr", "qsr",
                           "usr", "psim", "psit", "psiq", "psim_ref", "psit_ref",
                           "psiq_ref", "u10n", "t10n", "q10n", "zo", "zot", "zoq",
                           "uref", "tref", "qref", "dter", "dqer", "dtwl", "tkt",
                           "qair", "qsea", "Rl", "Rs", "Rnl", "ug", "usrGF",
                           "GustFact", "Rb", "rh", "rho", "cp", "lv", "theta",
                           "itera")
696 697 698
            the "limited" pandas array is:
                out_var = ("tau", "sensible", "latent", "uref", "tref", "qref")
            the user can define a custom pandas array of variables to  output.
699
        L : str
sbiri's avatar
sbiri committed
700
           Monin-Obukhov length definition options
701 702 703
           "tsrv"  : default
           "Rb" : following ecmwf (IFS Documentation cy46r1)

sbiri's avatar
sbiri committed
704 705 706
    Returns
    -------
        res : array that contains
sbiri's avatar
sbiri committed
707 708 709 710
                       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
711
                       5. drag coefficient (cd)
sbiri's avatar
sbiri committed
712
                       6. neutral drag coefficient (cd10n)
sbiri's avatar
sbiri committed
713
                       7. heat exchange coefficient (ct)
sbiri's avatar
sbiri committed
714
                       8. neutral heat exchange coefficient (ct10n)
sbiri's avatar
sbiri committed
715
                       9. moisture exhange coefficient (cq)
sbiri's avatar
sbiri committed
716
                       10. neutral moisture exchange coefficient (cq10n)
sbiri's avatar
sbiri committed
717
                       11. star virtual temperatcure (tsrv)
sbiri's avatar
sbiri committed
718
                       12. star temperature (tsr)
719 720
                       13. star specific humidity (qsr)
                       14. star wind speed (usr)
sbiri's avatar
sbiri committed
721
                       15. momentum stability function (psim)
722 723
                       16. heat stability function (psit)
                       17. moisture stability function (psiq)
724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755
                       18. momentum stability function at hout (psim_ref)
                       19. heat stability function at hout (psit_ref)
                       20. moisture stability function at hout (psiq_ref)
                       21. 10m neutral wind speed (u10n)
                       22. 10m neutral temperature (t10n)
                       23. 10m neutral specific humidity (q10n)
                       24. surface roughness length (zo)
                       25. heat roughness length (zot)
                       26. moisture roughness length (zoq)
                       27. wind speed at reference height (uref)
                       28. temperature at reference height (tref)
                       29. specific humidity at reference height (qref)
                       30. cool-skin temperature depression (dter)
                       31. cool-skin humidity depression (dqer)
                       32. warm layer correction (dtwl)
                       33. thickness of the viscous layer (delta)
                       34. specific humidity of air (qair)
                       35. specific humidity at sea surface (qsea)
                       36. downward longwave radiation (Rl)
                       37. downward shortwave radiation (Rs)
                       38. downward net longwave radiation (Rnl)
                       39. gust wind speed (ug)
                       40. star wind speed/GustFact (usrGF)
                       41. Gustiness Factor (GustFact)
                       42. Bulk Richardson number (Rb)
                       43. relative humidity (rh)
                       44. air density (rho)
                       45. specific heat of moist air (cp)
                       46. lv latent heat of vaporization (Jkg−1)
                       47. potential temperature (theta)
                       48. number of iterations until convergence
                       49. flag ("n": normal, "o": out of nominal range,
sbiri's avatar
sbiri committed
756
                                 "u": u10n<0, "q":q10n<0 or q>0.04
757
                                 "m": missing,
758
                                 "l": Rib<-0.5 or Rib>0.2 or z/L>1000,
sbiri's avatar
sbiri committed
759
                                 "r" : rh>100%,
760
                                 "t" : t10n<173K or t10n>373K
761
                                 "i": convergence fail at n)
762

763
    2021 / Author S. Biri
764 765
    2021 / Restructured by R. Cornes
    2021 / Simplified by E. Kent
sbiri's avatar
sbiri committed
766
    """
sbiri's avatar
sbiri committed
767
    logging.basicConfig(filename='flux_calc.log', filemode="w",
768
                        format='%(asctime)s %(message)s', level=logging.INFO)
769
    logging.captureWarnings(True)
sbiri's avatar
sbiri committed
770 771 772 773 774 775

    iclass = globals()[meth]()
    iclass.add_gust(gust=gust)
    iclass.add_variables(spd, T, SST, lat=lat, hum=hum, P=P, L=L)
    iclass.get_heights(hin, hout)
    iclass.get_specHumidity(qmeth=qmeth)
sbiri's avatar
sbiri committed
776 777
    iclass.set_coolskin_warmlayer(wl=wl, cskin=cskin, skin=skin, Rl=Rl, Rs=Rs)
    iclass.iterate(tol=tol, maxiter=maxiter)
778
    resAll = iclass.get_output(out_var=out_var, out=out)
sbiri's avatar
sbiri committed
779

780
    return resAll