def kde2d(X: Union[np.ndarray, Series, List, Tuple], Y: Union[np.ndarray, Series, List, Tuple], c: str = "red", ax: mpl.axes.Axes = None, fill: bool = False, with_scatter: bool = False, **contour_kwargs): """TODO: Generates a 2D KDE using contours.""" instance_check((X, Y), (list, tuple, np.ndarray, Series)) instance_check(c, str) instance_check((fill, with_scatter), bool) instance_check(ax, mpl.axes.Axes) arrays_equal_size(X, Y) # calculate density _X, _Y = remove_na(np.asarray(X), np.asarray(Y), paired=True) H = density(_X, _Y) offx = np.abs(_X.max() - _X.min()) / 15.0 offy = np.abs(_Y.max() - _Y.min()) / 15.0 _alpha = 0.5 if with_scatter else 1.0 if ax is None: fig, ax = plt.subplots(figsize=(8, 5)) if fill: ax.contourf( H, extent=(_X.min() - offx, _X.max() + offx, _Y.min() - offy, _Y.max() + offy), color=c, alpha=_alpha, ) else: cset = ax.contour(H, extent=(_X.min() - offx, _X.max() + offx, _Y.min() - offy, _Y.max() + offy), color=c, **contour_kwargs) ax.clabel(cset, inline=1, fontsize=10) if with_scatter: ax.scatter(_X, _Y, c=c, alpha=_alpha) return ax
def draw_3d_plot(ax: mpl.axes.Axes, x: np.ndarray, y: np.ndarray, z: np.ndarray, plot_type: str, marker: str = 'X', marker_size: int = 50, marker_color: str = 'red', interpolation: str = 'linear', cmap: str = 'viridis') -> None: '''Draw a 3d plot. See XYZData class for explanation of arguments >>> points = np.random.rand(1000, 2) >>> x = np.random.rand(10) >>> y = np.random.rand(10) >>> z = x ** 2 + y ** 2 >>> if has_display(): ... fig, ax = plt.subplots() ... draw_3d_plot(ax, x = x, y = y, z = z, plot_type = 'contour', interpolation = 'linear') ''' xi = np.linspace(min(x), max(x)) yi = np.linspace(min(y), max(y)) X, Y = np.meshgrid(xi, yi) Z = griddata((x, y), z, (xi[None, :], yi[:, None]), method=interpolation) Z = np.nan_to_num(Z) if plot_type == 'surface': ax.plot_surface(X, Y, Z, cmap=cmap) if marker is not None: ax.scatter(x, y, z, marker=marker, s=marker_size, c=marker_color) elif plot_type == 'contour': cs = ax.contour(X, Y, Z, linewidths=0.5, colors='k') ax.clabel(cs, cs.levels[::2], fmt="%.3g", inline=1) ax.contourf(X, Y, Z, cmap=cmap) if marker is not None: ax.scatter(x, y, marker=marker, s=marker_size, c=marker_color, zorder=10) else: raise Exception(f'unknown plot type: {plot_type}') m = cm.ScalarMappable(cmap=cmap) m.set_array(Z) plt.colorbar(m, ax=ax)