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)
