inver_sens2pv_sens2a.py

import multiprocessing as mp
from multiprocessing import Manager

import numpy as np
import xarray as xr
from netCDF4 import Dataset
from wrf import getvar, vinterp, get_cartopy, destagger, latlon_coords
import numpy.ma as ma
import time
from glob import glob

import metpy.constants as mpconst
import metpy.calc as mpcalc
import standard_atmosphere as sa
from qgpv_invert import qgpv_invert
# from calculate_vorticity_and_divergence import vorticity

def vorticity(u,v,msfm,dx,dy):
#    Note all data must be on T-pts
    import numpy as np
    
    nz = u.shape[0]
    ny = msfm.shape[0]
    nx = msfm.shape[1]
    mm = msfm*msfm

    dvdx = np.zeros([nz,ny,nx])
    dudy = np.zeros([nz,ny,nx])
    vor = np.zeros([nz,ny,nx])
    for j in np.arange(1,ny-1):
        jp1=min(j+1,ny)
        jm1=max(j-1,1)
        for i in np.arange(1,nx-1):
            ip1=min(i+1,nx)
            im1=max(i-1,1)   
            dsx=(ip1-im1)*dx
            dsy=(jp1-jm1)*dy
# Careful with map factors...
            dvdx[:,j,i]=mm[j,i]*(v[:,j,ip1]/msfm[j,ip1]-v[:,j,im1]/msfm[j,im1])/dsx
            dudy[:,j,i]=mm[j,i]*(u[:,jp1,i]/msfm[jp1,i]-u[:,jm1,i]/msfm[jm1,i])/dsy
    vor = dvdx - dudy
    return vor

def create_terrain_mask(terrain, zlevels):
    nz = len(zlevels)
    ny, nx = terrain.shape
    terrain_mask = np.ones((nz, ny, nx))
    for k in range(6):
        terrain_mask[k,:,:] = np.invert(terrain>zlevels[k])
        
    return terrain_mask

def repetition(au, av, at, BC, p_levels, f0, ds, msfm, SS, gamma, terrain_mask, dp=5000, iterations=300, threshold=10):
                
    """
    au: sens to u
    av: snes to v
    at: sens to theta
    BC: initial condition/old inverted field
    p_levels: in hPa
    dp: delta p, integer, not array
    """
    time_block_start = time.time()
    
    ## calculate forcing (sensitivity to streamfunctin*1/f0)
    
    nz, ny, nx = au.shape
    
    #ds = 1./wrfout.variables['RDX'][0]
    vort = np.zeros(au.shape)     # vort is defined on the mass points
    vort = vorticity(au, av, msfm, ds, ds)

    dtheta_hat_dp = np.zeros(au.shape)

    for k in np.arange(1,nz-1):
        dtheta_hat_dp[k,:] = ((at[k-1,:]/gamma[k-1])-(at[k+1,:]/gamma[k+1]))/(2.*dp)
    
    ### Forcing is sensitivity to streamfunction psi
    ### Following manuscript equation 7
    ### Is there a 1/f0 ?
    forcing = (-1./f0)*(vort-f0*dtheta_hat_dp)   
    forcing = forcing * terrain_mask
    
    print("start inversion",time.time())
    
    ## Inversion

    res = np.zeros(au.shape)

    BC_flag = 1

    Tb = at[0,:]
    Tt = at[-1,:]
    
    ### QA is sensitivity to QGPV
    ### following manuscript equation 5
    ### there should be a -1/f0 in the front. It is, however, fulfilled in the forcing step
    QA = np.zeros(au.shape)
    QA, res, num_iter = qgpv_invert(iterations,threshold, 
                                    nz, ny, nx, p_levels, f0, ds,
                                    dp, msfm, SS,gamma,
                                    forcing, Tb, Tt, BC, BC_flag)



    time_block_end = time.time()


    aug = mpcalc.first_derivative(QA, delta=ds, axis=1) * msfm[None,:,:]
    avg = - mpcalc.first_derivative(QA, delta=ds, axis=2) * msfm[None,:,:]
    aug = aug.magnitude
    avg = avg.magnitude
    
    atg = mpcalc.first_derivative(QA, delta=-dp, axis=0)
    SS_stagger = np.append(SS, SS[-1])
    atg = atg * f0 * gamma[:,None,None] / SS_stagger[:,None,None]
    
    # sensitivity to geopotential height
    aphg = -mpcalc.first_derivative(atg/gamma[:,None,None], delta=-dp, axis=0)
    print((time_block_end - time_block_start)/60,'min')
    
    atg = atg.magnitude
    aphg = aphg.magnitude
    
    return QA, aug, avg, atg, aphg, forcing, res, num_iter

