AirSeaFluxCode.py 36.1 KB
Newer Older
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)
Richard Cornes's avatar
Richard Cornes committed
6 7
from util_subs import *
from flux_subs import *
8
from cs_wl_subs import *
sbiri's avatar
sbiri committed
9

10
class S88:
Richard Cornes's avatar
Richard Cornes committed
11 12 13 14 15 16 17 18 19 20 21 22

    def _wind_firstguess(self):
        if (self.gust[0] == 1):
            self.wind = np.sqrt(np.power(np.copy(self.spd), 2)+np.power(0.5, 2))
        else:
            self.wind = np.copy(self.spd)

    def _wind_iterate(self,ind):
        if(self.gust[0] == 1):
            self.wind[ind] = np.sqrt(np.power(np.copy(self.spd[ind]), 2) +
                                np.power(get_gust(self.gust[1], self.Ta[ind], self.usr[ind],
                                                  self.tsrv[ind], self.gust[2],
23
                                                  self.g[ind]), 2))
Richard Cornes's avatar
Richard Cornes committed
24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60
    def get_heights(self, hin, hout=10):
        self.hout = hout
        self.hin = hin
        self.h_in = get_heights(hin, len(self.spd))
        self.h_out = get_heights(self.hout,1)

    def get_specHumidity(self,qmeth="Buck2"):
        self.qair, self.qsea = get_hum(self.hum, self.T, self.SST, self.P, qmeth)
        if (np.all(np.isnan(self.qsea)) or np.all(np.isnan(self.qair))):
            raise ValueError("qsea and qair cannot be nan")
        self.dq = self.qair - self.qsea
        
        # Set lapse rate and Potential Temperature (now we have humdity)
        self._get_potentialT()
        self._get_lapse()

    def _get_potentialT(self):
        self.cp = 1004.67*(1+0.00084*self.qsea)
        self.th = np.where(self.T < 200, (np.copy(self.T)+CtoK) *
                  np.power(1000/self.P, 287.1/self.cp),
                  np.copy(self.T)*np.power(1000/self.P, 287.1/self.cp))  # potential T

    def _get_lapse(self):
        self.tlapse = gamma("dry", self.SST, self.T, self.qair/1000, self.cp)
        self.Ta = np.where(self.T < 200, np.copy(self.T)+CtoK+self.tlapse*self.h_in[1],
                      np.copy(self.T)+self.tlapse*self.h_in[1])  # convert to Kelvin if needed
        self.dt = self.Ta - self.SST

    def _fix_coolskin_warmlayer(self, wl, cskin, skin, Rl, Rs):
        assert wl in [0,1], "wl not valid"
        assert cskin in [0,1], "cskin not valid"
        assert skin in ["C35", "ecmwf" or "Beljaars"], "Skin value not valid"

        if ((cskin == 1 or wl == 1) and (np.all(Rl == None) or np.all(np.isnan(Rl)))
            and ((np.all(Rs == None) or np.all(np.isnan(Rs))))):
            print("Cool skin/warm layer is switched ON; Radiation input should not be empty")
            raise 
61
        
Richard Cornes's avatar
Richard Cornes committed
62 63
        self.wl = wl            
        self.cskin = cskin
64
        self.skin = skin
Richard Cornes's avatar
Richard Cornes committed
65 66 67 68 69 70 71
        self.Rs = np.ones(self.spd.shape)*np.nan if Rs is None else Rs
        self.Rl = np.ones(self.spd.shape)*np.nan if Rl is None else Rl

    def set_coolskin_warmlayer(self, wl=0, cskin=0, skin="C35", Rl=None, Rs=None):
        wl = 0 if wl is None else wl
        self._fix_coolskin_warmlayer(wl, cskin, skin, Rl, Rs)

72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136
    def _update_coolskin_warmlayer(self,ind):
        if ((self.cskin == 1) and (self.wl == 0)):
            if (self.skin == "C35"):
                self.dter[ind], self.dqer[ind], self.tkt[ind] = cs_C35(np.copy(self.SST[ind]),
                                                                       self.qsea[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.lat[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.lat[ind])
                self.dqer[ind] = (self.dter[ind]*0.622*self.lv[ind]*self.qsea[ind] /
                                  (287.1*np.power(self.SST[ind], 2)))
            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.lat[ind],
                                                           np.copy(self.Qs[ind]))
                self.dqer = self.dter*0.622*self.lv*self.qsea/(287.1*np.power(self.SST, 2))
            self.skt = np.copy(self.SST)+self.dter
        elif ((self.cskin == 1) and (self.wl == 1)):
            if (self.skin == "C35"):
                self.dter[ind], self.dqer[ind], self.tkt[ind] = cs_C35(self.SST[ind], self.qsea[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.lat[ind])
                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.lat[ind])
                self.skt = np.copy(self.SST)+self.dter+self.dtwl
            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],
                                          self.sst[ind], self.lat[ind])
                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.lat[ind])
                self.skt = np.copy(self.SST)+self.dter+self.dtwl
                self.dqer[ind] = (self.dter[ind]*0.622*self.lv[ind]*self.qsea[ind] /
                             (287.1*np.power(self.skt[ind], 2)))
            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.lat[ind],
                                                           np.copy(self.Qs[ind]))
                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.lat[ind])
                self.skt = np.copy(self.SST)+self.dter+self.dtwl
                self.dqer = self.dter*0.622*self.lv*self.qsea/(287.1*np.power(self.skt, 2))
        else:
            self.dter[ind] = np.zeros(self.SST[ind].shape)
            self.dqer[ind] = np.zeros(self.SST[ind].shape)
            self.tkt[ind] = 0.001*np.ones(self.T[ind].shape)

