def plot_ellipse(Sigma,
mu,
ax,
n_std=3.0,
facecolor='none',
edgecolor='k',
plot_center='true',
**kwargs):
cov = Sigma
pearson = cov[0, 1] / np.sqrt(cov[0, 0] * cov[1, 1])
ell_radius_x = np.sqrt(1 + pearson)
ell_radius_y = np.sqrt(1 - pearson)
ellipse = Ellipse((0, 0),
width=ell_radius_x * 2,
height=ell_radius_y * 2,
facecolor=facecolor,
edgecolor=edgecolor,
**kwargs)
scale_x = np.sqrt(cov[0, 0]) * n_std
mean_x = mu[0]
scale_y = np.sqrt(cov[1, 1]) * n_std
mean_y = mu[1]
transf = (transforms.Affine2D().rotate_deg(45).scale(scale_x,
scale_y).translate(
mean_x, mean_y))
ellipse.set_transform(transf + ax.transData)
if plot_center:
ax.plot(mean_x, mean_y, '.')
return ax.add_patch(ellipse)
def transform_arrow(arrow, x, y, dx, dy):
# takes pyplot Arrow object and new coordinates for x, y, dx, and dy
# transforms the Arrow object
L = np.hypot(dx, dy)
if L != 0:
cx = dx / L
sx = dy / L
else:
# Account for division by zero
cx, sx = 0, 1
trans1 = transforms.Affine2D().scale(L, arrow.get_linewidth())
trans2 = transforms.Affine2D.from_values(cx, sx, -sx, cx, 0.0, 0.0)
trans3 = transforms.Affine2D().translate(x, y)
trans = trans1 + trans2 + trans3
arrow._patch_transform = trans.frozen()
def __init__(self, shape, xy, radius, **kwargs):
self.shape = shape
self.xy = xy
self.radius = radius
if shape == 'o': # circle
self.numVertices = None
self.orientation = 0
elif shape == '^': # triangle up
self.numVertices = 3
self.orientation = 0
elif shape == '<': # triangle left
self.numVertices = 3
self.orientation = np.pi * 0.5
elif shape == 'v': # triangle down
self.numVertices = 3
self.orientation = np.pi
elif shape == '>': # triangle right
self.numVertices = 3
self.orientation = np.pi * 1.5
elif shape == 's': # square
self.numVertices = 4
self.orientation = np.pi * 0.25
elif shape == 'd': # diamond
self.numVertices = 4
self.orientation = np.pi * 0.5
elif shape == 'p': # pentagon
self.numVertices = 5
self.orientation = 0
elif shape == 'h': # hexagon
self.numVertices = 6
self.orientation = 0
elif shape == '8': # octagon
self.numVertices = 8
self.orientation = 0
else:
raise ValueError(
"Node shape should be one of: 'so^>v<dph8'. Current shape:{}".
format(shape))
if self.shape == 'o': # circle
self.linewidth_correction = 2
self._path = Path.circle()
else: # regular polygon
self.linewidth_correction = 2 * np.sin(
np.pi / self.numVertices
) # derives from the ratio between a side and the radius in a regular polygon.
self._path = Path.unit_regular_polygon(self.numVertices)
self._patch_transform = transforms.Affine2D()
super().__init__(path=self._path, **kwargs)
def confidence_ellipse(x, y, ax, n_std=2.5, facecolor='none', **kwargs):
if x.size != y.size:
raise ValueError("x and y must be the same size")
cov = np.cov(x, y)
pearson = cov[0, 1]/np.sqrt(cov[0, 0] * cov[1, 1])
ell_radius_x = np.sqrt(1 + pearson)
ell_radius_y = np.sqrt(1 - pearson)
ellipse = Ellipse((0, 0), width=ell_radius_x * 2, height=ell_radius_y * 2,
facecolor=facecolor, **kwargs)
scale_x = np.sqrt(cov[0, 0]) * n_std
mean_x = np.mean(x)
scale_y = np.sqrt(cov[1, 1]) * n_std
mean_y = np.mean(y)
transf = transforms.Affine2D().rotate_deg(45).scale(scale_x, scale_y).translate(mean_x, mean_y)
ellipse.set_transform(transf + ax.transData)
return ax.add_patch(ellipse)