def invlap(iterations, threshold, ds, msfm, qgpv, IG,
           p_levels=None, dp=50e2, SS=None, Tb=None, Tt=None, BC_flag=0):
    """
    qgpv is the forcing term;
    IG is initial guess for the field to be inverted.
    """
    R = 287.04
    g = 9.8066
    nz,ny,nx = qgpv.shape
#BC_flag = 0 for Dirichlet and 1 for Neumann
    res = np.zeros_like(qgpv)
    phip = IG
    res_list = []

# define coefficients

    mfs = (ds*ds)/(msfm**2)

    AI_1 = 1.
    AI_2 = 1.
    AI_4 = 1.
    AI_5 = 1.
    AI_3v = -4.

#####
    time_block_start = time.time()
    print("start iteration")
    for num_iter in np.arange(iterations):
        res_itr = 0
        for k in np.arange(nz):
            for j in np.arange(1,ny-1):
                for i in np.arange(1,nx-1):
                    RES = AI_1*phip[k,j-1,i] + \
                    AI_2*phip[k,j,i-1] + \
                    (AI_3v)*phip[k,j,i] + \
                    AI_4*phip[k,j,i+1] + \
                    AI_5*phip[k,j+1,i] - qgpv[k,j,i]*mfs[j,i]
                
                    phip[k,j,i] = phip[k,j,i] - 1.8*RES/(AI_3v)
                    RES = abs(RES)
                    res[k,j,i]=RES
                    res_itr += RES
#             if(BC_flag==1):
#                 if(k==1): 
#                     phip[0,:] = phip[1,:] - (R * Tb[:]/(100*p_levels[0]))*dp
#                 if(k==nz-2): 
#                     phip[-1,:] = phip[-2,:] + (R * Tt[:]/(100.*p_levels[-1]))*dp
        res_list.append(res_itr)
        if((np.amax(res)<threshold) and (num_iter>0)): 
            print('stopping at iteration number: ', num_iter)
            print(f"time taken {(time.time()-time_block_start)/60} min")
            break
    print('exhausts the iterations with the largest residual: ', np.amax(res))
    print(f"time taken {(time.time()-time_block_start)/60} min")
    return phip,res_list,num_iter

def laplacian(s, dx, mf):
    # doesn't consider the boudnaries
    ret = np.zeros_like(s)
    ret[:,1:-1,1:-1] = (s[:,1:-1,2:] + s[:,1:-1,:-2] 
                        + s[:,2:,1:-1] + s[:,:-2,1:-1] 
                        - 4*s[:,1:-1,1:-1])
    ret = ret/(dx**2)*mf[None,:,:]**2
    return ret

def d2p(s, dp, stb, f0):
    dp = -dp 
    dsdp_c = (s[2:,:,:]-s[:-2,:,:])/dp/2
    dsdp_t = (s[-1,:,:]-s[-2,:,:])/dp
    dsdp_b = (s[1,:,:]-s[0,:,:])/dp
    dsdp = np.concatenate((dsdp_b[None,:,:], dsdp_c, dsdp_t[None,:,:]), axis=0)
    
    cs = f0/stb[:,None,None]*dsdp
    dcsdp_c = (cs[2:,:,:]-cs[:-2,:,:])/dp/2
    dcsdp_t = (cs[-1,:,:]-cs[-2,:,:])/dp
    dcsdp_b = (cs[1,:,:]-cs[0,:,:])/dp
    dcsdp = np.concatenate((dcsdp_b[None,:,:], dcsdp_c, dcsdp_t[None,:,:]), axis=0)
    return dcsdp

def ddp(s, dp, gamma):
    """dp=5000Pa, flip the sign inside the function"""
    dp = -dp
    sg = s/gamma[:,None,None]
    sg_c = (sg[2:,:,:]-sg[:-2,:,:])/dp/2
    sg_t = (sg[-1,:,:]-sg[-2,:,:])/dp
    sg_b = (sg[1,:,:]-sg[0,:,:])/dp
    dsgdp = np.concatenate((sg_b[None,:,:], sg_c, sg_t[None,:,:]), axis=0)
    return dsgdp

# ---------------------------------------------------------------------------------------------------------------------
# ------------------------------- start main ----------------------------

