conc2photdiag.py

conc2photdiag.py

__doc__ = """
INTEGRATE CMAQ Columns
======================

---
    author: Barron H. Henderson
    last-updated: 2026-07-17
---

Integrate columns from CMAQ using GRIDCRO2D, METCRO3D, and CONC files.
* GRIDCRO2D and METCRO3D are MCIP outputs and CMAQ inputs.
* CONC is an output from CMAQ that has 3D gas-phase concentrations.

More details are available in the integrate, atoverpass, and tocolumns
function doc strings.

Prerequisites:
python3 with pandas, numpy, and xarray

"""
import pandas as pd
import numpy as np
import xarray as xr


spcs = ['NO2', 'CO', 'O3', 'FORM']
colkeys = {
    'NO2': 'NO2_COLUMN', 'FORM': 'HCHO_COLUMN', 'CO': 'CO_COLUMN',
    'O3': 'TROPO_O3_COLUMN'
}
MWAIR = 0.0289628  # kg/mole
N_A = 6.02214076e+23  # molec/mole


def conc2photdiag(xf, mf, verbose=0):
    """
    Use mixing ratios (x) from xf, DENS from mf, and ZF from mf to integrate
    vertical columns from CMAQ results.

    Arguments
    ---------
    xf : xarray.Dataset
        Must contain NO2, CO, O3, and FORM variables and expects dimensions
        (TSTEP, LAY, ROW, COL)
    mf : xarray.Dataset
        Must contain DENS, ZF variables and expects dimensions
        (TSTEP, LAY, ROW, COL)

    Returns
    -------
    cf : xarray.Dataset
        Contains NO2_COLUMN, CO_COLUMN, O3_COLUMN, and HCHO_COLUMN variables
        mirroring the naming in PHOTDIAG1 has same dimensions as the inputs.

    Notes
    -----
    ```pseudocode
    COLUMN_i = MWAIR**-1 * N_A * 1e-4 sum_z(DENS_{z} * DZ_{z} * VMR_{i,z})
    [molec_i/cm2] = [mol/kg] * [molec/mole] * [m2 / cm2] * [kg/m3] * [m] * [mol_i / mol]
    * N_A = 6.02214076e+23 molec/mol2
    * MWAIR = 0.0289628 kg/mole
    * DZ = ZF_{z} - ZF_{z-1} m
    * VMR = species mixing ratio (ppm) divided by a million
    * DENS in METCRO3D
    * ZF in METCRO3D
    ```
    """
    dz = mf['ZF'][:].copy() * 1
    dz[:, 1:] -= mf['ZF'][:, :-1]
    molec_per_cm2 = mf['DENS'][:] * dz * 1e-4 / MWAIR * N_A
    # molec / cm2 [=]   kg / m3   * m * m2/cm2 * mol / kg * molec/mol
    spcn = len(spcs)
    s1 = slice(None, 1)  # first layer or time
    sv = slice(None, spcn)  # first n species
    # create a placeholder for results
    outf = xf[['TFLAG'] + spcs].isel(LAY=s1, VAR=sv).copy()
    outf = outf.rename_vars(colkeys)
    outf.attrs['NVARS'] = spcn  # output has n
    outf.attrs['NLAYS'] = 1     # output has 1 layer
    # VAR-LIST is concatenation of 16c names
    outf.attrs['VAR-LIST'] = ''.join([colkeys[spc].ljust(16) for spc in spcs])
    # VGLVLS = 1 at surface and 0 at model top
    outf.attrs['VGLVLS'] = np.array([1, 0], dtype='f')
    for spc in spcs:
        colkey = colkeys[spc]
        if verbose > 0:
            print(f'{spc:3s} avg 00-24utc=', end='')
        vmr_i = xf[spc][:] * 1e-6  # ppmv to mol/mol
        column_i = (vmr_i * molec_per_cm2).sum('LAY', keepdims=True)
        if verbose > 0:
            print(f'{column_i.mean():.6e} molecules/cm2')
        outv = outf[colkey]  # store result and update units
        if spc == 'O3':
            column_i /= 2.6879e16
            unit = 'DU'
        else:
            column_i *= 1e-15
            unit = 'petamolec cm-2'

        outv[:] = column_i
        outv.attrs['units'] = unit.ljust(16)[:16]
        vdesc = f'Instantaneous integrated {spc} column'.ljust(80)[:80]
        outv.attrs['var_desc'] = vdesc
        if verbose > 0:
            print(f'{spc:3s} avg hourly={outv[:].mean():.6e} molecules/cm2')
    return outf


