AirSeaFluxCode.py 35 KB
Newer Older
sbiri's avatar
sbiri committed
1
import numpy as np
2
import pandas as pd
sbiri's avatar
sbiri committed
3
import logging
4
from hum_subs import (get_hum, gamma)
Richard Cornes's avatar
Richard Cornes committed
5 6
from util_subs import *
from flux_subs import *
sbiri's avatar
sbiri committed
7

8
class S88:
Richard Cornes's avatar
Richard Cornes committed
9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46
    
    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 
47
        
Richard Cornes's avatar
Richard Cornes committed
48 49
        self.wl = wl            
        self.cskin = cskin
50
        self.skin = skin
Richard Cornes's avatar
Richard Cornes committed
51 52 53 54 55 56 57
        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)

58 59 60 61 62 63 64 65 66 67 68 69 70 71 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
    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
123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162
    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
        tv = self.th*(1+0.6077*self.qair)   # virtual potential T
        self.dtv = self.dt*(1+0.6077*self.qair)+0.6077*self.th*self.dq

        # 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)

        # Apply the gustiness adjustment if defined for this class
        try:
            self._adjust_gust()
        except AttributeError:
            pass

        self.u10n = self.wind*np.log(10/1e-4)/np.log(self.hin[0]/1e-4)
        self.usr = 0.035*self.u10n
        self.cd10n = cdn_calc(self.u10n, self.usr, self.Ta, self.lat, self.meth)
        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))
163

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

Richard Cornes's avatar
Richard Cornes committed
167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182
        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
    
    def iterate(self,n=10, tol=None):
        
        n = 5 if n < 5 else n
183 184

        # Decide which variables to use in tolerances based on tolerance specification
Richard Cornes's avatar
Richard Cornes committed
185 186 187
        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"

188 189 190 191 192 193 194 195
        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
196 197 198
        ind = np.where(self.spd > 0)
        it = 0

199 200 201 202 203 204 205 206 207 208
        # 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
209 210 211 212 213 214 215 216 217
        # Generate the first guess values
        self._first_guess()

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

218 219 220 221
            # Set the old variables (for comparison against "new")
            old = np.array([np.copy(getattr(self,i)) for i in old_vars])

            # Calculate cdn
Richard Cornes's avatar
Richard Cornes committed
222
            self.cd10n[ind] = cdn_calc(self.u10n[ind], self.usr[ind], self.Ta[ind], self.lat[ind], self.meth)
223
            
Richard Cornes's avatar
Richard Cornes committed
224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241
            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])

242 243 244 245 246 247
            # Some parameterizations set a minimum on parameters
            try:
                self._minimum_params()
            except AttributeError:
                pass
            
Richard Cornes's avatar
Richard Cornes committed
248 249 250 251 252 253 254
            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)

255 256 257

            # Update CS/WL parameters
            self._update_coolskin_warmlayer(ind)
258
            
Richard Cornes's avatar
Richard Cornes committed
259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277
            # 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])

            self.tsrv[ind], self.monob[ind], self.Rb[ind] = get_L(self.L, self.lat[ind], self.usr[ind], self.tsr[ind], self.qsr[ind], self.h_in[:, ind], self.Ta[ind],
                                                   (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
278
            # gust[0] is asserted to be either 0 or 1
Richard Cornes's avatar
Richard Cornes committed
279 280 281
            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], self.lat[ind]), 2))
Richard Cornes's avatar
Richard Cornes committed
282
            else:
Richard Cornes's avatar
Richard Cornes committed
283 284 285 286 287 288 289 290 291 292 293 294 295 296 297
                self.wind[ind] = np.copy(self.spd[ind])

                
            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

298 299 300
            # 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
301 302 303 304 305 306 307 308 309 310
            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]))

311
            self.ind = np.copy(ind)
Richard Cornes's avatar
Richard Cornes committed
312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327
            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))
328
            T = np.where(self.T < 200, np.copy(self.T)+CtoK, np.copy(self.T))
Richard Cornes's avatar
Richard Cornes committed
329 330 331 332 333 334 335 336 337 338 339
            #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'):
340 341 342 343
            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
344
        else:
345 346 347 348 349
            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
350
            flag = np.where(self.rh > 100, "r", flag)
351
        
Richard Cornes's avatar
Richard Cornes committed
352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373 374 375 376 377 378 379 380 381 382 383 384 385

        flag = np.where((self.u10n < 0) & (flag == "n"), "u",
                             np.where((self.u10n < 0) &
                                      (np.char.find(flag.astype(str), 'u') == -1),
                                      flag+[","]+["u"], flag))
        
        flag = np.where((self.q10n < 0) & (flag == "n"), "q",
                             np.where((self.q10n < 0) & (flag != "n"), flag+[","]+["q"],
                                      flag))
        
        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"
