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
