Adding large-scale blending step to regional_workflow.

scripts/da_blending_fv3.py

@@ -0,0 +1,173 @@
+import numpy as np
+from netCDF4 import Dataset
+import raymond
+import sys
+
+print("Starting blending code")
+
+Lx = float(sys.argv[1])  # BLENDING_LENGTHSCALE
+pi = np.pi
+nbdy = 40  # 20 on each side
+
+blend = str(sys.argv[5]) # TRUE:  Blend RRFS and GDAS EnKF
+                         # FALSE: Don't blend, activate cold2warm start only, and use either GDAS or RRFS
+if blend == "TRUE":
+    blend = True
+    print("Blending is activated")
+elif blend == "FALSE":
+    blend = False
+    print("Blending is **NOT** activated! Will perform cold2warm start conversion only.")
+else:
+    print("variable 'blend' not set correctly")
+    exit()
+
+if not blend:
+    use_host_EnKF = str(sys.argv[6])  # TRUE:  Final EnKF will be GDAS (no blending)
+                                      # FALSE: Final EnKF will be RRFS (no blending)
+    if use_host_EnKF == "TRUE":
+        use_host_EnKF = True
+        print("Using GDAS EnKF (no blending)")
+    elif use_host_EnKF == "FALSE":
+        use_host_EnKF = False
+        print("Using RRFS EnKF (no blending)")
+    else:
+        print("variable 'use_host_EnKF' not set correctly")
+        exit()
+
+
+# List of variables from the regional (fg) and global (bg) to blend respectively.
+vars_fg = ["u", "v", "T", "sphum", "delp"]
+vars_bg = ["u_cold2fv3", "v_cold2fv3", "t_cold2fv3", "sphum_cold2fv3", "delp_cold2fv3"]
+
+# GDAS EnKF file chgres_cube-ed from gaussian grid to ESG grid.
+# There is one more step to make sure the winds are on the same
+# grid staggering and have the same orientation as the RRFS winds.
+glb_fg = str(sys.argv[2])
+glb_fg_nc = Dataset(glb_fg)
+glb_nlon = glb_fg_nc.dimensions["lon"].size  # 1820   (lonp=1821)
+glb_nlat = glb_fg_nc.dimensions["lat"].size  # 1092   (latp=1093)
+glb_nlev = glb_fg_nc.dimensions["lev"].size  # 66     (levp=67)
+glb_Dx = 3.0
+
+# RRFS EnKF restart file fv_core.res.tile1 on ESG grid.
+reg_fg = str(sys.argv[3])
+# Open the blended file for updating the required vars (use a copy of the regional file)
+reg_fg_nc = Dataset(reg_fg, mode="a")
+nlon = reg_fg_nc.dimensions["xaxis_1"].size  # 1820   (xaxis_2=1821)
+nlat = reg_fg_nc.dimensions["yaxis_2"].size  # 1092   (yaxis_1=1093)
+nlev = reg_fg_nc.dimensions["zaxis_1"].size  # 65
+Dx = 3.0
+
+# RRFS EnKF restart file fv_tracer.res.tile1 on ESG grid.
+reg_fg_t = str(sys.argv[4])
+# Open the blended file for updating the required vars (use a copy of the regional file)
+reg_fg_t_nc = Dataset(reg_fg_t, mode="a")
+
+# Check matching grids
+# Note: global_hyblev_fcst_rrfsL65.txt has 0.000 0.0000000 as the 66th row, so
+# don't compare glb_nlev and nlev because glb_nlev will be 66 and nlev will be 65.
+# As a work around for now, we will just slice the top (or bottom?) 65 levels of
+# the global file and blend those with the regional file.
+if (glb_nlon != nlon or glb_nlat != nlat or glb_Dx != Dx):
+    print(f"glb_nlon:{glb_nlon} vs nlon:{nlon}")
+    print(f"glb_nlat:{glb_nlat} vs nlat:{nlat}")
+    print(f"glb_Dx:{glb_Dx}     vs Dx:{Dx}")
+    print("grids don't match")
+    exit()
+
+eps = (np.tan(pi*Dx/Lx))**-6  # 131319732.431162
+
+print(f"Input")
+print(f"  RRFS restart (core)           : {reg_fg}")
+print(f"  RRFS restart (tracer)         : {reg_fg_t}")
+print(f"  GDAS coldstart from chgres    : {glb_fg}")
+print(f"  Lx                            : {Lx}")
+print(f"  Dx                            : {Dx}")
+print(f"  NLON                          : {nlon}")
+print(f"  NLAT                          : {nlat}")
+print(f"  NLEV                          : {nlev}")
+print(f"  eps                           : {eps}")
+print(f"Output")
+print(f"  Blended background file       : {reg_fg}, {reg_fg_t}")
+
+# Step 1. blend.
+for (var_fg, var_bg) in zip(vars_fg, vars_bg):
+    i = vars_fg.index(var_fg)
+
+    if var_fg == "sphum":
+        reg_nc = reg_fg_t_nc
+    else:
+        reg_nc = reg_fg_nc
+    glb_nc = glb_fg_nc
+
+    dim = len(np.shape(reg_nc[var_fg]))-1
+    if dim == 2:  # 2D vars
+        glb = np.float64(glb_nc[var_bg][:, :])     # (1093 1820)
+        reg = np.float64(reg_nc[var_fg][:, :, :])  # (1, 1093, 1820)
+        ntim = np.shape(reg)[0]
+        nlat = np.shape(reg)[1]
+        nlon = np.shape(reg)[2]
+        nlev = 1
+        glb = np.reshape(glb, [ntim, nlat, nlon])  # add time dim bc missing from chgres
+        var_out = np.zeros(shape=(nlon, nlat, 1), dtype=np.float64)
+        field = np.zeros(shape=(nlon*nlat), dtype=np.float64)
+        var_work = np.zeros(shape=((nlon+nbdy), (nlat+nbdy), 1), dtype=np.float64)
+        field_work = np.zeros(shape=((nlon+nbdy)*(nlat+nbdy)), dtype=np.float64)
+    if dim == 3:  # 3D vars
+        glb = np.float64(glb_nc[var_bg][:, :, :])
+        reg = np.float64(reg_nc[var_fg][:, :, :, :])  # (1, 65, 1093, 1820)
+        ntim = np.shape(reg)[0]
+        nlev = np.shape(reg)[1]
+        nlat = np.shape(reg)[2]
+        nlon = np.shape(reg)[3]
+        glb = np.reshape(glb, [ntim, nlev, nlat, nlon])  # add time dim bc missing from chgres
+        var_out = np.zeros(shape=(nlon, nlat, nlev, 1), dtype=np.float64)
+        field = np.zeros(shape=(nlon*nlat, nlev), dtype=np.float64)
+        var_work = np.zeros(shape=((nlon+nbdy), (nlat+nbdy), nlev, 1), dtype=np.float64)
+        field_work = np.zeros(shape=((nlon+nbdy)*(nlat+nbdy), nlev), dtype=np.float64)
+    glbT = np.transpose(glb)        # (1820, 1093, 65)
+    regT = np.transpose(reg)        # (1820, 1093, 65, 1)
+
+    nlon_start = int(nbdy/2)
+    nlon_end = int(nlon+nbdy/2)
+    nlat_start = int(nbdy/2)
+    nlat_end = int(nlat+nbdy/2)
+
+    if blend:
+        print(f"Blending backgrounds for {var_fg}/{var_bg}")
+        var_work[nlon_start:nlon_end, nlat_start:nlat_end, :] = glbT - regT
+        field_work = var_work.reshape((nlon+nbdy)*(nlat+nbdy), nlev, order="F")  # order="F" (FORTRAN)
+        field_work = raymond.raymond(field_work, nlon+nbdy, nlat+nbdy, eps, nlev)
+        var_work = field_work.reshape(nlon+nbdy, nlat+nbdy, nlev, order="F")
+        var_out = var_work[nlon_start:nlon_end, nlat_start:nlat_end, :]
+        if dim == 2:  # 2D vars
+            var_out = var_out[:, :, 0] + regT[:, :, 0]
+            var_out = np.reshape(var_out, [nlon, nlat, 1])  # add the time ("1") dimension back
+        if dim == 3:  # 3D vars
+            var_out = var_out + regT[:, :, :, 0]
+            var_out = np.reshape(var_out, [nlon, nlat, nlev, 1])  # add the time ("1") dimension back
+    else:  # skip blending and use either host enkf (GDAS) or the RRFS enkf
+        print(f"Blending code is NOT executing blending!")
+        print(f"    This is used for finishing converting the cold start files into warm start format.")
+        if use_host_EnKF:
+            print(f"---> Use the GDAS EnKF")
+            var_out = glbT
+        else:
+            print(f"---> Use the RRFS EnKF")
+            var_out = regT
+
+    var_out = np.transpose(var_out)  # (1, 50, 834, 954)
+
+    # Overwrite blended fields to blended file.
+    if dim == 2:  # 2D vars
+        reg_nc.variables[var_fg][:, :, :] = var_out
+    if dim == 3:  # 3D vars
+        reg_nc.variables[var_fg][:, :, :, :] = var_out
+
+# Close nc files
+reg_nc.close()  # blended file
+glb_nc.close()
+
+print("Blending finished successfully.")
+
+exit(0)

