nldas_soil_moisture_ml.py 文件源码

def get_nldas_fora_X_and_vic_y(year, month, day, hour,
                           vic_or_fora, band_order=None,
                           prefix=None, data_arrs=None,
                           keep_columns=None):
    '''Load data from VIC for NLDAS Forcing A Grib files

    Parameters:
        year: year of forecast time
        month: month of forecast time
        day: day of forecast time
        vic_or_fora: string indicating which NLDAS data source
        band_order: list of DataArray names already loaded
        prefix: add a prefix to the DataArray name from Grib
        data_arrs: Add the DataArrays to an existing dict
        keep_columns: Retain only the DataArrays in this list, if given
    Returns:
        tuple or (data_arrs, band_order) where data_arrs is
        an OrderedDict of DataArrays and band_order is their
        order when they are flattened from rasters to a single
        2-D matrix
    '''
    data_arrs = data_arrs or OrderedDict()
    band_order = band_order or []
    path = get_file(year, month, day, hour, dset=vic_or_fora)
    dset = xr.open_dataset(path, engine='pynio')
    for k in dset.data_vars:
        if keep_columns and k not in keep_columns:
            continue
        arr = getattr(dset, k)
        if sorted(arr.dims) != ['lat_110', 'lon_110']:
            continue
        #print('Model: ',f, 'Param:', k, 'Detail:', arr.long_name)
        lon, lat = arr.lon_110, arr.lat_110
        geo_transform = [lon.Lo1, lon.Di, 0.0,
                         lat.La1, 0.0, lat.Dj]
        shp = arr.shape
        canvas = Canvas(geo_transform, shp[1], shp[0], arr.dims)
        arr.attrs['canvas'] = canvas
        if prefix:
            band_name = '{}_{}'.format(prefix, k)
        else:
            band_name = k
        data_arrs[band_name] = arr
        band_order.append(band_name)
    return data_arrs, band_order