def atoverpass(cf, gf, overpassh, verbose=0):
    """
    Use column integrals (c) from cf and LON from gf to create overpass average
    results. Expects cf to have 25-instantaneous hours.

    Arguments
    ---------
    cf : xarray.Dataset
        Contains NO2_COLUMN, CO_COLUMN, O3_COLUMN, and HCHO_COLUMN variables
        mirroring the naming in PHOTDIAG1 has same dimensions as the inputs.
    gf : xarray.Dataset
        Must contain LON variables and expects dimensions
        (TSTEP, LAY, ROW, COL) even though TSTEP and LAY should be unity
        dimensions

    Returns
    -------
    caf : xarray.Dataset
        Contains NO2_COLUMN, CO_COLUMN, O3_COLUMN, and HCHO_COLUMN variables
        mirroring the naming in PHOTDIAG1 has same dimensions as the inputs.

    Notes
    -----
    ```pseudocode
    ovph = 13.5
    LST = LON / 15 + UTC
    isoverpass = |LST % 24 - ovph|
    ```
    """
    utc_hour = np.arange(25)  # utc is cell independent
    lst_hour = (gf['LON'][:] / 15. + utc_hour[:, None, None, None])
    # within 1hour
    isoverpass = np.abs((lst_hour % 24) - overpassh) <= 1
    ovptag = f'{overpassh-0.5:02.0f}-{overpassh+0.5:02.0f}LST'
    # assuming CONUS, so negative hours are yesterday
    # and should not be included. A HEMI domain would require
    # better logic to separate days
    isoverpass = isoverpass & (lst_hour >= 0)
    if verbose > 1:
        hidx = np.where(isoverpass)[0]
        print('UTC H (avg)', utc_hour[hidx].mean())
        print('UTC H (std)', utc_hour[hidx].std())
        print('UTC H (min)', utc_hour[hidx].min())
        print('UTC H (max)', utc_hour[hidx].max())
    if verbose > 2:
        lst_houra = lst_hour.to_numpy()
        print('LST H (avg)', lst_houra[isoverpass].mean())
        print('LST H (std)', lst_houra[isoverpass].std())
        print('LST H (min)', lst_houra[isoverpass].min())
        print('LST H (max)', lst_houra[isoverpass].max())

    # check that you got two hours (i.e., this is a CONUS domain
    isoverpassn = isoverpass.sum('TSTEP')
    has2h = np.allclose(isoverpassn, 2)
    if not has2h:
        print(isoverpassn.mean(), isoverpassn.std())
    assert has2h
    outf = cf.isel(TSTEP=slice(None, 1))
    for spc in spcs:
        colkey = colkeys[spc]
        if verbose > 0:
            print(f'{spc:3s} avg 00-24utc=', end='')
        column_i = cf[colkey]
        column_avg_ovph_i = np.ma.masked_where(
            isoverpass is False, column_i
        ).mean(0, keepdims=True)
        outv = outf[colkey]  # store result and update units
        vdesc = f'{ovptag} average integrated {spc} in-CMAQ column'
        outv[:] = column_avg_ovph_i
        outv.attrs['var_desc'] = vdesc.ljust(80)[:80]
        if verbose > 0:
            print(f'{spc:3s} avg hourly={outv[:].mean():.6e} {unit}')
    return outf


