Python DataArray.squeeze方法代码示例

# 需要导入模块: from xarray import DataArray [as 别名]
# 或者: from xarray.DataArray import squeeze [as 别名]
def _resample_slice(arr_slice: xr.DataArray, w: int, h: int, ds_method: int, us_method: int,
                    parent_monitor: Monitor) -> xr.DataArray:
    """
    Resample a single time slice of a larger xr.DataArray

    :param arr_slice: xr.DataArray single slice
    :param w: The desired new width (amount of longitudes)
    :param h: The desired new height (amount of latitudes)
    :param ds_method: Downsampling method, see resampling.py
    :param us_method: Upsampling method, see resampling.py
    :param parent_monitor: the parent progress monitor.
    :return: resampled slice
    """
    monitor = parent_monitor.child(1)
    with monitor.observing("resample slice"):
        # In some cases the grouped dimension is not automatically squeezed out
        result = resampling.resample_2d(np.ma.masked_invalid(arr_slice.squeeze().values),
                                        w,
                                        h,
                                        ds_method,
                                        us_method)
        return xr.DataArray(result)

# 需要导入模块: from xarray import DataArray [as 别名]
# 或者: from xarray.DataArray import squeeze [as 别名]
def plot_correlation_maps_with_stfl(b: Basemap, xx, yy,
                                    hles_snfall, total_snfall, period_list, stfl_series: xarray.DataArray,
                                    label="", stfl_months_of_interest=tuple(range(1, 6)),
                                    hles_crit_fraction=0.1):

    hles_region = (hles_snfall.mean(axis=0) / total_snfall.mean(axis=0) >= hles_crit_fraction)
    i_arr, j_arr = np.where(hles_region)


    # get the coordinates of the streamflow point
    stfl_lon, stfl_lat = stfl_series.coords["lon"].values[0, 0], stfl_series.coords["lat"].values[0, 0]
    stfl_x, stfl_y = b(stfl_lon, stfl_lat)


    stfl = stfl_series.squeeze().to_series()
    stfl.index = stfl_series.coords["t"].values

    stfl = stfl[stfl.index.month.isin(stfl_months_of_interest)]
    stfl = stfl.groupby(stfl.index.year).max()

    stfl = stfl[[p.end.year for p in period_list]].values

    assert len(stfl) == len(period_list)
    assert len(stfl) == total_snfall.shape[0]


    hles_corr = np.ma.masked_all_like(xx)
    pval_hles_corr = np.ma.masked_all_like(xx)

    total_corr = np.ma.masked_all_like(yy)
    pval_total_corr = np.ma.masked_all_like(yy)

    for i, j in zip(i_arr, j_arr):
        hles_corr[i, j], pval_hles_corr[i, j] = pearsonr(hles_snfall[:, i, j], stfl)
        total_corr[i, j], pval_total_corr[i, j] = pearsonr(total_snfall[:, i, j], stfl)


    # plot correlations and pvalues
    corr_levs = list(np.arange(-1, 0, 0.1)) + list(np.arange(0.1, 1.1, 0.1))

    corr_cmap = get_with_white_added("coolwarm",
                                     white_start=0.45, white_end=0.55,
                                     ncolors_out=len(corr_levs) - 1)
    corr_norm = BoundaryNorm(corr_levs, len(corr_levs) - 1)


    pval_levs = [0, 0.01, 0.05, 0.1, 0.2]
    pval_cmap = cm.get_cmap("tab10", len(pval_levs) - 1)
    pval_norm = BoundaryNorm(pval_levs, len(pval_levs) - 1)



    gs = GridSpec(2, 2, wspace=0.05, hspace=0.)

    fig = plt.figure()

    axes_list = []
    ax = __add_subplot(fig, gs, 0, 0, ax_list=axes_list)
    cs = b.pcolormesh(xx, yy, total_corr, cmap=corr_cmap, norm=corr_norm, ax=ax)
    ax.set_title("Corr.")
    ax.set_ylabel("Total snowfall")
    cb = b.colorbar(cs, location="bottom")
    cb.ax.set_visible(False)


    ax = __add_subplot(fig, gs, 0, 1, ax_list=axes_list)
    cs = b.pcolormesh(xx, yy, pval_total_corr, cmap=pval_cmap, norm=pval_norm, ax=ax)
    ax.set_title("p-value")
    cb = b.colorbar(cs, location="bottom", extend="max")
    cb.ax.set_visible(False)


    # HLES
    ax = __add_subplot(fig, gs, 1, 0, ax_list=axes_list)
    cs = b.pcolormesh(xx, yy, hles_corr, cmap=corr_cmap, norm=corr_norm, ax=ax)
    ax.set_ylabel("HLES")
    b.colorbar(cs, location="bottom")

    ax = __add_subplot(fig, gs, 1, 1, ax_list=axes_list)
    cs = b.pcolormesh(xx, yy, pval_hles_corr, cmap=pval_cmap, norm=pval_norm, ax=ax)
    b.colorbar(cs, location="bottom", extend="max")



    # common adjustments
    for ax in axes_list:
        b.drawcoastlines(linewidth=0.5, ax=ax)
        b.scatter(stfl_x, stfl_y, 200, marker="*", zorder=10, color="m", ax=ax)

        # draw basin boundaries
        # add basin patches
        b.readshapefile(default_domains.GL_BASINS_FROM_MUSIC_BILJANA_PATH[:-4],
                        "basin",
                        linewidth=2,
                        color="m", zorder=12)

        patches = []

        for info, shape in zip(b.basin_info, b.basin):
            if info["LAKEBASIN"].lower().endswith("lawrence"): # Skip ottawa river basin
#.........这里部分代码省略.........