AirSeaFluxCode.py 32.8 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):
13
        if self.gust[0] in range(1, 5):
sbiri's avatar
sbiri committed
14
            self.wind[ind] = np.sqrt(np.power(np.copy(self.spd[ind]), 2) +
sbiri's avatar
sbiri committed
15 16 17
                                     np.power(get_gust(self.gust[1],
                                                       self.theta[ind],
                                                       self.usr[ind],
18 19
                                                       self.tsrv[ind],
                                                       self.gust[2],
20
                                                       self.grav[ind]), 2))
21
            # ratio of gusty to horizontal wind
22 23 24 25 26
            self.GustFact = apply_GF(self.gust, self.spd, self.wind, "u")
            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])
            if self.gust[0] == 3:
27
                # option to not remove GustFact
28
                self.u10n[ind] = self.wind[ind]-self.usr[ind]/kappa*(
29
                    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 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101
    def _update_coolskin_warmlayer(self, ind):
        if self.cskin == 1:
            if self.skin == "C35":
                self.dter[ind], self.tkt[ind] = cs_C35(np.copy(
                    self.SST[ind]), self.rho[ind], self.Rs[ind], self.Rnl[ind],
                    self.cp[ind], self.lv[ind], np.copy(self.tkt[ind]),
                    self.usr[ind], self.tsr[ind], self.qsr[ind], self.grav[ind])
            elif self.skin == "ecmwf":
                self.dter[ind] = cs_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]), self.grav[ind])
            elif self.skin == "Beljaars":
                self.Qs[ind], self.dter[ind] = cs_Beljaars(
                    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], np.copy(self.Qs[ind]))
            self.dqer[ind] = get_dqer(self.dter[ind], self.SST[ind],
                                      self.qsea[ind], self.lv[ind])
sbiri's avatar
sbiri committed
102 103
            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
104 105 106 107 108 109 110 111 112 113
            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
114 115 116 117 118 119 120 121
                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):
122
        # reference height
sbiri's avatar
sbiri committed
123 124 125 126 127 128 129 130 131 132 133 134
        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
135
        self.wind = np.sqrt(np.power(np.copy(self.spd), 2)+0.25)
sbiri's avatar
sbiri committed
136 137 138 139
        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
140
        Rb = self.grav*self.h_in[1]*self.dtv/(self.T*np.power(self.wind, 2))
141 142
        # eq. 12 Grachev & Fairall 1997   # DO.THIS
        self.monob = self.h_in[1]/12.0/Rb
sbiri's avatar
sbiri committed
143
        # where does 12.0 come from??
sbiri's avatar
sbiri committed
144 145 146

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

147 148
        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
149
        if self.cskin + self.wl > 0:
sbiri's avatar
sbiri committed
150 151
            self.dter, self.tkt, self.dtwl = [
                dummy_array(x) for x in (-0.3, 0.001, 0.3)]
152 153
            self.dqer = get_dqer(self.dter, self.SST, self.qsea,
                                 self.lv)
sbiri's avatar
sbiri committed
154 155
            self.Rnl = 0.97*(self.Rl-5.67e-8*np.power(
                self.SST-0.3*self.cskin, 4))
sbiri's avatar
sbiri committed
156 157
            self.Qs = 0.945*self.Rs
        else:
sbiri's avatar
sbiri committed
158 159 160 161
            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
162 163 164 165 166
        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
167 168
        self.cd10n, self.zo = cdn_calc(
            self.u10n, self.usr, self.theta, self.grav, self.meth)
sbiri's avatar
sbiri committed
169 170
        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)
171
        self.usr =np.sqrt(self.cd*np.power(self.wind, 2))
sbiri's avatar
sbiri committed
172 173 174 175
        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)]
176
        self.tv10n = self.zot
sbiri's avatar
sbiri committed
177 178 179 180 181 182

    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
183 184 185 186
        # 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
187 188
        assert tol[0] in ['flux', 'ref', 'all'], "unknown tolerance input"

sbiri's avatar
sbiri committed
189 190
        old_vars = {"flux": ["tau", "sensible", "latent"],
                    "ref": ["u10n", "t10n", "q10n"]}
sbiri's avatar
sbiri committed
191 192 193
        old_vars["all"] = old_vars["ref"] + old_vars["flux"]
        old_vars = old_vars[tol[0]]

sbiri's avatar
sbiri committed
194 195
        new_vars = {"flux": ["tau", "sensible", "latent"],
                    "ref": ["u10n", "t10n", "q10n"]}
sbiri's avatar
sbiri committed
196 197 198 199 200 201 202 203 204 205 206
        new_vars["all"] = new_vars["ref"] + new_vars["flux"]
        new_vars = new_vars[tol[0]]
        # I'm sure there are better ways of doing this
        # 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
207 208
        self.tsrv, self.psim, self.psit, self.psiq = [
            np.zeros(self.arr_shp)*self.msk for _ in range(4)]
sbiri's avatar
sbiri committed
209 210

        # extreme values for first comparison
211
        dummy_array = lambda val : np.full(self.T.shape, val)*self.msk
