Python DataArray.groupby方法代码示例

# 需要导入模块: from xarray import DataArray [as 别名]
# 或者: from xarray.DataArray import groupby [as 别名]
def _nested_groupby_apply(array: xr.DataArray,
                          groupby: list,
                          apply_fn: object,
                          kwargs: dict):
    """
    Perform a nested groupby over given dimensions and apply a function on the
    last 'slice'

    :param array: xr.DataArray to perform groupby on
    :param groupby: a list of coordinate labels over which to perform groupby
    :param apply_fn: The function to apply
    :return: groupby-split-appy result
    """
    if len(groupby) == 1:
        return array.groupby(groupby[0], squeeze=True).apply(apply_fn, **kwargs)
    else:
        return array.groupby(groupby[0], squeeze=True).apply(_nested_groupby_apply,
                                                             groupby=groupby[1:],
                                                             apply_fn=apply_fn,
                                                             kwargs=kwargs)

# 需要导入模块: from xarray import DataArray [as 别名]
# 或者: from xarray.DataArray import groupby [as 别名]
def ect_adjust_geometry(master: xr.DataArray, slave: xr.DataArray) -> xr.DataArray:
    import scipy.ndimage

    lat_factor = len(master.coords['latitude']) / len(slave.coords['lat'])
    lon_factor = len(master.coords['longitude']) / len(slave.coords['lon'])

    def resample(x):
        y = scipy.ndimage.zoom(x, [lat_factor, lon_factor])
        return xr.DataArray(y)

    # Help! This is soooo slow... few minutes on Norman's PC
    temp_da = slave.groupby('time').apply(resample)
    temp_lon = scipy.ndimage.zoom(slave.lon, [lon_factor])
    temp_lat = scipy.ndimage.zoom(slave.lat, [lat_factor])
    return xr.DataArray(temp_da,
                        name=slave.name,
                        dims=['time', 'lat', 'lon'],
                        coords=dict(time=slave.time, lat=temp_lat, lon=temp_lon),
                        attrs=slave.attrs)

# 需要导入模块: from xarray import DataArray [as 别名]
# 或者: from xarray.DataArray import groupby [as 别名]
def _is_seasonal(time: xr.DataArray):
    """
    Check if the given timestamp dataarray features consistent
    seasons. E.g. Each year has the same date-month values in it.
    """
    c = 0
    test = None
    for group in time.groupby('time.year'):
        # Test (month, day) dates of all years against
        # (month, day) dates of the first year, or second
        # year in case the first year is not full
        c = c + 1
        np_time = group[1].time.values
        months = pd.DatetimeIndex(np_time).month
        days = pd.DatetimeIndex(np_time).day
        if c == 1:
            first_months = months
            first_days = days
            continue
        elif c == 2:
            second_months = months
            second_days = days
            if len(second_months) > len(first_months):
                test = list(zip(second_months, second_days))
                for date in zip(first_months, first_days):
                    if date not in test:
                        return False
            else:
                test = list(zip(first_months, first_days))
                for date in zip(second_months, second_days):
                    if date not in test:
                        return False
            continue

        for date in zip(months, days):
            if date not in test:
                return False

    return True

# 需要导入模块: from xarray import DataArray [as 别名]
# 或者: from xarray.DataArray import groupby [as 别名]
def plot_monthly_panels(data:xarray.DataArray, basemap: Basemap, img_dir="default_img_dir",
                        data_label="",
                        color_levels=None, cmap=cm.get_cmap("jet")):
    """

    :param data:
    :param basemap:
    :param img_dir:
    :param data_label: should contain period and simulation and obs sources used to get the data
    :param color_levels:
    :param cmap:
    """

    plot_utils.apply_plot_params(font_size=14, width_cm=30)

    img_dir = Path(img_dir)
    img_dir.mkdir(parents=True, exist_ok=True)


    xx, yy = basemap(data["lon"].values, data["lat"].values)


    # calculate monthly means:
    monthly_data = data.groupby("t.month").mean(dim="t")

    nrows = 3
    ncols = 4

    gs = GridSpec(nrows=nrows, ncols=ncols, wspace=0.0)
    fig = plt.figure()


    norm = BoundaryNorm(color_levels, len(color_levels) - 1)
    cmap = cm.get_cmap(cmap, len(color_levels) - 1)

    month_with_colorbar = 2
    month_with_data_label = 5
    month_with_big_data_label = 2

    for ind, month in enumerate([12, ] + list(range(1, 12))):
        i, j = __get_ij_from_index(ind, ncols, nrows)
        ax = fig.add_subplot(gs[i, j])

        ax.set_title(calendar.month_abbr[month])

        im = basemap.pcolormesh(xx, yy,
                monthly_data.sel(month=month).to_masked_array(), cmap=cmap, norm=norm, ax=ax)
        basemap.drawcoastlines(ax=ax)
        basemap.drawstates(ax=ax, linewidth=0.5)
        basemap.drawcountries(ax=ax, linewidth=0.5)

        cb = basemap.colorbar(im, location="bottom")

        cb.ax.set_visible(month == month_with_colorbar) # show only the colorbar for October

        if month == month_with_colorbar:
            cb.ax.set_xticklabels(cb.ax.get_xticklabels(), rotation=45)


        if month == month_with_data_label:
            ax.set_xlabel("ndrw" + data_label.split("_ndrw")[1], ha="left", fontsize=10)

        if month == month_with_big_data_label:
            ax.set_ylabel(data_label.split("_ndrw")[0], ha="left")

    # save the plot to a file
    img_path = img_dir / f"{data_label}.png"
    fig.savefig(str(img_path), dpi=400, bbox_inches="tight")
    plt.close(fig)