def test_triangular_sample_compression_2d(): np.random.seed(0) n = 5000 x = np.random.rand(n) y = np.random.rand(n) w = np.random.rand(n) cov = np.identity(2) tri, W = triangular_sample_compression_2d(x, y, cov, w) assert len(W) == 1000 assert np.isclose(sum(W), sum(w), rtol=1e-1)
def kde_contour_plot_2d(ax, data_x, data_y, *args, **kwargs): """Plot a 2d marginalised distribution as contours. This functions as a wrapper around matplotlib.axes.Axes.tricontour, and matplotlib.axes.Axes.tricontourf with a kernel density estimation computation provided by scipy.stats.gaussian_kde in between. All remaining keyword arguments are passed onwards to both functions. Parameters ---------- ax: matplotlib.axes.Axes axis object to plot on. data_x, data_y: numpy.array x and y coordinates of uniformly weighted samples to generate kernel density estimator. weights: numpy.array, optional Sample weights. ncompress: int, optional Degree of compression. optional, Default 1000 xmin, xmax, ymin, ymax: float lower/upper prior bounds in x/y coordinates. optional, default None Returns ------- c: matplotlib.contour.QuadContourSet A set of contourlines or filled regions """ xmin = kwargs.pop('xmin', None) xmax = kwargs.pop('xmax', None) ymin = kwargs.pop('ymin', None) ymax = kwargs.pop('ymax', None) weights = kwargs.pop('weights', None) ncompress = kwargs.pop('ncompress', 1000) label = kwargs.pop('label', None) zorder = kwargs.pop('zorder', 1) linewidths = kwargs.pop('linewidths', 0.5) color = kwargs.pop('color', next(ax._get_lines.prop_cycler)['color']) if len(data_x) == 0 or len(data_y) == 0: return numpy.zeros(0), numpy.zeros(0), numpy.zeros((0, 0)) cov = numpy.cov(data_x, data_y, aweights=weights) tri, w = triangular_sample_compression_2d(data_x, data_y, cov, weights, ncompress) kde = gaussian_kde([tri.x, tri.y], weights=w) x, y = kde.resample(ncompress) x = numpy.concatenate([tri.x, x]) y = numpy.concatenate([tri.y, y]) w = numpy.concatenate([w, numpy.zeros(ncompress)]) tri = scaled_triangulation(x, y, cov) p = kde([tri.x, tri.y]) sigmax = numpy.sqrt(kde.covariance[0, 0]) p = cut_and_normalise_gaussian(tri.x, p, sigmax, xmin, xmax) sigmay = numpy.sqrt(kde.covariance[1, 1]) p = cut_and_normalise_gaussian(tri.y, p, sigmay, ymin, ymax) contours = iso_probability_contours_from_samples(p, weights=w) cmap = basic_cmap(color) cbar = ax.tricontourf(tri, p, contours, cmap=cmap, zorder=zorder, vmin=0, vmax=p.max(), *args, **kwargs) for c in cbar.collections: c.set_cmap(cmap) ax.tricontour(tri, p, contours, zorder=zorder, vmin=0, vmax=p.max(), linewidths=linewidths, colors='k', *args, **kwargs) ax.patches += [ plt.Rectangle((0, 0), 1, 1, fc=cmap(0.999), ec=cmap(0.32), lw=2, label=label) ] ax.set_xlim(*check_bounds(tri.x, xmin, xmax), auto=True) ax.set_ylim(*check_bounds(tri.y, ymin, ymax), auto=True) return cbar
def kde_contour_plot_2d(ax, data_x, data_y, *args, **kwargs): """Plot a 2d marginalised distribution as contours. This functions as a wrapper around matplotlib.axes.Axes.tricontour, and matplotlib.axes.Axes.tricontourf with a kernel density estimation computation provided by scipy.stats.gaussian_kde in between. All remaining keyword arguments are passed onwards to both functions. Parameters ---------- ax: matplotlib.axes.Axes axis object to plot on. data_x, data_y: np.array x and y coordinates of uniformly weighted samples to generate kernel density estimator. weights: np.array, optional Sample weights. levels: list, optional amount of mass within each iso-probability contour. Has to be ordered from outermost to innermost contour. optional, default [0.95, 0.68] ncompress: int, optional Degree of compression. optional, Default 1000 xmin, xmax, ymin, ymax: float lower/upper prior bounds in x/y coordinates. optional, default None Returns ------- c: matplotlib.contour.QuadContourSet A set of contourlines or filled regions """ kwargs = normalize_kwargs(kwargs, dict(linewidths=['linewidth', 'lw'], linestyles=['linestyle', 'ls'], color=['c'], facecolor=['fc'], edgecolor=['ec'])) xmin = kwargs.pop('xmin', None) xmax = kwargs.pop('xmax', None) ymin = kwargs.pop('ymin', None) ymax = kwargs.pop('ymax', None) weights = kwargs.pop('weights', None) ncompress = kwargs.pop('ncompress', 1000) label = kwargs.pop('label', None) zorder = kwargs.pop('zorder', 1) levels = kwargs.pop('levels', [0.95, 0.68]) color = kwargs.pop('color', next(ax._get_lines.prop_cycler)['color']) facecolor = kwargs.pop('facecolor', color) edgecolor = kwargs.pop('edgecolor', color) kwargs.pop('q', None) if len(data_x) == 0 or len(data_y) == 0: return np.zeros(0), np.zeros(0), np.zeros((0, 0)) if weights is not None: data_x = data_x[weights != 0] data_y = data_y[weights != 0] weights = weights[weights != 0] cov = np.cov(data_x, data_y, aweights=weights) tri, w = triangular_sample_compression_2d(data_x, data_y, cov, weights, ncompress) kde = gaussian_kde([tri.x, tri.y], weights=w) x, y = kde.resample(ncompress) x = np.concatenate([tri.x, x]) y = np.concatenate([tri.y, y]) w = np.concatenate([w, np.zeros(ncompress)]) tri = scaled_triangulation(x, y, cov) p = kde([tri.x, tri.y]) sigmax = np.sqrt(kde.covariance[0, 0]) p = cut_and_normalise_gaussian(tri.x, p, sigmax, xmin, xmax) sigmay = np.sqrt(kde.covariance[1, 1]) p = cut_and_normalise_gaussian(tri.y, p, sigmay, ymin, ymax) levels = iso_probability_contours_from_samples(p, contours=levels, weights=w) if facecolor not in [None, 'None', 'none']: linewidths = kwargs.pop('linewidths', 0.5) cmap = kwargs.pop('cmap', basic_cmap(facecolor)) contf = ax.tricontourf(tri, p, levels=levels, cmap=cmap, zorder=zorder, vmin=0, vmax=p.max(), *args, **kwargs) for c in contf.collections: c.set_cmap(cmap) ax.patches += [plt.Rectangle((0, 0), 0, 0, lw=2, label=label, fc=cmap(0.999), ec=cmap(0.32))] cmap = None else: linewidths = kwargs.pop('linewidths', plt.rcParams.get('lines.linewidth')) cmap = kwargs.pop('cmap', None) contf = None ax.patches += [ plt.Rectangle((0, 0), 0, 0, lw=2, label=label, fc='None' if cmap is None else cmap(0.999), ec=edgecolor if cmap is None else cmap(0.32)) ] edgecolor = edgecolor if cmap is None else None vmin, vmax = match_contour_to_contourf(levels, vmin=0, vmax=p.max()) cont = ax.tricontour(tri, p, levels=levels, zorder=zorder, vmin=vmin, vmax=vmax, linewidths=linewidths, colors=edgecolor, cmap=cmap, *args, **kwargs) ax.set_xlim(*check_bounds(tri.x, xmin, xmax), auto=True) ax.set_ylim(*check_bounds(tri.y, ymin, ymax), auto=True) return contf, cont