Source code for pycif.plugins.domains.wrfchem.read_domain

import numpy as np
import netCDF4 as nc
import os




def read_grid(domain, **kwargs):
    """Reads WRF domain definition from files

    Fills domain object with values read from WRF files (that must have
    been created prior to CIF run)

    Args:
        domain (domain object):
        Object defining the domain. For contents, see input_arguments in
        __init__.py

    Return:
        domain object with meshgrids for center lon/lat and corner lon/lat

    Notes:
        - Only able to work with one domain for now, no nesting.
    """

    if domain.max_dom > 1:
        raise NotImplementedError("Can't work with nested domains yet.")

    # For now, work on domain 1 (domain numbering in WRF starts at 1)
    ndom = 1

    # Get horizontal coordinates (grid cell centers and grid cell
    # corners)
    file_path = os.path.join(domain.dir_geo_em, "geo_em.d{:02d}.nc".format(ndom))
    ncf = nc.Dataset(file_path, "r")
    # Keep 2D-shaped variables in domain object, they are used elsewhere
    ncf.variables["XLONG_M"].set_auto_mask(False)
    domain.zlon2d = ncf.variables["XLONG_M"][0, :, :]
    ncf.variables["XLAT_M"].set_auto_mask(False)
    domain.zlat2d = ncf.variables["XLAT_M"][0, :, :]
    ncf.variables["XLONG_C"].set_auto_mask(False)
    domain.zlonc2d = ncf.variables["XLONG_C"][0, :, :]
    ncf.variables["XLAT_C"].set_auto_mask(False)
    domain.zlatc2d = ncf.variables["XLAT_C"][0, :, :]
    ncf.close()

    # Flatten centers: shape is (N, 1)
    domain.zlon = domain.zlon2d.flatten()[np.newaxis, :]
    domain.zlat = domain.zlat2d.flatten()[np.newaxis, :]

    # Flatten corners: shape is (4, N)
    domain.zlonc = np.ndarray((4, domain.zlon.size), dtype=domain.zlon.dtype)
    domain.zlatc = np.ndarray((4, domain.zlon.size), dtype=domain.zlon.dtype)
    
    domain.zlonc[0, :] = domain.zlonc2d[:-1, :-1].flatten()
    domain.zlonc[1, :] = domain.zlonc2d[:-1, 1:].flatten()
    domain.zlonc[2, :] = domain.zlonc2d[1:, :-1].flatten()
    domain.zlonc[3, :] = domain.zlonc2d[1:, 1:].flatten()
    
    domain.zlatc[0, :] = domain.zlatc2d[:-1, :-1].flatten()
    domain.zlatc[1, :] = domain.zlatc2d[:-1, 1:].flatten()
    domain.zlatc[2, :] = domain.zlatc2d[1:, :-1].flatten()
    domain.zlatc[3, :] = domain.zlatc2d[1:, 1:].flatten()

    # Saving ancillary information to domain attributes
    domain.nlat = domain.zlon.shape[0]
    domain.nlon = domain.zlon.shape[1]

    # Get vertical coordinates
    # From registry.hyb_coord:
    #       new dry pressure Pd(i,k,j)  = C3(k) * ( Pds(i,j) - Pt      ) + C4(k) + Pt

    # -> Translated to sigma coordinates p = a + b*psurf
    # a = C4 + p_top*(1-C3)
    # b = C3

    file_path = os.path.join(domain.dir_wrfinput, "wrfinput_d{:02d}".format(ndom))
    ncf = nc.Dataset(file_path, "r")
    domain.nlev = len(ncf.dimensions["bottom_top"])
    c3f = np.array(ncf["C3F"][:].squeeze())
    c4f = np.array(ncf["C4F"][:].squeeze())
    p_top = ncf["P_TOP"][:].data
    domain.sigma_a = c4f + p_top*(1-c3f)
    domain.sigma_b = c3f
    domain.pressure_unit = "Pa"
    ncf.close()