Richard Cornes's avatar
Richard Cornes committed
137 138 139 140 141 142 143 144 145 146
    def _first_guess(self):

        # reference height1
        self.ref_ht = 10

        #  first guesses
        self.t10n, self.q10n = np.copy(self.Ta), np.copy(self.qair)
        self.tv10n = self.t10n*(1+0.6077*self.q10n)

        #  Zeng et al. 1998
Richard Cornes's avatar
Richard Cornes committed
147
        self.tv = self.th*(1+0.6077*self.qair)   # virtual potential T
Richard Cornes's avatar
Richard Cornes committed
148 149
        self.dtv = self.dt*(1+0.6077*self.qair)+0.6077*self.th*self.dq

Richard Cornes's avatar
Richard Cornes committed
150 151 152
        # Set the wind array
        self._wind_firstguess()

Richard Cornes's avatar
Richard Cornes committed
153 154 155 156 157 158 159 160 161 162 163 164 165 166 167
        # Rb eq. 11 Grachev & Fairall 1997
        Rb = self.g*10*(self.dtv)/(np.where(self.T < 200, np.copy(self.T)+CtoK, np.copy(self.T)) * np.power(self.wind, 2))
        self.monob = 1/Rb  # eq. 12 Grachev & Fairall 1997

        # ------------
        self.rho = self.P*100/(287.1*self.tv10n)
        self.lv = (2.501-0.00237*(self.SST-CtoK))*1e6  # J/kg
        
        self.dter = np.full(self.T.shape, -0.3)*self.msk
        self.tkt = np.full(self.T.shape, 0.001)*self.msk
        self.dqer = self.dter*0.622*self.lv*self.qsea/(287.1*np.power(self.SST, 2))
        self.Rnl = 0.97*(self.Rl-5.67e-8*np.power(self.SST-0.3*self.cskin, 4))
        self.Qs = 0.945*self.Rs
        self.dtwl = np.full(self.T.shape,0.3)*self.msk
        self.skt = np.copy(self.SST)
Richard Cornes's avatar
Richard Cornes committed
168
        
Richard Cornes's avatar
Richard Cornes committed
169 170 171

        self.u10n = self.wind*np.log(10/1e-4)/np.log(self.hin[0]/1e-4)
        self.usr = 0.035*self.u10n
172
        self.cd10n = cdn_calc(self.u10n, self.usr, self.Ta, self.g, self.meth)
Richard Cornes's avatar
Richard Cornes committed
173 174
        self.cd = cd_calc(self.cd10n, self.h_in[0], self.ref_ht, self.psim)
        self.usr = np.sqrt(self.cd*np.power(self.wind, 2))
175

Richard Cornes's avatar
Richard Cornes committed
176
        self.zo = np.full(self.arr_shp,1e-4)*self.msk
177 178
        self.zot, self.zoq = np.copy(self.zo), np.copy(self.zo)

Richard Cornes's avatar
Richard Cornes committed
179 180 181 182 183 184 185 186 187 188 189 190 191
        self.ct10n = np.power(kappa, 2)/(np.log(self.h_in[0]/self.zo)*np.log(self.h_in[1]/self.zot))
        self.cq10n = np.power(kappa, 2)/(np.log(self.h_in[0]/self.zo)*np.log(self.h_in[2]/self.zoq))
        
        self.ct = np.power(kappa, 2)/((np.log(self.h_in[0]/self.zo)-self.psim) * (np.log(self.h_in[1]/self.zot)-self.psit))
        self.cq = np.power(kappa, 2)/((np.log(self.h_in[0]/self.zo)-self.psim) * (np.log(self.h_in[2]/self.zoq)-self.psiq))
    
        self.tsr = (self.dt-self.dter*self.cskin-self.dtwl*self.wl)*kappa/(np.log(self.h_in[1]/self.zot) - psit_calc(self.h_in[1]/self.monob, self.meth))
        self.qsr = (self.dq-self.dqer*self.cskin)*kappa/(np.log(self.h_in[2]/self.zoq) - psit_calc(self.h_in[2]/self.monob, self.meth))

    def _zo_calc(self, ref_ht, cd10n):
        zo = ref_ht/np.exp(kappa/np.sqrt(cd10n))
        return zo
    
192
    def iterate(self, n=30, tol=None):