def tocolumns(
    g2path, m3path, concpath, colpath, ovppath=None, overpassh=13.5, verbose=0
):
    """
    Thin wrapper around integrate and atoverpass with files opened from g2path,
    m3path, and concpath. Outputs integrate results averaged to 24h to colpath
    and overpass average to ovppath.

    Arguments
    ---------
    g2path : str
        Path to CMAQ input GRIDCRO2D file
    m3path : str
        Path to CMAQ input METCRO3D file
    concpath : str
        Path to CMAQ  output CONC file
    colpath : str
        Path for 24h column integral result.
    ovppath : str
        Path for overpass averaged column integral result.
    overpassh : float
        Local standard hour for the overpass time.
    verbose : int
        Level of verbosity.

    Results
    -------
    None
    """
    # open paths
    xf = xr.open_dataset(concpath, mode='rs', decode_cf=False)
    mf = xr.open_dataset(m3path, mode='rs', decode_cf=False)
    gf = xr.open_dataset(g2path, mode='rs', decode_cf=False)
    # integrate spcs
    cf = conc2photdiag(xf, mf, verbose=verbose)
    now = pd.to_datetime('now', utc=True)
    fdesc0 = 'Column integrals of species where: X_COLUMN ='
    fdesc0 += ' 1e-4 / MWAIR * N_A * sum_z(VMR_z * DENS_z * DZ_z * VMR_z)'
    fdesc0 += f' MWAIR={MWAIR:.7f}, N_A={N_A:.8e}. ZH and DENS come from'
    fdesc0 += f' {m3path} and mixing ratios (VMR) from {concpath}.'
    if ovppath is not None:
        ovpf = atoverpass(cf, gf, overpassh, verbose=verbose)
        fdesc = fdesc0 + f' Column integrals averaged around {overpassh}LST'
        fdesc += f' using ZH and DENS from {m3path},'
        fdesc += f' LST = UTC + LON / 15 where LON from {g2path}.'
        ovpf.attrs['CDATE'] = int(now.strftime('%Y%j'))
        ovpf.attrs['CTIME'] = int(now.strftime('%H%M%S'))
        ovpf.attrs['WDATE'] = int(now.strftime('%Y%j'))
        ovpf.attrs['WTIME'] = int(now.strftime('%H%M%S'))
        ovpf.attrs['UPNAM'] = 'conc2photdiag.py'.ljust(16)[:16]
        ovpf.attrs['FILEDESC'] = fdesc.ljust(4800)[:4800]
        # save to disk
        ovpf.to_netcdf(ovppath, format='NETCDF3_CLASSIC')
    fdesc = fdesc0 + ' Interpolated from 25h instantaneous to hourly averages.'
    cf.attrs['CDATE'] = int(now.strftime('%Y%j'))
    cf.attrs['CTIME'] = int(now.strftime('%H%M%S'))
    cf.attrs['WDATE'] = int(now.strftime('%Y%j'))
    cf.attrs['WTIME'] = int(now.strftime('%H%M%S'))
    cf.attrs['UPNAM'] = 'conc2photdiag.py'.ljust(16)[:16]
    cf.attrs['FILEDESC'] = fdesc.ljust(4800)[:4800]
    acf = cf.isel(TSTEP=slice(1, None)).copy()
    acf.attrs['STIME'] = int(cf['TFLAG'][1, 0, 1])
    for spc in spcs:
        colkey = colkeys[spc]
        acf[colkey][:] = (cf[colkey][1:].data + cf[colkey][:-1].data) / 2
    acf.to_netcdf(colpath, format='NETCDF3_CLASSIC')


if __name__ == '__main__':
    import argparse
    prsr = argparse.ArgumentParser(
        formatter_class=argparse.RawTextHelpFormatter
    )
    prsr.add_argument(
        '-v', '--verbose', action='count', default=0, help='Level of verbosity'
    )
    prsr.add_argument(
        '--overpassh', default=13.5, type=float,
        help='Overpass time in LST (default: 13.5)'
    )
    prsr.add_argument(
        '--daterange', default=None, help='2022-04-01,2022-09-30'
    )
    helpstr = 'Path to GRIDCRO2D; with daterange use %%y%%m%%d'
    prsr.add_argument('g2path', help=helpstr)
    helpstr = 'Path to METCRO3D; with daterange use %%y%%m%%d'
    prsr.add_argument('m3path', help=helpstr)
    helpstr = 'Path to CONC; with daterange use %%y%%m%%d'
    prsr.add_argument('concpath', help=helpstr)
    helpstr = 'Path to output; with daterange use %%y%%m%%d'
    prsr.add_argument('colpath', help=helpstr)
    helpstr = 'Path to output overpass averaged; with daterange use %%y%%m%%d'
    prsr.add_argument('ovppath', help=helpstr, nargs='?', default=None)
    g2path = 'inputs/GRIDCRO2D.36US3.35L.%y%m%d'
    m3path = 'inputs/METCRO3D.36US3.35L.%y%m%d'
    concpath = 'inputs/CCTM_CONC_36US3_%Y%m%d.nc'
    colpath = 'outputs/CCTM_COLUMN_36US3_%Y%m%d.nc'
    ovppath = 'outputs/CCTM_OVPCOLUMN_36US3_%Y%m%d.nc'
    inputs1 = [
        '--daterange=2022-04-21,2022-04-21', g2path, m3path, concpath, colpath
    ]
    inputstr1 = ' '.join(inputs1).replace('%', '%%')
    tmpdate = pd.to_datetime('2022-04-21')
    inputs2 = [
        tmpdate.strftime(tmpl)
        for tmpl in [g2path, m3path, concpath, colpath, ovppath]
    ]
    inputstr2 = ' '.join(inputs2).replace('%', '%%')
    prsr.description = __doc__
    prsr.epilog = f"""
Example 1:
    %(prog)s {inputstr1}
    
Example 2:
    %(prog)s {inputstr2}
    """
    args = vars(prsr.parse_args())
    if args['verbose'] > 0:
        print(args)
    if args['daterange'] is None:
        args.pop('daterange')
        tocolumns(**args)
    else:
        start, end = args.pop('daterange').split(',')
        dates = pd.date_range(start, end, freq='1D')
        for date in dates:
            dargs = {k: v for k, v in args.items()}
            dargs.update({
                k: date.strftime(v) for k, v in dargs.items()
                if isinstance(v, str)
            })
            tocolumns(**dargs)