示例#1
0
文件: postwrf.py 项目: cb678/met4ene
def _wrf2xarray(netcdf_data, query_variables):
    """
    Gets data from the netcdf wrfout file and uses wrf-python
    to create an xarray Dataset.

    Parameters
    ----------
    data: netcdf
        Data returned from UNIDATA NCSS query, or from your local forecast.
    query_variables: list
        The variables requested.
    start: Timestamp
        The start time
    end: Timestamp
        The end time

    Returns
    -------
    xarray.Dataset
    """
    first = True
    for key in query_variables:
        if key in ['wspd', 'wdir']:
            var = get_wspd_wdir(netcdf_data, key)
        else:
            var = wrf.getvar(netcdf_data, key, timeidx=wrf.ALL_TIMES)
        if first:
            data = var
            first = False
        else:
            with xr.set_options(keep_attrs=True):
                try:
                    data = xr.merge([data, var])
                except ValueError:
                    data = data.drop_vars('Time')
                    data = xr.merge([data, var])

    # Get global attributes from the NetCDF Dataset
    wrfattrs_names = netcdf_data.ncattrs()
    wrfattrs = wrf.extract_global_attrs(netcdf_data, wrfattrs_names)
    data = data.assign_attrs(wrfattrs)

    # Fix a bug in how wrfout data is read in -- attributes must be strings to be written to NetCDF
    for var in data.data_vars:
        try:
            data[var].attrs['projection'] = str(data[var].attrs['projection'])
        except KeyError:
            pass

    # Fix another bug that creates a conflict in the 'coordinates' attribute
    for var in data.data_vars:
        try:
            del data[var].attrs['coordinates']
        except KeyError:
            pass

    return data
示例#2
0
    def _wrf2xarray(self, netcdf_data, query_variables, start, end):
        """
        Transforms data from netcdf to xarray Dataset.

        Parameters
        ----------
        data: netcdf
            Data returned from UNIDATA NCSS query, or from your local forecast.
        query_variables: list
            The variables requested.
        start: Timestamp
            The start time
        end: Timestamp
            The end time

        Returns
        -------
        xarray.Dataset
        """
        first = True
        for key in query_variables:
            var = wrf.getvar(netcdf_data, key, timeidx=wrf.ALL_TIMES)
            if first:
                data = var
                first = False
            else:
                with xr.set_options(keep_attrs=True):
                    try:
                        data = xr.merge([data, var])
                    except ValueError:
                        data = data.drop_vars('Time')
                        data = xr.merge([data, var])

        # Get global attributes from the NetCDF Dataset
        wrfattrs_names = netcdf_data.ncattrs()
        wrfattrs = wrf.extract_global_attrs(netcdf_data, wrfattrs_names)
        data = data.assign_attrs(wrfattrs)

        # Fix a bug in how wrfout data is read in -- attributes must be strings to be written to NetCDF
        for var in data.data_vars:
            try:
                data[var].attrs['projection'] = str(
                    data[var].attrs['projection'])
            except KeyError:
                pass

        # Fix another bug that creates a conflict in the 'coordinates' attribute
        for var in data.data_vars:
            try:
                del data[var].attrs['coordinates']
            except KeyError:
                pass

        # Slice the dataset to only include specified time interval
        data = data.sel(Time=slice(start, end))
        return data
示例#3
0
def extract_proj_params(wrfnc):
    attrs = extract_global_attrs(
        wrfnc, ("MAP_PROJ", "TRUELAT1", "TRUELAT2", "STAND_LON", "POLE_LAT",
                "POLE_LON", "DX", "DY"))

    result = {key.lower(): val for key, val in viewitems(attrs)}

    if is_multi_file(wrfnc):
        wrfnc = wrfnc[0]

    result["known_x"] = 0
    result["known_y"] = 0
    result["ref_lat"] = wrfnc.variables["XLAT"][0, 0, 0]
    result["ref_lon"] = wrfnc.variables["XLONG"][0, 0, 0]

    return result
    def __init__(self, config):
        """ construct input wrf file names """

        self.ncfile = Dataset(config['INPUT']['input_fn'])
        self.resolution = float(
            wrf.extract_global_attrs(self.ncfile, 'DX')['DX'])

        n_radius = int(
            int(config['INPUT']['vis_radius']) / int(self.resolution))

        times_dt64 = wrf.getvar(self.ncfile,
                                'Times',
                                timeidx=wrf.ALL_TIMES,
                                method="cat").values
        self.tgt_timeframe = datetime.datetime.strptime(
            config['INPUT']['vis_timestamp'], '%Y%m%d%H%M%S')
        print(print_prefix + 'Selected Timeframe:%s' %
              self.tgt_timeframe.strftime('%Y-%m-%d_%H:%M:%S'))

        # get idx of the field will be extract
        tgt_times_dt64 = np.datetime64(self.tgt_timeframe)
        t_idx = np.where(times_dt64 == tgt_times_dt64)[0][0]

        print(print_prefix +
              'init from single input file for lat2d, lon2d, and hgt')
        self.xlat = wrf.getvar(self.ncfile, 'XLAT')
        self.xlon = wrf.getvar(self.ncfile, 'XLONG')

        print(print_prefix + 'init from single input file for Z4d')
        self.SLP = wrf.getvar(self.ncfile, 'slp', timeidx=t_idx)

        self.tc_ix, self.tc_iy = find_tc_center(self.SLP)

        print(print_prefix + 'init from single input file for Z4d')
        self.Z = wrf.getvar(self.ncfile, 'z', timeidx=t_idx)

        print(print_prefix + 'init from single input file for U4d')
        self.U = wrf.getvar(self.ncfile, 'ua', timeidx=t_idx)

        print(print_prefix + 'init from single input file for V4d')
        self.V = wrf.getvar(self.ncfile, 'va', timeidx=t_idx)

        print(print_prefix + 'init from single input file for W4d')
        self.W = wrf.getvar(self.ncfile, 'wa', timeidx=t_idx)

        print(print_prefix + 'init multi files successfully!')
