-
Notifications
You must be signed in to change notification settings - Fork 16
/
plot_points.py
146 lines (110 loc) · 3.96 KB
/
plot_points.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
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
import numpy as np
try:
import matplotlib.pyplot as plt
ENABLED_2D_PLOTTING = True
except ModuleNotFoundError:
ENABLED_2D_PLOTTING = False
try:
from mayavi import mlab
from mayavi.core.ui.api import MlabSceneModel
ENABLED_3D_PLOTTING = True
except ModuleNotFoundError:
ENABLED_3D_PLOTTING = True
__all__ = ['plot_points', 'plot_points_2d', 'plot_points_3d']
def plot_points_2d(coordinates, axis=None, colour='red', size=15, title=None):
"""
Utility function to plot 2D scattered points using matplotlib.
Parameters
----------
coordinates : 2-dimensional array
Coordinates of the points to be plotted, shape should be (n_points, 2).
axis : matplotlib axis, optional
Axis in which to make the plotting, defaults to new empty one.
colour : str
Colour to apply to the points, defaults to red.
size : float
Size of the plotted points, defaults to 15.
title : str, optional
Figure title, defaults to empty title.
Returns
-------
matplotlib axis
Generated matplotlib axis
"""
if not ENABLED_2D_PLOTTING:
return None
if axis is None:
figure, axis = plt.subplots(1, 1)
if len(coordinates.shape) == 1:
coordinates = coordinates.reshape((1, coordinates.shape[0]))
space_scale = 1e-3
im = axis.scatter(coordinates[:, 0]/space_scale, coordinates[:, 1]/space_scale,
s=size, c=colour)
if title is not None:
axis.set_title(title)
return axis
def plot_points_3d(coordinates, axis=None, colour='red', size=15, title=None):
"""
Utility function to plot 3D scattered points using MayaVi.
Parameters
----------
coordinates : 2-dimensional array
Coordinates of the points to be plotted, shape should be (n_points, 3).
axis : MayaVi axis, optional
Axis in which to make the plotting, defaults to new empty one.
colour : str
Colour to apply to the points, defaults to red.
size : float
Size of the plotted points, defaults to 15.
title : str, optional
Figure title, defaults to empty title.
Returns
-------
MayaVi figure
Generated MayaVi figure
"""
if not ENABLED_3D_PLOTTING:
return None
if axis is None:
axis = MlabSceneModel()
colour_map = {
'red': (1., 0., 0.),
'green': (0., 1., 0.),
'blue': (0., 0., 1.),
}
scale_factor = 100 * size / np.max(coordinates)
transducers = mlab.pipeline.scalar_scatter(coordinates[:, 0],
coordinates[:, 1],
coordinates[:, 2],
figure=axis.scene3d.mayavi_scene)
mlab.pipeline.glyph(transducers,
mode='sphere', color=colour_map[colour], scale_factor=scale_factor,
figure=axis.scene3d.mayavi_scene)
return axis
def plot_points(coordinates, axis=None, colour='red', size=15, title=None):
"""
Utility function to plot scattered points using matplotlib (2D) or MayaVi (3D).
Parameters
----------
coordinates : 2-dimensional array
Coordinates of the points to be plotted, shape should be (n_points, dimensions).
axis : axis, optional
Axis in which to make the plotting, defaults to new empty one.
colour : str
Colour to apply to the points, defaults to red.
size : float
Size of the plotted points, defaults to 15.
title : str, optional
Figure title, defaults to empty title.
Returns
-------
matplotlib or MayaVi axis
Generated axis
"""
if coordinates.shape[-1] > 2:
axis = plot_points_3d(coordinates,
axis=axis, colour=colour, size=size, title=title)
else:
axis = plot_points_2d(coordinates,
axis=axis, colour=colour, size=size, title=title)
return axis