212
        # you can use def instead of lambda
sbiri's avatar
sbiri committed
213 214 215
        # 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
216 217 218 219 220 221 222 223 224 225

        # 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
226
            old = np.array([np.copy(getattr(self, i)) for i in old_vars])
sbiri's avatar
sbiri committed
227 228

            # Calculate cdn
sbiri's avatar
sbiri committed
229 230 231 232 233
            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
234
                logging.info('break %s at iteration %s cd10n<0', meth, it)
sbiri's avatar
sbiri committed
235 236 237 238 239 240
                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])
241 242
            # Update the wind values
            self._wind_iterate(ind)
sbiri's avatar
sbiri committed
243 244

            # temperature
sbiri's avatar
sbiri committed
245 246 247 248 249 250 251 252
            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])
253
            # humidity
sbiri's avatar
sbiri committed
254 255 256 257 258 259 260 261
            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])
262

sbiri's avatar
sbiri committed
263 264 265 266 267
            # Some parameterizations set a minimum on parameters
            try:
                self._minimum_params()
            except AttributeError:
                pass
sbiri's avatar
sbiri committed
268 269 270

            self.dt_full[ind] = self.dt_in[ind] - \
                self.dter[ind]*self.cskin - self.dtwl[ind]*self.wl
sbiri's avatar
sbiri committed
271
            self.dq_full[ind] = self.dq_in[ind] - self.dqer[ind]*self.cskin
sbiri's avatar
sbiri committed
272 273 274 275 276
            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
277 278 279

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

sbiri's avatar
sbiri committed
281
            # Logging output
sbiri's avatar
sbiri committed
282 283 284 285 286 287 288 289
            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
290 291

            if self.cskin + self.wl > 0:
sbiri's avatar
sbiri committed
292 293 294
                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
295 296
            # not sure how to handle lapse/potemp
            # well-mixed in potential temperature ...
297
            self.t10n[ind] = self.theta[ind]-self.tlapse[ind]*self.ref10 - \
sbiri's avatar
sbiri committed
298
                self.tsr[ind]/kappa * \
299
                (np.log(self.h_in[1, ind]/self.ref10)-self.psit[ind])
300
            self.q10n[ind] = self.qair[ind]-self.qsr[ind]/kappa * \
sbiri's avatar
sbiri committed
301
                (np.log(self.h_in[2, ind]/self.ref10)-self.psiq[ind])
sbiri's avatar
sbiri committed
302 303

            # update stability info
sbiri's avatar
sbiri committed
304 305 306 307 308 309
            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
310
            if self.L == "tsrv":
sbiri's avatar
sbiri committed
311 312 313
                self.monob[ind] = get_Ltsrv(
                    self.tsrv[ind], self.grav[ind], self.tv[ind],
                    self.usr[ind])
sbiri's avatar
sbiri committed
314
            else:
sbiri's avatar
sbiri committed
315
                self.monob[ind] = get_LRb(
316
                    self.Rb[ind], self.h_in[1, ind], self.monob[ind],
sbiri's avatar
sbiri committed
317
                    self.zo[ind], self.zot[ind], self.meth)
sbiri's avatar
sbiri committed
318 319 320

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

sbiri's avatar
sbiri committed
322
            # make sure you allow small negative values convergence
sbiri's avatar
sbiri committed
323 324
            if it == 1:
                self.u10n = np.where(self.u10n < 0, 0.5, self.u10n)
sbiri's avatar
sbiri committed
325

sbiri's avatar
sbiri committed
326
            self.itera[ind] = np.full(1, it)
sbiri's avatar
sbiri committed
327
            # remove effect of gustiness following Fairall et al. (2003)
sbiri's avatar
sbiri committed
328 329
            # usr is divided by (GustFact)^0.5 (here applied to sensible and
            # latent as well as tau)
330
            # GustFact should be 1 if gust is OFF or gust[0]=2
331 332 333 334
            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
sbiri's avatar
sbiri committed
335 336

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

sbiri's avatar
sbiri committed
339
            if it > 2:  # force at least three iterations
340
                d = np.abs(new-old)  # change over this iteration
sbiri's avatar
sbiri committed
341 342 343 344
                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
345 346 347 348 349 350 351 352

            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
353 354
        logging.info('method %s | # of points that did not converge :%s \n',
                     self.meth, self.ind[0].size)
sbiri's avatar
sbiri committed
355 356

    def _get_humidity(self):
357
        """Calculate RH used for flagging purposes & output."""
sbiri's avatar
sbiri committed
358
        if self.hum[0] in ('rh', 'no'):
sbiri's avatar
sbiri committed
359
            self.rh = self.hum[1]
sbiri's avatar
sbiri committed
360
        elif self.hum[0] == 'Td':
sbiri's avatar
sbiri committed
361 362 363 364 365 366
            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
367 368 369 370
        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
371

372
    def _flag(self, out=0):
373 374 375 376
        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
377
        else:
378 379 380 381
            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
382 383
        self.flag = flag

384
    def get_output(self, out_var=None, out=0):