def function(d, timestring):
    # wrfout_fname = '/Users/morgan/march2020/em_adj/balanced/wrfout'+timestring  # circ bal
    wrfout_fname = '/Users/morgan/march2020/em_adj/normal/wrfout'+timestring     # circ trad
    wrfout = Dataset(wrfout_fname, 'r')  # Dataset is the class behavior to open the file

    # adjout_fname = '/Users/morgan/march2020/em_adj/balanced/adjout'+timestring  # circ bal
    adjout_fname = '/Users/morgan/march2020/em_adj/normal/adjout'+timestring     # circ trad

    adjout = Dataset(adjout_fname, 'r')
    ds_wrf = xr.open_dataset(wrfout_fname)
    ds_adj = xr.open_dataset(adjout_fname)

    itime = 0
    msfm = getvar(wrfout, "MAPFAC_M", timeidx=itime)   # Map scale factor on mass grid
    msfu = getvar(wrfout, "MAPFAC_U", timeidx=itime)    # Map scale factor on mass grid
    msfv = getvar(wrfout, "MAPFAC_V", timeidx=itime)    # Map scale factor on mass grid

    terrain = getvar(wrfout, 'ter', timeidx=itime)
    cart_proj = get_cartopy(msfm) # Get the cartopy mapping object
    lats, lons = latlon_coords(msfm)

    msfm = msfm.values
    cor = wrfout.variables['F'][0,:, :]     # Coriolis parameter
    f0 = np.average(cor)
    cor = cor*0.0 # need to calculate relative vorticity using wrf_python avo operator

    ds = 1./wrfout.variables['RDX'][0]     # grid spacing SAME IN ZONAL AND MERIDIONAL DIRECTIONS

    pb = 1000; pt = 50; 
    p_levels = np.linspace(pb,pt,20)
    dp_arr = p_levels[1:] - p_levels[:-1]
    p_half_levels = p_levels[:-1] + dp_arr/2
    dp = 50e2

    # define constants
    R = mpconst.dry_air_gas_constant.magnitude *1000
    cp = mpconst.Cp_d.magnitude
    gamma = (R/(p_levels*1.e2))*(p_levels/pb)**(R/cp)


    # find standard atmosphere for pressure levels
    ZS,phiS,TS,SS = sa.atmos_structure(p_levels)
    SS_stagger = np.append(SS, SS[-1])

    z = getvar(wrfout, "z", timeidx=itime)
    Z = vinterp(wrfout, z, 'pressure', p_levels, extrapolate=True, field_type='z', log_p=True, timeidx=itime,meta=False)

    # vinterp: valid entry ’pressure’, ‘pres’, ‘p’: pressure [hPa]
    au_staggered = ds_adj.A_U.values[0,:]
    au_ds = destagger(au_staggered, 2,meta=False)

    av_staggered = ds_adj.A_V.values[0,:]
    av_ds = destagger(av_staggered, 1,meta=False)

    at = ds_adj.A_T.values[0,:]

    AU = ma.filled(vinterp(wrfout, au_ds, 'pressure', p_levels, timeidx=itime,meta=False),fill_value=0.0)
    AV = ma.filled(vinterp(wrfout, av_ds, 'pressure', p_levels, timeidx=itime,meta=False),fill_value=0.0)
    AT = ma.filled(vinterp(wrfout, at, 'pressure', p_levels, timeidx=itime,meta=False),fill_value=0.0)

    terrain_mask = create_terrain_mask(terrain, ZS)
    nz, ny, nx = AU.shape  #nz=29

    # ---------------------------------------------------------------------------------------------------------------------
    ##### start inverting sensitivity to qgpv

    vort = np.zeros(au_ds.shape)     # vort is defined on the mass points
    vort = vorticity(au_ds, av_ds, msfm, ds, ds)
    #vor = to_np(1.e-5*avo(au_staggered, av_staggered, msfu, msfv, msfm, cor, ds, ds, meta=True))
    dtheta_hat_dp=np.zeros(AU.shape)

    for k in np.arange(1,nz-1):
        dtheta_hat_dp[k,:] = ((AT[k-1,:]/gamma[k-1])-(AT[k+1,:]/gamma[k+1]))/(2.*dp)
        
    curl_V_hat = ma.filled(vinterp(wrfout, vort, 'pressure', p_levels, timeidx=itime,meta=False),fill_value=0.0)
    forcing = (-1./f0)*(curl_V_hat-f0*dtheta_hat_dp)
    forcing_masked = forcing * terrain_mask 

    BC = np.zeros(AU.shape)
    QA, aug, avg, atg, aphg, forcing, res, num_iter = repetition(AU, AV, AT, BC, p_levels, f0, ds, msfm, SS, gamma, terrain_mask, threshold=20)  #50,60,60

    sens2q = QA.copy()
    savename = "./inverted_balance/sens2qgpv_"+timestring+"_20lev.npz"
    np.savez(savename, at=atg, au=aug, av=avg, aph=aphg, aq=QA, apsi=forcing)

    # ---------------------------------------------------------------------------------------------------------------------
    ##### start inverting sensitivity to ageostrophic components

    apsi = -vorticity(au_ds, av_ds, msfm, ds, ds)
    aphi = ddp(AT, dp, gamma) #5000 is dp

    apsi = ma.filled(vinterp(wrfout, apsi, 'pressure', p_levels, timeidx=itime,meta=False),fill_value=0.0)

    opapsi = d2p(apsi, dp, SS_stagger, f0)
    opaphi = laplacian(aphi, ds, msfm)

    forcing_ag = opapsi - opaphi
    forcing_ag_masked = forcing_ag * terrain_mask 

    BC = np.zeros(AU.shape)
    iterations=300
    threshold=1e-12

    res = np.zeros(AU.shape)
    BC_flag = 0
    Tb = np.zeros((ny, nx))
    Tt = np.zeros((ny, nx))

    print("start inversion",time.time())
    time_block_start = time.time()
        
    ## Inversion

    ### QA is sensitivity to ageostophic part
    ### following manuscript equation 5
    ### there should be a -1/f0 in the front. It is, however, fulfilled in the forcing step
    QA, res, num_iter = qgpv_invert(iterations,threshold, 
                                    nz, ny, nx, p_levels, f0, ds,
                                    dp, msfm, SS, gamma,
                                    forcing_ag, Tb, Tt, BC, BC_flag)

    time_block_end = time.time()
    print((time_block_end - time_block_start)/60,'min')

    sens2aq = QA.copy()
    # savename = "aq"+timestring+".npy"
    # np.save(savename, QA)


    daudx = mpcalc.first_derivative(AU, delta=ds, axis=2) * msfm[None,:,:]
    davdy = mpcalc.first_derivative(AV, delta=ds, axis=1) * msfm[None,:,:]
    achi = - (daudx + davdy)

    daudy = mpcalc.first_derivative(AU, delta=ds, axis=1) * msfm[None,:,:]
    davdx = mpcalc.first_derivative(AV, delta=ds, axis=2) * msfm[None,:,:]
    apsi = - (davdx - daudy)

    lap_sens2q = laplacian(sens2q, ds, msfm)
    apsi_m_lapaq = apsi - lap_sens2q

    BC_flag = 0
    Tb = np.zeros((ny, nx))
    Tt = np.zeros((ny, nx))
    IG = np.zeros(AU.shape)

    #### 
    dapsidy = mpcalc.first_derivative(apsi_m_lapaq, delta=ds, axis=1) * msfm[None,:,:]
    dachidx = mpcalc.first_derivative(achi, delta=ds, axis=2) * msfm[None,:,:]
    forcing_auag = dapsidy.magnitude - dachidx.magnitude 
        
    auag,res,num_iter = invlap(300, 1e-13, ds, msfm, forcing_auag, IG)

    #### 
    dapsidx = mpcalc.first_derivative(apsi_m_lapaq, delta=ds, axis=2) * msfm[None,:,:]
    dachidy = mpcalc.first_derivative(achi, delta=ds, axis=1) * msfm[None,:,:]
    forcing_avag = - (dachidy.magnitude + dapsidx.magnitude)

    IG = np.zeros(AU.shape)
    avag,res,num_iter = invlap(300, 1e-13, ds, msfm, forcing_avag, IG)

    savename = "./inverted_imbalance/ageo"+timestring+".npz"
    np.savez(savename, auag=auag, avag=avag, sens2aq=sens2aq)
    

def main():
    manager = Manager()
    results = manager.dict()
    
    pool = mp.Pool(8)

    jobs = []

    # stations = []

    # filelist = sorted(glob("/Users/morgan/march2020/em_adj/balanced/wrfout_d01_2020-03*"))
    filelist = sorted(glob("/Users/morgan/march2020/em_adj/normal/wrfout_d01_2020-03*"))
    print(filelist)
    # timelist = [file.replace("/Users/morgan/march2020/em_adj/balanced/wrfout","") for file in filelist]
    timelist = [file.replace("/Users/morgan/march2020/em_adj/normal/wrfout","") for file in filelist]
        
    for timestring in timelist[:-1]:
        print(timestring)
        job = pool.apply_async(function, (results, timestring))
        jobs.append(job)

    for job in jobs:
        job.get()

    pool.close()
    pool.join()

    print(results)

if __name__ == '__main__':
    main()