193 194 195 196

        if n < 5:
            warnings.warn("Iteration number <5 - resetting to 5.")
            n = 5
197 198

        # Decide which variables to use in tolerances based on tolerance specification
Richard Cornes's avatar
Richard Cornes committed
199 200 201
        tol = ['all', 0.01, 0.01, 1e-05, 1e-3, 0.1, 0.1] if tol is None else tol
        assert tol[0] in ['flux', 'ref', 'all'], "unknown tolerance input"

202 203 204 205 206 207 208 209
        old_vars = {"flux":["tau","sensible","latent"], "ref":["u10n","t10n","q10n"]}
        old_vars["all"] = old_vars["ref"] + old_vars["flux"]
        old_vars = old_vars[tol[0]]

        new_vars = {"flux":["tau","sensible","latent"], "ref":["utmp","t10n","q10n"]}
        new_vars["all"] = new_vars["ref"] + new_vars["flux"]
        new_vars = new_vars[tol[0]]

Richard Cornes's avatar
Richard Cornes committed
210 211 212
        ind = np.where(self.spd > 0)
        it = 0

213 214 215 216 217 218 219 220 221 222
        # Setup empty arrays
        self.itera = np.full(self.arr_shp,-1)*self.msk
        
        self.tsrv = np.zeros(self.arr_shp)*self.msk
        self.psim, self.psit, self.psiq = np.copy(self.tsrv), np.copy(self.tsrv), np.copy(self.tsrv)

        self.tau = np.full(self.arr_shp,0.05)*self.msk
        self.sensible = np.full(self.arr_shp,-5)*self.msk
        self.latent = np.full(self.arr_shp,-65)*self.msk

Richard Cornes's avatar
Richard Cornes committed
223 224 225 226 227 228 229 230 231
        # Generate the first guess values
        self._first_guess()

        #  iteration loop
        ii = True
        while ii:
            it += 1
            if it > n: break

232 233 234 235
            # Set the old variables (for comparison against "new")
            old = np.array([np.copy(getattr(self,i)) for i in old_vars])

            # Calculate cdn
236
            self.cd10n[ind] = cdn_calc(self.u10n[ind], self.usr[ind], self.Ta[ind], self.g[ind], self.meth)
237
            
Richard Cornes's avatar
Richard Cornes committed
238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255
            if (np.all(np.isnan(self.cd10n))):
                break
                logging.info('break %s at iteration %s cd10n<0', meth, it)
                
            self.zo[ind] = self.ref_ht/np.exp(kappa/np.sqrt(self.cd10n[ind]))
            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.ref_ht, self.psim[ind])

            self.ct10n[ind], self.cq10n[ind] = ctcqn_calc(self.h_in[1, ind]/self.monob[ind], self.cd10n[ind],
                                                          self.usr[ind], self.zo[ind], self.Ta[ind], self.meth)

            self.zot[ind] = self.ref_ht/(np.exp(np.power(kappa, 2) / (self.ct10n[ind]*np.log(self.ref_ht/self.zo[ind]))))
            self.zoq[ind] = self.ref_ht/(np.exp(np.power(kappa, 2) / (self.cq10n[ind]*np.log(self.ref_ht/self.zo[ind]))))
            self.psit[ind] = psit_calc(self.h_in[1, ind]/self.monob[ind], self.meth)
            self.psiq[ind] = psit_calc(self.h_in[2, ind]/self.monob[ind], self.meth)
            self.ct[ind], self.cq[ind] = ctcq_calc(self.cd10n[ind], self.cd[ind], self.ct10n[ind],self.cq10n[ind], self.h_in[:, ind],
                                         [self.ref_ht, self.ref_ht, self.ref_ht], self.psit[ind], self.psiq[ind])

256 257 258 259 260 261
            # Some parameterizations set a minimum on parameters
            try:
                self._minimum_params()
            except AttributeError:
                pass
            
Richard Cornes's avatar
Richard Cornes committed
262 263 264 265 266 267 268
            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[ind], self.dq[ind],
                                                self.dter[ind], self.dqer[ind],
                                                self.dtwl[ind], self.ct[ind], self.cq[ind],
                                                self.cskin, self.wl, self.meth)

269 270 271

            # Update CS/WL parameters
            self._update_coolskin_warmlayer(ind)
272
            
Richard Cornes's avatar
Richard Cornes committed
273 274 275 276 277 278 279 280 281 282 283
            # Logging output
            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))

            self.Rnl[ind] = 0.97*(self.Rl[ind]-5.67e-8*np.power(self.SST[ind] + self.dter[ind]*self.cskin, 4))
            self.t10n[ind] = (self.Ta[ind] - self.tsr[ind]/kappa*(np.log(self.h_in[1, ind]/self.ref_ht)-self.psit[ind]))
            self.q10n[ind] = (self.qair[ind] - self.qsr[ind]/kappa*(np.log(self.h_in[2, ind]/self.ref_ht)-self.psiq[ind]))
            self.tv10n[ind] = self.t10n[ind]*(1+0.6077*self.q10n[ind])