386 387 388

        self._get_humidity() # Get the Relative humidity
        self._flag(out=out)  # Get flags
Richard Cornes's avatar
Richard Cornes committed
389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418
        
        # 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
419 420 421 422 423
            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
424

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

428 429 430 431 432
        # Get class specific flags
        try:
            self._class_flag()
        except AttributeError:
            pass
Richard Cornes's avatar
Richard Cornes committed
433 434 435 436 437 438 439 440 441 442 443 444

        # 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
445 446
            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
447 448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465 466 467 468 469 470 471 472 473 474 475 476 477
        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

        # Set the wind array
Richard Cornes's avatar
Richard Cornes committed
478
        # gust[0] is asserted to be either 0 or 1
Richard Cornes's avatar
Richard Cornes committed
479 480
        if (self.gust[0] == 1):
            self.wind = np.sqrt(np.power(np.copy(self.spd), 2)+np.power(0.5, 2))
Richard Cornes's avatar
Richard Cornes committed
481
        else:
Richard Cornes's avatar
Richard Cornes committed
482 483 484 485 486 487 488 489 490 491 492 493 494
            self.wind = np.copy(spd)
        
    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
495
        assert gust[0] in [0,1], "gust at position 0 must be 0 or 1"
Richard Cornes's avatar
Richard Cornes committed
496 497
        self.gust = gust

498
    def _class_flag(self):