示例#5
0
def extract_proj_params(wrfnc, timeidx=0):
    attrs = extract_global_attrs(
        wrfnc, ("MAP_PROJ", "TRUELAT1", "TRUELAT2", "STAND_LON", "POLE_LAT",
                "POLE_LON", "DX", "DY"))

    result = {key.lower(): val for key, val in viewitems(attrs)}

    _timeidx = timeidx
    if is_multi_file(wrfnc):
        wrfnc0 = wrfnc[0]
        num_times_per_file = len(wrfnc0.dimensions["Time"])
        file_idx = timeidx // num_times_per_file
        _timeidx = timeidx % num_times_per_file

        wrfnc = wrfnc[file_idx]

    result["known_x"] = 0
    result["known_y"] = 0
    result["ref_lat"] = wrfnc.variables["XLAT"][_timeidx, 0, 0]
    result["ref_lon"] = wrfnc.variables["XLONG"][_timeidx, 0, 0]

    return result
示例#6
0
def open_mfdataset(fname,
                   convert_to_ppb=True,
                   mech="racm_esrl_vcp",
                   var_list=["o3"],
                   surf_only=False,
                   surf_only_nc=False,
                   **kwargs):
    """Method to open WRF-chem and RAP-chem netcdf files.

    Parameters
    ----------
    fname : string or list
        fname is the path to the file or files.  It will accept hot keys in
        strings as well.
    convert_to_ppb : boolean
        If true the units of the gas species will be converted to ppbv
    mech: str
        Mechanism to be used for calculating sums. Supported mechanisms
        include: "racm_esrl_vcp" and "redhc"
    var_list: list
        List of variables to include in output. MELODIES-MONET only reads in
        variables need to plot in order to save on memory and simulation cost
        especially for vertical data
    surf_only: boolean
        Whether to save only surface data to save on memory and computational
        cost (True) or not (False).
    surf_only_nc : bool
        Whether input data has been subsampled to only contain surface data and
        not extra pressure & temperature info (True) or not (False).

    Returns
    -------
    xarray.DataSet
        WRF-Chem or RAP-Chem model dataset in standard format for use
        in MELODIES-MONET


    """
    from netCDF4 import Dataset
    from wrf import ALL_TIMES, extract_global_attrs, getvar

    # Get dictionary of summed species for the mechanism of choice.
    dict_sum = dict_species_sums(mech=mech)

    list_calc_sum = []
    for var_sum in [
            "noy_gas",
            "noy_aer",
            "nox",
            "pm25_cl",
            "pm25_ec",
            "pm25_na",
            "pm25_nh4",
            "pm25_no3",
            "pm25_so4",
            "pm25_om",
    ]:
        if var_sum in var_list:
            var_list.extend(dict_sum[var_sum])
            var_list.remove(var_sum)
            list_calc_sum.append(var_sum)

    wrflist = []
    for files in fname:
        wrflist.append(Dataset(files))

    if not surf_only_nc:
        # Add some additional defaults needed for aircraft analysis
        # Turn this on also if need to convert aerosols
        var_list.append("pres")
        var_list.append("height")
        var_list.append("tk")
        var_list.append("height_agl")
        var_list.append("PSFC")
        # need to calculate surface pressure and dp and optionally dz here.

    var_wrf_list = []
    for var in var_list:
        if var == "pres":  # Insert special versions.
            var_wrf = getvar(wrflist,
                             var,
                             timeidx=ALL_TIMES,
                             method="cat",
                             squeeze=False,
                             units="Pa")
        elif var == "height":
            var_wrf = getvar(wrflist,
                             var,
                             timeidx=ALL_TIMES,
                             method="cat",
                             squeeze=False,
                             units="m")
        elif var == "height_agl":
            var_wrf = getvar(wrflist,
                             var,
                             timeidx=ALL_TIMES,
                             method="cat",
                             squeeze=False,
                             units="m")
        else:
            var_wrf = getvar(wrflist,
                             var,
                             timeidx=ALL_TIMES,
                             method="cat",
                             squeeze=False)
        var_wrf_list.append(var_wrf)

    dset = xr.merge(var_wrf_list)

    # Add global attributes needed
    a_truelat1 = extract_global_attrs(wrflist[0], "TRUELAT1")
    a_truelat2 = extract_global_attrs(wrflist[0], "TRUELAT2")
    a_moad_cen_lat = extract_global_attrs(wrflist[0], "MOAD_CEN_LAT")
    a_stand_lon = extract_global_attrs(wrflist[0], "STAND_LON")
    a_map_proj = extract_global_attrs(wrflist[0], "MAP_PROJ")
    a_cen_lat = extract_global_attrs(wrflist[0], "CEN_LAT")
    a_cen_lon = extract_global_attrs(wrflist[0], "CEN_LON")
    dset = dset.assign_attrs(a_truelat1)
    dset = dset.assign_attrs(a_truelat2)
    dset = dset.assign_attrs(a_moad_cen_lat)
    dset = dset.assign_attrs(a_stand_lon)
    dset = dset.assign_attrs(a_map_proj)
    dset = dset.assign_attrs(a_cen_lat)
    dset = dset.assign_attrs(a_cen_lon)

    for i in dset.variables:
        dset[i] = dset[i].assign_attrs(a_truelat1)
        dset[i] = dset[i].assign_attrs(a_truelat2)
        dset[i] = dset[i].assign_attrs(a_moad_cen_lat)
        dset[i] = dset[i].assign_attrs(a_stand_lon)
        dset[i] = dset[i].assign_attrs(a_map_proj)
        dset[i] = dset[i].assign_attrs(a_cen_lat)
        dset[i] = dset[i].assign_attrs(a_cen_lon)

    # Convert names to standards used in MONET
    dset = dset.rename({
        "Time": "time",
        "south_north": "y",
        "west_east": "x",
        "XLONG": "longitude",
        "XLAT": "latitude",
    })

    # These sums and conversions are quite expensive and memory intensive,
    # so add option to shrink dataset to just surface when needed
    if (not surf_only_nc) and surf_only:
        dset = dset.isel(bottom_top=0).expand_dims("bottom_top", axis=1)

    # convert all gas species to ppbv
    if convert_to_ppb:
        for i in dset.variables:
            if "units" in dset[i].attrs:
                if "ppmv" in dset[i].attrs["units"]:
                    dset[i] = dset[i] * 1000.0
                    dset[i].attrs["units"] = "ppbv"
    # convert "ug/kg-dryair -> ug/m3"
    for i in dset.variables:
        if "units" in dset[i].attrs:
            if "ug/kg-dryair" in dset[i].attrs["units"]:
                # ug/kg -> ug/m3 using dry air density
                dset[i] = dset[i] * dset["pressure"] / dset["temp"] / 287.05535
                dset[i].attrs["units"] = r"$\mu g m^{-3}$"

    # assign mapping table for airnow
    dset = _predefined_mapping_tables(dset)

    # add lazy diagnostic variables
    if "noy_gas" in list_calc_sum:
        dset = add_lazy_noy_g(dset, dict_sum)
    if "noy_aer" in list_calc_sum:
        dset = add_lazy_noy_a(dset, dict_sum)
    if "nox" in list_calc_sum:
        dset = add_lazy_nox(dset, dict_sum)
    if "pm25_cl" in list_calc_sum:
        dset = add_lazy_cl_pm25(dset, dict_sum)
    if "pm25_ec" in list_calc_sum:
        dset = add_lazy_ec_pm25(dset, dict_sum)
    if "pm25_na" in list_calc_sum:
        dset = add_lazy_na_pm25(dset, dict_sum)
    if "pm25_nh4" in list_calc_sum:
        dset = add_lazy_nh4_pm25(dset, dict_sum)
    if "pm25_no3" in list_calc_sum:
        dset = add_lazy_no3_pm25(dset, dict_sum)
    if "pm25_so4" in list_calc_sum:
        dset = add_lazy_so4_pm25(dset, dict_sum)
    if "pm25_om" in list_calc_sum:
        dset = add_lazy_om_pm25(dset, dict_sum)

    dset = dset.reset_index(["XTIME", "datetime"], drop=True)
    if not surf_only_nc:
        # Reset more variables
        dset = dset.rename({
            "bottom_top": "z",
            "temp": "temperature_k",
            "height": "alt_msl_m_mid",
            "height_agl": "alt_agl_m_mid",
            "PSFC": "surfpres_pa",
            "pressure": "pres_pa_mid",
        })
        dset2 = dset
    else:
        # Expand into z coordinate so that format is consistent.
        if "bottom_top" in dset.dims:
            # Workaround for tutorial version of the WRF-Chem example dataset
            # TODO: eliminate need
            dset2 = dset.rename({"bottom_top": "z"})
        else:
            dset2 = dset.expand_dims("z", axis=3).copy()

    dset2 = dset2.reset_coords()
    dset2 = dset2.set_coords(["latitude", "longitude"])

    return dset2