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)
def surface_plot_for_highlighting( ax: matplotlib.axes.Axes, X: np.ndarray, Y: np.ndarray, hist_array: np.ndarray, colormap: matplotlib.colors.ListedColormap, colorbar: bool = False) -> mpl_toolkits.mplot3d.art3d.Poly3DCollection: """ Plot a surface with the passed arguments for use with highlighting a region. This function could be generally useful for plotting surfaces. However, it should be noted that the arguments are specifically optimized for plotting a highlighted region. Args: X: X bin centers. Y: Y bin centers. hist_array: Histogram data as 2D array. colormap: Colormap used to map the data to colors. colorbar: True if a colorbar should be added. Returns: Value returned by the surface plot. """ # NOTE: {r,c}count tells the surf plot how many times to sample the data. By using len(), # we ensure that every point is plotted. # NOTE: Cannot just use norm and cmap as args to plot_surface, as this seems to return colors only based # on value, instead of based on (x, y). To work around this, we calculate the norm separately, # and then use that to set the facecolors manually, which appears to set the colors based on # (x, y) position. Inspired by https://stackoverflow.com/a/42927880 norm = matplotlib.colors.Normalize(vmin=np.min(hist_array), vmax=np.max(hist_array)) surf = ax.plot_surface( X, Y, hist_array.T, facecolors=colormap(norm(hist_array.T)), #norm = matplotlib.colors.Normalize(vmin = np.min(hist_array), vmax = np.max(hist_array)), #cmap = sns.cm.rocket, rcount=len(hist_array.T[:, 0]), ccount=len(hist_array.T[0])) if colorbar: # NOTE: Cannot use surf directly, because it doesn't contain the mapping from data to color # Instead, we basically handle it by hand. m = matplotlib.cm.ScalarMappable(cmap=colormap, norm=surf.norm) m.set_array(hist_array.T) ax.colorbar(m) # Need to manually update the 2d colors based on the 3d colors to be able to use `get_facecolors()`. # This workaround is from: https://github.com/matplotlib/matplotlib/issues/4067 surf._facecolors2d = surf._facecolors3d surf._edgecolors2d = surf._edgecolors3d return surf
def draw_3d_plot(ax: mpl.axes.Axes, x: np.ndarray, y: np.ndarray, z: np.ndarray, plot_type: str = 'contour', marker: str = 'X', marker_size: int = 50, marker_color: str = 'red', interpolation: str = 'linear', cmap: matplotlib.colors.Colormap = matplotlib.cm.RdBu_r, min_level: float = math.nan, max_level: float = math.nan) -> 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) m = cm.ScalarMappable(cmap=cmap) m.set_array(Z) plt.colorbar(m, ax=ax) elif plot_type == 'contour': # extract all colors from the map cmaplist = [cmap(i) for i in range(cmap.N)] # create the new map cmap = cmap.from_list('Custom cmap', cmaplist, cmap.N) Z = np.ma.masked_array(Z, mask=~np.isfinite(Z)) if math.isnan(min_level): min_level = np.min(Z) if math.isnan(max_level): max_level = np.max(Z) # define the bins and normalize and forcing 0 to be part of the colorbar! bounds = np.arange(min_level, max_level, (max_level - min_level) / cmap.N) idx = np.searchsorted(bounds, 0) bounds = np.insert(bounds, idx, 0) norm = BoundaryNorm(bounds, cmap.N) cs = ax.contourf(X, Y, Z, cmap=cmap, norm=norm) if marker is not None: x = x[np.isfinite(z)] y = y[np.isfinite(z)] ax.scatter(x, y, marker=marker, s=marker_size, c=z[np.isfinite(z)], zorder=10, cmap=cmap) LABEL_SIZE = 16 ax.tick_params(axis='both', which='major', labelsize=LABEL_SIZE) ax.tick_params(axis='both', which='minor', labelsize=LABEL_SIZE) cbar = plt.colorbar(cs, ax=ax) cbar.ax.tick_params(labelsize=LABEL_SIZE) else: raise Exception(f'unknown plot type: {plot_type}')
def plot_ellipsoid( fig: matplotlib.figure.Figure, ax: matplotlib.axes.Axes, covariance_matrix: np.array, mean: np.array, weight: np.array, color: any, print_properties: bool = False, additional_rotation: np.array = np.identity(3), ) -> Tuple[matplotlib.figure.Figure, matplotlib.axes.Axes]: """ A function for drawing 3d-multivariate guassians with method initially from: https://stackoverflow.com/questions/7819498/plotting-ellipsoid-with-matplotlib :param fig: The figure matplotlib.pyplot object :param ax: The axis matplotlib.pylot object with Axes3D extension :param covariance_matrix: A covariance matrix input from the multivariate guassian to be plotted. :param mean: ditto :param weight: ditto :param color: ditto :return: fig and ax so that they can be used by further plotting steps. """ # I arbitrarily choose some levels to in the multivariate Gaussian to plot. if print_properties: print("weight", weight) print("mean", mean) print("covariance matrix", covariance_matrix) for sigma, alpha in [ [3, 0.2 * weight], [2, 0.4 * weight], [1, 0.6 * weight], ]: # find the rotation matrix and radii of the axes _, s, rotation = la.svd(covariance_matrix) # Singular Value Decomposition from numpy.linalg # finds the variance vector s when the covariance # matrix has been rotated so that it is diagonal radii = np.sqrt(s) * sigma # s is the sigma*2 in each of the principal axes directions # now carry on with EOL's answer u = np.linspace(0.0, 2.0 * np.pi, 100) # AZIMUTHAL ANGLE (LONGITUDE) v = np.linspace(0.0, np.pi, 100) # POLAR ANGLE (LATITUDE) # COORDINATES OF THE SURFACE PRETENDING THAT THE # GAUSSIAN IS AT THE CENTRE & NON ROTATED x = radii[0] * np.outer(np.cos(u), np.sin(v)) # MESH FOR X y = radii[1] * np.outer(np.sin(u), np.sin(v)) # MESH FOR Y z = radii[2] * np.outer(np.ones_like(u), np.cos(v)) # MESH FOR Z # move so that the gaussian is actually rotated and on the right point. for i in range(len(x)): for j in range(len(x)): [x[i, j], y[i, j], z[i, j]] = (np.dot([x[i, j], y[i, j], z[i, j]], rotation) + mean) [x[i, j], y[i, j], z[i, j]] = np.dot(additional_rotation, [x[i, j], y[i, j], z[i, j]]) # plot the surface in a reasonable partially translucent way ax.plot_surface(x, y, z, rstride=4, cstride=4, color=color, alpha=alpha) return fig, ax