예제 #1
0
파일: merra.py 프로젝트: jenfly/atmos-read
    def get_data(nms, pres, year, month, day, concat_dim, subset_dict, verbose):
        # Lists of raw and extended variables
        ids = list(nms)
        ext = []
        for var in ['theta', 'theta_e']:
            if var in ids:
                ext.append(var)
                ids.remove(var)

        # Read raw data and calculate extended variables
        data = read_daily(ids, year, month, day, concat_dim=concat_dim,
                          subset_dict=subset_dict, verbose=verbose)
        if 'theta' in ext:
            print_if('Computing potential temperature', verbose)
            T = data[get_varname('T')]
            data['theta'] = atm.potential_temp(T, pres)
        if 'theta_e' in ext:
            print_if('Computing equivalent potential temperature', verbose)
            T = data[get_varname('T')]
            q = data[get_varname('q')]
            data['theta_e'] = atm.equiv_potential_temp(T, pres, q)

        return data
예제 #2
0
파일: merra.py 프로젝트: jenfly/atmos-read
def calc_fluxes(year, month,
                var_ids=['u', 'q', 'T', 'theta', 'theta_e', 'hgt'],
                concat_dim='TIME', scratchdir=None, keepscratch=False,
                verbose=True):
    """Return the monthly mean of MERRA daily fluxes.

    Reads MERRA daily data from OpenDAP urls, computes fluxes, and
    returns the monthly mean of the daily variable and its zonal and
    meridional fluxes.

    Parameters
    ----------
    year, month : int
        Numeric year and month (1-12).
    var_ids : list of str, optional
        IDs of variables to include.
    concat_dim : str, optional
        Name of dimension for concatenation.
    scratchdir : str, optional
        Directory path to store temporary files while processing data.
        If omitted, the current working directory is used.
    keepscratch : bool, optional
        If True, scratch files are kept in scratchdir. Otherwise they
        are deleted.
    verbose : bool, optional
        If True, print updates while processing files.

    Returns
    -------
    data : xray.Dataset
        Mean of daily data and the mean of the daily zonal fluxes
        (u * var) and meridional fluxes (v * var), for each variable
        in var_ids.
    """

    nms = [get_varname(nm) for nm in atm.makelist(var_ids)]
    u_nm, v_nm = get_varname('u'), get_varname('v')
    nms.extend([u_nm, v_nm])
    if 'theta' in nms:
        nms.append(get_varname('T'))
    if 'theta_e' in nms:
        nms.extend([get_varname('T'), get_varname('q')])
    nms = set(nms)

    days = range(1, atm.days_this_month(year, month) + 1)

    def scratchfile(nm, k, year, month, day):
        filestr = '%s_level%d_%d%02d%02d.nc' % (nm, k, year, month, day)
        if scratchdir is not None:
            filestr = scratchdir + '/' + filestr
        return filestr

    # Read metadata from one file to get pressure-level array
    dataset = 'p_daily'
    url = url_list(dataset, return_dict=False)[0]
    with xray.open_dataset(url) as ds:
        pname = atm.get_coord(ds, 'plev', 'name')
        plev = atm.get_coord(ds, 'plev')
        # Pressure levels in Pa for theta/theta_e calcs
        p_units = atm.pres_units(ds[pname].units)
        pres = atm.pres_convert(plev, p_units, 'Pa')

    # Get daily data (raw and calculate extended variables)
    def get_data(nms, pres, year, month, day, concat_dim, subset_dict, verbose):
        # Lists of raw and extended variables
        ids = list(nms)
        ext = []
        for var in ['theta', 'theta_e']:
            if var in ids:
                ext.append(var)
                ids.remove(var)

        # Read raw data and calculate extended variables
        data = read_daily(ids, year, month, day, concat_dim=concat_dim,
                          subset_dict=subset_dict, verbose=verbose)
        if 'theta' in ext:
            print_if('Computing potential temperature', verbose)
            T = data[get_varname('T')]
            data['theta'] = atm.potential_temp(T, pres)
        if 'theta_e' in ext:
            print_if('Computing equivalent potential temperature', verbose)
            T = data[get_varname('T')]
            q = data[get_varname('q')]
            data['theta_e'] = atm.equiv_potential_temp(T, pres, q)

        return data

    # Iterate over vertical levels
    for k, p in enumerate(plev):
        subset_dict = {pname : (p, p)}
        print_if('Pressure-level %.1f' % p, verbose)

        files = []

        for day in days:
            # Read data for this level and day
            ds = get_data(nms, pres[k], year, month, day, concat_dim,
                           subset_dict, verbose)

            # Compute fluxes
            print_if('Computing fluxes', verbose)
            u = ds[get_varname('u')]
            v = ds[get_varname('v')]
            for nm in var_ids:
                var = ds[get_varname(nm)]
                varname, attrs, _, _ = atm.meta(var)
                u_var = u * var
                v_var = v * var

                u_var.name = get_varname(u_nm) + '*' +  var.name
                units = var.attrs['units'] + ' * ' + u.attrs['units']
                u_var.attrs['units'] = units
                v_var.name = get_varname(v_nm) + '*' +  var.name
                v_var.attrs['units'] = units
                ds[u_var.name] = u_var
                ds[v_var.name] = v_var

            # Save to temporary scratch file
            filenm = scratchfile('fluxes', k, year, month, day)
            files.append(filenm)
            print_if('Saving to scratch file ' + filenm, verbose)
            ds.to_netcdf(filenm)

        # Concatenate daily scratch files
        ds = atm.load_concat(files)

        if not keepscratch:
            for f in files:
                os.remove(f)

        # Compute monthly means
        print_if('Computing monthly means', verbose)
        if k == 0:
            data = ds.mean(dim=concat_dim)
        else:
            data = xray.concat([data, ds.mean(dim=concat_dim)], dim=pname)

    for var in data.data_vars:
        data[var].attrs = ds[var].attrs

    return data