def countour_plot_for_highlighting(
ax: matplotlib.axes.Axes,
X: np.ndarray,
Y: np.ndarray,
hist_array: np.ndarray,
colormap: matplotlib.colors.ListedColormap,
level_step: float = 0.002) -> matplotlib.axes.Axes.contour:
""" Plot a contour for the region highlighting plot.
Included as an option because ROOT more or less includes one, but it doesn't really
have an impact on the output, so it can probably be ignored.
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.
level_step: Step in z between each contour
Returns:
Value returned by the contour plot.
"""
contour = ax.contour(X,
Y,
hist_array.T,
zdir='z',
alpha=1,
cmap=colormap,
levels=np.arange(np.min(hist_array),
np.max(hist_array), level_step))
return contour
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)
def plot_2d_gauss_fit(axes: mpl.axes.Axes, popt_2d, x_centers: np.ndarray,
y_centers: np.ndarray):
""" plot the contours of the Gaussian fit on the hitmap """
xx, yy = np.meshgrid(x_centers, y_centers)
axes.contour(xx, yy, ft.gauss_2d(*popt_2d)(xx, yy), cmap='coolwarm')
axes.scatter(popt_2d[0], popt_2d[1], color='darkred')
