-
Notifications
You must be signed in to change notification settings - Fork 0
/
plotting.py
95 lines (82 loc) · 3.75 KB
/
plotting.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
from typing import Optional, List, Tuple, Dict
import matplotlib
matplotlib.use("TkAgg")
from matplotlib.contour import QuadContourSet, ClabelText
from matplotlib.colorbar import Colorbar
from matplotlib.image import AxesImage
import matplotlib.pyplot as plt
import matplotlib.mlab as mlab
import warnings
def caldera_sim_function(x, y, map_size=(100,100)):
warnings.filterwarnings("ignore", category=matplotlib.cbook.MatplotlibDeprecationWarning)
x, y = x / (map_size[0] / 10.), y / (map_size[1] / 10.)
z0 = mlab.bivariate_normal(x, y, 10.0, 5.0, 5.0, 0.0)
z1 = mlab.bivariate_normal(x, y, 1.0, 2.0, 2.0, 5.0)
z2 = mlab.bivariate_normal(x, y, 1.7, 1.7, 8.0, 8.0)
return 50000.0 * z0 + 2500.0 * z1 + 5000.0 * z2
def draw_caldera_maxima(axes: plt.Axes, map_size=(100,100)):
maxima = list()
maxima.extend(axes.plot(20 * (map_size[0] / 100.), 46 * (map_size[1] / 100.), 'bx'))
maxima.extend(axes.plot(79 * (map_size[0] / 100.), 79 * (map_size[1] / 100.), 'bx'))
return maxima
class RobotPlot:
def __init__(self, ax, title):
self.ax: plt.Axes = ax
self.pos: Dict[int, Optional[Tuple[int, int]]] = dict()
self.robot_marker: Dict[int, List[plt.Line2D]] = dict()
self.path: Dict[int, List[List[plt.Line2D]]] = dict()
if title is not None:
self.ax.title.set_text(title)
def draw_robot(self, new_pos, color='m', connect=True, index=0):
for marker in self.robot_marker.get(index, list()):
marker.remove()
changed = list()
if connect and self.pos.get(index, None) is not None:
if index not in self.path:
self.path[index] = list()
self.path[index].append(self.ax.plot([self.pos[index][0], new_pos[0]],
[self.pos[index][1], new_pos[1]], color='k'))
changed.extend(self.path[index][-1])
self.robot_marker[index] = self.ax.plot(*new_pos, '*' + color)
self.pos[index] = new_pos
changed.extend(self.robot_marker[index])
return changed
class ContourPlot(RobotPlot):
def __init__(self, ax, title):
super().__init__(ax, title)
self.contours: Optional[QuadContourSet] = None
self.contour_labels: List[ClabelText] = list()
self.cbar: Optional[Colorbar] = None
def draw_contours(self, X, Y, Z, label=True, colorbar=False, **contour_kw):
if self.contours is not None:
for coll in self.contours.collections:
coll.remove()
for label in self.contour_labels:
label.remove()
if self.cbar is not None:
self.cbar.remove()
self.cbar = None
self.contours = self.ax.contour(X, Y, Z, **contour_kw)
changed = [self.contours]
if label:
self.contour_labels = self.ax.clabel(self.contours, inline=1, fontsize=8, fmt='%.3g')
changed.append(self.contour_labels)
if colorbar and len(self.contour_labels) > 0:
self.cbar = plt.gcf().colorbar(self.contours, ax=self.ax, fraction=0.046, pad=0.04)
changed.append(self.cbar)
return changed
class HeatmapPlot(RobotPlot):
def __init__(self, ax, title):
super().__init__(ax, title)
self.im: Optional[AxesImage] = None
self.cbar: Optional[Colorbar] = None
def draw_heatmap(self, map, colorbar=True, **heatmap_kw):
if self.cbar is not None:
self.cbar.remove()
self.im = self.ax.imshow(map, interpolation='nearest', **heatmap_kw)
self.ax.invert_yaxis()
changed = [self.im]
if colorbar:
self.cbar = plt.gcf().colorbar(self.im, ax=self.ax, fraction=0.046, pad=0.04)
# changed.append(self.cbar)
return changed