284
            self.tsrv[ind], self.monob[ind], self.Rb[ind] = get_L(self.L, self.g[ind], self.usr[ind], self.tsr[ind], self.qsr[ind], self.h_in[:, ind], self.Ta[ind],
Richard Cornes's avatar
Richard Cornes committed
285 286 287 288 289 290 291
                                                   (self.SST[ind]+self.dter[ind]*self.cskin + self.dtwl[ind]*self.wl), self.qair[ind], self.qsea[ind], self.wind[ind],
                                                   np.copy(self.monob[ind]), self.zo[ind], self.zot[ind], self.psim[ind], self.meth)

            self.psim[ind] = psim_calc(self.h_in[0, ind]/self.monob[ind], self.meth)
            self.psit[ind] = psit_calc(self.h_in[1, ind]/self.monob[ind], self.meth)
            self.psiq[ind] = psit_calc(self.h_in[2, ind]/self.monob[ind], self.meth)

Richard Cornes's avatar
Richard Cornes committed
292 293
            # Update the wind values
            self._wind_iterate(ind)
Richard Cornes's avatar
Richard Cornes committed
294 295 296 297 298 299 300 301 302 303 304 305 306
                
            self.u10n[ind] = self.wind[ind]-self.usr[ind]/kappa*(np.log(self.h_in[0, ind]/10) - self.psim[ind])

            # make sure you allow small negative values convergence
            if (it < 4):  self.u10n = np.where(self.u10n < 0, 0.5, self.u10n)

            self.utmp = np.copy(self.u10n)
            self.utmp = np.where(self.utmp < 0, np.nan, self.utmp)
            self.itera[ind] = np.ones(1)*it
            self.tau = self.rho*np.power(self.usr, 2)*(self.spd/self.wind)
            self.sensible = self.rho*self.cp*self.usr*self.tsr
            self.latent = self.rho*self.lv*self.usr*self.qsr

307 308 309
            # Set the new variables (for comparison against "old")
            new = np.array([np.copy(getattr(self,i)) for i in new_vars])

Richard Cornes's avatar
Richard Cornes committed
310 311 312 313 314 315 316 317 318 319
            if (it > 2):  # force at least two iterations
                d = np.abs(new-old)
                if (tol[0] == 'flux'):
                    ind = np.where((d[0, :] > tol[1])+(d[1, :] > tol[2]) + (d[2, :] > tol[3]))
                elif (tol[0] == 'ref'):
                    ind = np.where((d[0, :] > tol[1])+(d[1, :] > tol[2]) + (d[2, :] > tol[3]))
                elif (tol[0] == 'all'):
                    ind = np.where((d[0, :] > tol[1])+(d[1, :] > tol[2]) + (d[2, :] > tol[3])+(d[3, :] > tol[4]) +
                                   (d[4, :] > tol[5])+(d[5, :] > tol[6]))

320
            self.ind = np.copy(ind)
Richard Cornes's avatar
Richard Cornes committed
321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336
            ii = False if (ind[0].size == 0) else True
            # End of iteration loop

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


    def _get_humidity(self):
        "RH only used for flagging purposes"
        if ((self.hum[0] == 'rh') or (self.hum[0] == 'no')):
            self.rh = self.hum[1]
        elif (self.hum[0] == 'Td'):
            Td = self.hum[1]  # dew point temperature (K)
            Td = np.where(Td < 200, np.copy(Td)+CtoK, np.copy(Td))
337
            T = np.where(self.T < 200, np.copy(self.T)+CtoK, np.copy(self.T))
Richard Cornes's avatar
Richard Cornes committed
338 339 340 341 342 343 344 345 346 347 348
            #T = 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
        
    def _flag(self,out=0):
        "Set the general flags"
        
        flag = np.full(self.arr_shp, "n",dtype="object")

        if (self.hum[0] == 'no'):
349 350 351 352
            if (self.cskin == 1):
                flag = np.where(np.isnan(self.spd+self.T+self.SST+self.P+self.Rs+self.Rl), "m", flag)
            else:
                flag = np.where(np.isnan(self.spd+self.T+self.SST+self.P), "m", flag)
Richard Cornes's avatar
Richard Cornes committed
353
        else:
354 355 356 357 358
            if (self.cskin == 1):
                flag = np.where(np.isnan(self.spd+self.T+self.SST+self.hum[1]+self.P+self.Rs+self.Rl), "m", flag)
            else:
                flag = np.where(np.isnan(self.spd+self.T+self.SST+self.hum[1]+self.P), "m", flag)

Richard Cornes's avatar
Richard Cornes committed
359
            flag = np.where(self.rh > 100, "r", flag)
360
        
Richard Cornes's avatar
Richard Cornes committed
361

362 363
        # u10n flag
        flag = np.where(((self.u10n < 0)  | (self.u10n > 999)) & (flag == "n"), "u",
Richard Cornes's avatar
Richard Cornes committed
364 365 366
                             np.where((self.u10n < 0) &
                                      (np.char.find(flag.astype(str), 'u') == -1),
                                      flag+[","]+["u"], flag))