Richard Cornes's avatar
Richard Cornes committed
499 500 501
        "A flag specific to this class - only used for certain classes where utmp_lo and utmp_hi are defined"
        self.flag = np.where(((self.utmp < self.utmp_lo) | (self.utmp > self.utmp_hi)) & (self.flag == "n"), "o",
                             np.where(((self.utmp < self.utmp_lo) | (self.utmp > self.utmp_hi)) &
502 503 504 505
                                      ((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
506

Richard Cornes's avatar
Richard Cornes committed
507 508 509 510 511
    def __init__(self):
        self.meth = "S88"

class S80(S88):

Richard Cornes's avatar
Richard Cornes committed
512
    def __init__(self):
513
        self.meth = "S80"
Richard Cornes's avatar
Richard Cornes committed
514 515
        self.utmp_lo = 6
        self.utmp_hi = 22
516
        
517
class YT96(S88):
Richard Cornes's avatar
Richard Cornes committed
518 519 520

    def __init__(self):
        self.meth = "YT96"
Richard Cornes's avatar
Richard Cornes committed
521 522
        self.utmp_lo = 0
        self.utmp_hi = 26
Richard Cornes's avatar
Richard Cornes committed
523

524
class LP82(S88):
Richard Cornes's avatar
Richard Cornes committed
525 526 527
    
    def __init__(self):
        self.meth = "LP82"
Richard Cornes's avatar
Richard Cornes committed
528 529 530
        self.utmp_lo = 3
        self.utmp_hi = 25

Richard Cornes's avatar
Richard Cornes committed
531

532
class NCAR(S88):
Richard Cornes's avatar
Richard Cornes committed
533

534 535 536 537 538 539
    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
540 541 542 543 544 545 546 547
    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
548 549
        self.utmp_lo = 0.5
        self.utmp_hi = 999
Richard Cornes's avatar
Richard Cornes committed
550 551


Richard Cornes's avatar
Richard Cornes committed
552
class UA(S88):
Richard Cornes's avatar
Richard Cornes committed
553 554 555 556 557 558 559 560 561 562
    
    def _adjust_gust(self):
        # 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)))

    def __init__(self):
        self.meth = "UA"
        self.default_gust = [1,1,1000]
Richard Cornes's avatar
Richard Cornes committed
563 564 565
        self.utmp_lo = 18
        self.utmp_hi = 999

Richard Cornes's avatar
Richard Cornes committed
566

567
class C30(S88):
Richard Cornes's avatar
Richard Cornes committed
568 569 570 571 572 573
    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
574

Richard Cornes's avatar
Richard Cornes committed
575 576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593 594
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]
595
        
596 597
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,
598
                   meth="S88", qmeth="Buck2", tol=None, n=10, out=0, L=None):
599 600 601
    """
    Calculates turbulent surface fluxes using different parameterizations
    Calculates height adjusted values for spd, T, q
sbiri's avatar
sbiri committed
602 603 604 605 606 607 608 609 610 611 612

    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
613 614 615
            latitude (deg), default 45deg
        hum : float
            humidity input switch 2x1 [x, values] default is relative humidity
616 617 618
            x='rh' : relative humidity in %
            x='q' : specific humidity (g/kg)
            x='Td' : dew point temperature (K)
sbiri's avatar
sbiri committed
619
        P : float
620
            air pressure (hPa), default 1013hPa
sbiri's avatar
sbiri committed
621
        hin : float
622
            sensor heights in m (array 3x1 or 3xn), default 18m
sbiri's avatar
sbiri committed
623 624 625 626 627 628
        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
629
        cskin : int
630 631
            0 switch cool skin adjustment off, else 1
            default is 1
632 633 634 635
        skin : str
            cool skin method option "C35", "ecmwf" or "Beljaars"
        wl : int
            warm layer correction default is 0, to switch on set to 1
636
        gust : int
637 638
            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
639
            zi PBL height (m) 600 for COARE, 1000 for UA and ecmwf, 800 default
640
            default for COARE [1, 1.2, 600]
641
            default for UA, ecmwf [1, 1, 1000]
642 643
            default else [1, 1.2, 800]
        meth : str
sbiri's avatar
sbiri committed
644
            "S80", "S88", "LP82", "YT96", "UA", "NCAR", "C30", "C35",
645
            "ecmwf", "Beljaars"
646 647
        qmeth : str
            is the saturation evaporation method to use amongst
648 649
            "HylandWexler","Hardy","Preining","Wexler","GoffGratch","WMO",
            "MagnusTetens","Buck","Buck2","WMO2018","Sonntag","Bolton",
650 651 652 653 654 655 656
            "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
657 658
                    adjustment lim1-6
           default is tol=['all', 0.01, 0.01, 1e-05, 1e-3, 0.1, 0.1]
sbiri's avatar
sbiri committed
659
        n : int
660 661
            number of iterations (defautl = 10)
        out : int
662 663
            set 0 to set points that have not converged, negative values of
                  u10n and q10n to missing (default)
664
            set 1 to keep points
665
        L : str
sbiri's avatar
sbiri committed
666
           Monin-Obukhov length definition options
sbiri's avatar
sbiri committed
667
           "tsrv"  : default for "S80", "S88", "LP82", "YT96", "UA", "NCAR",
sbiri's avatar
sbiri committed
668 669 670
                     "C30", "C35"
           "Rb" : following ecmwf (IFS Documentation cy46r1), default for
                  "ecmwf", "Beljaars"
sbiri's avatar
sbiri committed
671 672 673
    Returns
    -------
        res : array that contains
sbiri's avatar
sbiri committed
674 675 676 677
                       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
678 679
                       5. drag coefficient (cd)
                       6. neutral drag coefficient (cdn)
sbiri's avatar
sbiri committed
680 681
                       7. heat exchange coefficient (ct)
                       8. neutral heat exchange coefficient (ctn)
sbiri's avatar
sbiri committed
682
                       9. moisture exhange coefficient (cq)
sbiri's avatar
sbiri committed
683 684
                       10. neutral moisture exchange coefficient (cqn)
                       11. star virtual temperatcure (tsrv)
sbiri's avatar
sbiri committed
685
                       12. star temperature (tsr)
686 687
                       13. star specific humidity (qsr)
                       14. star wind speed (usr)
sbiri's avatar
sbiri committed
688
                       15. momentum stability function (psim)
689 690
                       16. heat stability function (psit)
                       17. moisture stability function (psiq)
691
                       18. 10m neutral wind speed (u10n)
692 693 694 695 696 697
                       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
698
                       25. wind speed at reference height (uref)
699 700
                       26. temperature at reference height (tref)
                       27. specific humidity at reference height (qref)
701
                       28. number of iterations until convergence
702 703
                       29. cool-skin temperature depression (dter)
                       30. cool-skin humidity depression (dqer)
704 705 706 707 708 709
                       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)
710 711 712
                       37. gust wind speed (ug)
                       38. Bulk Richardson number (Rib)
                       39. relative humidity (rh)
713 714
                       40. thickness of the viscous layer (delta)
                       41. lv latent heat of vaporization (Jkg−1)
715
                       42. flag ("n": normal, "o": out of nominal range,
716
                                 "u": u10n<0, "q":q10n<0
717
                                 "m": missing,
718
                                 "l": Rib<-0.5 or Rib>0.2 or z/L>1000,
sbiri's avatar
sbiri committed
719
                                 "r" : rh>100%,
720
                                 "i": convergence fail at n)
721

722
    2021 / Author S. Biri
sbiri's avatar
sbiri committed
723
    """
sbiri's avatar
sbiri committed
724
    logging.basicConfig(filename='flux_calc.log', filemode="w",
725
                        format='%(asctime)s %(message)s', level=logging.INFO)
726
    logging.captureWarnings(True)
Richard Cornes's avatar
Richard Cornes committed
727 728

    iclass = globals()[meth]()
729
    iclass.add_gust(gust=gust)
Richard Cornes's avatar
Richard Cornes committed
730 731 732 733
    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
734
    iclass.iterate(tol=tol,n=n)
735
    resAll = iclass.get_output(out=out)
Richard Cornes's avatar
Richard Cornes committed
736

737
    return resAll