scripts/da_chgres_cold2fv3.py

@@ -0,0 +1,193 @@
+import numpy as np
+from netCDF4 import Dataset
+import remap_dwinds
+import remap_scalar
+import chgres_winds   # might need to rename in the future
+import sys
+
+print("Reading in NETCDF4 Files... ", end="\r")
+warm = str(sys.argv[1])
+cold = str(sys.argv[2])
+grid = str(sys.argv[3])
+akbk = str(sys.argv[4])
+akbkcold = str(sys.argv[5])
+orog = str(sys.argv[6])
+
+warmnc = Dataset(warm)
+coldnc = Dataset(cold, mode="a")
+akbknc = Dataset(akbk)
+gridnc = Dataset(grid)
+akbkcoldnc = Dataset(akbkcold)
+orognc = Dataset(orog)
+print("Reading in NETCDF4 Files... Done.")
+
+ColdStartWinds = True
+VertRemapScalar = True
+VertRemapWinds = True
+WriteData = True
+
+# STEP 1. ROTATE THE WINDS FROM CHGRES
+if ColdStartWinds:
+    print("Starting ColdStartWinds.... ", end="\r")
+    # Data from warm restarts
+    u = np.float64(warmnc["u"][0, :, :, :])
+    v = np.float64(warmnc["v"][0, :, :, :])
+    nlev = np.shape(u)[0]  # 127,z
+    nlat = np.shape(u)[1]  # 769,y
+    nlon = np.shape(u)[2]  # 768,x
+
+    km = nlev
+    nlev = coldnc.createDimension("nlev", nlev)  # 127
+
+    # Data from cold chgres
+    u_s = np.float64(coldnc["u_s"][:, :, :])   # (128, 769, 768)
+    v_s = np.float64(coldnc["v_s"][:, :, :])   # (128, 769, 768)
+    u_w = np.float64(coldnc["u_w"][:, :, :])   # (128, 768, 769)
+    v_w = np.float64(coldnc["v_w"][:, :, :])   # (128, 768, 769)
+
+    # grid data - usually has 2x as many grid points, so we need every other value.
+    gridx = np.float64(gridnc["x"][0:-1:2, 0:-1:2])
+    gridy = np.float64(gridnc["y"][0:-1:2, 0:-1:2])
+
+    # Fortran wants everything transposed and in fortran array type
+    gridx = np.asfortranarray(gridx.transpose())
+    gridy = np.asfortranarray(gridy.transpose())
+    u_s = np.asfortranarray(u_s.transpose())
+    v_s = np.asfortranarray(v_s.transpose())
+    u_w = np.asfortranarray(u_w.transpose())
+    v_w = np.asfortranarray(v_w.transpose())
+
+    # Initialize some computed fields to zero
+    ud = np.float64(0.0*u_s)  # initialize to zero
+    vd = np.float64(0.0*u_w)  # initialize to zero
+
+    # rotate winds to model d-grid (~30s; nodes=2; cpus=128)
+    chgres_winds.main(gridx, gridy, u_s, v_s, u_w, v_w, ud, vd)
+
+    print("Starting ColdStartWinds.... Done.")
+
+# STEP 2. VERTICAL REMAPPING OF SCALARS
+if VertRemapScalar:
+    print("Starting VertRemapScalar... ", end="\r")
+
+    # Data from cold restarts
+    ak0 = np.float64(akbkcoldnc["vcoord"][0, :])    # ( lev,         ) == (128,         )
+    bk0 = np.float64(akbkcoldnc["vcoord"][1, :])    # ( lev,         ) == (128,         )
+    ak = ak0[1:]
+    bk = bk0[1:]
+    ps = np.float64(coldnc["ps"][:, :])            # (      lat, lon) == (     768, 768)
+    zh = np.float64(coldnc["zh"][:, :, :])         # (levp, lat, lon) == (129, 768, 768)
+    omga = np.float64(coldnc["w"][:, :, :])          # ( lev, lat, lon) == (128, 768, 768)
+    delp_cold = np.float64(coldnc["delp"][:, :, :])  # (127, 768, 768)
+    t_cold = np.float64(coldnc["t"][:, :, :])        # (128, 768, 768)
+    Atm_phis = np.float64(orognc["orog_filt"][:, :])*9.80665
+
+    ak0[0] = 1.000000000000000E-009
+    bk0[0] = 1.000000000000000E-009
+
+    sphum = np.float64(coldnc["sphum"][:, :, :])  # ( lev, lat, lon) == (128, 768, 768)
+    liq_wat = np.float64(coldnc["liq_wat"][:, :, :])
+    o3mr = np.float64(coldnc["o3mr"][:, :, :])
+    ice_wat = np.float64(coldnc["ice_wat"][:, :, :])
+    rainwat = np.float64(coldnc["rainwat"][:, :, :])
+    snowwat = np.float64(coldnc["snowwat"][:, :, :])
+    graupel = np.float64(coldnc["graupel"][:, :, :])
+    ntracers = 7
+    qa = np.array([sphum, liq_wat, o3mr, ice_wat, rainwat, snowwat, graupel])
+
+    # Fortran wants everything transposed and in fortran array type
+    ak = np.asfortranarray(ak)  # Don't transpose 1D array
+    bk = np.asfortranarray(bk)  # Don't transpose 1D array
+    ak0 = np.asfortranarray(ak0)
+    bk0 = np.asfortranarray(bk0)
+    Atm_phis = np.asfortranarray(np.transpose(Atm_phis))
+    ps = np.asfortranarray(ps.transpose())
+    zh = np.asfortranarray(zh.transpose())
+    omga = np.asfortranarray(omga.transpose())
+    qa = np.asfortranarray(qa.transpose())
+    delp_cold = np.asfortranarray(delp_cold.transpose())
+    t_cold = np.asfortranarray(t_cold.transpose())
+
+    isrt = 1
+    jsrt = 1
+    iend = np.shape(t_cold)[0]
+    jend = np.shape(t_cold)[1]
+    npz = np.shape(t_cold)[2]-1
+    levp = npz + 1  # (km)
+
+    # Initialize some computed fields to zero
+    Atm_delp = 1.0*delp_cold[:, :, 1:]  # initialize to zero (delp for sfcp)
+    Atm_q = 1.0*qa[:, :, 1:, :]         # initialize to zero (tracers... sphum=1)
+    Atm_pt = 1.0*t_cold[:, :, 1:]      # initialize to zero (temperature)
+    Atm_ps = 1.0*ps[:, :]               # initialize to zero (need for remap_dwinds)
+
+    remap_scalar.main(levp, npz, ntracers, ak0, bk0, ak, bk, ps, qa, zh, omga, t_cold,
+                      isrt, iend, jsrt, jend, Atm_pt, Atm_q, Atm_delp, Atm_phis, Atm_ps)
+
+    print("Starting VertRemapScalar... Done.")
+
+
+# STEP 3. VERTICAL REMAPPING OF WINDS
+if VertRemapWinds:
+    print("Starting VertRemapWinds....", end="\r")
+
+    Atm_u = 1.0*ud[:, :, 1:]  # Atm_u has levs=npz, ud has levs km
+    Atm_v = 1.0*vd[:, :, 1:]
+
+    # vertically remap the dwinds
+    remap_dwinds.main(levp, npz, ak0, bk0, ak, bk, ps, ud, vd,
+                      isrt, iend, jsrt, jend, Atm_u, Atm_v, Atm_ps)
+
+    print("Starting VertRemapWinds.... Done.")
+
+else:
+    Atm_u = ud[:, :, 1:]
+    Atm_v = vd[:, :, 1:]
+
+
+# STEP 4. WRITE OUT DATA
+if WriteData:
+    # tranpose ud, vd back to original shape, cutoff one of levels (there is an extra level),
+    # add a new variable to the nc file by duplicating the corresponding u/v variable and
+    # redefining the shape of the array, finally, assign ud/vd into the u/v variable in nc file.
+
+    if ColdStartWinds:
+        # For ud
+        new_var = "u_cold2fv3"
+        ud = np.transpose(Atm_u)
+        var_to_duplicate = coldnc.variables["u_s"]
+        coldnc.createVariable(new_var, var_to_duplicate.datatype, ('nlev', 'latp', 'lon'))
+        coldnc.variables[new_var][:, :, :] = ud
+
+        # For vd
+        new_var = "v_cold2fv3"
+        vd = np.transpose(Atm_v)
+        var_to_duplicate = coldnc.variables["v_w"]
+        coldnc.createVariable(new_var, var_to_duplicate.datatype, ('nlev', 'lat', 'lonp'))
+        coldnc.variables[new_var][:, :, :] = vd
+
+    if VertRemapScalar:
+        # For Temperature
+        new_var = "t_cold2fv3"
+        t_cold = np.transpose(Atm_pt)
+        var_to_duplicate = coldnc.variables["t"]
+        coldnc.createVariable(new_var, var_to_duplicate.datatype, ('nlev', 'lat', 'lon'))
+        coldnc.variables[new_var][:, :, :] = t_cold
+
+        # For delp
+        new_var = "delp_cold2fv3"
+        delp = Atm_delp.T
+        var_to_duplicate = coldnc.variables["delp"]
+        coldnc.createVariable(new_var, var_to_duplicate.datatype, ('nlev', 'lat', 'lon'))
+        coldnc.variables[new_var][:, :, :] = delp
+
+        # For sphum
+        new_var = "sphum_cold2fv3"
+        sphum = Atm_q[:, :, :, 0].T
+        var_to_duplicate = coldnc.variables["sphum"]
+        coldnc.createVariable(new_var, var_to_duplicate.datatype, ('nlev', 'lat', 'lon'))
+        coldnc.variables[new_var][:, :, :] = sphum
+
+# close the nc files
+warmnc.close()
+coldnc.close()