367 368
        # q10n flag
        flag = np.where(((self.q10n < 0) | (self.q10n > 999)) & (flag == "n"), "q",
Richard Cornes's avatar
Richard Cornes committed
369 370
                             np.where((self.q10n < 0) & (flag != "n"), flag+[","]+["q"],
                                      flag))
371 372 373 374 375

        # t10n flag (not currently used)
        flag = np.where((self.t10n < -999) & (flag == "n"), "t",
                             np.where((self.t10n < 0) & (flag != "n"), flag+[","]+["t"],
                                      flag))
Richard Cornes's avatar
Richard Cornes committed
376 377 378 379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400
        
        flag = np.where(((self.Rb < -0.5) | (self.Rb > 0.2) | ((self.hin[0]/self.monob) > 1000)) &
                             (flag == "n"), "l",
                             np.where(((self.Rb < -0.5) | (self.Rb > 0.2) |
                                       ((self.hin[0]/self.monob) > 1000)) &
                                      (flag != "n"), flag+[","]+["l"], flag))

        if (out == 1):
            flag = np.where((self.itera == -1) & (flag == "n"), "i",
                                 np.where((self.itera == -1) &
                                          ((flag != "n") &
                                           (np.char.find(flag.astype(str), 'm') == -1)),
                                          flag+[","]+["i"], flag))
        else:
            flag = np.where((self.itera == -1) & (flag == "n"), "i",
                                 np.where((self.itera == -1) &
                                          ((flag != "n") &
                                           (np.char.find(flag.astype(str), 'm') == -1) &
                                           (np.char.find(flag.astype(str), 'u') == -1)),
                                          flag+[","]+["i"], flag))
        self.flag = flag
            
    def get_output(self,out=0):

        assert out in [0,1], "out must be either 0 or 1"
401 402 403

        self._get_humidity() # Get the Relative humidity
        self._flag(out=out)  # Get flags
Richard Cornes's avatar
Richard Cornes committed
404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433
        
        # calculate output parameters
        rho = (0.34838*self.P)/(self.tv10n)
        self.t10n = self.t10n-(273.16+self.tlapse*self.ref_ht)
        
        # solve for zo from cd10n
        zo = self.ref_ht/np.exp(kappa/np.sqrt(self.cd10n))
        
        # adjust neutral cdn at any output height
        self.cdn = np.power(kappa/np.log(self.hout/zo), 2)
        self.cd = cd_calc(self.cdn, self.h_out[0], self.h_out[0], self.psim)

        # solve for zot, zoq from ct10n, cq10n
        zot = self.ref_ht/(np.exp(kappa**2/(self.ct10n*np.log(self.ref_ht/zo))))
        zoq = self.ref_ht/(np.exp(kappa**2/(self.cq10n*np.log(self.ref_ht/zo))))
        
        # adjust neutral ctn, cqn at any output height
        self.ctn = np.power(kappa, 2)/(np.log(self.h_out[0]/zo)*np.log(self.h_out[1]/zot))
        self.cqn = np.power(kappa, 2)/(np.log(self.h_out[0]/zo)*np.log(self.h_out[2]/zoq))
        self.ct, self.cq = ctcq_calc(self.cdn, self.cd, self.ctn, self.cqn, self.h_out, self.h_out, self.psit, self.psiq)
        self.uref = (self.spd-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)))
        tref = (self.Ta-self.tsr/kappa*(np.log(self.h_in[1]/self.h_out[1])-self.psit + psit_calc(self.h_out[0]/self.monob, self.meth)))
        self.tref = tref-(CtoK+self.tlapse*self.h_out[1])
        self.qref = (self.qair-self.qsr/kappa*(np.log(self.h_in[2]/self.h_out[2]) - self.psit+psit_calc(self.h_out[2]/self.monob, self.meth)))

        if (self.wl == 0): self.dtwl = np.zeros(self.T.shape)*self.msk  # reset to zero if not used

        # Do not calculate lhf if a measure of humidity is not input
        if (self.hum[0] == 'no'):
            self.latent = np.ones(self.SST.shape)*np.nan
434 435 436 437 438
            self.qsr = np.copy(self.latent)
            self.q10n = np.copy(self.latent)
            self.qref = np.copy(self.latent)
            self.qair = np.copy(self.latent)
            self.rh =  np.copy(self.latent)
Richard Cornes's avatar
Richard Cornes committed
439

440
        # Set the final wind speed values
Richard Cornes's avatar
Richard Cornes committed
441 442
        self.wind_spd = np.sqrt(np.power(self.wind, 2)-np.power(self.spd, 2))

Richard Cornes's avatar
Richard Cornes committed
443
        # Get class specific flags (will only work if self.utmp_hi and self.utmp_lo have been set in the class)
444 445 446 447
        try:
            self._class_flag()
        except AttributeError:
            pass