sbiri's avatar
sbiri committed
385 386

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

388
        self._get_humidity()  # Get the Relative humidity
sbiri's avatar
sbiri committed
389 390
        self._flag(out=out)  # Get flags

391
        self.GustFact = apply_GF(self.gust, self.spd, self.wind, "u")
sbiri's avatar
sbiri committed
392 393
        # remove effect of gustiness following Fairall et al. (2003)
        # usr is divided by (GustFact)^0.5
394
        self.uref = self.spd-self.usr/kappa/self.GustFact * \
sbiri's avatar
sbiri committed
395
            (np.log(self.h_in[0]/self.h_out[0])-self.psim +
396
             psim_calc(self.h_out[0]/self.monob, self.meth))
sbiri's avatar
sbiri committed
397
        # include lapse rate adjustment as theta is well-mixed
sbiri's avatar
sbiri committed
398 399
        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 +
400
             psit_calc(self.h_out[1]/self.monob, self.meth))
sbiri's avatar
sbiri committed
401 402
        self.qref = self.qair-self.qsr/kappa * \
            (np.log(self.h_in[2]/self.h_out[2])-self.psiq +
403
             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
sbiri's avatar
sbiri committed
410 411 412 413 414
        # This gets filled into a pd dataframe and so no need to specify y
        # 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
        if out_var == "limited":
            out_var = ("tau", "sensible", "latent", "uref", "tref", "qref")
sbiri's avatar
sbiri committed
430 431 432
        res_vars = ("tau", "sensible", "latent", "monob", "cd", "cd10n", "ct",
                    "ct10n", "cq", "cq10n", "tsrv", "tsr", "qsr", "usr",
                    "psim", "psit", "psiq", "u10n", "t10n", "q10n", "zo",
433 434 435 436
                    "zot", "zoq", "uref", "tref", "qref", "dter", "dqer",
                    "dtwl", "qair", "qsea", "Rl", "Rs", "Rnl", "ug", "Rb",
                    "rh", "rho", "cp", "tkt", "lv",
                    "itera") if out_var is None else out_var
sbiri's avatar
sbiri committed
437 438

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

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

459 460
        resAll = pd.DataFrame(data=res.T, index=range(self.nlen),
                              columns=res_vars)
sbiri's avatar
sbiri committed
461 462 463 464 465 466 467 468

        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
469 470 471 472
        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
473 474 475 476
        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
477
        self.hum = ['no', np.full(SST.shape, 80)] if hum is None else hum
sbiri's avatar
sbiri committed
478
        self.SST = np.where(SST < 200, np.copy(SST)+CtoK, np.copy(SST))
sbiri's avatar
sbiri committed
479
        self.lat = np.full(self.arr_shp, 45) if lat is None else lat
sbiri's avatar
sbiri committed
480
        self.grav = gc(self.lat)
481
        self.P = np.full(self.nlen, 1013) if P is None else P
sbiri's avatar
sbiri committed
482 483

        # mask to preserve missing values when initialising variables
sbiri's avatar
sbiri committed
484
        self.msk = np.empty(SST.shape)
sbiri's avatar
sbiri committed
485 486 487
        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
488 489
    def add_gust(self, gust=None):
        if np.all(gust is None):
sbiri's avatar
sbiri committed
490 491 492
            try:
                gust = self.default_gust
            except AttributeError:
sbiri's avatar
sbiri committed
493
                gust = [0, 0, 0]  # gustiness OFF
sbiri's avatar
sbiri committed
494 495
        elif ((np.size(gust) < 3) and (gust == 0)):
            gust = [0, 0, 0]
sbiri's avatar
sbiri committed
496

sbiri's avatar
sbiri committed
497
        assert np.size(gust) == 3, "gust input must be a 3x1 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 = [3, 1, 1000]
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 = [3, 1.2, 600]
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 = [3, 1.2, 600]
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 = [3, 1, 1000]
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
        self.default_gust = [3, 1, 1000]
sbiri's avatar
sbiri committed
608
        self.skin = "Beljaars"
sbiri's avatar
sbiri committed
609 610


611 612 613 614
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):
615
    """
sbiri's avatar
sbiri committed
616 617 618
    Calculate turbulent surface fluxes using different parameterizations.

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

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

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

754
    2021 / Author S. Biri
755 756
    2021 / Restructured by R. Cornes
    2021 / Simplified by E. Kent
sbiri's avatar
sbiri committed
757
    """
sbiri's avatar
sbiri committed
758
    logging.basicConfig(filename='flux_calc.log', filemode="w",
759
                        format='%(asctime)s %(message)s', level=logging.INFO)
760
    logging.captureWarnings(True)
sbiri's avatar
sbiri committed
761 762 763 764 765 766

    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
767 768
    iclass.set_coolskin_warmlayer(wl=wl, cskin=cskin, skin=skin, Rl=Rl, Rs=Rs)
    iclass.iterate(tol=tol, maxiter=maxiter)
769
    resAll = iclass.get_output(out_var=out_var, out=out)
sbiri's avatar
sbiri committed
770

771
    return resAll