CMAQ to WRF

cmaq2wrf.py

__doc__ = """
# Make CMAQ WRF-Like

     author: Barron H. Henderson
     date: 2020-07-30
     updated: 2024-06-04

# Overview

* Need: Many tools exist for WRFChem that we would like to use with CMAQ.
* Question: How make a CMAQ file WRF-Like?
* Answer: This script peforms the following steps:
    1. Copy important variables from WRF (e.g., map scale, time, etc)
    2. Window the WRF file (normally remove 6 cells from each side)
    3. Create a place to keep CMAQ concentrations.
    4. Populate the variables.

# Prerequisites

* Python3
* PseudoNetCDF (`python -m pip install pseudonetcdf`)

# How-to:

Assuming that:
1. ./wrfout_d01 is the path to a wrfout file
2. ./CCTM_CONC is the path to a CMAQ conc file on the same grid.

The command line below will create wrf_cmaq.nc.

$ python %(prog)s --btrim=6 ./wrfout_d01 ./CCTM_CONC  wrf_cmaq.nc 

This can also be done in a python interpretter as follows:

from camq2wrf import cmaq2wrf
cmaq2wrf('wrfout_d01', 'CCTM_CONC', 'wrf_cmaq.nc')

"""

import numpy as np
import PseudoNetCDF as pnc
import warnings


def cmaq2wrf(wrfoutpath, concpath, outpath=None, btrim=6, cmaq_variables=None):
    """
    Arguments
    ---------
    wrfoutpath : str
        wrfout path with same vertical and horizontal structure as CMAQ.
    concpath : str
        CONC (3d) file with variables to copy with WRF
    outpath : str or None
        If not None, save as outpath
    btrim : int
        boundary triming from WRF to CMAQ
    cmaq_variables : list or None
        If None, copy all variables from concpath. Otherwise, just
        those in cmaq_variables

    Returns
    -------
    outf : netcdf-like
        File with WRF meta-data and structure, but CMAQ contents
    """
    # Open input files
    if isinstance(wrfoutpath, str):
        wrff = pnc.pncopen(wrfoutpath, format='netcdf').copy()
    else:
        wrff = wrfoutpath
    if isinstance(concpath, str):
        concf = pnc.pncopen(concpath, format='ioapi').copy()
    else:
        concf = concpath

    # Set trim edge setting for WRF to match CMAQ
    trim_south = trim_north = trim_east = trim_west = btrim

    # Check that btrim results in aligned grids.
    chkrow = concf.NROWS + trim_south + trim_north
    wrfrow = len(wrff.dimensions['south_north'])
    chkcol = concf.NCOLS + trim_west + trim_east
    wrfcol = len(wrff.dimensions['west_east'])

    if chkrow !=  wrfrow:
        print(
            f'btrim does not result in aligned rows (CMAQ={concf.NROWS},'
            f'WRF={wrfrow}, WRF-2xbtrim={chkrow})'
        )
    if chkcol != wrfcol:
        print(
            f'btrim does not result in aligned cols (CMAQ={concf.NCOLS},'
            f'WRF={wrfcol}, WRF-2xbtrim={chkcol})'
        )

    if chkrow !=  wrfrow or chkcol != wrfcol:
        raise ValueError('Aborting for incorrect btrim')

    keepwrfkeys = [
        'Times', 'XLAT', 'XLONG', 'MAPFAC_M', 'MAPFAC_V', 'MAPFAC_U', 'MAPFAC_MX',
        'MAPFAC_MY', 'ZNU', 'P', 'PH', 'PHB', 'PB', 'T', 'PSFC', 'U', 'V',
        'QVAPOR'
    ]

    keepwrfkeys = [k for k in keepwrfkeys if k in wrff.variables]
    exkey = 'P'
    addexpr = ""
    cmaqmapping = {}
    wrffactor = {}
    for key in list(concf.variables):
        outkey = key.lower()
        cmaqunit = concf.variables[key].units.strip().lower()
        if key != 'TFLAG' and (cmaq_variables is None or key in cmaq_variables):
            cmaqmapping[outkey] = key
            if cmaqunit in ('ppm', 'ppmv'):
                # 1e-6 for ppmv to ppv
                wrffactor[key] = 1e-6
                wrfunit = 'ppv'
            else:
                wrffactor[key] = 1
                wrfunit = cmaqunit
            addexpr += f"""{outkey} = {exkey} * 0
{outkey}.units = "{wrfunit}"
{outkey}.description = "{key} Concentration original unit {cmaqunit}"
"""

    # Keep only a subset of variables
    # Must include P, which is the default template variable.
    swrff = wrff.subset(keepwrfkeys)

    # Trim boundaries off of wrfout file to match CMAQ domain
    outf = swrff.slice(
        west_east=slice(trim_west, -trim_east),
        south_north=slice(trim_south, -trim_north),
    )

    # Define spatial meta-keys and how to adjust them
    adjustkeys = {
        'WEST-EAST_PATCH_END_UNSTAG': trim_west + trim_east,
        'WEST-EAST_PATCH_END_STAG': trim_west + trim_east,
        'SOUTH-NORTH_PATCH_END_UNSTAG': trim_south + trim_north,
        'SOUTH-NORTH_PATCH_END_STAG': trim_south + trim_north
    }

    # Adjust spatial properties
    for propk, adjval in adjustkeys.items():
        n = getattr(outf, propk)
        setattr(outf, propk, n - adjval)

    # Add an CMAQ concentration variables
    outf.eval(addexpr, inplace=True)

    # Populate CMAQ variables
    qnt = len(concf.dimensions['TSTEP'])
    wnt = len(wrff.dimensions['Time'])
    nt = min(qnt, wnt)
    for wrfk, cmaqk in cmaqmapping.items():
        print(' - Populating', wrfk, 'times:', end='')
        cmaqv = concf.variables[cmaqk]
        wrfv = outf.variables[wrfk]
        factor = wrffactor[cmaqk]
        # Iterate over times
        for ti in range(nt):
            print(ti, end=',', flush=True)
            wrfv[ti] = cmaqv[ti] * factor
        del cmaqv, wrfv
        print()

    # Overwrite time from CMAQ file
    times = concf.getTimes()
    for ti in range(nt):
        outf.variables['Times'][ti, :] = np.array(
            times[ti].strftime('%Y-%m-%d_%H:%M:%S'), dtype='c'
        )
    if outpath is not None:
        outf = outf.save(outpath, verbose=0)
    return outf


if __name__ == '__main__':
    import argparse

    # Set paths to input and output files
    prsr = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawTextHelpFormatter
    )
    prsr.add_argument(
        '--btrim', default=6, type=int, help='Boundary cells to trim (e.g, 6)'
    )
    prsr.add_argument(
        '--cmaq-variables', default=None, type=lambda x: x.split(','),
        help='Subset of CMAQ variables to process (default all)'
    )
    prsr.add_argument('wrfoutpath', help='wrfout path to use for meta-data')
    prsr.add_argument('concpath', help='CMAQ concentration path for data')
    prsr.add_argument('outpath', help='Path for WRF-like output with CMAQ')
    args = prsr.parse_args()
    kwds = vars(args)
    cmaq2wrf(**kwds)