Richard Cornes's avatar
Richard Cornes committed
448 449 450 451 452 453 454 455 456 457 458 459

        # Combine all output variables into a pandas array
        res_vars = ("tau","sensible","latent","monob","cd","cdn","ct","ctn","cq","cqn","tsrv","tsr","qsr","usr","psim","psit",
                    "psiq","u10n","t10n","tv10n","q10n","zo","zot","zoq","uref","tref","qref","itera","dter","dqer","dtwl",
                    "qair","qsea","Rl","Rs","Rnl","wind_spd","Rb","rh","tkt","lv")

        res = np.zeros((len(res_vars), len(self.spd)))
        for i, value in enumerate(res_vars): res[i][:] = getattr(self, value)

        if (out == 0):
            res[:, self.ind] = np.nan
            # set missing values where data have non acceptable values
460 461
            if (self.hum[0] != 'no'): res = np.asarray([np.where(self.q10n < 0, np.nan, res[i][:]) for i in range(len(res_vars))]) # FIXME: why 41?
            res = np.asarray([np.where(self.u10n < 0, np.nan, res[i][:]) for i in range(len(res_vars))])
Richard Cornes's avatar
Richard Cornes committed
462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502
        else:
            warnings.warn("Warning: the output will contain values for points that have not converged and negative values (if any) for u10n/q10n")

        resAll = pd.DataFrame(data=res.T, index=range(self.nlen), columns=res_vars)
    
        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
        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
        self.arr_shp = spd.shape
        self.nlen = len(spd)
        self.spd = spd
        self.T = T
        self.hum = ['no', np.full(SST.shape,80)] if hum is None else hum
        self.SST = np.where(SST < 200, np.copy(SST)+CtoK, np.copy(SST))
        self.lat = np.full(self.arr_shp,45) if lat is None else lat
        self.g = gc(self.lat)
        self.P = np.full(n, 1013) if P is None else P

        # mask to preserve missing values when initialising variables
        self.msk=np.empty(SST.shape)
        self.msk = np.where(np.isnan(spd+T+SST), np.nan, 1)
        self.Rb = np.empty(SST.shape)*self.msk
        self.dtwl = np.full(T.shape,0.3)*self.msk
        
    def add_gust(self,gust=None):

        if np.all(gust == None):
            try:
                gust = self.default_gust
            except AttributeError:
                gust = [1,1.2,800]
        elif ((np.size(gust) < 3) and (gust == 0)):
            gust = [0, 0, 0]
            
        assert np.size(gust) == 3, "gust input must be a 3x1 array"
Richard Cornes's avatar
Richard Cornes committed
503
        assert gust[0] in [0,1], "gust at position 0 must be 0 or 1"
Richard Cornes's avatar
Richard Cornes committed
504 505
        self.gust = gust

506
    def _class_flag(self):
Richard Cornes's avatar
Richard Cornes committed
507
        "A flag specific to this class - only used for certain classes where utmp_lo and utmp_hi are defined"
Richard Cornes's avatar
Richard Cornes committed
508 509
        self.flag = np.where(((self.utmp < self.utmp_lo[0]) | (self.utmp > self.utmp_hi[0])) & (self.flag == "n"), "o",
                             np.where(((self.utmp < self.utmp_lo[1]) | (self.utmp > self.utmp_hi[1])) &
510 511 512 513
                                      ((self.flag != "n") &
                                       (np.char.find(self.flag.astype(str), 'u') == -1) &
                                       (np.char.find(self.flag.astype(str), 'q') == -1)),
                                      self.flag+[","]+["o"], self.flag))
Richard Cornes's avatar
Richard Cornes committed
514

Richard Cornes's avatar
Richard Cornes committed
515 516 517 518 519
    def __init__(self):
        self.meth = "S88"

class S80(S88):

Richard Cornes's avatar
Richard Cornes committed
520
    def __init__(self):
521
        self.meth = "S80"
Richard Cornes's avatar
Richard Cornes committed
522 523
        self.utmp_lo = [6,6]
        self.utmp_hi = [22,22]
524
        
525
class YT96(S88):
Richard Cornes's avatar
Richard Cornes committed
526 527 528

    def __init__(self):
        self.meth = "YT96"
Richard Cornes's avatar
Richard Cornes committed
529 530
        self.utmp_lo = [0,3]
        self.utmp_hi = [26,26]
Richard Cornes's avatar
Richard Cornes committed
531

532
class LP82(S88):
Richard Cornes's avatar
Richard Cornes committed
533 534 535
    
    def __init__(self):
        self.meth = "LP82"
Richard Cornes's avatar
Richard Cornes committed
536 537
        self.utmp_lo = [3,3]
        self.utmp_hi = [25,25]
Richard Cornes's avatar
Richard Cornes committed
538

Richard Cornes's avatar
Richard Cornes committed
539

540
class NCAR(S88):
Richard Cornes's avatar
Richard Cornes committed
541

