Source code for pycif.plugins.models.wrfchem.io.outputs.read_sim

import os
import logging
import glob
import bisect
import copy
import pandas as pd
import numpy as np
import netCDF4 as nc
import datetime as dtm
from ......utils.check.errclass import CifFileNotFoundError, CifRuntimeError




def read_sim(self, data2dump, runsubdir, mode, ddi, ddf):
    """Read simulated observations from WRF-Chem output"""
    # This borrows from CTDAS-WRF
    
    # First, check whether there is anything to do.
    if mode != "fwd":
        raise CifRuntimeError("Only works with  mode 'fwd'")
    
    # Prepare a bit for multiple domains, but work only with one domain
    # for now
    wrfout_files = dict()
    wrfout_start_dtm_ids = dict()
    dom = 1
    
    # Get wrfout files
    fp_pattern = os.path.join(runsubdir, f"wrfout_d{dom:02d}_*00")
    wrfout_files[dom] = glob.glob(fp_pattern)
    if len(wrfout_files[dom]) == 0:
        raise CifFileNotFoundError("No WRF output found.")
    
    wrfout_files[dom] = np.sort(wrfout_files[dom]).tolist()
    _, wrfout_start_dtm_ids[dom] = \
        self.wrfhelper.wrf_times(wrfout_files[dom])
    
    dataout = {}
    for trid in data2dump:
        spec = trid[1]
        dat = self.dataobs[ddi][spec]
        
        # Get spatial indices for sampling
        iv = np.round(np.array(dat[("metadata", "i")])).astype(int)
        jv = np.round(np.array(dat[("metadata", "j")])).astype(int)
        # During development, the wrf domain changed from structured to unstructured.
        # Here, we unravel the indices in case we have the unstructured version.
        if self.domain.unstructured_domain:
            iv, jv = np.unravel_index(iv, self.domain.zlon2d.shape)
        levelv = np.round(np.array(dat[("metadata", "level")])).astype(int)

        # Get files and time index in file from tstep for sampling
        ind_all_t = dat["metadata"]["tstep"].values
        ind_f = [bisect.bisect_right(wrfout_start_dtm_ids[dom], ind_t) - 1 for ind_t in ind_all_t]
        files = [wrfout_files[dom][ind_f_] for ind_f_ in ind_f]
        ind_t_file = [ind_t - wrfout_start_dtm_ids[dom][ind_f_] for ind_t, ind_f_ in zip(ind_all_t, ind_f)]
        loc_ll = zip(files, ind_t_file)
        
        # This line is slow:
        loc = pd.DataFrame(loc_ll, columns=["file", "ind_t"])
        
        # Now that all the information for sampling is collected, read
        # data from wrfout files.
        # For vectorized indexing, have to read full netcdf output.
        # Therefore looping over timesteps with data, as some may not
        # have any. Perhaps this can be sped up, e.g. with xarray.
        loc["value"] = np.nan

        if trid[0] in ["pressure", "dpressure"]:
            ptop = float(self.namelist["domains"]["p_top_requested"]) # Pa
            loc["psfc"] = -999.
        if trid[0] in ["concs"]:
            varname = getattr(getattr(self.chemistry.acspecies, spec), "varname")
        UF = list(set(loc["file"]))

        for f in UF:
            ind_f = np.where(loc["file"] == f)[0]
            ncf = nc.Dataset(f, "r")
            UIT = list(set(loc.loc[ind_f, "ind_t"]))
            for it in UIT:
                ind = ind_f[np.where(loc.loc[ind_f, "ind_t"] == it)[0]]
                if trid[0] in ["concs"]:
                    field = ncf[varname][it, ...]
                    loc.loc[ind, "value"] = field[levelv[ind], iv[ind], jv[ind]]

                # Get surface pressure (for pressure
                # midpoints/layer thicknesses)
                if trid[0] in ["pressure", "dpressure"]:
                    psfc = ncf["PSFC"][it, ...] # Pa
                    loc.loc[ind, "psfc"] = psfc[iv[ind], jv[ind]]

            ncf.close()

        # Compute pressure midpoints/layer thicknesses from surface pressure
        if trid[0] in ["pressure", "dpressure"]:
            if trid[0] in ["pressure"]:
                loc["value"] = \
                    self.domain.sigma_a_mid[levelv] + loc["psfc"].values*self.domain.sigma_b_mid[levelv]
            if trid[0] in ["dpressure"]:
                loc["value"] = \
                    (self.domain.sigma_a[levelv]   + loc["psfc"].values*self.domain.sigma_b[levelv]) \
                   -(self.domain.sigma_a[levelv+1] + loc["psfc"].values*self.domain.sigma_b[levelv+1])
            del loc["psfc"]


        # Write values to output variable
        dat[("maindata", "spec")] = loc["value"].values
        
        # This is copied from chimere:
        if trid[0] in ["pressure", "dpressure"]:
            dat[("maindata", "incr")] = 0.

        dataout[trid] = dat

    return dataout