542 543 544 545 546 547
    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)
        self.zo = np.minimum(np.copy(self.zo), 0.0025)

Richard Cornes's avatar
Richard Cornes committed
548 549 550 551 552 553 554 555
    def _zo_calc(self, ref_ht, cd10n):
        "Special z0 calculation for NCAR"
        zo = ref_ht/np.exp(kappa/np.sqrt(cd10n))
        zo = np.minimum(np.copy(zo), 0.0025)
        return zo
    
    def __init__(self):
        self.meth = "NCAR"
Richard Cornes's avatar
Richard Cornes committed
556 557
        self.utmp_lo = 0.5
        self.utmp_hi = 999
Richard Cornes's avatar
Richard Cornes committed
558 559


Richard Cornes's avatar
Richard Cornes committed
560
class UA(S88):
Richard Cornes's avatar
Richard Cornes committed
561
    
Richard Cornes's avatar
Richard Cornes committed
562
    def _wind_firstguess(self):
Richard Cornes's avatar
Richard Cornes committed
563 564 565 566
        # gustiness adjustment
        if (self.gust[0] == 1):
            self.wind = np.where(self.dtv >= 0, np.where(self.spd > 0.1, self.spd, 0.1),
                                 np.sqrt(np.power(np.copy(self.spd), 2)+np.power(0.5, 2)))
Richard Cornes's avatar
Richard Cornes committed
567 568 569 570 571 572 573 574 575
        else:
            self.wind = np.copy(self.spd)
    def _wind_iterate(self,ind):
        if(self.gust[0] == 1):
            self.wind[ind] = np.where(self.dtv[ind] >= 0, np.where(self.spd[ind] > 0.1,
                                                         self.spd[ind], 0.1),
                                      np.sqrt(np.power(np.copy(self.spd[ind]), 2) +
                                         np.power(get_gust(self.gust[1], self.tv[ind],
                                                           self.usr[ind], self.tsrv[ind],
576
                                                           self.gust[2], self.g[ind]),
Richard Cornes's avatar
Richard Cornes committed
577
                                                  2)))  # Zeng et al. 1998 (20)
Richard Cornes's avatar
Richard Cornes committed
578 579 580
    def __init__(self):
        self.meth = "UA"
        self.default_gust = [1,1,1000]
Richard Cornes's avatar
Richard Cornes committed
581 582
        self.utmp_lo = [-999,-999]
        self.utmp_hi = [18,18]
Richard Cornes's avatar
Richard Cornes committed
583

Richard Cornes's avatar
Richard Cornes committed
584

585
class C30(S88):
Richard Cornes's avatar
Richard Cornes committed
586 587 588 589 590 591
    def set_coolskin_warmlayer(self, wl=0, cskin=1, skin="C35", Rl=None, Rs=None):
        self._fix_coolskin_warmlayer(wl, cskin, skin, Rl, Rs)

    def __init__(self):
        self.meth = "C30"
        self.default_gust = [1,1.2,600]
sbiri's avatar
sbiri committed
592

Richard Cornes's avatar
Richard Cornes committed
593 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611 612
class C35(C30):
    def __init__(self):
        self.meth = "C35"
        self.default_gust = [1,1.2,600]

class ecmwf(C30):
    def set_coolskin_warmlayer(self, wl=0, cskin=1, skin="ecmwf", Rl=None, Rs=None):
        self._fix_coolskin_warmlayer(wl, cskin, skin, Rl, Rs)

    def __init__(self):
        self.meth = "ecmwf"
        self.default_gust = [1,1,1000]

class Beljaars(C30):
    def set_coolskin_warmlayer(self, wl=0, cskin=1, skin="Beljaars", Rl=None, Rs=None):
        self._fix_coolskin_warmlayer(wl, cskin, skin, Rl, Rs)
        
    def __init__(self):
        self.meth = "Beljaars"
        self.default_gust = [1,1,1000]
613
        
614 615
def AirSeaFluxCode(spd, T, SST, lat=None, hum=None, P=None, hin=18, hout=10,
                   Rl=None, Rs=None, cskin=None, skin="C35", wl=0, gust=None,
616
                   meth="S88", qmeth="Buck2", tol=None, n=30, out=0, L=None):
617 618 619
    """
    Calculates turbulent surface fluxes using different parameterizations
    Calculates height adjusted values for spd, T, q
sbiri's avatar
sbiri committed
620 621 622 623 624 625 626 627 628 629 630

    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
631 632 633
            latitude (deg), default 45deg
        hum : float
            humidity input switch 2x1 [x, values] default is relative humidity
634 635 636
            x='rh' : relative humidity in %
            x='q' : specific humidity (g/kg)
            x='Td' : dew point temperature (K)
sbiri's avatar
sbiri committed
637
        P : float
638
            air pressure (hPa), default 1013hPa
sbiri's avatar
sbiri committed
639
        hin : float
640
            sensor heights in m (array 3x1 or 3xn), default 18m
sbiri's avatar
sbiri committed
641 642 643 644 645 646
        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
647
        cskin : int
648 649
            0 switch cool skin adjustment off, else 1
            default is 1
650 651 652 653
        skin : str
            cool skin method option "C35", "ecmwf" or "Beljaars"
        wl : int
            warm layer correction default is 0, to switch on set to 1
654
        gust : int
655 656
            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
657
            zi PBL height (m) 600 for COARE, 1000 for UA and ecmwf, 800 default
658
            default for COARE [1, 1.2, 600]
659
            default for UA, ecmwf [1, 1, 1000]
660 661
            default else [1, 1.2, 800]
        meth : str
sbiri's avatar
sbiri committed
662
            "S80", "S88", "LP82", "YT96", "UA", "NCAR", "C30", "C35",
663
            "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
        n : int
678 679
            number of iterations (defautl = 10)
        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
        L : str
sbiri's avatar
sbiri committed
684
           Monin-Obukhov length definition options
sbiri's avatar
sbiri committed
685
           "tsrv"  : default for "S80", "S88", "LP82", "YT96", "UA", "NCAR",
sbiri's avatar
sbiri committed
686 687 688
                     "C30", "C35"
           "Rb" : following ecmwf (IFS Documentation cy46r1), default for
                  "ecmwf", "Beljaars"
sbiri's avatar
sbiri committed
689 690 691
    Returns
    -------
        res : array that contains
sbiri's avatar
sbiri committed
692 693 694 695
                       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
696 697
                       5. drag coefficient (cd)
                       6. neutral drag coefficient (cdn)
sbiri's avatar
sbiri committed
698 699
                       7. heat exchange coefficient (ct)
                       8. neutral heat exchange coefficient (ctn)
sbiri's avatar
sbiri committed
700
                       9. moisture exhange coefficient (cq)
sbiri's avatar
sbiri committed
701 702
                       10. neutral moisture exchange coefficient (cqn)
                       11. star virtual temperatcure (tsrv)
sbiri's avatar
sbiri committed
703
                       12. star temperature (tsr)
704 705
                       13. star specific humidity (qsr)
                       14. star wind speed (usr)
sbiri's avatar
sbiri committed
706
                       15. momentum stability function (psim)
707 708
                       16. heat stability function (psit)
                       17. moisture stability function (psiq)
709
                       18. 10m neutral wind speed (u10n)
710 711 712 713 714 715
                       19. 10m neutral temperature (t10n)
                       20. 10m neutral virtual temperature (tv10n)
                       21. 10m neutral specific humidity (q10n)
                       22. surface roughness length (zo)
                       23. heat roughness length (zot)
                       24. moisture roughness length (zoq)
sbiri's avatar
sbiri committed
716
                       25. wind speed at reference height (uref)
717 718
                       26. temperature at reference height (tref)
                       27. specific humidity at reference height (qref)
719
                       28. number of iterations until convergence
720 721
                       29. cool-skin temperature depression (dter)
                       30. cool-skin humidity depression (dqer)
722 723 724 725 726 727
                       31. warm layer correction (dtwl)
                       32. specific humidity of air (qair)
                       33. specific humidity at sea surface (qsea)
                       34. downward longwave radiation (Rl)
                       35. downward shortwave radiation (Rs)
                       36. downward net longwave radiation (Rnl)
728 729 730
                       37. gust wind speed (ug)
                       38. Bulk Richardson number (Rib)
                       39. relative humidity (rh)
731 732
                       40. thickness of the viscous layer (delta)
                       41. lv latent heat of vaporization (Jkg−1)
733
                       42. flag ("n": normal, "o": out of nominal range,
734
                                 "u": u10n<0, "q":q10n<0
735
                                 "m": missing,
736
                                 "l": Rib<-0.5 or Rib>0.2 or z/L>1000,
sbiri's avatar
sbiri committed
737
                                 "r" : rh>100%,
738
                                 "i": convergence fail at n)
739

740
    2021 / Author S. Biri
sbiri's avatar
sbiri committed
741
    """
sbiri's avatar
sbiri committed
742
    logging.basicConfig(filename='flux_calc.log', filemode="w",
743
                        format='%(asctime)s %(message)s', level=logging.INFO)
744
    logging.captureWarnings(True)
Richard Cornes's avatar
Richard Cornes committed
745 746

    iclass = globals()[meth]()
747
    iclass.add_gust(gust=gust)
Richard Cornes's avatar
Richard Cornes committed
748 749 750 751
    iclass.add_variables(spd, T, SST, lat=lat, hum=hum, P=P, L=L)
    iclass.get_heights(hin, hout)
    iclass.get_specHumidity(qmeth=qmeth)
    iclass.set_coolskin_warmlayer(wl=wl, cskin=cskin,skin=skin,Rl=Rl,Rs=Rs)
Richard Cornes's avatar
Richard Cornes committed
752
    iclass.iterate(tol=tol,n=n)
753
    resAll = iclass.get_output(out=out)
Richard Cornes's avatar
Richard Cornes committed
754

755
    return resAll